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