示例#1
0
parser = argparse.ArgumentParser()
# parser.add_argument('--input', '-i', help='Output of PostFitShapes or PostFitShapesFromWorkspace, specified as FILE:BIN')
parser.add_argument('--output', '-o', default=None, help='Output name')
parser.add_argument('--typet', '-t', default='cat', help='category or channel')
# parser.add_argument('--channel', '-c', default='mt', choices=['mt', 'et', 'em', 'tt', 'mm'], help='Channel')
# parser.add_argument('--x-title', default='m_{ll}', help='x-axis variable, without GeV')
# parser.add_argument('--logy', action='store_true')
# parser.add_argument('--y-min', type=float, default=1)

args = parser.parse_args()

Type = args.typet

# Canvas and pads
canv = ROOT.TCanvas(args.output, args.output)
pads = plot.OnePad()
pads[0].SetTicks(1, -1)

axis = ROOT.TH2F('axis', '', 1, -.5, 3, 10, 0, 10)

plot.Set(axis.GetYaxis(), LabelSize=0)
plot.Set(axis.GetXaxis(), Title='Best fit #mu = #sigma/#sigma_{SM}')
axis.Draw()

y_pos = 8.5
x_text = 1.85

latex = ROOT.TLatex()
latexNum = ROOT.TLatex()
plot.Set(latex, TextAlign=12, TextSize=0.035)
plot.Set(latexNum, TextAlign=12, TextSize=0.025)
def plot_hadronic(input_file, triggers, working_points, file_types, era,
                  per_dm, output, draw_options, title, y_range, ratio_y_range,
                  binned_in, x_title, ratio_to, plot_dir, label_pos):

    ROOT.PyConfig.IgnoreCommandLineOptions = True
    ROOT.gROOT.SetBatch(ROOT.kTRUE)
    ROOT.TH1.AddDirectory(0)
    plot.ModTDRStyle()
    plot_dir = "plots"
    hists = {}
    graphs = {}
    if plot_dir != '':
        os.system('mkdir -p %s' % plot_dir)
    if per_dm:
        decayModes = ["inclusive", 0, 1, 10]
    else:
        decayModes = ["inclusive"]
    file = ROOT.TFile(input_file)
    for trg in triggers:
        hists[trg] = {}
        graphs[trg] = {}
        for decayMode in decayModes:
            dm_label = "" if decayMode == "inclusive" else "dm{}_".format(
                decayMode)
            hists[trg][decayMode] = {}
            graphs[trg][decayMode] = {}
            for wp in working_points:
                hists[trg][decayMode][wp] = []
                graphs[trg][decayMode][wp] = []

                for ft in file_types:
                    hists[trg][decayMode][wp].append(
                        file.Get(
                            "hist_{}TriggerEfficiency_{}TauMVA_{}{}".format(
                                trg, wp, dm_label, ft)).Clone())
                    graphs[trg][decayMode][wp].append(
                        file.Get(
                            "hist_{}TriggerEfficiency_{}TauMVA_{}{}".format(
                                trg, wp, dm_label, ft)).Clone())

    file.Close()

    latex = ROOT.TLatex()
    latex.SetNDC()
    for trg in triggers:
        for decayMode in decayModes:
            dm_label = "" if decayMode == "inclusive" else "dm{}_".format(
                decayMode)
            for wp in working_points:
                bin_label = "{}".format(wp)
                canv = ROOT.TCanvas('%s_%s' % (output, wp), output)

                if ratio_to is not None:
                    pads = plot.TwoPadSplit(0.50, 0.01, 0.01)
                else:
                    pads = plot.OnePad()
                slices = []

                if label_pos == 1:
                    text = ROOT.TPaveText(0.55, 0.37, 0.9, 0.50, 'NDC')
                    legend = ROOT.TLegend(0.18, 0.37, 0.5, 0.50, '', 'NDC')
                elif label_pos == 2:
                    text = ROOT.TPaveText(0.55, 0.67, 0.9, 0.80, 'NDC')
                    legend = ROOT.TLegend(0.18, 0.67, 0.5, 0.80, '', 'NDC')
                else:
                    text = ROOT.TPaveText(0.55, 0.54, 0.9, 0.67, 'NDC')
                    legend = ROOT.TLegend(0.55, 0.67, 0.9, 0.80, '', 'NDC')
                text = ROOT.TPaveText(0.55, 0.54, 0.9, 0.67, 'NDC')
                legend = ROOT.TLegend(0.6, 0.54, 0.95, 0.74, '', 'NDC')
                #~ if 'ID' in splitsrc[1]:
                #~ legend = ROOT.TLegend(0.18, 0.67, 0.5, 0.85, '', 'NDC')
                for j, hist in enumerate(graphs[trg][decayMode][wp]):
                    #splitsrc = src.split(':')
                    # htgr = hists[j].ProjectionX('%s_projx_%i' % (hists[j].GetName(), j), i, i)
                    #if len(splitsrc) >= 3:
                    #    settings = {x.split('=')[0]: eval(x.split('=')[1]) for x in splitsrc[2].split(',')}
                    if draw_options[j] != None:
                        settings = draw_options[j]
                        plot.Set(hist, **settings)
                    hist.Draw('LP SAME')  #htgr.Draw('SAME')
                    legend.AddEntry(hist)
                for j, hist in enumerate(hists[trg][decayMode][wp]):
                    if draw_options[j] != None:
                        settings = draw_options[j]
                        plot.Set(hist, **settings)
                    slices.append(hist)
                latex.SetTextSize(0.06)
                #~ text.AddText(args.title)
                #~ text.AddText(bin_label)
                #~ text.SetTextAlign(13)
                #~ text.SetBorderSize(0)
                #~ text.Draw()
                legend.Draw()
                pads[0].SetLogx(True)
                axis = plot.GetAxisHist(pads[0])
                axis.GetYaxis().SetTitle('Efficiency')
                axis.GetXaxis().SetTitle(x_title)
                # axis.GetXaxis().SetRangeUser(1,1000)
                axis.SetMinimum(float(y_range[0]))
                axis.SetMaximum(float(y_range[1]))
                #~ pads[0].SetGrid(0, 1)
                pads[0].RedrawAxis('g')

                #~ plot.DrawCMSLogo(pads[0], args.title, bin_label, 0, 0.16, 0.035, 1.2, cmsTextSize=0.5)
                plot.DrawTitle(
                    pads[0], trg + ' ' + title +
                    dm_label.replace('_', '').replace('dm', ' - dm') + ' - ' +
                    bin_label, 1)

                #plot.DrawTitle(pads[0], '18.99 fb^{-1} (13 TeV)', 3)
                if era == "2016":
                    plot.DrawTitle(pads[0], '35.9 fb^{-1} (2016, 13 TeV)', 3)
                elif era == "2017":
                    plot.DrawTitle(pads[0], '41.5 fb^{-1} (2017, 13 TeV)', 3)
                elif era == "2018":
                    plot.DrawTitle(pads[0], '59.7 fb^{-1} (2018, 13 TeV)', 3)

                if ratio_to is not None:
                    pads[1].cd()
                    pads[1].SetLogx(True)
                    ratios = []
                    for slice in slices:
                        ratios.append(slice.Clone())
                        ratios[-1].Divide(slices[ratio_to])
                    ratios[0].Draw('AXIS')
                    plot.SetupTwoPadSplitAsRatio(pads,
                                                 plot.GetAxisHist(pads[0]),
                                                 plot.GetAxisHist(pads[1]),
                                                 'Ratio to data', True,
                                                 float(ratio_y_range[0]),
                                                 float(ratio_y_range[1]))
                    for j, ratio in enumerate(ratios):
                        if j == ratio_to:
                            continue
                        ratio.Draw('LP SAME')  #('SAME E0')
                    pads[1].SetGrid(0, 1)
                    pads[1].RedrawAxis('g')

                outname = '{}_{}_{}_{}{}'.format(trg, era, output, dm_label,
                                                 wp)
                outname = outname.replace('(', '_')
                outname = outname.replace(')', '_')
                outname = outname.replace('.', '_')

                canv.Print('%s/%s.png' % (plot_dir, outname))
                canv.Print('%s/%s.pdf' % (plot_dir, outname))
