Completely package and refactor contribution plotting methods

packages/pidplots/pidplots/contribution.py

+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

packages/pidplots/pidplots/distributions.py

+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
+    )