OfflineValidation/scripts/plotBaryCentre_VS_BeamSpot.py

0001 #!/usr/bin/env python3
0002
0003 import sys, os
0004 from array import array
0005 import optparse
0006 from collections import OrderedDict
0007 import json
0008
0009 import ROOT
0010 ROOT.gSystem.Load("libFWCoreFWLite.so")
0011
0012 import CondCore.Utilities.conddblib as conddb
0013
0014 # 1/lumiScaleFactor to go from 1/pb to 1/fb
0015 lumiScaleFactor = 1000
0016
0017 grootargs = []
0018 def callback_rootargs(option, opt, value, parser):
0019     grootargs.append(opt)
0020
0021 def vararg_callback(option, opt_str, value, parser):
0022     assert value is None
0023     value = []
0024
0025     def floatable(str):
0026         try:
0027             float(str)
0028             return True
0029         except ValueError:
0030             return False
0031
0032     for arg in parser.rargs:
0033         # stop on --foo like options
0034         if arg[:2] == "--" and len(arg) > 2:
0035             break
0036         # stop on -a, but not on -3 or -3.0
0037         if arg[:1] == "-" and len(arg) > 1 and not floatable(arg):
0038             break
0039         value.append(arg)
0040
0041     del parser.rargs[:len(value)]
0042     setattr(parser.values, option.dest, value)
0043
0044 def parseOptions():
0045     usage = ('usage: %prog [options]\n'
0046              + '%prog -h for help')
0047     parser = optparse.OptionParser(usage)
0048
0049     parser.add_option("--inputFileName", dest="inputFileName", default="PixelBaryCentre.root",help="name of the ntuple file that contains the barycentre tree")
0050     parser.add_option("--plotConfigFile", dest="plotConfigFile", default="PixelBaryCentrePlotConfig.json",help="json file that configs the plotting")
0051
0052     parser.add_option("--usePixelQuality",action="store_true", dest="usePixelQuality", default=False,help="whether use SiPixelQuality")
0053     parser.add_option("--showLumi",action="store_true", dest="showLumi", default=False,help="whether use integrated lumi as x-axis")
0054     parser.add_option("--years", dest="years", default = [2017], action="callback", callback=vararg_callback, help="years to plot")
0055
0056     parser.add_option("-l",action="callback",callback=callback_rootargs)
0057     parser.add_option("-q",action="callback",callback=callback_rootargs)
0058     parser.add_option("-b",action="callback",callback=callback_rootargs)
0059
0060     return parser
0061
0062
0063 def findRunIndex(run, runs) :
0064     #runs has to be sorted
0065     if(len(runs)==0) :
0066       print("Empty run list!")
0067       return -1
0068     elif(len(runs)==1) :
0069        if(run>=runs[0]) :
0070          return 0
0071        else :
0072          print("Only one run but the run requested is before the run!")
0073          return -1
0074     else :
0075        # underflow
0076        if(run <= runs[0]) :
0077           return 0
0078        # overflow
0079        elif(run >= runs[len(runs)-1]) :
0080           return len(runs)-1
0081        else :
0082           return ROOT.TAxis(len(runs)-1,array('d',runs)).FindBin(run) - 1
0083
0084
0085 def readBaryCentreAnalyzerTree(t, branch_names, accumulatedLumiPerRun, showLumi, isEOY) :
0086     # to store lumi sections info for each run
0087     run_maxlumi = {}
0088     run_lumis = {}
0089
0090     # y-axis of TGraph
0091     # runs for all years // integrated luminosity as a function of run
0092     runs = list(accumulatedLumiPerRun.keys())
0093     runs.sort()
0094
0095     current_run = 0
0096     for iov in t :
0097         # skip runs out-of-range
0098         if(iov.run>runs[len(runs)-1] or iov.run<runs[0]):
0099           continue
0100
0101         if(iov.run!=current_run) : # a new run, initialize lumi sections
0102           run_lumis[iov.run] = [iov.ls]
0103           run_maxlumi[iov.run] = iov.ls
0104         else : # current run, append lumi sections
0105           run_lumis[iov.run].append(iov.ls)
0106           if(run_maxlumi[iov.run]<iov.ls):
0107              run_maxlumi[iov.run] = iov.ls
0108         # update current run
0109         current_run = iov.run
0110
0111     # initialize store barycentre
0112     pos = {}
0113     for branch_name in branch_names :
0114         for coord in ["x","y","z"] :
0115             pos[coord+"_"+branch_name] = array('d',[])
0116             # max and min to determine the plotting range
0117             pos[coord+"max_"+branch_name] = -9999
0118             pos[coord+"min_"+branch_name] = 9999
0119
0120     # y-errors
0121     zeros = array('d',[])
0122
0123     # x-axis of TGraph
0124     runlumi = array('d',[])
0125     runlumiplot = array('d',[])
0126     runlumiplot_error = array('d',[])
0127
0128     max_run = 0
0129     # loop over IOVs
0130     for iov in t :
0131         # skip runs out-of-range
0132         if(iov.run>runs[len(runs)-1] or iov.run<runs[0]):
0133           continue
0134         # exclude 2018D for EOY rereco
0135         if(isEOY and iov.run>=320413 and iov.run<=325175):
0136           continue
0137
0138         # if x-axis is luminosity
0139         if(showLumi) :
0140           run_index = findRunIndex(iov.run,runs)
0141           instLumi = 0
0142           if(run_index==0) :
0143             instLumi = accumulatedLumiPerRun[ runs[run_index] ]
0144           if(run_index>0) :
0145             instLumi = accumulatedLumiPerRun[ runs[run_index] ] - accumulatedLumiPerRun[ runs[run_index-1] ]
0146           # remove runs with zero luminosity if x-axis is luminosity
0147           if( instLumi==0 ) : #and accumulatedLumiPerRun[ runs[run_index] ]==0 ) :
0148             continue
0149
0150           if(len(run_lumis[iov.run])>1) :  # lumi-based conditions
0151             if(run_index==0) :
0152               runlumi.append(0.0+instLumi*iov.ls*1.0/run_maxlumi[iov.run])
0153             else :
0154               runlumi.append(accumulatedLumiPerRun[ runs[run_index-1] ]+instLumi*iov.ls*1.0/run_maxlumi[iov.run])
0155
0156           else : # run-based or only one-IOV in the run
0157               runlumi.append(accumulatedLumiPerRun[ runs[run_index] ])
0158
0159         else: # else x-axis is run number
0160           if(len(run_lumis[iov.run])>1) :#lumi-based conditions
0161                runlumi.append(iov.run+iov.ls*1.0/run_maxlumi[iov.run])
0162
0163           else : # run-based or only one-IOV in the run
0164                runlumi.append(iov.run)
0165
0166         zeros.append(0)
0167
0168         #10000 is to translate cm to micro-metre
0169         for branch_name in branch_names :
0170             pos_ = {"x":10000*getattr(iov, branch_name).x(),
0171                     "y":10000*getattr(iov, branch_name).y(),
0172                     "z":10000*getattr(iov, branch_name).z()}
0173
0174             for coord in ["x","y","z"] :
0175                 pos[coord+"_"+branch_name].append(pos_[coord])
0176                 # max/min
0177                 if(pos_[coord]>pos[coord+"max_"+branch_name]) :
0178                    pos[coord+"max_"+branch_name] = pos_[coord]
0179                    max_run = iov.run
0180                 if(pos_[coord]<pos[coord+"min_"+branch_name]) :
0181                    pos[coord+"min_"+branch_name] = pos_[coord]
0182
0183     # x-axis : run/lumi or integrtated luminosity
0184     for iov in range(len(runlumi)-1) :
0185         runlumiplot.append(0.5*(runlumi[iov]+runlumi[iov+1]))
0186         runlumiplot_error.append(0.5*(runlumi[iov+1]-runlumi[iov]))
0187
0188     runlumiplot.append(runlumiplot[len(runlumiplot_error)-1]+2*runlumiplot_error[len(runlumiplot_error)-1])
0189     runlumiplot_error.append(runlumiplot_error[len(runlumiplot_error)-1])
0190
0191     v_runlumiplot = ROOT.TVectorD(len(runlumiplot),runlumiplot)
0192     v_runlumiplot_error = ROOT.TVectorD(len(runlumiplot_error),runlumiplot_error)
0193
0194     # y-axis error
0195     v_zeros = ROOT.TVectorD(len(zeros),zeros)
0196
0197     # store barycentre into a dict
0198     barryCentre = {}
0199     v_pos = {}
0200     for branch_name in branch_names :
0201         for coord in ["x","y","z"] :
0202             v_pos[coord] = ROOT.TVectorD(len(pos[coord+"_"+branch_name]),pos[coord+"_"+branch_name])
0203
0204             barryCentre[coord+'_'+branch_name] = ROOT.TGraphErrors(v_runlumiplot, v_pos[coord], v_runlumiplot_error, v_zeros)
0205             barryCentre['a_'+coord+'_'+branch_name] = pos[coord+"_"+branch_name]
0206
0207             barryCentre[coord+'max_'+branch_name] = pos[coord+"max_"+branch_name]
0208             barryCentre[coord+'min_'+branch_name] = pos[coord+"min_"+branch_name]
0209
0210     barryCentre['v_runlumiplot'] = v_runlumiplot
0211     barryCentre['v_runlumierror'] = v_runlumiplot_error
0212     barryCentre['v_zeros'] = v_zeros
0213
0214     return barryCentre
0215
0216
0217 def blackBox(x1, y1, x2, y2):
0218     x = array('d',[x1, x2, x2, x1, x1])
0219     y = array('d',[y1, y1, y2, y2, y1])
0220     v_x = ROOT.TVectorD(len(x),x)
0221     v_y = ROOT.TVectorD(len(y),y)
0222
0223     gr = ROOT.TGraph(v_x,v_y)
0224     gr.SetLineColor(ROOT.kBlack)
0225
0226     return gr
0227
0228
0229 def plotbarycenter(bc,coord,plotConfigJson, substructure,runsPerYear,pixelLocalRecos,accumulatedLumiPerRun, withPixelQuality,showLumi) :
0230     runs = list(accumulatedLumiPerRun.keys())
0231     runs.sort()
0232     years = list(runsPerYear.keys())
0233     years.sort()
0234     labels = list(bc.keys())
0235
0236     can = ROOT.TCanvas("barycentre_"+substructure+"_"+coord, "", 2000, 900)
0237     can.cd()
0238
0239     range_ = 0
0240     width_ = 0
0241     upper = 0
0242     lower = 0
0243     xmax = 0
0244
0245     gr = {}
0246     firstGraph = True
0247     for label in labels :
0248         gr[label] = ROOT.TGraph()
0249         gr[label] = bc[label][coord+"_"+substructure]
0250
0251         gr[label].SetMarkerStyle(8)
0252         gr[label].SetMarkerSize(0)
0253         gr[label].SetMarkerStyle(8)
0254         gr[label].SetMarkerSize(0)
0255         gr[label].SetLineColor(plotConfigJson["colorScheme"][label])
0256
0257         width_ = gr[label].GetXaxis().GetXmax() - gr[label].GetXaxis().GetXmin()
0258         xmax   = gr[label].GetXaxis().GetXmax()
0259
0260         if firstGraph :
0261            upper  = bc[label][coord+"max_"+substructure]
0262            lower  = bc[label][coord+"min_"+substructure]
0263            firstGraph = False
0264         else :
0265            upper = max(upper, bc[label][coord+"max_"+substructure])
0266            lower = min(lower, bc[label][coord+"min_"+substructure])
0267
0268     scale = 1.1
0269     if(upper>0) :
0270       upper = upper * scale
0271     else :
0272       upper = upper / scale
0273     if(lower>0) :
0274       lower = lower / scale
0275     else :
0276       lower = lower * scale
0277     range_ = upper - lower
0278
0279     firstGraph = True
0280     for label in labels :
0281         if(firstGraph) :
0282             gr[label].GetYaxis().SetRangeUser(lower, upper)
0283             gr[label].GetYaxis().SetTitle(plotConfigJson["substructures"][substructure]+" barycentre ("+coord+") [#mum]")
0284             gr[label].GetXaxis().SetTitle("Run Number")
0285             gr[label].GetYaxis().CenterTitle(True)
0286             gr[label].GetXaxis().CenterTitle(True)
0287             gr[label].GetYaxis().SetTitleOffset(0.80)
0288             gr[label].GetYaxis().SetTitleSize(0.055)
0289             gr[label].GetXaxis().SetTitleOffset(0.80)
0290             gr[label].GetXaxis().SetTitleSize(0.055)
0291             gr[label].GetXaxis().SetMaxDigits(6)
0292             if(showLumi) :
0293                 gr[label].GetXaxis().SetTitle("Delivered luminosity [1/fb]")
0294
0295             gr[label].Draw("AP")
0296             firstGraph = False
0297         else :
0298             gr[label].Draw("P")
0299
0300     # dummy TGraph for pixel local reco changes and first-of-year Run
0301     gr_dummyFirstRunOfTheYear = blackBox(-999, 10000, -999, -10000)
0302     gr_dummyFirstRunOfTheYear.SetLineColor(ROOT.kBlack)
0303     gr_dummyFirstRunOfTheYear.SetLineStyle(1)
0304     gr_dummyFirstRunOfTheYear.Draw("L")
0305     gr_dummyPixelReco = blackBox(-999, 10000, -999, -10000)
0306     gr_dummyPixelReco.SetLineColor(ROOT.kGray+1)
0307     gr_dummyPixelReco.SetLineStyle(3)
0308     gr_dummyPixelReco.Draw("L")
0309     gr_dummyFirstRunOfTheYear.SetTitle("First run of the year")
0310     gr_dummyPixelReco.SetTitle("Pixel calibration update")
0311
0312     for label in labels :
0313         gr[label].SetTitle(plotConfigJson["baryCentreLabels"][label])
0314     legend = can.BuildLegend()#0.65, 0.65, 0.85, 0.85)
0315     legend.SetShadowColor(0)
0316     legend.SetFillColor(0)
0317     legend.SetLineColor(1)
0318
0319     for label in labels :
0320         gr[label].SetTitle("")
0321
0322     # Add legends
0323     # and vertical lines
0324     years_label = ""
0325     for year in years :
0326         years_label += str(year)
0327         years_label += "+"
0328     years_label = years_label.rstrip("+")
0329
0330     # CMS logo
0331     CMSworkInProgress = ROOT.TPaveText( xmax-0.3*width_, upper+range_*0.005,
0332                                         xmax, upper+range_*0.055, "nb")
0333     CMSworkInProgress.AddText("CMS #bf{#it{Preliminary} ("+years_label+" pp collisions)}")
0334     CMSworkInProgress.SetTextAlign(32) #right/bottom aligned
0335     CMSworkInProgress.SetTextSize(0.04)
0336     CMSworkInProgress.SetFillColor(10)
0337     CMSworkInProgress.Draw()
0338
0339     # vertical lines
0340     #pixel local reco
0341     line_pixels = {}
0342     for since in pixelLocalRecos :
0343         if showLumi :
0344            run_index = findRunIndex(since,runs)
0345            integrated_lumi = accumulatedLumiPerRun[runs[run_index]]
0346            line_pixels[since] = ROOT.TLine(integrated_lumi, lower, integrated_lumi, upper)
0347
0348         else :
0349            line_pixels[since] = ROOT.TLine(since, lower, since, upper)
0350
0351         line_pixels[since].SetLineColor(ROOT.kGray+1)
0352         line_pixels[since].SetLineStyle(3)
0353         line_pixels[since].Draw()
0354
0355     # years
0356     line_years = {}
0357     box_years = {}
0358     text_years = {}
0359     if(len(years)>1 or (not showLumi) ) : # indicate begining of the year if more than one year to show or use run number
0360       for year in years :
0361           if showLumi :
0362              #first run of the year
0363              run_index = findRunIndex(runsPerYear[year][0],runs)
0364              integrated_lumi = accumulatedLumiPerRun[runs[run_index]]
0365              line_years[year] = ROOT.TLine(integrated_lumi, lower, integrated_lumi, upper)
0366              text_years[year] = ROOT.TPaveText( integrated_lumi+0.01*width_, upper-range_*0.05,
0367                                               integrated_lumi+0.05*width_,  upper-range_*0.015, "nb")
0368              box_years[year] = blackBox(integrated_lumi+0.005*width_, upper-range_*0.01, integrated_lumi+0.055*width_, upper-range_*0.055)
0369           else :
0370              line_years[year] = ROOT.TLine(runsPerYear[year][0], lower, runsPerYear[year][0], upper)
0371              text_years[year] = ROOT.TPaveText( runsPerYear[year][0]+0.01*width_, upper-range_*0.05,
0372                                               runsPerYear[year][0]+0.05*width_,  upper-range_*0.015, "nb")
0373              box_years[year] = blackBox(runsPerYear[year][0]+0.01*width_, upper-range_*0.015, runsPerYear[year][0]+0.05*width_, upper-range_*0.05)
0374
0375
0376           box_years[year].Draw("L")
0377           line_years[year].Draw()
0378
0379           # Add TextBox at the beginning of each year
0380           text_years[year].AddText(str(year))
0381           text_years[year].SetTextAlign(22)
0382           text_years[year].SetTextSize(0.025)
0383           text_years[year].SetFillColor(10)
0384           text_years[year].Draw()
0385
0386     #legend.Draw()
0387     can.Update()
0388
0389     if(showLumi) :
0390       can.SaveAs("baryCentre"+withPixelQuality+"_"+coord+"_"+substructure+"_"+years_label+"_IntegratedLumi.pdf")
0391       can.SaveAs("baryCentre"+withPixelQuality+"_"+coord+"_"+substructure+"_"+years_label+"_IntegratedLumi.png")
0392     else :
0393       can.SaveAs("baryCentre"+withPixelQuality+"_"+coord+"_"+substructure+"_"+years_label+"_RunNumber.pdf")
0394       can.SaveAs("baryCentre"+withPixelQuality+"_"+coord+"_"+substructure+"_"+years_label+"_RunNumber.png")
0395
0396     #####################################################################################################################
0397
0398
0399 # main call
0400 def Run():
0401
0402     #ROOT.gSystem.Load("libFWCoreFWLite.so")
0403     parser=parseOptions()
0404     (options,args) = parser.parse_args()
0405     sys.argv = grootargs
0406
0407     inputFileName = options.inputFileName
0408     if os.path.isfile(inputFileName) == False :
0409        print ("File "+inputFileName+" not exist!")
0410        return -1
0411
0412     plotConfigFile = open(options.plotConfigFile)
0413     plotConfigJson = json.load(plotConfigFile)
0414     plotConfigFile.close()
0415
0416     usePixelQuality = options.usePixelQuality
0417     withPixelQuality = ""
0418     if(usePixelQuality) :
0419        withPixelQuality = "WithPixelQuality"
0420     showLumi = options.showLumi
0421     # order years from old to new
0422     years = options.years
0423     years.sort()
0424
0425     # runs per year
0426     runsPerYear = {}
0427     for year in years :
0428         runsPerYear[year] = []
0429     # integrated lumi vs run
0430     accumulatedLumiPerRun = {}
0431
0432     run_index = 0
0433     lastRun = 1
0434
0435     # get lumi per IOV
0436     CMSSW_Dir = os.getenv("CMSSW_BASE")
0437     for year in years :
0438         inputLumiFile = CMSSW_Dir + "/src/Alignment/OfflineValidation/data/lumiperrun"+str(year)+".txt"
0439         if os.path.isfile(inputLumiFile) == False :
0440            print ("File "+inputLumiFile+" not exist!")
0441            return -1
0442         lumiFile = open(inputLumiFile,'r')
0443         lines = lumiFile.readlines()
0444
0445         for line in lines :
0446             # line = "run inst_lumi"
0447             run = int(line.split()[0])
0448             integrated_lumi = float(line.split()[1])/lumiScaleFactor # 1/pb to 1/fb
0449
0450             # runs per year
0451             runsPerYear[year].append(run)
0452             # integrated luminosity per run
0453             # run number must be ordered from small to large in the text file
0454             if(run_index == 0) :
0455               accumulatedLumiPerRun[run] = integrated_lumi
0456             else :
0457               accumulatedLumiPerRun[run] = accumulatedLumiPerRun[lastRun]+integrated_lumi
0458
0459             run_index+=1
0460             lastRun = run
0461
0462         # close file
0463         lumiFile.close()
0464
0465     # order by key (year)
0466     runsPerYear = OrderedDict(sorted(runsPerYear.items(), key=lambda t: t[0]))
0467     # order by key (run number)
0468     accumulatedLumiPerRun = OrderedDict(sorted(accumulatedLumiPerRun.items(), key=lambda t: t[0]))
0469
0470     #pixel local reco update (IOVs/sinces)
0471     pixelLocalRecos = []
0472     # connnect to ProdDB to access pixel local reco condition change
0473     db = plotConfigJson["pixelDataBase"]
0474     pixel_template = plotConfigJson["pixelLocalReco"]
0475     db = db.replace("sqlite_file:", "").replace("sqlite:", "")
0476     db = db.replace("frontier://FrontierProd/CMS_CONDITIONS", "pro")
0477     db = db.replace("frontier://FrontierPrep/CMS_CONDITIONS", "dev")
0478
0479     con = conddb.connect(url = conddb.make_url(db))
0480     session = con.session()
0481     # get IOV table
0482     IOV = session.get_dbtype(conddb.IOV)
0483     iovs = set(session.query(IOV.since).filter(IOV.tag_name == pixel_template).all())
0484     session.close()
0485     pixelLocalRecos = sorted([int(item[0]) for item in iovs])
0486     #pixelLocalRecos = [1, 186500, 195360, 197749, 200961, 203368, 204601, 206446, 238341, 246866, 253914, 255655, 271866, 276315, 278271, 280928, 290543, 297281, 298653, 299443, 300389, 301046, 302131, 303790, 303998, 304911, 313041, 314881, 316758, 317475, 317485, 317527, 317661, 317664, 318227, 320377, 321831, 322510, 322603, 323232, 324245]
0487
0488     # substructures to plot
0489     substructures = list(plotConfigJson["substructures"].keys())
0490
0491     # start barycentre plotter
0492     bc = {}
0493     try:
0494        f = ROOT.TFile(inputFileName,"READ")
0495        # read TTrees
0496        for label in list(plotConfigJson["baryCentreLabels"].keys()) :
0497            isEOY = False
0498            t = ROOT.TTree()
0499            if label == "" :
0500               t = f.Get("PixelBaryCentreAnalyzer"+withPixelQuality+"/PixelBarycentre")
0501            else :
0502               t = f.Get("PixelBaryCentreAnalyzer"+withPixelQuality+"/PixelBarycentre_"+label)
0503               if(label=="EOY") :
0504                  isEOY = True
0505
0506            bc[label] = readBaryCentreAnalyzerTree(t, substructures, accumulatedLumiPerRun, showLumi, isEOY)
0507
0508     except IOError:
0509        print("File "+inputFileName+" not accessible")
0510
0511     # plot
0512     for substructure in substructures :
0513         for coord in ['x','y','z'] :
0514             plotbarycenter(bc,coord,plotConfigJson,substructure, runsPerYear,pixelLocalRecos,accumulatedLumiPerRun, withPixelQuality,showLumi)
0515
0516
0517 if __name__ == "__main__":
0518     Run()