Repository: PennyLaneAI/pennylane (python)

Issue:
Drawing `PrepSelPrep` shoud be less verbose
### Feature details

`PrepSelPrep`'s LCU argument is interpreted as a parameter, resulting in a verbose drawing. The feature we'd like to have is a custom label method that treats the LCU variable as a single array, which is displayed below the tape drawing like for other arrays.

### Implementation

_No response_

### How important would you say this feature is?

1: Not important. Would be nice to have.

### Additional information

```python
import numpy as np
import pennylane as qml

k = 4
num_terms = 2**k
ops = list(qml.pauli.pauli_group(k))[:num_terms]
lcu = qml.dot(np.random.random(num_terms), ops)
@qml.qnode(qml.device("default.qubit"))
def node():
    qml.PrepSelPrep(lcu, control=range(k, 2*k))
    return qml.state()
```

```
>>> print(qml.draw(node)())
0: ─╭PrepSelPrep(0.68,0.54,0.51,0.15,0.74,0.50,0.62,0.24,0.95,0.41,0.52,0.99,0.98,0.79,0.27,0.75)─┤
1: ─├PrepSelPrep(0.68,0.54,0.51,0.15,0.74,0.50,0.62,0.24,0.95,0.41,0.52,0.99,0.98,0.79,0.27,0.75)─┤
2: ─├PrepSelPrep(0.68,0.54,0.51,0.15,0.74,0.50,0.62,0.24,0.95,0.41,0.52,0.99,0.98,0.79,0.27,0.75)─┤
3: ─├PrepSelPrep(0.68,0.54,0.51,0.15,0.74,0.50,0.62,0.24,0.95,0.41,0.52,0.99,0.98,0.79,0.27,0.75)─┤
4: ─├PrepSelPrep(0.68,0.54,0.51,0.15,0.74,0.50,0.62,0.24,0.95,0.41,0.52,0.99,0.98,0.79,0.27,0.75)─┤
5: ─├PrepSelPrep(0.68,0.54,0.51,0.15,0.74,0.50,0.62,0.24,0.95,0.41,0.52,0.99,0.98,0.79,0.27,0.75)─┤
6: ─├PrepSelPrep(0.68,0.54,0.51,0.15,0.74,0.50,0.62,0.24,0.95,0.41,0.52,0.99,0.98,0.79,0.27,0.75)─┤
7: ─╰PrepSelPrep(0.68,0.54,0.51,0.15,0.74,0.50,0.62,0.24,0.95,0.41,0.52,0.99,0.98,0.79,0.27,0.75)─┤
```

Pull request description:
Create a custom label method for `PrepSelPrep`
**Context:**
`PrepSelPrep` prints the coefficients of the encoded LCU as parameters in its label.
That's very verbose and impractical, see #7159.

**Description of the Change:**
Implement a custom `label` method for `PrepSelPrep` that defers the coefficients to the standard matrix cache,
which is also used for all matrix-valued parameters.
Like requested in the linked issue.

**Benefits:**
More readable drawings.

**Possible Drawbacks:**
N/A

**Related GitHub Issues:**
Resolves #7159

[sc-87444]


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index ca3917198..43f94eb5e 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -194,6 +194,9 @@
 
 <h3>Improvements 🛠</h3>
 
