python/physicsutils/MuScleFitCorrector.py

0001 import os
0002 from ROOT import heppy, TLorentzVector
0003
0004 class MuScleFitCorr:
0005     def __init__(self,isMC,isReReco,isSync=False):
0006         #colin need to import muscle fit inputs, if tool still in use
0007         path = "%s/src/CMGTools/RootTools/data/musclefit/" % os.environ['CMSSW_BASE'];
0008         self.isMC = isMC
0009         if self.isMC:
0010             self.corr = heppy.MuScleFitCorrector(path+"MuScleFit_2012_MC_53X_smear%s.txt" % ("ReReco" if isReReco else "Prompt"))
0011             self.isSync = isSync
0012         else:
0013             self.corrABC = heppy.MuScleFitCorrector(path+"MuScleFit_2012ABC_DATA%s_53X.txt" % ("_ReReco" if isReReco else ""))
0014             self.corrD   = heppy.MuScleFitCorrector(path+"MuScleFit_2012D_DATA%s_53X.txt"   % ("_ReReco" if isReReco else ""))
0015     def corrected_p4(self, mu, run):
0016         p4 = TLorentzVector(mu.px(), mu.py(), mu.pz(), mu.energy())
0017         if self.isMC:
0018             self.corr.applyPtCorrection(p4, mu.charge())
0019             self.corr.applyPtSmearing(p4, mu.charge(), self.isSync)
0020         else:
0021             corr = self.corrD if run >= 203773 else self.corrABC
0022             corr.applyPtCorrection(p4, mu.charge())
0023         ## convert to the proper C++ class (but preserve the mass!)
0024         return ROOT.reco.Muon.PolarLorentzVector( p4.Pt(), p4.Eta(), p4.Phi(), mu.mass() )
0025
0026     def correct(self, mu, run):
0027         mu.setP4( self.corrected_p4(mu, run) )
0028
0029     def correct_all(self, mus, run):
0030         for mu in mus:
0031             mu.setP4( self.corrected_p4(mu, run) )
0032
0033 if __name__ == '__main__':
0034     muscle = MuScleFitCorr(True, True)