import numpy as np
import operator as op
from collections import defaultdict
from analyzer.utils.structure_tools import commonDict
import functools as ft
import matplotlib as mpl
import matplotlib.pyplot as plt
import mplhep
from analyzer.postprocessing.style import Styler
# from ..grouping import doFormatting
from .annotations import labelAxis
from .common import PlotConfiguration
from .utils import saveFig, saveFigVariants, scaleYAxis, addLegend
[docs]
def getRatioAndUnc(num, den, uncertainty_type="poisson-ratio"):
import hist.intervals as hinter
with np.errstate(divide="ignore", invalid="ignore"):
ratios = num / den
unc = hinter.ratio_uncertainty(
num=num, denom=den, uncertainty_type=uncertainty_type
)
return ratios, unc
[docs]
def plotOne(
histograms,
stacked_hists,
common_metadata,
output_path,
style_set,
scale="linear",
normalize=False,
plot_configuration=None,
show_stacked_unc=True,
):
stacked_hists = stacked_hists or []
pc = plot_configuration or PlotConfiguration()
styler = Styler(style_set)
fig, ax = plt.subplots()
h = None
if stacked_hists:
stacked_hists = sorted(
stacked_hists, key=lambda x: x.item.histogram.sum().value
)
style_kwargs = defaultdict(list)
hists = []
titles = []
for item, meta in stacked_hists:
hists.append(item.histogram)
titles.append(meta.get("title") or meta["dataset_title"])
style = styler.getStyle(meta)
for k, v in style.get(plottype="fill").items():
style_kwargs[k].append(v)
style_kwargs["histtype"] = style_kwargs["histtype"][0]
mplhep.histplot(
hists,
ax=ax,
stack=True,
**style_kwargs,
label=titles, # sort="yield"
)
if show_stacked_unc:
stacked_total = ft.reduce(op.add, [x.item.histogram for x in stacked_hists])
mplhep.histplot(
stacked_total,
ax=ax,
label="Stacked Unc.",
histtype="band",
)
for item, meta in histograms:
title = meta.get("title") or meta["dataset_title"]
h = item.histogram
style = styler.getStyle(meta)
h.plot1d(
ax=ax,
label=title,
density=normalize,
yerr=style.yerr,
flow="none",
**style.get(),
)
if h is None:
h = stacked_hists[0]
labelAxis(ax, "y", h.axes, label=pc.y_label)
labelAxis(ax, "x", h.axes, label=pc.x_label)
all_meta = [x.metadata for x in histograms] + [x.metadata for x in stacked_hists]
ax.set_yscale(scale)
addLegend(ax, pc)
scaleYAxis(ax)
# mplhep.yscale_anchored_text(ax, soft_fail=True)
if style.y_min:
ax.set_ylim(bottom=style.y_min)
saveFigVariants(
fig,
ax,
output_path,
all_meta,
plot_configuration=pc,
metadata=common_metadata,
extra_text=f"{common_metadata.get('pipeline', '')}",
)
plt.close(fig)
[docs]
def makeStrHist(data, ax_name):
import hist
ax = hist.axis.StrCategory([x[0] for x in data], name=ax_name)
h = hist.Hist(ax, storage="double")
h[:] = np.array([x[1] for x in data])
return h
[docs]
def plotDictAsBars(
items,
common_meta,
output_path,
getter,
style_set,
ax_name=None,
normalize=False,
scale="linear",
plot_configuration=None,
):
pc = plot_configuration or PlotConfiguration()
styler = Styler(style_set)
mpl.use("Agg")
fig, ax = plt.subplots(layout="constrained")
for item, meta in items:
title = meta.get("title") or meta["dataset_title"]
flow = getter(item)
style = styler.getStyle(meta)
h = makeStrHist([(x, y) for x, y in flow.items()], ax_name=ax_name)
h.plot1d(
ax=ax,
label=title,
density=normalize,
**style.get(),
)
ax.legend()
labelAxis(ax, "y", h.axes)
labelAxis(ax, "x", h.axes)
ax.tick_params(axis="x", rotation=90)
all_meta = [x.metadata for x in items]
ax.set_yscale(scale)
addLegend(ax, pc)
mplhep.sort_legend(ax=ax)
scaleYAxis(ax)
saveFigVariants(
fig,
ax,
output_path,
all_meta,
plot_configuration=pc,
metadata=common_meta,
)
plt.close(fig)
[docs]
def makeRatioAxes(ratio_height):
gs_kw = dict(height_ratios=[1, ratio_height])
fig, (ax, ratio_ax) = plt.subplots(2, 1, sharex=True, gridspec_kw=gs_kw)
return fig, ax, ratio_ax
[docs]
def computeRatio(n, d, normalize=False, ratio_type="poisson"):
ratio, unc = getRatioAndUnc(n, d, uncertainty_type=ratio_type)
if normalize:
with np.errstate(divide="ignore", invalid="ignore"):
ratio = (n / np.sum(n)) / (d / np.sum(d))
ratio[ratio == 0] = np.nan
ratio[np.isinf(ratio)] = np.nan
return ratio, unc
[docs]
def computeSignificance(n, d, normalize=False, ratio_type="poisson"):
with np.errstate(divide="ignore", invalid="ignore"):
significance = n / np.sqrt(d)
return significance, None
[docs]
def plotRatioErrorBars(ratio_ax, x_values, ratio, unc, style):
opts = {
**style.get("errorbar", include_type=False),
"linestyle": "none",
}
ratio_ax.errorbar(x_values, ratio, yerr=unc, **opts)
[docs]
def plotStackedDenominators(ax, denominators, styler, normalize=False):
den_to_plot = sorted(denominators, key=lambda x: x.item.histogram.sum().value)
hists = []
titles = []
style_kwargs = defaultdict(list)
for item, meta in den_to_plot:
hists.append(item.histogram)
titles.append(meta.get("title") or meta["dataset_title"])
style = styler.getStyle(meta)
for key, value in style.get().items():
style_kwargs[key].append(value)
style_kwargs["histtype"] = style_kwargs["histtype"][0]
mplhep.histplot(
hists,
ax=ax,
stack=True,
density=normalize,
label=titles,
**style_kwargs,
)
den_total = ft.reduce(op.add, (x.item.histogram for x in denominators))
mplhep.histplot(
den_total,
ax=ax,
label="Den. Stat. Unc.",
histtype="band",
)
return den_total
[docs]
def plotUnstackedDenominators(ax, denominators, styler, *, normalize):
den_to_plot = sorted(
denominators, key=lambda x: x.metadata.get("title", "").lower()
)
den_styles = []
for item, meta in den_to_plot:
style = styler.getStyle(meta)
den_styles.append(style)
item.histogram.plot1d(
ax=ax,
label=meta.get("title") or meta["dataset_title"],
density=normalize,
yerr=style.yerr,
**style.get(),
)
return den_to_plot, den_styles
[docs]
def plotMultiNumerators(
ax,
ratio_ax,
numerators,
den_total,
styler,
normalize,
ratio_type,
x_values,
ratio_func=computeRatio,
show_den_unc=True,
):
for item, meta in numerators:
hist = item.histogram
style = styler.getStyle(meta)
n_vals = hist.values()
d_vals = den_total.values()
ratio, unc = ratio_func(
n_vals,
d_vals,
normalize=normalize,
ratio_type=ratio_type,
)
hist.plot1d(
ax=ax,
label=meta.get("title") or meta["dataset_title"],
density=normalize,
yerr=True,
**style.get(),
)
plotRatioErrorBars(ratio_ax, x_values, ratio, unc, style)
if show_den_unc:
with np.errstate(divide="ignore", invalid="ignore"):
den_scaled_uncertainties = np.where(
den_total.values() != 0,
np.sqrt(den_total.variances()) / den_total.values(),
np.nan,
)
ratio_ax.bar(
x=den_total.axes[0].centers,
bottom=np.nan_to_num(1 - den_scaled_uncertainties, nan=0),
height=np.nan_to_num(2 * den_scaled_uncertainties, nan=0),
width=np.diff(den_total.axes[0].edges),
edgecolor="dimgrey",
hatch="////",
fill=False,
lw=0,
)
[docs]
def plotSingleNumeratorMultiDen(
ax,
ratio_ax,
numerator,
den_to_plot,
den_styles,
styler,
normalize,
ratio_type,
x_values,
ratio_func=computeRatio,
):
hist = numerator.item.histogram
n_style = styler.getStyle(numerator.metadata)
hist.plot1d(
ax=ax,
label=numerator.metadata.get("title") or numerator.metadata["dataset_title"],
density=normalize,
yerr=True,
**n_style.get(),
)
for den, style in zip(den_to_plot, den_styles):
n_vals = hist.values()
d_vals = den.item.histogram.values()
ratio, unc = ratio_func(
n_vals,
d_vals,
normalize=normalize,
ratio_type=ratio_type,
)
plotRatioErrorBars(ratio_ax, x_values, ratio, unc, style)
[docs]
def plotRatio(
denominator,
numerators,
output_path,
style_set,
normalize=False,
ratio_ylim=(0, 2),
ratio_type="poisson",
scale="linear",
plot_configuration=None,
no_stack=False,
ratio_hlines=(1.0,),
ratio_height=0.3,
):
pc = plot_configuration or PlotConfiguration()
styler = Styler(style_set)
fig, ax, ratio_ax = makeRatioAxes(ratio_height)
den_hist = denominator[0].item.histogram
x_values = den_hist.axes[0].centers
left_edge = den_hist.axes.edges[0][0]
right_edge = den_hist.axes.edges[-1][-1]
ratio_func = computeSignificance if ratio_type == "significance" else computeRatio
if no_stack:
den_to_plot, den_styles = plotUnstackedDenominators(
ax,
denominator,
styler,
normalize=normalize,
)
plotSingleNumeratorMultiDen(
ax,
ratio_ax,
numerators[0],
den_to_plot,
den_styles,
styler,
normalize=normalize,
ratio_type=ratio_type,
x_values=x_values,
ratio_func=ratio_func,
)
else:
den_total = plotStackedDenominators(
ax,
denominator,
styler,
normalize=normalize,
)
plotMultiNumerators(
ax,
ratio_ax,
numerators,
den_total,
styler,
normalize=normalize,
ratio_type=ratio_type,
x_values=x_values,
ratio_func=ratio_func,
)
for y in ratio_hlines:
ratio_ax.axhline(y, color="black", linestyle="dashed", linewidth=1.0)
ratio_ax.set_xlim(left_edge, right_edge)
ratio_ax.set_ylim(*ratio_ylim)
if ratio_type == "significance":
rylabel = "Significance"
else:
rylabel = "Ratio"
ratio_ax.set_ylabel(rylabel)
labelAxis(
ax,
"y",
den_hist.axes,
label="Normalized Events" if normalize else None,
)
labelAxis(ratio_ax, "x", den_hist.axes)
ax.set_xlabel("")
addLegend(ax, pc)
all_meta = [x.metadata for x in numerators] + [x.metadata for x in denominator]
ax.set_yscale(scale)
scaleYAxis(ax)
common_meta = commonDict(numerators + denominator, key=lambda x: x.metadata)
saveFigVariants(
fig,
ax,
output_path,
all_meta,
plot_configuration=pc,
metadata=common_meta,
)
plt.close(fig)
[docs]
def plotRatioOfRatios(
num_numerator,
num_denominator,
den_numerator,
den_denominator,
output_path,
style_set,
r1_label="Ratio 1",
r2_label="Ratio 2",
double_ratio_label="Ratio of Ratios",
normalize=False,
ratio_ylim=(0, 2),
ratio_type="poisson",
scale="linear",
plot_configuration=None,
ratio_hlines=(1.0,),
ratio_height=0.3,
):
pc = plot_configuration or PlotConfiguration()
styler = Styler(style_set)
fig, ax, ratio_ax = makeRatioAxes(ratio_height)
ref_hist = num_denominator.item.histogram
x_values = ref_hist.axes[0].centers
left_edge = ref_hist.axes.edges[0][0]
right_edge = ref_hist.axes.edges[-1][-1]
h_num_num = num_numerator.item.histogram
h_num_den = num_denominator.item.histogram
h_den_num = den_numerator.item.histogram
h_den_den = den_denominator.item.histogram
r1, r1_unc = computeRatio(
h_num_num.values(),
h_num_den.values(),
normalize=normalize,
ratio_type=ratio_type,
)
r2, r2_unc = computeRatio(
h_den_num.values(),
h_den_den.values(),
normalize=normalize,
ratio_type=ratio_type,
)
with np.errstate(divide="ignore", invalid="ignore"):
double_ratio = r1 / r2
double_ratio_unc = np.abs(double_ratio) * np.sqrt(
(r1_unc / r1) ** 2 + (r2_unc / r2) ** 2
)
double_ratio[np.isnan(double_ratio)] = 0
double_ratio[np.isinf(double_ratio)] = 0
double_ratio_unc[np.isnan(double_ratio_unc)] = 0
style1 = styler.getStyle(num_numerator.metadata)
style2 = styler.getStyle(den_numerator.metadata)
opts1 = {
**style1.get("errorbar", include_type=False),
"label": r1_label,
}
opts2 = {**style2.get("errorbar", include_type=False), "label": r2_label}
ax.errorbar(x_values, r1, yerr=r1_unc, **opts1)
ax.errorbar(x_values, r2, yerr=r2_unc, **opts2)
ratio_ax.errorbar(
x_values, double_ratio, yerr=double_ratio_unc, color="black", fmt="o"
)
for y in ratio_hlines:
ratio_ax.axhline(y, color="black", linestyle="dashed", linewidth=1.0)
ratio_ax.set_xlim(left_edge, right_edge)
ratio_ax.set_ylim(*ratio_ylim)
ratio_ax.set_ylabel(double_ratio_label)
labelAxis(ax, "y", ref_hist.axes, label="Ratio")
labelAxis(ratio_ax, "x", ref_hist.axes)
addLegend(ax, pc)
all_meta = [
x.metadata
for x in [num_numerator, num_denominator, den_numerator, den_denominator]
]
ax.set_yscale(scale)
common_meta = commonDict(all_meta, key=lambda x: x)
saveFigVariants(
fig,
ax,
output_path,
all_meta,
plot_configuration=pc,
metadata=common_meta,
)
plt.close(fig)
[docs]
def plotModel(
data,
backgrounds,
signal,
output_path,
style_set,
normalize=False,
ratio_ylim=(0, 2),
ratio_type="poisson",
scale="linear",
plot_configuration=None,
ratio_hlines=(1.0,),
ratio_height=0.3,
):
pc = plot_configuration or PlotConfiguration()
styler = Styler(style_set)
fig, ax, ratio_ax = makeRatioAxes(ratio_height)
bg_total = plotStackedDenominators(
ax,
backgrounds,
styler,
normalize=normalize,
)
largest_bg = max(backgrounds, key=lambda x: x.item.histogram.sum().value)
largest_bg_style = styler.getStyle(largest_bg.metadata)
style_kwargs = largest_bg_style.get()
largest_bg_color = style_kwargs.get("color")
bg_color = largest_bg_color
sig_item, sig_meta = signal
sig_style = styler.getStyle(sig_meta)
sig_hist = sig_item.histogram
sig_hist.plot1d(
ax=ax,
label=sig_meta.get("title") or sig_meta["dataset_title"],
density=normalize,
yerr=sig_style.yerr,
**sig_style.get(),
)
sig_color = sig_style.get().get("color")
bg_sig_hist = bg_total + sig_hist
bg_sig_hist.plot1d(
ax=ax,
label=f"B + {sig_meta.get('title') or sig_meta.get('dataset_title', '')}",
density=normalize,
color=sig_color or "red",
linestyle="--",
)
data_item, data_meta = data
data_style = styler.getStyle(data_meta)
data_hist = data_item.histogram
data_hist.plot1d(
ax=ax,
label=data_meta.get("title") or data_meta["dataset_title"],
density=normalize,
yerr=True,
**data_style.get(),
)
ratio_func = computeSignificance if ratio_type == "significance" else computeRatio
x_values = bg_total.axes[0].centers
left_edge = bg_total.axes.edges[0][0]
right_edge = bg_total.axes.edges[-1][-1]
n_vals = data_hist.values()
ratio, unc = ratio_func(
n_vals,
bg_total.values(),
normalize=normalize,
ratio_type=ratio_type,
)
plotRatioErrorBars(ratio_ax, x_values, ratio, unc, data_style)
bs_ratio, bs_unc = ratio_func(
n_vals,
bg_sig_hist.values(),
normalize=normalize,
ratio_type=ratio_type,
)
ratio_ax.errorbar(
x_values,
bs_ratio,
yerr=bs_unc,
linestyle="none",
marker="o",
color=sig_color or bg_color,
label=f"Ratio w/ {sig_meta.get('title', 'Signal')}",
markerfacecolor="none",
)
with np.errstate(divide="ignore", invalid="ignore"):
den_scaled_uncertainties = np.where(
bg_total.values() != 0,
np.sqrt(bg_total.variances()) / bg_total.values(),
np.nan,
)
ratio_ax.bar(
x=bg_total.axes[0].centers,
bottom=np.nan_to_num(1 - den_scaled_uncertainties, nan=0),
height=np.nan_to_num(2 * den_scaled_uncertainties, nan=0),
width=np.diff(bg_total.axes[0].edges),
edgecolor="dimgrey",
hatch="////",
fill=False,
lw=0,
)
for y in ratio_hlines:
ratio_ax.axhline(y, color="black", linestyle="dashed", linewidth=1.0)
ratio_ax.set_xlim(left_edge, right_edge)
ratio_ax.set_ylim(*ratio_ylim)
rylabel = "Significance" if ratio_type == "significance" else "Ratio"
ratio_ax.set_ylabel(rylabel)
labelAxis(
ax,
"y",
data_hist.axes,
label="Normalized Events" if normalize else None,
)
labelAxis(ratio_ax, "x", data_hist.axes)
ax.set_xlabel("")
addLegend(ax, pc)
all_meta = [data_meta, sig_meta] + [x.metadata for x in backgrounds]
ax.set_yscale(scale)
scaleYAxis(ax)
common_meta = commonDict(all_meta, key=lambda x: x)
saveFigVariants(
fig,
ax,
output_path,
all_meta,
plot_configuration=pc,
metadata=common_meta,
)
plt.close(fig)