You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

500 lines
17 KiB

# flake8: noqa
#
# Takes data from Recent_get_resolution_and_eff_data.py and calculates efficiencies
#
# python scripts/ResidualPrCheckerEfficiency.py
# --filename data/resolutions_and_effs_B.root
# --trackers Match --outfile data/compare_effs.root
#
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)
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_B.root:Track/...
def getOriginFolders():
basedict = {
"Velo": {},
"Upstream": {},
"Forward": {},
"Match": {},
"BestLong": {},
"Seed": {},
}
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["Forward"]["folder"] = "ForwardTrackChecker_7a0dbfa7/"
# basedict["Match"]["folder"] = "MatchTrackChecker_29e3152a/"
# basedict["ResidualMatch"]["folder"] = "ResidualMatchTrackChecker_955c7a21/"
# basedict["BestLong"]["folder"] = "BestLongTrackChecker_c163325d/"
# basedict["Seed"]["folder"] = "SeedTrackChecker_1b1d5575/"
return basedict
def getTrackNames():
basedict = {
"Velo": {},
"Upstream": {},
"Forward": {},
"Match": {},
"BestLong": {},
"Seed": {},
}
basedict["Velo"] = "Velo"
basedict["Upstream"] = "VeloUT"
basedict["Forward"] = "Forward"
basedict["Match"] = "Match"
basedict["BestLong"] = "BestLong"
basedict["Seed"] = "Seed"
return basedict
def get_colors():
return [kBlack, kGreen + 3, kAzure, kMagenta + 2, kOrange, kCyan + 2]
def get_markers():
return [20, 24, 21, 22, 23, 25]
def get_fillstyles():
return [1003, 3001, 3002, 3325, 3144, 3244, 3444]
def getGhostHistoNames():
basedict = {
"Velo": {},
"Upstream": {},
"Forward": {},
"Match": {},
"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["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_Bd2KstEE_MDmaster.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", "Seed"], # ---
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)
# numerator = tf[lab].Get(numeratorName)
denominatorName = histoName + "_reconstructible"
denominator = findRootObjByName(tf[lab], denominatorName)
# 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 = findRootObjByName(tf[lab], numeratorName)
denominator = findRootObjByName(tf[lab], denominatorName)
# 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], colors[i + 2], 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]
elif histo == "phi":
pos = [0.3, 0.3, 0.9, 0.6]
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.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
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.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))