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))