Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2024-11-25 02:29:49

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