Ejemplo n.º 1
0
def do_projection_plots(in_file, plot_dir, do_fit=True, skip_dirs=None):
    hist_name = "pt_jet_response"
    tfile = cu.open_root_file(in_file)
    dirs = cu.get_list_of_element_names(tfile)

    for mydir in dirs:
        if skip_dirs and mydir in skip_dirs:
            continue

        if hist_name not in cu.get_list_of_element_names(tfile.Get(mydir)):
            continue

        print("Doing", mydir)

        h2d = cu.grab_obj_from_file(in_file, "%s/%s" % (mydir, hist_name))

        ax = h2d.GetXaxis()
        bin_edges = [ax.GetBinLowEdge(i) for i in range(1, ax.GetNbins() + 2)]

        bin_centers, sigmas, sigmas_unc = [], [], []

        for pt_min, pt_max in zip(bin_edges[:-1], bin_edges[1:]):
            obj = qgg.get_projection_plot(h2d, pt_min, pt_max, cut_axis='x')
            if obj.GetEffectiveEntries() < 20:
                continue
            # obj.Rebin(rebin)
            obj.Scale(1. / obj.Integral())

            label = "%s < p_{T}^{Gen} < %s GeV" % (str(pt_min), str(pt_max))
            if do_fit:
                do_gaus_fit(obj)
                fit = obj.GetFunction("gausFit")
                label += "\n"
                label += fit_results_to_str(fit)
                # bin_centers.append(fit.GetParameter(1))
                bin_centers.append(0.5 * (pt_max + pt_min))
                sigmas.append(fit.GetParameter(2))
                sigmas_unc.append(fit.GetParError(2))

            # output_filename = os.path.join(plot_dir, "%s_%s_ptGen%sto%s.%s" % (mydir, hist_name, str(pt_min), str(pt_max), OUTPUT_FMT))

            # cont = Contribution(obj, label=label)
            # delta = pt_max - pt_min
            # # xlim = (pt_min - 10*delta, pt_max + 10*delta)
            # xlim = (obj.GetMean()-3*obj.GetRMS(), obj.GetMean()+3*obj.GetRMS())
            # ylim = (0, obj.GetMaximum()*1.1)
            # plot = Plot([cont], what='hist',
            #             xtitle="p_{T}^{Reco} [GeV]", xlim=xlim, ylim=ylim)
            # plot.plot()  # don't use histe as it wont draw the fit
            # plot.save(output_filename)

        gr = ROOT.TGraphErrors(len(bin_centers), array('d', bin_centers),
                               array('d', sigmas),
                               array('d', [0] * len(bin_centers)),
                               array('d', sigmas_unc))
        factor = 0.2
        gr_ideal = ROOT.TGraphErrors(
            len(bin_centers), array('d', bin_centers),
            array('d', [factor * pt for pt in bin_centers]),
            array('d', [0] * len(bin_centers)),
            array('d', [0] * len(bin_centers)))
        gr_cont = Contribution(gr, label='Measured')
        gr_ideal_cont = Contribution(gr_ideal,
                                     label=str(factor) + '*p_{T}',
                                     line_color=ROOT.kBlue,
                                     marker_color=ROOT.kBlue)
        plot = Plot([gr_cont, gr_ideal_cont],
                    what='graph',
                    xtitle="p_{T}^{Reco}",
                    ytitle="#sigma [GeV]",
                    ylim=[0, 100],
                    xlim=[10, 4000])
        plot.plot()
        plot.set_logx()
        output_filename = os.path.join(
            plot_dir, "%s_%s_sigma_plot.%s" % (mydir, hist_name, OUTPUT_FMT))
        plot.save(output_filename)
import argparse
import os


import ROOT
ROOT.PyConfig.IgnoreCommandLineOptions = True
ROOT.gROOT.SetBatch(1)
ROOT.TH1.SetDefaultSumw2()
# ROOT.gStyle.SetOptStat(0)
ROOT.gStyle.SetOptFit(1111)

# My stuff
import common_utils as cu


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input", help="Input ROOT file with response hist", required=True)
    parser.add_argument("--inputDir", help="TDirectory in ROOT file", required=True)
    parser.add_argument("--output", help="Output ROOT file to append hist to", required=True)
    args = parser.parse_args()

    hist = cu.grab_obj_from_file(args.input, "%s/pt_jet_response" % (args.inputDir))
    hist.SetName("pt_jet_response")
    hist_rebin = cu.make_normalised_TH2(hist, 'X', recolour=False, do_errors=False)
    hist_rebin.SetName("pt_jet_response")
    outf = cu.open_root_file(args.output, "UPDATE")
    hist_rebin.Write()
    outf.Close()
def main(in_args):
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", help="Input ROOT file with correction graphs", required=True)
    parser.add_argument("--outputDir", help="Output directory for plots", default=os.getcwd())
    parser.add_argument("--title", help="Title string for plots", default="")
    parser.add_argument("--sampleName", help="Sample name string for plots", default="")
    args = parser.parse_args(in_args)

    cu.check_dir_exists_create(args.outputDir)

    lw = 1
    entry_dicts = [
        {"flav": "AbsCorVsJetPt", "label": "All", "colour": ROOT.kBlack, "marker_style": ROOT.kFullCircle, "line_style": 2, "line_width": lw, "marker_size": 1.2},
        {"flav": "ud_AbsCorVsJetPt", "label": "ud", "colour": ROOT.kRed, "marker_style": ROOT.kFullSquare, "line_style": 1, "line_width": lw, "marker_size": 1.2},
        {"flav": "s_AbsCorVsJetPt", "label": "s", "colour": ROOT.kBlue, "marker_style": ROOT.kFullTriangleUp, "line_style": 1, "line_width": lw, "marker_size": 1.4},
        {"flav": "c_AbsCorVsJetPt", "label": "c", "colour": ROOT.kGreen+2, "marker_style": ROOT.kFullTriangleDown, "line_style": 1, "line_width": lw, "marker_size": 1.4},
        {"flav": "b_AbsCorVsJetPt", "label": "b", "colour": ROOT.kOrange-3, "marker_style": ROOT.kFullDiamond, "line_style": 1, "line_width": lw, "marker_size": 1.6},
        {"flav": "g_AbsCorVsJetPt", "label": "g", "colour": ROOT.kAzure+1, "marker_style": 29, "line_style": 1, "line_width": lw, "marker_size": 1.8},
    ]


    all_dirs = cu.get_list_of_element_names(cu.open_root_file(args.input))
    dirs = sorted(list(all_dirs))[:1]
    print("Doing: ", dirs)
    # Loop through all different ak4pfchs, etc
    for mydir in dirs:
        jec_text = ROOT.TPaveText(0.15, 0.91, 0.2, 0.92, "NDC")
        jec_text.AddText(args.title)
        jec_text.SetTextAlign(ROOT.kHAlignLeft + ROOT.kVAlignBottom)
        jec_text.SetTextFont(42)
        jec_text.SetTextSize(FONT_SIZE)
        jec_text.SetBorderSize(0)
        jec_text.SetFillStyle(0)

        dir_text = ROOT.TPaveText(0.17, 0.76, 0.2, 0.77, "NDC")
        dir_label = mydir.upper().replace("PFCHS", " PF CHS").replace("PUPPI", " PUPPI").replace("L1L2L3", " + L1L2L3")
        dir_text.AddText(dir_label)
        dir_text.SetTextAlign(ROOT.kHAlignLeft + ROOT.kVAlignBottom)
        dir_text.SetTextFont(42)
        dir_text.SetTextSize(FONT_SIZE)
        dir_text.SetBorderSize(0)
        dir_text.SetFillStyle(0)

        sample_text = ROOT.TPaveText(0.93, 0.91, 0.94, 0.92, "NDC")
        # sample_text.AddText("Flat QCD 13 TeV")
        sample_text.AddText(args.sampleName + " 13 TeV")
        sample_text.SetTextFont(42)
        sample_text.SetTextSize(FONT_SIZE)
        sample_text.SetTextAlign(ROOT.kHAlignRight + ROOT.kVAlignBottom)
        sample_text.SetBorderSize(0)
        sample_text.SetFillStyle(0)
        sample_text.Draw()

        other_elements = [jec_text, dir_text, sample_text]

        plot_dir = os.path.join(args.outputDir, mydir)
        cu.check_dir_exists_create(plot_dir)

        obj_list = cu.get_list_of_objects_in_dir(args.input, mydir)

        # ylimits = (float(args.ylim[0]), float(args.ylim[1])) if args.ylim else None
        ylimits = None

        X_MIN, X_MAX = 8, 5000
        # For limit protection:
        Y_MIN, Y_MAX = 0.8, 1.6

        # Do all flavs corr vs pt for given eta bin
        common_eta_bins = cu.sort_human(cu.get_common_eta_bins(obj_list))
        # Find the matching + and - eta bins
        eta_bin_pairs = make_pos_neg_eta_pairs(common_eta_bins)
        print(eta_bin_pairs)

        for eta_bin_pair in eta_bin_pairs:
            all_entries = []
            
            title = eta_bin_pair[0].replace("to", " < |#eta| < ").replace("JetEta", "")
            
            # Do a per-flavour comparison plot
            for fdict in entry_dicts:
                entries = []
                for ind, (eta_bin, label) in enumerate(zip(eta_bin_pair, ['+#eta', '-#eta'])):
                    entry = deepcopy(fdict)
                    entry["graph"] = cu.grab_obj_from_file(args.input, "%s/%s_%s" % (mydir, fdict['flav'], eta_bin))
                    entry['label'] += " [%s]" % label
                    entry["line_color"] = fdict['colour']+ind
                    entry["marker_color"] = fdict['colour']+ind
                    entry["fill_color"] = fdict['colour']+ind
                    entry["fill_style"] = 1001
                    entry["fill_alpha"] = 0.7
                    if ind == 1:
                        entry["marker_style"] = get_open_marker(entry['marker_style'])
                        entry["line_style"] += 1
                    if ind == 2:
                        entry["line_style"] += 1
                    if ind == 3:
                        entry["line_style"] += 1
                        entry["marker_style"] = get_open_marker(entry['marker_style'])
                    entries.append(entry)
                    all_entries.append(entry)

                do_comparison_graph(entries, title=title,
                                    xtitle="p_{T}^{Reco} [GeV]", ytitle="Correction", logx=True,
                                    xlimits=(X_MIN, X_MAX), y_limit_protection=(Y_MIN, Y_MAX), 
                                    ylimits=ylimits,
                                    other_elements=other_elements,
                                    output_filename=os.path.join(plot_dir, "compare_corr_vs_pt_%s_pos_neg_%s.pdf" % (eta_bin_pair[0], fdict['label'])))

            # Do a plot with all flavours
            do_comparison_graph(all_entries, title=title,
                                xtitle="p_{T}^{Reco} [GeV]", ytitle="Correction", logx=True,
                                xlimits=(X_MIN, X_MAX), y_limit_protection=(Y_MIN, Y_MAX),
                                other_elements=other_elements,
                                output_filename=os.path.join(plot_dir, "compare_corr_vs_pt_%s_pos_neg_allFlavs.pdf" % (eta_bin_pair[0])))
            
    return 0
