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
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
