Esempio n. 1
0
def main():

    ops = options()

    if not ops.plotter:
        fatal("Need --plotter for configuration.")
    plotter = yaml.load(open(ops.plotter))

    trees   = {}
    plots   = {}
    weights = {}
    colors  = {}
    labels  = {}
    stack   = {}
    overlay = {}
    is_data = {}

    # retrieve inputs
    for sample in plotter["samples"]:

        name = sample["name"]

        # input trees
        paths = sample["path"]
        if not glob.glob(paths):
            fatal("Found no files at %s" % (paths))

        trees[name] = ROOT.TChain(plotter["tree"])
        for path in glob.glob(paths):
            trees[name].Add(path)

        # misc
        weights[name] = " * ".join(sample["weights"]) if "weights" in sample else "1"
        colors[name]  = eval(sample["color"]) if sample["color"].startswith("ROOT") else sample["color"] 
        labels[name]  = sample["label"]
        stack[name]   = sample["stack"]
        overlay[name] = sample["overlay"]
        is_data[name] = sample["is_data"]

    # create output file
    output = ROOT.TFile.Open("%s/plots.%s.canv.root" % (plotter["directory"], timestamp), "recreate")

    # make money
    for plot in plotter["plots"]:

        hists = {}
        draw = {}
        if "bins" in plot:
            draw["bins"]      = array.array("d", plot["bins"])
        draw["title"]     = ";%s;%s" % (plot["xtitle"], plot["ytitle"])
        draw["variable"]  = plot["variable"]
        draw["selection"] = " && ".join(plotter["selection"])

        canv = ROOT.TCanvas(plot["name"], plot["name"], 800, 800)
        canv.Draw()
        canv.SetLogy(plot["logY"])

        stacks   = ROOT.THStack(plot["name"]+"stacks",   draw["title"])
        overlays = ROOT.THStack(plot["name"]+"overlays", draw["title"])
        do_stack = False
        do_overlay = False
        
        for sample in trees:

            draw["name"]   = plot["name"]+"__"+sample
            draw["weight"] = weights[sample]
            for option in ["weight"]:
                draw[option] = "(%s)" % (draw[option])


            if "bins" in plot:
                hists[sample] = ROOT.TH1F(draw["name"], draw["title"], len(draw["bins"])-1, draw["bins"])
            else:
                hists[sample] = ROOT.TH1F(draw["name"], draw["title"], plot["n_bins"], plot["bin_low"], plot["bin_high"])
            hists[sample].Sumw2()

            trees[sample].Draw("%(variable)s >> %(name)s" % draw, "(%(selection)s) * %(weight)s" % draw, "goff")
            output.cd()
            hists[sample].Write()
#            print "(%(selection)s) * %(weight)s" % draw
            
            #hists[sample].Scale(1/hists[sample].Integral(0, hists[sample].GetNbinsX()))

            if stack[sample]:
                do_stack = True
                hists[sample].SetFillColor(colors[sample])
                hists[sample].SetLineColor(ROOT.kBlack)
                hists[sample].SetLineWidth(2)
                stacks.Add(copy.copy(hists[sample]), ("ep" if is_data[sample] else "hist"))

            if overlay[sample]:
                do_overlay = True
                hists[sample].SetFillColor(0)
                hists[sample].SetLineColor(colors[sample])
                hists[sample].SetLineWidth(3)
                overlays.Add(copy.copy(hists[sample]), ("ep" if is_data[sample] else "hist"))

        # draw
        maximum = max([stacks.GetMaximum(), overlays.GetMaximum("nostack")])
        maximum = maximum*(100.0 if plot["logY"]     else 2.0)
        maximum = maximum*(1.2   if plotter["ratio"] else 1.0)

        if do_stack:
            stacks.SetMaximum(maximum)
            stacks.Draw()

        if do_overlay and do_stack:
            overlays.SetMaximum(maximum)
            overlays.Draw("nostack,same")
        elif do_overlay:
            overlays.SetMaximum(maximum)
            overlays.Draw("nostack")

        if plotter["data"]:
            pass

        if plotter["ratio"] and stacks.GetStack():

            # numerator definition is a placeholder.
            # only works if overlay[0]=data.
            ratio = helpers.ratio(name   = canv.GetName()+"_ratio",
                                  numer  = overlays.GetHists()[0],   # AHH KILL ME
                                  denom  = stacks.GetStack().Last(),
                                  min    = 0.45,
                                  max    = 1.55,
                                  ytitle = "Data / pred."
                                  )
            share = helpers.same_xaxis(name          = canv.GetName()+"_share",
                                       top_canvas    = canv,
                                       bottom_canvas = ratio,
                                       )
            canv.SetName(canv.GetName()+"_noratio")
            share.SetName(share.GetName().replace("_share", ""))
            canv = share

        elif plotter["ratio"] and not stacks.GetStack():
            warn("Want to make ratio plot but dont have stack. Skipping ratio.")

        # stack legend
        xleg, yleg = 0.6, 0.7
        legend = ROOT.TLegend(xleg, yleg, xleg+0.2, yleg+0.2)

        if do_stack:
            for hist in reversed(stacks.GetHists()):
                legend.AddEntry(hist, labels[hist.GetName().split("__")[1]], "f")

        if do_overlay:
            for hist in reversed(overlays.GetHists()):
                legend.AddEntry(hist, labels[hist.GetName().split("__")[1]], "l")
        
        legend.SetBorderSize(0)
        legend.SetFillColor(0)
        legend.SetMargin(0.3)
        legend.SetTextSize(0.04)
        legend.Draw()

        # watermarks
        xatlas, yatlas = 0.38, 0.87
        atlas = ROOT.TLatex(xatlas,      yatlas, "ATLAS Internal")
        hh4b  = ROOT.TLatex(xatlas, yatlas-0.06, "X #rightarrow HH #rightarrow 4b")
        lumi  = ROOT.TLatex(xatlas, yatlas-0.12, "#sqrt{s} = 13 TeV")
        watermarks = [atlas, hh4b, lumi]

        # KS, chi2
        if stacks.GetStack():
            if plotter.get("ks"):
                kolg, chi2, ndf = helpers.compare(overlays.GetHists()[0],
                                                  stacks.GetStack().Last(),
                                                  ) # AH KILL ME
                yks   = 0.975
                ychi2 = 0.975
                xks   = 0.27
                xchi2 = 0.55
                
                ks = ROOT.TLatex(xks,   yks,   "KS = %5.3f" % (kolg))
                ch = ROOT.TLatex(xchi2, ychi2, "#chi^{2} / ndf = %.1f / %i" % (chi2, ndf))
                watermarks += [ks, ch]

        # draw watermarks
        for wm in watermarks:
            wm.SetTextAlign(22)
            wm.SetTextSize(0.04)
            wm.SetTextFont(42)
            wm.SetNDC()
            wm.Draw()

        canv.SaveAs(plotter["directory"] + "/" + canv.GetName()+".pdf")

        output.Close()
