Beispiel #1
0
def test_stack():
    stack = HistStack()
    assert_equal(len(stack), 0)
    stack.Add(Hist(10, 0, 1, fillstyle='solid', color='red'))
    stack.Add(Hist(10, 0, 1, fillstyle='solid', color='blue'))
    stack.Add(Hist(10, 0, 1, fillstyle='solid', color='green'))
    assert_equal(len(stack), 3)
    stack2 = stack.Clone()
    assert_equal(stack2[2].linecolor, 'green')

    # test stacked=True
    a = Hist(2, 0, 1)
    b = Hist(2, 0, 1)
    a.Fill(0.2)
    b.Fill(0.2)
    b.Fill(0.8, 5)
    stack = HistStack([a, b])
    assert_equal(stack.min(), 2)
    assert_equal(stack.max(), 5)

    # test stacked=False
    a = Hist(2, 0, 20)
    b = Hist(2, 10, 20)  # binning may be different
    a.Fill(0.2)
    b.Fill(15, 5)
    stack = HistStack([a, b], stacked=False)
    assert_equal(stack.min(), 0)
    assert_equal(stack.max(), 5)
Beispiel #2
0
def test_stack():
    stack = HistStack()
    stack.Add(Hist(10, 0, 1, fillstyle='solid', color='red'))
    stack.Add(Hist(10, 0, 1, fillstyle='solid', color='blue'))
    stack.Add(Hist(10, 0, 1, fillstyle='solid', color='green'))
    assert_equal(len(stack), 3)
    stack2 = stack.Clone()
    assert_equal(stack2[2].linecolor, 'green')
Beispiel #3
0
def test_hist():
    from rootpy.plotting import root2matplotlib as rplt
    h = Hist(100, -5, 5)
    h.FillRandom('gaus')
    rplt.hist(h)

    # stack
    h1 = h.Clone()
    stack = HistStack([h, h1])
    rplt.hist(stack)
    rplt.hist([h, h1])
Beispiel #4
0
def test_bar():
    from rootpy.plotting import root2matplotlib as rplt
    h = Hist(100, -5, 5)
    h.FillRandom('gaus')
    rplt.bar(h)

    # stack
    h1 = h.Clone()
    stack = HistStack([h, h1])
    rplt.bar(stack)
    rplt.bar([h, h1], stacked=True)
    rplt.bar([h, h1], stacked=False)
    rplt.bar([h, h1], stacked=False, reverse=True)
Beispiel #5
0
def make_control_region_data_mc_comparision(
        histograms,
        histogram_name,
        variable,
        x_label=r'$m(\mathrm{t}\bar{\mathrm{t}})$ [GeV]',
        x_min=300,
        x_max=1800,
        y_label='Events/50 GeV'):
    qcd = histograms['QCD'][histogram_name]
    ttjet = histograms['TTJet'][histogram_name]
    wjets = histograms['WJets'][histogram_name]
    zjets = histograms['ZJets'][histogram_name]
    single_top = histograms['SingleTop'][histogram_name]
    other = ttjet + wjets + zjets + single_top
    data = histograms['data'][histogram_name]
    data.SetMarkerSize(3)
    qcd.SetTitle('QCD from data')
    other.SetTitle('Combined other background')
    data.SetTitle('Data')

    qcd.fillcolor = 'yellow'
    other.fillcolor = 'red'

    qcd.fillstyle = 'solid'
    other.fillstyle = 'solid'
    stack = HistStack()
    stack.Add(other)
    stack.Add(qcd)

    # plot with matplotlib
    plt.figure(figsize=(16, 12), dpi=200, facecolor='white')
    axes = plt.axes()

    rplt.hist(stack, stacked=True, axes=axes)
    rplt.errorbar(data, xerr=None, emptybins=False, axes=axes)

    plt.xlabel(x_label, CMS.x_axis_title)
    plt.ylabel(y_label, CMS.y_axis_title)
    plt.tick_params(**CMS.axis_label_major)
    plt.tick_params(**CMS.axis_label_minor)
    plt.title(get_title(variable), CMS.title)
    plt.legend(numpoints=1, loc='upper right', prop=CMS.legend_properties)
    axes.set_xlim(xmin=x_min, xmax=x_max)
    axes.set_ylim(ymin=0)
    plt.tight_layout()
    plt.savefig(variable + '.png')
    plt.savefig(variable + '.pdf')
Beispiel #6
0
map(h2.Fill, x2)
map(h3.Fill, x1_obs)
map(h3.Fill, x2_obs)

# set visual attributes
h1.fillstyle = 'solid'
h1.fillcolor = 'green'
h1.linecolor = 'green'
h1.linewidth = 0

h2.fillstyle = 'solid'
h2.fillcolor = 'red'
h2.linecolor = 'red'
h2.linewidth = 0

stack = HistStack()
stack.Add(h1)
stack.Add(h2)

ratio = h3 / (h1 + h2)
ratio = Hist.divide(h3, h1 + h2)

# plot with matplotlib
plt.figure(figsize=(16, 12), dpi=200)
gs = gridspec.GridSpec(2, 1, height_ratios=[5, 1])
axes = plt.axes()
ax0 = plt.subplot(gs[0])
ax0.minorticks_on()
rplt.bar(stack, stacked=True)
rplt.errorbar(h3, xerr=False, emptybins=False)
# plt.xlabel('Mass', position=(1., 0.), ha='right', fontsize = 24)
def make_data_mc_comparison_plot(
    histograms=[],
    histogram_lables=[],
    histogram_colors=[],
    histogram_properties=Histogram_properties(),
    data_index=0,
    save_folder='plots/',
    save_as=['pdf'],
    normalise=False,
    show_ratio=False,
    show_stat_errors_on_mc=False,
    draw_vertical_line=0,
    systematics_for_ratio=None,
    systematics_for_plot=None,
    histograms_to_compare=None,
):
    '''
    systematics_for_plot takes the same input as systematics_for_ratio. There may be some repition with reagrds to 
    mc_error and mc_relative_errors, but these only deal with a flat error up and down.
    '''
    save_folder = check_save_folder(save_folder)
    # make copies in order not to mess with existing histograms
    histograms_ = deepcopy(histograms)
    stack = HistStack()
    add_mc = stack.Add
    for index, histogram in enumerate(histograms_):
        label = histogram_lables[index]
        color = histogram_colors[index]

        histogram.SetTitle(label)
        if normalise:
            histogram.Sumw2()

        if not index == data_index:
            histogram.fillstyle = 'solid'
            histogram.fillcolor = color
            histogram.legendstyle = 'F'
            add_mc(histogram)

    data = histograms_[data_index]
    data.SetMarkerSize(CMS.data_marker_size)
    if normalise:
        n_events_data = data.Integral()
        n_events_mc = stack.Integral()
        data.Scale(1 / n_events_data)
        stack.Scale(1 / n_events_mc)

    # plot with matplotlib
    plt.figure(figsize=CMS.figsize, dpi=CMS.dpi, facecolor=CMS.facecolor)
    axes = None
    if show_ratio:
        ratio = data.Clone('ratio')
        sumHists = sum(stack.GetHists())
        for bin_i in range(1, sumHists.GetNbinsX()):
            sumHists.SetBinError(bin_i, 0)
        ratio.Divide(sum(stack.GetHists()))
        ratio.SetMarkerSize(3)
        gs = gridspec.GridSpec(2, 1, height_ratios=[5, 1])
        axes = plt.subplot(gs[0])
    else:
        axes = plt.axes()

    if histogram_properties.set_log_y:
        axes.set_yscale('log', nonposy="clip")
        axes.set_ylim(ymin=1e-2)

    if systematics_for_plot != None:
        plusErrors = [x + 1 for x in systematics_for_plot]
        minusErrors = [1 - x for x in systematics_for_plot]

        stack_lower = sum(stack.GetHists())
        stack_upper = stack_lower.Clone('upper')

        for bin_i in range(1, stack_lower.GetNbinsX() + 1):
            central_value = stack_lower.GetBinContent(bin_i)
            error_upper_bound = plusErrors[bin_i - 1] * central_value
            error_lower_bound = minusErrors[bin_i - 1] * central_value
            stack_upper.SetBinContent(bin_i, error_upper_bound)
            stack_lower.SetBinContent(bin_i, error_lower_bound)

        rplt.fill_between(
            stack_lower,
            stack_upper,
            axes,
            hatch='//',
            # facecolor = 'Black',
            facecolor='None',
            edgecolor='Grey',
            alpha=1.,
            linewidth=0.,
            zorder=len(histograms_) + 1)

    mc_error = histogram_properties.mc_error
    mc_relative_errors = histogram_properties.mc_relative_errors
    if mc_relative_errors:
        stack_lower = sum(stack.GetHists())
        stack_upper = stack_lower.Clone('upper')
        for bin_i in range(1, stack_lower.GetNbinsX()):
            central_value = stack_lower.GetBinContent(bin_i)
            relative_error = mc_relative_errors[bin_i - 1]

            error_upper_bound = central_value * (1 + relative_error)
            error_lower_bound = central_value * (1 - relative_error)

            stack_lower.SetBinContent(bin_i, error_upper_bound)
            stack_upper.SetBinContent(bin_i, error_lower_bound)
        rplt.fill_between(stack_upper,
                          stack_lower,
                          axes,
                          facecolor='0.75',
                          alpha=0.5,
                          hatch='/',
                          zorder=len(histograms_) + 1)
    else:
        if mc_error > 0:
            stack_lower = sum(stack.GetHists())
            stack_upper = stack_lower.Clone('upper')
            stack_lower.Scale(1 - mc_error)
            stack_upper.Scale(1 + mc_error)
            rplt.fill_between(stack_upper,
                              stack_lower,
                              axes,
                              facecolor='0.75',
                              alpha=0.5,
                              hatch='/',
                              zorder=len(histograms_) + 1)
        if not mc_error > 0 and show_stat_errors_on_mc:
            stack_lower = sum(stack.GetHists())
            mc_errors = list(stack_lower.yerravg())
            stack_upper = stack_lower.Clone('upper')
            for bin_i in range(1, stack_lower.GetNbinsX()):
                central_value = stack_lower.GetBinContent(bin_i)
                error = mc_errors[bin_i - 1]
                error_upper_bound = central_value + error
                error_lower_bound = central_value - error
                stack_lower.SetBinContent(bin_i, error_lower_bound)
                stack_upper.SetBinContent(bin_i, error_upper_bound)
            rplt.fill_between(stack_upper,
                              stack_lower,
                              axes,
                              facecolor='0.75',
                              alpha=0.5,
                              hatch='/',
                              zorder=len(histograms_) + 1)

    # a comment on zorder: the MC stack should be always at the very back (z = 1),
    # then the MC error (z = len(histograms_) + 1) and finally the data
    # (z = len(histograms_) + 2)
    rplt.hist(stack, stacked=True, axes=axes, zorder=1)
    rplt.errorbar(data,
                  emptybins=histogram_properties.emptybins,
                  axes=axes,
                  xerr=histogram_properties.xerr,
                  elinewidth=2,
                  capsize=10,
                  capthick=2,
                  zorder=len(histograms_) + 2)

    if histograms_to_compare:
        h_compare = {}
        for h, l, c in zip(histograms_to_compare['hists'],
                           histograms_to_compare['labels'],
                           histograms_to_compare['colours']):
            for histogram in histograms_:
                if histogram.GetTitle() not in [
                        histograms_to_compare['to_replace'], 'data'
                ]:
                    h += histogram
            h_compare[l] = [h, c]
            rplt.step(
                h,
                axes=axes,
                label=l,
                color=c,
                linewidth=4,
            )

    # put legend into the correct order (data is always first!)
    handles, labels = axes.get_legend_handles_labels()
    data_label_index = labels.index('data')
    data_handle = handles[data_label_index]
    labels.remove('data')
    handles.remove(data_handle)
    labels.insert(0, 'data')
    handles.insert(0, data_handle)
    if mc_error > 0 or (not mc_error > 0 and show_stat_errors_on_mc):
        p1 = Rectangle((0, 0), 1, 1, fc="0.75", alpha=0.5, hatch='/')
        handles.append(p1)
        labels.append(histogram_properties.mc_errors_label)

    l1 = axes.legend(handles,
                     labels,
                     numpoints=1,
                     frameon=histogram_properties.legend_color,
                     bbox_to_anchor=histogram_properties.legend_location,
                     bbox_transform=plt.gcf().transFigure,
                     prop=CMS.legend_properties,
                     ncol=histogram_properties.legend_columns)
    l1.set_zorder(102)

    set_labels(plt,
               histogram_properties,
               show_x_label=not show_ratio,
               axes=axes)

    x_limits = histogram_properties.x_limits
    y_limits = histogram_properties.y_limits

    if len(x_limits) >= 2:
        axes.set_xlim(xmin=x_limits[0], xmax=x_limits[-1])
    if len(y_limits) >= 2:
        axes.set_ylim(ymin=y_limits[0], ymax=y_limits[-1])
    else:
        y_max = get_best_max_y(
            histograms_, x_limits=x_limits) * histogram_properties.y_max_scale
        print("Chosen limits : ", 0, y_max)
        axes.set_ylim(ymin=0, ymax=y_max)
    if histogram_properties.set_log_y:
        if not len(y_limits) == 2:  # if not user set y-limits, set default
            axes.set_ylim(ymin=1e-1)

    #draw a red vertical line if needed:
    if draw_vertical_line != 0:
        plt.axvline(x=draw_vertical_line, color='red', linewidth=3)

    if show_ratio:
        plt.setp(axes.get_xticklabels(), visible=False)
        ax1 = plt.subplot(gs[1])
        ax1.minorticks_on()
        ax1.grid(True, 'major', linewidth=1)
        ax1.axhline(y=1, linewidth=1)
        set_labels(plt,
                   histogram_properties,
                   show_x_label=True,
                   show_title=False)
        plt.ylabel(r'$\frac{\mathrm{data}}{\mathrm{pred.}}$', CMS.y_axis_title)
        ax1.yaxis.set_label_coords(-0.115, 0.8)
        rplt.errorbar(ratio,
                      emptybins=histogram_properties.emptybins,
                      axes=ax1,
                      xerr=histogram_properties.xerr,
                      elinewidth=1.5,
                      capsize=5,
                      capthick=1.5)
        if histograms_to_compare:
            for l, h in h_compare.iteritems():
                r = data.Clone(l).Divide(h[0])
                rplt.step(
                    r,
                    axes=ax1,
                    label='',
                    colour=h[1],
                    linewidth=2,
                )

        if len(x_limits) >= 2:
            ax1.set_xlim(xmin=x_limits[0], xmax=x_limits[-1])
        if len(histogram_properties.ratio_y_limits) >= 2:
            ax1.set_ylim(ymin=histogram_properties.ratio_y_limits[0],
                         ymax=histogram_properties.ratio_y_limits[-1])

        # dynamic tick placement
        adjust_ratio_ticks(ax1.yaxis,
                           n_ticks=3,
                           y_limits=histogram_properties.ratio_y_limits)

        if histogram_properties.integerXVariable:
            ax1.tick_params(axis='x', which='minor', bottom='off', top='off')

        if systematics_for_ratio != None:
            plusErrors = [x + 1 for x in systematics_for_ratio]
            minusErrors = [1 - x for x in systematics_for_ratio]

            ratioPlusError = ratio.Clone('plus')
            ratioMinusError = ratio.Clone('minus')
            for bin_i in range(1, ratioPlusError.GetNbinsX() + 1):
                ratioPlusError.SetBinContent(bin_i, plusErrors[bin_i - 1])
                ratioMinusError.SetBinContent(bin_i, minusErrors[bin_i - 1])
            rplt.fill_between(
                ratioPlusError,
                ratioMinusError,
                axes,
                hatch='//',
                # facecolor = 'Black',
                facecolor='None',
                edgecolor='Grey',
                alpha=1.,
                linewidth=0.,
                # zorder = len(histograms_) + 1
                zorder=0)

    if CMS.tight_layout:
        plt.tight_layout()

    for save in save_as:
        if save == 'root':
            saveHistogramsToROOTFile(
                data, stack,
                save_folder + histogram_properties.name + '.' + save)
        else:
            plt.savefig(save_folder + histogram_properties.name + '.' + save)

    plt.close()
