Repository: PennyLaneAI/pennylane
Pull request #4957 (2023-12-18 16:54:26)

MPL indicates returned statistics of mid circuit measurements
[sc-52256]


Here's what iterative QPE looks like.

<img width="1184" alt="Screenshot 2023-12-18 at 2 46 30 PM" src="https://github.com/PennyLaneAI/pennylane/assets/6364575/dfd9eee3-7858-49ee-a25c-d2ef6f6036f6">

Looks like I accidentally branched off the wrong branch, so the commit history is a little messed up, but I don't think its worth creating a fresh PR to tidy that up.


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4957?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`f496a63`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/f496a638f74a87a8b941e4b9cdbfaa55b2e6ea5c?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.50% compared to head [(`a80f6da`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4957?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.49%.
> Report is 1 commits behind head on master.


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


```diff
@@            Coverage Diff             @@
##           master    #4957      +/-   ##
==========================================
- Coverage   99.50%   99.49%   -0.01%     
==========================================
  Files         390      390              
  Lines       35496    35309     -187     
==========================================
- Hits        35320    35132     -188     
- Misses        176      177       +1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4957?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).


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 571b60080..eb143426c 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -15,6 +15,7 @@
   [(#4850)](https://github.com/PennyLaneAI/pennylane/pull/4850)
   [(#4917)](https://github.com/PennyLaneAI/pennylane/pull/4917)
   [(#4930)](https://github.com/PennyLaneAI/pennylane/pull/4930)
+  [(#4957)](https://github.com/PennyLaneAI/pennylane/pull/4957)
 
   Drawing of mid-circuit measurement capabilities including qubit reuse and reset,
   postselection, conditioning, and collecting statistics is supported.
diff --git a/pennylane/drawer/mpldrawer.py b/pennylane/drawer/mpldrawer.py
index e96acfae8..88f947242 100644
--- a/pennylane/drawer/mpldrawer.py
+++ b/pennylane/drawer/mpldrawer.py
@@ -16,6 +16,7 @@ This module contains the MPLDrawer class for creating circuit diagrams with matp
 """
 from collections.abc import Iterable
 import warnings
+from typing import Sequence
 
 has_mpl = True
 try:
@@ -267,13 +268,13 @@ class MPLDrawer:
         ## Creating figure and ax
 
         if figsize is None:
-            figsize = (self.n_layers + 3, self.n_wires + self._cwire_scaling * c_wires + 1)
+            figsize = (self.n_layers + 3, self.n_wires + self._cwire_scaling * c_wires + 1.5)
 
         self._fig = plt.figure(figsize=figsize)
         self._ax = self._fig.add_axes(
             [0, 0, 1, 1],
             xlim=(-2, self.n_layers + 1),
-            ylim=(-1, self.n_wires + self._cwire_scaling * c_wires),
+            ylim=(-1, self.n_wires + self._cwire_scaling * c_wires + 0.5),
             xticks=[],
             yticks=[],
         )
@@ -847,20 +848,29 @@ class MPLDrawer:
         if "zorder" not in lines_options:
             lines_options["zorder"] = 3
 
+        if not isinstance(wires, Sequence):
+            wires = (wires,)
+
+        wires = tuple(self._y(w) for w in wires)
+
+        box_min = min(wires)
+        box_max = max(wires)
+        box_center = (box_max + box_min) / 2.0
+
         x_loc = layer - self._box_length / 2.0 + self._pad
-        y_loc = wires - self._box_length / 2.0 + self._pad
+        y_loc = box_min - self._box_length / 2.0 + self._pad
 
         box = patches.FancyBboxPatch(
             (x_loc, y_loc),
             self._box_length - 2 * self._pad,
-            self._box_length - 2 * self._pad,
+            box_max - box_min + self._box_length - 2 * self._pad,
             boxstyle=self._boxstyle,
             **box_options,
         )
         self._ax.add_patch(box)
 
         arc = patches.Arc(
-            (layer, wires + 0.15 * self._box_length),
+            (layer, box_center + 0.15 * self._box_length),
             0.6 * self._box_length,
             0.55 * self._box_length,
             theta1=180,
@@ -871,7 +881,7 @@ class MPLDrawer:
 
         # can experiment with the specific numbers to make it look decent
         arrow_start_x = layer - 0.15 * self._box_length
-        arrow_start_y = wires + 0.3 * self._box_length
+        arrow_start_y = box_center + 0.3 * self._box_length
         arrow_width = 0.3 * self._box_length
         arrow_height = -0.5 * self._box_length
 
@@ -886,18 +896,21 @@ class MPLDrawer:
         )
         if text:
             self._ax.text(
-                layer + 0.05 * self._box_length, wires + 0.225, text, fontsize=(self.fontsize - 2)
+                layer + 0.05 * self._box_length,
+                box_center + 0.225,
+                text,
+                fontsize=(self.fontsize - 2),
             )
 
     def _y(self, wire):
         """Used for determining the correct y coordinate for classical wires.
         Classical wires should be enumerated starting at the number of quantum wires the drawer has.
         For example, if the drawer has ``3`` quantum wires, the first classical wire should be located at ``3``
-        which corresponds to a ``y`` coordinate of ``2.6``.
+        which corresponds to a ``y`` coordinate of ``2.9``.
         """
         if wire < self.n_wires:
             return wire
-        return self.n_wires + self._cwire_scaling * (wire - self.n_wires) - 0.4
+        return self.n_wires + self._cwire_scaling * (wire - self.n_wires)
 
     def classical_wire(self, layers, wires) -> None:
         """Draw a classical control line.
diff --git a/pennylane/drawer/tape_mpl.py b/pennylane/drawer/tape_mpl.py
index 463e8241d..4495de82f 100644
--- a/pennylane/drawer/tape_mpl.py
+++ b/pennylane/drawer/tape_mpl.py
@@ -172,18 +172,23 @@ def _(op: qml.ops.op_math.Conditional, drawer, layer, config) -> None:
 def _get_measured_wires(measurements, wires) -> set:
     measured_wires = set()
     for m in measurements:
-        # state and probs
-        if len(m.wires) == 0:
-            return wires
+        if not m.mv:
+            # state and probs
+            if len(m.wires) == 0:
+                return wires
 
-        for wire in m.wires:
-            measured_wires.add(wire)
+            for wire in m.wires:
+                measured_wires.add(wire)
     return measured_wires
 
 
 def _add_classical_wires(drawer, layers, wires):
     for cwire, (cwire_layers, layer_wires) in enumerate(zip(layers, wires), start=drawer.n_wires):
         xs, ys = [], []
+
+        len_diff = len(cwire_layers) - len(layer_wires)
+        if len_diff > 0:
+            layer_wires += [cwire] * len_diff
         for l, w in zip(cwire_layers, layer_wires):
             xs.extend([l, l, l])
             ys.extend([cwire, w, cwire])
@@ -191,6 +196,17 @@ def _add_classical_wires(drawer, layers, wires):
         drawer.classical_wire(xs, ys)
 
 
+def _get_measured_bits(measurements, bit_map, offset):
+    measured_bits = []
+    for m in measurements:
+        if isinstance(m.mv, list):
+            for mv in m.mv:
+                measured_bits += [bit_map[mcm] + offset for mcm in mv.measurements]
+        elif m.mv:
+            measured_bits += [bit_map[mcm] + offset for mcm in m.mv.measurements]
+    return measured_bits
+
+
 def _tape_mpl(tape, wire_order=None, show_all_wires=False, decimals=None, **kwargs):
     """Private function wrapped with styling."""
     wire_options = kwargs.get("wire_options", None)
@@ -211,7 +227,7 @@ def _tape_mpl(tape, wire_order=None, show_all_wires=False, decimals=None, **kwar
     n_layers = len(layers)
     n_wires = len(wire_map)
 
-    cwire_layers, cwire_wires = cwire_connections(layers, bit_map)
+    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
@@ -241,6 +257,10 @@ def _tape_mpl(tape, wire_order=None, show_all_wires=False, decimals=None, **kwar
     for wire in _get_measured_wires(tape.measurements, list(range(n_wires))):
         drawer.measure(n_layers, wire)
 
+    measured_bits = _get_measured_bits(tape.measurements, bit_map, drawer.n_wires)
+    if measured_bits:
+        drawer.measure(n_layers, measured_bits)
+
     return drawer.fig, drawer.ax
 
 


Test patch:
diff --git a/tests/drawer/test_mpldrawer.py b/tests/drawer/test_mpldrawer.py
index bf15be15f..c17a98f39 100644
--- a/tests/drawer/test_mpldrawer.py
+++ b/tests/drawer/test_mpldrawer.py
@@ -42,7 +42,7 @@ class TestInitialization:
         drawer = MPLDrawer(n_wires=n_wires, n_layers=n_layers)
 
         assert drawer.fig.get_figwidth() == (n_layers + 3)
-        assert drawer.fig.get_figheight() == (n_wires + 1)
+        assert drawer.fig.get_figheight() == (n_wires + 1.5)
 
         drawer = MPLDrawer(n_wires=n_wires, n_layers=n_layers)
 
@@ -644,6 +644,54 @@ class TestMeasure:
         assert len(drawer.ax.texts) == 0
         plt.close()
 
+    def test_measure_multiple_wires(self):
+        """Tests the measure method when multiple wires are provided."""
+
+        drawer = MPLDrawer(n_layers=1, n_wires=2)
+        drawer.measure(layer=0, wires=(0, 1))
+
+        box = drawer.ax.patches[0]
+        assert box.get_x() == -drawer._box_length / 2.0 + drawer._pad
+        assert box.get_y() == -drawer._box_length / 2.0 + drawer._pad
+        assert box.get_width() == drawer._box_length - 2 * drawer._pad
+        assert box.get_height() == drawer._box_length - 2 * drawer._pad + 1
+
+        arc = drawer.ax.patches[1]
+        assert arc.center == (0, 0.5 + 0.15 * drawer._box_length)
+        assert arc.theta1 == 180
+        assert arc.theta2 == 0
+        assert allclose(arc.height, 0.55 * drawer._box_length)
+        assert arc.width == 0.6 * drawer._box_length
+
+        arrow = drawer.ax.patches[2]
+        assert isinstance(arrow, FancyArrow)
+        assert len(drawer.ax.texts) == 0
+        plt.close()
+
+    def test_measure_classical_wires(self):
+        """Tests the measure method when multiple wires are provided."""
+
+        drawer = MPLDrawer(n_layers=1, n_wires=2, c_wires=2)
+        drawer.measure(layer=0, wires=(2, 3))
+
+        box = drawer.ax.patches[0]
+        assert box.get_x() == -drawer._box_length / 2.0 + drawer._pad
+        assert box.get_y() == 2 - drawer._box_length / 2.0 + drawer._pad
+        assert box.get_width() == drawer._box_length - 2 * drawer._pad
+        assert box.get_height() == drawer._box_length - 2 * drawer._pad + drawer._cwire_scaling
+
+        arc = drawer.ax.patches[1]
+        assert arc.center == (0, 2 + drawer._cwire_scaling / 2 + 0.15 * drawer._box_length)
+        assert arc.theta1 == 180
+        assert arc.theta2 == 0
+        assert allclose(arc.height, 0.55 * drawer._box_length)
+        assert arc.width == 0.6 * drawer._box_length
+
+        arrow = drawer.ax.patches[2]
+        assert isinstance(arrow, FancyArrow)
+        assert len(drawer.ax.texts) == 0
+        plt.close()
+
     def test_measure_text(self):
         """Test adding a postselection label to a measure box."""
         drawer = MPLDrawer(1, 1)
@@ -799,7 +847,7 @@ class TestClassicalWires:
 
         [_, cwire] = drawer.ax.lines
         assert cwire.get_xdata() == layers
-        assert cwire.get_ydata() == [0, 0.6, 0.6, 0]  # cwires are scaledc
+        assert cwire.get_ydata() == [0, 1, 1, 0]  # cwires are scaledc
 
         [pe1, pe2] = cwire.get_path_effects()
 
@@ -824,7 +872,7 @@ class TestClassicalWires:
         [_, eraser] = drawer.ax.lines
 
         assert eraser.get_xdata() == (0.8, 1)
-        assert eraser.get_ydata() == (0.85, 0.85)
+        assert eraser.get_ydata() == (1.25, 1.25)
         assert eraser.get_color() == plt.rcParams["figure.facecolor"]
         assert eraser.get_linewidth() == 3 * plt.rcParams["lines.linewidth"]
         plt.close()
@@ -838,7 +886,7 @@ class TestClassicalWires:
         [_, eraser] = drawer.ax.lines
 
         assert eraser.get_xdata() == (0.8, 1.2)
-        assert eraser.get_ydata() == (0.6, 0.6)
+        assert eraser.get_ydata() == (1, 1)
         assert eraser.get_color() == plt.rcParams["figure.facecolor"]
         assert eraser.get_linewidth() == 3 * plt.rcParams["lines.linewidth"]
         plt.close()
diff --git a/tests/drawer/test_tape_mpl.py b/tests/drawer/test_tape_mpl.py
index bd08f9115..7016e8ac4 100644
--- a/tests/drawer/test_tape_mpl.py
+++ b/tests/drawer/test_tape_mpl.py
@@ -830,7 +830,7 @@ class TestClassicalControl:
         [_, cwire] = ax.lines
 
         assert cwire.get_xdata() == [0, 0, 0, 1, 1, 1, 2, 2, 2]
-        assert cwire.get_ydata() == [0.6, 0, 0.6, 0.6, 0, 0.6, 0.6, 0, 0.6]
+        assert cwire.get_ydata() == [1, 0, 1, 1, 0, 1, 1, 0, 1]
 
         [pe1, pe2] = cwire.get_path_effects()
 
@@ -859,10 +859,10 @@ class TestClassicalControl:
         [_, _, cwire1, cwire2, eraser] = ax.lines
 
         assert cwire1.get_xdata() == [0, 0, 0, 2, 2, 2]
-        assert cwire1.get_ydata() == [1.6, 0, 1.6, 1.6, 0, 1.6]
+        assert cwire1.get_ydata() == [2, 0, 2, 2, 0, 2]
 
         assert cwire2.get_xdata() == [1, 1, 1, 2, 2, 2]
-        assert cwire2.get_ydata() == [1.85, 1, 1.85, 1.85, 0, 1.85]
+        assert cwire2.get_ydata() == [2.25, 1, 2.25, 2.25, 0, 2.25]
 
         for cwire in [cwire1, cwire2]:
             [pe1, pe2] = cwire.get_path_effects()
@@ -878,7 +878,7 @@ class TestClassicalControl:
             }
 
         assert eraser.get_xdata() == (1.8, 2)
-        assert eraser.get_ydata() == (1.6, 1.6)
+        assert eraser.get_ydata() == (2, 2)
         assert eraser.get_color() == plt.rcParams["figure.facecolor"]
         assert eraser.get_linewidth() == 3 * plt.rcParams["lines.linewidth"]
 
@@ -899,10 +899,10 @@ class TestClassicalControl:
         [_, _, cwire1, cwire2, eraser] = ax.lines
 
         assert cwire1.get_xdata() == [0, 0, 0, 2, 2, 2, 3, 3, 3]
-        assert cwire1.get_ydata() == [1.6, 0, 1.6, 1.6, 0, 1.6, 1.6, 0, 1.6]
+        assert cwire1.get_ydata() == [2, 0, 2, 2, 0, 2, 2, 0, 2]
 
         assert cwire2.get_xdata() == [1, 1, 1, 2, 2, 2, 4, 4, 4]
-        assert cwire2.get_ydata() == [1.85, 1, 1.85, 1.85, 0, 1.85, 1.85, 1, 1.85]
+        assert cwire2.get_ydata() == [2.25, 1, 2.25, 2.25, 0, 2.25, 2.25, 1, 2.25]
 
         for cwire in [cwire1, cwire2]:
             [pe1, pe2] = cwire.get_path_effects()
@@ -918,6 +918,72 @@ class TestClassicalControl:
             }
 
         assert eraser.get_xdata() == (1.8, 2.2)
-        assert eraser.get_ydata() == (1.6, 1.6)
+        assert eraser.get_ydata() == (2, 2)
         assert eraser.get_color() == plt.rcParams["figure.facecolor"]
         assert eraser.get_linewidth() == 3 * plt.rcParams["lines.linewidth"]
+
+    def test_single_mcm_measure(self):
+        """Test a final measurement of a mid circuit measurement."""
+
+        with qml.queuing.AnnotatedQueue() as q:
+            m0 = qml.measure(0)
+            qml.expval(m0)
+        _, ax = tape_mpl(qml.tape.QuantumScript.from_queue(q))
+
+        assert len(ax.patches) == 6  # two measurement boxes
+        assert ax.patches[3].get_x() == 1 - 0.75 / 2 + 0.2  # 1 - box_length/2 + pad
+        assert qml.math.allclose(ax.patches[3].get_y(), 1 - 0.75 / 2 + 0.2)  # 1- box_length/2 + pad
+
+        assert ax.patches[4].center == (
+            1,
+            1 + 0.15 * 0.75,
+        )  # 1 +0.15 *box_length
+        assert isinstance(ax.patches[5], mpl.patches.FancyArrow)
+
+        [_, cwire] = ax.lines
+        assert cwire.get_xdata() == [0, 0, 0, 1, 1, 1]
+        assert cwire.get_ydata() == [1, 0, 1, 1, 1, 1]
+
+        [pe1, pe2] = cwire.get_path_effects()
+
+        # probably not a good way to test this, but the best I can figure out
+        assert pe1._gc == {
+            "linewidth": 5 * plt.rcParams["lines.linewidth"],
+            "foreground": "black",  # lines.color for black white style
+        }
+        assert pe2._gc == {
+            "linewidth": 3 * plt.rcParams["lines.linewidth"],
+            "foreground": "white",  # figure.facecolor for black white sytle
+        }
+
+    def test_multiple_mcm_measure(self):
+        """Test final measurements of multiple mid circuit measurements"""
+        with qml.queuing.AnnotatedQueue() as q:
+            m0 = qml.measure(0)
+            m1 = qml.measure(0)
+            _ = qml.measure(0)
+            m2 = qml.measure(0)
+            _ = qml.measure(0)
+            qml.sample([m0, m1])
+            qml.expval(m2)
+        _, ax = qml.drawer.tape_mpl(qml.tape.QuantumScript.from_queue(q))
+
+        [_, cwire0, cwire1, cwire2] = ax.lines
+        assert cwire0.get_xdata() == [0, 0, 0, 5, 5, 5]
+        assert cwire0.get_ydata() == [1, 0, 1, 1, 1, 1]
+        assert cwire1.get_xdata() == [1, 1, 1, 5, 5, 5]
+        assert cwire1.get_ydata() == [1.25, 0, 1.25, 1.25, 1.25, 1.25]
+        assert cwire2.get_xdata() == [3, 3, 3, 5, 5, 5]
+        assert cwire2.get_ydata() == [1.5, 0, 1.5, 1.5, 1.5, 1.5]
+
+        assert len(ax.patches) == 18  # 6 * 3
+
+        final_measure_box = ax.patches[15]
+        assert final_measure_box.get_x() == 5 - 0.75 / 2 + 0.2  # 5 - box_length/2 + pad
+        assert qml.math.allclose(
+            final_measure_box.get_y(), 1 - 0.75 / 2 + 0.2
+        )  # 1- box_length/2 + pad
+
+        assert (
+            final_measure_box.get_height() == 0.75 - 2 * 0.2 + 2 * 0.25
+        )  # box_length - 2 * pad + 2 *cwire_scaling