def main(in_args):
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input",
        help=
        "Input ROOT file with response & resolution graphs (from jet_response_and_resolution_x)",
        required=True)
    parser.add_argument("--outputDir",
                        help="Output directory for plots",
                        default=os.getcwd())
    parser.add_argument("--title", help="Title string for plots", default="")
    parser.add_argument("--sampleName",
                        help="Sample name string for plots",
                        default="")
    args = parser.parse_args(in_args)

    cu.check_dir_exists_create(args.outputDir)

    lw = 1
    entry_dicts = [
        {
            "flav": "AbsCorVsJetPt",
            "label": "All",
            "colour": ROOT.kBlack,
            "marker_style": ROOT.kFullCircle,
            "line_style": 1,
            "line_width": lw,
            "marker_size": 1.2
        },
        {
            "flav": "ud_AbsCorVsJetPt",
            "label": "ud",
            "colour": ROOT.kRed,
            "marker_style": ROOT.kFullSquare,
            "line_style": 1,
            "line_width": lw,
            "marker_size": 1.2
        },
        {
            "flav": "s_AbsCorVsJetPt",
            "label": "s",
            "colour": ROOT.kBlue,
            "marker_style": ROOT.kFullTriangleUp,
            "line_style": 1,
            "line_width": lw,
            "marker_size": 1.4
        },
        {
            "flav": "c_AbsCorVsJetPt",
            "label": "c",
            "colour": ROOT.kGreen + 2,
            "marker_style": ROOT.kFullTriangleDown,
            "line_style": 1,
            "line_width": lw,
            "marker_size": 1.4
        },
        {
            "flav": "b_AbsCorVsJetPt",
            "label": "b",
            "colour": ROOT.kOrange - 3,
            "marker_style": ROOT.kFullDiamond,
            "line_style": 1,
            "line_width": lw,
            "marker_size": 1.6
        },
        {
            "flav": "g_AbsCorVsJetPt",
            "label": "g",
            "colour": ROOT.kAzure + 1,
            "marker_style": 29,
            "line_style": 1,
            "line_width": lw,
            "marker_size": 1.8
        },
    ]

    # Loop through all different ak4pfchs, etc
    dirs = cu.get_list_of_element_names(cu.open_root_file(args.input))
    for mydir in dirs[:]:
        jec_text = ROOT.TPaveText(0.14, 0.91, 0.2, 0.92, "NDC")
        jec_text.AddText(args.title)
        jec_text.SetTextAlign(ROOT.kHAlignLeft + ROOT.kVAlignBottom)
        jec_text.SetTextFont(42)
        jec_text.SetTextSize(FONT_SIZE)
        jec_text.SetBorderSize(0)
        jec_text.SetFillStyle(0)

        dir_text = ROOT.TPaveText(0.17, 0.71, 0.2, 0.74, "NDC")
        dir_label = mydir.upper().replace("PFCHS", " PF CHS").replace(
            "PUPPI", " PUPPI").replace("L1L2L3",
                                       " + L1L2L3").replace("L1", " + L1")
        dir_text.AddText(dir_label)
        dir_text.SetTextAlign(ROOT.kHAlignLeft + ROOT.kVAlignBottom)
        dir_text.SetTextFont(42)
        dir_text.SetTextSize(FONT_SIZE)
        dir_text.SetBorderSize(0)
        dir_text.SetFillStyle(0)

        sample_text = ROOT.TPaveText(0.93, 0.91, 0.94, 0.92, "NDC")
        # sample_text.AddText("Flat QCD 13 TeV")
        sample_text.AddText(args.sampleName + " 13 TeV")
        sample_text.SetTextFont(42)
        sample_text.SetTextSize(FONT_SIZE)
        sample_text.SetTextAlign(ROOT.kHAlignRight + ROOT.kVAlignBottom)
        sample_text.SetBorderSize(0)
        sample_text.SetFillStyle(0)
        sample_text.Draw()

        other_elements = [jec_text, dir_text, sample_text]

        plot_dir = os.path.join(args.outputDir, mydir)
        cu.check_dir_exists_create(plot_dir)

        obj_list = cu.get_list_of_objects_in_dir(args.input, mydir)

        # Do all flavs rsp vs pt for given eta bin
        common_eta_bins = cu.get_common_eta_bins(obj_list)
        for eta_bin in common_eta_bins:
            entries = []

            # repsonse
            # for fdict in entry_dicts:
            #     entry = deepcopy(fdict)
            #     entry["graph"] = cu.grab_obj_from_file(args.input, "%s/%s_%s" % (mydir, fdict['flav'], eta_bin))
            #     entry["line_color"] = fdict['colour']
            #     entry["marker_color"] = fdict['colour']
            #     entries.append(entry)
            # bin_title = eta_bin.replace("to", " < |#eta| < ").replace("JetEta", "")
            # do_comparison_graph(entries, bin_title=bin_title,
            #                     xtitle="p_{T}^{Gen} [GeV]", ytitle="Response", logx=True,
            #                     xlimits=(10, 5000), y_limit_protection=(0.5, 1.5),
            #                     other_elements=other_elements,
            #                     output_filename=os.path.join(plot_dir, "rsp_vs_pt_%s.pdf" % (eta_bin)))

            # correction
            entries = []
            bin_title = eta_bin.replace("to",
                                        " < #eta < ").replace("JetEta", "")
            for fdict in entry_dicts:
                entry = deepcopy(fdict)
                entry["graph"] = cu.grab_obj_from_file(
                    args.input, "%s/%s_%s" % (mydir, fdict['flav'], eta_bin))
                entry["line_color"] = fdict['colour']
                entry["marker_color"] = fdict['colour']

                entries.append(entry)
                do_comparison_graph([entry],
                                    bin_title=bin_title,
                                    xtitle="p_{T}^{Reco} [GeV]",
                                    ytitle="Correction",
                                    logx=True,
                                    y_limit_protection=(0.7, 1.8),
                                    other_elements=other_elements,
                                    draw_fit_graph_ratio=True,
                                    output_filename=os.path.join(
                                        plot_dir, "%s_%s_logX.pdf" %
                                        (fdict['flav'], eta_bin)))

                do_comparison_graph([entry],
                                    bin_title=bin_title,
                                    xtitle="p_{T}^{Reco} [GeV]",
                                    ytitle="Correction",
                                    logx=False,
                                    y_limit_protection=(0.7, 1.8),
                                    other_elements=other_elements,
                                    draw_fit_graph_ratio=True,
                                    output_filename=os.path.join(
                                        plot_dir, "%s_%s_linX.pdf" %
                                        (fdict['flav'], eta_bin)))

            do_comparison_graph(entries,
                                bin_title=bin_title,
                                xtitle="p_{T}^{Reco} [GeV]",
                                ytitle="Correction",
                                logx=True,
                                y_limit_protection=(0.7, 1.8),
                                other_elements=other_elements,
                                draw_fit_graph_ratio=True,
                                output_filename=os.path.join(
                                    plot_dir,
                                    "corr_vs_pt_%s_logX.pdf" % (eta_bin)))

            do_comparison_graph(entries,
                                bin_title=bin_title,
                                xtitle="p_{T}^{Reco} [GeV]",
                                ytitle="Correction",
                                logx=False,
                                y_limit_protection=(0.7, 1.8),
                                other_elements=other_elements,
                                draw_fit_graph_ratio=True,
                                output_filename=os.path.join(
                                    plot_dir,
                                    "corr_vs_pt_%s_linX.pdf" % (eta_bin)))

    return 0