Beispiel #8
0
    def plot(self):

        evaluator = Evaluator(self.model, self.transformation, self.features)
        makedirs(self.plt_dir, exist_ok=True)
        makedirs('/'.join([self.plt_dir, 'lin']), exist_ok=True)
        makedirs('/'.join([self.plt_dir, 'log']), exist_ok=True)
        makedirs('/'.join([self.plt_dir, 'lin', 'png']), exist_ok=True)
        makedirs('/'.join([self.plt_dir, 'lin', 'root']), exist_ok=True)
        makedirs('/'.join([self.plt_dir, 'log', 'png']), exist_ok=True)
        makedirs('/'.join([self.plt_dir, 'log', 'root']), exist_ok=True)
        makedirs('/'.join([self.plt_dir, 'lnt_region', 'lin']), exist_ok=True)
        makedirs('/'.join([self.plt_dir, 'lnt_region', 'log']), exist_ok=True)
        makedirs('/'.join([self.plt_dir, 'lnt_region', 'lin', 'png']),
                 exist_ok=True)
        makedirs('/'.join([self.plt_dir, 'lnt_region', 'lin', 'root']),
                 exist_ok=True)
        makedirs('/'.join([self.plt_dir, 'lnt_region', 'log', 'png']),
                 exist_ok=True)
        makedirs('/'.join([self.plt_dir, 'lnt_region', 'log', 'root']),
                 exist_ok=True)
        makedirs('/'.join([self.plt_dir, 'shapes', 'lin']), exist_ok=True)
        makedirs('/'.join([self.plt_dir, 'shapes', 'log']), exist_ok=True)
        makedirs('/'.join([self.plt_dir, 'shapes', 'lin', 'png']),
                 exist_ok=True)
        makedirs('/'.join([self.plt_dir, 'shapes', 'lin', 'root']),
                 exist_ok=True)
        makedirs('/'.join([self.plt_dir, 'shapes', 'log', 'png']),
                 exist_ok=True)
        makedirs('/'.join([self.plt_dir, 'shapes', 'log', 'root']),
                 exist_ok=True)

        # NN evaluator
        print('============> starting reading the trees')
        print('Plots will be stored in: ', self.plt_dir)
        now = time()
        signal = []
        if self.process_signals:
            signal = self.get_signal_samples(self.channel,
                                             self.base_dir,
                                             self.post_fix,
                                             self.selection_data,
                                             mini=self.mini_signals)
        else:
            signal = []
        data = self.get_data_samples(self.channel, self.base_dir,
                                     self.post_fix, self.selection_data)
        mc = self.get_mc_samples(self.channel, self.base_dir, self.post_fix,
                                 self.selection_mc)
        print('============> it took %.2f seconds' % (time() - now))

        # evaluate FR
        for isample in (mc + data):  #+signal):
            isample.df['fr'] = evaluator.evaluate(isample.df)
            # already corrected, ready to be applied in lnt-not-tight
            isample.df['fr_corr'] = isample.df['fr'] / (1. - isample.df['fr'])

        # apply an extra selection to the pandas dataframes
        if len(self.pandas_selection):
            for isample in (mc + data + signal):
                isample.df = isample.df.query(self.pandas_selection)

        # split the dataframe in tight and lnt-not-tight (called simply lnt for short)
        print('============> splitting dataframe in tight and loose not tight')
        for isample in (mc + data + signal):
            isample.df_tight = isample.df.query(self.selection_tight)
            if isample not in signal:
                isample.df_lnt = isample.df.query(self.selection_lnt)
            # free some mem
            del isample.df
            gc.collect()
        print('============> ... done')

        # sort depending on their position in the stack
        mc.sort(key=lambda x: x.position_in_stack)

        # now we plot
        self.create_canvas(self.do_ratio)

        for ivar in variables:

            variable, bins, label, xlabel, ylabel, extra_sel = ivar.var, ivar.bins, ivar.label, ivar.xlabel, ivar.ylabel, ivar.extra_selection

            print('plotting', label)

            ######################################################################################
            # plot MC stacks, in tight and lnt
            ######################################################################################

            stack_prompt = []
            stack_nonprompt = []
            stack_nonprompt_control = []

            for imc in mc:

                if extra_sel:
                    mc_df_tight = imc.df_tight.query(extra_sel)
                    mc_df_lnt = imc.df_lnt.query(extra_sel)
                else:
                    mc_df_tight = imc.df_tight
                    mc_df_lnt = imc.df_lnt

                histo_tight = Hist(bins,
                                   title=imc.label,
                                   markersize=0,
                                   legendstyle='F',
                                   name=imc.datacard_name + '#' + label + '#t')
                weights = self.total_weight_calculator(
                    mc_df_tight, ['weight'] + imc.extra_signal_weights,
                    [self.lumi, imc.lumi_scaling])
                histo_tight.fill_array(mc_df_tight[variable],
                                       weights=weights *
                                       self.relaxed_mc_scaling)

                histo_tight.fillstyle = 'solid'
                histo_tight.fillcolor = 'steelblue' if self.data_driven else imc.colour
                histo_tight.linewidth = 0

                stack_prompt.append(histo_tight)

                # optionally remove the MC subtraction in loose-not-tight
                # may help if MC stats is terrible (and it often is)
                if self.data_driven:
                    if self.mc_subtraction:
                        histo_lnt = Hist(bins,
                                         title=imc.label,
                                         markersize=0,
                                         legendstyle='F',
                                         name=imc.datacard_name + '#' + label +
                                         '#lnt')
                        weights = self.total_weight_calculator(
                            mc_df_lnt,
                            ['weight', 'fr_corr'] + imc.extra_signal_weights,
                            [-1., self.lumi, imc.lumi_scaling])
                        histo_lnt.fill_array(mc_df_lnt[variable],
                                             weights=weights *
                                             self.relaxed_mc_scaling)

                        histo_lnt.fillstyle = 'solid'
                        histo_lnt.fillcolor = 'skyblue' if self.data_driven else imc.colour
                        histo_lnt.linewidth = 0
                        stack_nonprompt.append(histo_lnt)

                    histo_lnt_control = Hist(bins,
                                             title=imc.label,
                                             markersize=0,
                                             legendstyle='F',
                                             name=imc.datacard_name + '#' +
                                             label + '#lntcontrol')
                    weights_control = self.total_weight_calculator(
                        mc_df_lnt, ['weight'] + imc.extra_signal_weights,
                        [self.lumi, imc.lumi_scaling])
                    histo_lnt_control.fill_array(mc_df_lnt[variable],
                                                 weights=weights_control *
                                                 self.relaxed_mc_scaling)

                    histo_lnt_control.fillstyle = 'solid'
                    histo_lnt_control.fillcolor = imc.colour
                    histo_lnt_control.linewidth = 0

                    #                     print(histo_lnt_control)
                    #                     print(histo_lnt_control.fillcolor)
                    #                     print(imc.name, imc.colour)
                    #                     print(histo_lnt_control.integral())
                    stack_nonprompt_control.append(histo_lnt_control)

            # merge different samples together (add the histograms)
            # prepare two temporary containers for the post-grouping histograms
            stack_prompt_tmp = []
            stack_nonprompt_tmp = []
            stack_nonprompt_control_tmp = []
            for ini, fin in [(stack_prompt, stack_prompt_tmp),
                             (stack_nonprompt, stack_nonprompt_tmp),
                             (stack_nonprompt_control,
                              stack_nonprompt_control_tmp)]:
                for k, v in groups.items():
                    grouped = []
                    for ihist in ini:
                        if ihist.name.split('#')[0] in v:
                            grouped.append(ihist)
                        elif ihist.name.split('#')[0] not in togroup:
                            fin.append(ihist)
                    if len(grouped):
                        group = sum(grouped)
                        group.title = k
                        group.name = '#'.join([k] + ihist.name.split('#')[1:])
                        group.fillstyle = grouped[0].fillstyle
                        group.fillcolor = grouped[0].fillcolor
                        group.linewidth = grouped[0].linewidth
                    fin.append(group)

            stack_prompt = stack_prompt_tmp
            stack_nonprompt = stack_nonprompt_tmp
            stack_nonprompt_control = stack_nonprompt_control_tmp

            ######################################################################################
            # plot the signals
            ######################################################################################

            all_signals = []
            signals_to_plot = []

            for isig in signal:

                if variable not in self.datacards:
                    if not isig.toplot:
                        continue

                if variable == 'fr' or variable == 'fr_corr':
                    continue

                if extra_sel:
                    isig_df_tight = isig.df_tight.query(extra_sel)
                else:
                    isig_df_tight = isig.df_tight

                histo_tight = Hist(
                    bins,
                    title=isig.label,
                    markersize=0,
                    legendstyle='L',
                    name=isig.datacard_name + '#' + label
                )  # the "#" thing is a trick to give hists unique name, else ROOT complains
                weights = self.total_weight_calculator(
                    isig_df_tight, ['weight'] + isig.extra_signal_weights,
                    [self.lumi, isig.lumi_scaling])
                histo_tight.fill_array(isig_df_tight[variable],
                                       weights=weights)
                histo_tight.color = isig.colour
                histo_tight.fillstyle = 'hollow'
                histo_tight.linewidth = 2
                histo_tight.linestyle = 'dashed'
                histo_tight.drawstyle = 'HIST'

                all_signals.append(histo_tight)
                if isig.toplot: signals_to_plot.append(histo_tight)

            ######################################################################################
            # plot the data
            ######################################################################################

            data_prompt = []
            data_nonprompt = []
            data_nonprompt_control = []

            for idata in data:

                if extra_sel:
                    idata_df_tight = idata.df_tight.query(extra_sel)
                    idata_df_lnt = idata.df_lnt.query(extra_sel)
                else:
                    idata_df_tight = idata.df_tight
                    idata_df_lnt = idata.df_lnt

                histo_tight = Hist(bins,
                                   title=idata.label,
                                   markersize=1,
                                   legendstyle='LEP')
                histo_tight.fill_array(idata_df_tight[variable])

                data_prompt.append(histo_tight)

                if self.data_driven:
                    histo_lnt = Hist(bins,
                                     title=idata.label,
                                     markersize=0,
                                     legendstyle='F')
                    histo_lnt.fill_array(idata_df_lnt[variable],
                                         weights=idata_df_lnt.fr_corr)

                    histo_lnt.fillstyle = 'solid'
                    histo_lnt.fillcolor = 'skyblue'
                    histo_lnt.linewidth = 0

                    histo_lnt_control = Hist(bins,
                                             title=idata.label,
                                             markersize=1,
                                             legendstyle='LEP')
                    histo_lnt_control.fill_array(idata_df_lnt[variable])

                    data_nonprompt.append(histo_lnt)
                    data_nonprompt_control.append(histo_lnt_control)

            if self.data_driven:
                # put the prompt backgrounds together
                all_exp_prompt = sum(stack_prompt)
                all_exp_prompt.title = 'prompt'

                # put the nonprompt backgrounds together
                all_exp_nonprompt = sum(stack_nonprompt + data_nonprompt)
                all_exp_nonprompt.fillstyle = 'solid'
                all_exp_nonprompt.fillcolor = 'skyblue'
                all_exp_nonprompt.linewidth = 0
                all_exp_nonprompt.title = 'nonprompt'

                # create the stacks
                stack = HistStack([all_exp_prompt, all_exp_nonprompt],
                                  drawstyle='HIST',
                                  title='')
                stack_control = HistStack(stack_nonprompt_control,
                                          drawstyle='HIST',
                                          title='')

            else:
                stack = HistStack(stack_prompt, drawstyle='HIST', title='')

            # stat uncertainty
            hist_error = stack.sum  #sum([all_exp_prompt, all_exp_nonprompt])
            hist_error.drawstyle = 'E2'
            hist_error.fillstyle = '/'
            hist_error.color = 'gray'
            hist_error.title = 'stat. unc.'
            hist_error.legendstyle = 'F'

            if self.data_driven:
                hist_error_control = stack_control.sum
                hist_error_control.drawstyle = 'E2'
                hist_error_control.fillstyle = '/'
                hist_error_control.color = 'gray'
                hist_error_control.title = 'stat. unc.'
                hist_error_control.legendstyle = 'F'

            # put the data together
            all_obs_prompt = sum(data_prompt)
            all_obs_prompt.title = 'observed'

            if self.data_driven:
                all_obs_nonprompt_control = sum(data_nonprompt_control)
                all_obs_nonprompt_control.title = 'observed'
                all_obs_nonprompt_control.drawstyle = 'EP'

            # prepare the legend
            print(signals_to_plot)
            for jj in signals_to_plot:
                print(jj.name, jj.integral())
            if len(signals_to_plot):
                legend = Legend([all_obs_prompt, stack, hist_error],
                                pad=self.main_pad,
                                leftmargin=0.,
                                rightmargin=0.,
                                topmargin=0.,
                                textfont=42,
                                textsize=0.025,
                                entrysep=0.01,
                                entryheight=0.04)
                legend_signals = Legend(signals_to_plot,
                                        pad=self.main_pad,
                                        leftmargin=0.,
                                        rightmargin=0.,
                                        topmargin=0.,
                                        textfont=42,
                                        textsize=0.025,
                                        entrysep=0.01,
                                        entryheight=0.04)
                legend_signals.SetBorderSize(0)
                legend_signals.x1 = 0.42
                legend_signals.y1 = 0.74
                legend_signals.x2 = 0.88
                legend_signals.y2 = 0.90
                legend_signals.SetFillColor(0)
                legend.SetBorderSize(0)
                legend.x1 = 0.2
                legend.y1 = 0.74
                legend.x2 = 0.45
                legend.y2 = 0.90
                legend.SetFillColor(0)
            else:
                legend = Legend([all_obs_prompt, stack, hist_error],
                                pad=self.main_pad,
                                leftmargin=0.,
                                rightmargin=0.,
                                topmargin=0.,
                                textfont=42,
                                textsize=0.03,
                                entrysep=0.01,
                                entryheight=0.04)
                legend.SetBorderSize(0)
                legend.x1 = 0.55
                legend.y1 = 0.74
                legend.x2 = 0.88
                legend.y2 = 0.90
                legend.SetFillColor(0)

            # plot with ROOT, linear and log scale
            for islogy in [False, True]:

                things_to_plot = [stack, hist_error]
                if not self.blinded:
                    things_to_plot.append(all_obs_prompt)

                # plot signals, as an option
                if self.plot_signals:
                    things_to_plot += signals_to_plot

                # set the y axis range
                # FIXME! setting it by hand to each object as it doesn't work if passed to draw
                if islogy:
                    yaxis_max = 40. * max(
                        [ithing.max() for ithing in things_to_plot])
                else:
                    yaxis_max = 1.65 * max(
                        [ithing.max() for ithing in things_to_plot])
                if islogy: yaxis_min = 0.01
                else: yaxis_min = 0.

                for ithing in things_to_plot:
                    ithing.SetMaximum(yaxis_max)
                draw(things_to_plot,
                     xtitle=xlabel,
                     ytitle=ylabel,
                     pad=self.main_pad,
                     logy=islogy)

                # expectation uncertainty in the ratio pad
                ratio_exp_error = Hist(bins)
                ratio_data = Hist(bins)
                for ibin in hist_error.bins_range():
                    ratio_exp_error.set_bin_content(ibin, 1.)
                    ratio_exp_error.set_bin_error(
                        ibin,
                        hist_error.get_bin_error(ibin) /
                        hist_error.get_bin_content(ibin)
                        if hist_error.get_bin_content(ibin) != 0. else 0.)
                    ratio_data.set_bin_content(
                        ibin,
                        all_obs_prompt.get_bin_content(ibin) /
                        hist_error.get_bin_content(ibin)
                        if hist_error.get_bin_content(ibin) != 0. else 0.)
                    ratio_data.set_bin_error(
                        ibin,
                        all_obs_prompt.get_bin_error(ibin) /
                        hist_error.get_bin_content(ibin)
                        if hist_error.get_bin_content(ibin) != 0. else 0.)

                ratio_data.drawstyle = 'EP'
                ratio_data.title = ''

                ratio_exp_error.drawstyle = 'E2'
                ratio_exp_error.markersize = 0
                ratio_exp_error.title = ''
                ratio_exp_error.fillstyle = '/'
                ratio_exp_error.color = 'gray'

                for ithing in [ratio_data, ratio_exp_error]:
                    ithing.xaxis.set_label_size(
                        ithing.xaxis.get_label_size() * 3.
                    )  # the scale should match that of the main/ratio pad size ratio
                    ithing.yaxis.set_label_size(
                        ithing.yaxis.get_label_size() * 3.
                    )  # the scale should match that of the main/ratio pad size ratio
                    ithing.xaxis.set_title_size(
                        ithing.xaxis.get_title_size() * 3.
                    )  # the scale should match that of the main/ratio pad size ratio
                    ithing.yaxis.set_title_size(
                        ithing.yaxis.get_title_size() * 3.
                    )  # the scale should match that of the main/ratio pad size ratio
                    ithing.yaxis.set_ndivisions(405)
                    ithing.yaxis.set_title_offset(0.4)

                things_to_plot = [ratio_exp_error]
                if not self.blinded:
                    things_to_plot.append(ratio_data)

                draw(things_to_plot,
                     xtitle=xlabel,
                     ytitle='obs/exp',
                     pad=self.ratio_pad,
                     logy=False,
                     ylimits=(0.5, 1.5))

                line = ROOT.TLine(min(bins), 1., max(bins), 1.)
                line.SetLineColor(ROOT.kBlack)
                line.SetLineWidth(1)
                self.ratio_pad.cd()
                line.Draw('same')

                #                 chi2_score_text = '\chi^{2}/NDF = %.1f' %(all_obs_prompt.Chi2Test(hist_error, 'UW CHI2/NDF'))
                chi2_score_text = 'p-value = %.2f' % (all_obs_prompt.Chi2Test(
                    hist_error, 'UW'))
                chi2_score = ROOT.TLatex(0.7, 0.81, chi2_score_text)
                chi2_score.SetTextFont(43)
                chi2_score.SetTextSize(15)
                chi2_score.SetNDC()
                chi2_score.Draw('same')

                self.canvas.cd()
                # FIXME! add SS and OS channels
                if self.full_channel == 'mmm': channel = '\mu\mu\mu'
                elif self.full_channel == 'eee': channel = 'eee'
                elif self.full_channel == 'mem_os':
                    channel = '\mu^{\pm}\mu^{\mp}e'
                elif self.full_channel == 'mem_ss':
                    channel = '\mu^{\pm}\mu^{\pm}e'
                elif self.full_channel == 'eem_os':
                    channel = 'e^{\pm}e^{\mp}\mu'
                elif self.full_channel == 'eem_ss':
                    channel = 'e^{\pm}e^{\pm}\mu'
                else:
                    assert False, 'ERROR: Channel not valid.'
                finalstate = ROOT.TLatex(0.68, 0.68, channel)
                finalstate.SetTextFont(43)
                finalstate.SetTextSize(25)
                finalstate.SetNDC()
                finalstate.Draw('same')

                self.canvas.cd()
                # remove old legend
                for iprim in self.canvas.primitives:
                    if isinstance(iprim, Legend):
                        self.canvas.primitives.remove(iprim)
                legend.Draw('same')
                if self.plot_signals:
                    legend_signals.Draw('same')
                CMS_lumi(self.main_pad,
                         4,
                         0,
                         lumi_13TeV="%d, L = %.1f fb^{-1}" %
                         (self.year, self.lumi / 1000.))
                self.canvas.Modified()
                self.canvas.Update()
                for iformat in ['pdf', 'png', 'root']:
                    self.canvas.SaveAs('/'.join([
                        self.plt_dir, 'log' if islogy else 'lin',
                        iformat if iformat != 'pdf' else '',
                        '%s%s.%s' %
                        (label, '_log' if islogy else '_lin', iformat)
                    ]))

                # plot distributions in loose not tight
                # check MC contamination there
                if self.data_driven and variable not in ['fr', 'fr_corr']:
                    things_to_plot = [
                        stack_control, hist_error_control,
                        all_obs_nonprompt_control
                    ]
                    # set the y axis range
                    # FIXME! setting it by hand to each object as it doesn't work if passed to draw
                    if islogy:
                        yaxis_max = 40. * max(
                            [ithing.max() for ithing in things_to_plot])
                    else:
                        yaxis_max = 1.65 * max(
                            [ithing.max() for ithing in things_to_plot])
                    if islogy: yaxis_min = 0.01
                    else: yaxis_min = 0.

                    for ithing in things_to_plot:
                        ithing.SetMaximum(yaxis_max)
                        ithing.SetMinimum(yaxis_min)

                    draw(things_to_plot,
                         xtitle=xlabel,
                         ytitle=ylabel,
                         pad=self.main_pad,
                         logy=islogy,
                         ylimits=(yaxis_min, yaxis_max))

                    new_legend = Legend(
                        stack_control.hists +
                        [hist_error_control, all_obs_nonprompt_control],
                        pad=self.main_pad,
                        leftmargin=0.,
                        rightmargin=0.,
                        topmargin=0.,
                        textfont=42,
                        textsize=0.03,
                        entrysep=0.01,
                        entryheight=0.04)
                    new_legend.SetBorderSize(0)
                    new_legend.x1 = 0.55
                    new_legend.y1 = 0.71
                    new_legend.x2 = 0.88
                    new_legend.y2 = 0.90
                    new_legend.SetFillColor(0)

                    # divide MC to subtract by data
                    stack_nonprompt_control_scaled_list = []
                    for ihist in stack_control.hists:
                        new_hist = copy(ihist)
                        for ibin in new_hist.bins_range():
                            new_hist.SetBinContent(
                                ibin,
                                np.nan_to_num(
                                    np.divide(
                                        new_hist.GetBinContent(ibin),
                                        all_obs_nonprompt_control.
                                        GetBinContent(ibin))))
                            new_hist.SetBinError(
                                ibin,
                                np.nan_to_num(
                                    np.divide(
                                        new_hist.GetBinError(ibin),
                                        all_obs_nonprompt_control.
                                        GetBinContent(ibin))))
                        stack_nonprompt_control_scaled_list.append(new_hist)

                    stack_control_scaled = HistStack(
                        stack_nonprompt_control_scaled_list,
                        drawstyle='HIST',
                        title='')
                    stack_control_scaled_err = stack_control_scaled.sum
                    stack_control_scaled_err.drawstyle = 'E2'
                    stack_control_scaled_err.fillstyle = '/'
                    stack_control_scaled_err.color = 'gray'
                    stack_control_scaled_err.title = 'stat. unc.'
                    stack_control_scaled_err.legendstyle = 'F'

                    draw([stack_control_scaled, stack_control_scaled_err],
                         xtitle=xlabel,
                         ytitle='MC/data',
                         pad=self.ratio_pad,
                         logy=False)

                    stack_control_scaled.xaxis.set_label_size(
                        stack_control_scaled.xaxis.get_label_size() * 3.
                    )  # the scale should match that of the main/ratio pad size ratio
                    stack_control_scaled.yaxis.set_label_size(
                        stack_control_scaled.yaxis.get_label_size() * 3.
                    )  # the scale should match that of the main/ratio pad size ratio
                    stack_control_scaled.xaxis.set_title_size(
                        stack_control_scaled.xaxis.get_title_size() * 3.
                    )  # the scale should match that of the main/ratio pad size ratio
                    stack_control_scaled.yaxis.set_title_size(
                        stack_control_scaled.yaxis.get_title_size() * 3.
                    )  # the scale should match that of the main/ratio pad size ratio
                    stack_control_scaled.yaxis.set_ndivisions(405)
                    stack_control_scaled.yaxis.set_title_offset(0.4)
                    stack_control_scaled.SetMinimum(0.)
                    stack_control_scaled.SetMaximum(1.5)

                    CMS_lumi(self.main_pad,
                             4,
                             0,
                             lumi_13TeV="%d, L = %.1f fb^{-1}" %
                             (self.year, self.lumi / 1000.))

                    self.canvas.cd()
                    # remove old legend
                    for iprim in self.canvas.primitives:
                        if isinstance(iprim, Legend):
                            self.canvas.primitives.remove(iprim)

                    # draw new legend
                    new_legend.Draw('same')

                    self.canvas.Modified()
                    self.canvas.Update()

                    for iformat in ['pdf', 'png', 'root']:
                        self.canvas.SaveAs('/'.join([
                            self.plt_dir, 'lnt_region',
                            'log' if islogy else 'lin',
                            iformat if iformat != 'pdf' else '',
                            '%s%s.%s' %
                            (label, '_log' if islogy else '_lin', iformat)
                        ]))

                    # compare shapes in tight and loose not tight

                    # data in tight
                    all_obs_prompt_norm = copy(all_obs_prompt)
                    if all_obs_prompt_norm.integral() != 0:
                        all_obs_prompt_norm.Scale(
                            np.nan_to_num(
                                np.divide(1., all_obs_prompt_norm.integral())))
                    #import pdb; pdb.set_trace()
                    all_obs_prompt_norm.drawstyle = 'hist e'
                    all_obs_prompt_norm.linecolor = 'black'
                    all_obs_prompt_norm.markersize = 0
                    all_obs_prompt_norm.legendstyle = 'LE'
                    all_obs_prompt_norm.title = ''
                    all_obs_prompt_norm.label = 'data - tight'

                    # data MC subtracted in loose
                    all_obs_prompt_mc_sub_norm = copy(all_obs_prompt)
                    all_obs_prompt_mc_sub_norm.add(all_exp_prompt, -1)
                    all_obs_prompt_mc_sub_norm.Scale(
                        np.nan_to_num(
                            np.divide(1.,
                                      all_obs_prompt_mc_sub_norm.integral())))
                    all_obs_prompt_mc_sub_norm.drawstyle = 'hist e'
                    all_obs_prompt_mc_sub_norm.linecolor = 'green'
                    all_obs_prompt_mc_sub_norm.markersize = 0
                    all_obs_prompt_mc_sub_norm.legendstyle = 'LE'
                    all_obs_prompt_mc_sub_norm.title = ''
                    all_obs_prompt_mc_sub_norm.label = '(data-MC) - tight'

                    # data in loose
                    all_obs_nonprompt_control_norm = copy(
                        all_obs_nonprompt_control)
                    all_obs_nonprompt_control_norm.Scale(
                        np.nan_to_num(
                            np.divide(
                                1.,
                                all_obs_nonprompt_control_norm.integral())))
                    all_obs_nonprompt_control_norm.drawstyle = 'hist e'
                    all_obs_nonprompt_control_norm.linecolor = 'red'
                    all_obs_nonprompt_control_norm.markersize = 0
                    all_obs_nonprompt_control_norm.legendstyle = 'LE'
                    all_obs_nonprompt_control_norm.title = ''
                    all_obs_nonprompt_control_norm.label = 'data - l-n-t'

                    # data MC subtracted in loose
                    all_obs_nonprompt_control_mc_sub_norm = copy(
                        all_obs_nonprompt_control)
                    all_obs_nonprompt_control_mc_sub_norm.add(
                        stack_control.sum, -1)
                    all_obs_nonprompt_control_mc_sub_norm.Scale(
                        np.nan_to_num(
                            np.divide(
                                1.,
                                all_obs_nonprompt_control_mc_sub_norm.integral(
                                ))))
                    all_obs_nonprompt_control_mc_sub_norm.drawstyle = 'hist e'
                    all_obs_nonprompt_control_mc_sub_norm.linecolor = 'blue'
                    all_obs_nonprompt_control_mc_sub_norm.markersize = 0
                    all_obs_nonprompt_control_mc_sub_norm.legendstyle = 'LE'
                    all_obs_nonprompt_control_mc_sub_norm.title = ''
                    all_obs_nonprompt_control_mc_sub_norm.label = '(data-MC) - l-n-t'

                    things_to_plot = [
                        all_obs_prompt_norm,
                        all_obs_prompt_mc_sub_norm,
                        all_obs_nonprompt_control_norm,
                        all_obs_nonprompt_control_mc_sub_norm,
                    ]

                    yaxis_max = max([ii.GetMaximum() for ii in things_to_plot])

                    draw(things_to_plot,
                         xtitle=xlabel,
                         ytitle=ylabel,
                         pad=self.main_pad,
                         logy=islogy,
                         ylimits=(yaxis_min, 1.55 * yaxis_max))

                    self.canvas.cd()
                    # remove old legend
                    for iprim in self.canvas.primitives:
                        if isinstance(iprim, Legend):
                            self.canvas.primitives.remove(iprim)

                    shape_legend = Legend([],
                                          pad=self.main_pad,
                                          leftmargin=0.,
                                          rightmargin=0.,
                                          topmargin=0.,
                                          textfont=42,
                                          textsize=0.03,
                                          entrysep=0.01,
                                          entryheight=0.04)
                    shape_legend.AddEntry(all_obs_prompt_norm,
                                          all_obs_prompt_norm.label,
                                          all_obs_prompt_norm.legendstyle)
                    shape_legend.AddEntry(
                        all_obs_prompt_mc_sub_norm,
                        all_obs_prompt_mc_sub_norm.label,
                        all_obs_prompt_mc_sub_norm.legendstyle)
                    shape_legend.AddEntry(
                        all_obs_nonprompt_control_norm,
                        all_obs_nonprompt_control_norm.label,
                        all_obs_nonprompt_control_norm.legendstyle)
                    shape_legend.AddEntry(
                        all_obs_nonprompt_control_mc_sub_norm,
                        all_obs_nonprompt_control_mc_sub_norm.label,
                        all_obs_nonprompt_control_mc_sub_norm.legendstyle)
                    shape_legend.SetBorderSize(0)
                    shape_legend.x1 = 0.50
                    shape_legend.y1 = 0.71
                    shape_legend.x2 = 0.88
                    shape_legend.y2 = 0.90
                    shape_legend.SetFillColor(0)
                    shape_legend.Draw('same')

                    # plot ratios
                    all_obs_prompt_norm_ratio = copy(all_obs_prompt_norm)
                    all_obs_prompt_mc_sub_norm_ratio = copy(
                        all_obs_prompt_mc_sub_norm)
                    all_obs_nonprompt_control_norm_ratio = copy(
                        all_obs_nonprompt_control_norm)
                    all_obs_nonprompt_control_mc_sub_norm_ratio = copy(
                        all_obs_nonprompt_control_mc_sub_norm)

                    all_obs_prompt_norm_ratio.Divide(
                        all_obs_prompt_mc_sub_norm_ratio)
                    all_obs_nonprompt_control_norm_ratio.Divide(
                        all_obs_prompt_mc_sub_norm_ratio)
                    all_obs_nonprompt_control_mc_sub_norm_ratio.Divide(
                        all_obs_prompt_mc_sub_norm_ratio)

                    things_to_plot_ratio = [
                        all_obs_prompt_norm_ratio,
                        all_obs_nonprompt_control_norm_ratio,
                        all_obs_nonprompt_control_mc_sub_norm_ratio,
                    ]

                    for ithing in things_to_plot_ratio:
                        ithing.xaxis.set_label_size(
                            ithing.xaxis.get_label_size() * 3.
                        )  # the scale should match that of the main/ratio pad size ratio
                        ithing.yaxis.set_label_size(
                            ithing.yaxis.get_label_size() * 3.
                        )  # the scale should match that of the main/ratio pad size ratio
                        ithing.xaxis.set_title_size(
                            ithing.xaxis.get_title_size() * 3.
                        )  # the scale should match that of the main/ratio pad size ratio
                        ithing.yaxis.set_title_size(
                            ithing.yaxis.get_title_size() * 3.
                        )  # the scale should match that of the main/ratio pad size ratio
                        ithing.yaxis.set_ndivisions(405)
                        ithing.yaxis.set_title_offset(0.4)
                        ithing.SetMinimum(0.)
                        ithing.SetMaximum(2.)

                    draw(things_to_plot_ratio,
                         xtitle=xlabel,
                         ytitle='1/(data-MC)_{tight}',
                         pad=self.ratio_pad,
                         logy=False,
                         ylimits=(0., 2.))
                    self.ratio_pad.cd()
                    line.Draw('same')

                    CMS_lumi(self.main_pad,
                             4,
                             0,
                             lumi_13TeV="%d, L = %.1f fb^{-1}" %
                             (self.year, self.lumi / 1000.))

                    self.canvas.Modified()
                    self.canvas.Update()

                    for iformat in ['pdf', 'png', 'root']:
                        self.canvas.SaveAs('/'.join([
                            self.plt_dir, 'shapes', 'log' if islogy else 'lin',
                            iformat if iformat != 'pdf' else '',
                            '%s%s.%s' %
                            (label, '_log' if islogy else '_lin', iformat)
                        ]))

            # save only the datacards you want, don't flood everything
            if len(self.datacards) and label not in self.datacards:
                continue

            # FIXME! allow it to save datacards even for non data driven bkgs
            if self.data_driven:
                self.create_datacards(data=all_obs_prompt,
                                      bkgs={
                                          'prompt': all_exp_prompt,
                                          'nonprompt': all_exp_nonprompt
                                      },
                                      signals=all_signals,
                                      label=label)
    histpad.cd()
    histpad.SetFrameBorderSize(2)
    histpad.SetFrameLineWidth(2);

    canvas.cd()
    #ratiopad    = Pad(leftmargins[1],0.00,1 - rightmargins[1],bottommargins[1]-0.02)
    ratiopad    = Pad(leftmargins[1],0.00,1 - rightmargins[1],bottommargins[1]-0.02)
    ratiopad.SetBottomMargin(0.33)
    ratiopad.SetTopMargin(0.03)
    ratiopad.SetFrameLineWidth(2);

    if(doRatio):
        ratiopad.Draw()
    histpad.cd()

    stack   = HistStack()
    datastack = HistStack() 
    
    yminimum, ymaximum = RegionYrange(Region)
    stack.SetMinimum(yminimum)
    stack.SetMaximum(ymaximum)

    total = Cut("1")
    for cut in Cutlist:
        if (variable in cut) and donminus1:continue 
        total = total & cut 


    backgroundstacks = OrderedDict()
    background = []
    for sample in sampleDictionary:
