Commit 8a4646d9 authored by connor.hainje@pnnl.gov's avatar connor.hainje@pnnl.gov
Browse files

Add a unified architecture for plots of distributions

parent 4fe24853
......@@ -14,66 +14,63 @@ from .utils import (
)
def subplot_contribution_distribution(ax, df, color_list=None, det="max"):
color_list = colors.fill_color_list(color_list)
contrib_cols = [f"contrib_{det}" for det in const.DETECTORS]
if det == "all":
dist.all(
ax,
df[contrib_cols].values,
labels=const.DETECTORS,
colors=color_list,
bin_lims=(-1, 1),
)
elif det == "min":
dist.smallest(
ax,
df[contrib_cols].values,
labels=const.DETECTORS,
colors=color_list,
bin_lims=(-1, 1),
)
elif det == "max":
dist.largest(
ax,
df[contrib_cols].values,
labels=const.DETECTORS,
colors=color_list,
bin_lims=(-1, 1),
)
elif det in const.DETECTORS:
dist.single(
ax,
df[f"contrib_{det}"].values,
label=det,
color=color_list[const.DETECTORS.index(det)],
bin_lims=(-1, 1),
)
def subplot_contribution_distribution(
ax, df, det="max", color_list=None, bin_lims=(-1, 1), n_bins=30
):
if det in const.PARTICLES:
data = df[f"contrib_{det}"].values
elif det in ["all", "min", "max"]:
contrib_cols = [f"contrib_{det}" for det in const.DETECTORS]
data = df[contrib_cols].values
else:
raise ValueError("`det` not understood.")
raise ValueError("argument det not understood")
dist.subplot_distribution(
ax,
data,
kind=det,
detectors=True,
labels=const.DETECTORS,
colors=color_list,
bin_lims=bin_lims,
n_bins=n_bins,
)
def plot_contribution_distribution(
df, det="max", color_list=None, yscale="log", figsize=None
df,
det="max",
figsize=None,
title=None,
yscale="log",
bin_lims=(-1, 1),
n_bins=30,
):
fig, ax = plt.subplots(figsize=figsize)
subplot_contribution_distribution(ax, df, det=det, color_list=color_list)
ax.set_yscale(yscale)
if det == "all":
xlabel = "Detector contribution"
elif det == "min":
xlabel = "Smallest detector contribution"
elif det == "max":
xlabel = "Largest detector contribution"
if det in ["all", "min", "max"]:
if det == "all":
xlabel = "Detector contribution"
elif det == "min":
xlabel = "Smallest detector contribution"
elif det == "max":
xlabel = "Largest detector contribution"
contrib_cols = [f"contrib_{det}" for det in const.DETECTORS]
data = df[contrib_cols].values
elif det in const.DETECTORS:
xlabel = f"{det} contribution"
ax.set_xlabel(xlabel)
fig.tight_layout()
if det not in const.DETECTORS:
outer_legend(fig, loc="bottom", y=0.02)
return fig
data = df[f"contrib_{det}"].values
else:
raise ValueError("argument det not understood")
return dist.plot_distribution(
data,
kind=det,
figsize=figsize,
xlabel=xlabel,
title=title,
yscale=yscale,
bin_lims=bin_lims,
n_bins=n_bins,
)
def subplot_avg_contribution(ax, df, color_list=None):
......
import numpy as np
from . import colors, const
# There are a few different ways that we can plot the distributions
# We can plot histograms of...
......@@ -7,38 +8,140 @@ import numpy as np
# - the largest/smallest (which we'll split by distribution of origin)
def all(ax, data, bin_lims=(0, 1), n_bins=30, labels=None, colors=None):
bins = np.linspace(*bin_lims, num=n_bins + 1)
for i in range(data.shape[1]):
label = labels[i] if labels else ""
color = colors[i] if colors else None
ax.hist(data[:, i], bins=bins, histtype="step", label=label, color=color)
# single(ax, data, i, bin_lims=bin_lims, n_bins=n_bins, label=label)
def single(ax, data, bin_lims=(0, 1), n_bins=30, label="", color=None):
bins = np.linspace(*bin_lims, num=n_bins + 1)
ax.hist(data, bins=bins, histtype="step", label=label, color=color)
def _min_or_max(ax, data, idx, bin_lims=(0, 1), n_bins=30, labels=None, colors=None):
def all(ax, data, bin_lims=(0, 1), n_bins=30, labels=None, color_list=None):
for i in range(data.shape[1]):
label = labels[i] if labels else ""
color = color_list[i] if color_list else None
single(
ax, data[:, i], bin_lims=bin_lims, n_bins=n_bins, label=label, color=color
)
def _min_or_max(
ax, data, idx, bin_lims=(0, 1), n_bins=30, labels=None, color_list=None
):
val = data[np.arange(len(idx)), idx]
bins = np.linspace(*bin_lims, num=n_bins + 1)
for i in range(data.shape[1]):
label = labels[i] if labels else ""
color = colors[i] if colors else None
ax.hist(val[idx == i], bins=bins, histtype="step", label=label, color=color)
color = color_list[i] if color_list else None
single(
ax,
val[idx == i],
bin_lims=bin_lims,
n_bins=n_bins,
label=label,
color=color,
)
def largest(ax, data, bin_lims=(0, 1), n_bins=30, labels=None, colors=None):
def largest(ax, data, bin_lims=(0, 1), n_bins=30, labels=None, color_list=None):
idx = np.argmax(data, axis=1)
_min_or_max(
ax, data, idx, bin_lims=bin_lims, n_bins=n_bins, labels=labels, colors=colors
ax,
data,
idx,
bin_lims=bin_lims,
n_bins=n_bins,
labels=labels,
color_list=color_list,
)
def smallest(ax, data, bin_lims=(0, 1), n_bins=30, labels=None, colors=None):
def smallest(ax, data, bin_lims=(0, 1), n_bins=30, labels=None, color_list=None):
idx = np.argmin(data, axis=1)
_min_or_max(
ax, data, idx, bin_lims=bin_lims, n_bins=n_bins, labels=labels, colors=colors
ax,
data,
idx,
bin_lims=bin_lims,
n_bins=n_bins,
labels=labels,
color_list=color_list,
)
def subplot_distribution(
ax,
data,
kind="all",
detectors=False,
labels=None,
color_list=None,
bin_lims=None,
n_bins=30,
):
color_list = colors.fill_color_list(color_list, detectors=detectors)
if bin_lims is None:
if detectors:
bin_lims = (-1, 1)
else:
bin_lims = (0, 1)
if kind in ["all", "min", "max"]:
if kind == "all":
func = all
elif kind == "min":
func = smallest
else:
func = largest
func(
ax,
data,
labels=labels,
color_list=color_list,
bin_lims=bin_lims,
n_bins=n_bins,
)
else:
if detectors and kind in const.DETECTORS:
d_index = const.DETECTORS.index(kind)
single(
ax,
data,
label=kind,
color=color_list[d_index],
bin_lims=bin_lims,
n_bins=n_bins,
)
elif (not detectors) and kind in const.PARTICLES:
p_index = const.PARTICLES.index(kind)
single(
ax,
data,
label=const.PART_PLOT_LABELS[p_index],
color=color_list[p_index],
bin_lims=bin_lims,
n_bins=n_bins,
)
else:
raise ValueError("")
def plot_distribution(
data,
kind="all",
detectors=False,
figsize=None,
xlabel=None,
title=None,
yscale="log",
):
from matplotlib.pyplot import subplots
from .utils import outer_legend
fig, ax = subplots(figsize=figsize)
subplot_distribution(ax, data, kind=kind, detectors=detectors)
ax.set_yscale(yscale)
ax.set_xlabel(xlabel)
ax.set_title(title)
fig.tight_layout()
if kind not in const.PARTICLES and kind not in const.DETECTORS:
outer_legend(fig, loc="bottom", y=0.02)
return fig
from . import const, distributions as dist
def subplot_likelihood_ratio(
ax, df, particle="all", color_list=None, bin_lims=(0, 1), n_bins=30
):
if particle in const.PARTICLES:
data = df[f"lr_{particle}"].values
elif particle in ["all", "min", "max"]:
lr_cols = [f"lr_{h}" for h in const.PARTICLES]
data = df[lr_cols].values
else:
raise ValueError("argument particle not understood")
dist.subplot_distribution(
ax,
data,
kind=particle,
detectors=False,
labels=const.PART_PLOT_LABELS,
colors=color_list,
bin_lims=bin_lims,
n_bins=n_bins,
)
def plot_likelihood_ratio(
df,
particle="all",
figsize=None,
title=None,
yscale="log",
bin_lims=(0, 1),
n_bins=30,
):
if particle in ["all", "min", "max"]:
if particle == "all":
xlabel = "Likelihood ratios"
elif particle == "min":
xlabel = "Smallest likelihood ratio"
elif particle == "max":
xlabel = "Largest likelihood ratio"
lr_cols = [f"lr_{h}" for h in const.PARTICLES]
data = df[lr_cols].values
elif particle in const.PARTICLES:
p_index = const.PARTICLES.index(particle)
xlabel = f"{const.PART_PLOT_LABELS[p_index]} likelihood ratio"
data = df[f"lr_{particle}"].values
else:
raise ValueError("argument particle not understood")
return dist.plot_distribution(
data,
kind=particle,
figsize=figsize,
xlabel=xlabel,
title=title,
yscale=yscale,
bin_lims=bin_lims,
n_bins=n_bins,
)
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment