Merge pull request #23 from arizzi/jpata-heppycore-tth-vars

Improvements on gen analyzers from ttH-Hbb  people

PhysicsTools/Heppy/python/analyzers/core/AutoFillTreeProducer.py

@@ -170,4 +170,33 @@ def fillTree(self, event, resetFirst=True):
                 c.fillBranchesVector(self.tree, getattr(event, cn), isMC)
 
         self.tree.tree.Fill()      
+
+    def getPythonWrapper(self):
+        """
+        This function produces a string that contains a Python wrapper for the event.
+        The wrapper is automatically generated based on the collections and allows the full
+        event contents to be accessed from subsequent Analyzers using e.g.
+        
+        leps = event.selLeptons #is of type selLeptons
+        pt0 = leps[0].pt
+
+        One just needs to add the EventAnalyzer to the sequence.
+        """
+
+        isMC = self.cfg_comp.isMC 
+
+        classes = ""
+        anclass = ""
+        anclass += "from PhysicsTools.HeppyCore.framework.analyzer import Analyzer\n"
+        anclass += "class EventAnalyzer(Analyzer):\n"
+        anclass += "    def __init__(self, cfg_ana, cfg_comp, looperName):\n"
+        anclass += "        super(EventAnalyzer, self).__init__(cfg_ana, cfg_comp, looperName)\n"
+
+        anclass += "    def process(self, event):\n"
+
+        for cname, coll in self.collections.items():
+            classes += coll.get_py_wrapper_class(isMC)
+            anclass += "        event.{0} = {0}.make_array(event)\n".format(coll.name)
+
+        return classes + "\n" + anclass
 

PhysicsTools/Heppy/python/analyzers/core/TreeAnalyzerNumpy.py

@@ -27,8 +27,8 @@ def beginLoop(self, setup) :
             print 'Compression', isCompressed
             self.file = TFile( fileName, 'recreate', '', isCompressed )
         self.file.cd()
-	if self.file.Get(self.treename) :
-		raise RuntimeError, "You are booking two Trees with the same name in the same file"
+        if self.file.Get(self.treename) :
+            raise RuntimeError, "You are booking two Trees with the same name in the same file"
         self.tree = Tree(self.treename, self.name)
         self.tree.setDefaultFloatType(getattr(self.cfg_ana, 'defaultFloatType','D')); # or 'F'
         self.declareVariables(setup)

PhysicsTools/Heppy/python/analyzers/core/autovars.py

@@ -146,4 +146,32 @@ def fillBranchesVector(self,treeNumpy,collection,isMC):
     def __repr__(self):
         return "<NTupleCollection[%s]>" % self.name
 
+    def get_cpp_declaration(self, isMC):
+        s = []
+        for v in self.objectType.allVars(isMC):
+            s += ["{0} {1}__{2}[{3}];".format(v.type.__name__, self.name, v.name, self.maxlen)]
+        return "\n".join(s)
+
+    def get_cpp_wrapper_class(self, isMC):
+        s = "class %s {\n" % self.name
+        s += "public:\n"
+        for v in self.objectType.allVars(isMC):
+            s += "    {0} {1};\n".format(v.type.__name__, v.name)
+        s += "};\n"
+        return s
+
+    def get_py_wrapper_class(self, isMC):
+        s = "class %s:\n" % self.name
+        s += "    def __init__(self, tree, n):\n"
+        for v in self.objectType.allVars(isMC):
+            if len(v.name)>0:
+                s += "        self.{0} = tree.{1}_{2}[n];\n".format(v.name, self.name, v.name)
+            else:
+                s += "        self.{0} = tree.{0}[n];\n".format(self.name)
+
+        s += "    @staticmethod\n"
+        s += "    def make_array(event):\n"
+        s += "        return [{0}(event.input, i) for i in range(event.input.n{0})]\n".format(self.name)
+        return s
+
 

PhysicsTools/Heppy/python/analyzers/gen/GeneratorAnalyzer.py

@@ -40,6 +40,7 @@ class GeneratorAnalyzer( Analyzer ):
             event.genwzquarks   = [] # quarks from W,Z decays
             event.genbquarksFromTop = []
             event.genbquarksFromH   = []
+            event.genlepsFromTop = [] #mu/ele that have a t->W chain as ancestor, also contained in event.genleps
        event.genwzquarks and event.genbquarks, might have overlaps 
        event.genbquarksFromTop and event.genbquarksFromH are all contained in event.genbquarks
        
@@ -180,6 +181,7 @@ def makeMCInfo(self, event):
             event.genwzquarks    = []
             event.genbquarksFromTop = []
             event.genbquarksFromH   = []
+            event.genlepsFromTop = []
             for p in event.generatorSummary:
                 id = abs(p.pdgId())
                 if id == 25: 
@@ -189,10 +191,23 @@ def makeMCInfo(self, event):
                 elif id in {12,14,16}:
                     event.gennus.append(p)
                 elif id in {11,13}:
+                    #taus to separate vector
                     if abs(p.motherId) == 15:
                         event.gentauleps.append(p)
+                    #all muons and electrons
                     else:
                         event.genleps.append(p)
+                        momids = [(m, abs(m.pdgId())) for m in realGenMothers(p)]
+
+                        #have a look at the lepton mothers
+                        for mom, momid in momids:
+                            #lepton from W
+                            if momid == 24:
+                                wmomids = [abs(m.pdgId()) for m in realGenMothers(mom)]
+                                #W from t
+                                if 6 in wmomids:
+                                    #save mu,e from t->W->mu/e
+                                    event.genlepsFromTop.append(p)
                 elif id == 15:
                     if self.allGenTaus or not any([abs(d.pdgId()) in {11,13} for d in realGenDaughters(p)]):
                         event.gentaus.append(p)