Esempio n. 2
0
def main():

    # create output file
    output = ROOT.TFile.Open(CONF.outplotpath + "Trig_eff.root", "recreate")
    outputdir = CONF.outplotpath
    # setup canvas
    canv = ROOT.TCanvas("Trig", "Trig", 800, 800)
    # load input MC file
    #input_mc = ROOT.TFile.Open("../../test_ttbarnonhad/hist-hh4b-00-06-03q.root")
    eff_lst = ["j_Pt_t", "j_Eta", "j_Phi"]
    for plot_name in eff_lst:

        input_mc = ROOT.TFile.Open(CONF.inputpath + "trig_ttbarnonhad/hist-MiniNTuple.root ")
        allevt_mc = input_mc.Get("PassLep_%s" % plot_name)
        passevt_mc = input_mc.Get("PassTrig_%s" % plot_name)
        eff_mc = ROOT.TEfficiency(passevt_mc, allevt_mc)
        eff_mc.SetLineColor(ROOT.kRed)
        eff_mc.SetMarkerColor(ROOT.kRed)
        hst_eff_mc = eff_mc.GetCopyTotalHisto()
        hst_eff_mc.GetYaxis().SetTitle("HLT_j360_a10r Trigger Efficiency")
        hst_eff_mc.SetMaximum(1.5)
        hst_eff_mc.SetMinimum(0)
        hst_eff_mc.Draw("")
        eff_mc.Draw("same")
        ROOT.gPad.Update();
        eff_mc.GetPaintedGraph().SetMaximum(1.5)
        eff_mc.GetPaintedGraph().SetMinimum(0)
        # load input data file
        input_data = ROOT.TFile.Open(CONF.inputpath + "trig_data/hist-MiniNTuple.root")
        allevt_data = input_data.Get("PassLep_%s" % plot_name)
        passevt_data = input_data.Get("PassTrig_%s" % plot_name)
        eff_data = ROOT.TEfficiency(passevt_data, allevt_data)
        eff_data.SetMarkerStyle(20)
        eff_data.SetMarkerSize(1)
        eff_data.SetLineColor(ROOT.kBlack)
        eff_data.SetMarkerColor(ROOT.kBlack)
        eff_data.Draw("same")
        ROOT.gPad.Update();
        eff_data.GetPaintedGraph().SetMaximum(1.5)
        eff_data.GetPaintedGraph().SetMinimum(0)
        # hst_eff_data.SetMaximum(1.5)
        # hst_eff_data.SetMinimum(0)
        # hst_eff_data.Draw("same")

        # stack legend
        xleg, yleg = 0.6, 0.7
        legend = ROOT.TLegend(xleg, yleg, xleg+0.2, yleg+0.2)
        legend.AddEntry(eff_mc, "t#bar{t} MC", "apl")
        legend.AddEntry(eff_data, "data", "apl")
        legend.SetBorderSize(0)
        legend.SetMargin(0.3)
        legend.SetTextSize(0.04)
        legend.Draw()


        # watermarks
        xatlas, yatlas = 0.38, 0.87
        atlas = ROOT.TLatex(xatlas, yatlas, "ATLAS Internal")
        hh4b  = ROOT.TLatex(xatlas, yatlas-0.06, "Semi-leptonic t#bar{t}")
        lumi  = ROOT.TLatex(xatlas, yatlas-0.12, "#sqrt{s} = 13 TeV")
        watermarks = [atlas, hh4b, lumi]

        # draw watermarks
        for wm in watermarks:
            wm.SetTextAlign(22)
            wm.SetTextSize(0.04)
            wm.SetTextFont(42)
            wm.SetNDC()
            wm.Draw()

        # draw data/MC ratio
        ratioMin = 0
        ratioMax = 2

        # for i in range(eff_data.GetPaintedGraph().GetGlobalBin()):
        #     print eff_data.GetName(), " ", eff_data.GetEfficiency(i)
        #passevt_data.Divide(allevt_data)    
        #passevt_mc.Divide(allevt_mc)
        for i in range(0, passevt_data.GetNbinsX() + 2):
            passevt_data.SetBinContent(i, eff_data.GetEfficiency(i))
            passevt_mc.SetBinContent(i, eff_mc.GetEfficiency(i))
            passevt_data.SetBinError(i, max(eff_data.GetEfficiencyErrorLow(i), eff_data.GetEfficiencyErrorUp(i)))
            passevt_mc.SetBinError(i, max(eff_mc.GetEfficiencyErrorLow(i), eff_mc.GetEfficiencyErrorUp(i)))
    
        passevt_data.SetMarkerStyle(20)
        passevt_data.SetMarkerSize(1)
        passevt_data.SetLineColor(ROOT.kBlack)
        passevt_data.SetMarkerColor(ROOT.kBlack)


        ratio = helpers.ratio(name   = canv.GetName()+"_ratio",
                                 numer  = passevt_data,   # AHH KILL ME
                                 denom  = passevt_mc,
                                 min    = ratioMin,
                                 max    = ratioMax,
                                 ytitle = "Data / pred."
                                )
        share,top,bot = helpers.same_xaxis(name          = canv.GetName()+"_share",
                                   top_canvas    = canv,
                                   bottom_canvas = ratio,
                                   )
        canv.SetName(canv.GetName()+"_noratio")
        share.SetName(share.GetName().replace("_share", ""))
        canv = share

        # finish up
        output.cd()
        canv.SaveAs(outputdir + eff_mc.GetName() + ".pdf")


    canv.Clear()
    #check basic distributions in mc makes sense
    plot_lst = ["j_Pt", "j_Eta", "j_Phi", "j_M", "mu_Pt", "mu_Eta", "mu_Phi",\
     "met", "met_Phi", "bj_Pt", "bj_Eta", "bj_Phi", "bj_M", "bj_MV2", \
     "w_MT", "l_dPhi_mu_met", "l_dEta_mu_bj", "l_dPhi_mu_bj", "l_dR_mu_bj", \
     "l_Pt", "l_Eta", "l_Phi", "l_M", "l_MT", \
     "leptop_dR", "leptop_deta", "leptop_dphi", "leptop_m", "leptop_mT"]

    plot_2d_lst = ["mu_Pt_met", "l_Pt_met", "l_Pt_mu_pt", "l_Pt_j_pt", "l_Pt_bj_pt", "l_Pt_dR_mu_bj"]

    ## validate in MC
    # for plot_name in plot_lst:
    #     plt_before = input_mc.Get("AllLep_%s" % plot_name)
    #     plt_after  = input_mc.Get("PassLep_%s" % plot_name)
    #     plt_before.SetLineColor(ROOT.kRed)
    #     plt_after.SetLineColor(ROOT.kBlue)
    #     maxvalue = max(plt_before.GetMaximum(), plt_after.GetMaximum())
    #     plt_before.SetMaximum(1.5 * maxvalue)
    #     plt_after.SetMaximum(1.5 * maxvalue)
    #     plt_before.Draw()
    #     plt_after.Draw("same")

    #     xatlas, yatlas = 0.4, 0.9
    #     atlas = ROOT.TLatex(xatlas, yatlas, "ATLAS Internal")
    #     hh4b  = ROOT.TLatex(xatlas, yatlas-0.06, "Semi-leptonic t#bar{t} MC")
    #     lumi  = ROOT.TLatex(xatlas, yatlas-0.12, "#sqrt{s} = 13 TeV")
    #     watermarks = [atlas, hh4b, lumi]

    #     xleg, yleg = 0.6, 0.7
    #     legend = ROOT.TLegend(xleg, yleg, xleg+0.2, yleg+0.2)
    #     legend.AddEntry(plt_before, "before selection", "apl")
    #     legend.AddEntry(plt_after,  "after selection", "apl")
    #     legend.SetBorderSize(0)
    #     legend.SetMargin(0.3)
    #     legend.SetTextSize(0.04)
    #     legend.Draw()
    #      # draw watermarks
    #     for wm in watermarks:
    #         wm.SetTextAlign(22)
    #         wm.SetTextSize(0.04)
    #         wm.SetTextFont(42)
    #         wm.SetNDC()
    #         wm.Draw()

    #     output.cd()
    #     canv.SaveAs(outputdir + plt_before.GetName()+".pdf")
    #     canv.Clear()

    # ## validate in data
    for plot_name in plot_lst:
        plt_before = input_mc.Get("PassLep_%s" % plot_name)
        plt_after  = input_data.Get("PassLep_%s" % plot_name)
        plt_before.SetLineColor(ROOT.kRed)
        plt_after.SetLineColor(ROOT.kBlack)
        plt_after.SetMarkerStyle(20)
        plt_after.SetMarkerSize(1)
        plt_after.GetYaxis().SetTitleOffset(1.5)
        plt_before.GetYaxis().SetTitleOffset(1.5)
        maxvalue = max(plt_before.GetMaximum(), plt_after.GetMaximum())
        plt_before.SetMaximum(1.5 * maxvalue)
        plt_after.SetMaximum(1.5 * maxvalue)
        plt_before.SetMinimum(0)
        plt_after.SetMinimum(0)
        plt_after.Draw("ep")
        plt_before.Draw("hist same")

        xatlas, yatlas = 0.38, 0.87
        atlas = ROOT.TLatex(xatlas, yatlas, "ATLAS Internal")
        hh4b  = ROOT.TLatex(xatlas, yatlas-0.06, "Semi-leptonic t#bar{t} MC")
        lumi  = ROOT.TLatex(xatlas, yatlas-0.12, "#sqrt{s} = 13 TeV")
        watermarks = [atlas, hh4b, lumi]

        xleg, yleg = 0.6, 0.7
        legend = ROOT.TLegend(xleg, yleg, xleg+0.2, yleg+0.2)
        legend.AddEntry(plt_before, "MC", "apl")
        legend.AddEntry(plt_after,  "data", "apl")
        legend.SetBorderSize(0)
        legend.SetMargin(0.3)
        legend.SetTextSize(0.04)
        legend.Draw()
         # draw watermarks
        for wm in watermarks:
            wm.SetTextAlign(22)
            wm.SetTextSize(0.04)
            wm.SetTextFont(42)
            wm.SetNDC()
            wm.Draw()


        # draw data/MC ratio
        ratioMin = 0
        ratioMax = 2 
        ratio = helpers.ratio(name   = canv.GetName()+"_ratio",
                              numer  = plt_after,   # AHH KILL ME
                              denom  = plt_before,
                              min    = ratioMin,
                              max    = ratioMax,
                              ytitle = "Data / pred."
                              )
        share,top,bot = helpers.same_xaxis(name          = canv.GetName()+"_share",
                                   top_canvas    = canv,
                                   bottom_canvas = ratio,
                                   )
        canv.SetName(canv.GetName()+"_noratio")
        share.SetName(share.GetName().replace("_share", ""))
        canv = share

        output.cd()
        canv.SaveAs(outputdir + plt_before.GetName()+".pdf")
        plt_before.Write()
        plt_after.Write()
        canv.Clear()


    # for plot_name in plot_2d_lst:
    #     plt_after  = input_mc.Get("PassLep_%s" % plot_name)
    #     plt_after.Draw("colz")

    #     xatlas, yatlas = 0.4, 0.9
    #     atlas = ROOT.TLatex(xatlas, yatlas, "ATLAS Internal")
    #     hh4b  = ROOT.TLatex(xatlas, yatlas-0.06, "Semi-leptonic t#bar{t} MC")
    #     lumi  = ROOT.TLatex(xatlas, yatlas-0.12, "#sqrt{s} = 13 TeV")
    #     watermarks = [atlas, hh4b, lumi]

    #     # draw watermarks
    #     for wm in watermarks:
    #         wm.SetTextAlign(22)
    #         wm.SetTextSize(0.04)
    #         wm.SetTextFont(42)
    #         wm.SetNDC()
    #         wm.Draw()

    #     output.cd()
    #     canv.SaveAs(outputdir + plt_after.GetName()+".pdf")
    #     canv.Clear()

    eff_mc.Write()
    eff_data.Write()
    output.Close()
