From 5a0d9068647e91b4ebfcf3d6e68888f090575923 Mon Sep 17 00:00:00 2001
From: MarcSerraPeralta <43704266+MarcSerraPeralta@users.noreply.github.com>
Date: Fri, 13 Sep 2024 11:18:52 +0200
Subject: [PATCH] Improve plot parameter options (#25)

---
 qec_util/performance/plots.py | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/qec_util/performance/plots.py b/qec_util/performance/plots.py
index bedff38..be886c2 100644
--- a/qec_util/performance/plots.py
+++ b/qec_util/performance/plots.py
@@ -10,8 +10,7 @@ def plot_line_threshold(
     log_prob: np.ndarray,
     log_prob_lower: Optional[np.ndarray] = None,
     log_prob_upper: Optional[np.ndarray] = None,
-    label: Optional[str] = None,
-    color: Optional[str] = "blue",
+    **kargs,
 ) -> plt.Axes:
     """Plots the logical error probability as a function of the physiscal
     error probability, including its upper and lower error bars (if given).
@@ -28,17 +27,19 @@ def plot_line_threshold(
         Lower bound on the logical error probability uncertainty.
     log_prob_upper
         Upper bound on the logical error probability uncertainty.
-    label
-        Label for the data. By default, no label.
-    color
-        Color for the data. By default, blue.
+    **kargs
+        Arguments for ``ax.plot``.
 
     Returns
     -------
     ax
         Matplotlib axis.
     """
-    ax.plot(phys_prob, log_prob, ".", color=color, label=label)
+    kargs_ = dict(marker=".", linestyle="none")
+    kargs_.update(kargs)
+
+    p = ax.plot(phys_prob, log_prob, **kargs_)
+    color = p[0].get_color()
     if (log_prob_lower is not None) and (log_prob_upper is not None):
         ax.fill_between(
             phys_prob, log_prob_lower, log_prob_upper, color=color, alpha=0.1
@@ -48,7 +49,8 @@ def plot_line_threshold(
     ax.set_yscale("log")
     ax.set_xlabel("physical error probability, $p$")
     ax.set_ylabel("logical error probability, $p_L$")
-    if label is not None:
+
+    if ("label" in kargs_) and (kargs_["label"] is not None):
         ax.legend(loc="best")
 
     return ax