python/tools/jetCollectionTools.py

0001 import FWCore.ParameterSet.Config as cms
0002
0003 from PhysicsTools.PatAlgos.tools.ConfigToolBase import *
0004
0005 from CommonTools.ParticleFlow.pfCHS_cff import pfCHS
0006
0007 from CommonTools.PileupAlgos.Puppi_cff      import puppi
0008 from CommonTools.PileupAlgos.softKiller_cfi import softKiller
0009
0010 from RecoJets.JetProducers.PFJetParameters_cfi         import PFJetParameters
0011 from RecoJets.JetProducers.GenJetParameters_cfi        import GenJetParameters
0012 from RecoJets.JetProducers.GenJetParameters_cfi        import GenJetParameters
0013 from RecoJets.JetProducers.AnomalousCellParameters_cfi import AnomalousCellParameters
0014
0015 from RecoJets.JetProducers.ak4GenJets_cfi  import ak4GenJets
0016 from RecoJets.JetProducers.ak4PFJets_cfi   import ak4PFJets, ak4PFJetsCHS, ak4PFJetsPuppi, ak4PFJetsSK, ak4PFJetsCS
0017 from RecoJets.JetProducers.ak4CaloJets_cfi import ak4CaloJets
0018
0019 from PhysicsTools.PatAlgos.producersLayer1.jetUpdater_cfi import updatedPatJets
0020 from PhysicsTools.PatAlgos.recoLayer0.jetCorrFactors_cfi  import patJetCorrFactors
0021 from PhysicsTools.PatAlgos.mcMatchLayer0.jetFlavourId_cff import patJetFlavourAssociation
0022
0023 from PhysicsTools.PatAlgos.tools.jetTools import supportedJetAlgos, addJetCollection, updateJetCollection
0024 from PhysicsTools.PatAlgos.tools.helpers  import getPatAlgosToolsTask, addToProcessAndTask
0025
0026 import re
0027
0028 #============================================
0029 #
0030 # GenJetInfo
0031 #
0032 #============================================
0033 class GenJetInfo(object):
0034   """
0035   Class to hold information of a genjet collection
0036   """
0037   def __init__(self, jet, inputCollection):
0038     self.jet = jet
0039     self.jetLower = jet.lower()
0040     self.jetUpper = jet.upper()
0041     self.jetTagName = self.jetUpper
0042     self.inputCollection = inputCollection
0043     algoKey     = 'algo'
0044     sizeKey     = 'size'
0045     recoKey     = 'reco'
0046     jetRegex = re.compile(
0047       r'(?P<{algo}>({algoList}))(?P<{size}>[0-9]+)gen'.format(
0048         algo     = algoKey,
0049         algoList = '|'.join(supportedJetAlgos.keys()),
0050         size     = sizeKey,
0051       )
0052     )
0053     jetMatch = jetRegex.match(jet.lower())
0054     if not jetMatch:
0055       raise RuntimeError('Invalid jet collection: %s' % jet)
0056     self.jetAlgo    = jetMatch.group(algoKey)
0057     self.jetSize    = jetMatch.group(sizeKey)
0058     self.jetSizeNr  = float(self.jetSize) / 10.
0059
0060 #============================================
0061 #
0062 # GenJetAdder
0063 #
0064 #============================================
0065 class GenJetAdder(object):
0066   """
0067   Tool to schedule modules for building a genjet collection with input MiniAODs
0068   """
0069   def __init__(self):
0070     self.prerequisites = []
0071     self.main = []
0072     self.gpLabel = "prunedGenParticles"
0073
0074   def addProcessAndTask(self, proc, label, module):
0075     task = getPatAlgosToolsTask(proc)
0076     addToProcessAndTask(label, module, proc, task)
0077
0078   def addGenJetCollection(self,
0079       proc,
0080       jet,
0081       inputCollection    = "",
0082       minPtFastjet       = None,
0083     ):
0084     print("jetCollectionTools::GenJetAdder::addGenJetCollection: Adding Gen Jet Collection: {}".format(jet))
0085
0086     #
0087     # Decide which genJet collection we are dealing with
0088     #
0089     genJetInfo = GenJetInfo(jet,inputCollection)
0090     jetLower = genJetInfo.jetLower
0091     jetUpper = genJetInfo.jetUpper
0092
0093     #=======================================================
0094     #
0095     # If genJet collection in MiniAOD is not
0096     # specified, build the genjet collection.
0097     #
0098     #========================================================
0099     if not inputCollection:
0100       print("jetCollectionTools::GenJetAdder::addGenJetCollection: inputCollection not specified. Building genjet collection now")
0101       #
0102       # Setup GenParticles
0103       #
0104       packedGenPartNoNu = "packedGenParticlesForJetsNoNu"
0105       if packedGenPartNoNu not in self.prerequisites:
0106         self.addProcessAndTask(proc, packedGenPartNoNu, cms.EDFilter("CandPtrSelector",
0107             src = cms.InputTag("packedGenParticles"),
0108             cut = cms.string("abs(pdgId) != 12 && abs(pdgId) != 14 && abs(pdgId) != 16"),
0109           )
0110         )
0111         self.prerequisites.append(packedGenPartNoNu)
0112       #
0113       # Create the GenJet collection
0114       #
0115       genJetsCollection = "{}{}{}".format(genJetInfo.jetAlgo.upper(), genJetInfo.jetSize, 'GenJetsNoNu')
0116       self.addProcessAndTask(proc, genJetsCollection, ak4GenJets.clone(
0117           src           = packedGenPartNoNu,
0118           jetAlgorithm  = cms.string(supportedJetAlgos[genJetInfo.jetAlgo]),
0119           rParam        = cms.double(genJetInfo.jetSizeNr),
0120         )
0121       )
0122       #
0123       # Set minimum pt threshold of gen jets to be saved after fastjet clustering
0124       #
0125       if minPtFastjet != None:
0126         getattr(proc, genJetsCollection).jetPtMin = minPtFastjet
0127       self.prerequisites.append(genJetsCollection)
0128
0129     return genJetInfo
0130 #============================================
0131 #
0132 # RecoJetInfo
0133 #
0134 #============================================
0135 class RecoJetInfo(object):
0136   """
0137   Class to hold information of a recojet collection
0138   """
0139   def __init__(self, jet, inputCollection):
0140     self.jet = jet
0141     self.jetLower   = jet.lower()
0142     self.jetUpper   = jet.upper()
0143     self.jetTagName = self.jetUpper
0144     self.inputCollection = inputCollection
0145     algoKey     = 'algo'
0146     sizeKey     = 'size'
0147     recoKey     = 'reco'
0148     puMethodKey = 'puMethod'
0149     jetRegex = re.compile(
0150       r'(?P<{algo}>({algoList}))(?P<{size}>[0-9]+)(?P<{reco}>(pf|calo))(?P<{puMethod}>(chs|puppi|sk|cs|))'.format(
0151         algo     = algoKey,
0152         algoList = '|'.join(supportedJetAlgos.keys()),
0153         size     = sizeKey,
0154         reco     = recoKey,
0155         puMethod = puMethodKey,
0156       )
0157     )
0158     jetMatch = jetRegex.match(jet.lower())
0159     if not jetMatch:
0160       raise RuntimeError('Invalid jet collection: %s' % jet)
0161
0162     self.jetAlgo     = jetMatch.group(algoKey)
0163     self.jetSize     = jetMatch.group(sizeKey)
0164     self.jetReco     = jetMatch.group(recoKey)
0165     self.jetPUMethod = jetMatch.group(puMethodKey)
0166
0167     self.jetSizeNr = float(self.jetSize) / 10.
0168
0169     self.doCalo = self.jetReco == "calo"
0170     self.doPF   = self.jetReco == "pf"
0171
0172     self.doCS   = self.jetPUMethod == "cs"
0173     self.skipUserData = self.doCalo or (self.jetPUMethod in [ "puppi", "sk" ] and inputCollection == "")
0174
0175     self.jetCorrPayload = "{}{}{}".format(
0176       self.jetAlgo.upper(), self.jetSize, "Calo" if self.doCalo else self.jetReco.upper()
0177     )
0178
0179     if self.jetPUMethod == "puppi":
0180       self.jetCorrPayload += "Puppi"
0181     elif self.jetPUMethod in [ "cs", "sk" ]:
0182       self.jetCorrPayload += "chs"
0183     else:
0184       self.jetCorrPayload += self.jetPUMethod.lower()
0185
0186     self.patJetFinalCollection = ""
0187
0188 #============================================
0189 #
0190 # RecoJetAdder
0191 #
0192 #============================================
0193 class RecoJetAdder(object):
0194   """
0195   Tool to schedule modules for building a patJet collection from MiniAODs
0196   """
0197   def __init__(self,runOnMC=True):
0198     self.prerequisites = []
0199     self.main = []
0200     self.pfLabel = "packedPFCandidates"
0201     self.pvLabel = "offlineSlimmedPrimaryVertices"
0202     self.svLabel = "slimmedSecondaryVertices"
0203     self.muLabel = "slimmedMuons"
0204     self.elLabel = "slimmedElectrons"
0205     self.gpLabel = "prunedGenParticles"
0206     self.runOnMC = runOnMC
0207     self.patJetsInMiniAOD = ["slimmedJets", "slimmedJetsAK8", "slimmedJetsPuppi", "slimmedCaloJets"]
0208
0209   def addProcessAndTask(self, proc, label, module):
0210     task = getPatAlgosToolsTask(proc)
0211     addToProcessAndTask(label, module, proc, task)
0212
0213   def addRecoJetCollection(self,
0214     proc,
0215     jet,
0216     inputCollection    = "",
0217     minPtFastjet       = None,
0218     genJetsCollection  = "",
0219     bTagDiscriminators = ["None"],
0220     JETCorrLevels      = ["L1FastJet", "L2Relative", "L3Absolute", "L2L3Residual"],
0221     ):
0222     print("jetCollectionTools::RecoJetAdder::addRecoJetCollection: Adding Reco Jet Collection: {}".format(jet))
0223
0224     currentTasks = []
0225
0226     if inputCollection and inputCollection not in self.patJetsInMiniAOD:
0227       raise RuntimeError("Invalid input collection: %s" % inputCollection)
0228
0229     #=======================================================
0230     #
0231     # Figure out which jet collection we're dealing with
0232     #
0233     #=======================================================
0234     recoJetInfo = RecoJetInfo(jet, inputCollection)
0235     jetLower = recoJetInfo.jetLower
0236     jetUpper = recoJetInfo.jetUpper
0237     tagName  = recoJetInfo.jetTagName
0238
0239     if inputCollection == "slimmedJets":
0240       assert(jetLower == "ak4pfchs")
0241     elif inputCollection == "slimmedJetsAK8":
0242       assert(jetLower == "ak8pfpuppi")
0243     elif inputCollection == "slimmedJetsPuppi":
0244       assert(jetLower == "ak4pfpuppi")
0245     elif inputCollection == "slimmedCaloJets":
0246       assert(jetLower == "ak4calo")
0247
0248     #=======================================================
0249     #
0250     # If the patJet collection in MiniAOD is not specified,
0251     # we have to build the patJet collection from scratch.
0252     #
0253     #========================================================
0254     if not inputCollection or recoJetInfo.doCalo:
0255       print("jetCollectionTools::RecoJetAdder::addRecoJetCollection: inputCollection not specified. Building recojet collection now")
0256
0257       #=======================================================
0258       #
0259       # Prepare the inputs to jet clustering
0260       #
0261       #========================================================
0262       #
0263       # Specify PF candidates
0264       #
0265       pfCand = self.pfLabel
0266       #
0267       # Setup PU method for PF candidates
0268       #
0269       if recoJetInfo.jetPUMethod not in [ "", "cs" ]:
0270         pfCand += recoJetInfo.jetPUMethod
0271
0272
0273       #
0274       # Setup modules to perform PU mitigation for
0275       # PF candidates
0276       #
0277       if pfCand not in self.prerequisites:
0278         #
0279         # Skip if no PU Method or CS specified
0280         #
0281         if recoJetInfo.jetPUMethod in [ "", "cs" ]:
0282           pass
0283         #
0284         # CHS
0285         #
0286         elif recoJetInfo.jetPUMethod == "chs":
0287           self.addProcessAndTask(proc, pfCand, pfCHS.clone(
0288               src = self.pfLabel
0289              )
0290            )
0291           self.prerequisites.append(pfCand)
0292         #
0293         # PUPPI
0294         #
0295         elif recoJetInfo.jetPUMethod == "puppi":
0296           self.addProcessAndTask(proc, pfCand, puppi.clone(
0297               candName = self.pfLabel,
0298               vertexName = self.pvLabel,
0299               clonePackedCands = True,
0300               useExistingWeights = True,
0301             )
0302           )
0303           self.prerequisites.append(pfCand)
0304         #
0305         # Softkiller
0306         #
0307         elif recoJetInfo.jetPUMethod == "sk":
0308           self.addProcessAndTask(proc, pfCand, softKiller.clone(
0309               PFCandidates = self.pfLabel,
0310               rParam = recoJetInfo.jetSizeNr,
0311             )
0312           )
0313           self.prerequisites.append(pfCand)
0314         else:
0315           raise RuntimeError("Currently unsupported PU method: '%s'" % recoJetInfo.jetPUMethod)
0316
0317       #============================================
0318       #
0319       # Create the recojet collection
0320       #
0321       #============================================
0322       if not recoJetInfo.doCalo:
0323         jetCollection = '{}Collection'.format(jetUpper)
0324
0325         if jetCollection in self.main:
0326           raise ValueError("Step '%s' already implemented" % jetCollection)
0327         #
0328         # Cluster new jet
0329         #
0330         if recoJetInfo.jetPUMethod == "chs":
0331           self.addProcessAndTask(proc, jetCollection, ak4PFJetsCHS.clone(
0332               src = pfCand,
0333             )
0334           )
0335         elif recoJetInfo.jetPUMethod == "puppi":
0336           self.addProcessAndTask(proc, jetCollection, ak4PFJetsPuppi.clone(
0337               src = self.pfLabel,
0338               srcWeights = pfCand
0339             )
0340           )
0341         elif recoJetInfo.jetPUMethod == "sk":
0342           self.addProcessAndTask(proc, pfCand, ak4PFJetsSK.clone(
0343               src = pfCand,
0344             )
0345           )
0346         elif recoJetInfo.jetPUMethod == "cs":
0347           self.addProcessAndTask(proc, jetCollection, ak4PFJetsCS.clone(
0348             src = pfCand,
0349           )
0350         )
0351         else:
0352           self.addProcessAndTask(proc, jetCollection, ak4PFJets.clone(
0353             src = pfCand,
0354           )
0355         )
0356         getattr(proc, jetCollection).jetAlgorithm = supportedJetAlgos[recoJetInfo.jetAlgo]
0357         getattr(proc, jetCollection).rParam = recoJetInfo.jetSizeNr
0358         #
0359         # Set minimum pt threshold of reco jets to be saved after fastjet clustering
0360         #
0361         if minPtFastjet != None:
0362           getattr(proc, jetCollection).jetPtMin = minPtFastjet
0363         currentTasks.append(jetCollection)
0364       else:
0365         jetCollection = inputCollection
0366
0367
0368       #=============================================
0369       #
0370       # Make patJet collection
0371       #
0372       #=============================================
0373       #
0374       # Jet correction
0375       #
0376       if recoJetInfo.jetPUMethod == "puppi":
0377         jetCorrLabel = "Puppi"
0378       elif recoJetInfo.jetPUMethod in [ "cs", "sk" ]:
0379         jetCorrLabel = "chs"
0380       else:
0381         jetCorrLabel = recoJetInfo.jetPUMethod
0382
0383       jetCorrections = (
0384         "{}{}{}{}".format(
0385           recoJetInfo.jetAlgo.upper(),
0386           recoJetInfo.jetSize,
0387           "Calo" if recoJetInfo.doCalo else recoJetInfo.jetReco.upper(),
0388           jetCorrLabel
0389         ),
0390         JETCorrLevels,
0391         "None",
0392       )
0393
0394       postfix = "Recluster" if inputCollection == "" else ""
0395       addJetCollection(
0396         proc,
0397         labelName          = jetUpper,
0398         postfix            = postfix,
0399         jetSource          = cms.InputTag(jetCollection),
0400         algo               = recoJetInfo.jetAlgo,
0401         rParam             = recoJetInfo.jetSizeNr,
0402         pvSource           = cms.InputTag(self.pvLabel),
0403         pfCandidates       = cms.InputTag(self.pfLabel),
0404         svSource           = cms.InputTag(self.svLabel),
0405         muSource           = cms.InputTag(self.muLabel),
0406         elSource           = cms.InputTag(self.elLabel),
0407         genJetCollection   = cms.InputTag(genJetsCollection),
0408         genParticles       = cms.InputTag(self.gpLabel),
0409         jetCorrections     = jetCorrections,
0410       )
0411
0412       #
0413       # Need to set this explicitly for PUPPI jets
0414       #
0415       if recoJetInfo.jetPUMethod == "puppi":
0416         getattr(proc, "patJetFlavourAssociation{}{}".format(jetUpper,postfix)).weights = cms.InputTag(pfCand)
0417
0418       getJetMCFlavour = not recoJetInfo.doCalo and recoJetInfo.jetPUMethod != "cs"
0419       if not self.runOnMC: #Remove modules for Gen-level object matching
0420         delattr(proc, 'patJetGenJetMatch{}{}'.format(jetUpper,postfix))
0421         delattr(proc, 'patJetPartonMatch{}{}'.format(jetUpper,postfix))
0422         getJetMCFlavour = False
0423       setattr(getattr(proc, "patJets{}{}".format(jetUpper,postfix)), "getJetMCFlavour", cms.bool(getJetMCFlavour))
0424
0425       selectedPatJets = "selectedPatJets{}{}".format(jetUpper,postfix)
0426       #=============================================
0427       #
0428       # Update the patJet collection.
0429       # This is where we setup
0430       # -  JEC
0431       # -  b-tagging discriminators
0432       #
0433       #=============================================
0434       updateJetCollection(
0435         proc,
0436         labelName          = jetUpper,
0437         postfix            = "Final",
0438         jetSource          = cms.InputTag(selectedPatJets),
0439         jetCorrections     = jetCorrections,
0440         btagDiscriminators = bTagDiscriminators,
0441       )
0442
0443       recoJetInfo.patJetFinalCollection = "selectedUpdatedPatJets{}{}".format(jetUpper,"Final")
0444     else:
0445       recoJetInfo.patJetFinalCollection = inputCollection
0446
0447     self.main.extend(currentTasks)
0448
0449     return recoJetInfo