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.softKiller_cfi import softKiller
0008
0009 from RecoJets.JetProducers.PFJetParameters_cfi         import PFJetParameters
0010 from RecoJets.JetProducers.GenJetParameters_cfi        import GenJetParameters
0011 from RecoJets.JetProducers.GenJetParameters_cfi        import GenJetParameters
0012 from RecoJets.JetProducers.AnomalousCellParameters_cfi import AnomalousCellParameters
0013
0014 from RecoJets.JetProducers.ak4GenJets_cfi  import ak4GenJets
0015 from RecoJets.JetProducers.ak4PFJets_cfi   import ak4PFJets, ak4PFJetsCHS, ak4PFJetsPuppi, ak4PFJetsSK, ak4PFJetsCS
0016 from RecoJets.JetProducers.ak4CaloJets_cfi import ak4CaloJets
0017
0018 from PhysicsTools.PatAlgos.producersLayer1.jetUpdater_cfi import updatedPatJets
0019 from PhysicsTools.PatAlgos.recoLayer0.jetCorrFactors_cfi  import patJetCorrFactors
0020 from PhysicsTools.PatAlgos.mcMatchLayer0.jetFlavourId_cff import patJetFlavourAssociation
0021
0022 from PhysicsTools.PatAlgos.tools.jetTools import supportedJetAlgos, addJetCollection, updateJetCollection
0023 from PhysicsTools.PatAlgos.tools.jetTools import setupPuppiForPackedPF
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           pfCandWeight = setupPuppiForPackedPF(proc)[0]
0297           self.prerequisites.append(pfCandWeight)
0298         #
0299         # Softkiller
0300         #
0301         elif recoJetInfo.jetPUMethod == "sk":
0302           self.addProcessAndTask(proc, pfCand, softKiller.clone(
0303               PFCandidates = self.pfLabel,
0304               rParam = recoJetInfo.jetSizeNr,
0305             )
0306           )
0307           self.prerequisites.append(pfCand)
0308         else:
0309           raise RuntimeError("Currently unsupported PU method: '%s'" % recoJetInfo.jetPUMethod)
0310
0311       #============================================
0312       #
0313       # Create the recojet collection
0314       #
0315       #============================================
0316       if not recoJetInfo.doCalo:
0317         jetCollection = '{}Collection'.format(jetUpper)
0318
0319         if jetCollection in self.main:
0320           raise ValueError("Step '%s' already implemented" % jetCollection)
0321         #
0322         # Cluster new jet
0323         #
0324         if recoJetInfo.jetPUMethod == "chs":
0325           self.addProcessAndTask(proc, jetCollection, ak4PFJetsCHS.clone(
0326               src = pfCand,
0327             )
0328           )
0329         elif recoJetInfo.jetPUMethod == "puppi":
0330           self.addProcessAndTask(proc, jetCollection, ak4PFJetsPuppi.clone(
0331               src = self.pfLabel,
0332               srcWeights = pfCandWeight
0333             )
0334           )
0335         elif recoJetInfo.jetPUMethod == "sk":
0336           self.addProcessAndTask(proc, pfCand, ak4PFJetsSK.clone(
0337               src = pfCand,
0338             )
0339           )
0340         elif recoJetInfo.jetPUMethod == "cs":
0341           self.addProcessAndTask(proc, jetCollection, ak4PFJetsCS.clone(
0342             src = pfCand,
0343           )
0344         )
0345         else:
0346           self.addProcessAndTask(proc, jetCollection, ak4PFJets.clone(
0347             src = pfCand,
0348           )
0349         )
0350         getattr(proc, jetCollection).jetAlgorithm = supportedJetAlgos[recoJetInfo.jetAlgo]
0351         getattr(proc, jetCollection).rParam = recoJetInfo.jetSizeNr
0352         #
0353         # Set minimum pt threshold of reco jets to be saved after fastjet clustering
0354         #
0355         if minPtFastjet != None:
0356           getattr(proc, jetCollection).jetPtMin = minPtFastjet
0357         currentTasks.append(jetCollection)
0358       else:
0359         jetCollection = inputCollection
0360
0361
0362       #=============================================
0363       #
0364       # Make patJet collection
0365       #
0366       #=============================================
0367       #
0368       # Jet correction
0369       #
0370       if recoJetInfo.jetPUMethod == "puppi":
0371         jetCorrLabel = "Puppi"
0372       elif recoJetInfo.jetPUMethod in [ "cs", "sk" ]:
0373         jetCorrLabel = "chs"
0374       else:
0375         jetCorrLabel = recoJetInfo.jetPUMethod
0376
0377       jetCorrections = (
0378         "{}{}{}{}".format(
0379           recoJetInfo.jetAlgo.upper(),
0380           recoJetInfo.jetSize,
0381           "Calo" if recoJetInfo.doCalo else recoJetInfo.jetReco.upper(),
0382           jetCorrLabel
0383         ),
0384         JETCorrLevels,
0385         "None",
0386       )
0387
0388       postfix = "Recluster" if inputCollection == "" else ""
0389       addJetCollection(
0390         proc,
0391         labelName          = jetUpper,
0392         postfix            = postfix,
0393         jetSource          = cms.InputTag(jetCollection),
0394         algo               = recoJetInfo.jetAlgo,
0395         rParam             = recoJetInfo.jetSizeNr,
0396         pvSource           = cms.InputTag(self.pvLabel),
0397         pfCandidates       = cms.InputTag(self.pfLabel),
0398         svSource           = cms.InputTag(self.svLabel),
0399         muSource           = cms.InputTag(self.muLabel),
0400         elSource           = cms.InputTag(self.elLabel),
0401         genJetCollection   = cms.InputTag(genJetsCollection),
0402         genParticles       = cms.InputTag(self.gpLabel),
0403         jetCorrections     = jetCorrections,
0404       )
0405
0406       #
0407       # Need to set this explicitly for PUPPI jets
0408       #
0409       if recoJetInfo.jetPUMethod == "puppi":
0410         getattr(proc, "patJetFlavourAssociation{}{}".format(jetUpper,postfix)).weights = cms.InputTag(pfCandWeight)
0411
0412       getJetMCFlavour = not recoJetInfo.doCalo and recoJetInfo.jetPUMethod != "cs"
0413       if not self.runOnMC: #Remove modules for Gen-level object matching
0414         delattr(proc, 'patJetGenJetMatch{}{}'.format(jetUpper,postfix))
0415         delattr(proc, 'patJetPartonMatch{}{}'.format(jetUpper,postfix))
0416         getJetMCFlavour = False
0417       setattr(getattr(proc, "patJets{}{}".format(jetUpper,postfix)), "getJetMCFlavour", cms.bool(getJetMCFlavour))
0418
0419       selectedPatJets = "selectedPatJets{}{}".format(jetUpper,postfix)
0420       #=============================================
0421       #
0422       # Update the patJet collection.
0423       # This is where we setup
0424       # -  JEC
0425       # -  b-tagging discriminators
0426       #
0427       #=============================================
0428       updateJetCollection(
0429         proc,
0430         labelName          = jetUpper,
0431         postfix            = "Final",
0432         jetSource          = cms.InputTag(selectedPatJets),
0433         jetCorrections     = jetCorrections,
0434         btagDiscriminators = bTagDiscriminators,
0435       )
0436
0437       recoJetInfo.patJetFinalCollection = "selectedUpdatedPatJets{}{}".format(jetUpper,"Final")
0438     else:
0439       recoJetInfo.patJetFinalCollection = inputCollection
0440
0441     self.main.extend(currentTasks)
0442
0443     return recoJetInfo