From 3d50c77d56212995cdf9cb7e544644bcd037f19c Mon Sep 17 00:00:00 2001
From: MarcSerraPeralta <43704266+MarcSerraPeralta@users.noreply.github.com>
Date: Fri, 13 Sep 2024 09:45:17 +0200
Subject: [PATCH] Add plot for threshold lines (#24)

---
 qec_util/performance/__init__.py |  2 ++
 qec_util/performance/plots.py    | 54 ++++++++++++++++++++++++++++++++
 tests/conftest.py                | 11 +++++++
 tests/layouts/test_plotter.py    |  6 +++-
 tests/performance/test_plots.py  | 28 +++++++++++++++++
 5 files changed, 100 insertions(+), 1 deletion(-)
 create mode 100644 qec_util/performance/plots.py
 create mode 100644 tests/conftest.py
 create mode 100644 tests/performance/test_plots.py

diff --git a/qec_util/performance/__init__.py b/qec_util/performance/__init__.py
index b3dc607..eb5e4be 100644
--- a/qec_util/performance/__init__.py
+++ b/qec_util/performance/__init__.py
@@ -6,6 +6,7 @@
     confidence_interval_binomial,
 )
 from .sampler import sample_failures
+from . import plots
 
 __all__ = [
     "logical_error_prob",
@@ -14,4 +15,5 @@
     "lmfit_par_to_ufloat",
     "confidence_interval_binomial",
     "sample_failures",
+    "plots",
 ]
diff --git a/qec_util/performance/plots.py b/qec_util/performance/plots.py
new file mode 100644
index 0000000..bedff38
--- /dev/null
+++ b/qec_util/performance/plots.py
@@ -0,0 +1,54 @@
+from typing import Optional
+
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+def plot_line_threshold(
+    ax: plt.Axes,
+    phys_prob: np.ndarray,
+    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",
+) -> plt.Axes:
+    """Plots the logical error probability as a function of the physiscal
+    error probability, including its upper and lower error bars (if given).
+
+    Parameters
+    ----------
+    ax
+        Matplotlib axis in which to plot.
+    phys_prob
+        Physical error probabilities.
+    log_prob
+        Logical error probability.
+    log_prob_lower
+        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.
+
+    Returns
+    -------
+    ax
+        Matplotlib axis.
+    """
+    ax.plot(phys_prob, log_prob, ".", color=color, label=label)
+    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
+        )
+
+    ax.set_xscale("log")
+    ax.set_yscale("log")
+    ax.set_xlabel("physical error probability, $p$")
+    ax.set_ylabel("logical error probability, $p_L$")
+    if label is not None:
+        ax.legend(loc="best")
+
+    return ax
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 0000000..aa606c6
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,11 @@
+def pytest_addoption(parser):
+    parser.addoption("--show-figures", action="store_true")
+
+
+def pytest_generate_tests(metafunc):
+    # This is called for every test. Only get/set command line arguments
+    # if the argument is specified in the list of test "fixturenames".
+    # The hyphens in the arguments are substutited by underscores.
+    option_value = metafunc.config.option.show_figures
+    if "show_figures" in metafunc.fixturenames and option_value is not None:
+        metafunc.parametrize("show_figures", [bool(option_value)])
diff --git a/tests/layouts/test_plotter.py b/tests/layouts/test_plotter.py
index 85618a0..35cb50f 100644
--- a/tests/layouts/test_plotter.py
+++ b/tests/layouts/test_plotter.py
@@ -4,10 +4,14 @@
 from qec_util.layouts import plot
 
 
-def test_plot():
+def test_plot(show_figures):
     _, ax = plt.subplots()
     layout = rot_surf_code(3)
 
     plot(ax, layout)
 
+    if show_figures:
+        plt.show()
+    plt.close()
+
     return
diff --git a/tests/performance/test_plots.py b/tests/performance/test_plots.py
new file mode 100644
index 0000000..e5dc941
--- /dev/null
+++ b/tests/performance/test_plots.py
@@ -0,0 +1,28 @@
+import numpy as np
+import matplotlib.pyplot as plt
+
+from qec_util.performance.plots import plot_line_threshold
+
+
+def test_plot_line_threshold(show_figures):
+    phys_prob = np.linspace(1, 4, 10)
+    log_prob = np.linspace(1, 6, 10)
+    log_prob_lower = log_prob - 1
+    log_prob_upper = log_prob + 2
+
+    _, ax = plt.subplots()
+    ax = plot_line_threshold(
+        ax,
+        phys_prob,
+        log_prob,
+        log_prob_lower,
+        log_prob_upper,
+        color="red",
+        label="example",
+    )
+
+    if show_figures:
+        plt.show()
+    plt.close()
+
+    assert isinstance(ax, plt.Axes)