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

Make color, linestyle args for efficiency subplot

parent 49942609
......@@ -115,23 +115,52 @@ def compute_purity_in_bins(df, threshold=0.5, bin_by="p"):
return _comp_in_bins(compute_purity, df, bin_by=bin_by, threshold=threshold)
def subplot_efficiency(ax, df, part, threshold=0.5, bin_by="p"):
def subplot_efficiency(
ax,
df,
part,
threshold=0.5,
bin_by="p",
ls=None,
color=None,
ls_list=None,
color_list=None,
):
_validate_binary()
_validate_bin_by_1D(bin_by)
effs = compute_efficiency_in_bins(df, threshold=threshold, bin_by=bin_by)
if ls is not None:
ls_list = [ls, ls]
else:
if ls_list is None:
ls_list = [None, None]
else:
ls_list = ls_list
if color is not None:
color_list = [color, color]
else:
if color_list is None:
if part == "pi":
color_list = [colors.PARTICLES["pi"], colors.PARTICLES[BINARY]]
else:
color_list = [colors.PARTICLES[BINARY], colors.PARTICLES["pi"]]
else:
color_list = color_list
if bin_by == "p":
bincs = const.P_BIN_CENTERS
else:
bincs = const.THETA_BIN_CENTERS
if part == "pi":
ax.plot(bincs, effs[:, 0, 0], "-")
ax.plot(bincs, effs[:, 1, 0], "-")
ax.plot(bincs, effs[:, 0, 0], ls=ls_list[0], color=color_list[0])
ax.plot(bincs, effs[:, 1, 0], ls=ls_list[1], color=color_list[1])
else:
ax.plot(bincs, effs[:, 1, 1], "-")
ax.plot(bincs, effs[:, 0, 1], "-")
ax.plot(bincs, effs[:, 1, 1], ls=ls_list[0], color=color_list[0])
ax.plot(bincs, effs[:, 0, 1], ls=ls_list[1], color=color_list[1])
ax.set_ylim(0, 1)
ax.grid(ls=":")
......@@ -184,6 +213,9 @@ def subplot_purity(ax, df, part, threshold=0.5, bin_by="p", color=None, ls="-"):
else:
bincs = const.THETA_BIN_CENTERS
if color is None:
color = colors.PARTICLES[part]
if part == "pi":
ax.plot(bincs, purs[:, 0, 0], ls=ls, c=color)
else:
......
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