Back to home page

Project CMSSW displayed by LXR

 
 

    


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)