tracking-parametrisation-tuner/scripts/BakPrCheckerEfficiency.py

# flake8: noqa


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)


def getEfficiencyHistoNames():
    return ["p", "pt", "phi", "eta", "nPV"]


def getTrackers(trackers):
    return trackers


def getOriginFolders():
    basedict = {
        "Velo": {},
        "Upstream": {},
        "Forward": {},
        "Match": {},
        "BestLong": {},
    }

    basedict["Velo"]["folder"] = "VeloTrackChecker/"
    basedict["Upstream"]["folder"] = "UpstreamTrackChecker/"
    basedict["Forward"]["folder"] = "ForwardTrackChecker/"
    basedict["Match"]["folder"] = "MatchTrackChecker/"
    basedict["BestLong"]["folder"] = "BestLongTrackChecker/"

    return basedict


def getTrackNames():
    basedict = {
        "Velo": {},
        "Upstream": {},
        "Forward": {},
        "Match": {},
        "BestLong": {},
    }

    basedict["Velo"] = "Velo"
    basedict["Upstream"] = "VeloUT"
    basedict["Forward"] = "Forward"
    basedict["Match"] = "Match"
    basedict["BestLong"] = "BestLong"
    return basedict


def get_colors():
    return [kBlack, kAzure, kGreen + 3, kMagenta + 2, kOrange, kCyan + 2]


def get_markers():
    return [20, 24, 21, 22, 23, 25]


def get_fillstyles():
    return [1003, 1002, 3002, 3325, 3144, 3244, 3444]


def getGhostHistoNames():
    basedict = {
        "Velo": {},
        "Upstream": {},
        "Forward": {},
        "Match": {},
        "BestLong": {},
    }

    basedict["Velo"] = ["eta", "nPV"]
    basedict["Upstream"] = ["eta", "p", "pt", "nPV"]
    basedict["Forward"] = ["eta", "p", "pt", "nPV"]
    basedict["Match"] = basedict["Forward"]
    basedict["BestLong"] = basedict["Forward"]

    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_Bs2PhiPhi_MD_default.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"],
        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",
        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 = tf[lab].Get(numeratorName)
        denominatorName = histoName + "_reconstructible"
        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 = 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], kBlue - 7, 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]
                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.045)
                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.01, 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]
            else:
                pos = [0.4, 0.37, 0.88, 0.68]
            legend = place_legend(canvas, *pos, header="LHCb Simulation", option="LPE")
            legend.SetTextFont(132)
            legend.SetTextSize(0.045)
            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))