Beispiel #10
0
map(h2.Fill, signal)
h3.FillRandom('landau', 1000)
map(h3.Fill, signal_obs)

# set visual attributes
h1.fillstyle = 'solid'
h1.fillcolor = 'green'
h1.linecolor = 'green'
h1.linewidth = 0

h2.fillstyle = 'solid'
h2.fillcolor = 'red'
h2.linecolor = 'red'
h2.linewidth = 0

stack = HistStack([h1, h2], drawstyle='HIST E1 X0')

# plot with ROOT
canvas = Canvas(width=700, height=500)
draw([stack, h3], xtitle='Mass', ytitle='Events', pad=canvas)
# set the number of expected legend entries
legend = Legend([h1, h2, h3], leftmargin=0.45, margin=0.3)
legend.Draw()
label = ROOT.TText(0.3, 0.8, 'ROOT')
label.SetTextFont(43)
label.SetTextSize(25)
label.SetNDC()
label.Draw()
canvas.Modified()
canvas.Update()
Beispiel #11
0
          # rebin?
          rebin = plots_path.get('rebin', plots_config.get('rebin', None))
          if rebin is not None:
            hist.rebin(rebin)

          # exclusion, so we don't need to plot it
          if hist.title in plots_path.get('exclude', []): continue
          if group.get('stack it', False):
            # overwrite with solid when stacking
            hist.fillstyle = 'solid'
            stackHists.append(hist)
          else:
            soloHists.append(hist)


        hstack = HistStack(name=h.path)
        # for some reason, this causes noticable slowdowns
        hstack.drawstyle = 'hist'
        map(hstack.Add, stackHists[::-1])

        # this is where we would set various parameters of the min, max and so on?
        # need to set things like min, max, change to log, etc for hstack and soloHists
        normalizeTo = plots_path.get('normalize', plots_config.get('normalize', None))
        if normalizeTo is not None:
          dataScale = 0.
          if "unity" in normalizeTo:
            for hist in hstack:
              if (hist.integral() != 0):
                hist.scale(1.0/hist.integral())
            for hist in soloHists:
              if (hist.integral() != 0):
