analyzers/core/AutoFillTreeProducer.py

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