Repository: PennyLaneAI/pennylane
Pull request #4956 (2023-12-17 21:20:04)

Allow specification of matplotlib figure for circuit drawing

**Context:**

The current implementation of `qmpl.draw_mpl` always creates a new figure window. In this PR we allow the figure window to be specified. For example

```
import pennylane
import pennylane as qml
import matplotlib.pyplot as plt

dev = qml.device("default.qubit", wires=2)

@qml.qnode(dev)
def circuit(params):
    qml.RY(params[0], wires=0)
    qml.RY(params[1], wires=1)
    qml.CNOT(wires=[0, 1])
    return qml.expval(qml.PauliZ(1))

params = [0.1, 0.3]

qml.draw_mpl(circuit, decimals=2, fig=plt.figure(10))(params) # set figure window number 10
# ...
qml.draw_mpl(circuit, decimals=None, fig=plt.figure(10))(params)
```

**Description of the Change:**

* We have chosen to pass the matplotlib figure and not the matplotlib axes as figure properties are used in the current implementation. See for example `MPLDrawer._text_dims`.
* Since `draw_mpl` sets the figure size explicitely we also do this when using an existing figure
* Use `fig` as a keyword argument only

@josh146 

**Benefits:**

Re-using figure windows prevents many open windows in interactive environments such as Spyder

**Possible Drawbacks:**

PR is fully backwards compatible.

**Related GitHub Issues:**

