NanoAOD/test/prepareDQM.py

0001 #!/usr/bin/env python3
0002
0003 from __future__ import print_function
0004 import sys, os
0005 from collections import defaultdict
0006 from math import ceil, pi, log
0007 import ROOT
0008 ROOT.PyConfig.IgnoreCommandLineOptions = True
0009 ROOT.gROOT.SetBatch(True)
0010
0011 from PhysicsTools.NanoAOD.nanoDQM_cfi import nanoDQM, NoPlot, Plot1D, Count1D, shortDump
0012
0013 from optparse import OptionParser
0014 parser = OptionParser(usage="%prog [options] inputFile")
0015 parser.add_option("-u", dest="update", action="store_true", default=False, help="Update nanoDQM_cfi.py directly")
0016 parser.add_option("-d", dest="delete", action="store_true", default=False, help="Delete lines corresponding to missing plots, instead of commenting them out")
0017 parser.add_option("-q", dest="verbose", action="store_false", default=True, help="Quiet mode")
0018 parser.add_option("-o", dest="out", type="string", default="newDQM.py", help="Output file (unless -u is used)")
0019 (options, args) = parser.parse_args()
0020 if options.update: options.out = "%s/src/PhysicsTools/NanoAOD/python/nanoDQM_cfi.py" % os.environ["CMSSW_BASE"]
0021
0022 infile = args[0]
0023 if not os.path.isfile(infile): raise RuntimeError
0024
0025 class Branch:
0026     def __init__(self, branch):
0027         self.name = branch.GetName()
0028         self.title = branch.GetTitle()
0029         self.counter = None
0030         self.good = True
0031         if branch.GetNleaves() != 1:
0032             sys.stderr.write("Cannot parse branch '%s' in tree %s (%d leaves)\n", tree.GetName(), branch.GetName(), branch.GetNleaves())
0033             self.good = False
0034             return
0035         self.leaf = branch.FindLeaf(branch.GetName())
0036         if not self.leaf:
0037             sys.stderr.write("Cannot parse branch '%s' in tree %s (no leaf)\n", tree.GetName(), branch.GetName())
0038             self.good = False
0039             return
0040         self.kind = self.leaf.GetTypeName()
0041         if self.leaf.GetLen() == 0 and self.leaf.GetLeafCount() != None:
0042             self.counter = self.leaf.GetLeafCount().GetName()
0043
0044 class BranchGroup:
0045     def __init__(self, name):
0046         self.name = name
0047         self.subs = []
0048     def append(self, sub):
0049         self.subs.append(sub)
0050
0051 tfile = ROOT.TFile.Open(infile)
0052 tree = tfile.Get("Events")
0053 allbranches = [ Branch(b) for b in tree.GetListOfBranches() ]
0054 allbranches = [ b for b in allbranches if b.good ]
0055 branchmap = dict((b.name,b) for b in allbranches )
0056 branchgroups = defaultdict(dict)
0057 iscounter = {}
0058 for b in allbranches:
0059     if b.counter: iscounter[b.counter] = True
0060 for b in allbranches:
0061     if b.name in iscounter: continue
0062     if "_" in b.name:
0063         head, tail = b.name.split("_",1)
0064         if head == "HLT": continue # no monitoring for these
0065     else:
0066         head, tail = b.name, ''
0067     branchgroups[head][tail] = b
0068     if b.counter and ('@size' not in branchgroups[head]):
0069         branchgroups[head]['@size'] = branchmap[b.counter]
0070
0071 pyout = open(options.out, "w")
0072 pyout.write("""# automatically generated by prepareDQM.py
0073 import FWCore.ParameterSet.Config as cms
0074 from PhysicsTools.NanoAOD.nanoDQM_tools_cff import *
0075
0076 from DQMServices.Core.DQMEDAnalyzer import DQMEDAnalyzer
0077 nanoDQM = DQMEDAnalyzer("NanoAODDQM",
0078     vplots = cms.PSet(
0079 """);
0080
0081 c1 = ROOT.TCanvas("c1","c1")
0082 def smartRound(x):
0083     if abs(x) < 1e-7: return 0
0084     if x < 0: return -smartRound(-x)
0085     shift = pow(10,ceil(log(x,10))-1)
0086     xu = x/shift;
0087     if (xu > 30): xu = 5*ceil(xu/5)
0088     else: xu = ceil(xu)
0089     return xu*shift;
0090 def autoPlot1D(name, col, branch):
0091     if branch.kind == "Bool_t":
0092         return Plot1D(name, col, 2, -0.5, 1.5)
0093     xmin, xmax = tree.GetMinimum(branch.name), tree.GetMaximum(branch.name)
0094     if col == "@size":
0095         return Count1D(name, int(xmax+1 if xmax < 40 else 40), -0.5, xmax+0.5)
0096     if branch.kind == "Int_t" and "Idx" in name:
0097         return NoPlot(name)
0098     if branch.kind != "Float_t":
0099         xmin, xmax = map(int, (xmin,xmax))
0100         if xmax-xmin < 20:
0101             return Plot1D(name, col, xmax-xmin+1, xmin-0.5, xmax+0.5)
0102     elif name == "phi":
0103         return Plot1D(name, col, 20, -pi, pi)
0104     if xmin < 0 and xmax > 0:
0105         xmin, xmax = min(xmin, -xmax), max(xmax, -xmin) # symmetrize
0106     elif xmax > 0 and xmin/xmax < 0.03:
0107         xmin = 0
0108     xmin, xmax = map(smartRound, (xmin,xmax))
0109     return Plot1D(name, col, 20, xmin, xmax)
0110
0111 for head in sorted(branchgroups.iterkeys()):
0112     pset = getattr(nanoDQM.vplots, head, None)
0113     if not pset:
0114         if options.verbose: print("%s <skipped as it's not in vplots>" % head)
0115         continue
0116     if options.verbose: print("%s" % head)
0117     vpset = pset.plots
0118     allplots = [ (x.name.value(),x) for x in vpset ]
0119     existing = set(x[0] for x in allplots)
0120     found  = set()
0121     title  = dict( (n,x.title.value()) for (n,x) in allplots if x.kind.value() != "none" and x.title.value() )
0122     for (t,branch) in sorted(branchgroups[head].items()):
0123         t_noat = t.replace("@","_")
0124         found.add(t_noat)
0125         if t_noat not in title: title[t_noat] = branch.title
0126         if t_noat not in existing: allplots.append( (t_noat,autoPlot1D(t_noat,t,branch)) )
0127     # update titles
0128     for k,v in allplots:
0129         if k in title and title[k] and v.kind.value() != "none": v.title = title[k]
0130     allplots.sort()
0131     pyout.write("        %s = cms.PSet(\n" % head)
0132     seldump = pset.sels.dumpPython().replace("\n","\n            ") if len(pset.sels.parameterNames_()) else "cms.PSet()"
0133     pyout.write("            sels = %s,\n" % seldump)
0134     pyout.write("            plots = cms.VPSet(\n")
0135     for k,v in allplots:
0136         if k not in found and options.delete: continue
0137         pyout.write("              %s %s,\n" % (" " if k in found else "#", shortDump(v)));
0138     pyout.write("            )\n")
0139     pyout.write("        ),\n")
0140 pyout.write("    )\n)\n");
0141