tracking-parametrisation-tuner/scripts/MyPrCheckerEfficiency.py

# flake8: noqa

"""
Takes data from Recent_get_resolution_and_eff_data.py and calculates efficiencies


python scripts/MyPrCheckerEfficiency.py --filename data/resolutions_and_effs_B_normal_weights.root data/resolutions_and_effs_B_electron_weights.root
--outfile data_results/PrCheckerNormalElectron.root --label normal electron --trackers Match BestLong

python scripts/MyPrCheckerEfficiency.py --filename data/resolutions_and_effs_D_electron_weights.root --outfile data_results/PrCheckerDElectron.root

"""

import argparse
from ROOT import TMultiGraph, TLatex, TCanvas, TFile, TGaxis
from ROOT import (
    kGreen,
    kBlue,
    kBlack,
    kAzure,
    kOrange,
    kMagenta,
    kCyan,
    kGray,
    kViolet,
    kTeal,
)
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_Bd2KstEE_MDmaster.root:Track/...
def getOriginFolders():
    basedict = {
        "Velo": {},
        "Upstream": {},
        "Forward": {},
        "Match": {},
        "BestLong": {},
        "Seed": {},
        "MergedMatch": {},
        "DefaultMatch": {},
    }
    # evtl anpassen wenn die folders anders heissen
    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["MergedMatch"]["folder"] = (
        "MergedMatchTrackChecker" + unique_name_ext_re() + "/"
    )
    basedict["DefaultMatch"]["folder"] = (
        "DefaultMatchTrackChecker" + unique_name_ext_re() + "/"
    )
    # basedict["Forward"]["folder"] = "ForwardTrackChecker_7a0dbfa7/"
    # basedict["Match"]["folder"] = "MatchTrackChecker_29e3152a/"
    # basedict["BestLong"]["folder"] = "BestLongTrackChecker_4ddacce1/"
    # basedict["Seed"]["folder"] = "SeedTrackChecker_1b1d5575/"

    return basedict


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

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


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


def get_elec_colors():
    # [kBlack, kGreen + 3, kAzure, kMagenta + 2, kOrange, kCyan + 2]
    return [kBlue - 7, kGray + 2, kGreen + 1, kViolet, kOrange - 3, kTeal - 1]


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


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


