tracking-parametrisation-tuner/scripts/BakPrCheckerTrackResolution.py

310 lines
9.6 KiB
Python
Raw Normal View History

2023-12-19 13:00:59 +01:00
import argparse
from ROOT import TLatex, TCanvas, TFile, TGaxis
from ROOT import kOrange, kGray, kMagenta, kCyan, kGreen, kBlue, kBlack, gPad, TF1
from ROOT import gROOT
from ROOT import TObjArray
gROOT.SetBatch(True)
def get_colors():
return [kBlack, kBlue, kGreen + 3, kMagenta + 2, kOrange, kCyan + 2]
def get_markers():
return [20, 24, 21, 22, 23, 25]
def get_fillstyles():
return [3004, 3003, 3325, 3144, 3244, 3444]
def get_files(tf, filename, label):
for i, f in enumerate(filename):
tf[label[i]] = TFile(f, "read")
return tf
def argument_parser():
parser = argparse.ArgumentParser(description="location of the histogram file")
parser.add_argument("--filename", nargs="+", default=[], help="name of input files")
parser.add_argument(
"--label",
nargs="+",
default=["TrackRes"],
help="name of input tuple files",
)
parser.add_argument(
"--trackers",
type=str,
nargs="+",
default=["Forward", "BestLong", "BestForward"],
help="Trackers to plot.",
)
parser.add_argument(
"--outfile",
type=str,
default="checks/TrackResolution_plots.root",
help="name of output files",
)
parser.add_argument(
"--savepdf",
default=False,
action="store_true",
help="save plots in pdf format",
)
return parser
def PrCheckerTrackResolution(filename, label, trackers, outfile, savepdf):
from utils.LHCbStyle import setLHCbStyle, set_style
from utils.Legend import place_legend
setLHCbStyle()
latex = TLatex()
latex.SetNDC()
latex.SetTextSize(0.05)
markers = get_markers()
colors = get_colors()
styles = get_fillstyles()
tf = {}
tf = get_files(tf, filename, label)
outputfile = TFile(outfile, "recreate")
outputfile.cd()
for tracker in trackers:
hres_p = {}
hres_eta = {}
canvas1 = TCanvas("res_p", "res v.s. p")
canvas1.cd()
arrp = TObjArray()
gaus = TF1("gaus", "gaus", -1, 1)
for idx, lab in enumerate(label):
hdp_p = tf[lab].Get(
f"Track/TrackResChecker{tracker}/ALL/vertex/dpoverp_vs_p",
)
hdp_p.SetName("dpoverp_p_" + lab)
hmom = hdp_p.ProjectionX()
hmom.SetTitle("p distribution Long Tracks")
hdp_p.FitSlicesY(gaus, 0, -1, 0, "Q", arrp)
hres_p[lab] = arrp[2]
hres_p[lab].GetYaxis().SetTitle("dp/p [%]")
hres_p[lab].GetXaxis().SetTitle("p [GeV]")
hres_p[lab].GetYaxis().SetRangeUser(0, 1.2)
hres_p[lab].SetTitle(lab)
set_style(hres_p[lab], colors[idx], markers[idx], 0)
if idx == 0:
hres_p[lab].Draw("E1 p1")
set_style(hmom, kGray + 1, markers[idx], styles[idx])
hmom.Scale(gPad.GetUymax() / hmom.GetMaximum())
hmom.Draw("hist same")
else:
hres_p[lab].Draw("E1 p1 same")
set_style(hmom, colors[idx] - 10, markers[idx], styles[idx])
# hmom.Scale(gPad.GetUymax() / hmom.GetMaximum())
# hmom.Draw("hist same")
for i in range(1, hres_p[lab].GetNbinsX() + 1):
hres_p[lab].SetBinContent(i, hres_p[lab].GetBinContent(i) * 100)
hres_p[lab].SetBinError(i, hres_p[lab].GetBinError(i) * 100)
print(
lab
+ ": Track resolution (dp/p) in p region: ("
+ format(hres_p[lab].GetBinLowEdge(i), ".2f")
+ ", "
+ format(
hres_p[lab].GetBinLowEdge(i) + hres_p[lab].GetBinWidth(i),
".2f",
)
+ ") [GeV/c]"
+ " --- ("
+ format(hres_p[lab].GetBinContent(i), ".2f")
+ "+-"
+ format(hres_p[lab].GetBinError(i), ".2f")
+ ")%",
)
print("-----------------------------------------------------")
if "Best" not in tracker:
legend = place_legend(
canvas1,
0.45,
0.33,
0.93,
0.63,
header="LHCb Simulation",
option="LPE",
)
else:
legend = place_legend(
canvas1,
0.45,
0.53,
0.93,
0.83,
header="LHCb Simulation",
option="LPE",
)
for le in legend.GetListOfPrimitives():
if "distribution" in le.GetLabel():
le.SetOption("LF")
legend.Draw()
for lab in label:
hres_p[lab].Draw("E1 p1 same")
canvas1.SetRightMargin(0.1)
if "Best" not in tracker:
latex.DrawLatex(
legend.GetX1() + 0.01,
legend.GetY1() - 0.05,
"without Kalman Filter",
)
low = 0
high = 1.2
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 p distribution [a.u.]")
axis.SetLabelSize(0)
axis.Draw()
canvas1.Write()
if savepdf:
filestypes = ["pdf"] # , "png", "eps", "C", "ps", "tex"]
for ftype in filestypes:
canvas1.SaveAs(f"checks/{tracker}_trackres_p." + ftype)
gaus.Delete()
arrp.Delete()
canvas2 = TCanvas("res_eta", "res v.s. eta")
canvas2.cd()
arreta = TObjArray()
gaus = TF1("gaus", "gaus", -1, 1)
for idx, lab in enumerate(label):
hdp_eta = tf[lab].Get(
f"Track/TrackResChecker{tracker}/ALL/vertex/dpoverp_vs_eta",
)
hdp_eta.SetName("dpoverp_eta_" + lab)
hdp_eta.FitSlicesY(gaus, 0, -1, 0, "Q", arreta)
heta = hdp_eta.ProjectionX()
heta.SetTitle("#eta distribution Long Tracks")
hres_eta[lab] = arreta[2]
hres_eta[lab].GetYaxis().SetTitle("dp/p [%]")
hres_eta[lab].GetXaxis().SetTitle("#eta")
hres_eta[lab].GetYaxis().SetRangeUser(0, 1.2)
hres_eta[lab].SetTitle(lab)
set_style(hres_eta[lab], colors[idx], markers[idx], 0)
if idx == 0:
hres_eta[lab].Draw("E1 p1")
set_style(heta, kGray + 1, markers[idx], styles[idx])
heta.Scale(gPad.GetUymax() / heta.GetMaximum())
heta.Draw("hist same")
else:
hres_eta[lab].Draw("E1 p1 same")
set_style(heta, colors[idx] - 10, markers[idx], styles[idx])
# heta.Scale(gPad.GetUymax() / heta.GetMaximum())
# heta.Draw("hist same")
for i in range(1, hres_eta[lab].GetNbinsX() + 1):
hres_eta[lab].SetBinContent(i, hres_eta[lab].GetBinContent(i) * 100)
hres_eta[lab].SetBinError(i, hres_eta[lab].GetBinError(i) * 100)
print(
lab
+ ": Track resolution (dp/p) in eta region: ("
+ format(hres_eta[lab].GetBinLowEdge(i), ".2f")
+ ", "
+ format(
hres_eta[lab].GetBinLowEdge(i) + hres_eta[lab].GetBinWidth(i),
".2f",
)
+ ")"
+ " --- ("
+ format(hres_eta[lab].GetBinContent(i), ".2f")
+ "+-"
+ format(hres_eta[lab].GetBinError(i), ".2f")
+ ")%",
)
print("-----------------------------------------------------")
legend = place_legend(
canvas2,
0.41,
0.27,
0.89,
0.57,
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:
hres_eta[lab].Draw("E1 p1 same")
canvas2.SetRightMargin(0.1)
if "Best" not in tracker:
latex.DrawLatex(
legend.GetX1() + 0.01,
legend.GetY1() - 0.05,
"without Kalman Filter",
)
low = 0
high = 1.2
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 #eta distribution [a.u.]")
axis.SetLabelSize(0)
axis.Draw()
canvas2.Write()
if savepdf:
filestypes = ["pdf"] # , "png", "eps", "C", "ps", "tex"]
for ftype in filestypes:
canvas2.SaveAs(f"checks/{tracker}_trackres_eta." + ftype)
gaus.Delete()
arreta.Delete()
outputfile.Write()
outputfile.Close()
if __name__ == "__main__":
parser = argument_parser()
args = parser.parse_args()
PrCheckerTrackResolution(**vars(args))