Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2025-02-07 23:29:40

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