def getGhostHistoNames():
    basedict = {
        "Velo": {},
        "Upstream": {},
        "Forward": {},
        "Match": {},
        "MergedMatch": {},
        "DefaultMatch": {},
        "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["MergedMatch"] = basedict["Match"]
    basedict["DefaultMatch"] = basedict["Match"]
    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_B.root"],
        nargs="+",
        help="input files, including path",
    )
    parser.add_argument(
        "--outfile",
        type=str,
        default="data_results/efficiency_plots.root",
        help="output file",
    )
    parser.add_argument(
        "--trackers",
        type=str,
        nargs="+",
        default=["Forward", "Match", "BestLong", "Seed"],  # DefaultMatch
        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",
    )
    parser.add_argument(
        "--plot-velo",
        action="store_true",
        help="plot using momentum at EndVelo",
    )
    parser.add_argument(
        "--plot-velo-only",
        action="store_true",
        help="plot using only momentum at EndVelo",
    )
    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)
        denominatorName = histoName + "_reconstructible"
        denominator = findRootObjByName(tf[lab], 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("Forward") != -1:
            if histoName.find("electron") != -1:
                eff[lab].SetTitle(lab + " Forward, e^{-}")
            else:
                eff[lab].SetTitle(lab + " Forward")
        if histoName.find("Merged") != -1:
            if histoName.find("electron") != -1:
                eff[lab].SetTitle(lab + " MergedMatch, e^{-}")
            else:
                eff[lab].SetTitle(lab + " MergedMatch")
        if histoName.find("DefaultMatch") != -1:
            if histoName.find("electron") != -1:
                eff[lab].SetTitle(lab + " DefaultMatch, e^{-}")
            else:
                eff[lab].SetTitle(lab + " DefaultMatch")
        if histoName.find("Match") != -1:
            if histoName.find("electron") != -1:
                eff[lab].SetTitle(lab + " Match, e^{-}")
            else:
                eff[lab].SetTitle(lab + " Match")
        if histoName.find("Seed") != -1:
            if histoName.find("electron") != -1:
                eff[lab].SetTitle(lab + " Seed, e^{-}")
            else:
                eff[lab].SetTitle(lab + " Seed")
        if histoName.find("BestLong") != -1:
            if histoName.find("electron") != -1:
                eff[lab].SetTitle(lab + " BestLong, e^{-}")
            else:
                eff[lab].SetTitle(lab + " BestLong")
        if histoName.find("EndVelo") != -1:
            eff[lab].SetTitle(eff[lab].GetTitle() + " EndVelo")
        # if histoName.find("EndUT") != -1:
        #     eff[lab].SetTitle(eff[lab].GetTitle() + " EndUT")

        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^{-}")
        if histoName.find("EndVelo") != -1:
            hist[lab].SetTitle(hist[lab].GetTitle() + ", EndVelo")
        # if histoName.find("EndUT") != -1:
        #     hist[lab].SetTitle(hist[lab].GetTitle() + ", EndUT")

    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)
        print("Numerator = " + numeratorName.replace(unique_name_ext_re(), ""))
        print("Denominator = " + denominatorName.replace(unique_name_ext_re(), ""))
        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,
    plot_velo,
    plot_velo_only,
):
    from utils.LHCbStyle import setLHCbStyle, set_style
    from utils.ConfigHistos import (
        efficiencyHistoDict,
        ghostHistoDict,
        categoriesDict,
        getCuts,
    )
    from utils.Legend import place_legend

    # from utils.components import unique_name_ext_re, findRootObjByName

    setLHCbStyle()

    markers = get_markers()
    colors = get_colors()
    elec_colors = get_elec_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: " + folder.replace(unique_name_ext_re(), ""))
            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.replace(unique_name_ext_re(), ""))
                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.replace(unique_name_ext_re(), ""))
                    eff_elec = {}
                    hist_elec = {}
                    eff_elec, hist_elec = get_eff(
                        eff_elec,
                        hist_elec,
                        tf,
                        histoName_e,
                        label,
                        histo,
                    )
                    if (
                        categories[tracker][cut]["plotEndVelo"]
                        and plot_velo
                        and (histo == "p" or histo == "pt")
                    ):
                        histoNameEndVelo = (
                            "Track/"
                            + folder
                            + tracker
                            + "/"
                            + categories[tracker][cut]["EndVelo"]
                        )
                        histoName_v = (
                            histoNameEndVelo
                            + "_"
                            + efficiencyHistoDict[histo]["variable"]
                        )
                        print(
                            "EndVelo: " + histoName_v.replace(unique_name_ext_re(), "")
                        )
                        eff_velo = {}
                        hist_velo = {}
                        eff_velo, hist_velo = get_eff(
                            eff_velo,
                            hist_velo,
                            tf,
                            histoName_v,
                            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 and not plot_velo_only:
                        mg.Add(eff[lab])
                        set_style(eff[lab], colors[i], markers[i], styles[i])
                    if categories[tracker][cut]["plotElectrons"] and plot_electrons:
                        if not plot_velo_only:
                            mg.Add(eff_elec[lab])
                            set_style(
                                eff_elec[lab], elec_colors[i], markers[i], styles[i]
                            )
                        if (
                            categories[tracker][cut]["plotEndVelo"]
                            and plot_velo
                            and (histo == "p" or histo == "pt")
                        ):
                            mg.Add(eff_velo[lab])
                            set_style(
                                eff_velo[lab], kMagenta + 1, markers[i], 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 = "a.u."
                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 - 10
                mypurple = kMagenta - 10
                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 (
                            categories[tracker][cut]["plotEndVelo"]
                            and plot_velo
                            and (histo == "p" or histo == "pt")
                        ):
                            rightmax = 1.05 * hist_velo[lab].GetMaximum()
                            scale = gPad.GetUymax() / rightmax
                            hist_velo[lab].Scale(scale)

                    if i == 0:
                        if not plot_electrons_only and not plot_velo_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:
                            if not plot_velo_only:
                                set_style(hist_elec[lab], myblue, markers[i], styles[i])
                                # hist_elec[lab].SetFillColor(myblue)
                                hist_elec[lab].Draw("HIST PLC SAME")
                            if (
                                categories[tracker][cut]["plotEndVelo"]
                                and plot_velo
                                and (histo == "p" or histo == "pt")
                            ):
                                set_style(
                                    hist_velo[lab], mypurple, markers[i], styles[i]
                                )
                                hist_velo[lab].SetFillColorAlpha(mypurple, 0.5)
                                hist_velo[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.5, 0.3, 1.0, 0.5]  # [0.53, 0.4, 1.01, 0.71]
                elif histo == "pt":
                    pos = [0.5, 0.3, 0.99, 0.5]  # [0.5, 0.4, 0.98, 0.71]
                elif histo == "phi":
                    pos = [0.4, 0.3, 0.9, 0.5]
                elif histo == "eta":
                    pos = [0.5, 0.25, 1.0, 0.45]
                else:
                    pos = [0.35, 0.25, 0.85, 0.45]
                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 and not plot_velo_only:
                        eff[lab].Draw("P SAME")
                    if categories[tracker][cut]["plotElectrons"] and plot_electrons:
                        if not plot_velo_only:
                            eff_elec[lab].Draw("P SAME")
                        if (
                            categories[tracker][cut]["plotEndVelo"]
                            and plot_velo
                            and (histo == "p" or histo == "pt")
                        ):
                            eff_velo[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
        print("\ncalculate 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[2 * 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))