|
|
# flake8: noqa
"""
Takes data from Recent_get_resolution_and_eff_data.py and calculates efficiencies
python scripts/MyPrCheckerEfficiency.py --filename data/resolutions_and_effs_B_normal_weights.root data/resolutions_and_effs_B_electron_weights.root--outfile data_results/PrCheckerNormalElectron.root --label normal electron --trackers Match BestLong
python scripts/MyPrCheckerEfficiency.py --filename data/resolutions_and_effs_D_electron_weights.root --outfile data_results/PrCheckerDElectron.root
python scripts/PrCheckerEfficiency.py --filename data/resolutions_and_effs_B_thesis.root--outfile data_results/PrCheckerTestThesis.root --label particle --plot-velo --plot-velo-only --trackers Match Seed --savepdf
"""
import argparsefrom ROOT import TMultiGraph, TLatex, TCanvas, TFile, TGaxisfrom ROOT import ( kGreen, kBlue, kBlack, kAzure, kOrange, kMagenta, kCyan, kGray, kViolet, kTeal,)from ROOT import gROOT, gStyle, gPadfrom ROOT import TEfficiencyfrom array import array
gROOT.SetBatch(True)
from utils.components import unique_name_ext_re, findRootObjByName
def getEfficiencyHistoNames(): return ["p", "pt", "phi", "eta", "nPV"]
def getTrackers(trackers): return trackers
# data/resolutions_and_effs_Bd2KstEE_MDmaster.root:Track/...def getOriginFolders(): basedict = { "Velo": {}, "Upstream": {}, "Forward": {}, "Match": {}, "BestLong": {}, "Seed": {}, "MergedMatch": {}, "DefaultMatch": {}, } # evtl anpassen wenn die folders anders heissen basedict["Velo"]["folder"] = "VeloTrackChecker/" basedict["Upstream"]["folder"] = "UpstreamTrackChecker/" basedict["Forward"]["folder"] = "ForwardTrackChecker" + unique_name_ext_re() + "/" basedict["Match"]["folder"] = "MatchTrackChecker" + unique_name_ext_re() + "/" basedict["BestLong"]["folder"] = "BestLongTrackChecker" + unique_name_ext_re() + "/" basedict["Seed"]["folder"] = "SeedTrackChecker" + unique_name_ext_re() + "/" basedict["MergedMatch"]["folder"] = ( "MergedMatchTrackChecker" + unique_name_ext_re() + "/" ) basedict["DefaultMatch"]["folder"] = ( "DefaultMatchTrackChecker" + unique_name_ext_re() + "/" ) # basedict["Forward"]["folder"] = "ForwardTrackChecker_7a0dbfa7/" # basedict["Match"]["folder"] = "MatchTrackChecker_29e3152a/" # basedict["BestLong"]["folder"] = "BestLongTrackChecker_4ddacce1/" # basedict["Seed"]["folder"] = "SeedTrackChecker_1b1d5575/"
return basedict
def getTrackNames(): basedict = { "Velo": {}, "Upstream": {}, "Forward": {}, "Match": {}, "BestLong": {}, "Seed": {}, "MergedMatch": {}, "DefaultMatch": {}, }
basedict["Velo"] = "Velo" basedict["Upstream"] = "VeloUT" basedict["Forward"] = "Forward" basedict["Match"] = "Match" basedict["BestLong"] = "BestLong" basedict["Seed"] = "Seed" basedict["MergedMatch"] = "MergedMatch" basedict["DefaultMatch"] = "DefaultMatch" return basedict
def get_colors(): # [kBlack, kGreen + 3, kAzure, kMagenta + 2, kOrange, kCyan + 2] return [kBlack, kAzure, kGreen + 3, kMagenta + 2, kOrange, kCyan + 2]
def get_elec_colors(): # [kBlack, kGreen + 3, kAzure, kMagenta + 2, kOrange, kCyan + 2] return [kBlue - 7, kGray + 2, kGreen + 1, kViolet, kOrange - 3, kTeal - 1]
def get_markers(): # [20, 24, 21, 22, 23, 25] return [20, 21, 22, 23, 24, 25, 26, 32]
def get_elec_markers(): # [20, 24, 21, 22, 23, 25] return [24, 25, 26, 32, 22, 23, 26, 32]
def get_fillstyles(): return [1003, 3001, 3002, 3325, 3144, 3244, 3444]
def getGhostHistoNames(): basedict = { "Velo": {}, "Upstream": {}, "Forward": {}, "Match": {}, "MergedMatch": {}, "DefaultMatch": {}, "BestLong": {}, "Seed": {}, }
basedict["Velo"] = ["eta", "nPV"] basedict["Upstream"] = ["eta", "p", "pt", "nPV"] basedict["Forward"] = ["eta", "p", "pt", "nPV"] basedict["Match"] = ["eta", "p", "pt", "nPV"] basedict["MergedMatch"] = basedict["Match"] basedict["DefaultMatch"] = basedict["Match"] basedict["BestLong"] = ["eta", "p", "pt", "nPV"] basedict["Seed"] = ["eta", "p", "pt", "nPV"]
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_B.root"], nargs="+", help="input files, including path", ) parser.add_argument( "--outfile", type=str, default="data_results/efficiency_plots.root", help="output file", ) parser.add_argument( "--trackers", type=str, nargs="+", default=["Forward", "Match", "BestLong", "Seed"], # DefaultMatch 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", default=True, help="plot electrons", ) parser.add_argument( "--plot-electrons-only", action="store_true", help="plot only electrons", ) parser.add_argument( "--plot-velo", action="store_true", help="plot using momentum at EndVelo", ) parser.add_argument( "--plot-velo-only", action="store_true", help="plot using only momentum at EndVelo", ) 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 = findRootObjByName(tf[lab], numeratorName) denominatorName = histoName + "_reconstructible" denominator = findRootObjByName(tf[lab], 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("Forward") != -1: if histoName.find("electron") != -1: eff[lab].SetTitle(lab + " Forward, e^{-}") else: eff[lab].SetTitle(lab + " Forward, not e^{-}") if histoName.find("Merged") != -1: if histoName.find("electron") != -1: eff[lab].SetTitle(lab + " MergedMatch, e^{-}") else: eff[lab].SetTitle(lab + " MergedMatch, not e^{-}") if histoName.find("DefaultMatch") != -1: if histoName.find("electron") != -1: eff[lab].SetTitle(lab + " DefaultMatch, e^{-}") else: eff[lab].SetTitle(lab + " DefaultMatch, not e^{-}") if histoName.find("Match") != -1: if histoName.find("electron") != -1: eff[lab].SetTitle(lab + " Match, e^{-}") else: eff[lab].SetTitle(lab + " Match, not e^{-}") if histoName.find("Seed") != -1: if histoName.find("electron") != -1: eff[lab].SetTitle(lab + " Seed, e^{-}") else: eff[lab].SetTitle(lab + " Seed, not e^{-}") if histoName.find("BestLong") != -1: if histoName.find("electron") != -1: eff[lab].SetTitle(lab + " BestLong, e^{-}") else: eff[lab].SetTitle(lab + " BestLong, not e^{-}") if histoName.find("EndVelo") != -1: eff[lab].SetTitle(eff[lab].GetTitle() + " EndVelo") # if histoName.find("EndUT") != -1: # eff[lab].SetTitle(eff[lab].GetTitle() + " EndUT")
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^{-}") if histoName.find("EndVelo") != -1: hist[lab].SetTitle(hist[lab].GetTitle() + ", EndVelo") # if histoName.find("EndUT") != -1: # hist[lab].SetTitle(hist[lab].GetTitle() + ", EndUT")
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 = findRootObjByName(tf[lab], numeratorName) denominator = findRootObjByName(tf[lab], denominatorName) print("Numerator = " + numeratorName.replace(unique_name_ext_re(), "")) print("Denominator = " + denominatorName.replace(unique_name_ext_re(), "")) 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, plot_velo, plot_velo_only,): from utils.LHCbStyle import setLHCbStyle, set_style from utils.ConfigHistos import ( efficiencyHistoDict, ghostHistoDict, categoriesDict, getCuts, ) from utils.Legend import place_legend
# from utils.components import unique_name_ext_re, findRootObjByName
setLHCbStyle()
markers = get_markers() elec_markers = get_elec_markers() colors = get_colors() elec_colors = get_elec_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: " + folder.replace(unique_name_ext_re(), "")) 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.replace(unique_name_ext_re(), "")) 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.replace(unique_name_ext_re(), "")) eff_elec = {} hist_elec = {} eff_elec, hist_elec = get_eff( eff_elec, hist_elec, tf, histoName_e, label, histo, ) if ( categories[tracker][cut]["plotEndVelo"] and plot_velo and (histo == "p" or histo == "pt") ): histoNameEndVelo = ( "Track/" + folder + tracker + "/" + categories[tracker][cut]["EndVelo"] ) histoName_v = ( histoNameEndVelo + "_" + efficiencyHistoDict[histo]["variable"] ) print( "EndVelo: " + histoName_v.replace(unique_name_ext_re(), "") ) eff_velo = {} hist_velo = {} eff_velo, hist_velo = get_eff( eff_velo, hist_velo, tf, histoName_v, label, histo, ) name = "efficiency_" + histo canvas = TCanvas(name, canvastitle) canvas.SetRightMargin(0.1) mg = TMultiGraph() for i, lab in enumerate(label): if ( categories[tracker][cut]["plotEndVelo"] and plot_velo and (histo == "p" or histo == "pt") ): # and not plot_velo_only: mg.Add(eff_velo[lab]) set_style(eff_velo[lab], colors[i], markers[i], styles[i]) if categories[tracker][cut]["plotElectrons"] and plot_electrons: if (not plot_velo_only) or ( histo == "phi" or histo == "eta" or histo == "nPV" ): mg.Add(eff_elec[lab]) set_style( eff_elec[lab], elec_colors[i], elec_markers[i], 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 = "a.u." 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 - 10 mypurple = kMagenta - 10 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 ( categories[tracker][cut]["plotEndVelo"] and plot_velo and (histo == "p" or histo == "pt") ): rightmax = 1.05 * hist_velo[lab].GetMaximum() scale = gPad.GetUymax() / rightmax hist_velo[lab].Scale(scale)
if i == 0: if ( categories[tracker][cut]["plotEndVelo"] and plot_velo and (histo == "p" or histo == "pt") ): set_style(hist_velo[lab], mygray, markers[i], styles[i]) gStyle.SetPalette(2, array("i", [mygray - 1, myblue + 1])) hist_velo[lab].Draw("HIST PLC SAME") if categories[tracker][cut]["plotElectrons"] and plot_electrons: if not plot_velo_only or ( histo == "phi" or histo == "eta" or histo == "nPV" ): set_style( hist_elec[lab], myblue, elec_markers[i], styles[i] ) hist_elec[lab].SetFillColorAlpha(myblue, 0.5) 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.5, 0.35, 1.0, 0.7] # [0.53, 0.4, 1.01, 0.71] elif histo == "pt": pos = [0.5, 0.3, 0.99, 0.5] # [0.5, 0.4, 0.98, 0.71] elif histo == "phi": pos = [0.4, 0.3, 0.9, 0.5] elif histo == "eta": pos = [0.5, 0.25, 1.0, 0.45] else: pos = [0.35, 0.25, 0.85, 0.45] if histo == "p": 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.04) legend.Draw() for lab in label: if ( categories[tracker][cut]["plotEndVelo"] and plot_velo and (histo == "p" or histo == "pt") ): eff_velo[lab].Draw("P SAME") if categories[tracker][cut]["plotElectrons"] and plot_electrons: if not plot_velo_only or ( histo == "phi" or histo == "eta" or histo == "nPV" ): 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 print("\ncalculate 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[2 * 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.00, 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] elif histo == "phi": pos = [0.3, 0.3, 0.9, 0.6] else: pos = [0.4, 0.37, 0.80, 0.68] legend = place_legend(canvas, *pos, header="LHCb Simulation", option="LPE") legend.SetTextFont(132) legend.SetTextSize(0.04) 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))
|
