tracking-parametrisation-tuner/scripts/BakPrCheckerTrackResolution.py

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