Merge branch 'v0.40.0-docs' into v0.40.0-changelog

doc/releases/changelog-0.40.0.md

@@ -800,6 +800,11 @@ three-mode PES.
   [(#6530)](https://github.com/PennyLaneAI/pennylane/pull/6530)
   [(#6561)](https://github.com/PennyLaneAI/pennylane/pull/6561)
 
+* The developer-facing ``qml.drawer.MPLDrawer`` argument `n_wires` has been replaced with `wire_map`,
+  which contains more complete information about wire labels and order. This allows the new functionality 
+  to specify `wire_options` for specific wires when using string wire labels or non-sequential wire ordering.
+  [(#6805)](https://github.com/PennyLaneAI/pennylane/pull/6805)
+
 <h3>Deprecations 👋</h3>
 
 * The `tape` and `qtape` properties of `QNode` have been deprecated. Instead, use the 

pennylane/drawer/mpldrawer.py

@@ -25,6 +25,8 @@
 except (ModuleNotFoundError, ImportError) as e:  # pragma: no cover
     has_mpl = False
 
+# pylint: disable=too-many-positional-arguments
+
 
 def _to_tuple(a):
     """Converts int or iterable to tuple"""
@@ -63,20 +65,20 @@ class MPLDrawer:
 
     Args:
         n_layers (int): the number of layers
-        n_wires (int): the number of wires
+        wire_map (dict): the wires to be drawn. A dict mapping wire label to index (from top to bottom) in the figure
 
     Keyword Args:
         c_wires=0 (int): the number of classical wires to leave space for.
         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``.
+            to scale with the size of the circuit via ``n_layers`` and ``len(wire_map)``.
         fig=None (matplotlib Figure): Allows users to specify the figure window to plot to.
 
     **Example**
 
     .. code-block:: python
 
-        drawer = qml.drawer.MPLDrawer(n_wires=5, n_layers=6)
+        drawer = qml.drawer.MPLDrawer(wire_map={i: i for i in range(5)}, n_layers=6)
 
         drawer.label(["0", "a", r"$|\Psi\rangle$", r"$|\theta\rangle$", "aux"])
 
@@ -165,7 +167,7 @@ class MPLDrawer:
     .. code-block:: python
 
         wire_options = {"color": "indigo", "linewidth": 4}
-        drawer = MPLDrawer(n_wires=2, n_layers=4, wire_options=wire_options)
+        drawer = MPLDrawer(wire_map={0: 0, 1: 1}, n_layers=4, wire_options=wire_options)
 
         label_options = {"fontsize": "x-large", 'color': 'indigo'}
         drawer.label(["0", "a"], text_options=label_options)
@@ -203,7 +205,7 @@ class MPLDrawer:
 
     .. code-block:: python
 
-        drawer = MPLDrawer(2, 2)
+        drawer = MPLDrawer(2, {0:0, 1:1})
         drawer.box_gate(layer=0, wires=1, text="X")
         drawer.box_gate(layer=1, wires=1, text="Y")
 
@@ -255,15 +257,15 @@ 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, fig=None):
+    def __init__(self, n_layers, wire_map, 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. "
                 "You can install matplotlib via \n\n   pip install matplotlib"
             )
 
         self.n_layers = n_layers
-        self.n_wires = n_wires
+        self.n_wires = len(wire_map)
 
         ## Creating figure and ax
 
@@ -300,14 +302,14 @@ def __init__(self, n_layers, n_wires, c_wires=0, wire_options=None, figsize=None
 
         # Adding wire lines with individual styles based on wire_options
         self._wire_lines = []
-        for wire in range(self.n_wires):
-            specific_options = wire_specific_options.get(wire, {})
+        for wire_label, idx in wire_map.items():
+            specific_options = wire_specific_options.get(wire_label, {})
             line_options = {**global_options, **specific_options}
 
             # Create Line2D with the combined options
             line = plt.Line2D(
                 (-1, self.n_layers),
-                (wire, wire),
+                (idx, idx),
                 zorder=1,
                 **line_options,
             )
@@ -349,7 +351,7 @@ def label(self, labels, text_options=None):
 
         .. code-block:: python
 
-            drawer = MPLDrawer(n_wires=2, n_layers=1)
+            drawer = MPLDrawer(wire_map={0:0, 1:1}, n_layers=1)
             drawer.label(["a", "b"])
 
         .. figure:: ../../_static/drawer/labels.png
@@ -363,7 +365,7 @@ def label(self, labels, text_options=None):
 
         .. code-block:: python
 
-            drawer = MPLDrawer(n_wires=2, n_layers=1)
+            drawer = MPLDrawer(wire_map={0:0, 1:1}, n_layers=1)
             drawer.label(["a", "b"], text_options={"color": "indigo", "fontsize": "xx-large"})
 
         .. figure:: ../../_static/drawer/labels_formatted.png
@@ -419,7 +421,7 @@ def box_gate(self, layer, wires, text="", box_options=None, text_options=None, *
 
         .. code-block:: python
 
-            drawer = MPLDrawer(n_wires=2, n_layers=1)
+            drawer = MPLDrawer(wire_map={0:0, 1:1}, n_layers=1)
 
             drawer.box_gate(layer=0, wires=(0, 1), text="CY")
 
@@ -441,7 +443,7 @@ def box_gate(self, layer, wires, text="", box_options=None, text_options=None, *
             box_options = {'facecolor': 'lightcoral', 'edgecolor': 'maroon', 'linewidth': 5}
             text_options = {'fontsize': 'xx-large', 'color': 'maroon'}
 
-            drawer = MPLDrawer(n_wires=2, n_layers=1)
+            drawer = MPLDrawer(wire_map={0:0, 1:1}, n_layers=1)
 
             drawer.box_gate(layer=0, wires=(0, 1), text="CY",
                 box_options=box_options, text_options=text_options)
@@ -456,7 +458,7 @@ def box_gate(self, layer, wires, text="", box_options=None, text_options=None, *
 
         .. code-block:: python
 
-            drawer = MPLDrawer(n_layers=4, n_wires=2)
+            drawer = MPLDrawer(n_layers=4, wire_map={0:0, 1:1})
 
             drawer.box_gate(layer=0, wires=0, text="A longer label")
             drawer.box_gate(layer=0, wires=1, text="Label")
@@ -615,7 +617,7 @@ def ctrl(self, layer, wires, wires_target=None, control_values=None, options=Non
 
         .. code-block:: python
 
-            drawer = MPLDrawer(n_wires=2, n_layers=3)
+            drawer = MPLDrawer(wire_map={0:0, 1:1}, n_layers=3)
 
             drawer.ctrl(layer=0, wires=0, wires_target=1)
             drawer.ctrl(layer=1, wires=(0, 1), control_values=[0, 1])
@@ -714,7 +716,7 @@ def CNOT(self, layer, wires, control_values=None, options=None):
 
         .. code-block:: python
 
-            drawer = MPLDrawer(n_wires=2, n_layers=2)
+            drawer = MPLDrawer(wire_map={0:0, 1:1}, n_layers=2)
 
             drawer.CNOT(0, (0, 1))
 
@@ -778,7 +780,7 @@ def SWAP(self, layer, wires, options=None):
 
         .. code-block:: python
 
-            drawer = MPLDrawer(n_wires=2, n_layers=2)
+            drawer = MPLDrawer(wire_map={0:0, 1:1}, n_layers=2)
 
             drawer.SWAP(0, (0, 1))
 
@@ -848,7 +850,7 @@ def measure(self, layer, wires, text=None, box_options=None, lines_options=None)
 
         .. code-block:: python
 
-            drawer = MPLDrawer(n_wires=2, n_layers=1)
+            drawer = MPLDrawer(wire_map={0:0, 1:1}, n_layers=1)
             drawer.measure(layer=0, wires=0)
 
             measure_box = {'facecolor': 'white', 'edgecolor': 'indigo'}
@@ -996,7 +998,7 @@ def cond(self, layer, measured_layer, wires, wires_target, options=None):
 
         .. code-block:: python
 
-            drawer = MPLDrawer(n_wires=3, n_layers=4)
+            drawer = MPLDrawer(wire_map={0:0, 1:1, 2:2}, n_layers=4)
 
             drawer.cond(layer=1, measured_layer=0, wires=[0], wires_target=[1])
 

pennylane/drawer/tape_mpl.py

@@ -241,7 +241,7 @@ def _tape_mpl(tape, wire_order=None, show_all_wires=False, decimals=None, *, fig
 
     drawer = MPLDrawer(
         n_layers=n_layers,
-        n_wires=n_wires,
+        wire_map=wire_map,
         c_wires=len(bit_map),
         wire_options=wire_options,
         fig=fig,

tests/drawer/test_draw_mpl.py

@@ -319,7 +319,7 @@ def test_wire_order_not_on_device(self):
         assert ax.texts[2].get_text() == "0"
         plt.close()
 
-    def test_wire_options(self):
+    def test_uniform_wire_options(self):
         """Test wire options modifies wire styling"""
 
         _, ax = qml.draw_mpl(circuit1, wire_options={"color": "black", "linewidth": 4})(1.23, 2.34)
@@ -328,6 +328,11 @@ def test_wire_options(self):
             assert w.get_color() == "black"
             assert w.get_linewidth() == 4
 
+        plt.close()
+
+    def test_individual_wire_options(self):
+        """Test wire option styling when individual wires have their own options specified"""
+
         @qml.qnode(dev)
         def f_circ(x):
             """Circuit on ten qubits."""
@@ -399,6 +404,58 @@ def f_circ(x):
 
         plt.close()
 
+    def test_individual_wire_options_with_string_labels(self):
+        """Test that individual wire options work with string wire labels"""
+
+        @qml.qnode(qml.device("default.qubit"))
+        def circuit():
+            qml.X("a")
+            qml.Y("b")
+            return qml.expval(qml.Z("a"))
+
+        wire_options = {
+            "color": "teal",
+            "linewidth": 5,
+            "b": {"color": "orange", "linestyle": "--"},
+        }
+        _, ax = qml.draw_mpl(circuit, wire_options=wire_options)()
+
+        for i, w in enumerate(ax.lines):
+            assert w.get_linewidth() == 5
+            if i == 0:
+                assert w.get_color() == "teal"
+                assert w.get_linestyle() == "-"
+            if i == 1:
+                assert w.get_color() == "orange"
+                assert w.get_linestyle() == "--"
+
+    def test_wire_options_and_wire_order(self):
+        """Test that individual wire options work with specifying a wire_order"""
+
+        device = qml.device("default.qubit", wires=4)
+
+        @qml.qnode(device)
+        def circuit():
+            for w in device.wires:
+                qml.X(w)
+            return qml.expval(qml.Z(0))
+
+        wire_options = {
+            "color": "teal",
+            "linewidth": 5,
+            3: {"color": "orange", "linestyle": "--"},  # wire 3 should be orange and dashed
+        }
+        _, ax = qml.draw_mpl(circuit, wire_order=[1, 3, 0, 2], wire_options=wire_options)()
+
+        for i, w in enumerate(ax.lines):
+            assert w.get_linewidth() == 5
+            if i == 1:
+                assert w.get_color() == "orange"
+                assert w.get_linestyle() == "--"
+            else:
+                assert w.get_color() == "teal"
+                assert w.get_linestyle() == "-"
+
 
 class TestMPLIntegration:
     """Test using matplotlib styling to modify look of graphic."""