def main():

    qcd = ROOT.TFile.Open("hist/qcd/qcd.root")

    for fullpath, dirnames, objnames, _ in walk.walk(qcd):

        filepath, dirpath = fullpath.split(":/")

        if not "2tag" in fullpath:
            continue
        if dirpath not in ["sb_2tag77", "cr_2tag77", "sr_2tag77"]:
            continue

        for name in objnames:

            histpath = os.path.join(dirpath, name)

            qcd_nomin = qcd.Get(histpath)
            qcd_tight = qcd.Get(histpath.replace("_2tag77", "_2tag77_4tag97"))
            qcd_loose = qcd.Get(histpath.replace("_2tag77", "_2tag77_N4tag97"))

            if isinstance(qcd_nomin, ROOT.TH2):
                continue

            qcd_nomin.Scale(mu_qcd_4b_nomin)
            qcd_tight.Scale(mu_qcd_4b_tight)
            qcd_loose.Scale(mu_qcd_4b_loose)

            if "m_JJ" in histpath:
                bins = qcd_nomin.GetNbinsX()+1
                error = ROOT.Double(0)
                print
                print " nomin: %9.3f (%.3f)" % (qcd_nomin.IntegralAndError(0, bins, error), error)
                print " tight: %9.3f (%.3f)" % (qcd_tight.IntegralAndError(0, bins, error), error)
                print " loose: %9.3f (%.3f)" % (qcd_loose.IntegralAndError(0, bins, error), error)
                print

            for hist in [qcd_nomin, qcd_tight, qcd_loose]:
                hist.Rebin(rebin(name))
                helpers.show_overflow(hist)
                hist.SetMinimum(0.0)
                hist.SetMaximum(ymax(histpath))
                hist.GetXaxis().SetTitle(xtitle(name))
                hist.GetYaxis().SetTitle("Events")
                hist.SetLineWidth(3)
                hist.SetLabelSize(0.05, "xyz")
                hist.SetTitleSize(0.05, "xyz")
                hist.GetXaxis().SetNdivisions(505)

            for bin in xrange(0, qcd_nomin.GetNbinsX()+1):
                qcd_nomin.SetBinError(bin, 0.00001)

            qcd_nomin.SetLineColor(ROOT.kBlack)
            qcd_tight.SetLineColor(ROOT.kBlue)
            qcd_loose.SetLineColor(ROOT.kRed)

            canvas = ROOT.TCanvas(name, name, 800, 800)
            canvas.Draw()

            for hist in [qcd_nomin, qcd_tight, qcd_loose]:
                hist.Draw("histesame" if hist!=qcd_nomin else "histsame")

            if True:

                ratio = helpers.ratio(name   = canvas.GetName()+"_ratio",
                                      numers = [qcd_tight, qcd_loose],
                                      denom  = qcd_nomin,
                                      min    = 0.1,
                                      max    = 1.9,
                                      ytitle = "varied / nom."
                                      )
                share = helpers.same_xaxis(name          = canvas.GetName()+"_share",
                                           top_canvas    = canvas,
                                           bottom_canvas = ratio,
                                           )
                canvas.SetName(canvas.GetName()+"_noratio")
                share.SetName(share.GetName().replace("_share", ""))
                canvas = share

            ks = qcd_tight.KolmogorovTest(qcd_loose)

            xatlas, yatlas = 0.19, 0.96
            xleg, yleg     = 0.68, 0.84
            atlas         = ROOT.TLatex(xatlas,     yatlas,    "ATLAS Internal")
            kolmo         = ROOT.TLatex(xatlas+0.4, yatlas,    "KS(tight, loose): %.3f" % ks)
            legend_nomin  = ROOT.TLatex(xleg,       yleg,      "nominal")
            legend_tight  = ROOT.TLatex(xleg,       yleg-0.05, "tight 2-tag")
            legend_loose  = ROOT.TLatex(xleg,       yleg-0.10, "loose 2-tag")
            legend_nomin.SetTextColor(ROOT.kBlack)
            legend_tight.SetTextColor(ROOT.kBlue)
            legend_loose.SetTextColor(ROOT.kRed)
            for logo in [atlas, kolmo, legend_nomin, legend_tight, legend_loose]:
                logo.SetTextSize(0.035 if logo in [atlas, kolmo] else 0.04)
                logo.SetTextFont(42)
                logo.SetNDC()
                logo.Draw()

            outdir = os.path.join("plot", dirpath)
            if not os.path.isdir(outdir):
                os.makedirs(outdir)
            canvas.SaveAs(os.path.join(outdir, canvas.GetName()+".pdf"))

            canvas.Close()
