File indexing completed on 2024-04-06 12:23:26
0001 from builtins import range
0002 from PhysicsTools.Heppy.analyzers.core.TreeAnalyzerNumpy import TreeAnalyzerNumpy
0003 from PhysicsTools.Heppy.analyzers.core.AutoHandle import AutoHandle
0004
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
0013
0014 def __init__(self, cfg_ana, cfg_comp, looperName):
0015 super(AutoFillTreeProducer,self).__init__(cfg_ana, cfg_comp, looperName)
0016
0017
0018 self.scalar = not self.cfg_ana.vectorTree
0019
0020
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
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
0054
0055
0056
0057
0058
0059
0060
0061
0062 if not isMC:
0063 tr.var('intLumi', int, storageType="i")
0064
0065 if isMC:
0066
0067 tr.var('xsec', float)
0068
0069 tr.var("puWeight")
0070
0071 tr.var("nTrueInt")
0072
0073 tr.var("genWeight")
0074
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
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
0117
0118
0119
0120 if not isMC:
0121 tr.fill('intLumi', getattr(self.cfg_comp,'intLumi',1.0))
0122
0123 if isMC:
0124
0125 tr.fill('xsec', getattr(self.cfg_comp,'xSection',1.0))
0126
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
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
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