Beispiel #12
0
        samples = []
        backgrounds = []
        signals = []
        data = []
        stackyield = 0
        errorsum = 0

        backgroundstacks = {}
        # print "pre sampleDictionary"
        for sample in sampleDictionary:
            if sample != 'Zjets1516': continue
            if Region in BLINDEDLIST and sample is 'Data': continue
            LegendEntry = sampleDictionary[sample]['legend']
            Type = sampleDictionary[sample]['type']
            backgroundstacks[LegendEntry] = HistStack()
        #    print "post initialisation of backgroundstakcs"

        #"pre second sample dictionary"

        for (variable2, xmin2, xmax2, nbins2, xtitle2, ytitle2,
             variable_bool2) in zip(variable_list, xmin_list, xmax_list,
                                    nbins_list, xtitle_list, ytitle_list,
                                    variable_bool_list):
            if variable_bool != 'Plot': continue

            canvas = Canvas(width=canvaswidth,
                            height=int(
                                (1 - (1 - doRatio) * 0.2) * canvasheight))
            #canvas.SetFrameBorderMode(0)
Beispiel #13
0
                    # rebin?
                    rebin = plots_path.get('rebin',
                                           plots_config.get('rebin', None))
                    if rebin is not None:
                        hist.rebin(rebin)

                    # exclusion, so we don't need to plot it
                    if hist.title in plots_path.get('exclude', []): continue
                    if group.get('stack it', False):
                        # overwrite with solid when stacking
                        hist.fillstyle = 'solid'
                        stackHists.append(hist)
                    else:
                        soloHists.append(hist)

                hstack = HistStack(name=h.path)
                # for some reason, this causes noticable slowdowns
                hstack.drawstyle = 'hist'
                map(hstack.Add, stackHists[::-1])

                # this is where we would set various parameters of the min, max and so on?
                # need to set things like min, max, change to log, etc for hstack and soloHists
                normalizeTo = plots_path.get(
                    'normalize', plots_config.get('normalize', None))
                if normalizeTo is not None:
                    dataScale = 0.
                    if normalizeTo not in [hist.title for hist in soloHists]:
                        raise ValueError(
                            "Could not find %s as a solo hist for normalizing to."
                            % normalizeTo)
                    for hist in soloHists:
