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.simpleGenJetFlatTableProducer_cfi import simpleGenJetFlatTableProducer
0005 from PhysicsTools.NanoAOD.simplePATJetFlatTableProducer_cfi import simplePATJetFlatTableProducer
0006
0007 from RecoJets.JetProducers.hfJetShowerShape_cfi import hfJetShowerShape
0008 from RecoJets.JetProducers.PileupJetID_cfi import pileupJetIdCalculator, pileupJetId
0009
0010 from PhysicsTools.NanoAOD.common_cff import Var, P4Vars
0011 from PhysicsTools.NanoAOD.jetsAK4_CHS_cff import jetTable, jetCorrFactorsNano, updatedJets, finalJets, qgtagger
0012 from PhysicsTools.NanoAOD.jetsAK4_Puppi_cff import jetPuppiTable, jetPuppiCorrFactorsNano, updatedJetsPuppi, updatedJetsPuppiWithUserData
0013 from PhysicsTools.NanoAOD.jetMC_cff  import genJetTable, genJetFlavourAssociation, genJetFlavourTable
0014
0015 from PhysicsTools.PatAlgos.tools.jetCollectionTools import GenJetAdder, RecoJetAdder
0016 from PhysicsTools.PatAlgos.tools.jetTools import supportedJetAlgos
0017 from PhysicsTools.PatAlgos.tools.jetTools import updateJetCollection
0018
0019 bTagCSVV2    = ['pfCombinedInclusiveSecondaryVertexV2BJetTags']
0020 bTagDeepCSV  = ['pfDeepCSVJetTags:probb','pfDeepCSVJetTags:probbb','pfDeepCSVJetTags:probc','pfDeepCSVJetTags:probudsg']
0021 bTagDeepJet  = [
0022   'pfDeepFlavourJetTags:probb','pfDeepFlavourJetTags:probbb','pfDeepFlavourJetTags:problepb',
0023   'pfDeepFlavourJetTags:probc','pfDeepFlavourJetTags:probuds','pfDeepFlavourJetTags:probg'
0024 ]
0025 from RecoBTag.ONNXRuntime.pfParticleNetAK4_cff import _pfParticleNetAK4JetTagsAll
0026 from RecoBTag.ONNXRuntime.pfParticleNetFromMiniAODAK4_cff import _pfParticleNetFromMiniAODAK4PuppiCentralJetTagsAll
0027 from RecoBTag.ONNXRuntime.pfParticleNetFromMiniAODAK4_cff import _pfParticleNetFromMiniAODAK4PuppiForwardJetTagsAll
0028 from RecoBTag.ONNXRuntime.pfUnifiedParticleTransformerAK4_cff import _pfUnifiedParticleTransformerAK4JetTagsAll
0029 from RecoBTag.ONNXRuntime.pfUnifiedParticleTransformerAK4V1_cff import _pfUnifiedParticleTransformerAK4V1JetTagsAll
0030 bTagDiscriminatorsForAK4 = cms.PSet(foo = cms.vstring(
0031   bTagDeepJet+
0032   _pfParticleNetFromMiniAODAK4PuppiCentralJetTagsAll+_pfParticleNetFromMiniAODAK4PuppiForwardJetTagsAll+
0033   _pfUnifiedParticleTransformerAK4JetTagsAll+_pfUnifiedParticleTransformerAK4V1JetTagsAll
0034 ))
0035 run2_nanoAOD_ANY.toModify(
0036   bTagDiscriminatorsForAK4,
0037   foo = bTagCSVV2+bTagDeepCSV+bTagDeepJet+_pfParticleNetAK4JetTagsAll
0038 )
0039 bTagDiscriminatorsForAK4 = bTagDiscriminatorsForAK4.foo.value()
0040
0041
0042 #
0043 # By default, these collections are saved in NanoAODs:
0044 # - ak4gen (GenJet in NanoAOD), slimmedGenJets in MiniAOD
0045 # - ak8gen (GenJetAK8 in NanoAOD), slimmedGenJetsAK8 in MiniAOD
0046 # Below is a list of genjets that we can save in NanoAOD. Set
0047 # "enabled" to true if you want to store the jet collection
0048 config_genjets = [
0049   {
0050     "jet"     : "ak6gen",
0051     "enabled" : False,
0052   },
0053 ]
0054 config_genjets = list(filter(lambda k: k['enabled'], config_genjets))
0055 #
0056 # GenJets info in NanoAOD
0057 #
0058 nanoInfo_genjets = {
0059   "ak6gen"  : {
0060     "name" : "GenJetAK6",
0061     "doc"  : "AK6 Gen jets (made with visible genparticles) with pt > 3 GeV", # default genjets pt cut after clustering is 3 GeV
0062   },
0063 }
0064 #
0065 # By default, these collections are saved in the main NanoAODs:
0066 # - ak4pfpuppi  (Jet   in NanoAOD), slimmedJetsPuppi in MiniAOD
0067 # - ak8pfpuppi (FatJet in NanoAOD), slimmedJetsAK8 in MiniAOD
0068 # Below is a list of recojets that we can save in NanoAOD. Set
0069 # "enabled" to true if you want to store the recojet collection.
0070 #
0071 config_recojets = [
0072   {
0073     "jet" : "ak4calo",
0074     "enabled" : True,
0075     "inputCollection"  : "slimmedCaloJets", #Exist in MiniAOD
0076     "genJetsCollection": "AK4GenJetsNoNu",# Setup in ConfigureAK4GenJets() function
0077   },
0078   {
0079     "jet" : "ak4pf",
0080     "enabled" : False,
0081     "inputCollection" : "",
0082     "genJetsCollection": "AK4GenJetsNoNu",# Setup in ConfigureAK4GenJets() function
0083     "minPtFastjet" : 0.,
0084   },
0085   {
0086     "jet" : "ak8pf",
0087     "enabled" : False,
0088     "inputCollection" : "",
0089     "genJetsCollection": "AK8GenJetsNoNu", # Setup in AddNewAK8GenJetsForJEC() function
0090     "minPtFastjet" : 0.,
0091   },
0092 ]
0093 config_recojets = list(filter(lambda k: k['enabled'], config_recojets))
0094 #
0095 # RecoJets info in NanoAOD
0096 #
0097 nanoInfo_recojets = {
0098   "ak4calo" : {
0099     "name": "JetCalo",
0100     "doc" : "AK4 Calo jets (slimmedCaloJets)",
0101     "ptcutGenMatch": 5
0102   },
0103   "ak4pf" : {
0104     "name"  : "JetPF",
0105     "doc"   : "AK4 PF jets",
0106     "ptcut" : "",
0107   },
0108   "ak8pf" : {
0109     "name"  : "FatJetPF",
0110     "doc"   : "AK8 PF jets",
0111     "ptcut" : "",
0112   },
0113 }
0114
0115 GENJETVARS = cms.PSet(P4Vars,
0116   nConstituents = Var("numberOfDaughters()","uint8",doc="Number of particles in the jet"),
0117   chHEF  = Var("chargedHadronEnergy()/energy()", float, doc="charged Hadron Energy Fraction", precision=10),
0118   neHEF  = Var("neutralHadronEnergy()/energy()", float, doc="neutral Hadron Energy Fraction", precision=10),
0119   chEmEF = Var("chargedEmEnergy()/energy()", float, doc="charged EM Energy Fraction", precision=10),
0120   neEmEF = Var("neutralEmEnergy()/energy()", float, doc="neutral EM Energy Fraction", precision=10),
0121   muEF   = Var("muonEnergy()/energy()", float, doc="muon Energy", precision=10),
0122   chHadMultiplicity = Var("chargedHadronMultiplicity", "int16", doc="number of charged hadrons in the jet"),
0123   neHadMultiplicity = Var("neutralHadronMultiplicity", "int16", doc="number of neutral hadrons in the jet"),
0124   chEmMultiplicity  = Var("chargedEmMultiplicity", "int16", doc="number of charged EM particles in the jet"),
0125   neEmMultiplicity  = Var("neutralEmMultiplicity", "int16", doc="number of neutral EM particles in the jet"),
0126   muMultiplicity    = Var("muonMultiplicity", "int16", doc="number of muons in the jet"),
0127 )
0128 PFJETVARS = cms.PSet(P4Vars,
0129   rawFactor     = jetPuppiTable.variables.rawFactor,
0130   area          = jetPuppiTable.variables.area,
0131   chHEF         = jetPuppiTable.variables.chHEF,
0132   neHEF         = jetPuppiTable.variables.neHEF,
0133   chEmEF        = jetPuppiTable.variables.chEmEF,
0134   neEmEF        = jetPuppiTable.variables.neEmEF,
0135   muEF          = jetPuppiTable.variables.muEF,
0136   hfHEF         = jetPuppiTable.variables.hfHEF,
0137   hfEmEF        = jetPuppiTable.variables.hfEmEF,
0138   nMuons        = jetPuppiTable.variables.nMuons,
0139   nElectrons    = jetPuppiTable.variables.nElectrons,
0140   nConstituents = jetPuppiTable.variables.nConstituents,
0141   chHadMultiplicity = Var("chargedHadronMultiplicity()","int16",doc="(Puppi-weighted) number of charged hadrons in the jet"),
0142   neHadMultiplicity = Var("neutralHadronMultiplicity()","int16",doc="(Puppi-weighted) number of neutral hadrons in the jet"),
0143   hfHadMultiplicity = Var("HFHadronMultiplicity()", "int16",doc="(Puppi-weighted) number of HF hadrons in the jet"),
0144   hfEMMultiplicity  = Var("HFEMMultiplicity()","int16",doc="(Puppi-weighted) number of HF EMs in the jet"),
0145   muMultiplicity    = Var("muonMultiplicity()","int16",doc="(Puppi-weighted) number of muons in the jet"),
0146   elMultiplicity    = Var("electronMultiplicity()","int16",doc="(Puppi-weighted) number of electrons in the jet"),
0147   phoMultiplicity   = Var("photonMultiplicity()","int16",doc="(Puppi-weighted) number of photons in the jet"),
0148 )
0149 #MC-only variables
0150 PFJETVARS_MCOnly = cms.PSet(
0151   partonFlavour     = Var("partonFlavour()", "int16", doc="flavour from parton matching"),
0152   hadronFlavour     = Var("hadronFlavour()", "uint8", doc="flavour from hadron ghost clustering"),
0153   genJetIdx         = Var("?genJetFwdRef().backRef().isNonnull()&& genJetFwdRef().backRef().pt() > 5?genJetFwdRef().backRef().key():-1", "int16", doc="index of matched gen jet")
0154 )
0155 PUIDVARS = cms.PSet(
0156   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),
0157   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),
0158   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),
0159   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),
0160   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),
0161   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),
0162   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),
0163   puId_ptD        = Var("?(pt>=10)?userFloat('puId_ptD'):-1",float,doc="pT-weighted average pT of constituents (PileUp ID BDT input variable)", precision=14),
0164   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),
0165   puId_pull       = Var("?(pt>=10)?userFloat('puId_pull'):-1",float,doc="magnitude of pull vector (PileUp ID BDT input variable)", precision=14),
0166   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),
0167   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),
0168   puId_nCharged   = Var("?(pt>=10)?userInt('puId_nCharged'):-1","int16",doc="number of charged constituents (PileUp ID BDT input variable)"),
0169 )
0170 QGLVARS = cms.PSet(
0171   qgl_axis2       =  Var("?(pt>=10)?userFloat('qgl_axis2'):-1",float,doc="ellipse minor jet axis (Quark vs Gluon likelihood input variable)", precision=14),
0172   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),
0173   qgl_mult        =  Var("?(pt>=10)?userInt('qgl_mult'):-1", "int16",doc="PF candidates multiplicity (Quark vs Gluon likelihood input variable)"),
0174 )
0175 BTAGVARS = cms.PSet(
0176   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),
0177   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),
0178   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),
0179 )
0180 DEEPJETVARS = cms.PSet(
0181   btagDeepFlavB   = Var("?(pt>=15)?bDiscriminator('pfDeepFlavourJetTags:probb')+bDiscriminator('pfDeepFlavourJetTags:probbb')+bDiscriminator('pfDeepFlavourJetTags:problepb'):-1",float,doc="DeepJet b+bb+lepb tag discriminator",precision=10),
0182   btagDeepFlavC   = Var("?(pt>=15)?bDiscriminator('pfDeepFlavourJetTags:probc'):-1",float,doc="DeepFlavour charm tag raw score",precision=10),
0183   btagDeepFlavG   = Var("?(pt>=15)?bDiscriminator('pfDeepFlavourJetTags:probg'):-1",float,doc="DeepFlavour gluon tag raw score",precision=10),
0184   btagDeepFlavUDS = Var("?(pt>=15)?bDiscriminator('pfDeepFlavourJetTags:probuds'):-1",float,doc="DeepFlavour uds tag raw score",precision=10),
0185   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),
0186   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),
0187   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),
0188 )
0189 ROBUSTPARTAK4VARS = cms.PSet(
0190   btagRobustParTAK4B   = Var("?(pt>=15)?bDiscriminator('pfParticleTransformerAK4JetTags:probb')+bDiscriminator('pfParticleTransformerAK4JetTags:probbb')+bDiscriminator('pfParticleTransformerAK4JetTags:problepb'):-1",float,doc="RobustParTAK4 b+bb+lepb tag discriminator",precision=10),
0191   btagRobustParTAK4C   = Var("?(pt>=15)?bDiscriminator('pfParticleTransformerAK4JetTags:probc'):-1",float,doc="RobustParTAK4 charm tag raw score",precision=10),
0192   btagRobustParTAK4G   = Var("?(pt>=15)?bDiscriminator('pfParticleTransformerAK4JetTags:probg'):-1",float,doc="RobustParTAK4 gluon tag raw score",precision=10),
0193   btagRobustParTAK4UDS = Var("?(pt>=15)?bDiscriminator('pfParticleTransformerAK4JetTags:probuds'):-1",float,doc="RobustParTAK4 uds tag raw score",precision=10),
0194   btagRobustParTAK4CvL = Var("?(pt>=15)&&(bDiscriminator('pfParticleTransformerAK4JetTags:probc')+bDiscriminator('pfParticleTransformerAK4JetTags:probuds')+bDiscriminator('pfParticleTransformerAK4JetTags:probg'))>0?bDiscriminator('pfParticleTransformerAK4JetTags:probc')/(bDiscriminator('pfParticleTransformerAK4JetTags:probc')+bDiscriminator('pfParticleTransformerAK4JetTags:probuds')+bDiscriminator('pfParticleTransformerAK4JetTags:probg')):-1",float,doc="RobustParTAK4 c vs uds+g discriminator",precision=10),
0195   btagRobustParTAK4CvB = Var("?(pt>=15)&&(bDiscriminator('pfParticleTransformerAK4JetTags:probc')+bDiscriminator('pfParticleTransformerAK4JetTags:probb')+bDiscriminator('pfParticleTransformerAK4JetTags:probbb')+bDiscriminator('pfParticleTransformerAK4JetTags:problepb'))>0?bDiscriminator('pfParticleTransformerAK4JetTags:probc')/(bDiscriminator('pfParticleTransformerAK4JetTags:probc')+bDiscriminator('pfParticleTransformerAK4JetTags:probb')+bDiscriminator('pfParticleTransformerAK4JetTags:probbb')+bDiscriminator('pfParticleTransformerAK4JetTags:problepb')):-1",float,doc="RobustParTAK4 c vs b+bb+lepb discriminator",precision=10),
0196   btagRobustParTAK4QG  = Var("?(pt>=15)&&(bDiscriminator('pfParticleTransformerAK4JetTags:probg')+bDiscriminator('pfParticleTransformerAK4JetTags:probuds'))>0?bDiscriminator('pfParticleTransformerAK4JetTags:probg')/(bDiscriminator('pfParticleTransformerAK4JetTags:probg')+bDiscriminator('pfParticleTransformerAK4JetTags:probuds')):-1",float,doc="RobustParTAK4 g vs uds discriminator",precision=10),
0197 )
0198 UNIFIEDPARTAK4VARS = cms.PSet(
0199   btagUParTAK4B       = Var("?pt>=10 && bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:BvsAll')>0?bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:BvsAll'):-1",float,precision=12,doc="UnifiedParTAK4 b vs. udscg"),
0200   btagUParTAK4CvL     = Var("?pt>=10 && bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:CvsL')>0?bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:CvsL'):-1",float,precision=12,doc="UnifiedParTAK4 c vs. udsg"),
0201   btagUParTAK4CvB     = Var("?pt>=10 && bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:CvsB')>0?bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:CvsB'):-1",float,precision=12,doc="UnifiedParTAK4 c vs. b"),
0202   btagUParTAK4CvNotB  = Var("?pt>=10 && ((bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:probb')+bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:probbb')+bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:problepb')))>0?((bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:probc'))/(1.-bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:probb')-bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:probbb')-bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:problepb'))):-1",float,precision=12,doc="UnifiedParT c vs. not b"),
0203   btagUParTAK4SvCB    = Var("?pt>=10 && bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:SvsBC')>0?bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:SvsBC'):-1",float,precision=12,doc="UnifiedParTAK4 s vs. bc"),
0204   btagUParTAK4SvUDG   = Var("?pt>=10 && bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:SvsUDG')>0?bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:SvsUDG'):-1",float,precision=12,doc="UnifiedParTAK4 s vs. udg"),
0205   btagUParTAK4UDG     = Var("?pt>=10? bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:probu')+bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:probd')+bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:probg'):-1",float,precision=12,doc="UnifiedParTAK4 u+d+g raw score"),
0206   btagUParTAK4QvG     = Var("?pt>=10 && bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:QvsG')>0?bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:QvsG'):-1",float,precision=12,doc="UnifiedParTAK4 q (uds) vs. g"),
0207   btagUParTAK4TauVJet = Var("?pt>=10 && bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:TauVsJet')>0?bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:TauVsJet'):-1",float,precision=12,doc="UnifiedParTAK4 tau vs. jet"),
0208   btagUParTAK4Ele     = Var("?pt>=10?bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:probele'):-1",float,precision=12,doc="UnifiedParT electron raw score"),
0209   btagUParTAK4Mu      = Var("?pt>=10?bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:probmu'):-1",float,precision=12,doc="UnifiedParT muon raw score"),
0210   btagUParTAK4probb   = Var("?pt>=10?bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:probb'):-1",float,precision=12,doc="UnifiedParT b raw score"),
0211   btagUParTAK4probbb  = Var("?pt>=10?bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:probbb'):-1",float,precision=12,doc="UnifiedParT bb raw score"),
0212   UParTAK4RegPtRawCorr = jetPuppiTable.variables.UParTAK4RegPtRawCorr,
0213   UParTAK4RegPtRawCorrNeutrino = jetPuppiTable.variables.UParTAK4RegPtRawCorr,
0214   UParTAK4RegPtRawRes = jetPuppiTable.variables.UParTAK4RegPtRawCorr,
0215   UParTAK4V1RegPtRawCorr = jetPuppiTable.variables.UParTAK4V1RegPtRawCorr,
0216   UParTAK4V1RegPtRawCorrNeutrino = jetPuppiTable.variables.UParTAK4V1RegPtRawCorrNeutrino,
0217   UParTAK4V1RegPtRawRes = jetPuppiTable.variables.UParTAK4V1RegPtRawRes,
0218 )
0219 PARTICLENETAK4VARS = cms.PSet(
0220   particleNetAK4_B = Var("?(pt>=15)?bDiscriminator('pfParticleNetAK4DiscriminatorsJetTags:BvsAll'):-1",float,doc="ParticleNetAK4 tagger b vs all (udsg, c) discriminator",precision=10),
0221   particleNetAK4_CvsL = Var("?(pt>=15)?bDiscriminator('pfParticleNetAK4DiscriminatorsJetTags:CvsL'):-1",float,doc="ParticleNetAK4 tagger c vs udsg discriminator",precision=10),
0222   particleNetAK4_CvsB = Var("?(pt>=15)?bDiscriminator('pfParticleNetAK4DiscriminatorsJetTags:CvsB'):-1",float,doc="ParticleNetAK4 tagger c vs b discriminator",precision=10),
0223   particleNetAK4_QvsG = Var("?(pt>=15)?bDiscriminator('pfParticleNetAK4DiscriminatorsJetTags:QvsG'):-1",float,doc="ParticleNetAK4 tagger uds vs g discriminator",precision=10),
0224   particleNetAK4_G = Var("?(pt>=15)?bDiscriminator('pfParticleNetAK4JetTags:probg'):-1",float,doc="ParticleNetAK4 tagger g raw score",precision=10),
0225   particleNetAK4_puIdDisc = Var("?(pt>=15)?1-bDiscriminator('pfParticleNetAK4JetTags:probpu'):-1",float,doc="ParticleNetAK4 tagger pileup jet discriminator",precision=10),
0226 )
0227
0228 CALOJETVARS = cms.PSet(P4Vars,
0229   area      = jetPuppiTable.variables.area,
0230   rawFactor = jetPuppiTable.variables.rawFactor,
0231   emf       = Var("emEnergyFraction()", float, doc = "electromagnetic energy fraction", precision = 10),
0232 )
0233
0234 #******************************************
0235 #
0236 #
0237 # Reco Jets related functions
0238 #
0239 #
0240 #******************************************
0241 def AddPileUpJetIDVars(proc, jetName="", jetSrc="", jetTableName="", jetTaskName=""):
0242   """
0243   Setup modules to calculate pileup jet ID input variables for PF jet
0244   """
0245
0246   #
0247   # Calculate pileup jet ID variables
0248   #
0249   puJetIdVarsCalculator = "puJetIdCalculator{}".format(jetName)
0250   setattr(proc, puJetIdVarsCalculator, pileupJetIdCalculator.clone(
0251       jets = jetSrc,
0252       vertexes  = "offlineSlimmedPrimaryVertices",
0253       inputIsCorrected = True,
0254       applyJec  = False,
0255       srcConstituentWeights = "packedpuppi" if "PUPPI" in jetName.upper() else ""
0256     )
0257   )
0258   getattr(proc,jetTaskName).add(getattr(proc, puJetIdVarsCalculator))
0259
0260   #
0261   # Get the variables
0262   #
0263   puJetIDVar = "puJetIDVar{}".format(jetName)
0264   setattr(proc, puJetIDVar, cms.EDProducer("PileupJetIDVarProducer",
0265       srcJet = cms.InputTag(jetSrc),
0266       srcPileupJetId = cms.InputTag(puJetIdVarsCalculator)
0267     )
0268   )
0269   getattr(proc,jetTaskName).add(getattr(proc, puJetIDVar))
0270
0271   #
0272   # Save variables as userFloats and userInts for each jet
0273   #
0274   patJetWithUserData = "{}WithUserData".format(jetSrc)
0275   getattr(proc,patJetWithUserData).userFloats.puId_dR2Mean  = cms.InputTag("{}:dR2Mean".format(puJetIDVar))
0276   getattr(proc,patJetWithUserData).userFloats.puId_majW     = cms.InputTag("{}:majW".format(puJetIDVar))
0277   getattr(proc,patJetWithUserData).userFloats.puId_minW     = cms.InputTag("{}:minW".format(puJetIDVar))
0278   getattr(proc,patJetWithUserData).userFloats.puId_frac01   = cms.InputTag("{}:frac01".format(puJetIDVar))
0279   getattr(proc,patJetWithUserData).userFloats.puId_frac02   = cms.InputTag("{}:frac02".format(puJetIDVar))
0280   getattr(proc,patJetWithUserData).userFloats.puId_frac03   = cms.InputTag("{}:frac03".format(puJetIDVar))
0281   getattr(proc,patJetWithUserData).userFloats.puId_frac04   = cms.InputTag("{}:frac04".format(puJetIDVar))
0282   getattr(proc,patJetWithUserData).userFloats.puId_ptD      = cms.InputTag("{}:ptD".format(puJetIDVar))
0283   getattr(proc,patJetWithUserData).userFloats.puId_beta     = cms.InputTag("{}:beta".format(puJetIDVar))
0284   getattr(proc,patJetWithUserData).userFloats.puId_pull     = cms.InputTag("{}:pull".format(puJetIDVar))
0285   getattr(proc,patJetWithUserData).userFloats.puId_jetR     = cms.InputTag("{}:jetR".format(puJetIDVar))
0286   getattr(proc,patJetWithUserData).userFloats.puId_jetRchg  = cms.InputTag("{}:jetRchg".format(puJetIDVar))
0287   getattr(proc,patJetWithUserData).userInts.puId_nCharged   = cms.InputTag("{}:nCharged".format(puJetIDVar))
0288
0289   #
0290   # Specfiy variables in the jet table to save in NanoAOD
0291   #
0292   getattr(proc,jetTableName).variables.puId_dR2Mean  = PUIDVARS.puId_dR2Mean
0293   getattr(proc,jetTableName).variables.puId_majW     = PUIDVARS.puId_majW
0294   getattr(proc,jetTableName).variables.puId_minW     = PUIDVARS.puId_minW
0295   getattr(proc,jetTableName).variables.puId_frac01   = PUIDVARS.puId_frac01
0296   getattr(proc,jetTableName).variables.puId_frac02   = PUIDVARS.puId_frac02
0297   getattr(proc,jetTableName).variables.puId_frac03   = PUIDVARS.puId_frac03
0298   getattr(proc,jetTableName).variables.puId_frac04   = PUIDVARS.puId_frac04
0299   getattr(proc,jetTableName).variables.puId_ptD      = PUIDVARS.puId_ptD
0300   getattr(proc,jetTableName).variables.puId_beta     = PUIDVARS.puId_beta
0301   getattr(proc,jetTableName).variables.puId_pull     = PUIDVARS.puId_pull
0302   getattr(proc,jetTableName).variables.puId_jetR     = PUIDVARS.puId_jetR
0303   getattr(proc,jetTableName).variables.puId_jetRchg  = PUIDVARS.puId_jetRchg
0304   getattr(proc,jetTableName).variables.puId_nCharged = PUIDVARS.puId_nCharged
0305
0306   return proc
0307
0308 def AddQGLTaggerVars(proc, jetName="", jetSrc="", jetTableName="", jetTaskName="", calculateQGLVars=False):
0309   """
0310   Schedule the QGTagger module to calculate input variables to the QG likelihood
0311   """
0312
0313   isPUPPIJet = True if "PUPPI" in jetName.upper() else False
0314
0315   QGLTagger="qgtagger{}".format(jetName)
0316   patJetWithUserData="{}WithUserData".format(jetSrc)
0317
0318   if calculateQGLVars:
0319     setattr(proc, QGLTagger, qgtagger.clone(
0320         srcJets = jetSrc,
0321         computeLikelihood = False,
0322       )
0323     )
0324     if isPUPPIJet:
0325       getattr(proc,QGLTagger).srcConstituentWeights = cms.InputTag("packedpuppi")
0326
0327   #
0328   # Save variables as userFloats and userInts for each jet
0329   #
0330   getattr(proc,patJetWithUserData).userFloats.qgl_axis2 = cms.InputTag(QGLTagger+":axis2")
0331   getattr(proc,patJetWithUserData).userFloats.qgl_ptD   = cms.InputTag(QGLTagger+":ptD")
0332   getattr(proc,patJetWithUserData).userInts.qgl_mult    = cms.InputTag(QGLTagger+":mult")
0333
0334   #
0335   # Specfiy variables in the jet table to save in NanoAOD
0336   #
0337   getattr(proc,jetTableName).variables.qgl_axis2 =  QGLVARS.qgl_axis2
0338   getattr(proc,jetTableName).variables.qgl_ptD   =  QGLVARS.qgl_ptD
0339   getattr(proc,jetTableName).variables.qgl_mult  =  QGLVARS.qgl_mult
0340
0341   if calculateQGLVars:
0342     getattr(proc,jetTaskName).add(getattr(proc, QGLTagger))
0343
0344   return proc
0345
0346 def AddBTaggingScores(proc, jetTableName=""):
0347   """
0348   Store b-tagging scores from various algortihm
0349   """
0350
0351   getattr(proc, jetTableName).variables.btagDeepFlavB   = DEEPJETVARS.btagDeepFlavB
0352   getattr(proc, jetTableName).variables.btagDeepFlavCvL = DEEPJETVARS.btagDeepFlavCvL
0353   getattr(proc, jetTableName).variables.btagDeepFlavCvB = DEEPJETVARS.btagDeepFlavCvB
0354
0355   run2_nanoAOD_ANY.toModify(
0356     getattr(proc, jetTableName).variables,
0357     btagCSVV2 = Var("bDiscriminator('pfCombinedInclusiveSecondaryVertexV2BJetTags')",float,doc=" pfCombinedInclusiveSecondaryVertexV2 b-tag discriminator (aka CSVV2)",precision=10),
0358     btagDeepB = Var("?(bDiscriminator('pfDeepCSVJetTags:probb')+bDiscriminator('pfDeepCSVJetTags:probbb'))>=0?bDiscriminator('pfDeepCSVJetTags:probb')+bDiscriminator('pfDeepCSVJetTags:probbb'):-1",float,doc="DeepCSV b+bb tag discriminator",precision=10),
0359     btagDeepCvL = Var("?bDiscriminator('pfDeepCSVJetTags:probc')>=0?bDiscriminator('pfDeepCSVJetTags:probc')/(bDiscriminator('pfDeepCSVJetTags:probc')+bDiscriminator('pfDeepCSVJetTags:probudsg')):-1", float,doc="DeepCSV c vs udsg discriminator",precision=10),
0360     btagDeepCvB = Var("?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)
0361   )
0362
0363   return proc
0364
0365 def AddDeepJetGluonLQuarkScores(proc, jetTableName=""):
0366   """
0367   Store DeepJet raw score in jetTable for gluon and light quark
0368   """
0369
0370   getattr(proc, jetTableName).variables.btagDeepFlavG   = DEEPJETVARS.btagDeepFlavG
0371   getattr(proc, jetTableName).variables.btagDeepFlavUDS = DEEPJETVARS.btagDeepFlavUDS
0372   getattr(proc, jetTableName).variables.btagDeepFlavQG  = DEEPJETVARS.btagDeepFlavQG
0373
0374   return proc
0375
0376 def AddRobustParTAK4Scores(proc, jetTableName=""):
0377   """
0378   Store RobustParTAK4 scores in jetTable
0379   """
0380
0381   getattr(proc, jetTableName).variables.btagRobustParTAK4B = ROBUSTPARTAK4VARS.btagRobustParTAK4B
0382   getattr(proc, jetTableName).variables.btagRobustParTAK4CvL = ROBUSTPARTAK4VARS.btagRobustParTAK4CvL
0383   getattr(proc, jetTableName).variables.btagRobustParTAK4CvB = ROBUSTPARTAK4VARS.btagRobustParTAK4CvB
0384
0385   return proc
0386
0387 def AddUnifiedParTAK4Scores(proc, jetTableName=""):
0388   """
0389   Store RobustParTAK4 scores in jetTable
0390   """
0391
0392   getattr(proc, jetTableName).variables.btagUParTAK4B = UNIFIEDPARTAK4VARS.btagUParTAK4B
0393   getattr(proc, jetTableName).variables.btagUParTAK4CvL = UNIFIEDPARTAK4VARS.btagUParTAK4CvL
0394   getattr(proc, jetTableName).variables.btagUParTAK4CvB = UNIFIEDPARTAK4VARS.btagUParTAK4CvB
0395   getattr(proc, jetTableName).variables.btagUParTAK4CvNotB = UNIFIEDPARTAK4VARS.btagUParTAK4CvNotB
0396   getattr(proc, jetTableName).variables.btagUParTAK4SvCB = UNIFIEDPARTAK4VARS.btagUParTAK4SvCB
0397   getattr(proc, jetTableName).variables.btagUParTAK4SvUDG = UNIFIEDPARTAK4VARS.btagUParTAK4SvUDG
0398   getattr(proc, jetTableName).variables.btagUParTAK4UDG = UNIFIEDPARTAK4VARS.btagUParTAK4UDG
0399   getattr(proc, jetTableName).variables.btagUParTAK4QvG = UNIFIEDPARTAK4VARS.btagUParTAK4QvG
0400   getattr(proc, jetTableName).variables.btagUParTAK4TauVJet = UNIFIEDPARTAK4VARS.btagUParTAK4TauVJet
0401   getattr(proc, jetTableName).variables.btagUParTAK4Ele = UNIFIEDPARTAK4VARS.btagUParTAK4Ele
0402   getattr(proc, jetTableName).variables.btagUParTAK4Mu = UNIFIEDPARTAK4VARS.btagUParTAK4Mu
0403   getattr(proc, jetTableName).variables.btagUParTAK4probb = UNIFIEDPARTAK4VARS.btagUParTAK4probb
0404   getattr(proc, jetTableName).variables.btagUParTAK4probbb = UNIFIEDPARTAK4VARS.btagUParTAK4probbb
0405   getattr(proc, jetTableName).variables.UParTAK4RegPtRawCorr = UNIFIEDPARTAK4VARS.UParTAK4RegPtRawCorr
0406   getattr(proc, jetTableName).variables.UParTAK4RegPtRawCorrNeutrino = UNIFIEDPARTAK4VARS.UParTAK4RegPtRawCorrNeutrino
0407   getattr(proc, jetTableName).variables.UParTAK4RegPtRawRes = UNIFIEDPARTAK4VARS.UParTAK4RegPtRawRes
0408   getattr(proc, jetTableName).variables.UParTAK4V1RegPtRawCorr = UNIFIEDPARTAK4VARS.UParTAK4RegPtRawCorr
0409   getattr(proc, jetTableName).variables.UParTAK4V1RegPtRawCorrNeutrino = UNIFIEDPARTAK4VARS.UParTAK4RegPtRawCorrNeutrino
0410   getattr(proc, jetTableName).variables.UParTAK4V1RegPtRawRes = UNIFIEDPARTAK4VARS.UParTAK4RegPtRawRes
0411
0412   return proc
0413
0414 def AddParticleNetAK4Scores(proc, jetTableName=""):
0415   """
0416   Store ParticleNetAK4 scores in jetTable
0417   """
0418
0419   getattr(proc, jetTableName).variables.particleNetAK4_B = PARTICLENETAK4VARS.particleNetAK4_B
0420   getattr(proc, jetTableName).variables.particleNetAK4_CvsL = PARTICLENETAK4VARS.particleNetAK4_CvsL
0421   getattr(proc, jetTableName).variables.particleNetAK4_CvsB = PARTICLENETAK4VARS.particleNetAK4_CvsB
0422   getattr(proc, jetTableName).variables.particleNetAK4_QvsG = PARTICLENETAK4VARS.particleNetAK4_QvsG
0423   getattr(proc, jetTableName).variables.particleNetAK4_G = PARTICLENETAK4VARS.particleNetAK4_G
0424   getattr(proc, jetTableName).variables.particleNetAK4_puIdDisc = PARTICLENETAK4VARS.particleNetAK4_puIdDisc
0425
0426   return proc
0427
0428 def AddNewPatJets(proc, recoJetInfo, runOnMC):
0429   """
0430   Add patJet into custom nanoAOD
0431   """
0432
0433   jetName = recoJetInfo.jetUpper
0434   payload = recoJetInfo.jetCorrPayload
0435   doPF    = recoJetInfo.doPF
0436   doCalo  = recoJetInfo.doCalo
0437   patJetFinalColl = recoJetInfo.patJetFinalCollection
0438
0439   nanoInfoForJet = nanoInfo_recojets[recoJetInfo.jet]
0440   jetTablePrefix = nanoInfoForJet["name"]
0441   jetTableDoc    = nanoInfoForJet["doc"]
0442   ptcut          = nanoInfoForJet["ptcut"] if "ptcut" in nanoInfoForJet else 8
0443   ptcutGenMatch  = nanoInfoForJet["ptcutGenMatch"] if "ptcutGenMatch" in nanoInfoForJet else 5
0444   doPUIDVar      = nanoInfoForJet["doPUIDVar"] if "doPUIDVar" in nanoInfoForJet else False
0445   doQGL          = nanoInfoForJet["doQGL"] if "doQGL" in nanoInfoForJet else False
0446   doBTag         = nanoInfoForJet["doBTag"] if "doBTag" in nanoInfoForJet else False
0447
0448   SavePatJets(proc,
0449     jetName, payload, patJetFinalColl, jetTablePrefix, jetTableDoc, doPF, doCalo,
0450     ptcut=ptcut, ptcutGenMatch=ptcutGenMatch, doPUIDVar=doPUIDVar, doQGL=doQGL, doBTag=doBTag, runOnMC=runOnMC
0451   )
0452
0453   return proc
0454
0455 def SavePatJets(proc, jetName, payload, patJetFinalColl, jetTablePrefix, jetTableDoc,
0456                 doPF, doCalo, ptcut, ptcutGenMatch, doPUIDVar=False, doQGL=False, doBTag=False, runOnMC=False):
0457   """
0458   Schedule modules for a given patJet collection and save its variables into custom NanoAOD
0459   """
0460
0461   #
0462   # Setup jet correction factors
0463   #
0464   jetCorrFactors = "jetCorrFactorsNano{}".format(jetName)
0465   setattr(proc, jetCorrFactors, jetCorrFactorsNano.clone(
0466       src = patJetFinalColl,
0467       payload = payload,
0468     )
0469   )
0470
0471   #
0472   # Update jets
0473   #
0474   srcJets = "updatedJets{}".format(jetName)
0475   setattr(proc, srcJets, updatedJets.clone(
0476       jetSource = patJetFinalColl,
0477       jetCorrFactorsSource = [jetCorrFactors],
0478     )
0479   )
0480
0481   #
0482   # Setup UserDataEmbedder
0483   #
0484   srcJetsWithUserData = "updatedJets{}WithUserData".format(jetName)
0485   setattr(proc, srcJetsWithUserData, cms.EDProducer("PATJetUserDataEmbedder",
0486       src = cms.InputTag(srcJets),
0487       userFloats = cms.PSet(),
0488       userInts = cms.PSet(),
0489     )
0490   )
0491
0492   #
0493   # Filter jets with pt cut
0494   #
0495   finalJetsCut = "(pt >= {ptcut:.0f})".format(ptcut=ptcut)
0496   genJetMatchCut = "genJetFwdRef().backRef().isNonnull() && genJetFwdRef().backRef().pt() > {ptcutGenMatch:.0f}".format(ptcutGenMatch=ptcutGenMatch)
0497   if runOnMC:
0498     finalJetsCut  += "|| ((pt < {ptcut:.0f}) && ({genJetMatchCut}))".format(ptcut=ptcut,genJetMatchCut=genJetMatchCut)
0499
0500   finalJetsForTable = "finalJets{}".format(jetName)
0501   setattr(proc, finalJetsForTable, finalJets.clone(
0502       src = srcJetsWithUserData,
0503       cut = finalJetsCut
0504     )
0505   )
0506
0507   #
0508   # Save jets in table
0509   #
0510   tableContent = PFJETVARS
0511   if doCalo:
0512     tableContent = CALOJETVARS
0513
0514   jetTableCutDefault = "" #Don't apply any cuts for the table.
0515
0516   jetTableDocDefault = jetTableDoc + " with JECs applied. Jets with pt >= {ptcut:.0f} GeV are stored.".format(ptcut=ptcut)
0517   if runOnMC:
0518     jetTableDocDefault += "For jets with pt < {ptcut:.0f} GeV, only those matched to gen jets are stored.".format(ptcut=ptcut)
0519
0520   if doCalo:
0521     jetTableDocDefault = jetTableDoc
0522
0523   jetTableName = "jet{}Table".format(jetName)
0524   setattr(proc,jetTableName, simplePATJetFlatTableProducer.clone(
0525       src = cms.InputTag(finalJetsForTable),
0526       cut = cms.string(jetTableCutDefault),
0527       name = cms.string(jetTablePrefix),
0528       doc  = cms.string(jetTableDocDefault),
0529       variables = cms.PSet(tableContent)
0530     )
0531   )
0532   getattr(proc,jetTableName).variables.pt.precision=12
0533   getattr(proc,jetTableName).variables.mass.precision=12
0534   getattr(proc,jetTableName).variables.rawFactor.precision=10
0535
0536   #
0537   # Save MC-only jet variables in table
0538   #
0539   jetMCTableName = "jet{}MCTable".format(jetName)
0540   setattr(proc, jetMCTableName, simplePATJetFlatTableProducer.clone(
0541       src = cms.InputTag(finalJetsForTable),
0542       cut = getattr(proc,jetTableName).cut,
0543       name = cms.string(jetTablePrefix),
0544       extension = cms.bool(True), # this is an extension table
0545       variables = cms.PSet(
0546         partonFlavour = PFJETVARS_MCOnly.partonFlavour,
0547         hadronFlavour = PFJETVARS_MCOnly.hadronFlavour,
0548         genJetIdx = Var("?{genJetMatchCut}?genJetFwdRef().backRef().key():-1".format(genJetMatchCut=genJetMatchCut), "int16", doc="index of matched gen jet")
0549       )
0550     )
0551   )
0552   #
0553   # Define the jet modules Task first
0554   #
0555   jetTaskName = "jet{}Task".format(jetName)
0556   setattr(proc, jetTaskName, cms.Task(
0557      getattr(proc,jetCorrFactors),
0558      getattr(proc,srcJets),
0559      getattr(proc,srcJetsWithUserData),
0560      getattr(proc,finalJetsForTable)
0561    )
0562   )
0563   proc.nanoTableTaskCommon.add(getattr(proc,jetTaskName))
0564
0565   #
0566   # Define the jet tables Task
0567   #
0568   jetTableTaskName = "jet{}TablesTask".format(jetName)
0569   setattr(proc, jetTableTaskName, cms.Task(getattr(proc,jetTableName)))
0570   proc.nanoTableTaskCommon.add(getattr(proc,jetTableTaskName))
0571
0572   jetMCTableTaskName = "jet{}MCTablesTask".format(jetName)
0573   setattr(proc, jetMCTableTaskName, cms.Task(getattr(proc,jetMCTableName)))
0574   if runOnMC:
0575     proc.nanoTableTaskFS.add(getattr(proc,jetMCTableTaskName))
0576
0577   #
0578   # Schedule plugins to calculate Jet ID, PileUp Jet ID input variables, and Quark-Gluon Likehood input variables.
0579   #
0580   if doPF:
0581     if doPUIDVar:
0582       proc = AddPileUpJetIDVars(proc, jetName=jetName, jetSrc=srcJets, jetTableName=jetTableName, jetTaskName=jetTaskName)
0583     if doQGL:
0584       proc = AddQGLTaggerVars(proc,jetName=jetName, jetSrc=srcJets, jetTableName=jetTableName, jetTaskName=jetTaskName, calculateQGLVars=True)
0585
0586   #
0587   # Save b-tagging algorithm scores. Should only be done for jet collection with b-tagging
0588   # calculated when reclustered or collection saved with b-tagging info in MiniAOD
0589   #
0590   if doBTag:
0591     AddBTaggingScores(proc,jetTableName=jetTableName)
0592     AddDeepJetGluonLQuarkScores(proc,jetTableName=jetTableName)
0593     AddParticleNetAK4Scores(proc,jetTableName=jetTableName)
0594     AddRobustParTAK4Scores(proc,jetTableName=jetTableName)
0595     AddUnifiedParTAK4Scores(proc,jetTableName=jetTableName)
0596
0597   return proc
0598
0599
0600 def ReclusterAK4PuppiJets(proc, recoJA, runOnMC):
0601   """
0602   Recluster AK4 Puppi jets and replace slimmedJetsPuppi
0603   that is used as default to save AK4 Puppi jets in NanoAODs.
0604   """
0605   print("custom_jme_cff::ReclusterAK4PuppiJets: Recluster AK4 PF Puppi jets")
0606
0607   #
0608   # Recluster AK4 Puppi jets
0609   #
0610   cfg = {
0611     "jet" : "ak4pfpuppi",
0612     "inputCollection" : "",
0613     "genJetsCollection": "AK4GenJetsNoNu", # Setup in ConfigureAK4GenJets() function
0614     "bTagDiscriminators": bTagDiscriminatorsForAK4,
0615     "minPtFastjet" : 0.,
0616   }
0617   recoJetInfo = recoJA.addRecoJetCollection(proc, **cfg)
0618
0619   jetName = recoJetInfo.jetUpper
0620   patJetFinalColl = recoJetInfo.patJetFinalCollection
0621
0622   #
0623   # Change the input jet source for jetCorrFactorsNano
0624   # and updatedJets
0625   #
0626   proc.jetPuppiCorrFactorsNano.src=patJetFinalColl
0627   proc.updatedJetsPuppi.jetSource=patJetFinalColl
0628
0629   #
0630   # Change pt cut
0631   #
0632   finalJetsPuppiCut = "(pt >= 8)"
0633   if runOnMC:
0634     genJetMatchCut = "genJetFwdRef().backRef().isNonnull() && genJetFwdRef().backRef().pt() > 5"
0635     finalJetsPuppiCut += " || ((pt < 8) && ({genJetMatchCut}))".format(genJetMatchCut=genJetMatchCut)
0636   proc.finalJetsPuppi.cut = finalJetsPuppiCut
0637
0638   #
0639   # Add a minimum pt cut for corrT1METJets.
0640   #
0641   proc.corrT1METJetPuppiTable.cut = "pt>=8 && pt<15 && abs(eta)<9.9"
0642
0643   #
0644   # Jet table
0645   #
0646   # For Run-2 eras, the main AK4 jet collection in NanoAOD is the CHS collection
0647   run2_nanoAOD_ANY.toModify(
0648     proc.jetTable, name = "Jet"
0649   ).toModify(
0650     # So need to change the table name for AK4 puppi here
0651     proc.jetPuppiTable,
0652     name = "JetPuppi",
0653     src = cms.InputTag("finalJetsPuppi")
0654   )
0655
0656   #
0657   # Jet table documentation
0658   #
0659   jetPuppiTableDoc = "AK4 PF Puppi jets with JECs applied. Jets with pt >= 8 GeV are stored."
0660   if runOnMC:
0661     jetPuppiTableDoc += "For jets with pt < 8 GeV, only those matched to AK4 Gen jets are stored."
0662   proc.jetPuppiTable.doc = jetPuppiTableDoc
0663
0664   proc.jetPuppiTable.variables.pt.precision = 12
0665   proc.jetPuppiTable.variables.mass.precision = 12
0666   proc.jetPuppiTable.variables.rawFactor.precision = 10
0667
0668   #
0669   # Add variables
0670   #
0671   proc.jetPuppiTable.variables.chHadMultiplicity = PFJETVARS.chHadMultiplicity
0672   proc.jetPuppiTable.variables.neHadMultiplicity = PFJETVARS.neHadMultiplicity
0673   proc.jetPuppiTable.variables.hfHadMultiplicity = PFJETVARS.hfHadMultiplicity
0674   proc.jetPuppiTable.variables.hfEMMultiplicity  = PFJETVARS.hfEMMultiplicity
0675   proc.jetPuppiTable.variables.muMultiplicity    = PFJETVARS.muMultiplicity
0676   proc.jetPuppiTable.variables.elMultiplicity    = PFJETVARS.elMultiplicity
0677   proc.jetPuppiTable.variables.phoMultiplicity   = PFJETVARS.phoMultiplicity
0678
0679   #
0680   #
0681   #
0682   from PhysicsTools.NanoAOD.patJetPFConstituentVarProducer_cfi import patJetPFConstituentVarProducer
0683   jetPFConstituentVarName = "jetPFConstituentVar{}".format(jetName)
0684   setattr(proc, jetPFConstituentVarName, patJetPFConstituentVarProducer.clone(
0685       jets = "updatedJetsPuppi",
0686       puppi_value_map = 'packedpuppi',
0687       fallback_puppi_weight = False,
0688     )
0689   )
0690   proc.jetPuppiTask.add(getattr(proc, jetPFConstituentVarName))
0691   proc.updatedJetsPuppiWithUserData.userFloats.leadConstNeHadEF = cms.InputTag(jetPFConstituentVarName+':leadConstNeHadEF')
0692   proc.updatedJetsPuppiWithUserData.userFloats.leadConstChHadEF = cms.InputTag(jetPFConstituentVarName+':leadConstChHadEF')
0693   proc.updatedJetsPuppiWithUserData.userFloats.leadConstPhotonEF = cms.InputTag(jetPFConstituentVarName+':leadConstPhotonEF')
0694   proc.updatedJetsPuppiWithUserData.userFloats.leadConstElectronEF = cms.InputTag(jetPFConstituentVarName+':leadConstElectronEF')
0695   proc.updatedJetsPuppiWithUserData.userFloats.leadConstMuonEF = cms.InputTag(jetPFConstituentVarName+':leadConstMuonEF')
0696   proc.updatedJetsPuppiWithUserData.userFloats.leadConstHFHADEF = cms.InputTag(jetPFConstituentVarName+':leadConstHFHADEF')
0697   proc.updatedJetsPuppiWithUserData.userFloats.leadConstHFEMEF = cms.InputTag(jetPFConstituentVarName+':leadConstHFEMEF')
0698   proc.updatedJetsPuppiWithUserData.userFloats.leadConstNeHadPuppiWeight = cms.InputTag(jetPFConstituentVarName+':leadConstNeHadPuppiWeight')
0699   proc.updatedJetsPuppiWithUserData.userFloats.leadConstChHadPuppiWeight = cms.InputTag(jetPFConstituentVarName+':leadConstChHadPuppiWeight')
0700   proc.updatedJetsPuppiWithUserData.userFloats.leadConstPhotonPuppiWeight = cms.InputTag(jetPFConstituentVarName+':leadConstPhotonPuppiWeight')
0701   proc.updatedJetsPuppiWithUserData.userFloats.leadConstElectronPuppiWeight = cms.InputTag(jetPFConstituentVarName+':leadConstElectronPuppiWeight')
0702   proc.updatedJetsPuppiWithUserData.userFloats.leadConstMuonPuppiWeight = cms.InputTag(jetPFConstituentVarName+':leadConstMuonPuppiWeight')
0703   proc.updatedJetsPuppiWithUserData.userFloats.leadConstHFHADPuppiWeight = cms.InputTag(jetPFConstituentVarName+':leadConstHFHADPuppiWeight')
0704   proc.updatedJetsPuppiWithUserData.userFloats.leadConstHFEMPuppiWeight = cms.InputTag(jetPFConstituentVarName+':leadConstHFEMPuppiWeight')
0705   proc.jetPuppiTable.variables.leadConstNeHadEF = Var("userFloat('leadConstNeHadEF')",float,doc="Leading PF neutral hadron constituent energy fraction w.r.t jet raw energy",precision=10)
0706   proc.jetPuppiTable.variables.leadConstChHadEF = Var("userFloat('leadConstChHadEF')",float,doc="Leading PF charged hadron constituent energy fraction w.r.t jet raw energy",precision=10)
0707   proc.jetPuppiTable.variables.leadConstPhotonEF = Var("userFloat('leadConstPhotonEF')",float,doc="Leading PF photon constituent energy fraction w.r.t jet raw energy",precision=10)
0708   proc.jetPuppiTable.variables.leadConstElectronEF = Var("userFloat('leadConstElectronEF')",float,doc="Leading PF electron constituent energy fraction w.r.t jet raw energy",precision=10)
0709   proc.jetPuppiTable.variables.leadConstMuonEF = Var("userFloat('leadConstMuonEF')",float,doc="Leading PF muon constituent energy fraction w.r.t jet raw energy",precision=10)
0710   proc.jetPuppiTable.variables.leadConstHFHADEF = Var("userFloat('leadConstHFHADEF')",float,doc="Leading PF HF HAD constituent energy fraction w.r.t jet raw energy",precision=10)
0711   proc.jetPuppiTable.variables.leadConstHFEMEF = Var("userFloat('leadConstHFEMEF')",float,doc="Leading PF HF EM constituent energy fraction w.r.t jet raw energy",precision=10)
0712   proc.jetPuppiTable.variables.leadConstNeHadPuppiWeight = Var("userFloat('leadConstNeHadPuppiWeight')",float,doc="Leading PF neutral hadron constituent puppi weight",precision=10)
0713   proc.jetPuppiTable.variables.leadConstChHadPuppiWeight = Var("userFloat('leadConstChHadPuppiWeight')",float,doc="Leading PF charged hadron constituent puppi weight",precision=10)
0714   proc.jetPuppiTable.variables.leadConstPhotonPuppiWeight = Var("userFloat('leadConstPhotonPuppiWeight')",float,doc="Leading PF photon constituent puppi weight",precision=10)
0715   proc.jetPuppiTable.variables.leadConstElectronPuppiWeight = Var("userFloat('leadConstElectronPuppiWeight')",float,doc="Leading PF electron constituent puppi weight",precision=10)
0716   proc.jetPuppiTable.variables.leadConstMuonPuppiWeight = Var("userFloat('leadConstMuonPuppiWeight')",float,doc="Leading PF muon constituent puppi weight",precision=10)
0717   proc.jetPuppiTable.variables.leadConstHFHADPuppiWeight = Var("userFloat('leadConstHFHADPuppiWeight')",float,doc="Leading PF HF HAD constituent puppi weight",precision=10)
0718   proc.jetPuppiTable.variables.leadConstHFEMPuppiWeight = Var("userFloat('leadConstHFEMPuppiWeight')",float,doc="Leading PF HF EM constituent puppi weight",precision=10)
0719
0720   #
0721   # Add Pileup Jet ID for Puppi jets
0722   #
0723   from RecoJets.JetProducers.PileupJetID_cfi import pileupJetIdPuppi
0724   pileupJetIdName = "pileupJetId{}".format(jetName)
0725   setattr(proc, pileupJetIdName, pileupJetIdPuppi.clone(
0726       jets = "updatedJetsPuppi",
0727       srcConstituentWeights = "packedpuppi",
0728       vertexes  = "offlineSlimmedPrimaryVertices",
0729       inputIsCorrected=True,
0730       applyJec=False
0731     )
0732   )
0733   proc.jetPuppiTask.add(getattr(proc, pileupJetIdName))
0734   proc.updatedJetsPuppiWithUserData.userFloats.puIdDisc = cms.InputTag(pileupJetIdName+':fullDiscriminant')
0735   proc.jetPuppiTable.variables.puIdDisc = Var("userFloat('puIdDisc')", float, doc="Pileup ID BDT discriminant with 133X Winter24 PuppiV18 training",precision=10)
0736
0737
0738   #
0739   # Add variables for pileup jet ID studies.
0740   #
0741   proc = AddPileUpJetIDVars(proc,
0742     jetName = jetName,
0743     jetSrc = "updatedJetsPuppi",
0744     jetTableName = "jetPuppiTable",
0745     jetTaskName = "jetPuppiTask"
0746   )
0747   #
0748   # Add variables for quark guon likelihood tagger studies.
0749   # Save variables as userFloats and userInts in each jet
0750   #
0751   proc = AddQGLTaggerVars(proc,
0752                           jetName = jetName,
0753                           jetSrc = "updatedJetsPuppi",
0754                           jetTableName = "jetPuppiTable",
0755                           jetTaskName = "jetPuppiTask",
0756                           calculateQGLVars=True
0757                         )
0758   #
0759   # Save DeepJet b-tagging and c-tagging variables
0760   #
0761   proc.jetPuppiTable.variables.btagDeepFlavB   = DEEPJETVARS.btagDeepFlavB
0762   proc.jetPuppiTable.variables.btagDeepFlavCvL = DEEPJETVARS.btagDeepFlavCvL
0763   proc.jetPuppiTable.variables.btagDeepFlavCvB = DEEPJETVARS.btagDeepFlavCvB
0764   #
0765   # Save DeepJet raw score for gluon and light quarks
0766   #
0767   proc.jetPuppiTable.variables.btagDeepFlavG   = DEEPJETVARS.btagDeepFlavG
0768   proc.jetPuppiTable.variables.btagDeepFlavUDS = DEEPJETVARS.btagDeepFlavUDS
0769   proc.jetPuppiTable.variables.btagDeepFlavQG  = DEEPJETVARS.btagDeepFlavQG
0770   #
0771   # Save UnifiedParTAK4 b-tagging and c-tagging variables
0772   #
0773   proc.jetPuppiTable.variables.btagUParTAK4B = UNIFIEDPARTAK4VARS.btagUParTAK4B
0774   proc.jetPuppiTable.variables.btagUParTAK4CvL = UNIFIEDPARTAK4VARS.btagUParTAK4CvL
0775   proc.jetPuppiTable.variables.btagUParTAK4CvB = UNIFIEDPARTAK4VARS.btagUParTAK4CvB
0776   proc.jetPuppiTable.variables.btagUParTAK4CvNotB = UNIFIEDPARTAK4VARS.btagUParTAK4CvNotB
0777   proc.jetPuppiTable.variables.btagUParTAK4SvCB = UNIFIEDPARTAK4VARS.btagUParTAK4SvCB
0778   proc.jetPuppiTable.variables.btagUParTAK4SvUDG = UNIFIEDPARTAK4VARS.btagUParTAK4SvUDG
0779   proc.jetPuppiTable.variables.btagUParTAK4UDG = UNIFIEDPARTAK4VARS.btagUParTAK4UDG
0780   proc.jetPuppiTable.variables.btagUParTAK4QvG = UNIFIEDPARTAK4VARS.btagUParTAK4QvG
0781   proc.jetPuppiTable.variables.btagUParTAK4TauVJet = UNIFIEDPARTAK4VARS.btagUParTAK4TauVJet
0782   proc.jetPuppiTable.variables.btagUParTAK4Ele = UNIFIEDPARTAK4VARS.btagUParTAK4Ele
0783   proc.jetPuppiTable.variables.btagUParTAK4Mu = UNIFIEDPARTAK4VARS.btagUParTAK4Mu
0784   proc.jetPuppiTable.variables.btagUParTAK4probb = UNIFIEDPARTAK4VARS.btagUParTAK4probb
0785   proc.jetPuppiTable.variables.btagUParTAK4probbb = UNIFIEDPARTAK4VARS.btagUParTAK4probbb
0786   proc.jetPuppiTable.variables.UParTAK4RegPtRawCorr = UNIFIEDPARTAK4VARS.UParTAK4RegPtRawCorr
0787   proc.jetPuppiTable.variables.UParTAK4RegPtRawCorrNeutrino = UNIFIEDPARTAK4VARS.UParTAK4RegPtRawCorrNeutrino
0788   proc.jetPuppiTable.variables.UParTAK4RegPtRawRes = UNIFIEDPARTAK4VARS.UParTAK4RegPtRawRes
0789   proc.jetPuppiTable.variables.UParTAK4V1RegPtRawCorr = UNIFIEDPARTAK4VARS.UParTAK4V1RegPtRawCorr
0790   proc.jetPuppiTable.variables.UParTAK4V1RegPtRawCorrNeutrino = UNIFIEDPARTAK4VARS.UParTAK4V1RegPtRawCorrNeutrino
0791   proc.jetPuppiTable.variables.UParTAK4V1RegPtRawRes = UNIFIEDPARTAK4VARS.UParTAK4V1RegPtRawRes
0792
0793   #
0794   # For Run-2 eras, don't need to save the low pt AK4 Puppi jet table for MET
0795   #
0796   run2_nanoAOD_ANY.toReplaceWith(
0797     proc.jetPuppiForMETTask,
0798     proc.jetPuppiForMETTask.copyAndExclude([proc.corrT1METJetPuppiTable])
0799   )
0800
0801   #
0802   # Save MC-only jet variables in jet table
0803   #
0804   if runOnMC:
0805     jetMCTableName = "jet{}MCTable".format(jetName)
0806     setattr(proc, jetMCTableName, proc.jetMCTable.clone(
0807         src = proc.jetPuppiTable.src,
0808         name = proc.jetPuppiTable.name
0809       )
0810     )
0811     getattr(proc,jetMCTableName).variables.genJetIdx = Var("?{genJetMatchCut}?genJetFwdRef().backRef().key():-1".format(genJetMatchCut=genJetMatchCut), "int16", doc="index of matched gen jet")
0812
0813     jetMCTableTaskName = "jet{}MCTablesTask".format(jetName)
0814     setattr(proc, jetMCTableTaskName, cms.Task(getattr(proc,jetMCTableName)))
0815
0816     run2_nanoAOD_ANY.toReplaceWith(
0817       proc.nanoTableTaskFS,
0818       proc.nanoTableTaskFS.copyAndAdd( getattr(proc,jetMCTableTaskName))
0819     )
0820
0821   return proc
0822
0823 def ReclusterAK4CHSJets(proc, recoJA, runOnMC):
0824   """
0825   Recluster AK4 CHS jets and replace slimmedJets that is used as default to
0826   save AK4 CHS jets in NanoAODs (for Run-2).
0827   """
0828   print("custom_jme_cff::ReclusterAK4CHSJets: Recluster AK4 PF CHS jets")
0829
0830   from RecoBTag.ONNXRuntime.pfParticleNetFromMiniAODAK4_cff import _pfParticleNetFromMiniAODAK4CHSCentralJetTagsAll
0831   from RecoBTag.ONNXRuntime.pfParticleNetFromMiniAODAK4_cff import _pfParticleNetFromMiniAODAK4CHSForwardJetTagsAll
0832
0833   bTagDiscriminatorsForAK4CHS = cms.PSet(
0834     foo = cms.vstring(_pfParticleNetFromMiniAODAK4CHSCentralJetTagsAll+_pfParticleNetFromMiniAODAK4CHSForwardJetTagsAll)
0835   )
0836   bTagDiscriminatorsForAK4CHS = bTagDiscriminatorsForAK4CHS.foo.value()
0837
0838   #
0839   # Recluster AK4 CHS jets
0840   #
0841   cfg = {
0842     "jet" : "ak4pfchs",
0843     "inputCollection" : "",
0844     "genJetsCollection": "AK4GenJetsNoNu", # Setup in ConfigureAK4GenJets() function
0845     "bTagDiscriminators": bTagDiscriminatorsForAK4CHS,
0846     "minPtFastjet" : 0.,
0847   }
0848   recoJetInfo = recoJA.addRecoJetCollection(proc, **cfg)
0849
0850   jetName = recoJetInfo.jetUpper
0851   patJetFinalColl = recoJetInfo.patJetFinalCollection
0852
0853   #
0854   # Change the input jet source for jetCorrFactorsNano
0855   # and updatedJets
0856   #
0857   proc.jetCorrFactorsNano.src=patJetFinalColl
0858   proc.updatedJets.jetSource=patJetFinalColl
0859
0860   #
0861   # Change pt cut
0862   #
0863   finalJetsCut = "(pt >= 10)"
0864   genJetMatchCut = "genJetFwdRef().backRef().isNonnull() && genJetFwdRef().backRef().pt() > 5"
0865   if runOnMC:
0866     finalJetsCut += " || ((pt < 10) && ({genJetMatchCut}))".format(genJetMatchCut=genJetMatchCut)
0867
0868   proc.finalJets.cut = finalJetsCut
0869   #
0870   # Add a minimum pt cut for corrT1METJets.
0871   #
0872   proc.corrT1METJetTable.cut = "pt>=10 && pt<15 && abs(eta)<9.9"
0873
0874   #
0875   # Jet table cut
0876   #
0877   jetTableCut = "" # must not have any cut at the jetTable for AK4 CHS as it has been cross-cleaned
0878   proc.jetTable.src   = cms.InputTag("finalJets")
0879   proc.jetTable.cut   = jetTableCut
0880   proc.jetMCTable.cut = jetTableCut
0881   proc.jetTable.name  = "JetCHS"
0882
0883   #
0884   # For Run-2 eras, the main AK4 jet collection in NanoAOD is the CHS collection
0885   #
0886   run2_nanoAOD_ANY.toModify(
0887     proc.jetTable,
0888     src = cms.InputTag("linkedObjects","jets"),
0889     name = "Jet"
0890   )
0891
0892   #
0893   # Jet table documentation
0894   #
0895   jetTableDoc = "AK4 PF CHS jets with JECs applied. Jets with pt >= 10 GeV are stored."
0896   if runOnMC:
0897     jetTableDoc += "For jets with pt < 10 GeV, only those matched to AK4 Gen jets are stored."
0898   proc.jetTable.doc   = jetTableDoc
0899
0900   proc.jetTable.variables.pt = 12
0901   proc.jetTable.variables.mass = 12
0902   proc.jetTable.variables.rawFactor.precision = 10
0903
0904   #
0905   # Add variables
0906   #
0907   proc.jetTable.variables.chHadMultiplicity = PFJETVARS.chHadMultiplicity
0908   proc.jetTable.variables.neHadMultiplicity = PFJETVARS.neHadMultiplicity
0909   proc.jetTable.variables.hfHadMultiplicity = PFJETVARS.hfHadMultiplicity
0910   proc.jetTable.variables.hfEMMultiplicity  = PFJETVARS.hfEMMultiplicity
0911   proc.jetTable.variables.muMultiplicity    = PFJETVARS.muMultiplicity
0912   proc.jetTable.variables.elMultiplicity    = PFJETVARS.elMultiplicity
0913   proc.jetTable.variables.phoMultiplicity   = PFJETVARS.phoMultiplicity
0914
0915   #
0916   # Add charged energy fraction from other primary vertices
0917   #
0918   proc.updatedJetsWithUserData.userFloats.chFPV1EF = cms.InputTag("jercVars:chargedFromPV1EnergyFraction")
0919   proc.updatedJetsWithUserData.userFloats.chFPV2EF = cms.InputTag("jercVars:chargedFromPV2EnergyFraction")
0920   proc.updatedJetsWithUserData.userFloats.chFPV3EF = cms.InputTag("jercVars:chargedFromPV3EnergyFraction")
0921   proc.jetTable.variables.chFPV1EF = Var("userFloat('chFPV1EF')", float, doc="charged fromPV==1 Energy Fraction (component of the total charged Energy Fraction).", precision= 6)
0922   proc.jetTable.variables.chFPV2EF = Var("userFloat('chFPV2EF')", float, doc="charged fromPV==2 Energy Fraction (component of the total charged Energy Fraction).", precision= 6)
0923   proc.jetTable.variables.chFPV3EF = Var("userFloat('chFPV3EF')", float, doc="charged fromPV==3 Energy Fraction (component of the total charged Energy Fraction).", precision= 6)
0924
0925   #
0926   # Remove these tagger branches since for CHS, we just want to store
0927   # one tagger only.
0928   #
0929   for varNames in proc.jetTable.variables.parameterNames_():
0930     if "btagDeepFlav" in varNames or "btagRobustParT" in varNames or "btagUParT" in varNames:
0931       delattr(proc.jetTable.variables, varNames)
0932     if "UParTAK4Reg" in varNames:
0933       delattr(proc.jetTable.variables, varNames)
0934     if "svIdx" in varNames or "muonIdx" in varNames:
0935       delattr(proc.jetTable.variables, varNames)
0936     if "nElectrons" in varNames or "nMuons" in varNames or "nSVs" in varNames:
0937       delattr(proc.jetTable.variables, varNames)
0938
0939   proc.jetTable.variables.btagPNetB = Var("?pt>15 && bDiscriminator('pfParticleNetFromMiniAODAK4CHSCentralDiscriminatorsJetTags:BvsAll')>0?bDiscriminator('pfParticleNetFromMiniAODAK4CHSCentralDiscriminatorsJetTags:BvsAll'):-1",float,precision=10,doc="ParticleNet b vs. udscg")
0940   proc.jetTable.variables.btagPNetCvL = Var("?pt>15 && bDiscriminator('pfParticleNetFromMiniAODAK4CHSCentralDiscriminatorsJetTags:CvsL')>0?bDiscriminator('pfParticleNetFromMiniAODAK4CHSCentralDiscriminatorsJetTags:CvsL'):-1",float,precision=10,doc="ParticleNet c vs. udsg")
0941   proc.jetTable.variables.btagPNetCvB = Var("?pt>15 && bDiscriminator('pfParticleNetFromMiniAODAK4CHSCentralDiscriminatorsJetTags:CvsB')>0?bDiscriminator('pfParticleNetFromMiniAODAK4CHSCentralDiscriminatorsJetTags:CvsB'):-1",float,precision=10,doc="ParticleNet c vs. b")
0942   proc.jetTable.variables.btagPNetQvG = Var("?pt>15 && abs(eta())<2.5?bDiscriminator('pfParticleNetFromMiniAODAK4CHSCentralDiscriminatorsJetTags:QvsG'):bDiscriminator('pfParticleNetFromMiniAODAK4CHSForwardDiscriminatorsJetTags:QvsG')",float,precision=10,doc="ParticleNet q (udsbc) vs. g")
0943   proc.jetTable.variables.btagPNetTauVJet = Var("?pt>15 && bDiscriminator('pfParticleNetFromMiniAODAK4CHSCentralDiscriminatorsJetTags:TauVsJet')>0?bDiscriminator('pfParticleNetFromMiniAODAK4CHSCentralDiscriminatorsJetTags:TauVsJet'):-1",float,precision=10,doc="ParticleNet tau vs. jet")
0944   proc.jetTable.variables.PNetRegPtRawCorr = Var("?abs(eta())<2.5?bDiscriminator('pfParticleNetFromMiniAODAK4CHSCentralJetTags:ptcorr'):bDiscriminator('pfParticleNetFromMiniAODAK4CHSForwardJetTags:ptcorr')",float,precision=10,doc="ParticleNet universal flavor-aware visible pT regression (no neutrinos), correction relative to raw jet pT")
0945   proc.jetTable.variables.PNetRegPtRawCorrNeutrino = Var("?abs(eta())<2.5?bDiscriminator('pfParticleNetFromMiniAODAK4CHSCentralJetTags:ptnu'):bDiscriminator('pfParticleNetFromMiniAODAK4CHSForwardJetTags:ptnu')",float,precision=10,doc="ParticleNet universal flavor-aware pT regression neutrino correction, relative to visible. To apply full regression, multiply raw jet pT by both PNetRegPtRawCorr and PNetRegPtRawCorrNeutrino.")
0946   proc.jetTable.variables.PNetRegPtRawRes = Var("?abs(eta())<2.5?0.5*(bDiscriminator('pfParticleNetFromMiniAODAK4CHSCentralJetTags:ptreshigh')-bDiscriminator('pfParticleNetFromMiniAODAK4CHSCentralJetTags:ptreslow')):0.5*(bDiscriminator('pfParticleNetFromMiniAODAK4CHSForwardJetTags:ptreshigh')-bDiscriminator('pfParticleNetFromMiniAODAK4CHSForwardJetTags:ptreslow'))",float,precision=10,doc="ParticleNet universal flavor-aware jet pT resolution estimator, (q84 - q16)/2")
0947
0948   #Adding hf shower shape producer to the jet sequence. By default this producer is not automatically rerun at the NANOAOD step
0949   #The following lines make sure it is.
0950   hfJetShowerShapeforCustomNanoAOD = "hfJetShowerShapeforCustomNanoAOD"
0951   setattr(proc, hfJetShowerShapeforCustomNanoAOD, hfJetShowerShape.clone(jets="updatedJets", vertices="offlineSlimmedPrimaryVertices") )
0952   proc.jetUserDataTask.add(getattr(proc, hfJetShowerShapeforCustomNanoAOD))
0953   proc.updatedJetsWithUserData.userFloats.hfsigmaEtaEta = cms.InputTag('hfJetShowerShapeforCustomNanoAOD:sigmaEtaEta')
0954   proc.updatedJetsWithUserData.userFloats.hfsigmaPhiPhi = cms.InputTag('hfJetShowerShapeforCustomNanoAOD:sigmaPhiPhi')
0955   proc.updatedJetsWithUserData.userInts.hfcentralEtaStripSize = cms.InputTag('hfJetShowerShapeforCustomNanoAOD:centralEtaStripSize')
0956   proc.updatedJetsWithUserData.userInts.hfadjacentEtaStripsSize = cms.InputTag('hfJetShowerShapeforCustomNanoAOD:adjacentEtaStripsSize')
0957   proc.jetTable.variables.hfsigmaEtaEta = Var("userFloat('hfsigmaEtaEta')",float,doc="sigmaEtaEta for HF jets (noise discriminating variable)",precision=10)
0958   proc.jetTable.variables.hfsigmaPhiPhi = Var("userFloat('hfsigmaPhiPhi')",float,doc="sigmaPhiPhi for HF jets (noise discriminating variable)",precision=10)
0959   proc.jetTable.variables.hfcentralEtaStripSize = Var("userInt('hfcentralEtaStripSize')", int, doc="eta size of the central tower strip in HF (noise discriminating variable) ")
0960   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) ")
0961
0962   #
0963   # Since AK4 Puppi jet is the main AK4 jet collection for Run-3, disable
0964   # b-jets/c-jets NN-based mass regression for AK4 CHS.
0965   #
0966   (~run2_nanoAOD_ANY).toReplaceWith(
0967     proc.jetUserDataTask,
0968     proc.jetUserDataTask.copyAndExclude([proc.bJetVars])
0969   ).toReplaceWith(
0970     proc.jetTablesTask,
0971     proc.jetTablesTask.copyAndExclude([proc.bjetNN, proc.cjetNN])
0972   ).toModify(proc.updatedJetsWithUserData.userFloats,
0973     leadTrackPt = None,
0974     leptonPtRelv0 = None,
0975     leptonPtRelInvv0 = None,
0976     leptonDeltaR = None,
0977     vtxPt = None,
0978     vtxMass = None,
0979     vtx3dL = None,
0980     vtx3deL = None,
0981     ptD = None,
0982   ).toModify(
0983     proc.updatedJetsWithUserData.userInts,
0984     vtxNtrk = None,
0985     leptonPdgId = None
0986   ).toModify(
0987     proc.jetTable, externalVariables = cms.PSet()
0988   ).toReplaceWith(
0989   #
0990   # For Run-3, don't need to save the low pt AK4 CHS jet table for MET
0991   #
0992     proc.jetForMETTask,
0993     proc.jetForMETTask.copyAndExclude([proc.corrT1METJetTable])
0994   )
0995
0996   #
0997   # Save MC-only jet variables in jet table
0998   #
0999   if runOnMC:
1000     jetMCTableName = "jet{}MCTable".format(jetName)
1001     setattr(proc, jetMCTableName, proc.jetMCTable.clone(
1002         src = proc.jetTable.src,
1003         name = proc.jetTable.name
1004       )
1005     )
1006     getattr(proc,jetMCTableName).variables.genJetIdx = Var("?{genJetMatchCut}?genJetFwdRef().backRef().key():-1".format(genJetMatchCut=genJetMatchCut), "int16", doc="index of matched gen jet")
1007
1008     jetMCTableTaskName = "jet{}MCTablesTask".format(jetName)
1009     setattr(proc, jetMCTableTaskName, cms.Task(getattr(proc,jetMCTableName)))
1010
1011     (~run2_nanoAOD_ANY).toReplaceWith(
1012       proc.nanoTableTaskFS,
1013       proc.nanoTableTaskFS.copyAndAdd(getattr(proc,jetMCTableTaskName))
1014     )
1015
1016   return proc
1017
1018 def AddNewAK8PuppiJetsForJEC(proc, recoJA, runOnMC):
1019   """
1020   Store a separate AK8 Puppi jet collection for JEC studies.
1021   Only minimal info are stored
1022   """
1023   print("custom_jme_cff::AddNewAK8PuppiJetsForJEC: Make a new AK8 PF Puppi jet collection for JEC studies")
1024
1025   #
1026   # Recluster AK8 Puppi jets
1027   #
1028   cfg = {
1029     "jet" : "ak8pfpuppi",
1030     "inputCollection" : "",
1031     "genJetsCollection": "AK8GenJetsNoNu",
1032     "minPtFastjet" : 0., # Remove any pt threshold at the jet clustering stage.
1033   }
1034   recoJetInfo = recoJA.addRecoJetCollection(proc, **cfg)
1035
1036   jetName = recoJetInfo.jetUpper
1037   payload = recoJetInfo.jetCorrPayload
1038
1039   patJetFinalColl = recoJetInfo.patJetFinalCollection
1040   jetTablePrefix  = "FatJetForJEC"
1041   jetTableDoc     = "AK8 PF Puppi jets with JECs applied. Reclustered for JEC studies so only minimal info stored."
1042   ptcut           = 15
1043   ptcutGenMatch   = 10
1044
1045   SavePatJets(proc,
1046     jetName, payload, patJetFinalColl, jetTablePrefix, jetTableDoc, doPF=True,
1047     doCalo=False, ptcut=ptcut, ptcutGenMatch=ptcutGenMatch, doPUIDVar=False, doQGL=False, doBTag=False, runOnMC=runOnMC
1048   )
1049
1050   return proc
1051
1052 def AddNewAK8CHSJets(proc, recoJA, runOnMC):
1053   """
1054   Store an AK8 CHS jet collection for JEC studies.
1055   """
1056   print("custom_jme_cff::AddNewAK8CHSJets: Make a new AK8 PF CHS jet collection for JEC studies")
1057
1058   #
1059   # Recluster AK8 CHS jets
1060   #
1061   cfg = {
1062     "jet" : "ak8pfchs",
1063     "inputCollection" : "",
1064     "genJetsCollection": "AK8GenJetsNoNu",
1065     "minPtFastjet" : 0., # Remove any pt threshold at the jet clustering stage.
1066   }
1067   recoJetInfo = recoJA.addRecoJetCollection(proc, **cfg)
1068
1069   jetName = recoJetInfo.jetUpper
1070   payload = recoJetInfo.jetCorrPayload
1071
1072   patJetFinalColl = recoJetInfo.patJetFinalCollection
1073   jetTablePrefix  = "FatJetCHS"
1074   jetTableDoc     = "AK8 PF CHS jets with JECs applied. Reclustered for JEC studies so only minimal info stored."
1075   ptcut           = 15
1076   ptcutGenMatch   = 10
1077
1078   SavePatJets(proc,
1079     jetName, payload, patJetFinalColl, jetTablePrefix, jetTableDoc, doPF=True,
1080     doCalo=False, ptcut=ptcut, ptcutGenMatch=ptcutGenMatch, doPUIDVar=False, doQGL=False, doBTag=False, runOnMC=runOnMC
1081   )
1082
1083   return proc
1084
1085 def AddVariablesForAK8PuppiJets(proc):
1086   """
1087   Add more variables for AK8 PFPUPPI jets
1088   """
1089
1090   proc.fatJetTable.variables.rawFactor.precision=10
1091
1092   #
1093   #  These variables are not stored for AK8PFPUPPI (slimmedJetsAK8)
1094   #  in MiniAOD if their pt < 170 GeV. Hence the conditional fill.
1095   #
1096   proc.fatJetTable.variables.chHEF  = Var("?isPFJet()?chargedHadronEnergyFraction():-1",float,doc="charged Hadron Energy Fraction",precision=10)
1097   proc.fatJetTable.variables.neHEF  = Var("?isPFJet()?neutralHadronEnergyFraction():-1",float,doc="neutral Hadron Energy Fraction",precision=10)
1098   proc.fatJetTable.variables.chEmEF = Var("?isPFJet()?chargedEmEnergyFraction():-1",float,doc="charged Electromagnetic Energy Fraction",precision=10)
1099   proc.fatJetTable.variables.neEmEF = Var("?isPFJet()?neutralEmEnergyFraction():-1",float,doc="neutral Electromagnetic Energy Fraction",precision=10)
1100   proc.fatJetTable.variables.muEF   = Var("?isPFJet()?muonEnergyFraction():-1",float,doc="muon Energy Fraction",precision=10)
1101   proc.fatJetTable.variables.hfHEF  = Var("?isPFJet()?HFHadronEnergyFraction():-1", float,doc="energy fraction in forward hadronic calorimeter",precision=10)
1102   proc.fatJetTable.variables.hfEmEF = Var("?isPFJet()?HFEMEnergyFraction():-1",float,doc="energy fraction in forward EM calorimeter",precision=10)
1103   proc.fatJetTable.variables.chHadMultiplicity = Var("?isPFJet()?chargedHadronMultiplicity():-1","int16", doc="(Puppi-weighted) number of charged hadrons in the jet")
1104   proc.fatJetTable.variables.neHadMultiplicity = Var("?isPFJet()?neutralHadronMultiplicity():-1","int16", doc="(Puppi-weighted) number of neutral hadrons in the jet")
1105   proc.fatJetTable.variables.hfHadMultiplicity = Var("?isPFJet()?HFHadronMultiplicity():-1", "int16", doc="(Puppi-weighted) number of HF Hadrons in the jet")
1106   proc.fatJetTable.variables.hfEMMultiplicity  = Var("?isPFJet()?HFEMMultiplicity():-1", "int16", doc="(Puppi-weighted) number of HF EMs in the jet")
1107   proc.fatJetTable.variables.muMultiplicity    = Var("?isPFJet()?muonMultiplicity():-1", "int16", doc="(Puppi-weighted) number of muons in the jet")
1108   proc.fatJetTable.variables.elMultiplicity    = Var("?isPFJet()?electronMultiplicity():-1", "int16", doc="(Puppi-weighted) number of electrons in the jet")
1109   proc.fatJetTable.variables.phoMultiplicity   = Var("?isPFJet()?photonMultiplicity():-1", "int16", doc="(Puppi-weighted) number of photons in the jet")
1110
1111   return proc
1112 #******************************************
1113 #
1114 #
1115 # Gen Jets related functions
1116 #
1117 #
1118 #******************************************
1119 def AddNewGenJets(proc, genJetInfo):
1120   """
1121   Add genJet into custom nanoAOD
1122   """
1123
1124   genJetName         = genJetInfo.jetUpper
1125   genJetAlgo         = genJetInfo.jetAlgo
1126   genJetSize         = genJetInfo.jetSize
1127   genJetSizeNr       = genJetInfo.jetSizeNr
1128   genJetFinalColl    = "{}{}{}".format(genJetAlgo.upper(), genJetSize, "GenJetsNoNu")
1129   genJetTablePrefix  = nanoInfo_genjets[genJetInfo.jet]["name"]
1130   genJetTableDoc     = nanoInfo_genjets[genJetInfo.jet]["doc"]
1131
1132   SaveGenJets(proc, genJetName, genJetAlgo, genJetSizeNr, genJetFinalColl, genJetTablePrefix, genJetTableDoc)
1133
1134   return proc
1135
1136 def SaveGenJets(proc, genJetName, genJetAlgo, genJetSizeNr, genJetFinalColl, genJetTablePrefix, genJetTableDoc, genJetTableCut=""):
1137   """
1138   Schedule modules for a given genJet collection and save its variables into custom NanoAOD
1139   """
1140
1141   genJetTableName = "jet{}Table".format(genJetName)
1142   setattr(proc, genJetTableName, genJetTable.clone(
1143       src       = genJetFinalColl,
1144       cut       = genJetTableCut,
1145       name      = genJetTablePrefix,
1146       doc       = genJetTableDoc,
1147       variables = GENJETVARS
1148     )
1149   )
1150
1151   genJetFlavourAssociationName = "genJet{}FlavourAssociation".format(genJetName)
1152   setattr(proc, genJetFlavourAssociationName, genJetFlavourAssociation.clone(
1153       jets           = getattr(proc,genJetTableName).src,
1154       jetAlgorithm   = supportedJetAlgos[genJetAlgo],
1155       rParam         = genJetSizeNr,
1156     )
1157   )
1158
1159   genJetFlavourTableName = "genJet{}FlavourTable".format(genJetName)
1160   setattr(proc, genJetFlavourTableName, genJetFlavourTable.clone(
1161       name            = getattr(proc,genJetTableName).name,
1162       src             = getattr(proc,genJetTableName).src,
1163       cut             = getattr(proc,genJetTableName).cut,
1164       jetFlavourInfos = genJetFlavourAssociationName,
1165     )
1166   )
1167
1168   genJetTaskName = "genJet{}Task".format(genJetName)
1169   setattr(proc, genJetTaskName, cms.Task(
1170       getattr(proc,genJetTableName),
1171       getattr(proc,genJetFlavourAssociationName),
1172       getattr(proc,genJetFlavourTableName)
1173     )
1174   )
1175   proc.jetMCTask.add(getattr(proc,genJetTaskName))
1176
1177   return proc
1178
1179 def ConfigureAK4GenJets(proc, genJA):
1180   """
1181   Configure AK4GenJets to be saved in JMENano.
1182   Two options:
1183   1) Recluster AK4 Gen jets and replace slimmedGenJets that is
1184   used as default to save AK4 Gen jets in NanoAODs.
1185   2) Stick to slimmedGenJets and remove any AK4 Gen jet pt cut for genJetTable.
1186   """
1187
1188   #
1189   # Switch for AK4 genjet reclustering. depending on era
1190   #
1191   jmeNano_genjetRecluster_switch = cms.PSet(
1192     doAK4 = cms.untracked.bool(False),
1193   )
1194   run3_nanoAOD_pre142X.toModify(jmeNano_genjetRecluster_switch,
1195     doAK4 = True
1196   )
1197
1198   #
1199   # Recluster AK4 Gen jet
1200   #
1201   if jmeNano_genjetRecluster_switch.doAK4:
1202     print("custom_jme_cff::ReclusterAK4GenJets: Recluster AK4 Gen jets")
1203     cfg = {
1204       "jet" : "ak4gen",
1205       "minPtFastjet" : 5.
1206     }
1207     genJetInfo = genJA.addGenJetCollection(proc, **cfg)
1208     genJetName      = genJetInfo.jetUpper
1209     genJetAlgo      = genJetInfo.jetAlgo
1210     genJetSize      = genJetInfo.jetSize
1211     genJetSizeNr    = genJetInfo.jetSizeNr
1212     selectedGenJets = "{}{}{}".format(genJetAlgo.upper(), genJetSize, "GenJetsNoNu")
1213     #
1214     # Change jet source to the newly clustered jet collection. Set very low pt cut for jets
1215     # to be stored in the GenJet Table
1216     #
1217     proc.genJetTable.src = selectedGenJets
1218     proc.genJetTable.cut = "" # No cut specified here. Save all gen jets after clustering
1219     proc.genJetTable.doc = "AK4 Gen jets (made with visible genparticles) with pt > 5 GeV"
1220     proc.genJetFlavourTable.cut = proc.genJetTable.cut
1221
1222     genJetFlavourAssociationName = "genJet{}FlavourAssociation".format(genJetName)
1223     setattr(proc, genJetFlavourAssociationName, genJetFlavourAssociation.clone(
1224         jets           = proc.genJetTable.src,
1225         jetAlgorithm   = supportedJetAlgos[genJetAlgo],
1226         rParam         = genJetSizeNr,
1227       )
1228     )
1229     proc.jetMCTask.add(getattr(proc, genJetFlavourAssociationName))
1230   else:
1231     print("custom_jme_cff::ReclusterAK4GenJets: Use slimmedGenJets for AK4 Gen jets")
1232     proc.AK4GenJetsNoNu = cms.EDFilter("GenJetSelector",
1233       src = cms.InputTag("slimmedGenJets"),
1234       cut = cms.string('pt > 5'),
1235       filter = cms.bool(False)
1236     )
1237     proc.jetMCTask.add(proc.AK4GenJetsNoNu)
1238     proc.genJetTable.cut = ""
1239     proc.genJetTable.doc = "AK4 Gen jets (made with visible genparticles) with pt > 5 GeV. Sourced from slimmedGenJets"
1240     proc.genJetFlavourTable.cut = proc.genJetTable.cut
1241
1242   return proc
1243
1244 def AddNewAK8GenJetsForJEC(proc, genJA):
1245   """
1246   Make a separate AK8 Gen jet collection for JEC studies.
1247   """
1248   print("custom_jme_cff::AddNewAK8GenJetsForJEC: Add new AK8 Gen jets for JEC studies")
1249
1250   #
1251   # Recluster AK8 Gen jet
1252   #
1253   cfg = {
1254     "jet" : "ak8gen",
1255     "minPtFastjet" : 10.
1256   }
1257   genJetInfo = genJA.addGenJetCollection(proc, **cfg)
1258
1259   genJetName         = genJetInfo.jetUpper
1260   genJetAlgo         = genJetInfo.jetAlgo
1261   genJetSize         = genJetInfo.jetSize
1262   genJetSizeNr       = genJetInfo.jetSizeNr
1263   genJetFinalColl    = "{}{}{}".format(genJetAlgo.upper(), genJetSize, "GenJetsNoNu")
1264   genJetTablePrefix  = "GenJetAK8ForJEC"
1265   genJetTableDoc     = "AK8 Gen jets (made with visible genparticles) with pt > 10 GeV. Reclustered for JEC studies."
1266   genJetTableCut     = "pt > 10"
1267
1268   SaveGenJets(proc, genJetName, genJetAlgo, genJetSizeNr, genJetFinalColl, genJetTablePrefix, genJetTableDoc, genJetTableCut)
1269
1270   return proc
1271
1272 def AddVariablesForAK4GenJets(proc):
1273   proc.genJetTable.variables.pt.precision=12
1274   proc.genJetTable.variables.mass.precision=12
1275   proc.genJetTable.variables.nConstituents = GENJETVARS.nConstituents
1276   proc.genJetTable.variables.chHEF = GENJETVARS.chHEF
1277   proc.genJetTable.variables.neHEF = GENJETVARS.neHEF
1278   proc.genJetTable.variables.chEmEF = GENJETVARS.chEmEF
1279   proc.genJetTable.variables.neEmEF = GENJETVARS.neEmEF
1280   proc.genJetTable.variables.muEF = GENJETVARS.muEF
1281   proc.genJetTable.variables.chHadMultiplicity = GENJETVARS.chHadMultiplicity
1282   proc.genJetTable.variables.neHadMultiplicity = GENJETVARS.neHadMultiplicity
1283   proc.genJetTable.variables.chEmMultiplicity = GENJETVARS.chEmMultiplicity
1284   proc.genJetTable.variables.neEmMultiplicity = GENJETVARS.neEmMultiplicity
1285   proc.genJetTable.variables.muMultiplicity = GENJETVARS.muMultiplicity
1286
1287   return proc
1288
1289 def AddVariablesForAK8GenJets(proc):
1290   proc.genJetAK8Table.variables.pt.precision=12
1291   proc.genJetAK8Table.variables.mass.precision=12
1292   proc.genJetAK8Table.variables.nConstituents = GENJETVARS.nConstituents
1293   return proc
1294
1295 def ModifyAK4JetMCTable(proc):
1296   # Modify genjet pt selection when saving the genJetIdx
1297   proc.jetMCTable.variables.genJetIdx = Var("?genJetFwdRef().backRef().isNonnull() && genJetFwdRef().backRef().pt() > 5?genJetFwdRef().backRef().key():-1", "int16", doc="index of matched gen jet")
1298   return proc
1299
1300 #===========================================================================
1301 #
1302 # Misc. functions
1303 #
1304 #===========================================================================
1305 def RemoveAllJetPtCuts(proc):
1306   """
1307   Remove default pt cuts for all jets set in jets_cff.py
1308   """
1309
1310   proc.finalJets.cut             = "" # 15 -> 10
1311   proc.finalJetsPuppi.cut        = "" # 15 -> 10
1312   proc.finalJetsAK8.cut          = "" # 170 -> 170
1313   proc.genJetTable.cut           = "" # 10 -> 3
1314   proc.genJetFlavourTable.cut    = "" # 10 -> 3
1315   proc.genJetAK8Table.cut        = "" # 100 -> 80
1316   proc.genJetAK8FlavourTable.cut = "" # 100 -> 80
1317
1318   return proc
1319
1320 def RecomputePuppiWeights(proc):
1321   """
1322   Setup packedpuppi and packedpuppiNoLep to recompute puppi weights
1323   """
1324   if hasattr(proc,"packedpuppi"):
1325     proc.packedpuppi.useExistingWeights = False
1326   if hasattr(proc,"packedpuppiNoLep"):
1327     proc.packedpuppiNoLep.useExistingWeights = False
1328   return proc
1329
1330 def RecomputePuppiMET(proc):
1331   """
1332   Recompute PuppiMET. This is useful when puppi weights are recomputed.
1333   """
1334   runOnMC=True
1335   if hasattr(proc,"NANOEDMAODoutput") or hasattr(proc,"NANOAODoutput"):
1336     runOnMC = False
1337
1338   from PhysicsTools.PatUtils.tools.runMETCorrectionsAndUncertainties import runMetCorAndUncFromMiniAOD
1339   runMetCorAndUncFromMiniAOD(proc, isData=not(runOnMC),
1340     jetCollUnskimmed='updatedJetsPuppi',metType='Puppi',postfix='Puppi',jetFlavor='AK4PFPuppi',
1341     puppiProducerLabel='packedpuppi',puppiProducerForMETLabel='packedpuppiNoLep',
1342     recoMetFromPFCs=True
1343   )
1344   return proc
1345
1346 #===========================================================================
1347 #
1348 # Functions to setup recomputation of Puppi weights and recluster PuppiMET
1349 # and AK8 Puppi jets (slimmedJetsAK8)
1350 #
1351 #===========================================================================
1352 def RecomputePuppiWeightsAndMET(proc):
1353   """
1354   Recompute Puppi weights and PuppiMET.
1355   """
1356
1357   proc = RecomputePuppiWeights(proc)
1358   proc = RecomputePuppiMET(proc)
1359
1360   return proc
1361
1362 def RecomputePuppiWeightsMETAK8(proc):
1363   """
1364   Recompute Puppi weights and PuppiMET and rebuild slimmedJetsAK8.
1365   """
1366   runOnMC=True
1367   if hasattr(proc,"NANOEDMAODoutput") or hasattr(proc,"NANOAODoutput"):
1368     runOnMC = False
1369
1370   proc = RecomputePuppiWeightsAndMET(proc)
1371
1372   from RecoBTag.ONNXRuntime.pfParticleNet_cff import _pfParticleNetJetTagsAll as pfParticleNetJetTagsAll
1373   from RecoBTag.ONNXRuntime.pfParticleNet_cff import _pfParticleNetMassRegressionOutputs as pfParticleNetMassRegressionOutputs
1374   from RecoBTag.ONNXRuntime.pfParticleNet_cff import _pfParticleNetMassCorrelatedJetTagsAll as pfParticleNetMassCorrelatedJetTagsAll
1375   from RecoBTag.ONNXRuntime.pfParticleNetFromMiniAODAK8_cff import _pfParticleNetFromMiniAODAK8JetTagsAll as pfParticleNetFromMiniAODAK8JetTagsAll
1376   from RecoBTag.ONNXRuntime.pfGlobalParticleTransformerAK8_cff import _pfGlobalParticleTransformerAK8JetTagsAll as pfGlobalParticleTransformerAK8JetTagsAll
1377
1378   btagDiscriminatorsAK8 = cms.PSet(names = cms.vstring(
1379       pfParticleNetMassCorrelatedJetTagsAll+
1380       pfGlobalParticleTransformerAK8JetTagsAll+
1381       pfParticleNetFromMiniAODAK8JetTagsAll+
1382       pfParticleNetJetTagsAll+
1383       pfParticleNetMassRegressionOutputs
1384     )
1385   )
1386
1387   btagDiscriminatorsAK8Subjets = cms.PSet(names = cms.vstring(
1388       'pfDeepFlavourJetTags:probb',
1389       'pfDeepFlavourJetTags:probbb',
1390       'pfDeepFlavourJetTags:problepb',
1391       'pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:BvsAll'
1392     )
1393   )
1394
1395   from PhysicsTools.PatAlgos.tools.puppiJetMETReclusteringFromMiniAOD_cff import setupPuppiAK4AK8METReclustering
1396   proc = setupPuppiAK4AK8METReclustering(proc, runOnMC=runOnMC, useExistingWeights=False,
1397     reclusterAK4MET=False, # Already setup to recluster AK4 Puppi jets and PuppiMET
1398     reclusterAK8=True,
1399     btagDiscriminatorsAK8=btagDiscriminatorsAK8,
1400     btagDiscriminatorsAK8Subjets=btagDiscriminatorsAK8Subjets
1401   )
1402
1403   return proc
1404
1405 #===========================================================================
1406 #
1407 # CUSTOMIZATION function
1408 #
1409 #===========================================================================
1410 def PrepJMECustomNanoAOD(process):
1411   ## TODO : find a better way to handle data or MC by modifying the proper Tasks
1412   runOnMC=True
1413   if hasattr(process,"NANOEDMAODoutput") or hasattr(process,"NANOAODoutput"):
1414     runOnMC = False
1415   ############################################################################
1416   # Remove all default jet pt cuts from jets_cff.py
1417   ############################################################################
1418   process = RemoveAllJetPtCuts(process)
1419
1420   ###########################################################################
1421   #
1422   # Gen-level jets related functions. Only for MC.
1423   #
1424   ###########################################################################
1425   genJA = GenJetAdder()
1426   if runOnMC:
1427     ############################################################################
1428     # Save additional variables for AK8 GEN jets
1429     ############################################################################
1430     process = AddVariablesForAK8GenJets(process)
1431     ############################################################################
1432     # Recluster AK8 GEN jets
1433     ############################################################################
1434     process = AddNewAK8GenJetsForJEC(process, genJA)
1435     ############################################################################
1436     # Modify jetMCTable
1437     ############################################################################
1438     process = ModifyAK4JetMCTable(process)
1439     ###########################################################################
1440     # Recluster AK4 GEN jets
1441     ###########################################################################
1442     process = ConfigureAK4GenJets(process, genJA)
1443     process = AddVariablesForAK4GenJets(process)
1444     ###########################################################################
1445     # Add additional GEN jets to NanoAOD
1446     ###########################################################################
1447     for jetConfig in config_genjets:
1448       cfg = { k : v for k, v in jetConfig.items() if k != "enabled"}
1449       genJetInfo = genJA.addGenJetCollection(process, **cfg)
1450       AddNewGenJets(process, genJetInfo)
1451
1452   ###########################################################################
1453   #
1454   # Reco-level jets related functions. For both MC and data.
1455   #
1456   ###########################################################################
1457   recoJA = RecoJetAdder(runOnMC=runOnMC)
1458   ###########################################################################
1459   # Save additional variables for AK8Puppi jets
1460   ###########################################################################
1461   process = AddVariablesForAK8PuppiJets(process)
1462   ###########################################################################
1463   # Build a separate AK8Puppi jet collection for JEC studies
1464   ###########################################################################
1465   process = AddNewAK8PuppiJetsForJEC(process, recoJA, runOnMC)
1466   ###########################################################################
1467   # Recluster AK4 CHS jets and replace "slimmedJets"
1468   ###########################################################################
1469   process = ReclusterAK4CHSJets(process, recoJA, runOnMC)
1470   ###########################################################################
1471   # Recluster AK4 Puppi jets and replace "slimmedJets"
1472   ###########################################################################
1473   process = ReclusterAK4PuppiJets(process, recoJA, runOnMC)
1474   ###########################################################################
1475   # Add additional Reco jets to NanoAOD
1476   ###########################################################################
1477   for jetConfig in config_recojets:
1478     cfg = { k : v for k, v in jetConfig.items() if k != "enabled"}
1479
1480     recoJetInfo = recoJA.addRecoJetCollection(process, **cfg)
1481     AddNewPatJets(process, recoJetInfo, runOnMC)
1482
1483   ###########################################################################
1484   # Add jet tasks
1485   # By default for Run-3, add AK4 CHS jet tasks.
1486   # For Run-2 eras, add AK4 Puppi jet tasks
1487   ###########################################################################
1488   def addAK4JetTasks(proc, addAK4CHSJetTasks, addAK4PuppiJetTasks):
1489     if addAK4CHSJetTasks:
1490       proc.nanoTableTaskCommon.add(proc.jetTask)
1491       proc.nanoTableTaskCommon.add(proc.jetTablesTask)
1492       proc.nanoTableTaskCommon.add(proc.jetForMETTask)
1493     if addAK4PuppiJetTasks:
1494       proc.nanoTableTaskCommon.add(proc.jetPuppiTask)
1495       proc.nanoTableTaskCommon.add(proc.jetPuppiTablesTask)
1496       proc.nanoTableTaskCommon.add(proc.jetPuppiForMETTask)
1497     return proc
1498
1499   jmeNano_addAK4JetTasks_switch = cms.PSet(
1500     addAK4CHS_switch = cms.untracked.bool(True),
1501     addAK4Puppi_switch = cms.untracked.bool(False)
1502   )
1503   run2_nanoAOD_ANY.toModify(jmeNano_addAK4JetTasks_switch,
1504     addAK4CHS_switch = False,
1505     addAK4Puppi_switch = True
1506   )
1507   process = addAK4JetTasks(process,
1508     addAK4CHSJetTasks = jmeNano_addAK4JetTasks_switch.addAK4CHS_switch,
1509     addAK4PuppiJetTasks = jmeNano_addAK4JetTasks_switch.addAK4Puppi_switch,
1510   )
1511
1512   ###########################################################################
1513   # Fix ParticleNetFromMiniAOD input when slimmedTaus is updated
1514   ###########################################################################
1515   from PhysicsTools.NanoAOD.nano_cff import _fixPNetInputCollection
1516   (run2_nanoAOD_106Xv2 | run3_nanoAOD_pre142X | nanoAOD_rePuppi).toModify(
1517       process, lambda p: _fixPNetInputCollection(p)
1518   )
1519
1520   ###########################################################################
1521   # Save Maximum of Pt Hat Max
1522   ###########################################################################
1523   if runOnMC:
1524     process.puTable.savePtHatMax = True
1525
1526   ###########################################################################
1527   # Save all Parton-Shower weights
1528   ###########################################################################
1529   if runOnMC:
1530     process.genWeightsTable.keepAllPSWeights = True
1531
1532   return process
1533