# flake8: noqa import argparse from ROOT import TMultiGraph, TLatex, TCanvas, TFile, TGaxis from ROOT import kGreen, kBlue, kBlack, kAzure, kOrange, kMagenta, kCyan from ROOT import gROOT, gStyle, gPad from ROOT import TEfficiency from array import array gROOT.SetBatch(True) def getEfficiencyHistoNames(): return ["p", "pt", "phi", "eta", "nPV"] def getTrackers(trackers): return trackers def getOriginFolders(): basedict = { "Velo": {}, "Upstream": {}, "Forward": {}, "Match": {}, "BestLong": {}, } basedict["Velo"]["folder"] = "VeloTrackChecker/" basedict["Upstream"]["folder"] = "UpstreamTrackChecker/" basedict["Forward"]["folder"] = "ForwardTrackChecker/" basedict["Match"]["folder"] = "MatchTrackChecker/" basedict["BestLong"]["folder"] = "BestLongTrackChecker/" return basedict def getTrackNames(): basedict = { "Velo": {}, "Upstream": {}, "Forward": {}, "Match": {}, "BestLong": {}, } basedict["Velo"] = "Velo" basedict["Upstream"] = "VeloUT" basedict["Forward"] = "Forward" basedict["Match"] = "Match" basedict["BestLong"] = "BestLong" return basedict def get_colors(): return [kBlack, kAzure, kGreen + 3, kMagenta + 2, kOrange, kCyan + 2] def get_markers(): return [20, 24, 21, 22, 23, 25] def get_fillstyles(): return [1003, 1002, 3002, 3325, 3144, 3244, 3444] def getGhostHistoNames(): basedict = { "Velo": {}, "Upstream": {}, "Forward": {}, "Match": {}, "BestLong": {}, } basedict["Velo"] = ["eta", "nPV"] basedict["Upstream"] = ["eta", "p", "pt", "nPV"] basedict["Forward"] = ["eta", "p", "pt", "nPV"] basedict["Match"] = basedict["Forward"] basedict["BestLong"] = basedict["Forward"] return basedict def argument_parser(): parser = argparse.ArgumentParser(description="location of the tuple file") parser.add_argument( "--filename", type=str, default=["data/resolutions_and_effs_Bs2PhiPhi_MD_default.root"], nargs="+", help="input files, including path", ) parser.add_argument( "--outfile", type=str, default="data/efficiency_plots.root", help="output file", ) parser.add_argument( "--trackers", type=str, nargs="+", default=["Forward", "Match", "BestLong"], help="Trackers to plot.", ) parser.add_argument( "--label", nargs="+", default=["EffChecker"], help="label for files", ) parser.add_argument( "--savepdf", action="store_true", help="save plots in pdf format", ) parser.add_argument( "--plot-electrons", action="store_true", help="plot electrons") parser.add_argument( "--plot-electrons-only", action="store_true", help="plot only electrons", ) return parser def get_files(tf, filename, label): for i, f in enumerate(filename): tf[label[i]] = TFile(f, "read") return tf def get_nicer_var_string(var: str): nice_vars = dict(pt="p_{T}", eta="#eta", phi="#phi") try: return nice_vars[var] except KeyError: return var def get_eff(eff, hist, tf, histoName, label, var): eff = {} hist = {} var = get_nicer_var_string(var) for i, lab in enumerate(label): numeratorName = histoName + "_reconstructed" numerator = tf[lab].Get(numeratorName) denominatorName = histoName + "_reconstructible" denominator = tf[lab].Get(denominatorName) if numerator.GetEntries() == 0 or denominator.GetEntries() == 0: continue teff = TEfficiency(numerator, denominator) teff.SetStatisticOption(7) eff[lab] = teff.CreateGraph() eff[lab].SetName(lab) eff[lab].SetTitle(lab + " not e^{-}") if histoName.find("strange") != -1: eff[lab].SetTitle(lab + " from stranges") if histoName.find("electron") != -1: eff[lab].SetTitle(lab + " e^{-}") hist[lab] = denominator.Clone() hist[lab].SetName("h_numerator_notElectrons") hist[lab].SetTitle(var + " distribution, not e^{-}") if histoName.find("strange") != -1: hist[lab].SetTitle(var + " distribution, stranges") if histoName.find("electron") != -1: hist[lab].SetTitle(var + " distribution, e^{-}") return eff, hist def get_ghost(eff, hist, tf, histoName, label): ghost = {} for i, lab in enumerate(label): numeratorName = histoName + "_Ghosts" denominatorName = histoName + "_Total" numerator = tf[lab].Get(numeratorName) denominator = tf[lab].Get(denominatorName) print("Numerator = " + numeratorName) print("Denominator = " + denominatorName) teff = TEfficiency(numerator, denominator) teff.SetStatisticOption(7) ghost[lab] = teff.CreateGraph() print(lab) ghost[lab].SetName(lab) return ghost def PrCheckerEfficiency( filename, outfile, label, trackers, savepdf, plot_electrons, plot_electrons_only, ): from utils.LHCbStyle import setLHCbStyle, set_style from utils.ConfigHistos import ( efficiencyHistoDict, ghostHistoDict, categoriesDict, getCuts, ) from utils.Legend import place_legend setLHCbStyle() markers = get_markers() colors = get_colors() styles = get_fillstyles() tf = {} tf = get_files(tf, filename, label) outputfile = TFile(outfile, "recreate") latex = TLatex() latex.SetNDC() latex.SetTextSize(0.05) efficiencyHistoDict = efficiencyHistoDict() efficiencyHistos = getEfficiencyHistoNames() ghostHistos = getGhostHistoNames() ghostHistoDict = ghostHistoDict() categories = categoriesDict() cuts = getCuts() trackers = getTrackers(trackers) folders = getOriginFolders() # names = getTrackNames() for tracker in trackers: outputfile.cd() trackerDir = outputfile.mkdir(tracker) trackerDir.cd() for cut in cuts[tracker]: cutDir = trackerDir.mkdir(cut) cutDir.cd() folder = folders[tracker]["folder"] print(folder) histoBaseName = "Track/" + folder + tracker + "/" + cut + "_" # calculate efficiency for histo in efficiencyHistos: canvastitle = ( "efficiency_" + histo + ", " + categories[tracker][cut]["title"] ) # get efficiency for not electrons category histoName = histoBaseName + "" + efficiencyHistoDict[histo]["variable"] print("not electrons: " + histoName) eff = {} hist_den = {} eff, hist_den = get_eff(eff, hist_den, tf, histoName, label, histo) if categories[tracker][cut]["plotElectrons"] and plot_electrons: histoNameElec = ( "Track/" + folder + tracker + "/" + categories[tracker][cut]["Electrons"] ) histoName_e = ( histoNameElec + "_" + efficiencyHistoDict[histo]["variable"] ) print("electrons: " + histoName_e) eff_elec = {} hist_elec = {} eff_elec, hist_elec = get_eff( eff_elec, hist_elec, tf, histoName_e, label, histo, ) name = "efficiency_" + histo canvas = TCanvas(name, canvastitle) canvas.SetRightMargin(0.1) mg = TMultiGraph() for i, lab in enumerate(label): if not plot_electrons_only: mg.Add(eff[lab]) set_style(eff[lab], colors[i], markers[i], styles[i]) if categories[tracker][cut]["plotElectrons"] and plot_electrons: mg.Add(eff_elec[lab]) set_style(eff_elec[lab], kBlue - 7, markers[i + 1], styles[i]) mg.Draw("AP") mg.GetYaxis().SetRangeUser(0, 1.05) xtitle = efficiencyHistoDict[histo]["xTitle"] unit_l = xtitle.split("[") if "]" in unit_l[-1]: unit = unit_l[-1].replace("]", "") else: unit = "" print(unit) mg.GetXaxis().SetTitle(xtitle) mg.GetXaxis().SetTitleSize(0.06) mg.GetYaxis().SetTitle( "Efficiency of Long Tracks", ) # (" + str(round(hist_den[label[0]].GetBinWidth(1), 2)) + f"{unit})"+"^{-1}") mg.GetYaxis().SetTitleSize(0.06) mg.GetYaxis().SetTitleOffset(1.1) mg.GetXaxis().SetRangeUser(*efficiencyHistoDict[histo]["range"]) mg.GetXaxis().SetNdivisions(10, 5, 0) mygray = 18 myblue = kBlue - 9 for i, lab in enumerate(label): rightmax = 1.05 * hist_den[lab].GetMaximum() scale = gPad.GetUymax() / rightmax hist_den[lab].Scale(scale) if categories[tracker][cut]["plotElectrons"] and plot_electrons: rightmax = 1.05 * hist_elec[lab].GetMaximum() scale = gPad.GetUymax() / rightmax hist_elec[lab].Scale(scale) if i == 0: if not plot_electrons_only: set_style(hist_den[lab], mygray, markers[i], styles[i]) gStyle.SetPalette(2, array("i", [mygray - 1, myblue + 1])) hist_den[lab].Draw("HIST PLC SAME") if categories[tracker][cut]["plotElectrons"] and plot_electrons: set_style(hist_elec[lab], myblue, markers[i], styles[i]) hist_elec[lab].SetFillColorAlpha(myblue, 0.35) hist_elec[lab].Draw("HIST PLC SAME") else: print( "No distribution plotted for other labels.", "Can be added by uncommenting the code below this print statement.", ) # set_style(hist_den[lab], mygray, markers[i], styles[i]) # gStyle.SetPalette(2, array("i", [mygray - 1, myblue + 1])) # hist_den[lab].Draw("HIST PLC SAME") if histo == "p": pos = [0.53, 0.4, 1.01, 0.71] elif histo == "pt": pos = [0.5, 0.4, 0.98, 0.71] else: pos = [0.4, 0.37, 0.88, 0.68] legend = place_legend( canvas, *pos, header="LHCb Simulation", option="LPE" ) for le in legend.GetListOfPrimitives(): if "distribution" in le.GetLabel(): le.SetOption("LF") legend.SetTextFont(132) legend.SetTextSize(0.045) legend.Draw() for lab in label: if not plot_electrons_only: eff[lab].Draw("P SAME") if categories[tracker][cut]["plotElectrons"] and plot_electrons: eff_elec[lab].Draw("P SAME") cutName = categories[tracker][cut]["title"] latex.DrawLatex(legend.GetX1() + 0.01, legend.GetY1() - 0.05, cutName) low = 0 high = 1.05 gPad.Update() axis = TGaxis( gPad.GetUxmax(), gPad.GetUymin(), gPad.GetUxmax(), gPad.GetUymax(), low, high, 510, "+U", ) axis.SetTitleFont(132) axis.SetTitleSize(0.06) axis.SetTitleOffset(0.55) axis.SetTitle( "# Tracks " + get_nicer_var_string(histo) + " distribution [a.u.]", ) axis.SetLabelSize(0) axis.Draw() canvas.RedrawAxis() if savepdf: filestypes = ["pdf"] # , "png", "eps", "C", "ps", "tex"] for ftype in filestypes: if not plot_electrons_only: canvasName = tracker + "_" + cut + "_" + histo + "." + ftype else: canvasName = ( tracker + "Electrons_" + cut + "_" + histo + "." + ftype ) canvas.SaveAs("checks/" + canvasName) # canvas.SetRightMargin(0.05) canvas.Write() # calculate ghost rate histoBaseName = "Track/" + folder + tracker + "/" for histo in ghostHistos[tracker]: trackerDir.cd() title = "ghost_rate_vs_" + histo gPad.SetTicks() histoName = histoBaseName + ghostHistoDict[histo]["variable"] ghost = {} hist_den = {} ghost = get_ghost(ghost, hist_den, tf, histoName, label) canvas = TCanvas(title, title) mg = TMultiGraph() for i, lab in enumerate(label): mg.Add(ghost[lab]) set_style(ghost[lab], colors[i], markers[i], styles[i]) xtitle = ghostHistoDict[histo]["xTitle"] mg.GetXaxis().SetTitle(xtitle) mg.GetYaxis().SetTitle("Fraction of fake tracks") mg.Draw("ap") mg.GetXaxis().SetTitleSize(0.06) mg.GetYaxis().SetTitleSize(0.06) mg.GetYaxis().SetTitleOffset(1.1) mg.GetXaxis().SetRangeUser(*efficiencyHistoDict[histo]["range"]) mg.GetXaxis().SetNdivisions(10, 5, 0) # for lab in label: # ghost[lab].Draw("P SAME") if histo == "p": pos = [0.53, 0.4, 1.01, 0.71] elif histo == "pt": pos = [0.5, 0.4, 0.98, 0.71] elif histo == "eta": pos = [0.35, 0.6, 0.85, 0.9] else: pos = [0.4, 0.37, 0.88, 0.68] legend = place_legend(canvas, *pos, header="LHCb Simulation", option="LPE") legend.SetTextFont(132) legend.SetTextSize(0.045) legend.Draw() # if histo != "nPV": # latex.DrawLatex(0.7, 0.85, "LHCb simulation") # else: # latex.DrawLatex(0.2, 0.85, "LHCb simulation") # mg.GetYaxis().SetRangeUser(0, 0.4) if histo == "eta": mg.GetYaxis().SetRangeUser(0, 0.4) # track_name = names[tracker] + " tracks" # latex.DrawLatex(0.7, 0.75, track_name) # canvas.PlaceLegend() if savepdf: filestypes = ["pdf"] # , "png", "eps", "C", "ps", "tex"] for ftype in filestypes: canvas.SaveAs( "checks/" + tracker + "ghost_rate_" + histo + "." + ftype, ) canvas.Write() outputfile.Write() outputfile.Close() if __name__ == "__main__": parser = argument_parser() args = parser.parse_args() PrCheckerEfficiency(**vars(args))