Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2024-04-25 02:14:03

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