Beispiel #14
0
def stack(x, *args, **kwargs):

    ## parse arguments
    _data = kwargs.pop('data', None)
    _bkgs = kwargs.pop('bkgs', None)
    _sigs = kwargs.pop('sigs', None)
    _treename = kwargs.pop('treename', None)
    _datasearchpath = kwargs.pop('datasearchpath', None)
    _datadrivensearchpath = kwargs.pop('datadrivensearchpath', None)
    _bkgsearchpath = kwargs.pop('bkgsearchpath', None)
    _sigsearchpath = kwargs.pop('sigsearchpath', None)
    _lumi = kwargs.pop('lumi', None)

    global data
    global bkgs
    global sigs
    global treename
    global datasearchpath
    global datadrivensearchpath
    global bkgsearchpath
    global sigsearchpath
    global lumi

    data = _data or data
    bkgs = _bkgs or bkgs
    sigs = _sigs or sigs
    treename = _treename or treename
    datasearchpath = _datasearchpath or datasearchpath
    datadrivensearchpath = _datadrivensearchpath or datadrivensearchpath
    bkgsearchpath = _bkgsearchpath or bkgsearchpath
    sigsearchpath = _sigsearchpath or sigsearchpath
    if _lumi:
        lumi = float(_lumi)

    xtitle = kwargs.pop('xtitle', '')
    ytitle = kwargs.pop('ytitle', '')
    logx = bool(kwargs.pop('logx', False))
    logy = bool(kwargs.pop('logy', False))
    blind = kwargs.pop('blind', None)
    has_blinded_data = False

    ## save stuff to bookkeep and return
    stuff = dict()
    stuff['x'] = x

    ## get data histogram
    h_data = None
    if data:
        sp = datasearchpath  # HACK: just data to True!
        newx = '%s::%s::%s' % (sp, treename, x)
        h_data = ipyhep.tree.project(newx, *args, **kwargs)
        if h_data:
            stuff['h_data'] = h_data

    ## blind the data?
    if h_data and not blind is None:
        if isinstance(blind, tuple):
            blind1, blind2 = blind
            nbins = h_data.GetNbinsX()
            for i_bin in xrange(1, nbins +
                                2):  # skip underflow (but not overflow)
                xval1 = h_data.GetXaxis().GetBinLowEdge(i_bin)
                xval2 = h_data.GetXaxis().GetBinUpEdge(i_bin)
                if xval1 >= blind1 and xval2 <= blind2:
                    h_data.SetBinContent(i_bin, 0.0)
                    h_data.SetBinError(i_bin, 0.0)
                    has_blinded_data = True
        else:
            nbins = h_data.GetNbinsX()
            for i_bin in xrange(1, nbins +
                                2):  # skip underflow (but not overflow)
                xval = h_data.GetXaxis().GetBinLowEdge(i_bin)
                if xval >= blind:
                    h_data.SetBinContent(i_bin, 0.0)
                    h_data.SetBinError(i_bin, 0.0)
                    has_blinded_data = True

    ## get background histograms
    h_bkgs = list()
    n_bkgs = list()
    if bkgs:
        for bkg in bkgs:
            if isinstance(bkg, list):
                h_subtotal = None
                for dsid in bkg:
                    assert isinstance(dsid, str)
                    h_bkg = None
                    if dsid.isdigit():
                        ## mc backgrounds
                        sp = bkgsearchpath % int(dsid)
                        newx = '%s::%s::%s' % (sp, treename, x)
                        h_bkg = ipyhep.tree.project(newx, *args, **kwargs)
                    else:
                        ## data-driven backgrounds
                        assert dsid == 'fakes' or dsid == 'efakes'
                        sp = datadrivensearchpath % dsid
                        newx = '%s::%s::%s' % (sp, treename, x)
                        h_bkg = ipyhep.tree.project(newx, *args, **kwargs)
                    if h_bkg:
                        if h_subtotal:
                            h_subtotal.Add(h_bkg)
                        else:
                            h_subtotal = h_bkg.Clone()
                if h_subtotal:
                    h_bkgs.append(h_subtotal)
                    dsid = bkg[0]
                    n_bkgs.append(dsid)
            else:
                dsid = bkg
                assert isinstance(dsid, str)
                h_bkg = None
                if dsid.isdigit():
                    ## mc backgrounds
                    sp = bkgsearchpath % int(dsid)
                    newx = '%s::%s::%s' % (sp, treename, x)
                    h_bkg = ipyhep.tree.project(newx, *args, **kwargs)
                else:
                    ## data-driven backgrounds
                    assert dsid == 'fakes' or dsid == 'efakes'
                    sp = datadrivensearchpath % dsid
                    newx = '%s::%s::%s' % (sp, treename, x)
                    h_bkg = ipyhep.tree.project(newx, *args, **kwargs)
                if h_bkg:
                    h_bkgs.append(h_bkg)
                    n_bkgs.append(dsid)
        if h_bkgs:
            stuff['h_bkgs'] = h_bkgs

    ## get signal histograms
    h_sigs = list()
    n_sigs = list()
    if sigs:
        for dsid in sigs:
            sp = sigsearchpath % int(dsid)
            newx = '%s::%s::%s' % (sp, treename, x)
            h_sig = ipyhep.tree.project(newx, *args, **kwargs)
            if h_sig:
                h_sigs.append(h_sig)
                n_sigs.append(dsid)
        if h_sigs:
            stuff['h_sigs'] = h_sigs

    assert h_sigs

    ## style data
    if h_data:
        h_data.title = 'Data'
        h_data.linecolor = ipyhep.style.black
        h_data.linewidth = 2
        h_data.markercolor = ipyhep.style.black
        h_data.markerstyle = 20
        h_data.markersize = 1.2
        h_data.fillstyle = ipyhep.style.fill_hollow
        h_data.drawstyle = 'PE'
        h_data.legendstyle = 'LP'

    ## scale and style background histograms
    if h_bkgs:
        assert len(h_bkgs) == len(n_bkgs), '%s\n%s' % (h_bkgs, n_bkgs)

        for h, dsid in zip(h_bkgs, n_bkgs):
            sf = ipyhep.sampleops.get_sf(dsid)
            if dsid.isdigit():
                sf *= lumi / __ntuple_lumi
            h.Scale(sf)

            h.title = ipyhep.sampleops.get_label(dsid)
            h.linecolor = ipyhep.style.black
            h.linewidth = 1
            h.markercolor = ipyhep.sampleops.get_color(dsid)
            h.fillcolor = ipyhep.sampleops.get_color(dsid)
            h.fillstyle = ipyhep.style.fill_solid
            h.legendstyle = 'F'

    ## calculate stat error on total background
    h_bkg_total = None
    if h_bkgs:
        for h_bkg in h_bkgs:
            if h_bkg_total:
                h_bkg_total.Add(h_bkg)
            else:
                h_bkg_total = h_bkg.Clone()
        stuff['h_bkg_total'] = h_bkg_total

    ## style h_bkg_total
    if h_bkg_total:
        h_bkg_total.title = 'stat. uncert.'
        h_bkg_total.linecolor = ipyhep.style.black
        h_bkg_total.linewidth = 1
        h_bkg_total.markerstyle = 0
        h_bkg_total.fillcolor = ipyhep.style.dark_gray
        h_bkg_total.fillstyle = ipyhep.style.fill_lines
        h_bkg_total.drawstyle = 'E2'
        h_bkg_total.legendstyle = 'LF'

    ## scale and style signal histograms
    if h_sigs:
        assert len(h_sigs) == len(n_sigs)
        for h, dsid in zip(h_sigs, n_sigs):
            sf = ipyhep.sampleops.get_sf(dsid)
            sf *= lumi / __ntuple_lumi
            h.Scale(sf)

            h.title = ipyhep.sampleops.get_label(dsid)
            h.linecolor = ipyhep.sampleops.get_color(dsid)
            h.linewidth = 3
            h.fillstyle = ipyhep.style.fill_hollow
            h.markerstyle = 0
            h.drawstyle = 'HIST'
            h.legendstyle = 'L'

    ## build list of all_hists
    all_hists = list()
    main_hists = list()
    if h_data:
        all_hists.append(h_data)
        main_hists.append(h_data)
    if h_bkgs:
        all_hists.extend(h_bkgs)
        main_hists.extend(h_bkgs)
    if h_bkg_total:
        all_hists.append(h_bkg_total)
        main_hists.append(h_bkg_total)
    if h_sigs:
        all_hists.extend(h_sigs)

    ## get statistics
    if all_hists:
        stats_list = list()
        for h in all_hists:
            stats_list.extend(get_stats(h))
        html = convert_table_to_html(convert_stats_to_table(stats_list))
        stuff['html'] = html

    ## renormalize for bin widths
    bins = kwargs.pop('bins', None)
    if bins and isinstance(bins, list):
        for h in all_hists:
            renormalize_for_bin_widths(h, bins)

    ## stack background histograms
    if h_bkgs:
        assert len(h_bkgs) == len(n_bkgs), '%s\n%s' % (h_bkgs, n_bkgs)

        h_bkgs.reverse()
        n_bkgs.reverse()

        hstack = HistStack()
        for h in h_bkgs:
            hstack.Add(h)
        hstack.title = 'stack sum'
        hstack.drawstyle = 'HIST'
        stuff['stack'] = hstack

        h_bkgs.reverse()
        n_bkgs.reverse()