Esempio n. 4
0
                    ratioMin = ratioMin-.05
                    ratioMax = ratioMax+0.5

            else:
                ratioMin = 0
                ratioMax = 2 

            ratio = helpers.ratio(name   = canv.GetName()+"_ratio",
                                  numer  = overlays.GetHists()[0],   # AHH KILL ME
                                  denom  = stacks.GetStack().Last(),
                                  min    = ratioMin,
                                  max    = ratioMax,
                                  ytitle = "Data / pred."
                                  )
            share,top,bot = helpers.same_xaxis(name          = canv.GetName()+"_share",
                                       top_canvas    = canv,
                                       bottom_canvas = ratio,
                                       )
            canv.SetName(canv.GetName()+"_noratio")
            share.SetName(share.GetName().replace("_share", ""))
            canv = share

        elif plotter["ratio"] and not stacks.GetStack():
            warn("Want to make ratio plot but dont have stack. Skipping ratio.")

        # stack legend
        xleg, yleg = 0.6, 0.7
        legend = ROOT.TLegend(xleg, yleg, xleg+0.2, yleg+0.2)

        if do_stack:
            for hist in reversed(stacks.GetHists()):
                legend.AddEntry(hist, labels[hist.GetName().split("__")[1]], "f")
Esempio n. 5
0
def main():

    ops = options()

    if not ops.plotter:
        fatal("Need --plotter for configuration.")
    elif ".yml" in ops.plotter:
        print "use yml"
        plotter = yaml.load(open(ops.plotter))
    else:
        plotter = plotConfig.getConfig(ops.plotter)

    if not "directory" in plotter:
        plotter["directory"] = "."
        warn("No directory given. Writing output to cwd.")

    if not "tree" in plotter:
        print "No input Tree, assuming hists"
        files = {}
        folders = {}
        useTree = False
    else:
        useTree = True

    trees = {}
    plots = {}
    weights = {}
    colors = {}
    labels = {}
    stack = {}
    overlay = {}
    is_data = {}
    samples = []
    prename = plotter["prename"]
    inputpath = ""
    if ops.inputdir:
        inputpath = ops.inputdir + "/"
        prename = ops.inputdir + "_" + prename

    # retrieve inputs
    print time.strftime("%Y-%m-%d-%Hh%Mm%Ss"), "Retrieving input files"
    for sample in plotter["samples"]:
        print sample
        name = sample["name"]
        samples.append(name)

        # input files
        paths = sample["path"]
        #set the full path directory here
        paths = "../Output/" + inputpath + paths

        if not glob.glob(paths):
            fatal("Found no files at %s" % (paths))

        if useTree:
            trees[name] = ROOT.TChain(plotter["tree"])
            for path in glob.glob(paths):
                if len(path) > 100:
                    print "adding file:", path[0:50] + "..." + path[len(path) -
                                                                    50:]
                else:
                    print "adding file:", path
                trees[name].Add(path)

        else:
            path = paths
            files[name] = ROOT.TFile(path, "READ")
            folders[name] = sample["folder"]

        # misc
        if useTree:
            weights[name] = " * ".join(
                sample["weights"]) if "weights" in sample else "1"
        else:
            weights[name] = sample["weights"]
        print sample["weights"]
        colors[name] = eval(sample["color"]) if sample["color"].startswith(
            "ROOT") else sample["color"]
        labels[name] = sample["label"]
        stack[name] = sample["stack"]
        overlay[name] = sample["overlay"]
        is_data[name] = sample["is_data"]

    # create output file
    if not os.path.isdir(plotter["directory"]):
        os.makedirs(plotter["directory"])
    output = ROOT.TFile.Open(
        "%s/%splots.root" % (plotter["directory"], prename), "recreate")

    # make money
    for plot in plotter["plots"]:

        hists = {}
        draw = {}
        if "bins" in plot:
            draw["bins"] = array.array("d", [float(x) for x in plot["bins"]])
        draw["title"] = ";%s;%s" % (plot["xtitle"], plot["ytitle"])
        draw["variable"] = plot["variable"]
        if "selection" in plotter:
            draw["selection"] = " && ".join(plotter["selection"])

        canv = ROOT.TCanvas(plot["name"], plot["name"], 800, 800)
        canv.Draw()
        canv.SetLogy(plot["logY"])

        stacks = ROOT.THStack(plot["name"] + "stacks", draw["title"])
        overlays = ROOT.THStack(plot["name"] + "overlays", draw["title"])
        do_stack = False
        do_overlay = False

        for sample in reversed(sorted(samples)):
            draw["name"] = plot["name"] + "__" + sample
            draw["weight"] = weights[sample]
            if useTree:
                for option in ["weight"]:
                    draw[option] = "(%s)" % (draw[option])

            if useTree:
                if "bins" in plot:
                    hists[sample] = ROOT.TH1F(draw["name"], draw["title"],
                                              len(draw["bins"]) - 1,
                                              draw["bins"])
                else:
                    hists[sample] = ROOT.TH1F(draw["name"], draw["title"],
                                              plot["n_bins"], plot["bin_low"],
                                              plot["bin_high"])
                hists[sample].Sumw2()

            else:
                if folders[sample] != "":
                    hists[sample] = files[sample].Get(folders[sample] + "/" +
                                                      draw["variable"])
                else:
                    hists[sample] = files[sample].Get(draw["variable"])

                hists[sample].SetName(draw["name"])
                hists[sample].SetTitle(draw["title"])
                hists[sample].Rebin(plot["rebin"])
                if "bin_high" in plot.keys():
                    bin_low = float(plot["bin_low"])
                    bin_high = float(plot["bin_high"])
                    hists[sample].GetXaxis().SetRangeUser(bin_low, bin_high)

            if is_data[sample]:
                hists[sample].SetMarkerStyle(20)
                hists[sample].SetMarkerSize(1)

            if useTree:
                trees[sample].Draw("%(variable)s >> %(name)s" % draw,
                                   "(%(selection)s) * %(weight)s" % draw,
                                   "goff")
            else:
                hists[sample].Sumw2()
                hists[sample].Scale(float(draw["weight"]))
                hists[sample].Draw()

            # hists[sample].Scale(1/hists[sample].Integral(0, hists[sample].GetNbinsX()))

            if stack[sample]:
                do_stack = True
                hists[sample].SetFillColor(colors[sample])
                hists[sample].SetLineColor(ROOT.kBlack)
                hists[sample].SetLineWidth(2)
                stacks.Add(copy.copy(hists[sample]),
                           ("ep" if is_data[sample] else "hist"))
                if not useTree and "bin_high" in plot.keys(
                ):  #have to set stack xaxis range for zooming because ROOT SUCKS
                    stacks.Draw()
                    stacks.GetXaxis().SetRangeUser(bin_low, bin_high)

            if overlay[sample]:
                do_overlay = True
                hists[sample].SetFillColor(0)
                hists[sample].SetLineColor(colors[sample])
                hists[sample].SetMarkerColor(colors[sample])
                hists[sample].SetLineWidth(3)
                h = copy.copy(hists[sample])
                if plotter["norm"]:
                    print h.GetName()
                    scale = h.Integral(0, h.GetXaxis().GetNbins() + 1)
                    if scale != 0:
                        h.Scale(1.0 / scale)
                overlays.Add(h, ("ep" if is_data[sample] else "hist"))

            print sample
            print "Integral:", hists[sample].Integral(
                0, hists[sample].GetNbinsX() + 1)
            print " Entries:", hists[sample].GetEntries()

        # draw
        maximum = max([stacks.GetMaximum(), overlays.GetMaximum("nostack")])
        maximum = maximum * (100.0 if plot["logY"] else 2.0)
        maximum = maximum * (1.2 if plotter["ratio"] else 1.0)
        minimum = max([stacks.GetMinimum(), overlays.GetMinimum("nostack")])
        minimum = minimum / (2 if plot["logY"] else 1.2)

        if do_stack:
            stacks.SetMaximum(maximum)
            stacks.SetMinimum(minimum)
            stacks.Draw()
            h1stackerror = copy.copy(stacks.GetStack().Last())
            h1stackerror.SetName("stat. error")
            h1stackerror.SetFillColor(ROOT.kGray + 3)
            h1stackerror.SetFillStyle(3005)
            h1stackerror.SetMarkerStyle(0)
            h1stackerror.Draw("SAME,E2")

        if do_overlay and do_stack:
            overlays.SetMaximum(maximum)
            overlays.SetMinimum(minimum)
            overlays.Draw("nostack,same")
        elif do_overlay:
            overlays.SetMaximum(maximum)
            overlays.SetMinimum(minimum)
            overlays.Draw("nostack")

        if plotter["data"]:
            pass

        if plotter["ratio"] and stacks.GetStack():

            # numerator definition is a placeholder.
            # only works if overlay[0]=data.
            if plotter["autoRatio"]:
                top = overlays.GetHists()[0].Clone("hnew")
                bottom = stacks.GetStack().Last().Clone("hnew")
                top.Divide(bottom)
                ratioBins = []
                for bin in range(top.GetSize()):
                    if top.GetBinContent(bin) != 0.0 and top.GetBinContent(
                            bin) * top.GetBinError(bin) < 5.0:
                        ratioBins.append(top.GetBinContent(bin))
                if not ratioBins: ratioBins = [0]

                ratioMin = float(int(min(ratioBins) * 100)) / 100
                ratioMax = float(int(max(ratioBins) * 100)) / 100
                if ratioMax - ratioMin < 0.1:
                    ratioMin = ratioMin - .05
                    ratioMax = ratioMax + 0.5

            else:
                ratioMin = 0
                ratioMax = 2

            ratio = helpers.ratio(
                name=canv.GetName() + "_ratio",
                numer=overlays.GetHists()[0],  # AHH KILL ME
                denom=stacks.GetStack().Last(),
                min=ratioMin,
                max=ratioMax,
                ytitle="Data / pred.")
            share, top, bot = helpers.same_xaxis(
                name=canv.GetName() + "_share",
                top_canvas=canv,
                bottom_canvas=ratio,
            )
            canv.SetName(canv.GetName() + "_noratio")
            share.SetName(share.GetName().replace("_share", ""))
            canv = share

        elif plotter["ratio"] and not stacks.GetStack():
            warn(
                "Want to make ratio plot but dont have stack. Skipping ratio.")

        # stack legend
        xleg, yleg = 0.6, 0.7
        legend = ROOT.TLegend(xleg, yleg, xleg + 0.2, yleg + 0.2)

        if do_stack:
            for hist in reversed(stacks.GetHists()):
                legend.AddEntry(hist, labels[hist.GetName().split("__")[1]],
                                "f")

        if do_overlay:
            for hist in reversed(overlays.GetHists()):
                legend.AddEntry(hist, labels[hist.GetName().split("__")[1]],
                                "l")

        legend.SetBorderSize(0)
        legend.SetFillColor(0)
        legend.SetMargin(0.3)
        legend.SetTextSize(0.04)
        legend.Draw()

        # watermarks
        xatlas, yatlas = 0.38, 0.87
        atlas = ROOT.TLatex(xatlas, yatlas, "ATLAS Internal")
        hh4b = ROOT.TLatex(xatlas, yatlas - 0.06,
                           "X #rightarrow HH #rightarrow 4b")
        lumi = ROOT.TLatex(xatlas, yatlas - 0.12, "13 TeV")
        watermarks = [atlas, hh4b, lumi]

        # KS, chi2
        if stacks.GetStack():
            if plotter.get("ks"):
                kolg, chi2, ndf = helpers.compare(
                    overlays.GetHists()[0],
                    stacks.GetStack().Last(),
                )  # AH KILL ME
                yks = 0.975
                ychi2 = 0.975
                xks = 0.27
                xchi2 = 0.55

                ks = ROOT.TLatex(xks, yks, "KS = %5.3f" % (kolg))
                ch = ROOT.TLatex(xchi2, ychi2,
                                 "#chi^{2} / ndf = %.1f / %i" % (chi2, ndf))
                watermarks += [ks, ch]

        # draw watermarks
        for wm in watermarks:
            wm.SetTextAlign(22)
            wm.SetTextSize(0.04)
            wm.SetTextFont(42)
            wm.SetNDC()
            wm.Draw()

        canv.SaveAs(
            os.path.join(plotter["directory"],
                         prename + canv.GetName() + ".pdf"))

        output.cd()
        canv.Write()

    output.Close()
