Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2021-07-23 02:25:52

0001 from builtins import range
0002 from PhysicsTools.Heppy.analyzers.core.TreeAnalyzerNumpy import TreeAnalyzerNumpy
0003 from PhysicsTools.Heppy.analyzers.core.AutoHandle import AutoHandle
0004 #from ROOT import TriggerBitChecker
0005 from PhysicsTools.Heppy.analyzers.core.autovars import *
0006 from PhysicsTools.Heppy.analyzers.objects.autophobj  import *
0007 
0008 
0009 class AutoFillTreeProducer( TreeAnalyzerNumpy ):
0010 
0011     #-----------------------------------
0012     # BASIC TREE PRODUCER 
0013     #-----------------------------------
0014     def __init__(self, cfg_ana, cfg_comp, looperName):
0015         super(AutoFillTreeProducer,self).__init__(cfg_ana, cfg_comp, looperName)
0016 
0017         ## Read whether we want vectors or flat trees
0018         self.scalar = not self.cfg_ana.vectorTree
0019 
0020         ## Read whether we want 4-vectors 
0021         if not getattr(self.cfg_ana, 'saveTLorentzVectors', False):
0022             fourVectorType.removeVariable("p4")
0023 
0024 
0025         self.collections = {}      
0026         self.globalObjects = {}
0027         self.globalVariables = []
0028         if hasattr(cfg_ana,"collections"):
0029             self.collections.update(cfg_ana.collections)
0030         if hasattr(cfg_ana,"globalObjects"):
0031             self.globalObjects.update(cfg_ana.globalObjects)
0032         if hasattr(cfg_ana,"globalVariables"):
0033             self.globalVariables=cfg_ana.globalVariables[:]
0034 
0035     def beginLoop(self, setup) :
0036         super(AutoFillTreeProducer, self).beginLoop(setup)
0037 
0038     def declareHandles(self):
0039         super(AutoFillTreeProducer, self).declareHandles()
0040 #        self.handles['TriggerResults'] = AutoHandle( ('TriggerResults','','HLT'), 'edm::TriggerResults' )
0041         self.mchandles['GenInfo'] = AutoHandle( ('generator','',''), 'GenEventInfoProduct' )
0042         for k,v in self.collections.items():
0043             if isinstance(v, tuple) and isinstance(v[0], AutoHandle):
0044                 self.handles[k] = v[0]
0045 
0046     def declareCoreVariables(self, tr, isMC):
0047         """Here we declare the variables that we always want and that are hard-coded"""
0048         tr.var('run', int, storageType="i")
0049         tr.var('lumi', int, storageType="i")
0050         tr.var('evt', int, storageType="l")
0051         tr.var('isData', int)
0052 
0053  #       self.triggerBitCheckers = []
0054  #       if hasattr(self.cfg_ana, 'triggerBits'):
0055  #           for T, TL in self.cfg_ana.triggerBits.items():
0056  #               trigVec = ROOT.vector(ROOT.string)()
0057  #               for TP in TL:
0058  #                   trigVec.push_back(TP)
0059  #               tr.var( 'HLT_'+T, int )
0060 #                self.triggerBitCheckers.append( (T, TriggerBitChecker(trigVec)) )
0061 
0062         if not isMC:
0063             tr.var('intLumi', int, storageType="i")
0064 
0065         if isMC:
0066             ## cross section
0067             tr.var('xsec', float)
0068             ## PU weights
0069             tr.var("puWeight")
0070             ## number of true interactions
0071             tr.var("nTrueInt")
0072             ## generator weight
0073             tr.var("genWeight")
0074             ## PDF weights
0075             self.pdfWeights = []
0076             if hasattr(self.cfg_ana, "PDFWeights") and len(self.cfg_ana.PDFWeights) > 0:
0077                 self.pdfWeights = self.cfg_ana.PDFWeights
0078                 for (pdf,nvals) in self.pdfWeights:
0079                     if self.scalar:
0080                         for i in range(nvals): tr.var('pdfWeight_%s_%d' % (pdf,i))
0081                     else:
0082                         tr.vector('pdfWeight_%s' % pdf, nvals)
0083 
0084     def declareVariables(self,setup):
0085         isMC = self.cfg_comp.isMC 
0086         tree = self.tree
0087         self.declareCoreVariables(tree, isMC)
0088 
0089         if not hasattr(self.cfg_ana,"ignoreAnalyzerBookings") or not self.cfg_ana.ignoreAnalyzerBookings :
0090             #import variables declared by the analyzers
0091             if hasattr(setup,"globalVariables"):
0092                 self.globalVariables+=setup.globalVariables
0093             if hasattr(setup,"globalObjects"):
0094                 self.globalObjects.update(setup.globalObjects)
0095             if hasattr(setup,"collections"):
0096                 self.collections.update(setup.collections)
0097 
0098         for v in self.globalVariables:
0099             v.makeBranch(tree, isMC)
0100         for o in self.globalObjects.values(): 
0101             o.makeBranches(tree, isMC)
0102         for c in self.collections.values():
0103             if isinstance(c, tuple): c = c[-1]
0104             if self.scalar:
0105                 c.makeBranchesScalar(tree, isMC)
0106             else:
0107                 c.makeBranchesVector(tree, isMC)
0108 
0109     def fillCoreVariables(self, tr, event, isMC):
0110         """Here we fill the variables that we always want and that are hard-coded"""
0111         tr.fill('run', event.input.eventAuxiliary().id().run())
0112         tr.fill('lumi',event.input.eventAuxiliary().id().luminosityBlock())
0113         tr.fill('evt', event.input.eventAuxiliary().id().event())    
0114         tr.fill('isData', 0 if isMC else 1)
0115 
0116 #       triggerResults = self.handles['TriggerResults'].product()
0117 #       for T,TC in self.triggerBitCheckers:
0118 #           tr.fill("HLT_"+T, TC.check(event.object(), triggerResults))
0119 
0120         if not isMC:
0121             tr.fill('intLumi', getattr(self.cfg_comp,'intLumi',1.0))
0122 
0123         if isMC:
0124             ## xsection, if available
0125             tr.fill('xsec', getattr(self.cfg_comp,'xSection',1.0))
0126             ## PU weights, check if a PU analyzer actually filled it
0127             if hasattr(event,"nPU"):
0128                 tr.fill("nTrueInt", event.nPU)
0129                 tr.fill("puWeight", event.puWeight)
0130             else :
0131                 tr.fill("nTrueInt", -1)
0132                 tr.fill("puWeight", 1.0)
0133 
0134             tr.fill("genWeight", self.mchandles['GenInfo'].product().weight())
0135             ## PDF weights
0136             if hasattr(event,"pdfWeights") :
0137                 for (pdf,nvals) in self.pdfWeights:
0138                     if len(event.pdfWeights[pdf]) != nvals:
0139                         raise RuntimeError("PDF lenght mismatch for %s, declared %d but the event has %d" % (pdf,nvals,event.pdfWeights[pdf]))
0140                     if self.scalar:
0141                         for i,w in enumerate(event.pdfWeights[pdf]):
0142                             tr.fill('pdfWeight_%s_%d' % (pdf,i), w)
0143                     else:
0144                         tr.vfill('pdfWeight_%s' % pdf, event.pdfWeights[pdf])
0145 
0146     def process(self, event):
0147         if hasattr(self.cfg_ana,"filter") : 
0148             if not self.cfg_ana.filter(event) :
0149                 return True #do not stop processing, just filter myself
0150         self.readCollections( event.input)
0151         self.fillTree(event)
0152 
0153     def fillTree(self, event, resetFirst=True):
0154         isMC = self.cfg_comp.isMC 
0155         if resetFirst: self.tree.reset()
0156 
0157         self.fillCoreVariables(self.tree, event, isMC)
0158 
0159         for v in self.globalVariables:
0160             if not isMC and v.mcOnly: continue
0161             v.fillBranch(self.tree, event, isMC)
0162 
0163         for on, o in self.globalObjects.items(): 
0164             if not isMC and o.mcOnly: continue
0165             o.fillBranches(self.tree, getattr(event, on), isMC)
0166 
0167         for cn, c in self.collections.items():
0168             if isinstance(c, tuple) and isinstance(c[0], AutoHandle):
0169                 if not isMC and c[-1].mcOnly: continue
0170                 objects = self.handles[cn].product()
0171                 setattr(event, cn, [objects[i] for i in range(objects.size())])
0172                 c = c[-1]
0173             if not isMC and c.mcOnly: continue
0174             if self.scalar:
0175                 c.fillBranchesScalar(self.tree, getattr(event, cn), isMC)
0176             else:
0177                 c.fillBranchesVector(self.tree, getattr(event, cn), isMC)
0178 
0179         self.tree.tree.Fill()      
0180 
0181     def getPythonWrapper(self):
0182         """
0183         This function produces a string that contains a Python wrapper for the event.
0184         The wrapper is automatically generated based on the collections and allows the full
0185         event contents to be accessed from subsequent Analyzers using e.g.
0186 
0187         leps = event.selLeptons #is of type selLeptons
0188         pt0 = leps[0].pt
0189 
0190         One just needs to add the EventAnalyzer to the sequence.
0191         """
0192 
0193         isMC = self.cfg_comp.isMC 
0194 
0195         classes = ""
0196         anclass = ""
0197         anclass += "from PhysicsTools.HeppyCore.framework.analyzer import Analyzer\n"
0198         anclass += "class EventAnalyzer(Analyzer):\n"
0199         anclass += "    def __init__(self, cfg_ana, cfg_comp, looperName):\n"
0200         anclass += "        super(EventAnalyzer, self).__init__(cfg_ana, cfg_comp, looperName)\n"
0201 
0202         anclass += "    def process(self, event):\n"
0203 
0204         for cname, coll in self.collections.items():
0205             classes += coll.get_py_wrapper_class(isMC)
0206             anclass += "        event.{0} = {0}.make_array(event)\n".format(coll.name)
0207 
0208         return classes + "\n" + anclass
0209