postprocessing/framework/preskimming.py

0001 import json
0002 import re
0003 import ROOT
0004 ROOT.PyConfig.IgnoreCommandLineOptions = True
0005
0006
0007 class JSONFilter:
0008     def __init__(self, fname="", runsAndLumis={}):
0009         self.keep = {}
0010         if fname != "":
0011             self.runsAndLumis = json.load(open(fname, 'r'))
0012         else:
0013             self.runsAndLumis = runsAndLumis
0014         for _run, lumis in self.runsAndLumis.items():
0015             run = int(_run)
0016             if run not in self.keep:
0017                 self.keep[run] = []
0018             self.keep[run] += lumis
0019         for run in list(self.keep.keys()):
0020             if len(self.keep[run]) == 0:
0021                 del self.keep[run]
0022
0023     def filterRunLumi(self, run, lumi):
0024         try:
0025             for (l1, l2) in self.keep[run]:
0026                 if l1 <= lumi and lumi <= l2:
0027                     return True
0028             return False
0029         except KeyError:
0030             return False
0031
0032     def filterRunOnly(self, run):
0033         return (run in self.keep)
0034
0035     def runCut(self):
0036         return "%d <= run && run <= %s" % (min(self.keep.keys()), max(self.keep.keys()))
0037
0038     def filterEList(self, tree, elist):
0039         # FIXME this can be optimized for sure
0040         tree.SetBranchStatus("*", 0)
0041         tree.SetBranchStatus('run', 1)
0042         tree.SetBranchStatus('luminosityBlock', 1)
0043         filteredList = ROOT.TEntryList('filteredList', 'filteredList')
0044         if elist:
0045             for i in range(elist.GetN()):
0046                 entry = elist.GetEntry(0) if i == 0 else elist.Next()
0047                 tree.GetEntry(entry)
0048                 if self.filterRunLumi(tree.run, tree.luminosityBlock):
0049                     filteredList.Enter(entry)
0050         else:
0051             for entry in range(tree.GetEntries()):
0052                 tree.GetEntry(entry)
0053                 if self.filterRunLumi(tree.run, tree.luminosityBlock):
0054                     filteredList.Enter(entry)
0055         tree.SetBranchStatus("*", 1)
0056         return filteredList
0057
0058
0059 def preSkim(tree, jsonInput=None, cutstring=None, maxEntries=None, firstEntry=0):
0060     if jsonInput == None and cutstring == None:
0061         return None, None
0062     cut = None
0063     jsonFilter = None
0064     if jsonInput != None:
0065         if type(jsonInput) is dict:
0066             jsonFilter = JSONFilter(runsAndLumis=jsonInput)
0067         else:
0068             jsonFilter = JSONFilter(jsonInput)
0069         cut = jsonFilter.runCut()
0070     if cutstring != None:
0071         cut = "(%s) && (%s)" % (cutstring, cut) if cut else cutstring
0072     if maxEntries is None:
0073         maxEntries = ROOT.TVirtualTreePlayer.kMaxEntries
0074     while "AltBranch$" in cut:
0075         m = re.search(r"AltBranch\$\(\s*(\w+)\s*,\s*(\w+)\s*\)", cut)
0076         if not m:
0077             raise RuntimeError(
0078                 "Error, found AltBranch$ in cut string, but it doesn't comply with the syntax this code can support. The cut is %r" % cut)
0079         cut = cut.replace(m.group(0), m.group(
0080             1) if tree.GetBranch(m.group(1)) else m.group(2))
0081     tree.Draw('>>elist', cut, "entrylist", maxEntries, firstEntry)
0082     elist = ROOT.gDirectory.Get('elist')
0083     if jsonInput:
0084         elist = jsonFilter.filterEList(tree, elist)
0085     return elist, jsonFilter