+* `PrepSelPrep` now has a concise representation when drawn with `qml.draw` or `qml.draw_mpl`.
+  [(#7164)](https://github.com/PennyLaneAI/pennylane/pull/7164)
+
 * The decomposition of a single qubit `qml.QubitUnitary` now includes the global phase.
   [(#7143)](https://github.com/PennyLaneAI/pennylane/pull/7143)
 
diff --git a/pennylane/operation.py b/pennylane/operation.py
index fbe3b313b..405939437 100644
--- a/pennylane/operation.py
+++ b/pennylane/operation.py
@@ -1053,8 +1053,8 @@ class Operator(abc.ABC, metaclass=ABCCaptureMeta):
             return op_label if self._id is None else f'{op_label}("{self._id}")'
 
         params = self.parameters
-
-        if len(qml.math.shape(params[0])) != 0:
+        shape0 = qml.math.shape(params[0])
+        if len(shape0) != 0:
             # assume that if the first parameter is matrix-valued, there is only a single parameter
             # this holds true for all current operations and templates unless parameter broadcasting
             # is used
@@ -1067,23 +1067,15 @@ class Operator(abc.ABC, metaclass=ABCCaptureMeta):
                 return op_label if self._id is None else f'{op_label}("{self._id}")'
 
             for i, mat in enumerate(cache["matrices"]):
-                if qml.math.shape(params[0]) == qml.math.shape(mat) and qml.math.allclose(
-                    params[0], mat
-                ):
-                    return (
-                        f"{op_label}(M{i})"
-                        if self._id is None
-                        else f'{op_label}(M{i},"{self._id}")'
-                    )
-
-            # matrix not in cache
-            mat_num = len(cache["matrices"])
-            cache["matrices"].append(params[0])
-            return (
-                f"{op_label}(M{mat_num})"
-                if self._id is None
-                else f'{op_label}(M{mat_num},"{self._id}")'
-            )
+                if shape0 == qml.math.shape(mat) and qml.math.allclose(params[0], mat):
+                    str_wo_id = f"{op_label}(M{i})"
+                    break
+            else:
+                mat_num = len(cache["matrices"])
+                cache["matrices"].append(params[0])
+                str_wo_id = f"{op_label}(M{mat_num})"
+
+            return str_wo_id if self._id is None else f'{str_wo_id[:-1]},"{self._id}")'
 
         if decimals is None:
             return op_label if self._id is None else f'{op_label}("{self._id}")'
diff --git a/pennylane/templates/subroutines/prepselprep.py b/pennylane/templates/subroutines/prepselprep.py
index 6d8a7b660..bb100bfea 100644
--- a/pennylane/templates/subroutines/prepselprep.py
+++ b/pennylane/templates/subroutines/prepselprep.py
@@ -113,6 +113,24 @@ class PrepSelPrep(Operation):
     def decomposition(self):
         return self.compute_decomposition(self.lcu, self.control)
 
+    def label(self, decimals=None, base_label=None, cache=None) -> str:
+        op_label = base_label or self.__class__.__name__
+        if cache is None or not isinstance(cache.get("matrices", None), list):
+            return op_label if self._id is None else f'{op_label}("{self._id}")'
+
+        coeffs = qml.math.array(self.coeffs)
+        shape = qml.math.shape(coeffs)
+        for i, mat in enumerate(cache["matrices"]):
+            if shape == qml.math.shape(mat) and qml.math.allclose(coeffs, mat):
+                str_wo_id = f"{op_label}(M{i})"
+                break
+        else:
+            mat_num = len(cache["matrices"])
+            cache["matrices"].append(coeffs)
+            str_wo_id = f"{op_label}(M{mat_num})"
+
+        return str_wo_id if self._id is None else f'{str_wo_id[:-1]},"{self._id}")'
+
     @staticmethod
     def compute_decomposition(lcu, control):
         coeffs, ops = _get_new_terms(lcu)


Test patch:
diff --git a/tests/templates/test_subroutines/test_prepselprep.py b/tests/templates/test_subroutines/test_prepselprep.py
index a7037b116..fbf84d90c 100644
--- a/tests/templates/test_subroutines/test_prepselprep.py
+++ b/tests/templates/test_subroutines/test_prepselprep.py
@@ -101,6 +101,17 @@ def prepselprep_circuit(lcu, control):
     return qml.state()
 
 
+a_set_of_lcus = [
+    qml.ops.LinearCombination([0.25, 0.75], [qml.Z(2), qml.X(1) @ qml.X(2)]),
+    qml.dot([0.25, 0.75], [qml.Z(2), qml.X(1) @ qml.X(2)]),
+    qml.Hamiltonian([0.25, 0.75], [qml.Z(2), qml.X(1) @ qml.X(2)]),
+    0.25 * qml.Z(2) - 0.75 * qml.X(1) @ qml.X(2),
+    qml.Z(2) + qml.X(1) @ qml.X(2),
+    qml.ops.LinearCombination([-0.25, 0.75j], [qml.Z(3), qml.X(2) @ qml.X(3)]),
+    qml.ops.LinearCombination([-0.25 + 0.1j, 0.75j], [qml.Z(4), qml.X(4) @ qml.X(5)]),
+]
+
+
 class TestPrepSelPrep:
     """Test the correctness of the decomposition"""
 
@@ -266,18 +277,7 @@ class TestPrepSelPrep:
 
         qml.assert_equal(op, op_copy)
 
-    @pytest.mark.parametrize(
-        ("lcu"),
-        [
-            qml.ops.LinearCombination([0.25, 0.75], [qml.Z(2), qml.X(1) @ qml.X(2)]),
-            qml.dot([0.25, 0.75], [qml.Z(2), qml.X(1) @ qml.X(2)]),
-            qml.Hamiltonian([0.25, 0.75], [qml.Z(2), qml.X(1) @ qml.X(2)]),
-            0.25 * qml.Z(2) - 0.75 * qml.X(1) @ qml.X(2),
-            qml.Z(2) + qml.X(1) @ qml.X(2),
-            qml.ops.LinearCombination([-0.25, 0.75j], [qml.Z(3), qml.X(2) @ qml.X(3)]),
-            qml.ops.LinearCombination([-0.25 + 0.1j, 0.75j], [qml.Z(4), qml.X(4) @ qml.X(5)]),
-        ],
-    )
+    @pytest.mark.parametrize("lcu", a_set_of_lcus)
     def test_flatten_unflatten(self, lcu):
         """Test that the class can be correctly flattened and unflattened"""
 
@@ -305,6 +305,41 @@ class TestPrepSelPrep:
         assert op.target_wires == new_op.target_wires
         assert op is not new_op
 
+    @pytest.mark.parametrize("lcu", a_set_of_lcus)
+    def test_label(self, lcu):
+        """Test the custom label method of PrepSelPrep."""
+        op = qml.PrepSelPrep(lcu, control=0)
+        op_with_id = qml.PrepSelPrep(lcu, control=0, id="myID")
+
+        # Default
+        assert op.label() == "PrepSelPrep"
+        assert op_with_id.label() == 'PrepSelPrep("myID")'
+
+        # decimals do not affect label
+        assert op.label(decimals=3) == "PrepSelPrep"
+        assert op_with_id.label(decimals=3) == 'PrepSelPrep("myID")'
+
+        # use different base label
+        assert op.label(base_label="U(A)") == "U(A)"
+        assert op_with_id.label(base_label="U(A)") == 'U(A)("myID")'
+
+        # use cache without matrices
+        assert op.label(cache={}) == "PrepSelPrep"
+        assert op_with_id.label(cache={}) == 'PrepSelPrep("myID")'
+
+        # use cache with empty matrices
+        assert op.label(cache={"matrices": []}) == "PrepSelPrep(M0)"
+        assert op_with_id.label(cache={"matrices": []}) == 'PrepSelPrep(M0,"myID")'
+
+        # use cache with non-empty matrices
+        assert op.label(cache={"matrices": [0.1]}) == "PrepSelPrep(M1)"
+        assert op_with_id.label(cache={"matrices": [0.1, 0.6]}) == 'PrepSelPrep(M2,"myID")'
+
+        # use cache with same matrix existing
+        c = qml.math.array(op.coeffs)
+        assert op.label(cache={"matrices": [0.1, c]}) == "PrepSelPrep(M1)"
+        assert op_with_id.label(cache={"matrices": [c, 0.1, 0.6]}) == 'PrepSelPrep(M0,"myID")'
+
 
 def test_control_in_ops():
     """Test that using an operation wire as a control wire results in an error"""