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 import Alignment.OfflineValidation.TkAlAllInOneTool.findAndChange as fnc
0014
0015 # 1/lumiScaleFactor to go from 1/pb to 1/fb
0016 lumiScaleFactor = 1000
0017
0018 grootargs = []
0019 def callback_rootargs(option, opt, value, parser):
0020     grootargs.append(opt)
0021
0022 def vararg_callback(option, opt_str, value, parser):
0023     assert value is None
0024     value = []
0025
0026     def floatable(str):
0027         try:
0028             float(str)
0029             return True
0030         except ValueError:
0031             return False
0032
0033     for arg in parser.rargs:
0034         # stop on --foo like options
0035         if arg[:2] == "--" and len(arg) > 2:
0036             break
0037         # stop on -a, but not on -3 or -3.0
0038         if arg[:1] == "-" and len(arg) > 1 and not floatable(arg):
0039             break
0040         value.append(arg)
0041
0042     del parser.rargs[:len(value)]
0043     setattr(parser.values, option.dest, value)
0044
0045 def parseOptions():
0046     usage = ('usage: %prog [options]\n'
0047              + '%prog -h for help')
0048     parser = optparse.OptionParser(usage)
0049
0050     parser.add_option("--inputFileName", dest="inputFileName", default="PixelBaryCentre.root",help="name of the ntuple file that contains the barycentre tree")
0051     parser.add_option("--plotConfigFile", dest="plotConfigFile", default="PixelBaryCentrePlotConfig.json",help="json file that configs the plotting")
0052
0053     parser.add_option("--usePixelQuality",action="store_true", dest="usePixelQuality", default=False,help="whether use SiPixelQuality")
0054     parser.add_option("--showLumi",action="store_true", dest="showLumi", default=False,help="whether use integrated lumi as x-axis")
0055     parser.add_option("--years", dest="years", default = [2017], action="callback", callback=vararg_callback, help="years to plot")
0056
0057     parser.add_option("-l",action="callback",callback=callback_rootargs)
0058     parser.add_option("-q",action="callback",callback=callback_rootargs)
0059     parser.add_option("-b",action="callback",callback=callback_rootargs)
0060
0061     return parser
0062
0063
0064 def findRunIndex(run, runs) :
0065     #runs has to be sorted
0066     if(len(runs)==0) :
0067       print("Empty run list!")
0068       return -1
0069     elif(len(runs)==1) :
0070        if(run>=runs[0]) :
0071          return 0
0072        else :
0073          print("Only one run but the run requested is before the run!")
0074          return -1
0075     else :
0076        # underflow
0077        if(run <= runs[0]) :
0078           return 0
0079        # overflow
0080        elif(run >= runs[len(runs)-1]) :
0081           return len(runs)-1
0082        else :
0083           return ROOT.TAxis(len(runs)-1,array('d',runs)).FindBin(run) - 1
0084
0085
0086 def readBaryCentreAnalyzerTree(t, branch_names, accumulatedLumiPerRun, showLumi, isEOY) :
0087     # to store lumi sections info for each run
0088     run_maxlumi = {}
0089     run_lumis = {}
0090
0091     # y-axis of TGraph
0092     # runs for all years // integrated luminosity as a function of run
0093     runs = list(accumulatedLumiPerRun.keys())
0094     runs.sort()
0095
0096     current_run = 0
0097     for iov in t :
0098         # skip runs out-of-range
0099         if(iov.run>runs[len(runs)-1] or iov.run<runs[0]):
0100           continue
0101
0102         if(iov.run!=current_run) : # a new run, initialize lumi sections
0103           run_lumis[iov.run] = [iov.ls]
0104           run_maxlumi[iov.run] = iov.ls
0105         else : # current run, append lumi sections
0106           run_lumis[iov.run].append(iov.ls)
0107           if(run_maxlumi[iov.run]<iov.ls):
0108              run_maxlumi[iov.run] = iov.ls
0109         # update current run
0110         current_run = iov.run
0111
0112     # initialize store barycentre
0113     pos = {}
0114     for branch_name in branch_names :
0115         for coord in ["x","y","z"] :
0116             pos[coord+"_"+branch_name] = array('d',[])
0117             # max and min to determine the plotting range
0118             pos[coord+"max_"+branch_name] = -9999
0119             pos[coord+"min_"+branch_name] = 9999
0120
0121     # y-errors
0122     zeros = array('d',[])
0123
0124     # x-axis of TGraph
0125     runlumi = array('d',[])
0126     runlumiplot = array('d',[])
0127     runlumiplot_error = array('d',[])
0128
0129     max_run = 0
0130     # loop over IOVs
0131     for iov in t :
0132         # skip runs out-of-range
0133         if(iov.run>runs[len(runs)-1] or iov.run<runs[0]):
0134           continue
0135         # exclude 2018D for EOY rereco
0136         if(isEOY and iov.run>=320413 and iov.run<=325175):
0137           continue
0138
0139         # if x-axis is luminosity
0140         if(showLumi) :
0141           run_index = findRunIndex(iov.run,runs)
0142           instLumi = 0
0143           if(run_index==0) :
0144             instLumi = accumulatedLumiPerRun[ runs[run_index] ]
0145           if(run_index>0) :
0146             instLumi = accumulatedLumiPerRun[ runs[run_index] ] - accumulatedLumiPerRun[ runs[run_index-1] ]
0147           # remove runs with zero luminosity if x-axis is luminosity
0148           if( instLumi==0 ) : #and accumulatedLumiPerRun[ runs[run_index] ]==0 ) :
0149             continue
0150
0151           if(len(run_lumis[iov.run])>1) :  # lumi-based conditions
0152             if(run_index==0) :
0153               runlumi.append(0.0+instLumi*iov.ls*1.0/run_maxlumi[iov.run])
0154             else :
0155               runlumi.append(accumulatedLumiPerRun[ runs[run_index-1] ]+instLumi*iov.ls*1.0/run_maxlumi[iov.run])
0156
0157           else : # run-based or only one-IOV in the run
0158               runlumi.append(accumulatedLumiPerRun[ runs[run_index] ])
0159
0160         else: # else x-axis is run number
0161           if(len(run_lumis[iov.run])>1) :#lumi-based conditions
0162                runlumi.append(iov.run+iov.ls*1.0/run_maxlumi[iov.run])
0163
0164           else : # run-based or only one-IOV in the run
0165                runlumi.append(iov.run)
0166
0167         zeros.append(0)
0168
0169         #10000 is to translate cm to micro-metre
0170         for branch_name in branch_names :
0171             pos_ = {"x":10000*getattr(iov, branch_name).x(),
0172                     "y":10000*getattr(iov, branch_name).y(),
0173                     "z":10000*getattr(iov, branch_name).z()}
0174
0175             for coord in ["x","y","z"] :
0176                 pos[coord+"_"+branch_name].append(pos_[coord])
0177                 # max/min
0178                 if(pos_[coord]>pos[coord+"max_"+branch_name]) :
0179                    pos[coord+"max_"+branch_name] = pos_[coord]
0180                    max_run = iov.run
0181                 if(pos_[coord]<pos[coord+"min_"+branch_name]) :
0182                    pos[coord+"min_"+branch_name] = pos_[coord]
0183
0184     # x-axis : run/lumi or integrtated luminosity
0185     for iov in range(len(runlumi)-1) :
0186         runlumiplot.append(0.5*(runlumi[iov]+runlumi[iov+1]))
0187         runlumiplot_error.append(0.5*(runlumi[iov+1]-runlumi[iov]))
0188
0189     runlumiplot.append(runlumiplot[len(runlumiplot_error)-1]+2*runlumiplot_error[len(runlumiplot_error)-1])
0190     runlumiplot_error.append(runlumiplot_error[len(runlumiplot_error)-1])
0191
0192     v_runlumiplot = ROOT.TVectorD(len(runlumiplot),runlumiplot)
0193     v_runlumiplot_error = ROOT.TVectorD(len(runlumiplot_error),runlumiplot_error)
0194
0195     # y-axis error
0196     v_zeros = ROOT.TVectorD(len(zeros),zeros)
0197
0198     # store barycentre into a dict
0199     barryCentre = {}
0200     v_pos = {}
0201     for branch_name in branch_names :
0202         for coord in ["x","y","z"] :
0203             v_pos[coord] = ROOT.TVectorD(len(pos[coord+"_"+branch_name]),pos[coord+"_"+branch_name])
0204
0205             barryCentre[coord+'_'+branch_name] = ROOT.TGraphErrors(v_runlumiplot, v_pos[coord], v_runlumiplot_error, v_zeros)
0206             barryCentre['a_'+coord+'_'+branch_name] = pos[coord+"_"+branch_name]
0207
0208             barryCentre[coord+'max_'+branch_name] = pos[coord+"max_"+branch_name]
0209             barryCentre[coord+'min_'+branch_name] = pos[coord+"min_"+branch_name]
0210
0211     barryCentre['v_runlumiplot'] = v_runlumiplot
0212     barryCentre['v_runlumierror'] = v_runlumiplot_error
0213     barryCentre['v_zeros'] = v_zeros
0214
0215     return barryCentre
0216
0217
0218 def blackBox(x1, y1, x2, y2):
0219     x = array('d',[x1, x2, x2, x1, x1])
0220     y = array('d',[y1, y1, y2, y2, y1])
0221     v_x = ROOT.TVectorD(len(x),x)
0222     v_y = ROOT.TVectorD(len(y),y)
0223
0224     gr = ROOT.TGraph(v_x,v_y)
0225     gr.SetLineColor(ROOT.kBlack)
0226
0227     return gr
0228
0229
0230 def plotbarycenter(bc,coord,plotConfigJson, substructure,runsPerYear,pixelLocalRecos,accumulatedLumiPerRun, withPixelQuality,showLumi) :
0231     runs = list(accumulatedLumiPerRun.keys())
0232     runs.sort()
0233     years = list(runsPerYear.keys())
0234     years.sort()
0235     labels = list(bc.keys())
0236
0237     can = ROOT.TCanvas("barycentre_"+substructure+"_"+coord, "", 2000, 900)
0238     can.cd()
0239
0240     range_ = 0
0241     width_ = 0
0242     upper = 0
0243     lower = 0
0244     xmax = 0
0245
0246     gr = {}
0247     firstGraph = True
0248     for label in labels :
0249         gr[label] = ROOT.TGraph()
0250         gr[label] = bc[label][coord+"_"+substructure]
0251
0252         gr[label].SetMarkerStyle(8)
0253         gr[label].SetMarkerSize(0)
0254         gr[label].SetMarkerStyle(8)
0255         gr[label].SetMarkerSize(0)
0256         gr[label].SetLineColor(plotConfigJson["colorScheme"][label])
0257
0258         width_ = gr[label].GetXaxis().GetXmax() - gr[label].GetXaxis().GetXmin()
0259         xmax   = gr[label].GetXaxis().GetXmax()
0260
0261         if firstGraph :
0262            upper  = bc[label][coord+"max_"+substructure]
0263            lower  = bc[label][coord+"min_"+substructure]
0264            firstGraph = False
0265         else :
0266            upper = max(upper, bc[label][coord+"max_"+substructure])
0267            lower = min(lower, bc[label][coord+"min_"+substructure])
0268
0269     scale = 1.1
0270     if(upper>0) :
0271       upper = upper * scale
0272     else :
0273       upper = upper / scale
0274     if(lower>0) :
0275       lower = lower / scale
0276     else :
0277       lower = lower * scale
0278     range_ = upper - lower
0279
0280     firstGraph = True
0281     for label in labels :
0282         if(firstGraph) :
0283             gr[label].GetYaxis().SetRangeUser(lower, upper)
0284             gr[label].GetYaxis().SetTitle(plotConfigJson["substructures"][substructure]+" barycentre ("+coord+") [#mum]")
0285             gr[label].GetXaxis().SetTitle("Run Number")
0286             gr[label].GetYaxis().CenterTitle(True)
0287             gr[label].GetXaxis().CenterTitle(True)
0288             gr[label].GetYaxis().SetTitleOffset(0.80)
0289             gr[label].GetYaxis().SetTitleSize(0.055)
0290             gr[label].GetXaxis().SetTitleOffset(0.80)
0291             gr[label].GetXaxis().SetTitleSize(0.055)
0292             gr[label].GetXaxis().SetMaxDigits(6)
0293             if(showLumi) :
0294                 gr[label].GetXaxis().SetTitle("Delivered luminosity [1/fb]")
0295
0296             gr[label].Draw("AP")
0297             firstGraph = False
0298         else :
0299             gr[label].Draw("P")
0300
0301     # dummy TGraph for pixel local reco changes and first-of-year Run
0302     gr_dummyFirstRunOfTheYear = blackBox(-999, 10000, -999, -10000)
0303     gr_dummyFirstRunOfTheYear.SetLineColor(ROOT.kBlack)
0304     gr_dummyFirstRunOfTheYear.SetLineStyle(1)
0305     gr_dummyFirstRunOfTheYear.Draw("L")
0306     gr_dummyPixelReco = blackBox(-999, 10000, -999, -10000)
0307     gr_dummyPixelReco.SetLineColor(ROOT.kGray+1)
0308     gr_dummyPixelReco.SetLineStyle(3)
0309     gr_dummyPixelReco.Draw("L")
0310     gr_dummyFirstRunOfTheYear.SetTitle("First run of the year")
0311     gr_dummyPixelReco.SetTitle("Pixel calibration update")
0312
0313     for label in labels :
0314         gr[label].SetTitle(plotConfigJson["baryCentreLabels"][label])
0315     legend = can.BuildLegend()#0.65, 0.65, 0.85, 0.85)
0316     legend.SetShadowColor(0)
0317     legend.SetFillColor(0)
0318     legend.SetLineColor(1)
0319
0320     for label in labels :
0321         gr[label].SetTitle("")
0322
0323     # Add legends
0324     # and vertical lines
0325     years_label = ""
0326     for year in years :
0327         years_label += str(year)
0328         years_label += "+"
0329     years_label = years_label.rstrip("+")
0330
0331     # CMS logo
0332     CMSworkInProgress = ROOT.TPaveText( xmax-0.3*width_, upper+range_*0.005,
0333                                         xmax, upper+range_*0.055, "nb")
0334     CMSworkInProgress.AddText("CMS #bf{#it{Preliminary} ("+years_label+" pp collisions)}")
0335     CMSworkInProgress.SetTextAlign(32) #right/bottom aligned
0336     CMSworkInProgress.SetTextSize(0.04)
0337     CMSworkInProgress.SetFillColor(10)
0338     CMSworkInProgress.Draw()
0339
0340     # vertical lines
0341     #pixel local reco
0342     line_pixels = {}
0343     for since in pixelLocalRecos :
0344         if showLumi :
0345            run_index = findRunIndex(since,runs)
0346            integrated_lumi = accumulatedLumiPerRun[runs[run_index]]
0347            line_pixels[since] = ROOT.TLine(integrated_lumi, lower, integrated_lumi, upper)
0348
0349         else :
0350            line_pixels[since] = ROOT.TLine(since, lower, since, upper)
0351
0352         line_pixels[since].SetLineColor(ROOT.kGray+1)
0353         line_pixels[since].SetLineStyle(3)
0354         line_pixels[since].Draw()
0355
0356     # years
0357     line_years = {}
0358     box_years = {}
0359     text_years = {}
0360     if(len(years)>1 or (not showLumi) ) : # indicate begining of the year if more than one year to show or use run number
0361       for year in years :
0362           if showLumi :
0363              #first run of the year
0364              run_index = findRunIndex(runsPerYear[year][0],runs)
0365              integrated_lumi = accumulatedLumiPerRun[runs[run_index]]
0366              line_years[year] = ROOT.TLine(integrated_lumi, lower, integrated_lumi, upper)
0367              text_years[year] = ROOT.TPaveText( integrated_lumi+0.01*width_, upper-range_*0.05,
0368                                               integrated_lumi+0.05*width_,  upper-range_*0.015, "nb")
0369              box_years[year] = blackBox(integrated_lumi+0.005*width_, upper-range_*0.01, integrated_lumi+0.055*width_, upper-range_*0.055)
0370           else :
0371              line_years[year] = ROOT.TLine(runsPerYear[year][0], lower, runsPerYear[year][0], upper)
0372              text_years[year] = ROOT.TPaveText( runsPerYear[year][0]+0.01*width_, upper-range_*0.05,
0373                                               runsPerYear[year][0]+0.05*width_,  upper-range_*0.015, "nb")
0374              box_years[year] = blackBox(runsPerYear[year][0]+0.01*width_, upper-range_*0.015, runsPerYear[year][0]+0.05*width_, upper-range_*0.05)
0375
0376
0377           box_years[year].Draw("L")
0378           line_years[year].Draw()
0379
0380           # Add TextBox at the beginning of each year
0381           text_years[year].AddText(str(year))
0382           text_years[year].SetTextAlign(22)
0383           text_years[year].SetTextSize(0.025)
0384           text_years[year].SetFillColor(10)
0385           text_years[year].Draw()
0386
0387     #legend.Draw()
0388     can.Update()
0389
0390     if(showLumi) :
0391       can.SaveAs("baryCentre"+withPixelQuality+"_"+coord+"_"+substructure+"_"+years_label+"_IntegratedLumi.pdf")
0392       can.SaveAs("baryCentre"+withPixelQuality+"_"+coord+"_"+substructure+"_"+years_label+"_IntegratedLumi.png")
0393     else :
0394       can.SaveAs("baryCentre"+withPixelQuality+"_"+coord+"_"+substructure+"_"+years_label+"_RunNumber.pdf")
0395       can.SaveAs("baryCentre"+withPixelQuality+"_"+coord+"_"+substructure+"_"+years_label+"_RunNumber.png")
0396
0397     #####################################################################################################################
0398
0399
0400 # main call
0401 def Run():
0402
0403     #ROOT.gSystem.Load("libFWCoreFWLite.so")
0404     parser=parseOptions()
0405     (options,args) = parser.parse_args()
0406     sys.argv = grootargs
0407
0408     inputFileName = options.inputFileName
0409     if os.path.isfile(inputFileName) == False :
0410        print ("File "+inputFileName+" not exist!")
0411        return -1
0412
0413     plotConfigFile = open(options.plotConfigFile)
0414     plotConfigJson = json.load(plotConfigFile)
0415     plotConfigFile.close()
0416
0417     usePixelQuality = options.usePixelQuality
0418     withPixelQuality = ""
0419     if(usePixelQuality) :
0420        withPixelQuality = "WithPixelQuality"
0421     showLumi = options.showLumi
0422     # order years from old to new
0423     years = options.years
0424     years.sort()
0425
0426     # runs per year
0427     runsPerYear = {}
0428     for year in years :
0429         runsPerYear[year] = []
0430     # integrated lumi vs run
0431     accumulatedLumiPerRun = {}
0432
0433     run_index = 0
0434     lastRun = 1
0435
0436     # get lumi per IOV
0437     for year in years :
0438         inputLumiFile = fnc.digest_path("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()