Repository: PennyLaneAI/pennylane
Pull request #6805 (2025-01-09 22:39:36)

Bugfix for wire_options
**Context:**
When using wire labels that aren't integers ordered from largest to smallest (like strings, or a wire_order that isn't "in order"), the custom wire options are applied to the incorrect wire.

```
dev = qml.device('default.qubit', wires=4)

@qml.qnode(dev)
def circuit():
    for w in dev.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
}
fig, ax = qml.draw_mpl(circuit, wire_order=[1, 3, 0, 2], wire_options=wire_options)()
```

This should make wire 3 orange, but instead makes the wire at _index_ 3 orange:
![image](https://github.com/user-attachments/assets/193925cb-ee97-43e7-9d3a-22b0da700589)

**Description of the Change:**
We update it to fix the bug. This involves modifying the signature of `MPLDrawer`.

**Benefits:**
It works as expected now.

**Possible Drawbacks:**
While MPLDrawer is internal, and even other PL libraries should be using `draw_mpl` (whose signature is unchanged), this is technically a small breaking change. However, a search in Catalyst, Lightning and QML confirms that none of our other packages are calling this, and the only tests that needed modification were the ones directly testing the class, so it seems safe to say we've kept it very isolated and this shouldn't cause any problems.

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6805?dropdown=coverage&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:
> Please [upload](https://docs.codecov.com/docs/codecov-uploader) report for BASE (`v0.40.0-rc0@a5449d7`). [Learn more](https://docs.codecov.io/docs/error-reference?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#section-missing-base-commit) about missing BASE report.

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


```diff
@@              Coverage Diff               @@
##             v0.40.0-rc0    #6805   +/-   ##
==============================================
  Coverage               ?   99.60%           
==============================================
  Files                  ?      476           
  Lines                  ?    45222           
  Branches               ?        0           
==============================================
  Hits                   ?    45045           
  Misses                 ?      177           
  Partials               ?        0           
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6805?dropdown=coverage&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-0.40.0.md b/doc/releases/changelog-0.40.0.md
index 7d9707882..d8da41c0d 100644
--- a/doc/releases/changelog-0.40.0.md
+++ b/doc/releases/changelog-0.40.0.md
@@ -575,6 +575,11 @@ following 4 sets of functions have been either moved or removed[(#6588)](https:/
   [(#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.
diff --git a/pennylane/drawer/mpldrawer.py b/pennylane/drawer/mpldrawer.py
index 64e92a50c..ce559fc02 100644
--- a/pennylane/drawer/mpldrawer.py
+++ b/pennylane/drawer/mpldrawer.py
@@ -25,6 +25,8 @@ try:
 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"""
@@ -76,7 +78,7 @@ class MPLDrawer:
 
     .. 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,7 +257,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, 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. "
@@ -263,7 +265,7 @@ class MPLDrawer:
             )
 
         self.n_layers = n_layers
-        self.n_wires = n_wires
+        self.n_wires = len(wire_map)
 
         ## Creating figure and ax
 
@@ -300,14 +302,14 @@ class MPLDrawer:
 
         # 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 @@ class MPLDrawer:
 
         .. 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 @@ class MPLDrawer:
 
         .. 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 @@ class MPLDrawer:
 
         .. 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 @@ class MPLDrawer:
             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 @@ class MPLDrawer:
 
         .. 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 @@ class MPLDrawer:
 
         .. 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 @@ class MPLDrawer:
 
         .. 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 @@ class MPLDrawer:
 
         .. 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 @@ class MPLDrawer:
 
         .. 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 @@ class MPLDrawer:
 
         .. 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])
 
diff --git a/pennylane/drawer/tape_mpl.py b/pennylane/drawer/tape_mpl.py
index b1aa8f93a..a41104e57 100644
--- a/pennylane/drawer/tape_mpl.py
+++ b/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,


Test patch:
diff --git a/tests/drawer/test_draw_mpl.py b/tests/drawer/test_draw_mpl.py
index 23ca78d19..c8c162cce 100644
--- a/tests/drawer/test_draw_mpl.py
+++ b/tests/drawer/test_draw_mpl.py
@@ -319,7 +319,7 @@ class TestWireBehaviour:
         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 @@ class TestWireBehaviour:
             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 @@ class TestWireBehaviour:
 
         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."""
diff --git a/tests/drawer/test_mpldrawer.py b/tests/drawer/test_mpldrawer.py
index b28e6070b..1dd5f00a2 100644
--- a/tests/drawer/test_mpldrawer.py
+++ b/tests/drawer/test_mpldrawer.py
@@ -35,17 +35,19 @@ from pennylane.math import allclose
 class TestInitialization:
     """Tests drawer creation"""
 
-    @pytest.mark.parametrize("n_wires", [2, 3])
+    @pytest.mark.parametrize("wire_map", [{0: 0, 1: 1}, {"a": 0, "b": 1, "c": 2}])
     @pytest.mark.parametrize("n_layers", [2, 3])
-    def test_figsize_wires(self, n_wires, n_layers):
+    def test_figsize_wires(self, wire_map, n_layers):
         """Tests the figure is sized correctly."""
 
-        drawer = MPLDrawer(n_wires=n_wires, n_layers=n_layers)
+        n_wires = len(wire_map)
+
+        drawer = MPLDrawer(wire_map=wire_map, n_layers=n_layers)
 
         assert drawer.fig.get_figwidth() == (n_layers + 3)
         assert drawer.fig.get_figheight() == (n_wires + 1)
 
-        drawer = MPLDrawer(n_wires=n_wires, n_layers=n_layers)
+        drawer = MPLDrawer(wire_map=wire_map, n_layers=n_layers)
 
         lines = drawer.ax.lines
 
@@ -62,7 +64,9 @@ class TestInitialization:
         c_wires = 4
         n_layers = 1
 
-        drawer = MPLDrawer(n_wires=n_wires, n_layers=n_layers, c_wires=c_wires)
+        drawer = MPLDrawer(
+            wire_map={i: i for i in range(n_wires)}, n_layers=n_layers, c_wires=c_wires
+        )
 
         assert drawer.fig.get_figheight() == n_wires + 1 + 0.25 * 4 + 0.5
         assert drawer.fig.get_figwidth() == n_layers + 3
@@ -73,7 +77,7 @@ class TestInitialization:
     def test_customfigsize(self):
         """Tests a custom figsize alters the size"""
 
-        drawer = MPLDrawer(1, 1, figsize=(5, 5))
+        drawer = MPLDrawer(1, {0: 0}, figsize=(5, 5))
 
         assert drawer.fig.get_figwidth() == 5
         assert drawer.fig.get_figheight() == 5
@@ -83,7 +87,7 @@ class TestInitialization:
         """Tests a custom figure is used"""
 
         fig = plt.figure()
-        drawer = MPLDrawer(1, 1, fig=fig)
+        drawer = MPLDrawer(1, {0: 0}, fig=fig)
 
         assert drawer.fig == fig
         plt.close()
@@ -91,7 +95,7 @@ class TestInitialization:
     def test_config_params_set(self):
         """Tests sizing hidden variables are set."""
 
-        drawer = MPLDrawer(1, 1)
+        drawer = MPLDrawer(1, {0: 0})
 
         assert drawer._box_length == 0.75
         assert drawer._circ_rad == 0.3
@@ -111,7 +115,7 @@ class TestInitialization:
 
         rgba_red = (1, 0, 0, 1)
         options = {"linewidth": 3, "color": rgba_red}
-        drawer = MPLDrawer(n_wires=2, n_layers=2, wire_options=options)
+        drawer = MPLDrawer(wire_map={0: 0, 1: 1}, n_layers=2, wire_options=options)
 
         for line in drawer.ax.lines:
             assert line.get_linewidth() == 3
@@ -121,7 +125,7 @@ class TestInitialization:
 
     def test_fontsize(self):
         """Test fontsize can be get and set via property."""
-        drawer = MPLDrawer(1, 1)
+        drawer = MPLDrawer(1, {0: 0})
 
         assert drawer._fontsize == drawer.fontsize
         assert drawer.fontsize == 14
@@ -139,9 +143,11 @@ class TestLabels:
     def test_labels(self):
         """Tests labels are added"""
 
-        drawer = MPLDrawer(1, 3)
+        wire_map = {"a": 0, "b": 1, "c": 2}
+
+        drawer = MPLDrawer(1, wire_map)
 
-        labels = ("a", "b", "c")
+        labels = wire_map.keys()
         drawer.label(labels)
 
         drawn_labels = drawer.ax.texts
@@ -156,10 +162,12 @@ class TestLabels:
     def test_labels_formatting(self):
         """Test labels are formatted with text options."""
 
-        drawer = MPLDrawer(1, 3)
+        wire_map = {0: 0, 1: 1, 2: 2}
+
+        drawer = MPLDrawer(1, wire_map)
 
         rgba_red = (1, 0, 0, 1)
-        labels = (0, 1, 2)
+        labels = wire_map.keys()
         options = {"fontsize": 10, "color": rgba_red}
         drawer.label(labels, text_options=options)
 
@@ -172,9 +180,12 @@ class TestLabels:
     @pytest.mark.parametrize("n_layers", [1, 4])
     def test_crop_wire_labels(self, n_layers):
         """Test that cropping wire labels works."""
-        drawer = MPLDrawer(n_layers, 3)
 
-        labels = ("a", "b", "c")
+        wire_map = {"a": 0, "b": 1, "c": 2}
+
+        drawer = MPLDrawer(n_layers, wire_map)
+
+        labels = wire_map.keys()
         drawer.label(labels)
         old_width = drawer.fig.get_figwidth()
         assert old_width == n_layers + 3
@@ -190,7 +201,7 @@ class TestLabels:
 def test_erase_wire():
     """Test the erase wire method."""
 
-    drawer = MPLDrawer(5, 1)
+    drawer = MPLDrawer(5, {0: 0})
     drawer.erase_wire(1, 0, 3)
 
     assert len(drawer.ax.patches) == 1
@@ -208,7 +219,7 @@ class TestBoxGate:
     def test_simple_box(self):
         """tests basic functionality of box_gate."""
 
-        drawer = MPLDrawer(1, 1)
+        drawer = MPLDrawer(1, {0: 0})
 
         drawer.box_gate(0, 0, "X")
 
@@ -230,7 +241,7 @@ class TestBoxGate:
     def test_multiwire_box(self):
         """tests a gate spanning multiple wires."""
 
-        drawer = MPLDrawer(1, 3)
+        drawer = MPLDrawer(1, {0: 0, 1: 1, 2: 2})
         drawer.box_gate(0, (0, 2), text="Tall Gate")
 
         box = drawer.ax.patches[0]
@@ -251,7 +262,7 @@ class TestBoxGate:
     def test_notch_standard_styling(self):
         """Test notch styling is correct"""
 
-        drawer = MPLDrawer(n_layers=1, n_wires=3)
+        drawer = MPLDrawer(n_layers=1, wire_map={i: i for i in range(3)})
         drawer.box_gate(0, (0, 2))
 
         xs = [-0.415, 0.375, -0.415, 0.375]
@@ -273,7 +284,7 @@ class TestBoxGate:
     def test_active_wire_notches_number(self, wires, n_notches):
         """Tests the number of active wires drawn is the expected amount."""
 
-        drawer = MPLDrawer(n_layers=1, n_wires=4)
+        drawer = MPLDrawer(n_layers=1, wire_map={i: i for i in range(4)})
         drawer.box_gate(layer=0, wires=wires)
 
         assert len(drawer.ax.patches) == (n_notches + 1)
@@ -281,7 +292,7 @@ class TestBoxGate:
 
     def test_no_active_wire_notches(self):
         """Test active wire notches deactivated by keyword."""
-        drawer = MPLDrawer(n_layers=1, n_wires=3)
+        drawer = MPLDrawer(n_layers=1, wire_map={i: i for i in range(3)})
         drawer.box_gate(layer=0, wires=(0, 2), active_wire_notches=False)
 
         # only one patch for big box, no patches for notches
@@ -291,7 +302,7 @@ class TestBoxGate:
     def test_extra_width(self):
         """tests a box with added width."""
 
-        drawer = MPLDrawer(1, 3)
+        drawer = MPLDrawer(1, {i: i for i in range(3)})
         drawer.box_gate(0, (0, 2), text="Wide Gate", extra_width=0.4)
 
         box = drawer.ax.patches[0]
@@ -315,7 +326,7 @@ class TestBoxGate:
     def test_box_formatting(self):
         """Tests that box_options influences the rectangle"""
 
-        drawer = MPLDrawer(1, 3)
+        drawer = MPLDrawer(1, {i: i for i in range(3)})
         rgba_red = (1, 0, 0, 1)
         rgba_green = (0, 1, 0, 1)
         options = {"facecolor": rgba_red, "edgecolor": rgba_green, "zorder": 5}
@@ -335,7 +346,7 @@ class TestBoxGate:
     def test_text_formatting(self):
         """Tests rotated text"""
 
-        drawer = MPLDrawer(1, 1)
+        drawer = MPLDrawer(1, {0: 0})
         rgba_red = (1, 0, 0, 1)
         options = {"color": rgba_red, "rotation": "vertical"}
         drawer.box_gate(0, 0, text="X", text_options=options)
@@ -352,7 +363,7 @@ class TestCTRL:
     def test_ctrl_no_target(self):
         """Tests a single control with no target"""
 
-        drawer = MPLDrawer(1, 1)
+        drawer = MPLDrawer(1, {0: 0})
 
         drawer.ctrl(0, 0)
 
@@ -371,7 +382,7 @@ class TestCTRL:
     def test_ctrl_multi_wires(self):
         """Tests two control wires with no target."""
 
-        drawer = MPLDrawer(1, 3)
+        drawer = MPLDrawer(1, {i: i for i in range(3)})
 
         ctrl_wires = (0, 1)
         drawer.ctrl(0, ctrl_wires)
@@ -392,7 +403,7 @@ class TestCTRL:
     def test_ctrl_on_zero(self):
         """Tests a control on zero circle is open"""
 
-        drawer = MPLDrawer(1, 1)
+        drawer = MPLDrawer(1, {0: 0})
 
         drawer.ctrl(0, 0, control_values=False)
 
@@ -409,7 +420,7 @@ class TestCTRL:
     def test_ctrl_control_values_error(self):
         """Tests a ValueError is raised if different number of wires and control_values."""
 
-        drawer = MPLDrawer(1, 2)
+        drawer = MPLDrawer(1, {0: 0, 1: 1})
 
         with pytest.raises(ValueError, match="`control_values` must be the same length"):
             drawer.ctrl(0, (0, 1), control_values=True)
@@ -419,7 +430,7 @@ class TestCTRL:
     def test_ctrl_formatting(self):
         """Tests two control wires with no target."""
 
-        drawer = MPLDrawer(1, 3)
+        drawer = MPLDrawer(1, {i: i for i in range(3)})
 
         ctrl_wires = (0, 1)
         rgba_red = (1, 0, 0, 1)
@@ -443,7 +454,7 @@ class TestCTRL:
     def test_ctrl_circ(self):
         """Test only the ``_ctrl_circ`` private method."""
 
-        drawer = MPLDrawer(1, 1)
+        drawer = MPLDrawer(1, {0: 0})
         drawer._ctrl_circ(0, 0)
         circ = drawer.ax.patches[0]
 
@@ -457,7 +468,7 @@ class TestCTRL:
     def test_ctrlo_circ(self):
         """Test only the ``ctrlo_circ`` private method."""
 
-        drawer = MPLDrawer(1, 1)
+        drawer = MPLDrawer(1, {0: 0})
         drawer._ctrlo_circ(0, 0)
         circ = drawer.ax.patches[0]
 
@@ -469,7 +480,7 @@ class TestCTRL:
     def test_ctrl_target(self):
         """Tests target impacts line extent"""
 
-        drawer = MPLDrawer(1, 3)
+        drawer = MPLDrawer(1, {i: i for i in range(3)})
 
         drawer.ctrl(0, 0, 2)
 
@@ -497,14 +508,14 @@ class TestCTRL:
     )
     def test_ctrl_raises_warning_with_overlap(self, control_wires, target_wires):
         """Tests that a warning is raised if some control indicators are not visible."""
-        drawer = MPLDrawer(1, 4)
+        drawer = MPLDrawer(1, {i: i for i in range(4)})
         with pytest.warns(UserWarning, match="control indicators are hidden behind an operator"):
             drawer.ctrl(0, control_wires, target_wires)
         plt.close()
 
     @pytest.mark.parametrize("control_wires,target_wires", [((0,), (1, 2)), ((2,), (0, 1))])
     def test_ctrl_no_warning_without_overlap(self, control_wires, target_wires):
-        drawer = MPLDrawer(1, 3)
+        drawer = MPLDrawer(1, {i: i for i in range(3)})
         with warnings.catch_warnings(record=True) as w:
             drawer.ctrl(0, control_wires, target_wires)
         assert len(w) == 0
@@ -513,7 +524,7 @@ class TestCTRL:
     def test_target_x(self):
         """Tests hidden target_x drawing method"""
 
-        drawer = MPLDrawer(1, 3)
+        drawer = MPLDrawer(1, {i: i for i in range(3)})
 
         drawer._target_x(0, 0)
 
@@ -533,7 +544,7 @@ class TestCTRL:
 
     def test_target_x_color(self):
         """Test the color of target_x."""
-        drawer = MPLDrawer(1, 3)
+        drawer = MPLDrawer(1, {i: i for i in range(3)})
 
         rgba_red = (1, 0, 0, 1)
         drawer._target_x(0, 0, options={"color": rgba_red})
@@ -553,7 +564,7 @@ class TestCTRL:
     def test_CNOT(self):
         """Tests the CNOT method"""
 
-        drawer = MPLDrawer(1, 3)
+        drawer = MPLDrawer(1, {i: i for i in range(3)})
 
         drawer.CNOT(0, (0, 1))
 
@@ -578,7 +589,7 @@ class TestCTRL:
     def test_CNOT_control_values(self):
         """Tests the ``control_values`` keyword for CNOT."""
 
-        drawer = MPLDrawer(1, 3)
+        drawer = MPLDrawer(1, {i: i for i in range(3)})
 
         drawer.CNOT(0, (0, 1, 2), control_values=[True, False])
 
@@ -593,7 +604,7 @@ class TestCTRL:
 
     def test_CNOT_color(self):
         """Tests the color of CNOT."""
-        drawer = MPLDrawer(1, 3)
+        drawer = MPLDrawer(1, {i: i for i in range(3)})
         rgba_red = (1, 0, 0, 1)
         drawer.CNOT(0, (0, 1), options={"color": rgba_red})
 
@@ -618,7 +629,7 @@ class TestSWAP:
     def test_swap_x(self):
         """Tests the ``_swap_x`` private method."""
 
-        drawer = MPLDrawer(1, 1)
+        drawer = MPLDrawer(1, {0: 0})
         drawer._swap_x(0, 0)
 
         l1 = drawer.ax.lines[1]
@@ -631,7 +642,7 @@ class TestSWAP:
     def test_SWAP(self):
         """Tests the SWAP method."""
 
-        drawer = MPLDrawer(1, 3)
+        drawer = MPLDrawer(1, {i: i for i in range(3)})
         drawer.SWAP(0, (0, 2))
 
         connecting_line = drawer.ax.lines[3]
@@ -647,7 +658,7 @@ class TestSWAP:
     def test_SWAP_options(self):
         """Tests that SWAP can be colored."""
 
-        drawer = MPLDrawer(1, 3)
+        drawer = MPLDrawer(1, {i: i for i in range(3)})
         rgba_red = (1, 0, 0, 1)
         options = {"color": rgba_red, "linewidth": 3}
         drawer.SWAP(0, (0, 2), options=options)
@@ -665,7 +676,7 @@ class TestMeasure:
     def test_measure(self):
         """Tests the measure method."""
 
-        drawer = MPLDrawer(1, 1)
+        drawer = MPLDrawer(1, {0: 0})
         drawer.measure(0, 0)
 
         box = drawer.ax.patches[0]
@@ -688,7 +699,7 @@ class TestMeasure:
     def test_measure_multiple_wires(self):
         """Tests the measure method when multiple wires are provided."""
 
-        drawer = MPLDrawer(n_layers=1, n_wires=2)
+        drawer = MPLDrawer(n_layers=1, wire_map={0: 0, 1: 1})
         drawer.measure(layer=0, wires=(0, 1))
 
         box = drawer.ax.patches[0]
@@ -712,7 +723,7 @@ class TestMeasure:
     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 = MPLDrawer(n_layers=1, wire_map={0: 0, 1: 1}, c_wires=2)
         drawer.measure(layer=0, wires=(2, 3))
 
         box = drawer.ax.patches[0]
@@ -735,7 +746,7 @@ class TestMeasure:
 
     def test_measure_text(self):
         """Test adding a postselection label to a measure box."""
-        drawer = MPLDrawer(1, 1)
+        drawer = MPLDrawer(1, {0: 0})
         drawer.measure(0, 0, text="0")
         assert len(drawer.ax.texts) == 1
         assert drawer.ax.texts[0].get_text() == "0"
@@ -745,7 +756,7 @@ class TestMeasure:
     def test_measure_formatted(self):
         """Tests you can color the measure box"""
 
-        drawer = MPLDrawer(1, 1)
+        drawer = MPLDrawer(1, {0: 0})
         rgba_red = (1.0, 0, 0, 1.0)
         rgba_green = (0, 1, 0, 1)
         box_options = {"facecolor": rgba_red, "edgecolor": rgba_green}
@@ -794,7 +805,7 @@ class TestAutosize:
     def test_autosize_false(self):
         """Test that the text is unchanged if autosize is set to False."""
 
-        drawer = MPLDrawer(n_layers=1, n_wires=2)
+        drawer = MPLDrawer(n_layers=1, wire_map={0: 0, 1: 1})
         drawer.box_gate(0, (0, 1), text="very very long text", autosize=False)
 
         t = drawer.ax.texts[0]
@@ -807,7 +818,7 @@ class TestAutosize:
         """Test case where the box is on only one wire.  The text should still
         be inside the box, but not rotated."""
 
-        drawer = MPLDrawer(n_layers=1, n_wires=1)
+        drawer = MPLDrawer(n_layers=1, wire_map={0: 0})
         drawer.box_gate(0, 0, text="very very long text", autosize=True)
 
         t = drawer.ax.texts[0]
@@ -821,7 +832,7 @@ class TestAutosize:
         """Test case where the box is on multiple wires.  The text should
         be rotated 90deg, and still inside the box."""
 
-        drawer = MPLDrawer(n_layers=1, n_wires=2)
+        drawer = MPLDrawer(n_layers=1, wire_map={0: 0, 1: 1})
         drawer.box_gate(0, (0, 1), text="very very long text", autosize=True)
 
         t = drawer.ax.texts[0]
@@ -835,7 +846,7 @@ class TestAutosize:
         """Test case where the box is on one wire and the text is skinny and tall.
         If the text is too tall, it should still be shrunk to fit inside the box."""
 
-        drawer = MPLDrawer(n_layers=1, n_wires=1)
+        drawer = MPLDrawer(n_layers=1, wire_map={0: 0})
         drawer.box_gate(0, 0, text="text\nwith\nall\nthe\nlines\nyep", autosize=True)
 
         t = drawer.ax.texts[0]
@@ -850,7 +861,7 @@ class TestAutosize:
         If the text is just too tall and the width fits inside the box, the text should not
         be rotated."""
 
-        drawer = MPLDrawer(n_layers=1, n_wires=2)
+        drawer = MPLDrawer(n_layers=1, wire_map={0: 0, 1: 1})
         drawer.box_gate(
             0,
             (0, 1),
@@ -869,7 +880,7 @@ class TestAutosize:
         """Test case where the box is on multiple wires and text is fat, tall,
         and fatter than it is tall. It should be rotated."""
 
-        drawer = MPLDrawer(n_layers=1, n_wires=2)
+        drawer = MPLDrawer(n_layers=1, wire_map={0: 0, 1: 1})
         drawer.box_gate(0, (0, 1), text="very very long text\nall\nthe\nlines\nyep", autosize=True)
 
         assert self.text_in_box(drawer)
@@ -880,7 +891,7 @@ class TestAutosize:
 class TestClassicalWires:
     def test_classical_wire(self):
         """Test the addition of horiziontal classical wires."""
-        drawer = MPLDrawer(n_wires=1, n_layers=4, c_wires=3)
+        drawer = MPLDrawer(wire_map={0: 0}, n_layers=4, c_wires=3)
 
         layers = [0, 0, 1, 1]
         wires = [0, 1, 1, 0]
@@ -906,7 +917,7 @@ class TestClassicalWires:
 
     def test_cwire_join(self):
         """Test the cwire join method."""
-        drawer = MPLDrawer(n_wires=1, n_layers=4, c_wires=3)
+        drawer = MPLDrawer(wire_map={0: 0}, n_layers=4, c_wires=3)
 
         drawer.cwire_join(1, 2)
 
@@ -920,7 +931,7 @@ class TestClassicalWires:
 
     def test_cwire_join_erase_right(self):
         """Test the cwire join method."""
-        drawer = MPLDrawer(n_wires=1, n_layers=4, c_wires=3)
+        drawer = MPLDrawer(wire_map={0: 0}, n_layers=4, c_wires=3)
 
         drawer.cwire_join(1, 1, erase_right=True)
 
@@ -938,7 +949,7 @@ class TestCond:
 
     def test_cond_basic(self):
         """Tests cond from one wire to the next."""
-        drawer = MPLDrawer(n_wires=2, n_layers=2)
+        drawer = MPLDrawer(wire_map={0: 0, 1: 1}, n_layers=2)
         wire_data_before = [line.get_data() for line in drawer._wire_lines]
 
         drawer.cond(layer=1, measured_layer=0, wires=[0], wires_target=[1])
@@ -957,7 +968,7 @@ class TestCond:
 
     def test_cond_two_ctrl_wires(self):
         """Tests cond from two separated wires."""
-        drawer = MPLDrawer(n_wires=4, n_layers=2)
+        drawer = MPLDrawer(wire_map={i: i for i in range(4)}, n_layers=2)
         drawer.cond(layer=1, measured_layer=0, wires=[0, 2], wires_target=[3])
         actual_data = [line.get_data() for line in drawer.ax.lines]
 
@@ -978,7 +989,7 @@ class TestCond:
 
     def test_cond_two_ctrl_wires_upward(self):
         """Test cond when the conditional operation is above the control wires."""
-        drawer = MPLDrawer(n_wires=3, n_layers=2)
+        drawer = MPLDrawer(wire_map={i: i for i in range(3)}, n_layers=2)
         drawer.cond(layer=1, measured_layer=0, wires=[1, 2], wires_target=[0])
         actual_data = [line.get_data() for line in drawer.ax.lines]
         assert actual_data == [
@@ -1006,7 +1017,7 @@ class TestCond:
     )
     def test_cond_fail_with_bad_order(self, ctrl_wires, target_wires):
         """Tests cond raises an error when the wires aren't neatly separated."""
-        drawer = MPLDrawer(n_wires=4, n_layers=2)
+        drawer = MPLDrawer(wire_map={i: i for i in range(4)}, n_layers=2)
         with pytest.raises(ValueError, match="Cannot draw interspersed mid-circuit measurements"):
             drawer.cond(layer=1, measured_layer=0, wires=ctrl_wires, wires_target=target_wires)
         plt.close()