Heppy/scripts/heppy_report.py

0001 #!/usr/bin/env python3
0002 from __future__ import print_function
0003 from builtins import range
0004 from optparse import OptionParser
0005 import json
0006
0007 def root2map(dir,ana,treename):
0008     import ROOT
0009     tfile = ROOT.TFile.Open("%s/%s/%s.root"%(dir,ana,treename))
0010     if not tfile:
0011         print("Error: dir %s does not contain %s/%s.root" % (dir,ana,treename))
0012         return None
0013     tree = tfile.Get(treename)
0014     if not tree:
0015         print("Error: rootfile %s/%s/%s.root does not contain a TTree %s" % (dir,ana,treename,treename))
0016         return None
0017     jsonind = {}
0018     for e in range(tree.GetEntries()):
0019         tree.GetEntry(e)
0020         run,lumi = tree.run, tree.lumi
0021         if run not in jsonind:
0022             jsonind[run] = [lumi]
0023         else:
0024             jsonind[run].append(lumi)
0025     # remove duplicates
0026     for run in jsonind:
0027         jsonind[run] =  list(set(jsonind[run]))
0028
0029     nruns = len(jsonind)
0030     nlumis = sum(len(v) for v in jsonind.values())
0031     jsonmap = {}
0032     for r,lumis in jsonind.items():
0033         if len(lumis) == 0: continue # shouldn't happen
0034         lumis.sort()
0035         ranges = [ [ lumis[0], lumis[0] ] ]
0036         for lumi in lumis[1:]:
0037             if lumi == ranges[-1][1] + 1:
0038                 ranges[-1][1] = lumi
0039             else:
0040                 ranges.append([lumi,lumi])
0041         jsonmap[r] = ranges
0042     return (jsonmap, nruns, nlumis)
0043
0044 if __name__ == '__main__':
0045     parser = OptionParser(usage='%prog <target_directories> [options]',
0046                           description='Check the output of the JSONAnalyzer and produce a json file of the processed runs and lumisections')
0047     parser.add_option("-a", "--analyzer", dest="jsonAnalyzer", default="JSONAnalyzer", help="Name of the JSONAnalyzer")
0048     parser.add_option("-t", "--tree", dest="treeName", default="RLTInfo", help="Name of the TTree produced by the JSONAnalyzer")
0049     parser.add_option("-o", "--out", dest="outputFile", default="lumiSummary.json", help="Name of the output file")
0050     (options,args) = parser.parse_args()
0051     if len(args)==0:
0052         print('provide at least one directory in argument. Use -h to display help')
0053         exit()
0054     for a in args:
0055         summary = root2map(a,options.jsonAnalyzer,options.treeName)
0056         if summary:
0057             oname = "%s/%s" % (a,options.outputFile)
0058             jmap, runs, lumis = summary
0059             json.dump(jmap,open(oname,'w'))
0060             print("Saved %s (%d runs, %d lumis)" % (oname, runs, lumis))