def plot_lepton(files, label, era, output, draw_options, title, y_range,
                ratio_y_range, binned_in, x_title, ratio_to, plot_dir,
                label_pos):
    ROOT.PyConfig.IgnoreCommandLineOptions = True
    ROOT.gROOT.SetBatch(ROOT.kTRUE)
    ROOT.TH1.AddDirectory(0)
    plot.ModTDRStyle()

    # parser = argparse.ArgumentParser()

    # parser.add_argument(
    #     'input', nargs='+', help="""Input files""")
    # parser.add_argument(
    #     '--output', '-o', default='efficiency', help="""Name of the output
    #     plot without file extension""")
    # parser.add_argument('--title', default='Muon ID Efficiency')
    # parser.add_argument('--y-range', default='0,1')
    # parser.add_argument('--ratio-y-range', default='0.5,1.5')
    # parser.add_argument('--binned-in', default='#eta')
    # parser.add_argument('--x-title', default='p_{T} (GeV)')
    # parser.add_argument('--ratio-to', default=None, type=int)
    # parser.add_argument('--plot-dir', '-p', default='./')
    # parser.add_argument('--label-pos', default=1)
    # args = parser.parse_args()

    if plot_dir != '':
        os.system('mkdir -p %s' % plot_dir)

    hists = []

    # file = ROOT.TFile('%s.root' % target)

    # Process each input argument
    for src in files:
        #splitsrc = src.split(':')
        file = ROOT.TFile(src)
        if src.find("gen") >= 0:
            hists.append(file.Get(label + "_tot").Clone())
        else:
            hists.append(file.Get(label).Clone())
        file.Close()

    print hists

    hist = hists[0]

    latex = ROOT.TLatex()
    latex.SetNDC()

    for i in xrange(1, hist.GetNbinsY() + 1):
        bin_label = '%s: [%g,%g]' % (binned_in,
                                     hist.GetYaxis().GetBinLowEdge(i),
                                     hist.GetYaxis().GetBinUpEdge(i))
        canv = ROOT.TCanvas('%s_%i' % (output, i), output)

        if ratio_to is not None:
            pads = plot.TwoPadSplit(0.50, 0.01, 0.01)
        else:
            pads = plot.OnePad()
        slices = []

        if label_pos == 1:
            text = ROOT.TPaveText(0.55, 0.37, 0.9, 0.50, 'NDC')
            legend = ROOT.TLegend(0.18, 0.37, 0.5, 0.50, '', 'NDC')
        elif label_pos == 2:
            text = ROOT.TPaveText(0.55, 0.67, 0.9, 0.80, 'NDC')
            legend = ROOT.TLegend(0.18, 0.67, 0.5, 0.80, '', 'NDC')
        else:
            text = ROOT.TPaveText(0.55, 0.54, 0.9, 0.67, 'NDC')
            legend = ROOT.TLegend(0.55, 0.67, 0.9, 0.80, '', 'NDC')
        text = ROOT.TPaveText(0.55, 0.54, 0.9, 0.67, 'NDC')
        legend = ROOT.TLegend(0.6, 0.54, 0.95, 0.74, '', 'NDC')
        #~ if 'ID' in splitsrc[1]:
        #~ legend = ROOT.TLegend(0.18, 0.67, 0.5, 0.85, '', 'NDC')
        for j, src in enumerate(files):
            #splitsrc = src.split(':')
            htgr = hists[j].ProjectionX(
                '%s_projx_%i' % (hists[j].GetName(), j), i, i)
            #if len(splitsrc) >= 3:
            #    settings = {x.split('=')[0]: eval(x.split('=')[1]) for x in splitsrc[2].split(',')}
            if draw_options[j] != None:
                settings = draw_options[j]
                plot.Set(htgr, **settings)
            htgr.Draw('HIST LP SAME')  #htgr.Draw('SAME')
            legend.AddEntry(htgr)
            slices.append(htgr)
        latex.SetTextSize(0.06)
        #~ text.AddText(args.title)
        #~ text.AddText(bin_label)
        #~ text.SetTextAlign(13)
        #~ text.SetBorderSize(0)
        #~ text.Draw()
        legend.Draw()
        pads[0].SetLogx(True)
        axis = plot.GetAxisHist(pads[0])
        axis.GetYaxis().SetTitle('Efficiency')
        axis.GetXaxis().SetTitle(x_title)
        axis.GetXaxis().SetRangeUser(1, 1000)
        axis.SetMinimum(float(y_range[0]))
        axis.SetMaximum(float(y_range[1]))
        #~ pads[0].SetGrid(0, 1)
        pads[0].RedrawAxis('g')

        #~ plot.DrawCMSLogo(pads[0], args.title, bin_label, 0, 0.16, 0.035, 1.2, cmsTextSize=0.5)
        plot.DrawTitle(pads[0], title + ' - ' + bin_label, 1)

        #plot.DrawTitle(pads[0], '18.99 fb^{-1} (13 TeV)', 3)
        if era == "2016":
            plot.DrawTitle(pads[0], '35.9 fb^{-1} (2016, 13 TeV)', 3)
        elif era == "2017":
            plot.DrawTitle(pads[0], '41.5 fb^{-1} (2017, 13 TeV)', 3)
        elif era == "2018":
            plot.DrawTitle(pads[0], '59.7 fb^{-1} (2018, 13 TeV)', 3)

        if ratio_to is not None:
            pads[1].cd()
            pads[1].SetLogx(True)
            ratios = []
            for slice in slices:
                ratios.append(slice.Clone())
                ratios[-1].Divide(slices[ratio_to])
            ratios[0].Draw('AXIS')
            plot.SetupTwoPadSplitAsRatio(pads, plot.GetAxisHist(pads[0]),
                                         plot.GetAxisHist(pads[1]),
                                         'Ratio to data', True,
                                         float(ratio_y_range[0]),
                                         float(ratio_y_range[1]))
            for j, ratio in enumerate(ratios):
                if j == ratio_to:
                    continue
                ratio.Draw('HIST LP SAME')  #('SAME E0')
            pads[1].SetGrid(0, 1)
            pads[1].RedrawAxis('g')

        outname = '%s.%s.%g_%g' % (output, hist.GetYaxis().GetTitle(),
                                   hist.GetYaxis().GetBinLowEdge(i),
                                   hist.GetYaxis().GetBinUpEdge(i))
        outname = outname.replace('(', '_')
        outname = outname.replace(')', '_')
        outname = outname.replace('.', '_')

        canv.Print('%s/%s.png' % (plot_dir, outname))
        canv.Print('%s/%s.pdf' % (plot_dir, outname))
