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