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

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

parent 64602dd7
......@@ -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
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