NanoAOD/test/compare_branches.py

0001 #!/usr/bin/env python3
0002 from ROOT import TFile, TTree,TCut, gROOT,TH1F, TCanvas
0003 import glob, re
0004 from optparse import OptionParser
0005 parser = OptionParser(usage="%prog [options] job1 job2 ...")
0006 parser.add_option("--base", dest="base", default="step2.root", help="name of root file")
0007 parser.add_option("--ref", dest="ref", default="ref/", help="path to the reference files")
0008 parser.add_option("--png", dest="png", default="./", help="path to the plots")
0009 parser.add_option("--onlydiff", dest="diff", default=False, action="store_true", help="print only the histograms with differences")
0010 parser.add_option("--selection", dest="selection", default="", help="a selection of events to draw")
0011 parser.add_option("--branch", dest="branch", default=".*", help="a regexp for selecting branches")
0012 (options, args) = parser.parse_args()
0013
0014 def drawInOrder(t1,h1n, t2,h2n):
0015     N1 = t1.Draw(f'{vname} >> {h1n}', options.selection, '')
0016     h1=gROOT.Get(h1n)
0017     #print(h1)
0018     h2=TH1F(h2n,'',
0019             h1.GetNbinsX(),
0020             h1.GetXaxis().GetXmin(),
0021             h1.GetXaxis().GetXmax())
0022     h2=gROOT.Get(h2n)
0023     #print(h2)
0024     if t2:
0025         N2 = t2.Draw(f'{vname} >> {h2n}', options.selection, '')
0026     #print(f'created histograms {h1.GetName()} and {h2.GetName()}')
0027     #print(binContents(h1))
0028     #print(binContents(h2))
0029     return (h1,h2)
0030
0031 def binContents( h ):
0032     return [ h.GetBinContent(b) for b in range(-1,h.GetNbinsX()+1)]
0033 def emptyHist( h ):
0034     return all([ b==0 for b in binContents(h) ])
0035
0036 for job in args:
0037     fconn = glob.glob(f'{job}*/{options.base}')[0]
0038     fcon = TFile.Open(fconn)
0039     con = fcon.Get('Events')
0040     frefn=glob.glob(f'{options.ref}/{job}*/{options.base}')[0]
0041     fref = TFile.Open(frefn)
0042     ref = fref.Get('Events')
0043
0044     print(f'Comparing branch content from {fconn} and {frefn}')
0045     print(f'{ref.GetEntries()} events in reference tree and {con.GetEntries()} events in comparison tree')
0046
0047     ## get outside of the files
0048     gROOT.cd()
0049
0050     branches = [b.GetName() for b in con.GetListOfBranches()]
0051     ref_branches = [b.GetName() for b in ref.GetListOfBranches()]
0052
0053     for vname in sorted(set(branches)|set(ref_branches)):
0054         hrefn = f'ref_{vname}'
0055         hconn = f'con_{vname}'
0056
0057         if not re.match(options.branch, vname):
0058             #print(vname,"not in regexp")
0059             continue
0060
0061         no_base=no_ref=False
0062         if not vname in branches:
0063             print(f'Branch {vname} in not in file to compare')
0064             no_base=True
0065         if not vname in ref_branches:
0066             print(f'Branch {vname} in not in file to use as reference')
0067             no_ref=True
0068
0069         if no_ref:
0070             hcon,href=drawInOrder(con, hconn,
0071                         None, hrefn)
0072         elif no_base:
0073             href,hcon=drawInOrder(ref, hrefn,
0074                         None, hconn)
0075         else:
0076             href,hcon=drawInOrder(ref, hrefn,
0077                         con, hconn)
0078
0079         href.SetLineColor(1)
0080         hcon.SetLineColor(2)
0081
0082         ##ROOT is too stupid for me at this point. If I dont do this, the histograms are empty
0083         r = sum(binContents(href) + binContents(hcon))
0084
0085         ##make the difference
0086         hdiff= TH1F(f'diff_{vname}',
0087                     '',
0088                     href.GetNbinsX(),
0089                     href.GetXaxis().GetXmin(),
0090                     href.GetXaxis().GetXmax())
0091
0092         hdiff.Add(hcon)
0093         hdiff.Add(href, -1 )
0094
0095         ymax = max(href.GetMaximum(), hcon.GetMaximum())
0096         ymin = min(href.GetMinimum(), hcon.GetMinimum(), hdiff.GetMinimum())
0097
0098         hdiff.SetMarkerColor(4)
0099         hdiff.SetLineColor(4)
0100         hdiff.SetMarkerStyle(7)
0101
0102
0103         if options.diff and emptyHist(hdiff):
0104             #print(f'No changes for branch {vname}')
0105             continue
0106
0107         ## show them all
0108         fig = TCanvas(f'c_{vname}', f'plots for {vname}')
0109         href.Draw("he")
0110         hcon.Draw("same he")
0111         hdiff.Draw("same p e")
0112         fig.Print(f'{options.png}/{job}_{fig.GetName()}.png')
0113         fig = TCanvas(f'd_{vname}', f'Difference plot for {vname}')
0114         hdiff.Draw("hpe")
0115         fig.Print(f'{options.png}/{job}_{fig.GetName()}.png')