python/physicsutils/KalmanMuonCorrector.py

0001 import PhysicsTools.Heppy.loadlibs
0002 import ROOT
0003
0004 class KalmanMuonCorrector:
0005     def __init__(self, calibration, isMC, isSync=False, smearMode="none"):
0006         self.kamuca = ROOT.KalmanMuonCalibrator(calibration)
0007         self.isMC = isMC
0008         self.isSync = isSync
0009         self.smearMode = smearMode
0010     def correct(self, mu, run):
0011         newPt = self.kamuca.getCorrectedPt(mu.pt(), mu.eta(), mu.phi(), mu.charge())
0012         newPtErr = newPt * self.kamuca.getCorrectedError(newPt, mu.eta(), mu.ptErr()/newPt)
0013         if self.isMC: # new we do the smearing
0014             if self.isSync:
0015                 newPt = self.kamuca.smearForSync(newPt, mu.eta())
0016                 newPtErr = newPt * self.kamuca.getCorrectedErrorAfterSmearing(newPt, mu.eta(), newPtErr/newPt)
0017             elif self.smearMode == "none" or self.smearMode == None:
0018                 pass
0019             elif self.smearMode == "basic":
0020                 newPt = self.kamuca.smear(newPt, mu.eta())
0021                 newPtErr = newPt * self.kamuca.getCorrectedErrorAfterSmearing(newPt, mu.eta(), newPtErr/newPt)
0022             else:
0023                 newPt = self.kamuca.smearUsingEbE(newPt, mu.eta(), newPtErr/newPt)
0024                 newPtErr = newPt * self.kamuca.getCorrectedErrorAfterSmearing(newPt, mu.eta(), newPtErr/newPt)
0025         newP4 = ROOT.math.PtEtaPhiMLorentzVector(newPt, mu.eta(), mu.phi(), mu.mass())
0026         mu.setP4(newP4)
0027         mu._ptErr = newPtErr
0028
0029     def correct_all(self, mus, run):
0030         for mu in mus:
0031             self.correct(mu, run)
0032
0033 if __name__ == '__main__':
0034     kamuka = KalmanMuonCorrector("MC_76X_13TeV", True)