#    ## convert data to TGraphAsymmErrors
#    g_data = None
#    if h_data:
#        if __use_poissonize:
#            g_data = poissonize.GetPoissonizedGraph(h_data)
#        else:
#            g_data = ROOT.TGraphAsymmErrors()
#            i_g = 0
#            nbins = h_data.GetNbinsX()
#            for i_bin in xrange(1, nbins+1): # skip underflow/overflow
#                c = h_data.GetBinContent(i_bin)
#                e = h_data.GetBinError(i_bin)
#                if c != 0.0:
#                    g_data.SetPoint(i_g, h_data.GetBinCenter(i_bin), c)
#                    g_ratio.SetPointError(i_g,
#                            h_data.GetBinWidth(i_bin)/2.,
#                            h_data.GetBinWidth(i_bin)/2.,
#                            e,
#                            e)
#                i_g += 1

## build list of objects to draw
    objects = list()
    if h_bkgs:
        objects.append(stuff['stack'])
        objects.append(stuff['h_bkg_total'])
    if h_sigs:
        objects.extend(h_sigs)
    if h_data:
        objects.append(h_data)

    ## set xlimits and ylimits
    ypadding = 0.21
    logy_crop_value = 7e-3
    xmin, xmax, ymin, ymax = 0.0, 1.0, 0.0, 1.0
    if objects:
        xmin, xmax, ymin, ymax = get_limits(objects,
                                            logx=logx,
                                            logy=logy,
                                            ypadding=ypadding,
                                            logy_crop_value=logy_crop_value)
    if logy:
        ymin = 7e-3
    else:
        ymin = 0.0
    xlimits = (xmin, xmax)
    ylimits = (ymin, ymax)
    stuff['xlimits'] = xlimits
    stuff['ylimits'] = ylimits

    ## remove xtitle for do_ratio
    _xtitle = xtitle
    if h_data and h_bkg_total and kwargs.get('do_ratio'):
        _xtitle = ''

    ## make canvas
    canvas = Canvas(800, 600)
    stuff['canvas'] = canvas

    ## draw the objects
    if objects:
        canvas.cd()
        draw(objects,
             pad=canvas,
             xtitle=_xtitle,
             ytitle=ytitle,
             xlimits=xlimits,
             ylimits=ylimits)

    ## set log x/y, for some reason doesn't work before draw
    if logx or logy:
        if logx:
            canvas.SetLogx()
        if logy:
            canvas.SetLogy()
        canvas.Update()

    ## draw blind_line
    if has_blinded_data:
        if isinstance(blind, tuple):
            blind_list = list(blind)
        else:
            blind_list = [blind]
        blind_lines = list()
        for bl in blind_list:
            line_y1 = ymin
            line_y2 = ymax
            blind_line = ROOT.TLine(bl, line_y1, bl, line_y2)
            blind_line.SetLineColor(ROOT.kGray + 2)
            blind_line.SetLineStyle(7)
            blind_line.SetLineWidth(2)
            blind_line.Draw()
            blind_lines.append(blind_line)
        stuff['blind_lines'] = blind_lines
        canvas.Update()

    ## legend
    lefty = True
    if h_bkg_total:
        lefty = is_left_sided(h_bkg_total)
    elif h_data:
        lefty = is_left_sided(h_data)
    elif h_sigs:
        lefty = is_left_sided(h_sigs[0])

    if main_hists:
        header = '%.1f fb^{-1}, 13 TeV' % (lumi / 1000.0)
        if lefty:
            legend = Legend(main_hists,
                            pad=canvas,
                            header=header,
                            textsize=16,
                            topmargin=0.03,
                            leftmargin=0.60,
                            rightmargin=0.02,
                            entrysep=0.01,
                            entryheight=0.04)
        else:
            legend = Legend(main_hists,
                            pad=canvas,
                            header=header,
                            textsize=16,
                            topmargin=0.03,
                            leftmargin=0.03,
                            rightmargin=0.59,
                            entrysep=0.01,
                            entryheight=0.04)
        legend.Draw()
        stuff['legend'] = legend

    if h_sigs:
        #        header = 'ATLAS Internal'
        header = ''
        if lefty:
            legend2 = Legend(h_sigs,
                             pad=canvas,
                             header=header,
                             textsize=16,
                             topmargin=0.03,
                             leftmargin=0.37,
                             rightmargin=0.23,
                             entrysep=0.01,
                             entryheight=0.04)
        else:
            legend2 = Legend(h_sigs,
                             pad=canvas,
                             header=header,
                             textsize=16,
                             topmargin=0.03,
                             leftmargin=0.20,
                             rightmargin=0.40,
                             entrysep=0.01,
                             entryheight=0.04)
        legend2.Draw()
        stuff['legend2'] = legend2

    ## do_ratio
    if h_data and h_bkg_total and kwargs.get('do_ratio'):

        ## top canvas
        top_canvas = stuff.pop('canvas')
        stuff['top_canvas'] = top_canvas

        ## make SM/SM with error band: h_ratio_band
        i_sfratio = int(kwargs.get('sfratio', -1))
        if i_sfratio < 0:  # ratio plot of Data/Model
            h_ratio_band = h_bkg_total.Clone()
            nbins = h_ratio_band.GetNbinsX()
            for i_bin in xrange(nbins + 2):
                h_ratio_band.SetBinContent(i_bin, 1.0)
                c = h_bkg_total.GetBinContent(i_bin)
                e = h_bkg_total.GetBinError(i_bin) / c if c > 0.0 else 0.0
                h_ratio_band.SetBinError(i_bin, e)
            stuff['h_ratio_band'] = h_ratio_band
        else:  # ratio plot of Scale Factor for ith background
            hi = h_bkgs[i_sfratio]
            h_ratio_band = hi.Clone()
            nbins = h_ratio_band.GetNbinsX()
            for i_bin in xrange(nbins + 2):
                h_ratio_band.SetBinContent(i_bin, 1.0)
                c = hi.GetBinContent(i_bin)
                e = hi.GetBinError(i_bin) / c if c > 0.0 else 0.0
                h_ratio_band.SetBinError(i_bin, e)
            stuff['h_ratio_band'] = h_ratio_band

        ## make data/(SM) h_ratio
        if i_sfratio < 0:
            h_ratio = h_data.Clone()
            h_ratio.Divide(h_data, h_bkg_total, 1.0, 1.0)
            stuff['h_ratio'] = h_ratio
        else:
            ## SF1 = 1.0 + (data - MCtot) / MC1
            sfname = kwargs.get('sfname')
            sffile = kwargs.get('sffile')
            if not sfname:
                sfname = 'h_sf'
            hi = h_bkgs[i_sfratio]
            h_numer = h_data.Clone()
            h_numer.Add(h_bkg_total, -1.0)
            ## do the division
            h_ratio = h_data.Clone(sfname)
            h_ratio.Divide(h_numer, hi, 1.0, 1.0)
            ## add the 1.0
            nbins = h_ratio.GetNbinsX()
            for i_bin in xrange(nbins + 2):
                c = h_ratio.GetBinContent(i_bin)
                h_ratio.SetBinContent(i_bin, c + 1.0)
                h_ratio_band.SetBinContent(i_bin, c + 1.0)
            ## ignore bins with no data for SF
            for i_bin in xrange(nbins + 2):
                c = h_data.GetBinContent(i_bin)
                if c <= 0:
                    h_ratio.SetBinContent(i_bin, 0.0)
                    h_ratio.SetBinError(i_bin, 0.0)
                    h_ratio_band.SetBinError(i_bin, 0.0)
            stuff['h_ratio'] = h_ratio
            if sffile:
                f_out = ipyhep.file.write(h_ratio, sffile)
