File indexing completed on 2024-04-06 12:23:46
0001 from PhysicsTools.NanoAODTools.postprocessing.framework.datamodel import Event
0002 from PhysicsTools.NanoAODTools.postprocessing.framework.treeReaderArrayTools import clearExtraBranches
0003 import sys
0004 import time
0005 import ROOT
0006
0007
0008 class Module(object):
0009 def __init__(self):
0010 self.writeHistFile = False
0011
0012 def beginJob(self, histFile=None, histDirName=None):
0013 if histFile != None and histDirName != None:
0014 self.writeHistFile = True
0015 prevdir = ROOT.gDirectory
0016 self.histFile = histFile
0017 self.histFile.cd()
0018 self.dir = self.histFile.mkdir(histDirName, "", True)
0019 prevdir.cd()
0020 self.objs = []
0021
0022 def endJob(self):
0023 if hasattr(self, 'objs') and self.objs != None:
0024 prevdir = ROOT.gDirectory
0025 self.dir.cd()
0026 for obj in self.objs:
0027 obj.Write()
0028 prevdir.cd()
0029
0030 def beginFile(self, inputFile, outputFile, inputTree, wrappedOutputTree):
0031 pass
0032
0033 def endFile(self, inputFile, outputFile, inputTree, wrappedOutputTree):
0034 pass
0035
0036 def analyze(self, event):
0037 """process event, return True (go to next module) or False (fail, go to next event)"""
0038 pass
0039
0040 def addObject(self, obj):
0041 setattr(self, obj.GetName(), obj)
0042 self.objs.append(getattr(self, obj.GetName()))
0043
0044 def addObjectList(self, names, obj):
0045 objlist = []
0046 for iname, name in enumerate(names):
0047 setattr(self, obj.GetName() + '_' + name,
0048 obj.Clone(obj.GetName() + '_' + name))
0049 objlist.append(getattr(self, obj.GetName() + '_' + name))
0050 self.objs.append(getattr(self, obj.GetName() + '_' + name))
0051 setattr(self, obj.GetName(), objlist)
0052
0053
0054 def eventLoop(
0055 modules, inputFile, outputFile, inputTree, wrappedOutputTree,
0056 maxEvents=-1, eventRange=None, progress=(10000, sys.stdout),
0057 filterOutput=True
0058 ):
0059 for m in modules:
0060 m.beginFile(inputFile, outputFile, inputTree, wrappedOutputTree)
0061
0062 t0 = time.time()
0063 tlast = t0
0064 doneEvents = 0
0065 acceptedEvents = 0
0066 entries = inputTree.entries
0067 if eventRange:
0068 entries = len(eventRange)
0069 if maxEvents > 0:
0070 entries = min(entries, maxEvents)
0071
0072 for ie, i in enumerate(range(entries) if eventRange == None else eventRange):
0073 if maxEvents > 0 and ie >= maxEvents:
0074 break
0075 e = Event(inputTree, i)
0076 clearExtraBranches(inputTree)
0077 doneEvents += 1
0078 ret = True
0079 for m in modules:
0080 ret = m.analyze(e)
0081 if not ret:
0082 break
0083 if ret:
0084 acceptedEvents += 1
0085 if (ret or not filterOutput) and wrappedOutputTree != None:
0086 wrappedOutputTree.fill()
0087 if progress:
0088 if ie > 0 and ie % progress[0] == 0:
0089 t1 = time.time()
0090 progress[1].write("Processed %8d/%8d entries, %5.2f%% (elapsed time %7.1fs, curr speed %8.3f kHz, avg speed %8.3f kHz), accepted %8d/%8d events (%5.2f%%)\n" % (
0091 ie, entries, ie / float(0.01 * entries),
0092 t1 - t0, (progress[0] / 1000.) / (max(t1 - tlast, 1e-9)),
0093 ie / 1000. / (max(t1 - t0, 1e-9)),
0094 acceptedEvents, doneEvents,
0095 acceptedEvents / (0.01 * doneEvents)))
0096 tlast = t1
0097 for m in modules:
0098 m.endFile(inputFile, outputFile, inputTree, wrappedOutputTree)
0099 return (doneEvents, acceptedEvents, time.time() - t0)