def main(in_args):
    parser = argparse.ArgumentParser()
    parser.add_argument("--input",
                        help="Input ROOT file with correction graphs",
                        action='append')
    parser.add_argument("--label",
                        help="Label for input file",
                        action='append')
    parser.add_argument("--outputDir",
                        help="Output directory for plots",
                        default=os.getcwd())
    parser.add_argument("--title", help="Title string for plots")
    parser.add_argument("--chi2",
                        help="Do Chi2 comparison plot",
                        action='store_true')
    parser.add_argument("--ylim", help="Set explicit y limits", nargs=2)
    parser.add_argument(
        "--slides",
        help="Dump JSON snippet to file for use with beamer-plot-slides",
        default=None)
    parser.add_argument("--vertLine",
                        help="Add a vertical dashed line at this x value",
                        type=float)
    args = parser.parse_args(in_args)
    print(args)

    if not args.input or len(args.input) == 0:
        return 1

    cu.check_dir_exists_create(args.outputDir)

    if len(args.input) != len(args.label):
        raise RuntimeError("Need a --label for each --input")

    do_slides = args.slides not in (None, "")

    lw = 1
    entry_dicts = [
        # {"flav": "AbsCorVsJetPt", "label": "All", "colour": ROOT.kBlack, "marker_style": ROOT.kFullCircle, "line_style": 2, "line_width": lw, "marker_size": 1.2},
        {
            "flav": "ud_AbsCorVsJetPt",
            "label": "ud",
            "colour": ROOT.kRed,
            "marker_style": ROOT.kFullSquare,
            "line_style": 1,
            "line_width": lw,
            "marker_size": 0.8
        },
        {
            "flav": "g_AbsCorVsJetPt",
            "label": "g",
            "colour": ROOT.kAzure + 1,
            "marker_style": 29,
            "line_style": 1,
            "line_width": lw,
            "marker_size": 1.
        },
        {
            "flav": "s_AbsCorVsJetPt",
            "label": "s",
            "colour": ROOT.kBlue,
            "marker_style": ROOT.kFullTriangleUp,
            "line_style": 1,
            "line_width": lw,
            "marker_size": 1.
        },
        {
            "flav": "c_AbsCorVsJetPt",
            "label": "c",
            "colour": ROOT.kGreen + 2,
            "marker_style": ROOT.kFullTriangleDown,
            "line_style": 1,
            "line_width": lw,
            "marker_size": 1.
        },
        {
            "flav": "b_AbsCorVsJetPt",
            "label": "b",
            "colour": ROOT.kOrange - 3,
            "marker_style": ROOT.kFullDiamond,
            "line_style": 1,
            "line_width": lw,
            "marker_size": 1.2
        },
    ][:]

    sample_str = "QCD MC"

    all_dirs = [
        set(cu.get_list_of_element_names(cu.open_root_file(infile)))
        for infile in args.input
    ]
    dirs = all_dirs[0]
    for d in all_dirs[1:]:
        dirs = dirs & d
    dirs = sorted(list(dirs))[:1]
    print("Doing: ", dirs)
    # Loop through all different ak4pfchs, etc
    slides_contents = []
    for mydir in dirs:
        jec_text = ROOT.TPaveText(0.15, 0.93, 0.2, 0.94, "NDC")
        # jec_label = "Without JEC"
        # jec_label = "With JEC"
        # jec_label = "Summer16_23Sep2016V4"
        # jec_label = "Summer16_03Feb2017_V8"
        jec_text.AddText(args.title)
        jec_text.SetTextAlign(ROOT.kHAlignLeft + ROOT.kVAlignBottom)
        jec_text.SetTextFont(42)
        jec_text.SetTextSize(FONT_SIZE)
        jec_text.SetBorderSize(0)
        jec_text.SetFillStyle(0)

        dir_text = ROOT.TPaveText(0.17, 0.78, 0.2, 0.79, "NDC")
        dir_label = mydir.upper().replace("PFCHS", " PF CHS").replace(
            "PUPPI", " PUPPI").replace("L1L2L3",
                                       " + L1L2L3").replace("L1", " + L1")
        dir_text.AddText(dir_label)
        dir_text.SetTextAlign(ROOT.kHAlignLeft + ROOT.kVAlignBottom)
        dir_text.SetTextFont(42)
        dir_text.SetTextSize(FONT_SIZE)
        dir_text.SetBorderSize(0)
        dir_text.SetFillStyle(0)

        other_elements = [jec_text, dir_text]

        plot_dir = os.path.join(args.outputDir, mydir)
        cu.check_dir_exists_create(plot_dir)

        obj_list = cu.get_list_of_objects_in_dir(args.input[0], mydir)

        ylimits = (float(args.ylim[0]),
                   float(args.ylim[1])) if args.ylim else None

        X_MIN, X_MAX = 7, None
        # For limit protection:
        Y_MIN, Y_MAX = 0.8, 1.6

        py_herwig_ratio_title = "Parton / Hadron"
        py_herwig_ratio_title = "H++ / Py8"

        # Store chi2 for fits over pT
        # Each entry is for a flavour,
        # and is a dict of file label : chi2/dof values
        fit_chi2entries = {}
        for fdict in entry_dicts:
            fit_chi2entries[fdict['label']] = {}
            for label in args.label:
                fit_chi2entries[fdict['label']][label] = []

        # Do all flavs corr vs pt for given eta bin
        common_eta_bins = cu.sort_human(cu.get_common_eta_bins(obj_list))
        for eta_bin in common_eta_bins:
            title = eta_bin.replace("to", " < #eta < ").replace("JetEta", "")
            # Do a per-flavour comparison plot
            outputs = []
            for fdict in entry_dicts:
                entries = []
                chi2entries = []
                for ind, (input_filename,
                          label) in enumerate(zip(args.input, args.label)):
                    entry = deepcopy(fdict)
                    entry["graph"] = cu.grab_obj_from_file(
                        input_filename,
                        "%s/%s_%s" % (mydir, fdict['flav'], eta_bin))
                    entry['label'] += " [%s]" % label
                    new_colour = cu.get_alternate_colour(fdict['colour'], ind)
                    entry["line_color"] = new_colour
                    entry["marker_color"] = new_colour
                    entry["fill_color"] = new_colour
                    entry["fill_style"] = 1001
                    entry["fill_alpha"] = 0.7
                    entry["ratio"] = None
                    if ind == 1:
                        entry["marker_style"] = cu.get_open_marker(
                            entry['marker_style'])
                        entry["line_style"] += 1
                    if ind == 2:
                        entry["line_style"] += 1
                    if ind == 3:
                        entry["line_style"] += 1
                        entry["marker_style"] = cu.get_open_marker(
                            entry['marker_style'])
                    entries.append(entry)
                    if ind != 0:
                        entries[-1]['ratio'] = entries[0]['graph']

                    if args.chi2:
                        # chi2entry = deepcopy(entry)
                        # chi2entry["graph"] = cu.grab_obj_from_file(input_filename, "%s/%s_%s" % (mydir, fdict['flav'].replace("corrVs", "Chi2NDoFVs"), eta_bin))
                        # chi2entries.append(chi2entry)
                        if entry['graph'].GetListOfFunctions().GetSize() > 0:
                            fit = entry['graph'].GetListOfFunctions().Last()
                            # chi2 = entry['graph'].Chisquare(fit)  # can't do this as we don't want to use the plateau in chi2
                            chi2 = fit.GetChisquare()
                            ndof = fit.GetNDF()
                            fit_chi2entries[fdict['label']][label].append(
                                chi2 / ndof)

                output_filename = os.path.abspath(
                    os.path.join(
                        plot_dir,
                        "compare_corr_vs_pt_%s_%s_funcGraphRatio.pdf" %
                        (eta_bin, fdict['label'])))
                cu.do_comparison_graph(
                    entries,
                    title=title,
                    sample_title=sample_str,
                    xtitle="p_{T}^{Reco} [GeV]",
                    ytitle="Correction",
                    logx=True,
                    xlimits=(X_MIN, X_MAX),
                    # xlimits=None,
                    y_limit_protection=(Y_MIN, Y_MAX),
                    ylimits=ylimits,
                    other_elements=other_elements,
                    output_filename=output_filename,
                    vertical_line=args.vertLine,
                    do_fit_graph_ratio=True)
                outputs.append(output_filename)

                output_filename = os.path.abspath(
                    os.path.join(
                        plot_dir,
                        "compare_corr_vs_pt_%s_%s_pyHerwigRatio.pdf" %
                        (eta_bin, fdict['label'])))
                cu.do_comparison_graph(
                    entries,
                    title=title,
                    sample_title=sample_str,
                    xtitle="p_{T}^{Reco} [GeV]",
                    ytitle="Correction",
                    logx=True,
                    xlimits=(X_MIN, X_MAX),
                    # xlimits=None,
                    y_limit_protection=(Y_MIN, Y_MAX),
                    ylimits=ylimits,
                    other_elements=other_elements,
                    output_filename=output_filename,
                    vertical_line=args.vertLine,
                    draw_fits=True,
                    do_ratio_plots=True,
                    ratio_title=py_herwig_ratio_title)

                output_filename = os.path.abspath(
                    os.path.join(
                        plot_dir, "compare_corr_vs_pt_%s_%s_pull.pdf" %
                        (eta_bin, fdict['label'])))
                cu.do_comparison_graph(
                    entries,
                    title=title,
                    sample_title=sample_str,
                    xtitle="p_{T}^{Reco} [GeV]",
                    ytitle="Correction",
                    logx=True,
                    xlimits=(X_MIN, X_MAX),
                    # xlimits=None,
                    y_limit_protection=(Y_MIN, Y_MAX),
                    ylimits=ylimits,
                    other_elements=other_elements,
                    output_filename=output_filename,
                    vertical_line=args.vertLine,
                    do_fit_graph_pull=True)
                # if args.chi2:
                #     cu.do_comparison_graph(chi2entries, title=title, sample_title=sample_str,
                #                         xtitle="p_{T}^{Reco} [GeV]", ytitle="#chi^{2}/N_{DoF}", logx=True,
                #                         xlimits=(X_MIN, X_MAX),
                #                         other_elements=other_elements,
                #                         output_filename=os.path.join(plot_dir, "compare_corr_vs_pt_%s_%s_chi2.pdf" % (eta_bin, fdict['label'])))

            # Do a plot with all flavours
            entries = []
            ud_entries, g_entries = [], []
            for fdict in entry_dicts[:]:
                for ind, (input_filename,
                          label) in enumerate(zip(args.input, args.label)):
                    entry = deepcopy(fdict)
                    entry["graph"] = cu.grab_obj_from_file(
                        input_filename,
                        "%s/%s_%s" % (mydir, fdict['flav'], eta_bin))
                    entry['label'] += " [%s]" % label
                    new_colour = cu.get_alternate_colour(fdict['colour'], ind)
                    entry["line_color"] = new_colour
                    entry["marker_color"] = new_colour
                    entry["fill_color"] = new_colour
                    entry["fill_style"] = 1001
                    entry["fill_alpha"] = 0.8
                    entry['ratio'] = None
                    if ind == 1:
                        entry["marker_style"] = cu.get_open_marker(
                            entry['marker_style'])
                        entry["line_style"] += 1
                    if ind == 2:
                        entry["line_style"] += 1
                    if ind == 3:
                        entry["line_style"] += 1
                        entry["marker_style"] = cu.get_open_marker(
                            entry['marker_style'])
                    entries.append(entry)
                    if fdict['label'] == "ud":
                        ud_entries.append(entry)
                    elif fdict['label'] == "g":
                        g_entries.append(entry)
                    if ind > 0:
                        entries[-1]['ratio'] = entries[-2]['graph']

            title = eta_bin.replace("to", " < #eta < ").replace("JetEta", "")
            output_filename = os.path.abspath(
                os.path.join(
                    plot_dir,
                    "compare_corr_vs_pt_%s_allFlavs_funcGraphRatio.pdf" %
                    (eta_bin)))
            cu.do_comparison_graph(entries,
                                   title=title,
                                   sample_title=sample_str,
                                   xtitle="p_{T}^{Reco} [GeV]",
                                   ytitle="Correction",
                                   logx=True,
                                   xlimits=(X_MIN, X_MAX),
                                   y_limit_protection=(Y_MIN, Y_MAX),
                                   ylimits=ylimits,
                                   other_elements=other_elements,
                                   output_filename=output_filename,
                                   vertical_line=args.vertLine,
                                   do_fit_graph_ratio=True)
            outputs.insert(0, output_filename)

            output_filename = os.path.abspath(
                os.path.join(
                    plot_dir,
                    "compare_corr_vs_pt_%s_allFlavs_pyHerwigRatio.pdf" %
                    (eta_bin)))
            cu.do_comparison_graph(entries,
                                   title=title,
                                   sample_title=sample_str,
                                   xtitle="p_{T}^{Reco} [GeV]",
                                   ytitle="Correction",
                                   logx=True,
                                   xlimits=(X_MIN, X_MAX),
                                   y_limit_protection=(Y_MIN, Y_MAX),
                                   ylimits=ylimits,
                                   other_elements=other_elements,
                                   output_filename=output_filename,
                                   vertical_line=args.vertLine,
                                   draw_fits=True,
                                   do_ratio_plots=True,
                                   ratio_title=py_herwig_ratio_title)

            # make 1 slide with all flavours for this eta bin
            if do_slides:
                plots = ",\n".join([
                    '            ["{{{fname}}}{ext}", "{title}"]'.format(
                        fname=os.path.splitext(fname)[0],
                        ext=os.path.splitext(fname)[1],
                        title="") for fname in outputs
                ])
                text = """    {{
    "title": "${thistitle}$",
    "plots": [
{plots}
        ]
    }}""".format(thistitle=title.replace("#", r"\\"), plots=plots)

                slides_contents.append(text)

        if args.chi2:
            # for each flav, do a plot of chi2 of total fit vs eta for all files
            all_entries = []
            all_ymax = 0
            for fdict in entry_dicts:
                entries = []
                file_labels = list(fit_chi2entries[fdict['label']].keys())
                ymax = 0
                for ind, flabel in enumerate(file_labels):
                    eta_bins = [get_eta_mid(x) for x in common_eta_bins]
                    y = fit_chi2entries[fdict['label']][flabel]
                    ymax = max(ymax, max(y))
                    ey = [0] * len(y)
                    ex = [get_eta_half_width(x) for x in common_eta_bins]
                    gr = ROOT.TGraphErrors(len(y), array('d', eta_bins),
                                           array('d', y), array('d', ex),
                                           array('d', ey))
                    entry = deepcopy(fdict)
                    entry["graph"] = gr
                    entry['label'] = "%s [%s]" % (fdict['label'], flabel)
                    new_colour = cu.get_alternate_colour(fdict['colour'], ind)
                    entry["line_color"] = new_colour
                    entry["marker_color"] = new_colour
                    entry["fill_color"] = new_colour
                    entry["fill_style"] = 1001
                    entry["fill_alpha"] = 0.8
                    if ind == 1:
                        entry["marker_style"] = cu.get_open_marker(
                            entry['marker_style'])
                        entry["line_style"] += 1
                    if ind == 2:
                        entry["line_style"] += 1
                    if ind == 3:
                        entry["line_style"] += 1
                        entry["marker_style"] = cu.get_open_marker(
                            entry['marker_style'])
                    entries.append(entry)
                all_ymax = max(all_ymax, ymax)
                ylimits = [0, ymax * 1.3]
                # if ylimits[1] > 20:
                #     ylimits[1] = 20
                output_filename = os.path.abspath(
                    os.path.join(
                        plot_dir,
                        "compare_chi2_vs_eta_%s.pdf" % (fdict['label'])))
                cu.do_comparison_graph(entries,
                                       title="",
                                       sample_title=sample_str,
                                       xtitle="#eta^{Reco}",
                                       ytitle="#chi^{2} / N_{dof}",
                                       draw_opt="ALP",
                                       ylimits=ylimits,
                                       other_elements=other_elements,
                                       output_filename=output_filename)
                all_entries.extend(entries)

            # do a single plot with all flavs
            ylimits = [0, all_ymax * 1.3]
            if ylimits[1] > 20:
                ylimits[1] = 20
            output_filename = os.path.abspath(
                os.path.join(plot_dir, "compare_chi2_vs_eta_allFlavs.pdf"))
            cu.do_comparison_graph(all_entries,
                                   title="",
                                   sample_title=sample_str,
                                   xtitle="#eta^{Reco}",
                                   ytitle="#chi^{2} / N_{dof}",
                                   draw_opt="ALP",
                                   ylimits=ylimits,
                                   other_elements=other_elements,
                                   output_filename=output_filename)

            # diff_entries = []
            # for ind, (ud, g, label) in enumerate(zip(ud_entries, g_entries, args.label)):
            #     diff_graph = construct_difference_graph(ud['graph'], g['graph'])
            #     diff = deepcopy(ud)
            #     diff['graph'] = diff_graph
            #     diff['label'] = label
            #     diff_entries.append(diff)
            # cu.do_comparison_graph(diff_entries, title=title, sample_title=sample_str,
            #                     xtitle="p_{T}^{Reco} [GeV]", ytitle="Correction (ud) - Correction (g)", logx=True,
            #                     xlimits=(X_MIN, X_MAX), ylimits=(0, 0.08),
            #                     other_elements=other_elements,
            #                     output_filename=os.path.join(plot_dir, "compare_corr_vs_pt_%s_ud_g_diff.pdf" % (eta_bin)))

        # # Do all flavs corr vs eta for given pt bin
        # common_pt_bins = cu.sort_human(cu.get_common_pt_bins(obj_list))
        # for pt_bin in common_pt_bins:
        #     # Do a per-flavour comparison plot
        #     for fdict in entry_dicts:
        #         entries = []
        #         chi2entries = []
        #         for ind, (input_filename, label) in enumerate(zip(args.input, args.label)):
        #             entry = deepcopy(fdict)
        #             entry["graph"] = cu.grab_obj_from_file(input_filename, "%s/%s_%s" % (mydir, fdict['flav'].replace("RefPt", "JetEta"), pt_bin))
        #             entry['label'] += " [%s]" % label
        #             entry["line_color"] = fdict['colour']
        #             entry["marker_color"] = fdict['colour']
        #             if ind == 1:
        #                 entry["marker_style"] = cu.get_open_marker(entry['marker_style'])
        #             if ind == 2:
        #                 entry["line_style"] += 1
        #             if ind == 3:
        #                 entry["line_style"] += 1
        #                 entry["marker_style"] = cu.get_open_marker(entry['marker_style'])
        #             entries.append(entry)
        #             if args.chi2:
        #                 chi2entry = deepcopy(entry)
        #                 chi2entry["graph"] = cu.grab_obj_from_file(input_filename, "%s/%s_%s" % (mydir, fdict['flav'].replace("RefPt", "JetEta").replace("corrVs", "Chi2NDoFVs"), pt_bin))
        #                 chi2entries.append(chi2entry)

        #         title = pt_bin.replace("to", " < p_{T} < ").replace("RefPt", "")
        #         cu.do_comparison_graph(entries, title=title + " GeV", sample_title=sample_str,
        #                             xtitle="|#eta|", ytitle="Correction",
        #                             xlimits=(0, 5.2), y_limit_protection=(Y_MIN, Y_MAX),
        #                             other_elements=other_elements, ylimits=ylimits,
        #                             output_filename=os.path.join(plot_dir, "compare_corr_vs_eta_%s_%s.pdf" % (pt_bin, fdict['label'])))
        #         if args.chi2:
        #             cu.do_comparison_graph(entries, title=title, sample_title=sample_str,
        #                                 xtitle="|#eta|", ytitle="#chi^{2}/N_{DoF}",
        #                                 xlimits=(0, 5.2),
        #                                 other_elements=other_elements,
        #                                 output_filename=os.path.join(plot_dir, "compare_corr_vs_eta_%s_%s_chi2.pdf" % (pt_bin, fdict['label'])))

        #     # Do a plot with all flavours
        #     entries = []
        #     for fdict in entry_dicts:
        #         for ind, (input_filename, label) in enumerate(zip(args.input, args.label)):
        #             entry = deepcopy(fdict)
        #             entry["graph"] = cu.grab_obj_from_file(input_filename, "%s/%s_%s" % (mydir, fdict['flav'].replace("RefPt", "JetEta"), pt_bin))
        #             entry['label'] += " [%s]" % label
        #             entry["line_color"] = fdict['colour']
        #             entry["marker_color"] = fdict['colour']
        #             if ind == 1:
        #                 entry["marker_style"] = cu.get_open_marker(entry['marker_style'])
        #             if ind == 2:
        #                 entry["line_style"] += 1
        #             if ind == 3:
        #                 entry["line_style"] += 1
        #                 entry["marker_style"] = cu.get_open_marker(entry['marker_style'])
        #             entries.append(entry)
        #     title = pt_bin.replace("to", " < p_{T} < ").replace("RefPt", "")
        #     cu.do_comparison_graph(entries, title=title + " GeV", sample_title=sample_str,
        #                         xtitle="|#eta|", ytitle="Correction",
        #                         xlimits=(0, 5.2), y_limit_protection=(Y_MIN, Y_MAX),
        #                         other_elements=other_elements, ylimits=ylimits,
        #                         output_filename=os.path.join(plot_dir, "compare_corr_vs_eta_%s_allFlavs.pdf" % (pt_bin)))

    if do_slides:
        print("Writing JSON to", args.slides)
        with open(args.slides, "w") as fout:
            fout.write(",\n".join(slides_contents))
    return 0
