Add a unified architecture for plots of distributions

packages/pidplots/pidplots/contribution.py

@@ -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):

packages/pidplots/pidplots/distributions.py

 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

packages/pidplots/pidplots/likelihood_ratios.py

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