from rexart.objects import Sample, Histogram, Template
from rexart.utils import draw_ratio_with_line, draw_atlas_label, set_labels, shrink_pdf
import rexart.constants
from pathlib import PosixPath
import matplotlib.pyplot as plt
import numpy as np
import re
import yaml
import uproot
import logging
log = logging.getLogger(__name__)
[docs]def stackem(args, region, data, histograms, template_var, band=None, figsize=(6, 5.25)):
"""Create a stack plot
Parameters
----------
args : argparse.ArgumentParser
command line arguments
region : str
region from TRExFitter
data : rexart.objects.Sample
data sample
histograms : List[rexart.objects.Sample]
list of MC samples to stack
template_var : str
name of the template variable
band : Optional[uproot_methods.classes.TGraphAsymmErrors]
error band
figsize : tuple(float, float)
matplotlib figure size
Returns
-------
fig : matplotlib.figure.Figure
matplotlib fiture
axes : tuple(matplotlib.axes.Axes, matplotlib.axes.Axes)
matplotlib axes (main, ratio)
"""
# fmt: off
fig, (ax, axr) = plt.subplots(2, 1, sharex=True, figsize=figsize,
gridspec_kw=dict(height_ratios=[3.25, 1], hspace=0.025))
expected_sum = np.sum([h.content for h in histograms], axis=0)
expected_err = np.sqrt(np.sum([h.sumw2 for h in histograms], axis=0))
ax.hist([h.bin_centers for h in histograms], weights=[h.content for h in histograms],
bins=histograms[0].bins, histtype="stepfilled", stacked=True,
color=[h.sample.color for h in histograms], label=[h.sample.tex for h in histograms])
ax.errorbar(data.bin_centers, data.content, yerr=data.error, fmt="ko", label="Data", zorder=999)
set_labels(ax, histograms[0])
if rexart.constants.region_meta[template_var].logy:
ax.set_yscale("log")
ax.set_ylim([ax.get_ylim()[0], ax.get_ylim()[1] * 1000])
else:
ax.set_ylim([ax.get_ylim()[0], ax.get_ylim()[1] * 1.5])
if band is not None:
yerrlo = np.hstack([band.yerrorslow, band.yerrorslow[-1]])
yerrhi = np.hstack([band.yerrorshigh, band.yerrorshigh[-1]])
expected_sum4band = np.hstack([expected_sum, expected_sum[-1]])
if args.band_style == "hatch":
ax.fill_between(x=data.bins,y1=(expected_sum4band - yerrlo), y2=(expected_sum4band + yerrhi),
step="post", facecolor="none", hatch="////", edgecolor="cornflowerblue",
linewidth=0.0, zorder=50, label="Syst. Unc.")
axr.fill_between(x=data.bins, y1=1 - yerrlo / expected_sum4band, y2=1 + yerrhi / expected_sum4band,
step="post", facecolor="none", hatch="////", edgecolor="cornflowerblue",
linewidth=0.0, zorder=50, label="Syst. Unc.")
elif args.band_style == "shade":
axr.fill_between(x=data.bins, y1=1 - yerrlo / expected_sum4band, y2=1 + yerrhi / expected_sum4band,
step="post", facecolor=(0, 0, 0, 0.33), linewidth=0.0, zorder=50, label="Syst. Unc.")
# fmt: on
draw_ratio_with_line(axr, data, expected_sum, expected_err)
axr.set_ylim([0.8, 1.2])
axr.set_yticks([0.9, 1.0, 1.1])
axis_title = "{}{}".format(data.mpl_title, "" if data.unit == "" else f" [{data.unit}]")
axr.set_xlabel(axis_title, horizontalalignment="right", x=1.0)
if band is not None and args.band_style == "shade":
axr.legend(loc="upper center", frameon=True, fontsize=8)
ax.legend(loc="upper right")
handles, labels = ax.get_legend_handles_labels()
handles.insert(0, handles.pop())
labels.insert(0, labels.pop())
ax.legend(handles, labels, loc="upper right", ncol=args.legend_ncol)
raw_region = region.split("reg")[-1].split("_")[0]
extra_line1 = f"$\\sqrt{{s}}$ = 13 TeV, $L = {args.lumi}$ fb$^{{-1}}$"
extra_line2 = f"$pp\\rightarrow tW \\rightarrow e^{{\\pm}}\\mu^{{\\mp}} ({raw_region})$"
extra_line3 = "Pre-fit"
if histograms[0].postfit:
extra_line3 = "Post-fit"
draw_atlas_label(ax, extra_lines=[extra_line1, extra_line2, extra_line3])
fig.subplots_adjust(left=0.115, bottom=0.115, right=0.965, top=0.95)
return fig, (ax, axr)
[docs]def prefit_histograms(args, fit_name, region, samples):
"""Prepare prefit histogram objects
Parameters
---------
args : argparse.ArgumentParser
command line arguments
fit_name : str
TRExFitter fit name
region : str
TRExFitter region
samples : List[rexart.objects.Sample]
list of MC samples
Returns
-------
data : rexart.objects.Histogram
data histogram
histograms : List[rexart.objects.Histogram]
MC histograms
band : uproot_methods.classes.TGraphAsymmErrors
uncertainty band
"""
hfile = f"{args.workspace}/Histograms/{fit_name}_{region}_histos.root"
bfile = f"{args.workspace}/Histograms/{fit_name}_{region}_preFit.root"
band = uproot.open(bfile).get("g_totErr")
histograms = [Histogram(hfile, region, sample) for sample in samples]
data = Histogram(hfile, region, Sample("Data", "Data", "", "Data"))
return data, histograms, band
[docs]def postfit_histograms(args, fit_name, region, samples):
"""Prepare postfit histogram objects
Parameters
---------
args : argparse.ArgumentParser
command line arguments
fit_name : str
TRExFitter fit name
region : str
TRExFitter region
samples : List[rexart.objects.Sample]
list of MC samples
Returns
-------
data : rexart.objects.Histogram
data histogram
histograms : List[rexart.objects.Histogram]
MC histograms
band : uproot_methods.classes.TGraphAsymmErrors
uncertainty band
"""
hfile = f"{args.workspace}/Histograms/{region}_postFit.root"
bfile = f"{args.workspace}/Histograms/{region}_postFit.root"
band = uproot.open(bfile).get("g_totErr_postFit")
histograms = [Histogram(hfile, region, sample, postfit=True) for sample in samples]
return histograms, band
def split_region_str(region):
splits = region.split("_")
if len(splits) == 1:
return (region, "bdt_response")
else:
return (splits[0], "_".join(splits[1:]))
[docs]def run_stacks(args):
"""Given command line arguments generate stack plots
Parameters
----------
args : argparse.ArgumentParser
"""
samples = rexart.constants.samples
if args.out_dir is None:
outd = f"{args.workspace}/MPL"
else:
outd = args.out_dir
if outd != ".":
PosixPath(outd).mkdir(parents=True, exist_ok=True)
fit_name = PosixPath(args.workspace).stem
hfiledir = PosixPath(f"{fit_name}/Histograms")
regions = []
if args.skip_regions is not None:
skipregex = re.compile(args.skip_regions)
else:
skipregex = None
for hfile in hfiledir.iterdir():
if "_histos.root" in hfile.name:
region = hfile.name.split("_histos.root")[0].split(f"{fit_name}_")[-1]
if skipregex:
if re.search(skipregex, region):
continue
regions.append(region)
for region in regions:
raw_region, template_variable = split_region_str(region)
data, histograms, band = prefit_histograms(args, fit_name, region, samples)
data.unit = rexart.constants.region_meta[template_variable].unit
data.mpl_title = rexart.constants.region_meta[template_variable].title
fig, (ax, axr) = stackem(args, region, data, histograms, template_variable, band=band)
out_name = f"{outd}/preFit_{region}.pdf"
fig.savefig(out_name)
plt.close(fig)
if args.shrink:
shrink_pdf(out_name)
log.info(f"Done with {region} prefit")
if args.do_postfit:
histograms, band = postfit_histograms(args, fit_name, region, samples)
fig, (ax, axr) = stackem(args, region, data, histograms, template_variable, band=band)
axr.set_ylim([0.925, 1.075])
axr.set_yticks([0.95, 1.0, 1.05])
out_name = f"{outd}/postFit_{region}.pdf"
fig.savefig(out_name)
plt.close(fig)
if args.shrink:
shrink_pdf(out_name)
log.info(f"Done with {region} postfit")