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

Completely package and refactor contribution plotting methods

parent 236dc4f2
import numpy as np
import matplotlib.pyplot as plt
from . import const, colors
from . import distributions as dist
from .utils import (
_ax_xlabel,
connected_region_borders,
imshow_2D_axis_labels,
outer_legend,
better_xlabel,
text_array_values,
xticks_radians_to_degrees,
)
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),
)
else:
raise ValueError("`det` not understood.")
def plot_contribution_distribution(
df, det="max", color_list=None, yscale="log", figsize=None
):
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"
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
def subplot_avg_contribution(ax, df, color_list=None):
ctrbs = df[[f"contrib_{det}" for det in const.DETECTORS]].values
ctrbs = np.mean(ctrbs, axis=0)
x_vals = np.arange(len(const.DETECTORS))
bars = ax.bar(x_vals, ctrbs, width=0.8)
color_list = colors.fill_color_list(color_list)
for i in range(len(const.DETECTORS)):
bars[i].set_facecolor(color_list[i])
ax.set_xticks(x_vals)
ax.set_xticklabels(const.DETECTORS)
ax.set_xlim(x_vals[0] - 0.5, x_vals[-1] + 0.5)
def plot_avg_contribution(df, title=None, color_list=None, figsize=None):
fig, ax = plt.subplots(figsize=figsize)
subplot_avg_contribution(ax, df, color_list=color_list)
ax.set_xlabel("Detector")
ax.set_ylabel("Average contribution value")
ax.set_title(title)
fig.tight_layout()
return fig
def subplot_avg_contribution_1D(ax, df, bin_by="p", ls="-", color_list=None):
if bin_by == "p":
ctrbs = df.groupby(["p_bin"]).mean()
bin_centers = const.P_BIN_CENTERS
elif bin_by == "theta":
ctrbs = df.groupby(["theta_bin"]).mean()
bin_centers = const.THETA_BIN_CENTERS
color_list = colors.fill_color_list(color_list)
ctrbs = ctrbs[[f"contrib_{det}" for det in const.DETECTORS]]
for i, det in enumerate(const.DETECTORS):
vals = np.zeros(len(bin_centers))
vals[ctrbs.index] = ctrbs[f"contrib_{det}"]
ax.plot(bin_centers, vals, ls=ls, c=color_list[i], label=det)
def plot_avg_contribution_1D(df, bin_by="p", title=None, figsize=None, color_list=None):
fig, ax = plt.subplots(figsize=figsize)
subplot_avg_contribution_1D(ax, df, bin_by=bin_by, color_list=color_list)
fig.tight_layout()
_ax_xlabel(ax, bin_by)
if bin_by == "theta":
xticks_radians_to_degrees(ax)
ax.set_ylabel("Average contribution value")
outer_legend(fig, loc="top", ncol=len(const.DETECTORS), y=0.98)
better_xlabel(fig, title, loc="top", y=1.08)
return fig
def subplot_avg_contribution_2D(ax, df, det="max", color_list=None, cell_fontsize=10):
color_list = colors.fill_color_list(color_list)
gb = df.groupby(["p_bin", "theta_bin"])
# handle single detector avg contribution vs p, theta
if det in const.DETECTORS:
cmap = colors.make_linear_colormap(color_list[const.DETECTORS.index(det)])
val = np.zeros((const.N_P_BINS, const.N_THETA_BINS))
for (i, j), g in gb:
val[i, j] = g.mean()[f"contrib_{det}"]
im = ax.imshow(val, cmap=cmap, aspect="auto")
# add cell values
if cell_fontsize > 0:
text_array_values(ax, val, fontsize=cell_fontsize)
# handle min/max avg contributions vs p, theta
elif det == "min" or det == "max":
cmap = colors.make_colormap(color_list)
all_vals = np.zeros((const.N_P_BINS, const.N_THETA_BINS, len(const.DETECTORS)))
for (i, j), g in gb:
all_vals[i, j] = g.mean()[
[f"contrib_{det}" for det in const.DETECTORS]
].values
if det == "min":
idx = np.argmin(all_vals, axis=2)
else:
idx = np.argmax(all_vals, axis=2)
# TODO: make sure vlims are correct
im = ax.imshow(
idx, cmap=cmap, vmin=-0.5, vmax=len(const.DETECTORS) - 0.5, aspect="auto"
)
# add cell values
val = np.take_along_axis(all_vals, np.expand_dims(idx, 2), axis=2).squeeze()
if cell_fontsize > 0:
text_array_values(ax, val, fontsize=cell_fontsize)
connected_region_borders(ax, idx)
else:
raise ValueError("Value given for argument 'max' not understood.")
imshow_2D_axis_labels(ax)
return im
def plot_avg_contribution_2D(
df, det="max", figsize=None, color_list=None, cell_fontsize=10
):
fig, ax = plt.subplots(figsize=figsize)
color_list = colors.fill_color_list(color_list)
im = subplot_avg_contribution_2D(
ax, df, color_list=color_list, det=det, cell_fontsize=cell_fontsize
)
fig.tight_layout()
if det in const.DETECTORS:
cb = fig.colorbar(im)
cb.set_label(f"{det} average contribution value")
else:
from matplotlib.patches import Patch
handles = [Patch(color=color_list[i]) for i in range(len(const.DETECTORS))]
title = ("Largest" if det == "max" else "Smallest") + "\ncontributor"
outer_legend(
fig,
handles,
const.DETECTORS,
loc="right",
ncol=1,
title=title,
x=0.98,
)
return fig
import numpy as np
# There are a few different ways that we can plot the distributions
# We can plot histograms of...
# - all distributions
# - a single distribution
# - 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):
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)
def largest(ax, data, bin_lims=(0, 1), n_bins=30, labels=None, colors=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
)
def smallest(ax, data, bin_lims=(0, 1), n_bins=30, labels=None, colors=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
)
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