示例#4
0
def main():
    parser = argparse.ArgumentParser()

    parser.add_argument('main', help='Main input file for the scan')
    parser.add_argument('--y-cut',
                        type=float,
                        default=7.,
                        help='Remove points with y > y-cut')
    parser.add_argument('--y-max',
                        type=float,
                        default=8.,
                        help='y-axis maximum')
    parser.add_argument('--output',
                        '-o',
                        help='output name without file extension',
                        default='scan')
    parser.add_argument('--POI',
                        help='use this parameter of interest',
                        default='r')
    parser.add_argument('--translate',
                        default=None,
                        help='json file with POI name translation')
    parser.add_argument('--main-label',
                        default='Observed',
                        type=str,
                        help='legend label for the main scan')
    parser.add_argument('--main-color',
                        default=1,
                        type=int,
                        help='line and marker color for main scan')
    parser.add_argument(
        '--others',
        nargs='*',
        help='add secondary scans processed as main: FILE:LABEL:COLOR')
    parser.add_argument('--breakdown',
                        help='do quadratic error subtraction using --others')
    parser.add_argument('--logo', default='CMS')
    parser.add_argument('--logo-sub', default='Internal')
    args = parser.parse_args()

    # print '--------------------------------------'
    # print  args.output
    # print '--------------------------------------'

    fixed_name = args.POI
    xaxis_name = fixed_name
    label_name = fixed_name
    unit = ""
    if args.translate is not None:
        with open(args.translate) as jsonfile:
            name_translate = json.load(jsonfile)
        if args.POI in name_translate:
            fixed_name = name_translate[args.POI]

    # yvals = [1., 3.84]
    yvals = [1., 4.0]

    scalePOI = 9.2
    xaxis_name = "#sigma [fb]"
    label_name = "#sigma"
    unit = "fb"

    main_scan = BuildScan(args.output, "{0}*{1}".format(scalePOI, args.POI),
                          [args.main], args.main_color, yvals, args.y_cut)

    other_scans = []
    other_scans_opts = []
    if args.others is not None:
        for oargs in args.others:
            splitargs = oargs.split(':')
            other_scans_opts.append(splitargs)
            other_scans.append(
                BuildScan(args.output, args.POI, [splitargs[0]],
                          int(splitargs[2]), yvals, args.y_cut))

    canv = ROOT.TCanvas(args.output, args.output)
    pads = plot.OnePad()
    main_scan['graph'].SetMarkerColor(1)
    main_scan['graph'].Draw('AP')

    axishist = plot.GetAxisHist(pads[0])

    axishist.SetMaximum(args.y_max)
    axishist.GetYaxis().SetTitle("- 2 #Delta ln L")
    axishist.GetXaxis().SetTitle("%s" % xaxis_name)

    new_min = axishist.GetXaxis().GetXmin()
    new_max = axishist.GetXaxis().GetXmax()
    mins = []
    maxs = []
    for other in other_scans:
        mins.append(other['graph'].GetX()[0])
        maxs.append(other['graph'].GetX()[other['graph'].GetN() - 1])

    if len(other_scans) > 0:
        if min(mins) < main_scan['graph'].GetX()[0]:
            new_min = min(mins) - (main_scan['graph'].GetX()[0] - new_min)
        if max(maxs) > main_scan['graph'].GetX()[main_scan['graph'].GetN() -
                                                 1]:
            new_max = max(maxs) + (
                new_max -
                main_scan['graph'].GetX()[main_scan['graph'].GetN() - 1])
        axishist.GetXaxis().SetLimits(new_min, new_max)

    for other in other_scans:
        if args.breakdown is not None:
            other['graph'].SetMarkerSize(0.4)
        other['graph'].Draw('PSAME')

    line = ROOT.TLine()
    line.SetLineColor(16)
    # line.SetLineStyle(7)
    for yval in yvals:
        plot.DrawHorizontalLine(pads[0], line, yval)
        if (len(other_scans) == 0):
            for cr in main_scan['crossings'][yval]:
                if cr['valid_lo']: line.DrawLine(cr['lo'], 0, cr['lo'], yval)
                if cr['valid_hi']: line.DrawLine(cr['hi'], 0, cr['hi'], yval)

    main_scan['func'].Draw('same')
    for other in other_scans:
        if args.breakdown is not None:
            other['func'].SetLineStyle(2)
            other['func'].SetLineWidth(2)
        other['func'].Draw('SAME')

    box = ROOT.TBox(axishist.GetXaxis().GetXmin(), 0.625 * args.y_max,
                    axishist.GetXaxis().GetXmax(), args.y_max)
    box.Draw()
    pads[0].GetFrame().Draw()
    pads[0].RedrawAxis()

    crossings = main_scan['crossings']
    val_nom = main_scan['val']
    val_2sig = main_scan['val_2sig']

    textfit = '%s = %.1f{}^{#plus %.1f}_{#minus %.1f} %s' % (
        label_name, val_nom[0], val_nom[1], abs(val_nom[2]), unit)

    pt = ROOT.TPaveText(0.59, 0.82 - len(other_scans) * 0.08, 0.95, 0.91,
                        'NDCNB')
    pt.AddText(textfit)

    if args.breakdown is None:
        for i, other in enumerate(other_scans):
            textfit = '#color[%s]{%s = %.3f{}^{#plus %.3f}_{#minus %.3f}}' % (
                other_scans_opts[i][2], fixed_name, other['val'][0],
                other['val'][1], abs(other['val'][2]))
            pt.AddText(textfit)

    if args.breakdown is not None:
        pt.SetX1(0.50)
        if len(other_scans) >= 3:
            pt.SetX1(0.19)
            pt.SetX2(0.88)
            pt.SetY1(0.66)
            pt.SetY2(0.82)
        breakdown = args.breakdown.split(',')
        v_hi = [val_nom[1]]
        v_lo = [val_nom[2]]
        for other in other_scans:
            v_hi.append(other['val'][1])
            v_lo.append(other['val'][2])
        assert (len(v_hi) == len(breakdown))
        textfit = '%s = %.2f' % (label_name, val_nom[0])
        for i, br in enumerate(breakdown):
            if i < (len(breakdown) - 1):
                if (abs(v_hi[i + 1]) > abs(v_hi[i])):
                    print 'ERROR SUBTRACTION IS NEGATIVE FOR %s HI' % br
                    hi = 0.
                else:
                    hi = math.sqrt(v_hi[i] * v_hi[i] -
                                   v_hi[i + 1] * v_hi[i + 1])
                if (abs(v_lo[i + 1]) > abs(v_lo[i])):
                    print 'ERROR SUBTRACTION IS NEGATIVE FOR %s LO' % br
                    lo = 0.
                else:
                    lo = math.sqrt(v_lo[i] * v_lo[i] -
                                   v_lo[i + 1] * v_lo[i + 1])
            else:
                hi = v_hi[i]
                lo = v_lo[i]
            textfit += '{}^{#plus %.2f}_{#minus %.2f}(%s)' % (hi, abs(lo), br)
        pt.AddText(textfit)

    pt.SetTextAlign(11)
    pt.SetTextFont(42)
    pt.Draw()

    plot.DrawCMSLogo(pads[0],
                     args.logo,
                     args.logo_sub,
                     11,
                     0.045,
                     0.035,
                     1.2,
                     cmsTextSize=1.)

    legend_l = 0.69
    if len(other_scans) > 0:
        legend_l = legend_l - len(other_scans) * 0.04
    legend = ROOT.TLegend(0.15, legend_l, 0.45, 0.78, '', 'NBNDC')
    if len(other_scans) >= 3:
        legend = ROOT.TLegend(0.46, 0.83, 0.95, 0.93, '', 'NBNDC')
        legend.SetNColumns(2)

    legend.AddEntry(main_scan['func'], args.main_label, 'L')
    for i, other in enumerate(other_scans):
        legend.AddEntry(other['func'], other_scans_opts[i][1], 'L')
    legend.Draw()

    save_graph = main_scan['graph'].Clone()
    save_graph.GetXaxis().SetTitle(
        '%s = %.2f %+.2f/%+.2f' %
        (label_name, val_nom[0], val_nom[2], val_nom[1]))
    outfile = ROOT.TFile(args.output + '.root', 'RECREATE')
    outfile.WriteTObject(save_graph)
    outfile.Close()
    canv.Print('.pdf')
    canv.Print('.png')
    canv.Print('.C')