#4950


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4956?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: `4 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`1e22d48`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/1e22d48691e8ec8dd11817dff66f12cd6fcdafb1?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.66% compared to head [(`577812d`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4956?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.48%.
> Report is 18 commits behind head on master.

| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4956?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Patch % | Lines |
|---|---|---|
| [pennylane/drawer/mpldrawer.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4956?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RyYXdlci9tcGxkcmF3ZXIucHk=) | 42.85% | [4 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4956?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #4956      +/-   ##
==========================================
- Coverage   99.66%   99.48%   -0.19%     
==========================================
  Files         388      390       +2     
  Lines       35339    35275      -64     
==========================================
- Hits        35222    35094     -128     
- Misses        117      181      +64     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4956?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

josh146: Thanks @eendebakpt! I will organize a code review for this PR. Out of curiosity, how come you decided to pass `fig` as a keyword argument, as opposed to `ax`? My understanding was that passing axes would provide more flexibility/is more common
eendebakpt: > Thanks @eendebakpt! I will organize a code review for this PR. Out of curiosity, how come you decided to pass `fig` as a keyword argument, as opposed to `ax`? My understanding was that passing axes would provide more flexibility/is more common

Passing `ax` is indeed more flexible. The reason for passsing `fig` is that the `MPLDrawer` requires a matplotlib figure (the `.fig` property). It is used to set the figure size (here https://github.com/PennyLaneAI/pennylane/blob/master/pennylane/drawer/mpldrawer.py#L272) and some text properties
(here https://github.com/PennyLaneAI/pennylane/blob/master/pennylane/drawer/mpldrawer.py#L564).

I am not familiar with the design and history of the drawing functionality, but perhaps with a little refactoring passing the `ax` instead is also possible.

eendebakpt: > Thanks for opening this PR. rThis seems like it could be quite a nice feature 🎉
> 
> As you pointed out in the PR description, we do use the figure to autosize text. I do see a couple of options that would allow us to accept the axes object. We could accept both the figure _and_ the axes object. Then we could either:
> 
> 1. Error if axes are provided without a figure.
> 2. Turn off text auto-sizing and warn the user if axes are provided without a figure.
> 
> The autosizing is just to make the text look better if a gate has a particularily long label. Most of the time, it won't kick in and make a difference.
> 
> Minor housekeeping tasks:
> 
> * [ ]   Add a [changelog](https://github.com/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) entry and add your name to the list of contributors
> * [ ]  Add a short test to [test_draw_mpl.py](https://github.com/PennyLaneAI/pennylane/blob/master/tests/drawer/test_draw_mpl.py)
> * [ ]  Add a short test to [test_tape_mpl.py](https://github.com/PennyLaneAI/pennylane/blob/master/tests/drawer/test_tape_mpl.py)
> * [ ]  Add a test to [test_mpldrawer.py](https://github.com/PennyLaneAI/pennylane/blob/master/tests/drawer/test_mpldrawer.py)
> * [ ]  Add `fix` to the list of keyword arguments for `MPLDrawer`: https://github.com/PennyLaneAI/pennylane/blob/f496a638f74a87a8b941e4b9cdbfaa55b2e6ea5c/pennylane/drawer/mpldrawer.py#L71
> 
> Tests can simply be checking the outputted figure is the same object as the inputted figure.

Housekeeping tasks have been done. I decided to keep the `ax` argument out for now. If there is a use case, we could add it (with either one of your suggestions of handling the resizing)

Patch:
diff --git a/doc/releases/changelog-0.34.0.md b/doc/releases/changelog-0.34.0.md
index cadafb0c7..100820ab5 100644
--- a/doc/releases/changelog-0.34.0.md
+++ b/doc/releases/changelog-0.34.0.md
@@ -336,6 +336,9 @@
 * The function ``qml.Snapshot`` now supports arbitrary measurements of type ``StateMeasurement``.
   [(#4876)](https://github.com/PennyLaneAI/pennylane/pull/4908)
 
+* The function ``qml.draw_mpl`` now accept a keyword argument ``fig`` to specify the output figure window.
+  [(#4956)](https://github.com/PennyLaneAI/pennylane/pull/4956)
+
 <h4>Better support for batching</h4>
 
 * `default.qubit` now can evolve already batched states with `ParametrizedEvolution`
diff --git a/pennylane/drawer/draw.py b/pennylane/drawer/draw.py
index 66f0a7cf3..1fb8fecdf 100644
--- a/pennylane/drawer/draw.py
+++ b/pennylane/drawer/draw.py
@@ -315,6 +315,8 @@ def draw_mpl(
     decimals=None,
     expansion_strategy=None,
     style=None,
+    *,
+    fig=None,
     **kwargs,
 ):
     """Draw a qnode with matplotlib
@@ -348,7 +350,7 @@ def draw_mpl(
 
             - ``device``: The QNode will attempt to decompose the internal circuit
               such that all circuit operations are natively supported by the device.
-
+        fig (None or matplotlib.Figure): Matplotlib figure to plot onto. If None, then create a new figure
 
     Returns:
         A function that has the same argument signature as ``qnode``. When called,
@@ -523,6 +525,7 @@ def draw_mpl(
             decimals=decimals,
             expansion_strategy=expansion_strategy,
             style=style,
+            fig=fig,
             **kwargs,
         )
 
@@ -543,6 +546,7 @@ def draw_mpl(
             show_all_wires=show_all_wires,
             decimals=decimals,
             style=style,
+            fig=fig,
             **kwargs,
         )
 
@@ -556,6 +560,8 @@ def _draw_mpl_qnode(
     decimals=None,
     expansion_strategy=None,
     style="black_white",
+    *,
+    fig=None,
     **kwargs,
 ):
     @wraps(qnode)
@@ -587,6 +593,7 @@ def _draw_mpl_qnode(
             show_all_wires=show_all_wires,
             decimals=decimals,
             style=style,
+            fig=fig,
             **kwargs,
         )
 
diff --git a/pennylane/drawer/mpldrawer.py b/pennylane/drawer/mpldrawer.py
index 88f947242..f5f9894c9 100644
--- a/pennylane/drawer/mpldrawer.py
+++ b/pennylane/drawer/mpldrawer.py
@@ -71,6 +71,7 @@ class MPLDrawer:
         wire_options=None (dict): matplotlib configuration options for drawing the wire lines
         figsize=None (Iterable): Allows users to specify the size of the figure manually. Defaults
             to scale with the size of the circuit via ``n_layers`` and ``n_wires``.
+        fig=None (matplotlib Figure): Allows users to specify the figure window to plot to.
 
     **Example**
 
@@ -255,7 +256,7 @@ class MPLDrawer:
     _cwire_scaling = 0.25
     """The distance between successive control wires."""
 
-    def __init__(self, n_layers, n_wires, c_wires=0, wire_options=None, figsize=None):
+    def __init__(self, n_layers, n_wires, c_wires=0, wire_options=None, figsize=None, fig=None):
         if not has_mpl:  # pragma: no cover
             raise ImportError(
                 "Module matplotlib is required for ``MPLDrawer`` class. "
@@ -270,7 +271,14 @@ class MPLDrawer:
         if figsize is None:
             figsize = (self.n_layers + 3, self.n_wires + self._cwire_scaling * c_wires + 1.5)
 
-        self._fig = plt.figure(figsize=figsize)
+        if fig is None:
+            self._fig = plt.figure(figsize=figsize)
+        else:
+            fig.clear()
+            fig.set_figwidth(figsize[0])
+            fig.set_figheight(figsize[1])
+            self._fig = fig
+
         self._ax = self._fig.add_axes(
             [0, 0, 1, 1],
             xlim=(-2, self.n_layers + 1),
@@ -278,6 +286,7 @@ class MPLDrawer:
             xticks=[],
             yticks=[],
         )
+
         self._ax.axis("off")
 
         self._ax.invert_yaxis()
diff --git a/pennylane/drawer/tape_mpl.py b/pennylane/drawer/tape_mpl.py
index 4495de82f..8dd4b0b72 100644
--- a/pennylane/drawer/tape_mpl.py
+++ b/pennylane/drawer/tape_mpl.py
@@ -207,7 +207,7 @@ def _get_measured_bits(measurements, bit_map, offset):
     return measured_bits
 
 
-def _tape_mpl(tape, wire_order=None, show_all_wires=False, decimals=None, **kwargs):
+def _tape_mpl(tape, wire_order=None, show_all_wires=False, decimals=None, *, fig=None, **kwargs):
     """Private function wrapped with styling."""
     wire_options = kwargs.get("wire_options", None)
     label_options = kwargs.get("label_options", None)
@@ -230,7 +230,11 @@ def _tape_mpl(tape, wire_order=None, show_all_wires=False, decimals=None, **kwar
     cwire_layers, cwire_wires = cwire_connections(layers + [tape.measurements], bit_map)
 
     drawer = MPLDrawer(
-        n_layers=n_layers, n_wires=n_wires, c_wires=len(bit_map), wire_options=wire_options
+        n_layers=n_layers,
+        n_wires=n_wires,
+        c_wires=len(bit_map),
+        wire_options=wire_options,
+        fig=fig,
     )
 
     config = _Config(
@@ -264,7 +268,9 @@ def _tape_mpl(tape, wire_order=None, show_all_wires=False, decimals=None, **kwar
     return drawer.fig, drawer.ax
 
 
-def tape_mpl(tape, wire_order=None, show_all_wires=False, decimals=None, style=None, **kwargs):
+def tape_mpl(
+    tape, wire_order=None, show_all_wires=False, decimals=None, style=None, *, fig=None, **kwargs
+):
     """Produces a matplotlib graphic from a tape.
 
     Args:
@@ -287,6 +293,7 @@ def tape_mpl(tape, wire_order=None, show_all_wires=False, decimals=None, style=N
         label_options (dict): matplotlib formatting options for the wire labels
         active_wire_notches (bool): whether or not to add notches indicating active wires.
             Defaults to ``True``.
+        fig (None or matplotlib Figure): Matplotlib figure to plot onto. If None, then create a new figure.
 
     Returns:
         matplotlib.figure.Figure, matplotlib.axes._axes.Axes: The key elements for matplotlib's object oriented interface.
@@ -444,7 +451,12 @@ def tape_mpl(tape, wire_order=None, show_all_wires=False, decimals=None, style=N
         _set_style(style)
     try:
         return _tape_mpl(
-            tape, wire_order=wire_order, show_all_wires=show_all_wires, decimals=decimals, **kwargs
+            tape,
+            wire_order=wire_order,
+            show_all_wires=show_all_wires,
+            decimals=decimals,
+            fig=fig,
+            **kwargs,
         )
     finally:
         if update_style:


Test patch:
diff --git a/tests/drawer/test_draw_mpl.py b/tests/drawer/test_draw_mpl.py
index d6c2d68ec..f53e1511f 100644
--- a/tests/drawer/test_draw_mpl.py
+++ b/tests/drawer/test_draw_mpl.py
@@ -70,6 +70,15 @@ def test_standard_use():
     plt.close()
 
 
+def test_fig_argument():
+    """Tests figure argument is used correcly"""
+
+    fig = plt.figure()
+    output_fig, ax = qml.draw_mpl(circuit1, fig=fig)(1.23, 2.34)
+    assert ax.get_figure() == fig
+    assert output_fig == fig
+
+
 @pytest.mark.parametrize(
     "device",
     [qml.device("default.qubit.legacy", wires=3), qml.devices.DefaultQubit(wires=3)],
diff --git a/tests/drawer/test_mpldrawer.py b/tests/drawer/test_mpldrawer.py
index c17a98f39..cf2da7d60 100644
--- a/tests/drawer/test_mpldrawer.py
+++ b/tests/drawer/test_mpldrawer.py
@@ -64,6 +64,15 @@ class TestInitialization:
         assert drawer.fig.get_figheight() == 5
         plt.close()
 
+    def test_customfigure(self):
+        """Tests a custom figure is used"""
+
+        fig = plt.figure()
+        drawer = MPLDrawer(1, 1, fig=fig)
+
+        assert drawer.fig == fig
+        plt.close()
+
     def test_config_params_set(self):
         """Tests sizing hidden variables are set."""
 
diff --git a/tests/drawer/test_tape_mpl.py b/tests/drawer/test_tape_mpl.py
index 7016e8ac4..2baab80b4 100644
--- a/tests/drawer/test_tape_mpl.py
+++ b/tests/drawer/test_tape_mpl.py
@@ -58,6 +58,17 @@ def test_fontsize():
     plt.close()
 
 
+def test_fig_argument():
+    """Test figure argument is used correctly"""
+
+    fig = plt.figure()
+    output_fig, ax = tape_mpl(tape1, fontsize=20, fig=fig)
+
+    assert ax.get_figure() == fig
+    assert output_fig == fig
+    plt.close()
+
+
 label_data = [
     ({}, ["0", "a", "1.234"]),  # default behaviour
     ({"wire_order": [1.234, "a", 0]}, ["1.234", "a", "0"]),  # provide standard wire order