Adding new feature wire_options to draw_mpl function

doc/code/qml_drawer.rst

@@ -85,4 +85,4 @@ Currently Available Styles
 +|pls|+|plw|+|skd|+
 +-----+-----+-----+
 +|sol|+|sod|+|def|+
-+-----+-----+-----+
++-----+-----+-----+

doc/releases/changelog-dev.md

@@ -6,7 +6,9 @@
 
 <h3>Improvements 🛠</h3>
 
-<h4>Other Improvements</h4>
+* Added support for the `wire_options` dictionary to customize wire line formatting in `qml.draw_mpl` circuit
+  visualizations, allowing global and per-wire customization with options like `color`, `linestyle`, and `linewidth`.
+  [(#6486)](https://github.com/PennyLaneAI/pennylane/pull/6486)
 
 * Added `qml.devices.qubit_mixed` module for mixed-state qubit device support. This module introduces:
   - A new API for mixed-state operations
@@ -38,4 +40,5 @@
 
 This release contains contributions from (in alphabetical order):
 
-Andrija Paurevic
+Shiwen An
+Andrija Paurevic

pennylane/drawer/draw.py

@@ -385,7 +385,9 @@ def draw_mpl(
         fontsize (float or str): fontsize for text. Valid strings are
             ``{'xx-small', 'x-small', 'small', 'medium', large', 'x-large', 'xx-large'}``.
             Default is ``14``.
-        wire_options (dict): matplotlib formatting options for the wire lines
+        wire_options (dict): matplotlib formatting options for the wire lines. In addition to
+            standard options, options per wire can be specified with ``wire_label: options``
+            pairs, also see examples below.
         label_options (dict): matplotlib formatting options for the wire labels
         show_wire_labels (bool): Whether or not to show the wire labels.
         active_wire_notches (bool): whether or not to add notches indicating active wires.
@@ -458,7 +460,8 @@ def circuit2(x, y):
 
         **Wires:**
 
-        The keywords ``wire_order`` and ``show_all_wires`` control the location of wires from top to bottom.
+        The keywords ``wire_order`` and ``show_all_wires`` control the location of wires
+        from top to bottom.
 
         .. code-block:: python
 
@@ -470,8 +473,8 @@ def circuit2(x, y):
                 :width: 60%
                 :target: javascript:void(0);
 
-        If a wire is in ``wire_order``, but not in the ``tape``, it will be omitted by default.  Only by selecting
-        ``show_all_wires=True`` will empty wires be displayed.
+        If a wire is in ``wire_order``, but not in the ``tape``, it will be omitted by default.
+        Only by selecting ``show_all_wires=True`` will empty wires be displayed.
 
         .. code-block:: python
 
@@ -568,6 +571,26 @@ def circuit2(x, y):
                 :width: 60%
                 :target: javascript:void(0);
 
+
+        Additionally, ``wire_options`` may contain wire labels as keys and dictionaries with
+        keyword-value pairs of matplotlib options as values, which will control the line style for
+        the respective individual wire.
+
+        .. code-block:: python
+
+            wire_options = {
+                'color': 'teal', # all wires but wire 2 will be teal
+                'linewidth': 5, # all wires but wire 2 will be bold
+                2: {'color': 'orange', 'linestyle': '--'}, # wire 2 will be orange and dashed
+            }
+            fig, ax = qml.draw_mpl(circuit, wire_options=wire_options)(1.2345,1.2345)
+            fig.show()
+
+        .. figure:: ../../_static/draw_mpl/per_wire_options.png
+                :align: center
+                :width: 60%
+                :target: javascript:void(0);
+
         **Levels:**
 
         The ``level`` keyword argument allows one to select a subset of the transforms to apply on the ``QNode``

pennylane/drawer/mpldrawer.py

@@ -293,12 +293,30 @@ def __init__(self, n_layers, n_wires, c_wires=0, wire_options=None, figsize=None
 
         if wire_options is None:
             wire_options = {}
+            self._wire_lines = [
+                plt.Line2D((-1, self.n_layers), (wire, wire), zorder=1, **wire_options)
+                for wire in range(self.n_wires)
+            ]
+        else:
+            # Separate global options from per wire options
+            global_options = {k: v for k, v in wire_options.items() if not isinstance(v, dict)}
+            wire_specific_options = {k: v for k, v in wire_options.items() if isinstance(v, dict)}
+
+            # 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, {})
+                line_options = {**global_options, **specific_options}
+
+                # Create Line2D with the combined options
+                line = plt.Line2D(
+                    (-1, self.n_layers),
+                    (wire, wire),
+                    zorder=1,
+                    **line_options,
+                )
+                self._wire_lines.append(line)
 
-        # adding wire lines
-        self._wire_lines = [
-            plt.Line2D((-1, self.n_layers), (wire, wire), zorder=1, **wire_options)
-            for wire in range(self.n_wires)
-        ]
         for line in self._wire_lines:
             self._ax.add_line(line)
 

pennylane/drawer/tape_mpl.py

@@ -303,7 +303,9 @@ def tape_mpl(
         fontsize (float or str): fontsize for text. Valid strings are
             ``{'xx-small', 'x-small', 'small', 'medium', large', 'x-large', 'xx-large'}``.
             Default is ``14``.
-        wire_options (dict): matplotlib formatting options for the wire lines
+        wire_options (dict): matplotlib formatting options for the wire lines. In addition to
+            standard options, options per wire can be specified with ``wire_label: options``
+            pairs, also see examples below.
         label_options (dict): matplotlib formatting options for the wire labels
         show_wire_labels (bool): Whether or not to show the wire labels.
         active_wire_notches (bool): whether or not to add notches indicating active wires.
@@ -328,7 +330,7 @@ def tape_mpl(
         measurements = [qml.expval(qml.Z(0))]
         tape = qml.tape.QuantumTape(ops, measurements)
 
-        fig, ax = tape_mpl(tape)
+        fig, ax = qml.drawer.tape_mpl(tape)
         fig.show()
 
     .. figure:: ../../_static/tape_mpl/default.png
@@ -350,7 +352,7 @@ def tape_mpl(
         measurements = [qml.expval(qml.Z(0))]
         tape2 = qml.tape.QuantumTape(ops, measurements)
 
-        fig, ax = tape_mpl(tape2, decimals=2)
+        fig, ax = qml.drawer.tape_mpl(tape2, decimals=2)
 
     .. figure:: ../../_static/tape_mpl/decimals.png
         :align: center
@@ -363,7 +365,7 @@ def tape_mpl(
 
     .. code-block:: python
 
-        fig, ax = tape_mpl(tape, wire_order=[3,2,1,0])
+        fig, ax = qml.drawer.tape_mpl(tape, wire_order=[3,2,1,0])
 
     .. figure:: ../../_static/tape_mpl/wire_order.png
             :align: center
@@ -375,7 +377,7 @@ def tape_mpl(
 
     .. code-block:: python
 
-        fig, ax = tape_mpl(tape, wire_order=["aux"], show_all_wires=True)
+        fig, ax = qml.drawer.tape_mpl(tape, wire_order=["aux"], show_all_wires=True)
 
     .. figure:: ../../_static/tape_mpl/show_all_wires.png
             :align: center
@@ -389,7 +391,7 @@ def tape_mpl(
 
     .. code-block:: python
 
-        fig, ax = tape_mpl(tape)
+        fig, ax = qml.drawer.tape_mpl(tape)
         fig.suptitle("My Circuit", fontsize="xx-large")
 
         options = {'facecolor': "white", 'edgecolor': "#f57e7e", "linewidth": 6, "zorder": -1}
@@ -413,7 +415,7 @@ def tape_mpl(
 
     .. code-block:: python
 
-        fig, ax = tape_mpl(tape, style='sketch')
+        fig, ax = qml.drawer.tape_mpl(tape, style='sketch')
 
     .. figure:: ../../_static/tape_mpl/sketch_style.png
             :align: center
@@ -437,7 +439,7 @@ def tape_mpl(
         plt.rcParams['lines.linewidth'] = 5
         plt.rcParams['figure.facecolor'] = 'ghostwhite'
 
-        fig, ax = tape_mpl(tape, style="rcParams")
+        fig, ax = qml.drawer.tape_mpl(tape, style="rcParams")
 
     .. figure:: ../../_static/tape_mpl/rcparams.png
             :align: center
@@ -450,14 +452,31 @@ def tape_mpl(
 
     .. code-block:: python
 
-        fig, ax = tape_mpl(tape, wire_options={'color':'teal', 'linewidth': 5},
+        fig, ax = qml.drawer.tape_mpl(tape, wire_options={'color':'teal', 'linewidth': 5},
                     label_options={'size': 20})
 
     .. figure:: ../../_static/tape_mpl/wires_labels.png
             :align: center
             :width: 60%
             :target: javascript:void(0);
 
+    Additionally, ``wire_options`` may contain wire labels as keys and dictionaries with
+    keyword-value pairs of matplotlib options as values, which will control the line style for
+    the respective individual wire.
+
+    .. code-block:: python
+
+        wire_options = {
+            'color': 'teal', # all wires but wire 2 will be teal
+            'linewidth': 5, # all wires but wire 2 will be bold
+            2: {'color': 'orange', 'linestyle': '--'}, # wire 2 will be orange and dashed
+        }
+        fig, ax = qml.drawer.tape_mpl(tape, wire_options=wire_options)
+
+    .. figure:: ../../_static/tape_mpl/per_wire_options.png
+            :align: center
+            :width: 60%
+            :target: javascript:void(0);
     """
 
     restore_params = {}

tests/drawer/test_draw_mpl.py

@@ -328,6 +328,75 @@ def test_wire_options(self):
             assert w.get_color() == "black"
             assert w.get_linewidth() == 4
 
+        @qml.qnode(dev)
+        def f_circ(x):
+            """Circuit on ten qubits."""
+            qml.RX(x, wires=0)
+            for w in range(10):
+                qml.Hadamard(w)
+            return qml.expval(qml.PauliZ(0) @ qml.PauliY(1))
+
+        # All wires are orange
+        wire_options = {"color": "orange"}
+        _, ax = qml.draw_mpl(f_circ, wire_options=wire_options)(0.52)
+
+        for w in ax.lines:
+            assert w.get_color() == "orange"
+
+        # Wires are orange and cyan
+        wire_options = {0: {"color": "orange"}, 1: {"color": "cyan"}}
+        _, ax = qml.draw_mpl(f_circ, wire_options=wire_options)(0.52)
+
+        assert ax.lines[0].get_color() == "orange"
+        assert ax.lines[1].get_color() == "cyan"
+        assert ax.lines[2].get_color() == "black"
+
+        # Make all wires cyan and bold,
+        # except for wires 2 and 6, which are dashed and another color
+        wire_options = {
+            "color": "cyan",
+            "linewidth": 5,
+            2: {"linestyle": "--", "color": "red"},
+            6: {"linestyle": "--", "color": "orange", "linewidth": 1},
+        }
+        _, ax = qml.draw_mpl(f_circ, wire_options=wire_options)(0.52)
+
+        for i, w in enumerate(ax.lines):
+            if i == 2:
+                assert w.get_color() == "red"
+                assert w.get_linestyle() == "--"
+                assert w.get_linewidth() == 5
+            elif i == 6:
+                assert w.get_color() == "orange"
+                assert w.get_linestyle() == "--"
+                assert w.get_linewidth() == 1
+            else:
+                assert w.get_color() == "cyan"
+                assert w.get_linestyle() == "-"
+                assert w.get_linewidth() == 5
+
+        wire_options = {
+            "linewidth": 5,
+            2: {"linestyle": "--", "color": "red"},
+            6: {"linestyle": "--", "color": "orange"},
+        }
+
+        _, ax = qml.draw_mpl(f_circ, wire_options=wire_options)(0.52)
+
+        for i, w in enumerate(ax.lines):
+            if i == 2:
+                assert w.get_color() == "red"
+                assert w.get_linestyle() == "--"
+                assert w.get_linewidth() == 5
+            elif i == 6:
+                assert w.get_color() == "orange"
+                assert w.get_linestyle() == "--"
+                assert w.get_linewidth() == 5
+            else:
+                assert w.get_color() == "black"
+                assert w.get_linestyle() == "-"
+                assert w.get_linewidth() == 5
+
         plt.close()