Ejemplo n.º 6
0
def main(in_args):
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input",
        help=
        "Input ROOT file with response & resolution graphs (from jet_response_and_resolution_x)",
        action='append')
    parser.add_argument("--label",
                        help="Label for input file",
                        action='append')
    parser.add_argument("--outputDir",
                        help="Output directory for plots",
                        default=os.getcwd())
    parser.add_argument("--title", help="Title string for plots")
    parser.add_argument("--chi2",
                        help="Do Chi2 comparison plot",
                        action='store_true')
    parser.add_argument("--ylim", help="Set explicit y limits", nargs=2)
    args = parser.parse_args(in_args)

    cu.check_dir_exists_create(args.outputDir)

    if len(args.input) != len(args.label):
        raise RuntimeError("Need a --label for each --input")

    if args.input:

        lw = 2
        entry_dicts = [
            {
                "flav": "RelRspVsRefPt",
                "label": "All",
                "colour": ROOT.kBlack,
                "marker_style": ROOT.kFullCircle,
                "line_style": 2,
                "line_width": lw,
                "marker_size": 1.2
            },
            {
                "flav": "ud_RspVsRefPt_RelRsp",
                "label": "ud",
                "colour": ROOT.kRed,
                "marker_style": ROOT.kFullSquare,
                "line_style": 1,
                "line_width": lw,
                "marker_size": 1.2
            },
            {
                "flav": "s_RRspVsRefPt_RelRsp",
                "label": "s",
                "colour": ROOT.kBlue,
                "marker_style": ROOT.kFullTriangleUp,
                "line_style": 1,
                "line_width": lw,
                "marker_size": 1.4
            },
            {
                "flav": "c_RRspVsRefPt_RelRsp",
                "label": "c",
                "colour": ROOT.kGreen + 2,
                "marker_style": ROOT.kFullTriangleDown,
                "line_style": 1,
                "line_width": lw,
                "marker_size": 1.4
            },
            {
                "flav": "b_RRspVsRefPt_RelRsp",
                "label": "b",
                "colour": ROOT.kOrange - 3,
                "marker_style": ROOT.kFullDiamond,
                "line_style": 1,
                "line_width": lw,
                "marker_size": 1.6
            },
            {
                "flav": "g_RRspVsRefPt_RelRsp",
                "label": "g",
                "colour": ROOT.kAzure + 1,
                "marker_style": 29,
                "line_style": 1,
                "line_width": lw,
                "marker_size": 1.8
            },
        ]

        all_dirs = [
            set(cu.get_list_of_element_names(cu.open_root_file(infile)))
            for infile in args.input
        ]
        dirs = all_dirs[0]
        for d in all_dirs[1:]:
            dirs = dirs & d
        dirs = sorted(list(dirs))[:1]
        print("Doing: ", dirs)
        # Loop through all different ak4pfchs, etc
        for mydir in dirs:
            jec_text = ROOT.TPaveText(0.17, 0.91, 0.2, 0.92, "NDC")
            jec_label = "Without JEC"
            # jec_label = "With JEC"
            # jec_label = "Summer16_23Sep2016V4"
            # jec_label = "Summer16_03Feb2017_V8"
            jec_text.AddText(args.title)
            jec_text.SetTextAlign(ROOT.kHAlignLeft + ROOT.kVAlignBottom)
            jec_text.SetTextFont(42)
            jec_text.SetTextSize(FONT_SIZE)
            jec_text.SetBorderSize(0)
            jec_text.SetFillStyle(0)

            dir_text = ROOT.TPaveText(0.17, 0.76, 0.2, 0.77, "NDC")
            dir_label = mydir.upper().replace("PFCHS", " PF CHS").replace(
                "PUPPI", " PUPPI").replace("L1L2L3", " + L1L2L3")
            dir_text.AddText(dir_label)
            dir_text.SetTextAlign(ROOT.kHAlignLeft + ROOT.kVAlignBottom)
            dir_text.SetTextFont(42)
            dir_text.SetTextSize(FONT_SIZE)
            dir_text.SetBorderSize(0)
            dir_text.SetFillStyle(0)

            other_elements = [jec_text, dir_text]

            plot_dir = os.path.join(args.outputDir, mydir)
            cu.check_dir_exists_create(plot_dir)

            obj_list = cu.get_list_of_objects_in_dir(args.input[0], mydir)

            ylimits = (float(args.ylim[0]),
                       float(args.ylim[1])) if args.ylim else None

            # Do all flavs rsp vs pt for given eta bin
            common_eta_bins = cu.get_common_eta_bins(obj_list)
            for eta_bin in common_eta_bins:
                # Do a per-flavour comparison plot
                for fdict in entry_dicts:
                    entries = []
                    chi2entries = []
                    for ind, (input_filename,
                              label) in enumerate(zip(args.input, args.label)):
                        entry = deepcopy(fdict)
                        entry["graph"] = cu.grab_obj_from_file(
                            input_filename,
                            "%s/%s_%s" % (mydir, fdict['flav'], eta_bin))
                        entry['label'] += " [%s]" % label
                        entry["line_color"] = fdict['colour']
                        entry["marker_color"] = fdict['colour']
                        if ind == 1:
                            entry["marker_style"] = get_open_marker(
                                entry['marker_style'])
                        if ind == 2:
                            entry["line_style"] += 1
                        if ind == 3:
                            entry["line_style"] += 1
                            entry["marker_style"] = get_open_marker(
                                entry['marker_style'])
                        entries.append(entry)

                        if args.chi2:
                            chi2entry = deepcopy(entry)
                            chi2entry["graph"] = cu.grab_obj_from_file(
                                input_filename,
                                "%s/%s_%s" % (mydir, fdict['flav'].replace(
                                    "RspVs", "Chi2NDoFVs"), eta_bin))
                            chi2entries.append(chi2entry)

                    title = eta_bin.replace("to", " < |#eta| < ").replace(
                        "JetEta", "")
                    do_comparison_graph(entries,
                                        title=title,
                                        xtitle="p_{T}^{Gen} [GeV]",
                                        ytitle="Response",
                                        logx=True,
                                        xlimits=(10, 3000),
                                        y_limit_protection=(0.7, 1.4),
                                        ylimits=ylimits,
                                        other_elements=other_elements,
                                        output_filename=os.path.join(
                                            plot_dir,
                                            "compare_rsp_vs_pt_%s_%s.pdf" %
                                            (eta_bin, fdict['label'])))
                    if args.chi2:
                        do_comparison_graph(
                            chi2entries,
                            title=title,
                            xtitle="p_{T}^{Gen} [GeV]",
                            ytitle="#chi^{2}/N_{DoF}",
                            logx=True,
                            xlimits=(10, 3000),
                            other_elements=other_elements,
                            output_filename=os.path.join(
                                plot_dir, "compare_rsp_vs_pt_%s_%s_chi2.pdf" %
                                (eta_bin, fdict['label'])))

                # Do a plot with all flavours
                entries = []
                ud_entries, g_entries = [], []
                for fdict in entry_dicts:
                    for ind, (input_filename,
                              label) in enumerate(zip(args.input, args.label)):
                        entry = deepcopy(fdict)
                        entry["graph"] = cu.grab_obj_from_file(
                            input_filename,
                            "%s/%s_%s" % (mydir, fdict['flav'], eta_bin))
                        entry['label'] += " [%s]" % label
                        entry["line_color"] = fdict['colour']
                        entry["marker_color"] = fdict['colour']
                        if ind == 1:
                            entry["marker_style"] = get_open_marker(
                                entry['marker_style'])
                        if ind == 2:
                            entry["line_style"] += 1
                        if ind == 3:
                            entry["line_style"] += 1
                            entry["marker_style"] = get_open_marker(
                                entry['marker_style'])
                        entries.append(entry)
                        if fdict['label'] == "ud":
                            ud_entries.append(entry)
                        elif fdict['label'] == "g":
                            g_entries.append(entry)

                title = eta_bin.replace("to",
                                        " < |#eta| < ").replace("JetEta", "")
                do_comparison_graph(
                    entries,
                    title=title,
                    xtitle="p_{T}^{Gen} [GeV]",
                    ytitle="Response",
                    logx=True,
                    xlimits=(10, 3000),
                    y_limit_protection=(0.7, 1.4),
                    other_elements=other_elements,
                    output_filename=os.path.join(
                        plot_dir,
                        "compare_rsp_vs_pt_%s_allFlavs.pdf" % (eta_bin)))
                # diff_entries = []
                # for ind, (ud, g, label) in enumerate(zip(ud_entries, g_entries, args.label)):
                #     diff_graph = construct_difference_graph(ud['graph'], g['graph'])
                #     diff = deepcopy(ud)
                #     diff['graph'] = diff_graph
                #     diff['label'] = label
                #     diff_entries.append(diff)
                # do_comparison_graph(diff_entries, title=title,
                #                     xtitle="p_{T}^{Gen} [GeV]", ytitle="Response (ud) - response (g)", logx=True,
                #                     xlimits=(10, 3000), ylimits=(0, 0.08),
                #                     other_elements=other_elements,
                #                     output_filename=os.path.join(plot_dir, "compare_rsp_vs_pt_%s_ud_g_diff.pdf" % (eta_bin)))
            return

            # Do all flavs rsp vs eta for given pt bin
            common_pt_bins = cu.get_common_pt_bins(obj_list)
            for pt_bin in common_pt_bins:
                # Do a per-flavour comparison plot
                for fdict in entry_dicts:
                    entries = []
                    chi2entries = []
                    for ind, (input_filename,
                              label) in enumerate(zip(args.input, args.label)):
                        entry = deepcopy(fdict)
                        entry["graph"] = cu.grab_obj_from_file(
                            input_filename,
                            "%s/%s_%s" % (mydir, fdict['flav'].replace(
                                "RefPt", "JetEta"), pt_bin))
                        entry['label'] += " [%s]" % label
                        entry["line_color"] = fdict['colour']
                        entry["marker_color"] = fdict['colour']
                        if ind == 1:
                            entry["marker_style"] = get_open_marker(
                                entry['marker_style'])
                        if ind == 2:
                            entry["line_style"] += 1
                        if ind == 3:
                            entry["line_style"] += 1
                            entry["marker_style"] = get_open_marker(
                                entry['marker_style'])
                        entries.append(entry)
                        if args.chi2:
                            chi2entry = deepcopy(entry)
                            chi2entry["graph"] = cu.grab_obj_from_file(
                                input_filename,
                                "%s/%s_%s" % (mydir, fdict['flav'].replace(
                                    "RefPt", "JetEta").replace(
                                        "RspVs", "Chi2NDoFVs"), pt_bin))
                            chi2entries.append(chi2entry)

                    title = pt_bin.replace("to",
                                           " < p_{T} < ").replace("RefPt", "")
                    do_comparison_graph(entries,
                                        title=title + " GeV",
                                        xtitle="|#eta|",
                                        ytitle="Response",
                                        xlimits=(0, 5.2),
                                        y_limit_protection=(0.8, 1.4),
                                        other_elements=other_elements,
                                        ylimits=ylimits,
                                        output_filename=os.path.join(
                                            plot_dir,
                                            "compare_rsp_vs_eta_%s_%s.pdf" %
                                            (pt_bin, fdict['label'])))
                    if args.chi2:
                        do_comparison_graph(
                            entries,
                            title=title,
                            xtitle="|#eta|",
                            ytitle="#chi^{2}/N_{DoF}",
                            xlimits=(0, 5.2),
                            other_elements=other_elements,
                            output_filename=os.path.join(
                                plot_dir, "compare_rsp_vs_eta_%s_%s_chi2.pdf" %
                                (pt_bin, fdict['label'])))

                # Do a plot with all flavours
                entries = []
                for fdict in entry_dicts:
                    for ind, (input_filename,
                              label) in enumerate(zip(args.input, args.label)):
                        entry = deepcopy(fdict)
                        entry["graph"] = cu.grab_obj_from_file(
                            input_filename,
                            "%s/%s_%s" % (mydir, fdict['flav'].replace(
                                "RefPt", "JetEta"), pt_bin))
                        entry['label'] += " [%s]" % label
                        entry["line_color"] = fdict['colour']
                        entry["marker_color"] = fdict['colour']
                        if ind == 1:
                            entry["marker_style"] = get_open_marker(
                                entry['marker_style'])
                        if ind == 2:
                            entry["line_style"] += 1
                        if ind == 3:
                            entry["line_style"] += 1
                            entry["marker_style"] = get_open_marker(
                                entry['marker_style'])
                        entries.append(entry)
                title = pt_bin.replace("to",
                                       " < p_{T} < ").replace("RefPt", "")
                do_comparison_graph(
                    entries,
                    title=title + " GeV",
                    xtitle="|#eta|",
                    ytitle="Response",
                    xlimits=(0, 5.2),
                    y_limit_protection=(0.8, 1.4),
                    other_elements=other_elements,
                    ylimits=ylimits,
                    output_filename=os.path.join(
                        plot_dir,
                        "compare_rsp_vs_eta_%s_allFlavs.pdf" % (pt_bin)))

    return 0
