Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2024-04-06 12:23:28

0001 from __future__ import print_function
0002 from builtins import range
0003 import ROOT
0004 import os, types
0005 from math import *
0006 from PhysicsTools.HeppyCore.utils.deltar import *
0007 
0008 class JetReCalibrator:
0009     def __init__(self,globalTag,jetFlavour,doResidualJECs,jecPath,upToLevel=3,
0010                  calculateSeparateCorrections=False,
0011                  calculateType1METCorrection=False, type1METParams={'jetPtThreshold':15., 'skipEMfractionThreshold':0.9, 'skipMuons':True} ):
0012         """Create a corrector object that reads the payloads from the text dumps of a global tag under
0013             CMGTools/RootTools/data/jec  (see the getJec.py there to make the dumps).
0014            It will apply the L1,L2,L3 and possibly the residual corrections to the jets.
0015            If configured to do so, it will also compute the type1 MET corrections."""
0016         self.globalTag = globalTag
0017         self.jetFlavour = jetFlavour
0018         self.doResidualJECs = doResidualJECs
0019         self.jecPath = jecPath
0020         self.upToLevel = upToLevel
0021         self.calculateType1METCorr = calculateType1METCorrection
0022         self.type1METParams  = type1METParams
0023         # Make base corrections
0024         path = os.path.expandvars(jecPath) #"%s/src/CMGTools/RootTools/data/jec" % os.environ['CMSSW_BASE'];
0025         self.L1JetPar  = ROOT.JetCorrectorParameters("%s/%s_L1FastJet_%s.txt" % (path,globalTag,jetFlavour),"");
0026         self.L2JetPar  = ROOT.JetCorrectorParameters("%s/%s_L2Relative_%s.txt" % (path,globalTag,jetFlavour),"");
0027         self.L3JetPar  = ROOT.JetCorrectorParameters("%s/%s_L3Absolute_%s.txt" % (path,globalTag,jetFlavour),"");
0028         self.vPar = ROOT.vector(ROOT.JetCorrectorParameters)()
0029         self.vPar.push_back(self.L1JetPar);
0030         if upToLevel >= 2: self.vPar.push_back(self.L2JetPar);
0031         if upToLevel >= 3: self.vPar.push_back(self.L3JetPar);
0032         # Add residuals if needed
0033         if doResidualJECs : 
0034             self.ResJetPar = ROOT.JetCorrectorParameters("%s/%s_L2L3Residual_%s.txt" % (path,globalTag,jetFlavour))
0035             self.vPar.push_back(self.ResJetPar);
0036         #Step3 (Construct a FactorizedJetCorrector object) 
0037         self.JetCorrector = ROOT.FactorizedJetCorrector(self.vPar)
0038         if os.path.exists("%s/%s_Uncertainty_%s.txt" % (path,globalTag,jetFlavour)):
0039             self.JetUncertainty = ROOT.JetCorrectionUncertainty("%s/%s_Uncertainty_%s.txt" % (path,globalTag,jetFlavour));
0040         elif os.path.exists("%s/Uncertainty_FAKE.txt" % path):
0041             self.JetUncertainty = ROOT.JetCorrectionUncertainty("%s/Uncertainty_FAKE.txt" % path);
0042         else:
0043             print('Missing JEC uncertainty file "%s/%s_Uncertainty_%s.txt", so jet energy uncertainties will not be available' % (path,globalTag,jetFlavour))
0044             self.JetUncertainty = None
0045         self.separateJetCorrectors = {}
0046         if calculateSeparateCorrections or calculateType1METCorrection:
0047             self.vParL1 = ROOT.vector(ROOT.JetCorrectorParameters)()
0048             self.vParL1.push_back(self.L1JetPar)
0049             self.separateJetCorrectors["L1"] = ROOT.FactorizedJetCorrector(self.vParL1)
0050             if upToLevel >= 2 and calculateSeparateCorrections:
0051                 self.vParL2 = ROOT.vector(ROOT.JetCorrectorParameters)()
0052                 for i in [self.L1JetPar,self.L2JetPar]: self.vParL2.push_back(i)
0053                 self.separateJetCorrectors["L1L2"] = ROOT.FactorizedJetCorrector(self.vParL2)
0054             if upToLevel >= 3 and calculateSeparateCorrections:
0055                 self.vParL3 = ROOT.vector(ROOT.JetCorrectorParameters)()
0056                 for i in [self.L1JetPar,self.L2JetPar,self.L3JetPar]: self.vParL3.push_back(i)
0057                 self.separateJetCorrectors["L1L2L3"] = ROOT.FactorizedJetCorrector(self.vParL3)
0058             if doResidualJECs and calculateSeparateCorrections:
0059                 self.vParL3Res = ROOT.vector(ROOT.JetCorrectorParameters)()
0060                 for i in [self.L1JetPar,self.L2JetPar,self.L3JetPar,self.ResJetPar]: self.vParL3Res.push_back(i)
0061                 self.separateJetCorrectors["L1L2L3Res"] = ROOT.FactorizedJetCorrector(self.vParL3Res)
0062 
0063     def getCorrection(self,jet,rho,delta=0,corrector=None):
0064         if not corrector: corrector = self.JetCorrector
0065         if corrector != self.JetCorrector and delta!=0: raise RuntimeError('Configuration not supported')
0066         corrector.setJetEta(jet.eta())
0067         corrector.setJetPt(jet.pt()*jet.rawFactor())
0068         corrector.setJetA(jet.jetArea())
0069         corrector.setRho(rho)
0070         corr = corrector.getCorrection()
0071         if delta != 0:
0072             if not self.JetUncertainty: raise RuntimeError("Jet energy scale uncertainty shifts requested, but not available")
0073             self.JetUncertainty.setJetEta(jet.eta())
0074             self.JetUncertainty.setJetPt(corr * jet.pt() * jet.rawFactor())
0075             try:
0076                 jet.jetEnergyCorrUncertainty = self.JetUncertainty.getUncertainty(True) 
0077             except RuntimeError as r:
0078                 print("Caught %s when getting uncertainty for jet of pt %.1f, eta %.2f\n" % (r,corr * jet.pt() * jet.rawFactor(),jet.eta()))
0079                 jet.jetEnergyCorrUncertainty = 0.5
0080             #print "   jet with corr pt %6.2f has an uncertainty %.2f " % (jet.pt()*jet.rawFactor()*corr, jet.jetEnergyCorrUncertainty)
0081             corr *= max(0, 1+delta*jet.jetEnergyCorrUncertainty)
0082         return corr
0083 
0084     def rawP4forType1MET_(self, jet):
0085         """Return the raw 4-vector, after subtracting the muons (if requested),
0086            or None if the jet fails the EMF cut."""
0087         p4 = jet.p4() * jet.rawFactor()
0088         emf = ( jet.physObj.neutralEmEnergy() + jet.physObj.chargedEmEnergy() )/p4.E()
0089         if emf > self.type1METParams['skipEMfractionThreshold']:
0090             return None
0091         if self.type1METParams['skipMuons']:
0092             for idau in range(jet.numberOfDaughters()):
0093                 pfcand = jet.daughter(idau)
0094                 if pfcand.isGlobalMuon() or pfcand.isStandAloneMuon(): 
0095                     p4 -= pfcand.p4()
0096         return p4
0097 
0098     def correct(self,jet,rho,delta=0,addCorr=False,addShifts=False, metShift=[0,0],type1METCorr=[0,0,0]):
0099         """Corrects a jet energy (optionally shifting it also by delta times the JEC uncertainty)
0100 
0101            If addCorr, set jet.corr to the correction.
0102            If addShifts, set also the +1 and -1 jet shifts 
0103 
0104            The metShift vector will accumulate the x and y changes to the MET from the JEC, i.e. the 
0105            negative difference between the new and old jet momentum, for jets eligible for type1 MET 
0106            corrections, and after subtracting muons. The pt cut is applied to the new corrected pt.
0107            This shift can be applied on top of the *OLD TYPE1 MET*, but only if there was no change 
0108            in the L1 corrections nor in the definition of the type1 MET (e.g. jet pt cuts).
0109 
0110            The type1METCorr vector, will accumulate the x, y, sumEt type1 MET corrections, to be
0111            applied to the *RAW MET* (if the feature was turned on in the constructor of the class).
0112         """
0113         raw = jet.rawFactor()
0114         corr = self.getCorrection(jet,rho,delta)
0115         if addCorr: 
0116             jet.corr = corr
0117             for sepcorr in self.separateJetCorrectors.keys():
0118                 setattr(jet,"CorrFactor_"+sepcorr,self.getCorrection(jet,rho,delta=0,corrector=self.separateJetCorrectors[sepcorr]))
0119         if addShifts:
0120             for cdelta,shift in [(1.0, "JECUp"), (-1.0, "JECDown")]:
0121                 cshift = self.getCorrection(jet,rho,delta+cdelta)
0122                 setattr(jet, "corr"+shift, cshift)
0123         if corr <= 0:
0124             return False
0125         newpt = jet.pt()*raw*corr
0126         if newpt > self.type1METParams['jetPtThreshold']:
0127             rawP4forT1 = self.rawP4forType1MET_(jet)
0128             if rawP4forT1 and rawP4forT1.Pt()*corr > self.type1METParams['jetPtThreshold']:
0129                 metShift[0] -= rawP4forT1.Px() * (corr - 1.0/raw)
0130                 metShift[1] -= rawP4forT1.Py() * (corr - 1.0/raw)
0131                 if self.calculateType1METCorr:
0132                     l1corr = self.getCorrection(jet,rho,delta=0,corrector=self.separateJetCorrectors["L1"])
0133                     #print "\tfor jet with raw pt %.5g, eta %.5g, dpx = %.5g, dpy = %.5g" % (
0134                     #            jet.pt()*raw, jet.eta(), 
0135                     #            rawP4forT1.Px()*(corr - l1corr), 
0136                     #            rawP4forT1.Py()*(corr - l1corr))
0137                     type1METCorr[0] -= rawP4forT1.Px() * (corr - l1corr) 
0138                     type1METCorr[1] -= rawP4forT1.Py() * (corr - l1corr) 
0139                     type1METCorr[2] += rawP4forT1.Et() * (corr - l1corr) 
0140         jet.setCorrP4(jet.p4() * (corr * raw))
0141         return True
0142 
0143     def correctAll(self,jets,rho,delta=0, addCorr=False, addShifts=False, metShift=[0.,0.], type1METCorr=[0.,0.,0.]):
0144         """Applies 'correct' to all the jets, discard the ones that have bad corrections (corrected pt <= 0)"""
0145         badJets = []
0146         if metShift     != [0.,0.   ]: raise RuntimeError("input metShift tuple is not initialized to zeros")
0147         if type1METCorr != [0.,0.,0.]: raise RuntimeError("input type1METCorr tuple is not initialized to zeros")
0148         for j in jets:
0149             ok = self.correct(j,rho,delta,addCorr=addCorr,addShifts=addShifts,metShift=metShift,type1METCorr=type1METCorr)
0150             if not ok: badJets.append(j)
0151         if len(badJets) > 0:
0152             print("Warning: %d out of %d jets flagged bad by JEC." % (len(badJets), len(jets)))
0153         for bj in badJets:
0154             jets.remove(bj)
0155 
0156 class Type1METCorrector:
0157     def __init__(self, old74XMiniAODs):
0158         """Object to apply type1 corrections computed by the JetReCalibrator to the MET.
0159            old74XMiniAODs should be True if using inputs produced with CMSSW_7_4_11 or earlier."""
0160         self.oldMC = old74XMiniAODs
0161     def correct(self,met,type1METCorrections):
0162         oldpx, oldpy = met.px(), met.py()
0163         #print "old met: px %+10.5f, py %+10.5f" % (oldpx, oldpy)
0164         if self.oldMC:
0165             raw  = met.shiftedP2_74x(12,0);
0166             rawsumet =  met.shiftedSumEt_74x(12,0);
0167         else:
0168             raw = met.uncorP2()
0169             rawsumet =  met.uncorSumEt();
0170         rawpx, rawpy = raw.px, raw.py
0171         #print "raw met: px %+10.5f, py %+10.5f" % (rawpx, rawpy)
0172         corrpx = rawpx + type1METCorrections[0]
0173         corrpy = rawpy + type1METCorrections[1]
0174         corrsumet = rawsumet  + type1METCorrections[2]
0175         #print "done met: px %+10.5f, py %+10.5f\n" % (corrpx,corrpy)
0176         met.setP4(ROOT.reco.Particle.LorentzVector(corrpx,corrpy,0,hypot(corrpx,corrpy)))
0177         ## workaround for missing setSumEt in reco::MET and pat::MET
0178         met._sumEt = corrsumet
0179         met.sumEt = types.MethodType(lambda myself : myself._sumEt, met, met.__class__) 
0180         if not self.oldMC:
0181             met.setCorShift(rawpx, rawpy, rawsumet, met.Raw)
0182         else: 
0183             # to avoid segfauls in pat::MET, I need a more ugly solution
0184             setFakeRawMETOnOldMiniAODs(met, rawpx,rawpy, rawsumet)
0185 
0186 def setFakeRawMETOnOldMiniAODs(met, rawpx, rawpy, rawsumet):
0187         met._rawSumEt = rawsumet
0188         met._rawP4    = ROOT.reco.Particle.LorentzVector(rawpx,rawpy,0,hypot(rawpx,rawpy))
0189         met.uncorPt  = types.MethodType(lambda myself : myself._rawP4.Pt(), met, met.__class__)
0190         met.uncorPx  = types.MethodType(lambda myself : myself._rawP4.Px(), met, met.__class__)
0191         met.uncorPy  = types.MethodType(lambda myself : myself._rawP4.Py(), met, met.__class__)
0192         met.uncorPhi = types.MethodType(lambda myself : myself._rawP4.Phi(), met, met.__class__)
0193         met.uncorP4  = types.MethodType(lambda myself : myself._rawP4, met, met.__class__)
0194         met.uncorSumEt = types.MethodType(lambda myself : myself._rawSumEt, met, met.__class__)
0195         # the two below are a bit more tricky, but probably less needed, but something dummy
0196         met.uncorP2 = types.MethodType(lambda myself : None, met, met.__class__)
0197         met.uncorP3 = types.MethodType(lambda myself : None, met, met.__class__)
0198