NanoAOD/python/custom_jme_cff.py

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