def main(in_args):
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", help="Input ROOT file with response & resolution graphs (from jet_response_and_resolution_x)", required=True)
    parser.add_argument("--outputDir", help="Output directory for plots", default=os.getcwd())
    parser.add_argument("--title", help="Title string for plots", default="")
    parser.add_argument("--sampleName", help="Sample name string for plots", default="")
    args = parser.parse_args(in_args)

    cu.check_dir_exists_create(args.outputDir)

    lw = 2
    entry_dicts = [
        {"flav": "RelRspVsRefPt", "label": "All", "colour": ROOT.kBlack, "marker_style": ROOT.kFullCircle, "line_style": 2, "line_width": lw, "marker_size": 1.2},
        {"flav": "ud_RspVsRefPt_RelRsp", "label": "ud", "colour": ROOT.kRed, "marker_style": ROOT.kFullSquare, "line_style": 1, "line_width": lw, "marker_size": 1.2},
        {"flav": "s_RRspVsRefPt_RelRsp", "label": "s", "colour": ROOT.kBlue, "marker_style": ROOT.kFullTriangleUp, "line_style": 1, "line_width": lw, "marker_size": 1.4},
        {"flav": "c_RRspVsRefPt_RelRsp", "label": "c", "colour": ROOT.kGreen+2, "marker_style": ROOT.kFullTriangleDown, "line_style": 1, "line_width": lw, "marker_size": 1.4},
        {"flav": "b_RRspVsRefPt_RelRsp", "label": "b", "colour": ROOT.kOrange-3, "marker_style": ROOT.kFullDiamond, "line_style": 1, "line_width": lw, "marker_size": 1.6},
        {"flav": "g_RRspVsRefPt_RelRsp", "label": "g", "colour": ROOT.kAzure+1, "marker_style": 29, "line_style": 1, "line_width": lw, "marker_size": 1.8},
    ]

    # no JEC
    eta_bin_limits = {
        'JetEta0to0.783': (0.9, 1.15),
        'JetEta0.783to1.305': None,
        'JetEta1.305to1.93': (0.85, 1.05),
        'JetEta1.93to2.5': None,
        'JetEta2.5to2.964': None,
        'JetEta2.964to3.139': None,
        'JetEta3.139to5.191': (0.98, 1.8),
    }
    eta_bin_ratio_limits = {
        'JetEta0to0.783': (0.96, 1.07),
        'JetEta0.783to1.305': None,
        'JetEta1.305to1.93': (0.95, 1.07),
        'JetEta1.93to2.5': None,
        'JetEta2.5to2.964': None,
        'JetEta2.964to3.139': None,
        'JetEta3.139to5.191': (0.93, 1.1),
    }
    
    # with L1 JEC
    # eta_bin_limits = {
    #     'JetEta0to0.783': (0.76, 1.05),
    #     'JetEta0.783to1.305': (0.77, 1),
    #     'JetEta1.305to1.93': (0.7, 1.),
    #     'JetEta1.93to2.5': (0.68, 1.1),
    #     'JetEta2.5to2.964': (0.55, 1.1),
    #     'JetEta2.964to3.139': (0.6, 1.1),
    #     'JetEta3.139to5.191': (0.75, 1.1),
    # }
    # eta_bin_ratio_limits = {
    #     'JetEta0to0.783': (0.96, 1.08),
    #     'JetEta0.783to1.305': (0.96, 1.08),
    #     'JetEta1.305to1.93': (0.95, 1.09),
    #     'JetEta1.93to2.5': (0.95, 1.13),
    #     'JetEta2.5to2.964': (0.93, 1.15),
    #     'JetEta2.964to3.139': (0.9, 1.12),
    #     'JetEta3.139to5.191': (0.93, 1.1),
    # }

    # default
    # eta_bin_limits = {
    #     'JetEta0to0.783': None,
    #     'JetEta0.783to1.305': None,
    #     'JetEta1.305to1.93': None,
    #     'JetEta1.93to2.5': None,
    #     'JetEta2.5to2.964': None,
    #     'JetEta2.964to3.139': None,
    #     'JetEta3.139to5.191': None,
    # }
    # eta_bin_ratio_limits = {
    #     'JetEta0to0.783': None,
    #     'JetEta0.783to1.305': None,
    #     'JetEta1.305to1.93': None,
    #     'JetEta1.93to2.5': None,
    #     'JetEta2.5to2.964': None,
    #     'JetEta2.964to3.139': None,
    #     'JetEta3.139to5.191': None,
    # }


    # Loop through all different ak4pfchs, etc
    dirs = cu.get_list_of_element_names(cu.open_root_file(args.input))
    for mydir in dirs[:]:
        jec_text = ROOT.TPaveText(0.14, 0.93, 0.2, 0.96, "NDC")
        jec_text.AddText(args.title)
        jec_text.SetTextAlign(ROOT.kHAlignLeft + ROOT.kVAlignBottom)
        jec_text.SetTextFont(42)
        jec_text.SetTextSize(FONT_SIZE)
        jec_text.SetBorderSize(0)
        jec_text.SetFillStyle(0)

        dir_text = ROOT.TPaveText(0.17, 0.75, 0.2, 0.77, "NDC")
        dir_label = mydir.upper().replace("PFCHS", " PF CHS").replace("PUPPI", " PUPPI").replace("L1L2L3", " + L1L2L3")
        dir_text.AddText(dir_label)
        dir_text.SetTextAlign(ROOT.kHAlignLeft + ROOT.kVAlignBottom)
        dir_text.SetTextFont(42)
        dir_text.SetTextSize(FONT_SIZE)
        dir_text.SetBorderSize(0)
        dir_text.SetFillStyle(0)

        sample_text = ROOT.TPaveText(0.94, 0.93, 0.95, 0.96, "NDC")
        # sample_text.AddText("Flat QCD 13 TeV")
        sample_text.AddText(args.sampleName) # + " 13 TeV")
        sample_text.SetTextFont(42)
        sample_text.SetTextSize(FONT_SIZE)
        sample_text.SetTextAlign(ROOT.kHAlignRight + ROOT.kVAlignBottom)
        sample_text.SetBorderSize(0)
        sample_text.SetFillStyle(0)
        sample_text.Draw()

        other_elements = [jec_text, dir_text, sample_text]

        plot_dir = os.path.join(args.outputDir, mydir)
        cu.check_dir_exists_create(plot_dir)

        obj_list = cu.get_list_of_objects_in_dir(args.input, mydir)

        # Do all flavs rsp vs pt for given eta bin
        common_eta_bins = cu.sort_human(cu.get_common_eta_bins(obj_list))
        for eta_bin in common_eta_bins[:]:
            entries = []
            contributions = []
            for fdict in entry_dicts:
                entry = deepcopy(fdict)
                # entry["graph"] = cu.grab_obj_from_file(args.input, "%s/%s_%s" % (mydir, fdict['flav'], eta_bin))
                entry["obj"] = cu.grab_obj_from_file(args.input, "%s/%s_%s" % (mydir, fdict['flav'], eta_bin))
                entry["line_color"] = fdict['colour']
                entry["marker_color"] = fdict['colour']
                # entries.append(entry)

                del entry['flav']
                del entry['colour']
                this_contrib = Contribution(**entry)
                contributions.append(this_contrib)

            bin_title = eta_bin.replace("to", " < |#eta| < ").replace("JetEta", "")
            do_comparison_graph(contributions, bin_title=bin_title,
                                xtitle="p_{T}^{Gen} [GeV]", ytitle="Response", logx=True,
                                xlimits=(10, 5000), 
                                ylimits=eta_bin_limits.get(eta_bin, None),
                                y_limit_protection=(0.5, 1.7),
                                other_elements=other_elements,
                                output_filename=os.path.join(plot_dir, "rsp_vs_pt_%s.pdf" % (eta_bin)),
                                ratio_limits=eta_bin_ratio_limits.get(eta_bin, None)
                                )

            # Inverse response ie ~ correction
            # entries = []
            # for fdict in entry_dicts:
            #     entry = deepcopy(fdict)
            #     entry["graph"] = construct_inverse_graph(cu.grab_obj_from_file(args.input, "%s/%s_%s" % (mydir, fdict['flav'], eta_bin)))
            #     entry["line_color"] = fdict['colour']
            #     entry["marker_color"] = fdict['colour']
            #     entries.append(entry)
            # bin_title = eta_bin.replace("to", " < |#eta| < ").replace("JetEta", "")
            # do_comparison_graph(entries, bin_title=bin_title,
            #                     xtitle="p_{T}^{Gen} [GeV]", ytitle="1/Response", logx=True,
            #                     y_limit_protection=(0.8, 1.2),
            #                     other_elements=other_elements,
            #                     output_filename=os.path.join(plot_dir, "inv_rsp_vs_pt_%s.pdf" % (eta_bin)))

        # Do all flavs rsp vs eta for given pt bin
        # common_pt_bins = cu.get_common_pt_bins(obj_list)
        # for pt_bin in common_pt_bins:
        #     entries = []
        #     contributions = []
        #     for fdict in entry_dicts:
        #         entry = deepcopy(fdict)
        #         # entry["graph"] = cu.grab_obj_from_file(args.input, "%s/%s_%s" % (mydir, fdict['flav'].replace("RefPt", "JetEta"), pt_bin))
        #         entry["obj"] = cu.grab_obj_from_file(args.input, "%s/%s_%s" % (mydir, fdict['flav'].replace("RefPt", "JetEta"), pt_bin))
        #         entry["line_color"] = fdict['colour']
        #         entry["marker_color"] = fdict['colour']
        #         # entries.append(entry)

        #         del entry['flav']
        #         del entry['colour']
        #         this_contrib = Contribution(**entry)
        #         contributions.append(this_contrib)
            
        #     bin_title = pt_bin.replace("to", " < p_{T} < ").replace("RefPt", "")
        #     do_comparison_graph(contributions, bin_title=bin_title + " GeV",
        #                         xtitle="|#eta|", ytitle="Response",
        #                         xlimits=(0, 5.2), y_limit_protection=(0.5, 1.5),
        #                         other_elements=other_elements,
        #                         output_filename=os.path.join(plot_dir, "rsp_vs_eta_%s.pdf" % (pt_bin)))

            # Inverse response ie ~ correction
            # entries = []
            # for fdict in entry_dicts:
            #     entry = deepcopy(fdict)
            #     entry["graph"] = construct_inverse_graph(cu.grab_obj_from_file(args.input, "%s/%s_%s" % (mydir, fdict['flav'].replace("RefPt", "JetEta"), pt_bin)))
            #     entry["line_color"] = fdict['colour']
            #     entry["marker_color"] = fdict['colour']
            #     entries.append(entry)
            # bin_title = pt_bin.replace("to", " < p_{T} < ").replace("RefPt", "")
            # do_comparison_graph(entries, bin_title=bin_title + " GeV",
            #                     xtitle="|#eta|", ytitle="1/Response",
            #                     y_limit_protection=(0.8, 1.6),
            #                     other_elements=other_elements,
            #                     output_filename=os.path.join(plot_dir, "inv_rsp_vs_eta_%s.pdf" % (pt_bin)))


        # Do all flavs resolution vs pt for given eta bin
        # for eta_bin in common_eta_bins:
        #     entries = []
        #     for fdict in entry_dicts:
        #         entry = deepcopy(fdict)
        #         entry["graph"] = cu.grab_obj_from_file(args.input, "%s/%s_%s" % (mydir, fdict['flav'].replace("Rsp", "Res", 1), eta_bin))
        #         entry["line_color"] = fdict['colour']
        #         entry["marker_color"] = fdict['colour']
        #         entries.append(entry)
        #     bin_title = eta_bin.replace("to", " < |#eta| < ").replace("JetEta", "")
        #     do_comparison_graph(entries, bin_title=bin_title,
        #                         xtitle="p_{T}^{Gen} [GeV]", ytitle="Relative resolution", logx=True,
        #                         y_limit_protection=(0, 0.5), draw_fits=False,
        #                         ylimits=(0, 0.5),
        #                         xlimits=(10, 5000), other_elements=other_elements,
        #                         output_filename=os.path.join(plot_dir, "res_vs_pt_%s.pdf" % (eta_bin)))

        # Do all flavs resolution plots vs eta for given pt bin
        # for pt_bin in common_pt_bins:
        #     entries = []
        #     for fdict in entry_dicts:
        #         entry = deepcopy(fdict)
        #         entry["graph"] = cu.grab_obj_from_file(args.input, "%s/%s_%s" % (mydir, fdict['flav'].replace("Rsp", "Res", 1).replace("RefPt", "JetEta"), pt_bin))
        #         entry["line_color"] = fdict['colour']
        #         entry["marker_color"] = fdict['colour']
        #         entries.append(entry)
        #     bin_title = pt_bin.replace("to", " < p_{T} < ").replace("RefPt", "")
        #     do_comparison_graph(entries, bin_title=bin_title + " GeV",
        #                         xtitle="|#eta|", ytitle="Relative resolution", other_elements=other_elements,
        #                         y_limit_protection=(0, 0.5), draw_fits=True, xlimits=(0, 5.2),
        #                         output_filename=os.path.join(plot_dir, "res_vs_eta_%s.pdf" % (pt_bin)))

    return 0
