feat: option to disable legend (#20)

* feat: option to disable legend

* refactor: ensure param is keyword-only

* docs: add release note

* docs: fix RST formatting

* docs: close previous figures

* fix: suppress missing coverage

---------

Co-authored-by: Max Rossmannek <[email protected]>

Author: boonware

qiskit_addon_obp/utils/visualization.py

@@ -21,7 +21,7 @@
 from .metadata import OBPMetadata
 
 
-def plot_accumulated_error(metadata: OBPMetadata, axes: Axes) -> None:
+def plot_accumulated_error(metadata: OBPMetadata, axes: Axes, *, show_legend: bool = True) -> None:
     """Plot the accumulated error.
 
     This method populates the provided figure axes with a line-plot of the
@@ -36,7 +36,7 @@ def plot_accumulated_error(metadata: OBPMetadata, axes: Axes) -> None:
        >>> metadata = OBPMetadata.from_json("docs/_static/dummy_visualization_metadata.json")
 
     .. plot::
-       :context:
+       :context: close-figs
        :include-source:
 
        >>> from matplotlib import pyplot as plt
@@ -57,7 +57,7 @@ def plot_accumulated_error(metadata: OBPMetadata, axes: Axes) -> None:
     Args:
         metadata: the metadata to be visualized.
         axes: the matplotlib axes in which to plot.
-
+        show_legend: enable/disable showing the legend in the plot.
     """
     if not np.isinf(metadata.truncation_error_budget.max_error_total):
         axes.axhline(
@@ -79,10 +79,12 @@ def plot_accumulated_error(metadata: OBPMetadata, axes: Axes) -> None:
         )
     axes.set_xlabel("backpropagated slice number")
     axes.set_ylabel("accumulated error")
-    axes.legend()
+    _set_legend(axes, show_legend)
 
 
-def plot_left_over_error_budget(metadata: OBPMetadata, axes: Axes) -> None:
+def plot_left_over_error_budget(
+    metadata: OBPMetadata, axes: Axes, *, show_legend: bool = True
+) -> None:
     """Plot the left-over error budget.
 
     This method populates the provided figure axes with a line-plot of the
@@ -97,7 +99,7 @@ def plot_left_over_error_budget(metadata: OBPMetadata, axes: Axes) -> None:
        >>> metadata = OBPMetadata.from_json("docs/_static/dummy_visualization_metadata.json")
 
     .. plot::
-       :context:
+       :context: close-figs
        :include-source:
 
        >>> from matplotlib import pyplot as plt
@@ -113,7 +115,7 @@ def plot_left_over_error_budget(metadata: OBPMetadata, axes: Axes) -> None:
     Args:
         metadata: the metadata to be visualized.
         axes: the matplotlib axes in which to plot.
-
+        show_legend: enable/disable showing the legend in the plot.
     """
     for obs_idx in range(len(metadata.backpropagation_history[0].slice_errors)):
         axes.plot(
@@ -126,10 +128,10 @@ def plot_left_over_error_budget(metadata: OBPMetadata, axes: Axes) -> None:
         )
     axes.set_xlabel("backpropagated slice number")
     axes.set_ylabel("left-over error budget")
-    axes.legend()
+    _set_legend(axes, show_legend)
 
 
-def plot_slice_errors(metadata: OBPMetadata, axes: Axes) -> None:
+def plot_slice_errors(metadata: OBPMetadata, axes: Axes, *, show_legend: bool = True) -> None:
     """Plot the slice errors.
 
     This method populates the provided figure axes with a bar-plot of the truncation error incurred
@@ -144,7 +146,7 @@ def plot_slice_errors(metadata: OBPMetadata, axes: Axes) -> None:
        >>> metadata = OBPMetadata.from_json("docs/_static/dummy_visualization_metadata.json")
 
     .. plot::
-       :context:
+       :context: close-figs
        :include-source:
 
        >>> from matplotlib import pyplot as plt
@@ -163,7 +165,7 @@ def plot_slice_errors(metadata: OBPMetadata, axes: Axes) -> None:
     Args:
         metadata: the metadata to be visualized.
         axes: the matplotlib axes in which to plot.
-
+        show_legend: enable/disable showing the legend in the plot.
     """
     num_observables = len(metadata.backpropagation_history[0].slice_errors)
     width = 0.8 / num_observables
@@ -181,9 +183,10 @@ def plot_slice_errors(metadata: OBPMetadata, axes: Axes) -> None:
     axes.set_xlabel("backpropagated slice number")
     axes.set_ylabel("incurred slice error")
     axes.legend()
+    _set_legend(axes, show_legend)
 
 
-def plot_num_paulis(metadata: OBPMetadata, axes: Axes) -> None:
+def plot_num_paulis(metadata: OBPMetadata, axes: Axes, *, show_legend: bool = True) -> None:
     """Plot the number of Pauli terms.
 
     This method populates the provided figure axes with a line-plot of the number of Pauli terms at
@@ -198,7 +201,7 @@ def plot_num_paulis(metadata: OBPMetadata, axes: Axes) -> None:
        >>> metadata = OBPMetadata.from_json("docs/_static/dummy_visualization_metadata.json")
 
     .. plot::
-       :context:
+       :context: close-figs
        :include-source:
 
        >>> from matplotlib import pyplot as plt
@@ -217,7 +220,7 @@ def plot_num_paulis(metadata: OBPMetadata, axes: Axes) -> None:
     Args:
         metadata: the metadata to be visualized.
         axes: the matplotlib axes in which to plot.
-
+        show_legend: enable/disable showing the legend in the plot.
     """
     for obs_idx in range(len(metadata.backpropagation_history[0].slice_errors)):
         axes.plot(
@@ -227,10 +230,12 @@ def plot_num_paulis(metadata: OBPMetadata, axes: Axes) -> None:
         )
     axes.set_xlabel("backpropagated slice number")
     axes.set_ylabel("# Pauli terms")
-    axes.legend()
+    _set_legend(axes, show_legend)
 
 
-def plot_num_truncated_paulis(metadata: OBPMetadata, axes: Axes) -> None:
+def plot_num_truncated_paulis(
+    metadata: OBPMetadata, axes: Axes, *, show_legend: bool = True
+) -> None:
     """Plot the number of truncated Pauli terms.
 
     This method populates the provided figure axes with a bar-plot of the number of the truncated
@@ -245,7 +250,7 @@ def plot_num_truncated_paulis(metadata: OBPMetadata, axes: Axes) -> None:
        >>> metadata = OBPMetadata.from_json("docs/_static/dummy_visualization_metadata.json")
 
     .. plot::
-       :context:
+       :context: close-figs
        :include-source:
 
        >>> from matplotlib import pyplot as plt
@@ -264,7 +269,7 @@ def plot_num_truncated_paulis(metadata: OBPMetadata, axes: Axes) -> None:
     Args:
         metadata: the metadata to be visualized.
         axes: the matplotlib axes in which to plot.
-
+        show_legend: enable/disable showing the legend in the plot.
     """
     num_observables = len(metadata.backpropagation_history[0].slice_errors)
     width = 0.8 / num_observables
@@ -281,10 +286,10 @@ def plot_num_truncated_paulis(metadata: OBPMetadata, axes: Axes) -> None:
         offset += width
     axes.set_xlabel("backpropagated slice number")
     axes.set_ylabel("# truncated Pauli terms")
-    axes.legend()
+    _set_legend(axes, show_legend)
 
 
-def plot_sum_paulis(metadata: OBPMetadata, axes: Axes) -> None:
+def plot_sum_paulis(metadata: OBPMetadata, axes: Axes, *, show_legend: bool = True) -> None:
     """Plot the total number of all Pauli terms.
 
     This method populates the provided figure axes with a line-plot of the total number of all Pauli
@@ -299,7 +304,7 @@ def plot_sum_paulis(metadata: OBPMetadata, axes: Axes) -> None:
        >>> metadata = OBPMetadata.from_json("docs/_static/dummy_visualization_metadata.json")
 
     .. plot::
-       :context:
+       :context: close-figs
        :include-source:
 
        >>> from matplotlib import pyplot as plt
@@ -319,7 +324,7 @@ def plot_sum_paulis(metadata: OBPMetadata, axes: Axes) -> None:
     Args:
         metadata: the metadata to be visualized.
         axes: the matplotlib axes in which to plot.
-
+        show_legend: enable/disable showing the legend in the plot.
     """
     if metadata.operator_budget.max_paulis is not None:
         axes.axhline(
@@ -337,10 +342,10 @@ def plot_sum_paulis(metadata: OBPMetadata, axes: Axes) -> None:
     )
     axes.set_xlabel("backpropagated slice number")
     axes.set_ylabel("total # of Pauli terms")
-    axes.legend()
+    _set_legend(axes, show_legend)
 
 
-def plot_num_qwc_groups(metadata: OBPMetadata, axes: Axes) -> None:
+def plot_num_qwc_groups(metadata: OBPMetadata, axes: Axes, *, show_legend: bool = True) -> None:
     """Plot the number of qubit-wise commuting Pauli groups.
 
     This method populates the provided figure axes with a line-plot of the number of qubit-wise
@@ -355,7 +360,7 @@ def plot_num_qwc_groups(metadata: OBPMetadata, axes: Axes) -> None:
        >>> metadata = OBPMetadata.from_json("docs/_static/dummy_visualization_metadata.json")
 
     .. plot::
-       :context:
+       :context: close-figs
        :include-source:
 
        >>> from matplotlib import pyplot as plt
@@ -371,7 +376,7 @@ def plot_num_qwc_groups(metadata: OBPMetadata, axes: Axes) -> None:
     Args:
         metadata: the metadata to be visualized.
         axes: the matplotlib axes in which to plot.
-
+        show_legend: enable/disable showing the legend in the plot.
     """
     if metadata.operator_budget.max_qwc_groups is not None:
         axes.axhline(
@@ -389,4 +394,9 @@ def plot_num_qwc_groups(metadata: OBPMetadata, axes: Axes) -> None:
     )
     axes.set_xlabel("backpropagated slice number")
     axes.set_ylabel("# of qubit-wise commuting Pauli groups")
-    axes.legend()
+    _set_legend(axes, show_legend)
+
+
+def _set_legend(axes: Axes, show_legend: bool) -> None:
+    if show_legend:  # pragma: no cover
+        axes.legend()

releasenotes/notes/add-param-to-hide-legend-d62ff17eef5e80f4.yaml

@@ -0,0 +1,7 @@
+---
+features:
+  - |
+    A new parameter ``show_legend`` has been added to each function in the
+    :mod:`qiskit_addon_obp.utils.visualization` module that can show or hide the
+    legend on a plot. The legend is shown by default. This can be useful when
+    the legend becomes long and obstructs the plot.