Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2021-06-04 06:14:43

0001 import FWCore.ParameterSet.Config as cms
0002 
0003 from PhysicsTools.NanoAOD.nano_eras_cff import *
0004 
0005 from CommonTools.PileupAlgos.Puppi_cff import puppi
0006 
0007 from RecoJets.JetProducers.PileupJetID_cfi import pileupJetIdCalculator, pileupJetId
0008 from RecoJets.JetProducers.PileupJetID_cfi import _chsalgos_81x, _chsalgos_94x, _chsalgos_102x
0009 
0010 from PhysicsTools.NanoAOD.common_cff import Var, P4Vars
0011 from PhysicsTools.NanoAOD.jets_cff   import jetTable, jetCorrFactorsNano, updatedJets, finalJets, qgtagger, hfJetShowerShapeforNanoAOD
0012 from PhysicsTools.NanoAOD.jets_cff   import genJetTable, genJetFlavourAssociation, genJetFlavourTable
0013 
0014 from PhysicsTools.PatAlgos.tools.jetCollectionTools import GenJetAdder, RecoJetAdder
0015 from PhysicsTools.PatAlgos.tools.jetTools import supportedJetAlgos
0016 from PhysicsTools.PatAlgos.tools.jetTools import updateJetCollection
0017 
0018 import copy
0019 
0020 bTagCSVV2    = ['pfCombinedInclusiveSecondaryVertexV2BJetTags']
0021 bTagDeepCSV  = ['pfDeepCSVJetTags:probb','pfDeepCSVJetTags:probbb','pfDeepCSVJetTags:probc']
0022 bTagDeepJet  = [
0023   'pfDeepFlavourJetTags:probb','pfDeepFlavourJetTags:probbb','pfDeepFlavourJetTags:problepb',
0024   'pfDeepFlavourJetTags:probc','pfDeepFlavourJetTags:probuds','pfDeepFlavourJetTags:probg'
0025 ]
0026 from RecoBTag.ONNXRuntime.pfParticleNetAK4_cff import _pfParticleNetAK4JetTagsAll
0027 bTagDiscriminatorsForAK4 = bTagCSVV2+bTagDeepCSV+bTagDeepJet+_pfParticleNetAK4JetTagsAll
0028 
0029 from RecoBTag.ONNXRuntime.pfDeepBoostedJet_cff import _pfDeepBoostedJetTagsAll
0030 from RecoBTag.ONNXRuntime.pfParticleNet_cff import _pfParticleNetJetTagsAll
0031 
0032 btagHbb = ['pfBoostedDoubleSecondaryVertexAK8BJetTags']
0033 btagDDX = [
0034   'pfDeepDoubleBvLJetTags:probHbb',
0035   'pfDeepDoubleCvLJetTags:probHcc',
0036   'pfDeepDoubleCvBJetTags:probHcc',
0037   'pfMassIndependentDeepDoubleBvLJetTags:probHbb',
0038   'pfMassIndependentDeepDoubleCvLJetTags:probHcc',
0039   'pfMassIndependentDeepDoubleCvBJetTags:probHcc'
0040 ]
0041 btagDDXV2 = [
0042   'pfMassIndependentDeepDoubleBvLV2JetTags:probHbb',
0043   'pfMassIndependentDeepDoubleCvLV2JetTags:probHcc',
0044   'pfMassIndependentDeepDoubleCvBV2JetTags:probHcc'
0045 ]
0046 
0047 #
0048 # By default, these collections are saved in NanoAODs:
0049 # - ak4gen (GenJet in NanoAOD), slimmedGenJets in MiniAOD
0050 # - ak8gen (GenJetAK8 in NanoAOD), slimmedGenJetsAK8 in MiniAOD
0051 # Below is a list of genjets that we can save in NanoAOD. Set
0052 # "enabled" to true if you want to store the jet collection
0053 config_genjets = [
0054   {
0055     "jet"     : "ak6gen",
0056     "enabled" : False,
0057   },
0058 ]
0059 config_genjets = list(filter(lambda k: k['enabled'], config_genjets))
0060 #
0061 # GenJets info in NanoAOD
0062 #
0063 nanoInfo_genjets = {
0064   "ak6gen"  : {
0065     "name" : "GenJetAK6",
0066     "doc"  : "AK6 Gen jets (made with visible genparticles) with pt > 3 GeV", # default genjets pt cut after clustering is 3 GeV
0067   },
0068 }
0069 #
0070 # By default, these collections are saved in the main NanoAODs:
0071 # - ak4pfchs   (Jet    in NanoAOD), slimmedJets in MiniAOD
0072 # - ak8pfpuppi (FatJet in NanoAOD), slimmedJetsAK8 in MiniAOD
0073 # Below is a list of recojets that we can save in NanoAOD. Set
0074 # "enabled" to true if you want to store the recojet collection.
0075 #
0076 config_recojets = [
0077   {
0078     "jet" : "ak4calo",
0079     "enabled" : True,
0080     "inputCollection"  : "slimmedCaloJets", #Exist in MiniAOD
0081     "genJetsCollection": "AK4GenJetsNoNu",
0082   },
0083   {
0084     "jet" : "ak4pf",
0085     "enabled" : False,
0086     "inputCollection" : "",
0087     "genJetsCollection": "AK4GenJetsNoNu",
0088     "minPtFastjet" : 0.,
0089   },
0090   {
0091     "jet" : "ak4pfpuppi",
0092     "enabled" : True,
0093     "inputCollection" : "",
0094     "genJetsCollection": "AK4GenJetsNoNu",
0095     "bTagDiscriminators": bTagDiscriminatorsForAK4,
0096     "minPtFastjet" : 0.,
0097   },
0098   {
0099     "jet" : "ak8pf",
0100     "enabled" : False,
0101     "inputCollection" : "",
0102     "genJetsCollection": "AK8GenJetsNoNu",
0103     "minPtFastjet" : 0.,
0104   },
0105 ]
0106 config_recojets = list(filter(lambda k: k['enabled'], config_recojets))
0107 #
0108 # RecoJets info in NanoAOD
0109 #
0110 nanoInfo_recojets = {
0111   "ak4calo" : {
0112     "name": "JetCalo",
0113     "doc" : "AK4 Calo jets with JECs applied",
0114   },
0115   "ak4pf" : {
0116     "name"  : "JetPF",
0117     "doc"   : "AK4 PF jets",
0118     "ptcut" : "",
0119   },
0120   "ak4pfpuppi" : {
0121     "name"  : "JetPuppi",
0122     "doc"   : "AK4 PF Puppi",
0123     "ptcut" : "",
0124     "doQGL" : True,
0125     "doPUIDVar": True,
0126     "doBTag": True,
0127   },
0128   "ak8pf" : {
0129     "name"  : "FatJetPF",
0130     "doc"   : "AK8 PF jets",
0131     "ptcut" : "",
0132   },
0133 }
0134 
0135 GENJETVARS = cms.PSet(P4Vars,
0136   nConstituents   = jetTable.variables.nConstituents,
0137 )
0138 PFJETVARS = cms.PSet(P4Vars,
0139   rawFactor       = Var("1.-jecFactor('Uncorrected')",float,doc="1 - Factor to get back to raw pT",precision=6),
0140   area            = jetTable.variables.area,
0141   chHEF           = jetTable.variables.chHEF,
0142   neHEF           = jetTable.variables.neHEF,
0143   chEmEF          = jetTable.variables.chEmEF,
0144   neEmEF          = jetTable.variables.neEmEF,
0145   muEF            = jetTable.variables.muEF,
0146   hfHEF           = Var("HFHadronEnergyFraction()",float,doc = "hadronic energy fraction in HF",precision = 6),
0147   hfEmEF          = Var("HFEMEnergyFraction()",float,doc = "electromagnetic energy fraction in HF",precision = 6),
0148   nMuons          = jetTable.variables.nMuons,
0149   nElectrons      = jetTable.variables.nElectrons,
0150   nConstituents   = jetTable.variables.nConstituents,
0151   nConstChHads    = Var("chargedHadronMultiplicity()",int,doc="number of charged hadrons in the jet"),
0152   nConstNeuHads   = Var("neutralHadronMultiplicity()",int,doc="number of neutral hadrons in the jet"),
0153   nConstHFHads    = Var("HFHadronMultiplicity()", int,doc="number of HF hadrons in the jet"),
0154   nConstHFEMs     = Var("HFEMMultiplicity()",int,doc="number of HF EMs in the jet"),
0155   nConstMuons     = Var("muonMultiplicity()",int,doc="number of muons in the jet"),
0156   nConstElecs     = Var("electronMultiplicity()",int,doc="number of electrons in the jet"),
0157   nConstPhotons   = Var("photonMultiplicity()",int,doc="number of photons in the jet"),
0158 )
0159 PUIDVARS = cms.PSet(
0160   puId_dR2Mean    = Var("?(pt>10)?userFloat('puId_dR2Mean'):-1",float,doc="pT^2-weighted average square distance of jet constituents from the jet axis (PileUp ID BDT input variable)", precision= 6),
0161   puId_majW       = Var("?(pt>10)?userFloat('puId_majW'):-1",float,doc="major axis of jet ellipsoid in eta-phi plane (PileUp ID BDT input variable)", precision= 6),
0162   puId_minW       = Var("?(pt>10)?userFloat('puId_minW'):-1",float,doc="minor axis of jet ellipsoid in eta-phi plane (PileUp ID BDT input variable)", precision= 6),
0163   puId_frac01     = Var("?(pt>10)?userFloat('puId_frac01'):-1",float,doc="fraction of constituents' pT contained within dR <0.1 (PileUp ID BDT input variable)", precision= 6),
0164   puId_frac02     = Var("?(pt>10)?userFloat('puId_frac02'):-1",float,doc="fraction of constituents' pT contained within 0.1< dR <0.2 (PileUp ID BDT input variable)", precision= 6),
0165   puId_frac03     = Var("?(pt>10)?userFloat('puId_frac03'):-1",float,doc="fraction of constituents' pT contained within 0.2< dR <0.3 (PileUp ID BDT input variable)", precision= 6),
0166   puId_frac04     = Var("?(pt>10)?userFloat('puId_frac04'):-1",float,doc="fraction of constituents' pT contained within 0.3< dR <0.4 (PileUp ID BDT input variable)", precision= 6),
0167   puId_ptD        = Var("?(pt>10)?userFloat('puId_ptD'):-1",float,doc="pT-weighted average pT of constituents (PileUp ID BDT input variable)", precision= 6),
0168   puId_beta       = Var("?(pt>10)?userFloat('puId_beta'):-1",float,doc="fraction of pT of charged constituents associated to PV (PileUp ID BDT input variable)", precision= 6),
0169   puId_pull       = Var("?(pt>10)?userFloat('puId_pull'):-1",float,doc="magnitude of pull vector (PileUp ID BDT input variable)", precision= 6),
0170   puId_jetR       = Var("?(pt>10)?userFloat('puId_jetR'):-1",float,doc="fraction of jet pT carried by the leading constituent (PileUp ID BDT input variable)", precision= 6),
0171   puId_jetRchg    = Var("?(pt>10)?userFloat('puId_jetRchg'):-1",float,doc="fraction of jet pT carried by the leading charged constituent (PileUp ID BDT input variable)", precision= 6),
0172   puId_nCharged   = Var("?(pt>10)?userInt('puId_nCharged'):-1",int,doc="number of charged constituents (PileUp ID BDT input variable)"),
0173 )
0174 QGLVARS = cms.PSet(
0175   qgl_axis2       =  Var("?(pt>10)?userFloat('qgl_axis2'):-1",float,doc="ellipse minor jet axis (Quark vs Gluon likelihood input variable)", precision= 6),
0176   qgl_ptD         =  Var("?(pt>10)?userFloat('qgl_ptD'):-1",float,doc="pT-weighted average pT of constituents (Quark vs Gluon likelihood input variable)", precision= 6),
0177   qgl_mult        =  Var("?(pt>10)?userInt('qgl_mult'):-1", int,doc="PF candidates multiplicity (Quark vs Gluon likelihood input variable)"),
0178 )
0179 BTAGVARS = cms.PSet(
0180   btagDeepB = Var("?(pt>15)&&((bDiscriminator('pfDeepCSVJetTags:probb')+bDiscriminator('pfDeepCSVJetTags:probbb'))>=0)?bDiscriminator('pfDeepCSVJetTags:probb')+bDiscriminator('pfDeepCSVJetTags:probbb'):-1",float,doc="DeepCSV b+bb tag discriminator",precision=10),
0181   btagCSVV2 = Var("?pt>15?bDiscriminator('pfCombinedInclusiveSecondaryVertexV2BJetTags'):-1",float,doc=" pfCombinedInclusiveSecondaryVertexV2 b-tag discriminator (aka CSVV2)",precision=10),
0182   btagDeepCvL = Var("?(pt>15)&&(bDiscriminator('pfDeepCSVJetTags:probc')>=0)?bDiscriminator('pfDeepCSVJetTags:probc')/(bDiscriminator('pfDeepCSVJetTags:probc')+bDiscriminator('pfDeepCSVJetTags:probudsg')):-1", float,doc="DeepCSV c vs udsg discriminator",precision=10),
0183   btagDeepCvB = Var("?(pt>15)&&bDiscriminator('pfDeepCSVJetTags:probc')>=0?bDiscriminator('pfDeepCSVJetTags:probc')/(bDiscriminator('pfDeepCSVJetTags:probc')+bDiscriminator('pfDeepCSVJetTags:probb')+bDiscriminator('pfDeepCSVJetTags:probbb')):-1",float,doc="DeepCSV c vs b+bb discriminator",precision=10),
0184 )
0185 DEEPJETVARS = cms.PSet(
0186   btagDeepFlavB   = Var("?pt>15?bDiscriminator('pfDeepFlavourJetTags:probb')+bDiscriminator('pfDeepFlavourJetTags:probbb')+bDiscriminator('pfDeepFlavourJetTags:problepb'):-1",float,doc="DeepJet b+bb+lepb tag discriminator",precision=10),
0187   btagDeepFlavC   = Var("?pt>15?bDiscriminator('pfDeepFlavourJetTags:probc'):-1",float,doc="DeepFlavour charm tag raw score",precision=10),
0188   btagDeepFlavG   = Var("?pt>15?bDiscriminator('pfDeepFlavourJetTags:probg'):-1",float,doc="DeepFlavour gluon tag raw score",precision=10),
0189   btagDeepFlavUDS = Var("?pt>15?bDiscriminator('pfDeepFlavourJetTags:probuds'):-1",float,doc="DeepFlavour uds tag raw score",precision=10),
0190   btagDeepFlavCvL = Var("?(pt>15)&&(bDiscriminator('pfDeepFlavourJetTags:probc')+bDiscriminator('pfDeepFlavourJetTags:probuds')+bDiscriminator('pfDeepFlavourJetTags:probg'))>0?bDiscriminator('pfDeepFlavourJetTags:probc')/(bDiscriminator('pfDeepFlavourJetTags:probc')+bDiscriminator('pfDeepFlavourJetTags:probuds')+bDiscriminator('pfDeepFlavourJetTags:probg')):-1",float,doc="DeepJet c vs uds+g discriminator",precision=10),
0191   btagDeepFlavCvB = Var("?(pt>15)&&(bDiscriminator('pfDeepFlavourJetTags:probc')+bDiscriminator('pfDeepFlavourJetTags:probb')+bDiscriminator('pfDeepFlavourJetTags:probbb')+bDiscriminator('pfDeepFlavourJetTags:problepb'))>0?bDiscriminator('pfDeepFlavourJetTags:probc')/(bDiscriminator('pfDeepFlavourJetTags:probc')+bDiscriminator('pfDeepFlavourJetTags:probb')+bDiscriminator('pfDeepFlavourJetTags:probbb')+bDiscriminator('pfDeepFlavourJetTags:problepb')):-1",float,doc="DeepJet c vs b+bb+lepb discriminator",precision=10),
0192   btagDeepFlavQG  = Var("?(pt>15)&&(bDiscriminator('pfDeepFlavourJetTags:probg')+bDiscriminator('pfDeepFlavourJetTags:probuds'))>0?bDiscriminator('pfDeepFlavourJetTags:probg')/(bDiscriminator('pfDeepFlavourJetTags:probg')+bDiscriminator('pfDeepFlavourJetTags:probuds')):-1",float,doc="DeepJet g vs uds discriminator",precision=10),
0193 )
0194 PARTICLENETAK4VARS = cms.PSet(
0195   particleNetAK4_B = Var("?pt>15?bDiscriminator('pfParticleNetAK4DiscriminatorsJetTags:BvsAll'):-1",float,doc="ParticleNetAK4 tagger b vs all (udsg, c) discriminator",precision=10),
0196   particleNetAK4_CvsL = Var("?pt>15?bDiscriminator('pfParticleNetAK4DiscriminatorsJetTags:CvsL'):-1",float,doc="ParticleNetAK4 tagger c vs udsg discriminator",precision=10),
0197   particleNetAK4_CvsB = Var("?pt>15?bDiscriminator('pfParticleNetAK4DiscriminatorsJetTags:CvsB'):-1",float,doc="ParticleNetAK4 tagger c vs b discriminator",precision=10),
0198   particleNetAK4_QvsG = Var("?pt>15?bDiscriminator('pfParticleNetAK4DiscriminatorsJetTags:QvsG'):-1",float,doc="ParticleNetAK4 tagger uds vs g discriminator",precision=10),
0199   particleNetAK4_puIdDisc = Var("?pt>15?1-bDiscriminator('pfParticleNetAK4JetTags:probpu'):-1",float,doc="ParticleNetAK4 tagger pileup jet discriminator",precision=10),
0200 )
0201 
0202 CALOJETVARS = cms.PSet(P4Vars,
0203   area      = jetTable.variables.area,
0204   rawFactor = jetTable.variables.rawFactor,
0205   emf       = Var("emEnergyFraction()", float, doc = "electromagnetic energy fraction", precision = 10),
0206 )
0207 
0208 
0209 #******************************************
0210 #
0211 #
0212 # Reco Jets related functions
0213 #
0214 #
0215 #******************************************
0216 def AddJetID(proc, jetName="", jetSrc="", jetTableName="", jetSequenceName=""):
0217   """
0218   Setup modules to calculate PF jet ID
0219   """
0220 
0221   isPUPPIJet = True if "Puppi" in jetName else False
0222 
0223   looseJetId = "looseJetId{}".format(jetName)
0224   setattr(proc, looseJetId, proc.looseJetId.clone(
0225       src = jetSrc,
0226       filterParams = proc.looseJetId.filterParams.clone(
0227         version = "WINTER16"
0228       ),
0229     )
0230   )
0231 
0232   tightJetId = "tightJetId{}".format(jetName)
0233   setattr(proc, tightJetId, proc.tightJetId.clone(
0234       src = jetSrc,
0235       filterParams = proc.tightJetId.filterParams.clone(
0236         version = "RUN2UL{}".format("PUPPI" if isPUPPIJet else "CHS")
0237       ),
0238     )
0239   )
0240 
0241   tightJetIdLepVeto = "tightJetIdLepVeto{}".format(jetName)
0242   setattr(proc, tightJetIdLepVeto, proc.tightJetIdLepVeto.clone(
0243       src = jetSrc,
0244       filterParams = proc.tightJetIdLepVeto.filterParams.clone(
0245         version = "RUN2UL{}".format("PUPPI" if isPUPPIJet else "CHS")
0246       ),
0247     )
0248   )
0249 
0250   for modifier in run2_miniAOD_80XLegacy, run2_nanoAOD_94X2016:
0251     modifier.toModify(getattr(proc, tightJetId).filterParams, version = "WINTER16" )
0252     modifier.toModify(getattr(proc, tightJetIdLepVeto).filterParams, version = "WINTER16" )
0253   for modifier in run2_nanoAOD_94XMiniAODv1, run2_nanoAOD_94XMiniAODv2:
0254     modifier.toModify(getattr(proc, tightJetId).filterParams, version = "WINTER17{}".format("PUPPI" if isPUPPIJet else ""))
0255     modifier.toModify(getattr(proc, tightJetIdLepVeto).filterParams, version = "WINTER17{}".format("PUPPI" if isPUPPIJet else ""))
0256   run2_nanoAOD_102Xv1.toModify(getattr(proc, tightJetId).filterParams, version = "SUMMER18{}".format("PUPPI" if isPUPPIJet else "") )
0257   run2_nanoAOD_102Xv1.toModify(getattr(proc, tightJetIdLepVeto).filterParams, version = "SUMMER18{}".format("PUPPI" if isPUPPIJet else "") )
0258 
0259   #
0260   # Save variables as userInts in each jet
0261   #
0262   patJetWithUserData = "{}WithUserData".format(jetSrc)
0263   getattr(proc, patJetWithUserData).userInts.tightId = cms.InputTag(tightJetId)
0264   getattr(proc, patJetWithUserData).userInts.tightIdLepVeto = cms.InputTag(tightJetIdLepVeto)
0265   for modifier in run2_miniAOD_80XLegacy, run2_nanoAOD_94X2016:
0266     modifier.toModify(getattr(proc, patJetWithUserData).userInts, looseId = cms.InputTag(looseJetId))
0267 
0268   #
0269   # Specfiy variables in the jetTable to save in NanoAOD
0270   #
0271   getattr(proc, jetTableName).variables.jetId = Var("userInt('tightId')*2+4*userInt('tightIdLepVeto')",int,doc="Jet ID flags bit1 is loose (always false in 2017 since it does not exist), bit2 is tight, bit3 is tightLepVeto")
0272   for modifier in run2_miniAOD_80XLegacy, run2_nanoAOD_94X2016:
0273     modifier.toModify(getattr(proc, jetTableName).variables, jetId = Var("userInt('tightIdLepVeto')*4+userInt('tightId')*2+userInt('looseId')",int, doc="Jet ID flags bit1 is loose, bit2 is tight, bit3 is tightLepVeto"))
0274 
0275 
0276   getattr(proc,jetSequenceName).insert(getattr(proc,jetSequenceName).index(getattr(proc, jetSrc))+1, getattr(proc, tightJetId))
0277   getattr(proc,jetSequenceName).insert(getattr(proc,jetSequenceName).index(getattr(proc, tightJetId))+1, getattr(proc, tightJetIdLepVeto))
0278 
0279   setattr(proc,"_"+jetSequenceName+"_2016", getattr(proc,jetSequenceName).copy())
0280   getattr(proc,"_"+jetSequenceName+"_2016").insert(getattr(proc, "_"+jetSequenceName+"_2016").index(getattr(proc, tightJetId)), getattr(proc, looseJetId))
0281   for modifier in run2_miniAOD_80XLegacy, run2_nanoAOD_94X2016:
0282     modifier.toReplaceWith(getattr(proc,jetSequenceName), getattr(proc, "_"+jetSequenceName+"_2016"))
0283 
0284   return proc
0285 
0286 def AddPileUpJetIDVars(proc, jetName="", jetSrc="", jetTableName="", jetSequenceName=""):
0287   """
0288   Setup modules to calculate pileup jet ID input variables for PF jet
0289   """
0290 
0291   #
0292   # Calculate pileup jet ID variables
0293   #
0294   puJetIdVarsCalculator = "puJetIdCalculator{}".format(jetName)
0295   setattr(proc, puJetIdVarsCalculator, pileupJetIdCalculator.clone(
0296       jets = jetSrc,
0297       vertexes  = "offlineSlimmedPrimaryVertices",
0298       inputIsCorrected = True,
0299       applyJec  = False,
0300       usePuppi = True if "Puppi" in jetName else False
0301     )
0302   )
0303   getattr(proc,jetSequenceName).insert(getattr(proc,jetSequenceName).index(getattr(proc, jetSrc))+1, getattr(proc, puJetIdVarsCalculator))
0304 
0305   #
0306   # Get the variables
0307   #
0308   puJetIDVar = "puJetIDVar{}".format(jetName)
0309   setattr(proc, puJetIDVar, cms.EDProducer("PileupJetIDVarProducer",
0310       srcJet = cms.InputTag(jetSrc),
0311       srcPileupJetId = cms.InputTag(puJetIdVarsCalculator)
0312     )
0313   )
0314   getattr(proc,jetSequenceName).insert(getattr(proc,jetSequenceName).index(getattr(proc, puJetIdVarsCalculator))+1, getattr(proc, puJetIDVar))
0315 
0316   #
0317   # Save variables as userFloats and userInts for each jet
0318   #
0319   patJetWithUserData = "{}WithUserData".format(jetSrc)
0320   getattr(proc,patJetWithUserData).userFloats.puId_dR2Mean  = cms.InputTag("{}:dR2Mean".format(puJetIDVar))
0321   getattr(proc,patJetWithUserData).userFloats.puId_majW     = cms.InputTag("{}:majW".format(puJetIDVar))
0322   getattr(proc,patJetWithUserData).userFloats.puId_minW     = cms.InputTag("{}:minW".format(puJetIDVar))
0323   getattr(proc,patJetWithUserData).userFloats.puId_frac01   = cms.InputTag("{}:frac01".format(puJetIDVar))
0324   getattr(proc,patJetWithUserData).userFloats.puId_frac02   = cms.InputTag("{}:frac02".format(puJetIDVar))
0325   getattr(proc,patJetWithUserData).userFloats.puId_frac03   = cms.InputTag("{}:frac03".format(puJetIDVar))
0326   getattr(proc,patJetWithUserData).userFloats.puId_frac04   = cms.InputTag("{}:frac04".format(puJetIDVar))
0327   getattr(proc,patJetWithUserData).userFloats.puId_ptD      = cms.InputTag("{}:ptD".format(puJetIDVar))
0328   getattr(proc,patJetWithUserData).userFloats.puId_beta     = cms.InputTag("{}:beta".format(puJetIDVar))
0329   getattr(proc,patJetWithUserData).userFloats.puId_pull     = cms.InputTag("{}:pull".format(puJetIDVar))
0330   getattr(proc,patJetWithUserData).userFloats.puId_jetR     = cms.InputTag("{}:jetR".format(puJetIDVar))
0331   getattr(proc,patJetWithUserData).userFloats.puId_jetRchg  = cms.InputTag("{}:jetRchg".format(puJetIDVar))
0332   getattr(proc,patJetWithUserData).userInts.puId_nCharged   = cms.InputTag("{}:nCharged".format(puJetIDVar))
0333 
0334   #
0335   # Specfiy variables in the jet table to save in NanoAOD
0336   #
0337   getattr(proc,jetTableName).variables.puId_dR2Mean  = PUIDVARS.puId_dR2Mean
0338   getattr(proc,jetTableName).variables.puId_majW     = PUIDVARS.puId_majW
0339   getattr(proc,jetTableName).variables.puId_minW     = PUIDVARS.puId_minW
0340   getattr(proc,jetTableName).variables.puId_frac01   = PUIDVARS.puId_frac01
0341   getattr(proc,jetTableName).variables.puId_frac02   = PUIDVARS.puId_frac02
0342   getattr(proc,jetTableName).variables.puId_frac03   = PUIDVARS.puId_frac03
0343   getattr(proc,jetTableName).variables.puId_frac04   = PUIDVARS.puId_frac04
0344   getattr(proc,jetTableName).variables.puId_ptD      = PUIDVARS.puId_ptD
0345   getattr(proc,jetTableName).variables.puId_beta     = PUIDVARS.puId_beta
0346   getattr(proc,jetTableName).variables.puId_pull     = PUIDVARS.puId_pull
0347   getattr(proc,jetTableName).variables.puId_jetR     = PUIDVARS.puId_jetR
0348   getattr(proc,jetTableName).variables.puId_jetRchg  = PUIDVARS.puId_jetRchg
0349   getattr(proc,jetTableName).variables.puId_nCharged = PUIDVARS.puId_nCharged
0350 
0351   return proc
0352 
0353 def AddQGLTaggerVars(proc, jetName="", jetSrc="", jetTableName="", jetSequenceName="", calculateQGLVars=False):
0354   """
0355   Schedule the QGTagger module to calculate input variables to the QG likelihood
0356   """
0357 
0358   QGLTagger="qgtagger{}".format(jetName)
0359   patJetWithUserData="{}WithUserData".format(jetSrc)
0360 
0361   if calculateQGLVars:
0362     setattr(proc, QGLTagger, qgtagger.clone(
0363         srcJets = jetSrc
0364       )
0365     )
0366 
0367   #
0368   # Save variables as userFloats and userInts for each jet
0369   #
0370   getattr(proc,patJetWithUserData).userFloats.qgl_axis2 = cms.InputTag(QGLTagger+":axis2")
0371   getattr(proc,patJetWithUserData).userFloats.qgl_ptD   = cms.InputTag(QGLTagger+":ptD")
0372   getattr(proc,patJetWithUserData).userInts.qgl_mult    = cms.InputTag(QGLTagger+":mult")
0373 
0374   #
0375   # Specfiy variables in the jet table to save in NanoAOD
0376   #
0377   getattr(proc,jetTableName).variables.qgl_axis2 =  QGLVARS.qgl_axis2
0378   getattr(proc,jetTableName).variables.qgl_ptD   =  QGLVARS.qgl_ptD
0379   getattr(proc,jetTableName).variables.qgl_mult  =  QGLVARS.qgl_mult
0380 
0381   if calculateQGLVars:
0382     getattr(proc,jetSequenceName).insert(getattr(proc,jetSequenceName).index(getattr(proc, jetSrc))+1, getattr(proc, QGLTagger))
0383 
0384   return proc
0385 
0386 def AddBTaggingScores(proc, jetTableName=""):
0387   """
0388   Store b-tagging scores from various algortihm
0389   """
0390 
0391   getattr(proc, jetTableName).variables.btagDeepB       = BTAGVARS.btagDeepB
0392   getattr(proc, jetTableName).variables.btagCSVV2       = BTAGVARS.btagCSVV2
0393   getattr(proc, jetTableName).variables.btagDeepCvL     = BTAGVARS.btagDeepCvL
0394   getattr(proc, jetTableName).variables.btagDeepCvB     = BTAGVARS.btagDeepCvB
0395   getattr(proc, jetTableName).variables.btagDeepFlavB   = DEEPJETVARS.btagDeepFlavB
0396   getattr(proc, jetTableName).variables.btagDeepFlavCvL = DEEPJETVARS.btagDeepFlavCvL
0397   getattr(proc, jetTableName).variables.btagDeepFlavCvB = DEEPJETVARS.btagDeepFlavCvB
0398 
0399   return proc
0400 
0401 def AddDeepJetGluonLQuarkScores(proc, jetTableName=""):
0402   """
0403   Store DeepJet raw score in jetTable for gluon and light quark
0404   """
0405 
0406   getattr(proc, jetTableName).variables.btagDeepFlavG   = DEEPJETVARS.btagDeepFlavG
0407   getattr(proc, jetTableName).variables.btagDeepFlavUDS = DEEPJETVARS.btagDeepFlavUDS
0408   getattr(proc, jetTableName).variables.btagDeepFlavQG  = DEEPJETVARS.btagDeepFlavQG
0409 
0410   return proc
0411 
0412 def AddParticleNetAK4Scores(proc, jetTableName=""):
0413   """
0414   Store ParticleNetAK4 scores in jetTable
0415   """
0416 
0417   getattr(proc, jetTableName).variables.particleNetAK4_B = PARTICLENETAK4VARS.particleNetAK4_B
0418   getattr(proc, jetTableName).variables.particleNetAK4_CvsL = PARTICLENETAK4VARS.particleNetAK4_CvsL
0419   getattr(proc, jetTableName).variables.particleNetAK4_CvsB = PARTICLENETAK4VARS.particleNetAK4_CvsB
0420   getattr(proc, jetTableName).variables.particleNetAK4_QvsG = PARTICLENETAK4VARS.particleNetAK4_QvsG
0421   getattr(proc, jetTableName).variables.particleNetAK4_puIdDisc = PARTICLENETAK4VARS.particleNetAK4_puIdDisc
0422 
0423   return proc
0424 
0425 def AddNewPatJets(proc, recoJetInfo, runOnMC):
0426   """
0427   Add patJet into custom nanoAOD
0428   """
0429 
0430   jetName = recoJetInfo.jetUpper
0431   payload = recoJetInfo.jetCorrPayload
0432   doPF    = recoJetInfo.doPF
0433   doCalo  = recoJetInfo.doCalo
0434   patJetFinalColl = recoJetInfo.patJetFinalCollection
0435 
0436   nanoInfoForJet = nanoInfo_recojets[recoJetInfo.jet]
0437   jetTablePrefix = nanoInfoForJet["name"]
0438   jetTableDoc    = nanoInfoForJet["doc"]
0439   ptcut          = nanoInfoForJet["ptcut"] if "ptcut" in nanoInfoForJet else ""
0440   doPUIDVar      = nanoInfoForJet["doPUIDVar"] if "doPUIDVar" in nanoInfoForJet else False
0441   doQGL          = nanoInfoForJet["doQGL"] if "doQGL" in nanoInfoForJet else False
0442   doBTag         = nanoInfoForJet["doBTag"] if "doBTag" in nanoInfoForJet else False
0443 
0444   SavePatJets(proc,
0445     jetName, payload, patJetFinalColl, jetTablePrefix, jetTableDoc, doPF, doCalo,
0446     ptcut=ptcut, doPUIDVar=doPUIDVar, doQGL=doQGL, doBTag=doBTag, runOnMC=runOnMC
0447   )
0448 
0449   return proc
0450 
0451 def SavePatJets(proc, jetName, payload, patJetFinalColl, jetTablePrefix, jetTableDoc,
0452                 doPF, doCalo, ptcut="", doPUIDVar=False, doQGL=False, doBTag=False, runOnMC=False):
0453   """
0454   Schedule modules for a given patJet collection and save its variables into custom NanoAOD
0455   """
0456 
0457   #
0458   # Setup jet correction factors
0459   #
0460   jetCorrFactors = "jetCorrFactorsNano{}".format(jetName)
0461   setattr(proc, jetCorrFactors, jetCorrFactorsNano.clone(
0462       src = patJetFinalColl,
0463       payload = payload,
0464     )
0465   )
0466 
0467   #
0468   # Update jets
0469   #
0470   srcJets = "updatedJets{}".format(jetName)
0471   setattr(proc, srcJets, updatedJets.clone(
0472       jetSource = patJetFinalColl,
0473       jetCorrFactorsSource = [jetCorrFactors],
0474     )
0475   )
0476 
0477   #
0478   # Setup UserDataEmbedder
0479   #
0480   srcJetsWithUserData = "updatedJets{}WithUserData".format(jetName)
0481   setattr(proc, srcJetsWithUserData, cms.EDProducer("PATJetUserDataEmbedder",
0482       src = cms.InputTag(srcJets),
0483       userFloats = cms.PSet(),
0484       userInts = cms.PSet(),
0485     )
0486   )
0487 
0488   #
0489   # Filter jets with pt cut
0490   #
0491   finalJetsCutDefault = "(pt >= 8)"
0492   if runOnMC:
0493     finalJetsCutDefault = "(pt >= 8) || ((pt < 8) && (genJetFwdRef().backRef().isNonnull()))"
0494 
0495   finalJetsForTable = "finalJets{}".format(jetName)
0496   setattr(proc, finalJetsForTable, finalJets.clone(
0497       src = srcJetsWithUserData,
0498       cut = ptcut if ptcut != "" else finalJetsCutDefault
0499     )
0500   )
0501 
0502   #
0503   # Save jets in table
0504   #
0505   tableContent = PFJETVARS
0506   if doCalo:
0507     tableContent = CALOJETVARS
0508 
0509   jetTableCutDefault = "" #Don't apply any cuts for the table.
0510 
0511   jetTableDocDefault = jetTableDoc + " with JECs applied. Jets with pt > 8 GeV are stored."
0512   if runOnMC:
0513     jetTableDocDefault += "For jets with pt < 8 GeV, only those matched to gen jets are stored."
0514 
0515   jetTable = "jet{}Table".format(jetName)
0516   setattr(proc,jetTable, cms.EDProducer("SimpleCandidateFlatTableProducer",
0517       src = cms.InputTag(finalJetsForTable),
0518       cut = cms.string(jetTableCutDefault),
0519       name = cms.string(jetTablePrefix),
0520       doc  = cms.string(jetTableDocDefault),
0521       singleton = cms.bool(False), # the number of entries is variable
0522       extension = cms.bool(False), # this is the main table for the jets
0523       variables = cms.PSet(tableContent)
0524     )
0525   )
0526   getattr(proc,jetTable).variables.pt.precision=10
0527 
0528   #
0529   # Save MC-only jet variables in table
0530   #
0531   jetMCTable = "jet{}MCTable".format(jetName)
0532   setattr(proc, jetMCTable, cms.EDProducer("SimpleCandidateFlatTableProducer",
0533       src = cms.InputTag(finalJetsForTable),
0534       cut = getattr(proc,jetTable).cut,
0535       name = cms.string(jetTablePrefix),
0536       singleton = cms.bool(False),
0537       extension = cms.bool(True), # this is an extension table
0538       variables = cms.PSet(
0539         partonFlavour = Var("partonFlavour()", int, doc="flavour from parton matching"),
0540         hadronFlavour = Var("hadronFlavour()", int, doc="flavour from hadron ghost clustering"),
0541         genJetIdx = Var("?genJetFwdRef().backRef().isNonnull()?genJetFwdRef().backRef().key():-1", int, doc="index of matched gen jet"),
0542       )
0543     )
0544   )
0545 
0546   #
0547   # Define the jet modules sequence first
0548   #
0549   jetSequenceName = "jet{}Sequence".format(jetName)
0550   setattr(proc, jetSequenceName, cms.Sequence(
0551       getattr(proc,jetCorrFactors)+
0552       getattr(proc,srcJets)+
0553       getattr(proc,srcJetsWithUserData)+
0554       getattr(proc,finalJetsForTable)
0555     )
0556   )
0557 
0558   #
0559   # Define the jet table sequences
0560   #
0561   jetTableSequenceName = "jet{}TablesSequence".format(jetName)
0562   setattr(proc, jetTableSequenceName, cms.Sequence(getattr(proc,jetTable)))
0563 
0564   jetTableSequenceMCName = "jet{}MCTablesSequence".format(jetName)
0565   setattr(proc, jetTableSequenceMCName, cms.Sequence(getattr(proc,jetMCTable)))
0566 
0567   if runOnMC:
0568     proc.nanoSequenceMC += getattr(proc,jetSequenceName)
0569     proc.nanoSequenceMC += getattr(proc,jetTableSequenceName)
0570     proc.nanoSequenceMC += getattr(proc,jetTableSequenceMCName)
0571   else:
0572     proc.nanoSequence += getattr(proc,jetSequenceName)
0573     proc.nanoSequence += getattr(proc,jetTableSequenceName)
0574 
0575   #
0576   # Schedule plugins to calculate Jet ID, PileUp Jet ID input variables, and Quark-Gluon Likehood input variables.
0577   #
0578   if doPF:
0579     proc = AddJetID(proc, jetName=jetName, jetSrc=srcJets, jetTableName=jetTable, jetSequenceName=jetSequenceName)
0580     if doPUIDVar:
0581       proc = AddPileUpJetIDVars(proc, jetName=jetName, jetSrc=srcJets, jetTableName=jetTable, jetSequenceName=jetSequenceName)
0582     if doQGL:
0583       proc = AddQGLTaggerVars(proc,jetName=jetName, jetSrc=srcJets, jetTableName=jetTable, jetSequenceName=jetSequenceName, calculateQGLVars=True)
0584 
0585   #
0586   # Save b-tagging algorithm scores. Should only be done for jet collection with b-tagging
0587   # calculated when reclustered or collection saved with b-tagging info in MiniAOD
0588   #
0589   if doBTag:
0590     AddBTaggingScores(proc,jetTableName=jetTable)
0591     AddDeepJetGluonLQuarkScores(proc,jetTableName=jetTable)
0592     AddParticleNetAK4Scores(proc,jetTableName=jetTable)
0593 
0594   return proc
0595 
0596 def ReclusterAK4CHSJets(proc, recoJA, runOnMC):
0597   """
0598   Recluster AK4 CHS jets and replace slimmedJets
0599   that is used as default to save AK4 CHS jets
0600   in NanoAODs.
0601   """
0602   print("custom_jme_cff::ReclusterAK4CHSJets: Recluster AK4 PF CHS jets")
0603 
0604   #
0605   # Recluster AK4 CHS jets
0606   #
0607   cfg = {
0608     "jet" : "ak4pfchs",
0609     "inputCollection" : "",
0610     "genJetsCollection": "AK4GenJetsNoNu",
0611     "bTagDiscriminators": bTagDiscriminatorsForAK4,
0612     "minPtFastjet" : 0.,
0613   }
0614   recoJetInfo = recoJA.addRecoJetCollection(proc, **cfg)
0615 
0616   jetName = recoJetInfo.jetUpper
0617   patJetFinalColl = recoJetInfo.patJetFinalCollection
0618 
0619   #
0620   # Change the input jet source for jetCorrFactorsNano
0621   # and updatedJets
0622   #
0623   proc.jetCorrFactorsNano.src=patJetFinalColl
0624   proc.updatedJets.jetSource=patJetFinalColl
0625 
0626   #
0627   # Change pt cut
0628   #
0629   finalJetsCut = ""
0630   if runOnMC:
0631     finalJetsCut = "(pt >= 8) || ((pt < 8) && (genJetFwdRef().backRef().isNonnull()))"
0632   else:
0633     finalJetsCut = "(pt >= 8)"
0634 
0635   proc.finalJets.cut = finalJetsCut
0636   #
0637   # Add a minimum pt cut for corrT1METJets.
0638   #
0639   proc.corrT1METJetTable.cut = "pt>=8 && pt<15 && abs(eta)<9.9"
0640 
0641   #
0642   # Jet table cut
0643   #
0644   jetTableCut = "" # must not have any cut at the jetTable for AK4 CHS as it has been cross-cleaned
0645   proc.jetTable.cut   = jetTableCut
0646   proc.jetMCTable.cut = jetTableCut
0647 
0648   #
0649   # Jet table documentation
0650   #
0651   jetTableDoc = "AK4 PF CHS jets with JECs applied. Jets with pt > 8 GeV are stored."
0652   if runOnMC:
0653     jetTableDoc += "For jets with pt < 8 GeV, only those matched to AK4 Gen jets are stored."
0654   proc.jetTable.doc   = jetTableDoc
0655 
0656   #
0657   # Add variables
0658   #
0659   proc.jetTable.variables.hfHEF         = PFJETVARS.hfHEF
0660   proc.jetTable.variables.hfEmEF        = PFJETVARS.hfEmEF
0661   proc.jetTable.variables.nConstChHads  = PFJETVARS.nConstChHads
0662   proc.jetTable.variables.nConstNeuHads = PFJETVARS.nConstNeuHads
0663   proc.jetTable.variables.nConstHFHads  = PFJETVARS.nConstHFHads
0664   proc.jetTable.variables.nConstHFEMs   = PFJETVARS.nConstHFEMs
0665   proc.jetTable.variables.nConstMuons   = PFJETVARS.nConstMuons
0666   proc.jetTable.variables.nConstElecs   = PFJETVARS.nConstElecs
0667   proc.jetTable.variables.nConstPhotons = PFJETVARS.nConstPhotons
0668 
0669   #
0670   # Setup pileup jet ID with 80X training.
0671   #
0672   pileupJetId80X = "pileupJetId80X"
0673   setattr(proc, pileupJetId80X, pileupJetId.clone(
0674       jets = "updatedJets",
0675       algos = cms.VPSet(_chsalgos_81x),
0676       inputIsCorrected = True,
0677       applyJec = False,
0678       vertexes = "offlineSlimmedPrimaryVertices"
0679     )
0680   )
0681   proc.jetSequence.insert(proc.jetSequence.index(proc.pileupJetId94X), getattr(proc, pileupJetId80X))
0682 
0683   proc.updatedJetsWithUserData.userInts.puId80XfullId = cms.InputTag('pileupJetId80X:fullId')
0684   proc.updatedJetsWithUserData.userFloats.puId80XDisc = cms.InputTag("pileupJetId80X:fullDiscriminant")
0685 
0686   run2_nanoAOD_94X2016.toModify(proc.jetTable.variables, puIdDisc = Var("userFloat('puId80XDisc')",float,doc="Pilup ID discriminant with 80X (2016) training",precision=10))
0687 
0688   for modifier in run2_nanoAOD_94X2016, run2_nanoAOD_94XMiniAODv1, run2_nanoAOD_94XMiniAODv2, run2_nanoAOD_102Xv1:
0689     modifier.toModify(proc.jetTable.variables, puId = Var("userInt('puId80XfullId')", int, doc="Pileup ID flags with 80X (2016) training"))
0690 
0691   #
0692   # Add charged energy fraction from other primary vertices
0693   #
0694   proc.updatedJetsWithUserData.userFloats.chFPV1EF = cms.InputTag("jercVars:chargedFromPV1EnergyFraction")
0695   proc.updatedJetsWithUserData.userFloats.chFPV2EF = cms.InputTag("jercVars:chargedFromPV2EnergyFraction")
0696   proc.updatedJetsWithUserData.userFloats.chFPV3EF = cms.InputTag("jercVars:chargedFromPV3EnergyFraction")
0697   proc.jetTable.variables.chFPV1EF = Var("userFloat('chFPV1EF')", float, doc="charged fromPV==1 Energy Fraction (component of the total charged Energy Fraction).", precision= 6)
0698   proc.jetTable.variables.chFPV2EF = Var("userFloat('chFPV2EF')", float, doc="charged fromPV==2 Energy Fraction (component of the total charged Energy Fraction).", precision= 6)
0699   proc.jetTable.variables.chFPV3EF = Var("userFloat('chFPV3EF')", float, doc="charged fromPV==3 Energy Fraction (component of the total charged Energy Fraction).", precision= 6)
0700 
0701   #
0702   # Add variables for pileup jet ID studies.
0703   #
0704   proc = AddPileUpJetIDVars(proc,
0705     jetName = "",
0706     jetSrc = "updatedJets",
0707     jetTableName = "jetTable",
0708     jetSequenceName = "jetSequence"
0709   )
0710   #
0711   # Add variables for quark guon likelihood tagger studies.
0712   # Save variables as userFloats and userInts in each jet
0713   #
0714   proc.updatedJetsWithUserData.userFloats.qgl_axis2 = cms.InputTag("qgtagger:axis2")
0715   proc.updatedJetsWithUserData.userFloats.qgl_ptD   = cms.InputTag("qgtagger:ptD")
0716   proc.updatedJetsWithUserData.userInts.qgl_mult    = cms.InputTag("qgtagger:mult")
0717   #
0718   # Save quark gluon likelihood input variables variables
0719   #
0720   proc.jetTable.variables.qgl_axis2 =  QGLVARS.qgl_axis2
0721   proc.jetTable.variables.qgl_ptD   =  QGLVARS.qgl_ptD
0722   proc.jetTable.variables.qgl_mult  =  QGLVARS.qgl_mult
0723   #
0724   # Save standard b-tagging and c-tagging variables
0725   #
0726   proc.jetTable.variables.btagDeepB = BTAGVARS.btagDeepB
0727   proc.jetTable.variables.btagCSVV2 = BTAGVARS.btagCSVV2
0728   proc.jetTable.variables.btagDeepCvL = BTAGVARS.btagDeepCvL
0729   proc.jetTable.variables.btagDeepCvB = BTAGVARS.btagDeepCvB
0730   #
0731   # Save DeepJet b-tagging and c-tagging variables
0732   #
0733   proc.jetTable.variables.btagDeepFlavB    = DEEPJETVARS.btagDeepFlavB
0734   proc.jetTable.variables.btagDeepFlavCvL  = DEEPJETVARS.btagDeepFlavCvL
0735   proc.jetTable.variables.btagDeepFlavCvB  = DEEPJETVARS.btagDeepFlavCvB
0736   #
0737   # Save DeepJet raw score for gluon and light quarks
0738   #
0739   proc.jetTable.variables.btagDeepFlavG   = DEEPJETVARS.btagDeepFlavG
0740   proc.jetTable.variables.btagDeepFlavUDS = DEEPJETVARS.btagDeepFlavUDS
0741   proc.jetTable.variables.btagDeepFlavQG  = DEEPJETVARS.btagDeepFlavQG
0742   #
0743   # Add ParticleNetAK4 scores
0744   #
0745   proc.jetTable.variables.particleNetAK4_B          = PARTICLENETAK4VARS.particleNetAK4_B
0746   proc.jetTable.variables.particleNetAK4_CvsL       = PARTICLENETAK4VARS.particleNetAK4_CvsL
0747   proc.jetTable.variables.particleNetAK4_CvsB       = PARTICLENETAK4VARS.particleNetAK4_CvsB
0748   proc.jetTable.variables.particleNetAK4_QvsG       = PARTICLENETAK4VARS.particleNetAK4_QvsG
0749   proc.jetTable.variables.particleNetAK4_puIdDisc   = PARTICLENETAK4VARS.particleNetAK4_puIdDisc
0750 
0751   #Adding hf shower shape producer to the jet sequence. By default this producer is not automatically rerun at the NANOAOD step
0752   #The following lines make sure it is.
0753   hfJetShowerShapeforCustomNanoAOD = "hfJetShowerShapeforCustomNanoAOD"
0754   setattr(proc, hfJetShowerShapeforCustomNanoAOD, hfJetShowerShapeforNanoAOD.clone(jets="updatedJets",vertices="offlineSlimmedPrimaryVertices") )
0755   proc.jetSequence.insert(proc.jetSequence.index(proc.updatedJetsWithUserData), getattr(proc, hfJetShowerShapeforCustomNanoAOD))
0756   proc.updatedJetsWithUserData.userFloats.hfsigmaEtaEta = cms.InputTag('hfJetShowerShapeforCustomNanoAOD:sigmaEtaEta')
0757   proc.updatedJetsWithUserData.userFloats.hfsigmaPhiPhi = cms.InputTag('hfJetShowerShapeforCustomNanoAOD:sigmaPhiPhi')
0758   proc.updatedJetsWithUserData.userInts.hfcentralEtaStripSize = cms.InputTag('hfJetShowerShapeforCustomNanoAOD:centralEtaStripSize')
0759   proc.updatedJetsWithUserData.userInts.hfadjacentEtaStripsSize = cms.InputTag('hfJetShowerShapeforCustomNanoAOD:adjacentEtaStripsSize')
0760   proc.jetTable.variables.hfsigmaEtaEta = Var("userFloat('hfsigmaEtaEta')",float,doc="sigmaEtaEta for HF jets (noise discriminating variable)",precision=10)
0761   proc.jetTable.variables.hfsigmaPhiPhi = Var("userFloat('hfsigmaPhiPhi')",float,doc="sigmaPhiPhi for HF jets (noise discriminating variable)",precision=10)
0762   proc.jetTable.variables.hfcentralEtaStripSize = Var("userInt('hfcentralEtaStripSize')", int, doc="eta size of the central tower strip in HF (noise discriminating variable) ")
0763   proc.jetTable.variables.hfadjacentEtaStripsSize = Var("userInt('hfadjacentEtaStripsSize')", int, doc="eta size of the strips next to the central tower strip in HF (noise discriminating variable) ")
0764 
0765   return proc
0766 
0767 def AddNewAK8PuppiJetsForJEC(proc, recoJA, runOnMC):
0768   """
0769   Store a separate AK8 Puppi jet collection for JEC studies.
0770   Only minimal info are stored
0771   """
0772   print("custom_jme_cff::AddNewAK8PuppiJetsForJEC: Make a new AK8 PF Puppi jet collection for JEC studies")
0773 
0774   #
0775   # Recluster AK8 Puppi jets
0776   #
0777   cfg = {
0778     "jet" : "ak8pfpuppi",
0779     "inputCollection" : "",
0780     "genJetsCollection": "AK8GenJetsNoNu",
0781     "minPtFastjet" : 0., # Remove any pt threshold at the jet clustering stage.
0782   }
0783   recoJetInfo = recoJA.addRecoJetCollection(proc, **cfg)
0784 
0785   jetName = recoJetInfo.jetUpper
0786   payload = recoJetInfo.jetCorrPayload
0787 
0788   patJetFinalColl = recoJetInfo.patJetFinalCollection
0789   jetTablePrefix  = "FatJetForJEC"
0790   jetTableDoc     = "AK8 PF Puppi jets with JECs applied. Reclustered for JEC studies so only minimal info stored."
0791   ptcut           = ""# No need to specify ptcut. Use default in SavePatJets function
0792 
0793   SavePatJets(proc,
0794     jetName, payload, patJetFinalColl, jetTablePrefix, jetTableDoc, doPF=True,
0795     doCalo=False, ptcut=ptcut, doPUIDVar=False, doQGL=False, doBTag=False, runOnMC=runOnMC
0796   )
0797 
0798   return proc
0799 
0800 def AddNewAK8CHSJets(proc, recoJA, runOnMC):
0801   """
0802   Store an AK8 CHS jet collection for JEC studies.
0803   """
0804   print("custom_jme_cff::AddNewAK8CHSJets: Make a new AK8 PF CHS jet collection for JEC studies")
0805 
0806   #
0807   # Recluster AK8 CHS jets
0808   #
0809   cfg = {
0810     "jet" : "ak8pfchs",
0811     "inputCollection" : "",
0812     "genJetsCollection": "AK8GenJetsNoNu",
0813     "minPtFastjet" : 0., # Remove any pt threshold at the jet clustering stage.
0814   }
0815   recoJetInfo = recoJA.addRecoJetCollection(proc, **cfg)
0816 
0817   jetName = recoJetInfo.jetUpper
0818   payload = recoJetInfo.jetCorrPayload
0819 
0820   patJetFinalColl = recoJetInfo.patJetFinalCollection
0821   jetTablePrefix  = "FatJetCHS"
0822   jetTableDoc     = "AK8 PF CHS jets with JECs applied. Reclustered for JEC studies so only minimal info stored."
0823   ptcut           = ""# No need to specify ptcut. Use default in SavePatJets function
0824 
0825   SavePatJets(proc,
0826     jetName, payload, patJetFinalColl, jetTablePrefix, jetTableDoc, doPF=True,
0827     doCalo=False, ptcut=ptcut, doPUIDVar=False, doQGL=False, doBTag=False, runOnMC=runOnMC
0828   )
0829 
0830   return proc
0831 
0832 def AddVariablesForAK8PuppiJets(proc):
0833   """
0834   Add more variables for AK8 PFPUPPI jets
0835   """
0836 
0837   #
0838   #  These variables are not stored for AK8PFPUPPI (slimmedJetsAK8)
0839   #  in MiniAOD if their pt < 170 GeV. Hence the conditional fill.
0840   #
0841   proc.fatJetTable.variables.chHEF  = Var("?isPFJet()?chargedHadronEnergyFraction():-1", float, doc="charged Hadron Energy Fraction", precision = 6)
0842   proc.fatJetTable.variables.neHEF  = Var("?isPFJet()?neutralHadronEnergyFraction():-1", float, doc="neutral Hadron Energy Fraction", precision = 6)
0843   proc.fatJetTable.variables.chEmEF = Var("?isPFJet()?chargedEmEnergyFraction():-1", float, doc="charged Electromagnetic Energy Fraction", precision = 6)
0844   proc.fatJetTable.variables.neEmEF = Var("?isPFJet()?neutralEmEnergyFraction():-1", float, doc="neutral Electromagnetic Energy Fraction", precision = 6)
0845   proc.fatJetTable.variables.muEF   = Var("?isPFJet()?muonEnergyFraction():-1", float, doc="muon Energy Fraction", precision = 6)
0846   proc.fatJetTable.variables.hfHEF  = Var("?isPFJet()?HFHadronEnergyFraction():-1", float, doc="energy fraction in forward hadronic calorimeter", precision = 6)
0847   proc.fatJetTable.variables.hfEmEF = Var("?isPFJet()?HFEMEnergyFraction():-1", float, doc="energy fraction in forward EM calorimeter", precision = 6)
0848   proc.fatJetTable.variables.nConstChHads   = Var("?isPFJet()?chargedHadronMultiplicity():-1",int, doc="number of charged hadrons in the jet")
0849   proc.fatJetTable.variables.nConstNeuHads  = Var("?isPFJet()?neutralHadronMultiplicity():-1",int, doc="number of neutral hadrons in the jet")
0850   proc.fatJetTable.variables.nConstHFHads   = Var("?isPFJet()?HFHadronMultiplicity():-1", int, doc="number of HF Hadrons in the jet")
0851   proc.fatJetTable.variables.nConstHFEMs    = Var("?isPFJet()?HFEMMultiplicity():-1", int, doc="number of HF EMs in the jet")
0852   proc.fatJetTable.variables.nConstMuons    = Var("?isPFJet()?muonMultiplicity():-1", int, doc="number of muons in the jet")
0853   proc.fatJetTable.variables.nConstElecs    = Var("?isPFJet()?electronMultiplicity():-1", int, doc="number of electrons in the jet")
0854   proc.fatJetTable.variables.nConstPhotons  = Var("?isPFJet()?photonMultiplicity():-1", int, doc="number of photons in the jet")
0855 
0856   return proc
0857 #******************************************
0858 #
0859 #
0860 # Gen Jets related functions
0861 #
0862 #
0863 #******************************************
0864 def AddNewGenJets(proc, genJetInfo):
0865   """
0866   Add genJet into custom nanoAOD
0867   """
0868 
0869   genJetName         = genJetInfo.jetUpper
0870   genJetAlgo         = genJetInfo.jetAlgo
0871   genJetSize         = genJetInfo.jetSize
0872   genJetSizeNr       = genJetInfo.jetSizeNr
0873   genJetFinalColl    = "{}{}{}".format(genJetAlgo.upper(), genJetSize, "GenJetsNoNu")
0874   genJetTablePrefix  = nanoInfo_genjets[genJetInfo.jet]["name"]
0875   genJetTableDoc     = nanoInfo_genjets[genJetInfo.jet]["doc"]
0876 
0877   SaveGenJets(proc, genJetName, genJetAlgo, genJetSizeNr, genJetFinalColl, genJetTablePrefix, genJetTableDoc, runOnMC=False)
0878 
0879   return proc
0880 
0881 def SaveGenJets(proc, genJetName, genJetAlgo, genJetSizeNr, genJetFinalColl, genJetTablePrefix, genJetTableDoc, runOnMC=False):
0882   """
0883   Schedule modules for a given genJet collection and save its variables into custom NanoAOD
0884   """
0885 
0886   genJetTableThisJet = "jet{}Table".format(genJetName)
0887   setattr(proc, genJetTableThisJet, genJetTable.clone(
0888       src       = genJetFinalColl,
0889       cut       = "", # No cut specified here. Save all gen jets after clustering
0890       name      = genJetTablePrefix,
0891       doc       = genJetTableDoc,
0892       variables = GENJETVARS
0893     )
0894   )
0895 
0896   genJetFlavourAssociationThisJet = "genJet{}FlavourAssociation".format(genJetName)
0897   setattr(proc, genJetFlavourAssociationThisJet, genJetFlavourAssociation.clone(
0898       jets           = getattr(proc,genJetTableThisJet).src,
0899       jetAlgorithm   = supportedJetAlgos[genJetAlgo],
0900       rParam         = genJetSizeNr,
0901     )
0902   )
0903 
0904   genJetFlavourTableThisJet = "genJet{}FlavourTable".format(genJetName)
0905   setattr(proc, genJetFlavourTableThisJet, genJetFlavourTable.clone(
0906       name            = getattr(proc,genJetTableThisJet).name,
0907       src             = getattr(proc,genJetTableThisJet).src,
0908       cut             = getattr(proc,genJetTableThisJet).cut,
0909       jetFlavourInfos = genJetFlavourAssociationThisJet,
0910     )
0911   )
0912 
0913   genJetSequenceName = "genJet{}Sequence".format(genJetName)
0914   setattr(proc, genJetSequenceName, cms.Sequence(
0915       getattr(proc,genJetTableThisJet)+
0916       getattr(proc,genJetFlavourAssociationThisJet)+
0917       getattr(proc,genJetFlavourTableThisJet)
0918     )
0919   )
0920   proc.nanoSequenceMC.insert(proc.nanoSequenceMC.index(proc.jetMC)+1, getattr(proc,genJetSequenceName))
0921 
0922   return proc
0923 
0924 def ReclusterAK4GenJets(proc, genJA):
0925   """
0926   Recluster AK4 Gen jets and replace
0927   slimmedGenJets that is used as default
0928   to save AK4 Gen jets in NanoAODs.
0929   """
0930   print("custom_jme_cff::ReclusterAK4GenJets: Recluster AK4 Gen jets")
0931 
0932   #
0933   # Recluster AK4 Gen jet
0934   #
0935   cfg = {
0936     "jet" : "ak4gen",
0937   }
0938   genJetInfo = genJA.addGenJetCollection(proc, **cfg)
0939 
0940   genJetName      = genJetInfo.jetUpper
0941   genJetAlgo      = genJetInfo.jetAlgo
0942   genJetSize      = genJetInfo.jetSize
0943   genJetSizeNr    = genJetInfo.jetSizeNr
0944   selectedGenJets = "{}{}{}".format(genJetAlgo.upper(), genJetSize, "GenJetsNoNu")
0945 
0946   #
0947   # Change jet source to the newly clustered jet collection. Set very low pt cut for jets
0948   # to be stored in the GenJet Table
0949   #
0950   proc.genJetTable.src = selectedGenJets
0951   proc.genJetTable.cut = "" # No cut specified here. Save all gen jets after clustering
0952   proc.genJetTable.doc = "AK4 Gen jets (made with visible genparticles) with pt > 3 GeV" # default pt cut after clustering is 3 GeV
0953 
0954   genJetFlavourAssociationThisJet = "genJet{}FlavourAssociation".format(genJetName)
0955   setattr(proc, genJetFlavourAssociationThisJet, genJetFlavourAssociation.clone(
0956       jets           = proc.genJetTable.src,
0957       jetAlgorithm   = supportedJetAlgos[genJetAlgo],
0958       rParam         = genJetSizeNr,
0959     )
0960   )
0961   proc.jetMC.insert(proc.jetMC.index(proc.genJetFlavourTable), getattr(proc, genJetFlavourAssociationThisJet))
0962   return proc
0963 
0964 def AddNewAK8GenJetsForJEC(proc, genJA):
0965   """
0966   Make a separate AK8 Gen jet collection for JEC studies.
0967   """
0968   print("custom_jme_cff::AddNewAK8GenJetsForJEC: Add new AK8 Gen jets for JEC studies")
0969 
0970   #
0971   # Recluster AK8 Gen jet
0972   #
0973   cfg = {
0974     "jet" : "ak8gen",
0975   }
0976   genJetInfo = genJA.addGenJetCollection(proc, **cfg)
0977 
0978   genJetName         = genJetInfo.jetUpper
0979   genJetAlgo         = genJetInfo.jetAlgo
0980   genJetSize         = genJetInfo.jetSize
0981   genJetSizeNr       = genJetInfo.jetSizeNr
0982   genJetFinalColl    = "{}{}{}".format(genJetAlgo.upper(), genJetSize, "GenJetsNoNu")
0983   genJetTablePrefix  = "GenJetAK8ForJEC"
0984   genJetTableDoc     = "AK8 Gen jets (made with visible genparticles) with pt > 3 GeV. Reclustered for JEC studies."
0985 
0986   SaveGenJets(proc, genJetName, genJetAlgo, genJetSizeNr, genJetFinalColl, genJetTablePrefix, genJetTableDoc, runOnMC=False)
0987 
0988   return proc
0989 
0990 def AddVariablesForAK4GenJets(proc):
0991   proc.genJetTable.variables.nConstituents = GENJETVARS.nConstituents
0992   return proc
0993 
0994 def AddVariablesForAK8GenJets(proc):
0995   proc.genJetAK8Table.variables.nConstituents = GENJETVARS.nConstituents
0996   return proc
0997 
0998 #===========================================================================
0999 #
1000 # Misc. functions
1001 #
1002 #===========================================================================
1003 def RemoveAllJetPtCuts(proc):
1004   """
1005   Remove default pt cuts for all jets set in jets_cff.py
1006   """
1007 
1008   proc.finalJets.cut             = "" # 15 -> 10
1009   proc.finalJetsAK8.cut          = "" # 170 -> 170
1010   proc.genJetTable.cut           = "" # 10 -> 8
1011   proc.genJetFlavourTable.cut    = "" # 10 -> 8
1012   proc.genJetAK8Table.cut        = "" # 100 -> 80
1013   proc.genJetAK8FlavourTable.cut = "" # 100 -> 80
1014 
1015   return proc
1016 
1017 #===========================================================================
1018 #
1019 # CUSTOMIZATION function
1020 #
1021 #===========================================================================
1022 def PrepJMECustomNanoAOD(process,runOnMC):
1023 
1024   ############################################################################
1025   # Remove all default jet pt cuts from jets_cff.py
1026   ############################################################################
1027   process = RemoveAllJetPtCuts(process)
1028 
1029   ###########################################################################
1030   #
1031   # Gen-level jets related functions. Only for MC.
1032   #
1033   ###########################################################################
1034   genJA = GenJetAdder()
1035   if runOnMC:
1036     ############################################################################
1037     # Save additional variables for AK8 GEN jets
1038     ############################################################################
1039     process = AddVariablesForAK8GenJets(process)
1040     ############################################################################
1041     # Recluster AK8 GEN jets
1042     ############################################################################
1043     process = AddNewAK8GenJetsForJEC(process, genJA)
1044     ###########################################################################
1045     # Recluster AK4 GEN jets
1046     ###########################################################################
1047     process = ReclusterAK4GenJets(process, genJA)
1048     process = AddVariablesForAK4GenJets(process)
1049     ###########################################################################
1050     # Add additional GEN jets to NanoAOD
1051     ###########################################################################
1052     for jetConfig in config_genjets:
1053       cfg = { k : v for k, v in jetConfig.items() if k != "enabled"}
1054       genJetInfo = genJA.addGenJetCollection(process, **cfg)
1055       AddNewGenJets(process, genJetInfo)
1056 
1057   ###########################################################################
1058   #
1059   # Reco-level jets related functions. For both MC and data.
1060   #
1061   ###########################################################################
1062   recoJA = RecoJetAdder(runOnMC=runOnMC)
1063   ###########################################################################
1064   # Save additional variables for AK8Puppi jets
1065   ###########################################################################
1066   process = AddVariablesForAK8PuppiJets(process)
1067   ###########################################################################
1068   # Build a separate AK8Puppi jet collection for JEC studies
1069   ###########################################################################
1070   process = AddNewAK8PuppiJetsForJEC(process, recoJA, runOnMC)
1071   ###########################################################################
1072   # Build a AK8CHS jet collection for JEC studies
1073   ###########################################################################
1074   process = AddNewAK8CHSJets(process, recoJA, runOnMC)
1075   ###########################################################################
1076   # Recluster AK4 CHS jets and replace "slimmedJets"
1077   ###########################################################################
1078   process = ReclusterAK4CHSJets(process, recoJA, runOnMC)
1079   ###########################################################################
1080   # Add additional Reco jets to NanoAOD
1081   ###########################################################################
1082   for jetConfig in config_recojets:
1083     cfg = { k : v for k, v in jetConfig.items() if k != "enabled"}
1084     recoJetInfo = recoJA.addRecoJetCollection(process, **cfg)
1085     AddNewPatJets(process, recoJetInfo, runOnMC)
1086 
1087   ###########################################################################
1088   # Save Maximum of Pt Hat Max
1089   ###########################################################################
1090   if runOnMC:
1091     process.puTable.savePtHatMax = True
1092 
1093   ###########################################################################
1094   # Save all Parton-Shower weights
1095   ###########################################################################
1096   if runOnMC:
1097     process.genWeightsTable.keepAllPSWeights = True
1098 
1099   return process
1100 
1101 def PrepJMECustomNanoAOD_MC(process):
1102   PrepJMECustomNanoAOD(process,runOnMC=True)
1103   return process
1104 
1105 def PrepJMECustomNanoAOD_Data(process):
1106   PrepJMECustomNanoAOD(process,runOnMC=False)
1107   return process