#                f_out.Close()

## convert ratio to a TGraphErrors so that Draw('E0')
## shows error bars for points off the pad
        g_ratio = ROOT.TGraphErrors()
        i_g = 0
        for i_bin in xrange(1, nbins + 1):  # skip underflow/overflow
            ratio_content = h_ratio.GetBinContent(i_bin)
            if ratio_content != 0.0:
                g_ratio.SetPoint(i_g, h_ratio.GetBinCenter(i_bin),
                                 ratio_content)
                g_ratio.SetPointError(i_g,
                                      h_ratio.GetBinWidth(i_bin) / 2.,
                                      h_ratio.GetBinError(i_bin))
                i_g += 1
            else:
                h_ratio.SetBinError(i_bin, 0.0)
        stuff['g_ratio'] = g_ratio

        ## style ratio
        h_ratio_band.title = 'bkg uncert.'
        if i_sfratio < 0:
            h_ratio_band.linecolor = ipyhep.style.yellow
        else:
            h_ratio_band.linecolor = ipyhep.style.light_gray
        h_ratio_band.linewidth = 0
        h_ratio_band.markerstyle = 0
        if i_sfratio < 0:
            h_ratio_band.fillcolor = ipyhep.style.yellow
        else:
            h_ratio_band.linecolor = ipyhep.style.light_gray
        h_ratio_band.fillstyle = ipyhep.style.fill_solid
        h_ratio_band.drawstyle = 'E2'
        h_ratio_band.legendstyle = 'F'
        h_ratio.title = 'ratio'
        h_ratio.linecolor = ipyhep.style.black
        h_ratio.linewidth = 2
        h_ratio.markercolor = ipyhep.style.black
        h_ratio.markerstyle = 20
        h_ratio.markersize = 1.2
        h_ratio.fillstyle = ipyhep.style.fill_hollow
        h_ratio.drawstyle = 'PE'
        h_ratio.legendstyle = 'LP'

        ## bottom canvas
        bottom_canvas = Canvas(800, 600)
        bottom_canvas.cd()
        stuff['bottom_canvas'] = bottom_canvas

        ## set ratio ylimits
        ratio_min = kwargs.get('ratio_min', -0.2)
        ratio_max = kwargs.get('ratio_max', 2.2)
        ratio_ylimits = (ratio_min, ratio_max)

        ## draw ratio band
        if i_sfratio < 0:
            _ytitle = 'Data / Model'
        else:
            hi = h_bkgs[i_sfratio]
            _ytitle = 'SF(%s)' % hi.title
        draw([h_ratio_band],
             pad=bottom_canvas,
             xtitle=xtitle,
             ytitle=_ytitle,
             xlimits=xlimits,
             ylimits=ratio_ylimits)

        ## set log x/y, for some reason doesn't work before draw?
        if logx:
            bottom_canvas.SetLogx()
            bottom_canvas.Update()

        ### make horiz lines in ratio plot every 0.5:
        line_ys = [
            y / 10.0
            for y in range(10 *
                           int(round(ratio_min)), 10 * int(round(ratio_max)) +
                           5, 5)
        ]
        line_x1 = canvas.GetUxmin()
        line_x2 = canvas.GetUxmax()
        line_xwidth = abs(line_x2 - line_x1)
        lines = []
        for line_y in line_ys:
            line = ROOT.TLine(line_x1 + 0.02 * line_xwidth, line_y,
                              line_x2 - 0.02 * line_xwidth, line_y)
            line.SetLineWidth(1)
            line.SetLineStyle(7)
            if line_y == 1.0:
                line.SetLineColor(ROOT.kGray + 2)
            else:
                line.SetLineColor(ROOT.kGray + 0)
            line.Draw()
            lines.append(line)
        stuff['lines'] = lines

        ## draw blind_line
        if has_blinded_data:
            if isinstance(blind, tuple):
                blind_list = list(blind)
            else:
                blind_list = [blind]
            blind_lines = list()
            for bl in blind_list:
                line_y1 = ymin
                line_y2 = ymax
                blind_line = ROOT.TLine(bl, line_y1, bl, line_y2)
                blind_line.SetLineColor(ROOT.kGray + 2)
                blind_line.SetLineStyle(7)
                blind_line.SetLineWidth(2)
                blind_line.Draw()
                blind_lines.append(blind_line)
            stuff['blind_lines2'] = blind_lines
            canvas.Update()

        ## draw ratio
        g_ratio.Draw('PE0')
        #        h_ratio.GetYaxis().SetRangeUser(ratio_min, ratio_max)
        #        h_ratio.Draw('PE,SAME')

        ## shared canvas
        shared_canvas = Canvas(800, 800)
        shared_plot = plot_shared_axis(top_canvas,
                                       bottom_canvas,
                                       canvas=shared_canvas,
                                       split=0.35,
                                       axissep=0.01)
        stuff['canvas'] = shared_canvas
        canvas = shared_canvas

    ## save figures
    save = kwargs.get('save')

    if save is None:  # NOTE: save can be False to skip saving
        save = ['pdf', 'png']

    if save:
        ipyhep.file.save_figures(canvas, x, save)

    global results
    results = stuff
    return stuff