PhysicsTools/Heppy/python/analyzers/gen/LHEAnalyzer.py

@@ -24,6 +24,10 @@ def process(self, event):
             return True
         event.lheHT=0
         event.lheNj=0
+        event.lheNb=0
+        event.lheNc=0
+        event.lheNl=0
+        event.lheNg=0
         event.lheV_pt = 0
         try:
           event.input.getByLabel( 'externalLHEProducer',self.lheh)
@@ -46,6 +50,14 @@ def process(self, event):
           if status == 1 and ( ( idabs == 21 ) or (idabs > 0 and idabs < 7) ) : # gluons and quarks
               event.lheHT += sqrt( pup[i][0]**2 + pup[i][1]**2 ) # first entry is px, second py
               event.lheNj +=1
+              if idabs==5:
+                event.lheNb += 1
+              if idabs==4:
+                event.lheNc += 1
+              if idabs in [1,2,3]:
+                event.lheNl += 1
+              if idabs==21:
+                event.lheNg += 1
           if idabs in [12,14,16] :  
               if id > 0 :
                 nu = i

PhysicsTools/Heppy/python/analyzers/gen/PDFWeightsAnalyzer.py

@@ -8,6 +8,7 @@ class PDFWeightsAnalyzer( Analyzer ):
     def __init__(self, cfg_ana, cfg_comp, looperName ):
         super(PDFWeightsAnalyzer,self).__init__(cfg_ana,cfg_comp,looperName)
         self.doPDFWeights = hasattr(self.cfg_ana, "PDFWeights") and len(self.cfg_ana.PDFWeights) > 0
+        self.doPDFVars = hasattr(self.cfg_ana, "doPDFVars") and self.cfg_ana.doPDFVars == True
         if self.doPDFWeights:
             self.pdfWeightInit = False
     #---------------------------------------------
@@ -18,7 +19,7 @@ def __init__(self, cfg_ana, cfg_comp, looperName ):
     def declareHandles(self):
         super(PDFWeightsAnalyzer, self).declareHandles()
 
-        if self.doPDFWeights:
+        if self.doPDFVars or self.doPDFWeights:
             self.mchandles['pdfstuff'] = AutoHandle( 'generator', 'GenEventInfoProduct' )
 
     def beginLoop(self, setup):
@@ -34,7 +35,7 @@ def initPDFWeights(self):
 
     def makePDFWeights(self, event):
         if not self.pdfWeightInit: self.initPDFWeights()
-        self.pdfWeightTool.processEvent(self.mchandles['pdfstuff'].product())
+        self.pdfWeightTool.processEvent(self.genInfo)
         event.pdfWeights = {}
         for pdf in self.cfg_ana.PDFWeights:
             ws = self.pdfWeightTool.getWeights(pdf+".LHgrid")
@@ -47,8 +48,17 @@ def process(self, event):
         if not self.cfg_comp.isMC: 
             return True
 
+        if self.doPDFVars or self.doPDFWeights:
+            self.genInfo = self.mchandles['pdfstuff'].product()
         if self.doPDFWeights:
             self.makePDFWeights(event)
+        if self.doPDFVars:
+            event.pdf_x1 = self.genInfo.pdf().x.first
+            event.pdf_x2 = self.genInfo.pdf().x.second
+            event.pdf_id1 = self.genInfo.pdf().id.first
+            event.pdf_id2 = self.genInfo.pdf().id.second
+            event.pdf_scale = self.genInfo.pdf().scalePDF
+
         return True
 
 setattr(PDFWeightsAnalyzer,"defaultConfig",

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

@@ -14,6 +14,8 @@
 import ROOT
 import math
 
+from copy import deepcopy
+
 class METAnalyzer( Analyzer ):
     def __init__(self, cfg_ana, cfg_comp, looperName ):
         super(METAnalyzer,self).__init__(cfg_ana,cfg_comp,looperName)
@@ -96,6 +98,20 @@ def makeMETs(self, event):
         event.met = self.handles['met'].product()[0]
         event.metNoPU = self.handles['nopumet'].product()[0]
 
+        #Shifted METs
+        #Uncertainties defined in https://github.com/cms-sw/cmssw/blob/CMSSW_7_2_X/DataFormats/PatCandidates/interface/MET.h#L168
+        #event.met_shifted = []
+        for i in range(14):
+            m = deepcopy(event.met)
+            px  = m.shiftedPx(i);
+            py  = m.shiftedPy(i);
+            m.setP4(ROOT.reco.Particle.LorentzVector(px,py, 0, math.hypot(px,py)))
+            #event.met_shifted += [m]
+            setattr(event, "met_shifted_{0}".format(i), m)
+        event.met_sig = event.met.significance()
+        event.met_sumet = event.met.sumEt()
+        #event.met_sigm = event.met.getSignificanceMatrix()
+
         ###https://github.com/cms-sw/cmssw/blob/CMSSW_7_2_X/DataFormats/PatCandidates/interface/MET.h
         event.metraw = event.met.shiftedPt(12, 0)
         event.metType1chs = event.met.shiftedPt(12, 1)

PhysicsTools/Heppy/python/utils/cmsswRelease.py

@@ -2,7 +2,9 @@
 import re
 
 def cmsswRelease():
-    return os.environ['CMSSW_BASE'].split('/')[-1]
+    #return os.environ['CMSSW_BASE'].split('/')[-1]
+    #this also works when the CMSSW directory is renamed
+    return os.environ['CMSSW_VERSION']
 
 def cmsswIs44X():
     return cmsswRelease().find('CMSSW_4_4_') != -1