Ejemplo n.º 8
0
def do_pileup_plot(input_dir, trigger_names, output_filename):
    # get histograms
    hists = [
        cu.grab_obj_from_file(
            os.path.join(input_dir, 'MyDataPileupHistogram_%s.root' % t),
            'pileup') for t in trigger_names
    ]
    h_up = cu.grab_obj_from_file(
        os.path.join(input_dir, 'MyDataPileupHistogram_PFJet500_72383.root'),
        'pileup')
    h_down = cu.grab_obj_from_file(
        os.path.join(input_dir, 'MyDataPileupHistogram_PFJet500_66017.root'),
        'pileup')
    h_down3 = cu.grab_obj_from_file(
        os.path.join(input_dir, 'MyDataPileupHistogram_PFJet500_59650.root'),
        'pileup')
    ratio_up = h_up.Clone()
    ratio_up.Divide(hists[-1])
    ratio_down = h_down.Clone()
    ratio_down.Divide(hists[-1])
    # hists.append(h_up)
    # hists.append(h_down)
    # trigger_names.append('72.383')
    # trigger_names.append('66.017')
    # Create contributions
    mark = cu.Marker()
    n_hists = len(hists)
    conts = [
        Contribution(h,
                     label=t,
                     line_width=2,
                     line_color=cu.get_colour_seq(ind, n_hists),
                     marker_color=cu.get_colour_seq(ind, n_hists),
                     marker_style=m,
                     normalise_hist=True)
        for ind, (h, t, m) in enumerate(
            zip(hists, trigger_names, mark.cycle(cycle_filling=True)))
    ]

    conts.insert(
        -1,
        Contribution(
            h_up,
            label='72.383 (+4.6%)',
            line_width=2,
            line_color=ROOT.kRed,
            marker_color=ROOT.kRed,
            normalise_hist=True,
        ))
    conts.insert(
        -1,
        Contribution(
            h_down,
            label='66.017 (-4.6%)',
            line_width=2,
            line_color=ROOT.kMagenta,
            marker_color=ROOT.kMagenta,
            normalise_hist=True,
        ))
    conts.insert(
        -1,
        Contribution(
            h_down3,
            label='59.650 (-13.8%)',
            line_width=2,
            line_style=1,
            line_color=ROOT.kMagenta + 3,
            marker_color=ROOT.kMagenta + 3,
            normalise_hist=True,
        ))
    print([h.Integral() for h in hists])
    plot = Plot(
        conts,
        what='hist',
        xtitle='Pileup',
        lumi=cu.get_lumi_str(do_dijet=True, do_zpj=True),
        subplot_type='ratio',
        subplot=conts[-1],
    )
    plot.subplot_maximum_ceil = 4.5
    plot.plot("NOSTACK HISTE", "NOSTACK HIST")
    plot.subplot_pad.cd()

    # To shade region between up and down hists, need to create graph with
    # the error bars as the up/down variations

    x_low = np.array([
        ratio_up.GetBinLowEdge(i) for i in range(1,
                                                 ratio_up.GetNbinsX() + 1)
    ],
                     dtype=float)
    x_high = np.array([
        ratio_up.GetBinLowEdge(i) for i in range(2,
                                                 ratio_up.GetNbinsX() + 2)
    ],
                      dtype=float)
    x_mid = 0.5 * (x_low + x_high)
    x_err = x_high - x_low

    y_high = np.array([
        max(ratio_up.GetBinContent(i), ratio_down.GetBinContent(i))
        for i in range(1,
                       ratio_up.GetNbinsX() + 1)
    ])
    y_high -= 1
    y_low = np.array([
        min(ratio_up.GetBinContent(i), ratio_down.GetBinContent(i))
        for i in range(1,
                       ratio_up.GetNbinsX() + 1)
    ])
    y_low = 1 - y_low
    y_mid = np.array([1. for i in range(1, ratio_up.GetNbinsX() + 1)])

    gr = ROOT.TGraphAsymmErrors(len(x_mid), x_mid, y_mid, x_err, x_err, y_low,
                                y_high)
    gr.SetFillColor(ROOT.kGray + 2)
    gr.SetFillStyle(3254)
    gr.Draw('2 SAME')
    plot.subplot_container.Draw("SAME NOSTACK HIST")
    plot.subplot_line.Draw()
    plot.save(output_filename)