Beispiel #15
0
def draw(
        name,
        category,
        data=None,
        data_info=None,
        model=None,
        model_colors=None,
        signal=None,
        signal_scale=1.,
        signal_on_top=False,
        signal_linestyles=None,
        signal_colors=None,
        show_signal_error=False,
        fill_signal=False,
        stack_signal=True,
        units=None,
        plot_label=None,
        ylabel='Events',
        blind=False,
        show_ratio=False,
        ratio_range=(0, 2),
        ratio_height=0.15,
        ratio_margin=0.06,
        output_formats=None,
        systematics=None,
        #systematics_components=None,
        integer=False,
        textsize=22,
        logy=False,
        logy_min=None,
        separate_legends=False,
        ypadding=None,
        legend_position='right',
        range=None,
        output_name=None,
        output_dir=PLOTS_DIR,
        arrow_values=None,
        overflow=True,
        show_pvalue=False,
        top_label=None,
        poisson_errors=True):

    if model is None and data is None and signal is None:
        # insufficient input
        raise ValueError("at least one of model, data, "
                         "or signal must be specified")

    if model is not None:
        if not isinstance(model, (list, tuple)):
            model = [model]
        if overflow:
            for hist in model:
                fold_overflow(hist)
    if signal is not None:
        if not isinstance(signal, (list, tuple)):
            signal = [signal]
        if overflow:
            for hist in signal:
                fold_overflow(hist)
    if data is not None and overflow:
        fold_overflow(data)

    # objects will be populated with all histograms in the main pad
    objects = []
    legends = []

    if show_ratio and (data is None or model is None):
        # cannot show the ratio if data or model was not specified
        show_ratio = False

    if ypadding is None:
        # select good defaults for log or linear scales
        if logy:
            ypadding = (.6, .0)
        else:
            ypadding = (.35, .0)

    template = data or model[0]
    xdivisions = min(template.nbins(), 7) if integer else 507

    if show_ratio:
        fig = RatioPlot(logy=logy,
                        ratio_title='Data / Model',
                        ratio_limits=(0, 2),
                        offset=-72,
                        ratio_margin=22,
                        prune_ratio_ticks=True)
    else:
        fig = SimplePlot(logy=logy)

    if signal is not None:
        if signal_scale != 1.:
            scaled_signal = []
            for sig in signal:
                scaled_h = sig * signal_scale
                scaled_h.SetTitle(r'%g #times %s' %
                                  (signal_scale, sig.GetTitle()))
                scaled_signal.append(scaled_h)
        else:
            scaled_signal = signal
        if signal_colors is not None:
            set_colors(scaled_signal, signal_colors)
        for i, s in enumerate(scaled_signal):
            s.drawstyle = 'HIST'
            if fill_signal:
                s.fillstyle = 'solid'
                s.fillcolor = s.linecolor
                s.linewidth = 0
                s.linestyle = 'solid'
                alpha = .75
            else:
                s.fillstyle = 'hollow'
                s.linewidth = 3
                if signal_linestyles is not None:
                    s.linestyle = signal_linestyles[i]
                else:
                    s.linestyle = 'solid'
                alpha = 1.

    if model is not None:
        if model_colors is not None:
            set_colors(model, model_colors)
        # create the model stack
        model_stack = HistStack()
        for hist in model:
            hist.SetLineWidth(0)
            hist.drawstyle = 'hist'
            model_stack.Add(hist)
        if signal is not None and signal_on_top:
            for s in scaled_signal:
                model_stack.Add(s)
        objects.append(model_stack)

    if signal is not None and not signal_on_top:
        if stack_signal:
            # create the signal stack
            signal_stack = HistStack()
            for hist in scaled_signal:
                signal_stack.Add(hist)
            objects.append(signal_stack)
        else:
            objects.extend(scaled_signal)

    if model is not None:
        # draw uncertainty band
        total_model, high_band_model, low_band_model = uncertainty_band(
            model, systematics)  #, systematics_components)
        high = total_model + high_band_model
        low = total_model - low_band_model
        error_band_model = rootpy_utils.get_band(low,
                                                 high,
                                                 middle_hist=total_model)
        error_band_model.fillstyle = '/'
        error_band_model.fillcolor = 13
        error_band_model.linecolor = 10
        error_band_model.markersize = 0
        error_band_model.markercolor = 10
        error_band_model.drawstyle = 'e2'
        objects.append(error_band_model)

    if signal is not None and show_signal_error:
        total_signal, high_band_signal, low_band_signal = uncertainty_band(
            signal, systematics)  #, systematics_components)
        high = (total_signal + high_band_signal) * signal_scale
        low = (total_signal - low_band_signal) * signal_scale
        if signal_on_top:
            high += total_model
            low += total_model
        error_band_signal = rootpy_utils.get_band(low,
                                                  high,
                                                  middle_hist=total_signal *
                                                  signal_scale)
        error_band_signal.fillstyle = '\\'
        error_band_signal.fillcolor = 13
        error_band_signal.linecolor = 10
        error_band_signal.markersize = 0
        error_band_signal.markercolor = 10
        error_band_signal.drawstyle = 'e2'
        objects.append(error_band_signal)

    if data is not None and blind is not True:
        # create the data histogram
        if isinstance(blind, tuple):
            low, high = blind
            # zero out bins in blind region
            for bin in data.bins():
                if (low < bin.x.high <= high or low <= bin.x.low < high):
                    data[bin.idx] = (0., 0.)

        if poisson_errors:
            # convert data to TGraphAsymmErrors with Poisson errors
            data_poisson = data.poisson_errors()
            data_poisson.markersize = 1.2
            data_poisson.drawstyle = 'PZ'
            objects.append(data_poisson)
        else:
            # Gaussian errors
            data.drawstyle = 'E0'
            objects.append(data)

        # draw ratio plot
        if model is not None and show_ratio:
            fig.cd('ratio')
            total_model = sum(model)
            ratio_hist = Hist.divide(data, total_model)
            # remove bins where data is zero
            max_dev = 0
            for bin in data.bins():
                if bin.value <= 0:
                    ratio_hist[bin.idx] = (-100, 0)
                else:
                    ratio_value = ratio_hist[bin.idx].value
                    dev = abs(ratio_value - 1)
                    if dev > max_dev:
                        max_dev = dev

            if max_dev < 0.2:
                ratio_range = (0.8, 1.2)
            elif max_dev < 0.4:
                ratio_range = (0.6, 1.4)
            ruler_high = (ratio_range[1] + 1.) / 2.
            ruler_low = (ratio_range[0] + 1.) / 2.

            ratio_hist.linecolor = 'black'
            ratio_hist.linewidth = 2
            ratio_hist.fillstyle = 'hollow'
            ratio_hist.drawstyle = 'E0'
            """
            # draw empty copy of ratio_hist first so lines will show
            ratio_hist_tmp = ratio_hist.Clone()
            ratio_hist_tmp.Reset()
            ratio_hist_tmp.Draw()
            ratio_hist_tmp.yaxis.SetLimits(*ratio_range)
            ratio_hist_tmp.yaxis.SetRangeUser(*ratio_range)
            ratio_hist_tmp.yaxis.SetTitle('Data / Model')
            ratio_hist_tmp.yaxis.SetNdivisions(4)
            # not certain why the following is needed
            ratio_hist_tmp.yaxis.SetTitleOffset(style.GetTitleYOffset())

            ratio_xrange = range or ratio_hist.bounds()

            ratio_hist_tmp.xaxis.SetLimits(*ratio_xrange)
            #ratio_hist_tmp.xaxis.SetRangeUser(*ratio_xrange)

            ratio_hist_tmp.xaxis.SetTickLength(
                ratio_hist_tmp.xaxis.GetTickLength() * 2)

            # draw ratio=1 line
            line = Line(ratio_xrange[0], 1,
                        ratio_xrange[1], 1)
            line.linestyle = 'dashed'
            line.linewidth = 2
            line.Draw()

            # draw high ratio line
            line_up = Line(ratio_xrange[0], ruler_high,
                           ratio_xrange[1], ruler_high)
            line_up.linestyle = 'dashed'
            line_up.linewidth = 2
            line_up.Draw()

            # draw low ratio line
            line_dn = Line(ratio_xrange[0], ruler_low,
                           ratio_xrange[1], ruler_low)
            line_dn.linestyle = 'dashed'
            line_dn.linewidth = 2
            line_dn.Draw()
            """

            # draw band below points on ratio plot
            ratio_hist_high = Hist.divide(total_model + high_band_model,
                                          total_model)
            ratio_hist_low = Hist.divide(total_model - low_band_model,
                                         total_model)
            fig.cd('ratio')
            error_band = rootpy_utils.get_band(ratio_hist_high, ratio_hist_low)
            error_band.fillstyle = '/'
            error_band.fillcolor = '#858585'
            error_band.drawstyle = 'E2'
            fig.draw('ratio', [error_band, ratio_hist], xdivisions=xdivisions)

    if separate_legends:
        fig.cd('main')
        right_legend = Legend(len(signal) + 1 if signal is not None else 1,
                              pad=fig.pad('main'),
                              **legend_params('right', textsize))
        right_legend.AddEntry(data, style='lep')
        if signal is not None:
            for s in reversed(scaled_signal):
                right_legend.AddEntry(s, style='F' if fill_signal else 'L')
        legends.append(right_legend)
        if model is not None:
            n_entries = len(model)
            if systematics:
                n_entries += 1
            model_legend = Legend(n_entries,
                                  pad=fig.pad('main'),
                                  **legend_params('left', textsize))
            for hist in reversed(model):
                model_legend.AddEntry(hist, style='F')
            if systematics:
                model_err_band = error_band_model.Clone()
                model_err_band.linewidth = 0
                model_err_band.linecolor = 'white'
                model_err_band.fillcolor = '#858585'
                model_err_band.title = 'Uncert.'
                model_legend.AddEntry(model_err_band, style='F')
            legends.append(model_legend)
    else:
        n_entries = 1
        if signal is not None:
            n_entries += len(scaled_signal)
        if model is not None:
            n_entries += len(model)
            if systematics:
                n_entries += 1
        fig.cd('main')
        legend = Legend(n_entries,
                        pad=fig.pad('main'),
                        **legend_params(legend_position, textsize))
        if data is not None:
            legend.AddEntry(data, style='lep')
        if signal is not None:
            for s in reversed(scaled_signal):
                legend.AddEntry(s, style='F' if fill_signal else 'L')
        if model:
            for hist in reversed(model):
                legend.AddEntry(hist, style='F')
            model_err_band = error_band_model.Clone()
            model_err_band.linewidth = 0
            model_err_band.linecolor = 'white'
            model_err_band.fillcolor = '#858585'
            model_err_band.title = 'Uncert.'
            legend.AddEntry(model_err_band, style='F')
        legends.append(legend)

    # draw the objects
    bounds = fig.draw('main',
                      objects,
                      ypadding=ypadding,
                      logy_crop_value=1E-1,
                      xdivisions=xdivisions)
    xaxis, yaxis = fig.axes('main')

    base_xaxis = xaxis
    base_xaxis.range_user = template.bounds()
    base_xaxis.limits = template.bounds()
    xmin, xmax, ymin, ymax = bounds

    if show_ratio:
        base_xaxis = fig.axes('ratio')[0]
        base_xaxis.range_user = template.bounds()
        base_xaxis.limits = template.bounds()

    # draw the legends
    fig.cd('main')
    for legend in legends:
        legend.Draw()

    label_plot(
        fig.pad('main'),
        template=template,
        xaxis=base_xaxis,
        yaxis=yaxis,
        xlabel=name,
        ylabel=ylabel,
        units=units,
        category_label=category.label,
        extra_label=plot_label,
        extra_label_position='right' if legend_position == 'left' else 'left',
        data_info=data_info)

    if logy and logy_min is not None:
        yaxis.min = logy_min
        ymin = logy_min

    # draw arrows
    if arrow_values is not None:
        arrow_top = ymin + (ymax - ymin) / 2.
        fig.cd('main')
        for value in arrow_values:
            arrow = Arrow(value, arrow_top, value, ymin, 0.05, '|>')
            arrow.SetAngle(30)
            arrow.SetLineWidth(2)
            arrow.Draw()

    if show_pvalue and data is not None and model:
        fig.cd('main')
        total_model = sum(model)
        # show p-value and chi^2
        pvalue = total_model.Chi2Test(data, 'WW')
        pvalue_label = ROOT.TLatex(0.6, 0.97, "p-value={0:.2f}".format(pvalue))
        pvalue_label.SetNDC(True)
        pvalue_label.SetTextFont(43)
        pvalue_label.SetTextSize(16)
        pvalue_label.Draw()
        chi2 = total_model.Chi2Test(data, 'WW CHI2/NDF')
        chi2_label = ROOT.TLatex(0.78, 0.97,
                                 "#chi^{{2}}/ndf={0:.2f}".format(chi2))
        chi2_label.SetNDC(True)
        chi2_label.SetTextFont(43)
        chi2_label.SetTextSize(16)
        chi2_label.Draw()

    if top_label is not None:
        fig.cd('main')
        label = ROOT.TLatex(
            fig.pad('main').GetLeftMargin() + 0.08, 0.97, top_label)
        label.SetNDC(True)
        label.SetTextFont(43)
        label.SetTextSize(16)
        label.Draw()

    if output_name is not None:
        # create the output filename
        filename = 'var_{0}_{1}'.format(category.name,
                                        output_name.lower().replace(' ', '_'))
        if logy:
            filename += '_logy'
        filename += '_root'
        if output_formats is None:
            output_formats = ('png', )
        # save the figure
        save_canvas(fig, output_dir, filename, formats=output_formats)

    return fig
Beispiel #16
0
def plotHistos(histos, text="", option="", statbox=True):
    """ 
    Plots a list of histograms
    """
    log = logging.getLogger('pyroplot')
    import rootpy
    from rootpy.plotting import Hist, HistStack
    from ROOT import kRed, gPad, TPaveStats
    # only for 1D and 2D type histograms:
    # determine and set maximum for all histograms
    stack = HistStack()
    for hist in histos:
        if not hist.__class__.__name__ == "Hist3D":
            stack.Add(hist)
    # determine maximum value and set it for all histograms
    if stack.GetHists():
        maxplot = stack.GetMaximum()
        minplot = stack.GetMinimum()
        for hist in stack.GetHists():
            hist.SetMaximum(maxplot)
            # special treatment for log scale Y
            if gPad.GetLogy():
                hist.SetMinimum(1.)
            else:
                # if histogram minimum is positive, set to 0.
                if minplot > 0:
                    hist.SetMinimum(0.)
    for idx, hist in enumerate(histos):
        try:
            thisopt = option
            # if not first histo, add "same" to options so previous ones are not overwritten
            if idx:
                if statbox:
                    thisopt += "sames"
                else:
                    thisopt += "same"
            histcopy = hist.DrawCopy(thisopt)
            # on first go: print identifying text on pad
            if not idx and text:
                markPad(text=text, x=.14, y=.85, size=0.041)
                if statbox:
                    thisopt += "sames"
                else:
                    thisopt += "same"
                histcopy.Draw(thisopt)
            gPad.Update()
            if statbox:
                try:
                    statBox = histcopy.GetListOfFunctions().FindObject("stats")
                    # offset from last statbox
                    offset = .18
                    # needs to be larger for Profile & 2D histos
                    if (hist.__class__.__name__ == "Profile"
                            or hist.__class__.__name__ == "Hist2D"
                            or hist.__class__.__name__ == "Profile2D"
                            or hist.__class__.__name__ == "Hist3D"):
                        offset = .26
                    statBox.SetY1NDC(statBox.GetY1NDC() - offset * (idx))
                    statBox.SetY2NDC(statBox.GetY2NDC() - offset * (idx))
                    statBox.SetTextColor(hist.GetLineColor())
                    statBox.SetBorderSize(2)
                except AttributeError:
                    log.debug("Could not get statbox for histogram " +
                              hist.GetName())

        except rootpy.ROOTError, e:
            log.error(
                "Drawing histogram %s caused ROOTError exception ('%s')" %
                (hist.GetName(), e.msg))
            gPad.Clear()  # otherwise this happens again when drawing..
            markPad(text="Could not draw %s ('%s')" % (hist.GetName(), e.msg),
                    x=0.14,
                    y=0.4,
                    size=0.023,
                    color=kRed)
            return  # give up!