Make color, linestyle args for efficiency subplot

packages/pidplots/pidplots/binary.py

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