Add support for passing a length-one color or label list in certain cases

packages/pidplots/pidplots/distributions.py

@@ -65,6 +65,15 @@ def smallest(ax, data, bin_lims=(0, 1), n_bins=30, labels=None, color_list=None)
     )
 
 
+def _fill_label_list(label_list, detectors=False):
+    if label_list is not None:
+        return label_list
+    if detectors:
+        return const.DETECTORS
+    else:
+        return const.PART_PLOT_LABELS
+
+
 def subplot_distribution(
     ax,
     data,
@@ -76,6 +85,7 @@ def subplot_distribution(
     n_bins=30,
 ):
     color_list = colors.fill_color_list(color_list, detectors=detectors)
+    labels = _fill_label_list(labels, detectors=detectors)
 
     if bin_lims is None:
         if detectors:
@@ -101,22 +111,42 @@ def subplot_distribution(
         )
     else:
         if detectors and kind in const.DETECTORS:
-            d_index = const.DETECTORS.index(kind)
+            if len(color_list) == 1:
+                color = color_list[0]
+            else:
+                d_index = const.DETECTORS.index(kind)
+                color = color_list[d_index]
+
+            if len(labels) == 1:
+                label = labels[0]
+            else:
+                label = labels[d_index]
+
             single(
                 ax,
                 data,
-                label=kind,
-                color=color_list[d_index],
+                label=label,
+                color=color,
                 bin_lims=bin_lims,
                 n_bins=n_bins,
             )
         elif (not detectors) and kind in const.PARTICLES:
             p_index = const.PARTICLES.index(kind)
+            if len(color_list) == 1:
+                color = color_list[0]
+            else:
+                color = color_list[p_index]
+
+            if len(labels) == 1:
+                label = labels[0]
+            else:
+                label = labels[p_index]
+
             single(
                 ax,
                 data,
-                label=const.PART_PLOT_LABELS[p_index],
-                color=color_list[p_index],
+                label=label,
+                color=color,
                 bin_lims=bin_lims,
                 n_bins=n_bins,
             )
@@ -128,6 +158,8 @@ def plot_distribution(
     data,
     kind="all",
     detectors=False,
+    bin_lims=None,
+    n_bins=30,
     figsize=None,
     xlabel=None,
     title=None,
@@ -137,11 +169,14 @@ def plot_distribution(
     from .utils import outer_legend
 
     fig, ax = subplots(figsize=figsize)
-    subplot_distribution(ax, data, kind=kind, detectors=detectors)
+    subplot_distribution(
+        ax, data, kind=kind, detectors=detectors, bin_lims=bin_lims, n_bins=n_bins
+    )
     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)
+        labels = const.DETECTORS if detectors else const.PART_PLOT_LABELS
+        outer_legend(fig, labels, loc="bottom", y=0.02)
     return fig