Merge pull request cms-sw#621 from gpetruc/heppy_KaMuCa_76X

Heppy Kalman Muon Corrections for 76X

PhysicsTools/Heppy/python/analyzers/objects/LeptonAnalyzer.py

@@ -7,6 +7,7 @@
 from PhysicsTools.HeppyCore.utils.deltar import bestMatch
 from PhysicsTools.Heppy.physicsutils.RochesterCorrections import rochcor
 from PhysicsTools.Heppy.physicsutils.MuScleFitCorrector   import MuScleFitCorr
+from PhysicsTools.Heppy.physicsutils.KalmanMuonCorrector   import KalmanMuonCorrector
 from PhysicsTools.Heppy.physicsutils.ElectronCalibrator import Run2ElectronCalibrator
 #from CMGTools.TTHAnalysis.electronCalibrator import ElectronCalibrator
 import PhysicsTools.HeppyCore.framework.config as cfg
@@ -28,7 +29,13 @@ def __init__(self, cfg_ana, cfg_comp, looperName ):
             raise RuntimeError, "doRochesterCorrections is not supported. Please set instead doMuonScaleCorrections = ( 'Rochester', <name> )"
         if self.cfg_ana.doMuonScaleCorrections:
             algo, options = self.cfg_ana.doMuonScaleCorrections
-            if algo == "Rochester":
+            if algo == "Kalman":
+                corr = options['MC' if self.cfg_comp.isMC else 'Data']
+                self.muonScaleCorrector = KalmanMuonCorrector(corr, 
+                                                    self.cfg_comp.isMC,
+                                                    options['isSync'] if 'isSync' in options else False,
+                                                    options['smearMode'] if 'smearMode' in options else "ebe")
+            elif algo == "Rochester":
                 print "WARNING: the Rochester correction in heppy is still from Run 1"
                 self.muonScaleCorrector = RochesterCorrections()
             elif algo == "MuScleFit":

PhysicsTools/Heppy/python/physicsobjects/Lepton.py

@@ -8,7 +8,8 @@ def ip3D(self):
 
     def sip3D(self):
         '''3D impact parameter significance.'''
-        return abs(self.dB(self.PV3D) / self.edB(self.PV3D))
+        db, edb = self.dB(self.PV3D), self.edB(self.PV3D)
+        return abs(db/edb) if edb > 0 else 999.
 
     def absIsoFromEA(self, area='04'):
         '''Calculate Isolation using the effective area approach.'''

PhysicsTools/Heppy/python/physicsutils/KalmanMuonCorrector.py

@@ -0,0 +1,34 @@
+import PhysicsTools.Heppy.loadlibs
+import ROOT
+
+class KalmanMuonCorrector:
+    def __init__(self, calibration, isMC, isSync=False, smearMode="none"):
+        self.kamuca = ROOT.KalmanMuonCalibrator(calibration)
+        self.isMC = isMC
+        self.isSync = isSync
+        self.smearMode = smearMode
+    def correct(self, mu, run):
+        newPt = self.kamuca.getCorrectedPt(mu.pt(), mu.eta(), mu.phi(), mu.charge())
+        newPtErr = newPt * self.kamuca.getCorrectedError(newPt, mu.eta(), mu.ptErr()/newPt)
+        if self.isMC: # new we do the smearing
+            if self.isSync:
+                newPt = self.kamuca.smearForSync(newPt, mu.eta())
+                newPtErr = newPt * self.kamuca.getCorrectedErrorAfterSmearing(newPt, mu.eta(), newPtErr/newPt)
+            elif self.smearMode == "none" or self.smearMode == None:
+                pass
+            elif self.smearMode == "basic":
+                newPt = self.kamuca.smear(newPt, mu.eta())
+                newPtErr = newPt * self.kamuca.getCorrectedErrorAfterSmearing(newPt, mu.eta(), newPtErr/newPt)
+            else:
+                newPt = self.kamuca.smearUsingEbE(newPt, mu.eta(), newPtErr/newPt)
+                newPtErr = newPt * self.kamuca.getCorrectedErrorAfterSmearing(newPt, mu.eta(), newPtErr/newPt)
+        newP4 = ROOT.math.PtEtaPhiMLorentzVector(newPt, mu.eta(), mu.phi(), mu.mass())
+        mu.setP4(newP4)
+        mu._ptErr = newPtErr
+
+    def correct_all(self, mus, run):
+        for mu in mus:
+            self.correct(mu, run)
+
+if __name__ == '__main__':
+    kamuka = KalmanMuonCorrector("MC_76X_13TeV", True)