536 lines
19 KiB
Python
536 lines
19 KiB
Python
# 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
|
|
|
|
"""
|
|
|
|
import argparse
|
|
from ROOT import TMultiGraph, TLatex, TCanvas, TFile, TGaxis
|
|
from ROOT import (
|
|
kGreen,
|
|
kBlue,
|
|
kBlack,
|
|
kAzure,
|
|
kOrange,
|
|
kMagenta,
|
|
kCyan,
|
|
kGray,
|
|
kViolet,
|
|
kTeal,
|
|
)
|
|
from ROOT import gROOT, gStyle, gPad
|
|
from ROOT import TEfficiency
|
|
from 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 [kGray + 2, kBlue - 7, kGreen + 1, kViolet, kOrange - 3, kTeal - 1]
|
|
|
|
|
|
def get_markers():
|
|
# [20, 24, 21, 22, 23, 25]
|
|
return [20, 21, 24, 25, 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",
|
|
)
|
|
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("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 = 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,
|
|
):
|
|
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()
|
|
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,
|
|
)
|
|
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], elec_colors[i], 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 - 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.5, 0.3, 1.0, 0.5] # [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]
|
|
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 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
|
|
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))
|