Repository: PennyLaneAI/pennylane
Pull request #5576 (2024-04-24 19:27:12)

Draw sorts wires
**Context:**

Devices may no longer provide a wire order, and we can also draw quantum functions that do not have an intrinsic wire order.  This can lead to the wire order in the drawing being unexpected, like `[4,3,0,1]`.

**Description of the Change:**

Tries to sort the wires if no other wires are provided, but fallbacks to the tape wires if sorting fails.


Before:
![Screenshot 2024-04-24 at 3 25 56 PM](https://github.com/PennyLaneAI/pennylane/assets/6364575/d8c8507b-2dd1-4d8e-a057-7614e7b9ab5c)

After: 
![Screenshot 2024-04-24 at 3 25 20 PM](https://github.com/PennyLaneAI/pennylane/assets/6364575/b13534bf-6683-438b-b5ed-682396cf8670)

**Benefits:**

Nicer looking drawings.

**Possible Drawbacks:**

try-except as logical control.

**Related GitHub Issues:**
[sc-61203]

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5576?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: Patch coverage is `92.85714%` with `2 lines` in your changes are missing coverage. Please review.
> Project coverage is 99.68%. Comparing base [(`a8bba50`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/a8bba509e16a8582d57a0baa978ddb89a6d4d75a?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`08aff2b`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5576?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 11 commits behind head on master.

| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5576?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Patch % | Lines |
|---|---|---|
| [pennylane/drawer/draw.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5576?src=pr&el=tree&filepath=pennylane%2Fdrawer%2Fdraw.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RyYXdlci9kcmF3LnB5) | 92.85% | [2 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5576?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    #5576      +/-   ##
==========================================
- Coverage   99.69%   99.68%   -0.01%     
==========================================
  Files         410      410              
  Lines       38230    37995     -235     
==========================================
- Hits        38113    37875     -238     
- Misses        117      120       +3     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5576?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-dev.md b/doc/releases/changelog-dev.md
index 20bc1ae2a..dc82919ac 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -347,6 +347,10 @@
 
 <h4>Other improvements</h4>
 
+* `qml.draw` and `qml.draw_mpl` will now attempt to sort the wires if no wire order
+  is provided by the user or the device.
+  [(#5576)](https://github.com/PennyLaneAI/pennylane/pull/5576)
+
 * `qml.ops.Conditional` now stores the `data`, `num_params`, and `ndim_param` attributes of
   the operator it wraps.
   [(#5473)](https://github.com/PennyLaneAI/pennylane/pull/5473)
diff --git a/pennylane/drawer/draw.py b/pennylane/drawer/draw.py
index 38ecfb523..43adbe73d 100644
--- a/pennylane/drawer/draw.py
+++ b/pennylane/drawer/draw.py
@@ -47,7 +47,9 @@ def draw(
 
     Args:
         qnode (.QNode or Callable): the input QNode or quantum function that is to be drawn.
-        wire_order (Sequence[Any]): the order (from top to bottom) to print the wires of the circuit
+        wire_order (Sequence[Any]): the order (from top to bottom) to print the wires of the circuit.
+           If not provided, the wire order defaults to the device wires. If device wires are not
+           available, the circuit wires are sorted if possible.
         show_all_wires (bool): If True, all wires, including empty wires, are printed.
         decimals (int): How many decimal points to include when formatting operation parameters.
             ``None`` will omit parameters from operation labels.
@@ -230,7 +232,14 @@ def draw(
     @wraps(qnode)
     def wrapper(*args, **kwargs):
         tape = qml.tape.make_qscript(qnode)(*args, **kwargs)
-        _wire_order = wire_order or tape.wires
+
+        if wire_order:
+            _wire_order = wire_order
+        else:
+            try:
+                _wire_order = sorted(tape.wires)
+            except TypeError:
+                _wire_order = tape.wires
 
         return tape_text(
             tape,
@@ -273,7 +282,15 @@ def _draw_qnode(
             finally:
                 qnode.expansion_strategy = original_expansion_strategy
 
-        _wire_order = wire_order or qnode.device.wires or qnode.tape.wires
+        if wire_order:
+            _wire_order = wire_order
+        elif qnode.device.wires:
+            _wire_order = qnode.device.wires
+        else:
+            try:
+                _wire_order = sorted(tapes[0].wires)
+            except TypeError:
+                _wire_order = tapes[0].wires
 
         if tapes is not None:
             cache = {"tape_offset": 0, "matrices": []}
@@ -327,7 +344,9 @@ def draw_mpl(
         qnode (.QNode or Callable): the input QNode/quantum function that is to be drawn.
 
     Keyword Args:
-        wire_order (Sequence[Any]): the order (from top to bottom) to print the wires of the circuit
+        wire_order (Sequence[Any]): the order (from top to bottom) to print the wires of the circuit.
+           If not provided, the wire order defaults to the device wires. If device wires are not
+           available, the circuit wires are sorted if possible.
         show_all_wires (bool): If True, all wires, including empty wires, are printed.
         decimals (int): How many decimal points to include when formatting operation parameters.
             Default ``None`` will omit parameters from operation labels.
@@ -547,7 +566,13 @@ def draw_mpl(
     @wraps(qnode)
     def wrapper(*args, **kwargs):
         tape = qml.tape.make_qscript(qnode)(*args, **kwargs)
-        _wire_order = wire_order or tape.wires
+        if wire_order:
+            _wire_order = wire_order
+        else:
+            try:
+                _wire_order = sorted(tape.wires)
+            except TypeError:
+                _wire_order = tape.wires
 
         return tape_mpl(
             tape,
@@ -592,7 +617,15 @@ def _draw_mpl_qnode(
             finally:
                 qnode.expansion_strategy = original_expansion_strategy
 
-        _wire_order = wire_order or qnode.device.wires or tape.wires
+        if wire_order:
+            _wire_order = wire_order
+        elif qnode.device.wires:
+            _wire_order = qnode.device.wires
+        else:
+            try:
+                _wire_order = sorted(tape.wires)
+            except TypeError:
+                _wire_order = tape.wires
 
         return tape_mpl(
             tape,


Test patch:
diff --git a/tests/drawer/test_draw.py b/tests/drawer/test_draw.py
index b571ed4fd..f09757259 100644
--- a/tests/drawer/test_draw.py
+++ b/tests/drawer/test_draw.py
@@ -44,10 +44,21 @@ class TestLabelling:
         assert split_str[1][:6] == "    a:"
         assert split_str[2][:6] == "1.234:"
 
-    def test_wire_order(self):
+    @pytest.mark.parametrize("as_qnode", (True, False))
+    def test_wire_order(self, as_qnode):
         """Test wire_order keyword changes order of the wires."""
 
-        split_str = draw(circuit, wire_order=[1.234, "a", 0, "b"])(1.2, 2.3, 3.4).split("\n")
+        def f(x, y, z):
+            """A quantum circuit on three wires."""
+            qml.RX(x, wires=0)
+            qml.RY(y, wires="a")
+            qml.RZ(z, wires=1.234)
+            return qml.expval(qml.PauliZ(0))
+
+        if as_qnode:
+            f = qml.QNode(f, qml.device("default.qubit", wires=(0, "a", 1.234)))
+
+        split_str = draw(f, wire_order=[1.234, "a", 0, "b"])(1.2, 2.3, 3.4).split("\n")
         assert split_str[0][:6] == "1.234:"
         assert split_str[1][:6] == "    a:"
         assert split_str[2][:6] == "    0:"
@@ -397,8 +408,8 @@ class TestMidCircuitMeasurements:
         expected_drawing = (
             "0: ──RX(3.14)──┤↗│  │0⟩─╭●─────────────────────╭MultiRZ(0.50)─┤     \n"
             "1: ──RX(3.12)──┤↗├──────│─────────────╭●───────│──────────────┤     \n"
-            "3: ─────────────────────╰X──┤↗₀│  │0⟩─│────────│──────────────┤     \n"
-            "2: ──RY(0.46)─────────────────────────╰X──┤↗₁├─╰MultiRZ(0.50)─┤  <Z>"
+            "2: ─────────────────────│───RY(0.46)──╰X──┤↗₁├─╰MultiRZ(0.50)─┤  <Z>\n"
+            "3: ─────────────────────╰X──┤↗₀│  │0⟩─────────────────────────┤     "
         )
 
         assert drawing == expected_drawing
@@ -432,8 +443,8 @@ class TestMidCircuitMeasurements:
         drawing = qml.draw(circ)(np.pi, np.pi / 2)
         expected_drawing = (
             "0: ──RX(3.14)──┤↗├───────────╭X─┤  \n"
-            "2: ─────────────║───RY(1.57)─│──┤  \n"
             "1: ─────────────║────────────╰●─┤  \n"
+            "2: ─────────────║───RY(1.57)──║─┤  \n"
             "                ╚═════════════╝    "
         )
 
@@ -980,3 +991,38 @@ def test_draw_with_qfunc_warns_with_expansion_strategy():
 
     with pytest.warns(UserWarning, match="the expansion_strategy argument is ignored"):
         _ = qml.draw(qfunc, expansion_strategy="gradient")
+
+
+@pytest.mark.parametrize("use_qnode", [True, False])
+def test_sort_wires(use_qnode):
+    """Test that drawing a qnode with no wire order or device wires sorts the wires automatically."""
+
+    def func():
+        qml.X(4)
+        qml.X(2)
+        qml.X(0)
+        return qml.expval(qml.Z(0))
+
+    if use_qnode:
+        func = qml.QNode(func, qml.device("default.qubit"))
+
+    expected = "0: ──X─┤  <Z>\n2: ──X─┤     \n4: ──X─┤     "
+    assert qml.draw(func)() == expected
+
+
+@pytest.mark.parametrize("use_qnode", [True, False])
+def test_sort_wires_fallback(use_qnode):
+    """Test that drawing a qnode with no wire order or device wires falls back to tape wires if
+    sorting fails."""
+
+    def func():
+        qml.X(4)
+        qml.X("a")
+        qml.X(0)
+        return qml.expval(qml.Z(0))
+
+    if use_qnode:
+        func = qml.QNode(func, qml.device("default.qubit"))
+
+    expected = "4: ──X─┤     \na: ──X─┤     \n0: ──X─┤  <Z>"
+    assert qml.draw(func)() == expected
diff --git a/tests/drawer/test_draw_mpl.py b/tests/drawer/test_draw_mpl.py
index 010b38e23..b7a7e91a3 100644
--- a/tests/drawer/test_draw_mpl.py
+++ b/tests/drawer/test_draw_mpl.py
@@ -189,10 +189,21 @@ class TestKwargs:
 class TestWireBehaviour:
     """Tests that involve how wires are displayed"""
 
-    def test_wire_order(self):
+    @pytest.mark.parametrize("use_qnode", (True, False))
+    def test_wire_order(self, use_qnode):
         """Test wire_order changes order of wires"""
 
-        _, ax = qml.draw_mpl(circuit1, wire_order=(1.23, "a"))(1.23, 2.34)
+        def f(x, y):
+            """Circuit on three qubits."""
+            qml.RX(x, wires=0)
+            qml.CNOT(wires=(0, "a"))
+            qml.RY(y, wires=1.23)
+            return qml.expval(qml.PauliZ(0))
+
+        if use_qnode:
+            f = qml.QNode(f, qml.device("default.qubit", wires=(0, "a", 1.23)))
+
+        _, ax = qml.draw_mpl(f, wire_order=(1.23, "a"))(1.23, 2.34)
 
         assert len(ax.texts) == 5
 
@@ -426,3 +437,43 @@ def test_applied_transforms(device):
     assert len(ax.texts) == 2  # one wire label, one gate label
 
     plt.close()
+
+
+@pytest.mark.parametrize("use_qnode", (True, False))
+def test_wire_sorting_if_no_wire_order(use_qnode):
+    """Test that wires are automatically sorted if the device and user
+    dont provide a wire order."""
+
+    def f():
+        qml.X(4)
+        qml.X(2)
+        return qml.expval(qml.Z(0))
+
+    if use_qnode:
+        f = qml.QNode(f, qml.device("default.qubit"))
+
+    _, ax = qml.draw_mpl(f)()
+
+    assert ax.texts[0].get_text() == "0"
+    assert ax.texts[1].get_text() == "2"
+    assert ax.texts[2].get_text() == "4"
+
+
+@pytest.mark.parametrize("use_qnode", (True, False))
+def test_wire_sorting_fallback_if_no_wire_order(use_qnode):
+    """Test that wires are automatically sorted if the device and user
+    dont provide a wire order."""
+
+    def f():
+        qml.X(4)
+        qml.X("a")
+        return qml.expval(qml.Z(0))
+
+    if use_qnode:
+        f = qml.QNode(f, qml.device("default.qubit"))
+
+    _, ax = qml.draw_mpl(f)()
+
+    assert ax.texts[0].get_text() == "4"
+    assert ax.texts[1].get_text() == "a"
+    assert ax.texts[2].get_text() == "0"