Esempio n. 6
0
def main():

    ops = options()

    if not ops.plotter:
        fatal("Need --plotter for configuration.")
    plotter = yaml.load(open(ops.plotter))

    if not "directory" in plotter:
        plotter["directory"] = "."
        warn("No directory given. Writing output to cwd.")

    trees = {}
    plots = {}
    weights = {}
    colors = {}
    labels = {}
    stack = {}
    overlay = {}
    is_data = {}

    # retrieve inputs
    for sample in plotter["samples"]:

        name = sample["name"]

        # input trees
        paths = sample["path"]
        if not glob.glob(paths):
            fatal("Found no files at %s" % (paths))

        trees[name] = ROOT.TChain(plotter["tree"])
        for path in glob.glob(paths):
            trees[name].Add(path)

        # misc
        weights[name] = " * ".join(
            sample["weights"]) if "weights" in sample else "1"
        colors[name] = eval(sample["color"]) if sample["color"].startswith(
            "ROOT") else sample["color"]
        labels[name] = sample["label"]
        stack[name] = sample["stack"]
        overlay[name] = sample["overlay"]
        is_data[name] = sample["is_data"]

    # create output file
    if not os.path.isdir(plotter["directory"]):
        os.makedirs(plotter["directory"])
    output = ROOT.TFile.Open(
        "%s/plots.%s.canv.root" % (plotter["directory"], timestamp),
        "recreate")

    # make money
    for plot in plotter["plots"]:

        hists = {}
        draw = {}
        if "bins" in plot:
            draw["bins"] = array.array("d", plot["bins"])
        draw["title"] = ";%s;%s" % (plot["xtitle"], plot["ytitle"])
        draw["variable"] = plot["variable"]
        draw["selection"] = " && ".join(plotter["selection"])

        canv = ROOT.TCanvas(plot["name"], plot["name"], 800, 800)
        canv.Draw()
        canv.SetLogy(plot["logY"])

        stacks = ROOT.THStack(plot["name"] + "stacks", draw["title"])
        overlays = ROOT.THStack(plot["name"] + "overlays", draw["title"])
        do_stack = False
        do_overlay = False

        for sample in reversed(sorted(trees.keys())):

            draw["name"] = plot["name"] + "__" + sample
            draw["weight"] = weights[sample]
            for option in ["weight"]:
                draw[option] = "(%s)" % (draw[option])

            if "bins" in plot:
                hists[sample] = ROOT.TH1F(draw["name"], draw["title"],
                                          len(draw["bins"]) - 1, draw["bins"])
            else:
                hists[sample] = ROOT.TH1F(draw["name"], draw["title"],
                                          plot["n_bins"], plot["bin_low"],
                                          plot["bin_high"])
            hists[sample].Sumw2()

            if is_data[sample]:
                hists[sample].SetMarkerStyle(20)
                hists[sample].SetMarkerSize(1)

            trees[sample].Draw("%(variable)s >> %(name)s" % draw,
                               "(%(selection)s) * %(weight)s" % draw, "goff")

            # hists[sample].Scale(1/hists[sample].Integral(0, hists[sample].GetNbinsX()))

            if stack[sample]:
                do_stack = True
                hists[sample].SetFillColor(colors[sample])
                hists[sample].SetLineColor(ROOT.kBlack)
                hists[sample].SetLineWidth(2)
                stacks.Add(copy.copy(hists[sample]),
                           ("ep" if is_data[sample] else "hist"))

            if overlay[sample]:
                do_overlay = True
                hists[sample].SetFillColor(0)
                hists[sample].SetLineColor(colors[sample])
                hists[sample].SetLineWidth(3)
                overlays.Add(copy.copy(hists[sample]),
                             ("ep" if is_data[sample] else "hist"))

            print hists[sample].Integral(0, hists[sample].GetNbinsX() + 1)
            print hists[sample].GetEntries()

        # draw
        maximum = max([stacks.GetMaximum(), overlays.GetMaximum("nostack")])
        maximum = maximum * (100.0 if plot["logY"] else 2.0)
        maximum = maximum * (1.2 if plotter["ratio"] else 1.0)

        if do_stack:
            stacks.SetMaximum(maximum)
            stacks.Draw()
            h1stackerror = copy.copy(stacks.GetStack().Last())
            h1stackerror.SetName("stat. error")
            h1stackerror.SetFillColor(ROOT.kGray + 3)
            h1stackerror.SetFillStyle(3005)
            h1stackerror.SetMarkerStyle(0)
            h1stackerror.Draw("SAME,E2")

        if do_overlay and do_stack:
            overlays.SetMaximum(maximum)
            overlays.Draw("nostack,same")
        elif do_overlay:
            overlays.SetMaximum(maximum)
            overlays.Draw("nostack")

        if plotter["data"]:
            pass

        if plotter["ratio"] and stacks.GetStack():

            # numerator definition is a placeholder.
            # only works if overlay[0]=data.
            ratio = helpers.ratio(
                name=canv.GetName() + "_ratio",
                numer=overlays.GetHists()[0],  # AHH KILL ME
                denom=stacks.GetStack().Last(),
                min=0.45,
                max=2.0,
                ytitle="Data / pred.")
            share = helpers.same_xaxis(
                name=canv.GetName() + "_share",
                top_canvas=canv,
                bottom_canvas=ratio,
            )
            canv.SetName(canv.GetName() + "_noratio")
            share.SetName(share.GetName().replace("_share", ""))
            canv = share

        elif plotter["ratio"] and not stacks.GetStack():
            warn(
                "Want to make ratio plot but dont have stack. Skipping ratio.")

        # stack legend
        xleg, yleg = 0.6, 0.7
        legend = ROOT.TLegend(xleg, yleg, xleg + 0.2, yleg + 0.2)

        if do_stack:
            for hist in reversed(stacks.GetHists()):
                legend.AddEntry(hist, labels[hist.GetName().split("__")[1]],
                                "f")

        if do_overlay:
            for hist in reversed(overlays.GetHists()):
                legend.AddEntry(hist, labels[hist.GetName().split("__")[1]],
                                "l")

        legend.SetBorderSize(0)
        legend.SetFillColor(0)
        legend.SetMargin(0.3)
        legend.SetTextSize(0.04)
        legend.Draw()

        # watermarks
        xatlas, yatlas = 0.38, 0.87
        atlas = ROOT.TLatex(xatlas, yatlas, "ATLAS Internal")
        hh4b = ROOT.TLatex(xatlas, yatlas - 0.06,
                           "X #rightarrow HH #rightarrow 4b")
        lumi = ROOT.TLatex(xatlas, yatlas - 0.12, "13 TeV, 78.7 pb^{-1}")
        watermarks = [atlas, hh4b, lumi]

        # KS, chi2
        if stacks.GetStack():
            if plotter.get("ks"):
                kolg, chi2, ndf = helpers.compare(
                    overlays.GetHists()[0],
                    stacks.GetStack().Last(),
                )  # AH KILL ME
                yks = 0.975
                ychi2 = 0.975
                xks = 0.27
                xchi2 = 0.55

                ks = ROOT.TLatex(xks, yks, "KS = %5.3f" % (kolg))
                ch = ROOT.TLatex(xchi2, ychi2,
                                 "#chi^{2} / ndf = %.1f / %i" % (chi2, ndf))
                watermarks += [ks, ch]

        # draw watermarks
        for wm in watermarks:
            wm.SetTextAlign(22)
            wm.SetTextSize(0.04)
            wm.SetTextFont(42)
            wm.SetNDC()
            wm.Draw()

        canv.SaveAs(os.path.join(plotter["directory"],
                                 canv.GetName() + ".pdf"))

        output.cd()
        canv.Write()

    output.Close()