Repository: PennyLaneAI/pennylane

Issue:
Add a QutritChannel operation
### Feature details

Version 0.36 of PennyLane includes a [`"default.qutrit.mixed"`](https://docs.pennylane.ai/en/latest/code/api/pennylane.devices.default_qutrit_mixed.DefaultQutritMixed.html) device that is capable of simulating [qutrit systems](https://pennylane.ai/qml/demos/tutorial_qutrits_bernstein_vazirani/) in the mixed-state setting. This unlocks the ability to simulate qutrit-based noise, and we would like to add more qutrit noise channels to complement our existing range of [qutrit operators](https://docs.pennylane.ai/en/stable/introduction/operations.html#qutrit-operators).

This objective of this issue is to provide a qutrit analogue of [`QubitChannel`](https://docs.pennylane.ai/en/stable/code/api/pennylane.QubitChannel.html), which would allow noise to be specified via a collection of $(3\times 3)$-dimensional [Kraus operators](https://en.wikipedia.org/wiki/Quantum_operation#Kraus_operators).

### Implementation

It should be possible to execute the following code:

```python
import pennylane as qml
import numpy as np

kraus = [  # See https://github.com/PennyLaneAI/pennylane/pull/5503
    np.array([[1, 0, 0],
              [0, 0.70710678, 0],
              [0, 0, 0.8660254]]),
    np.array([[0, 0.70710678, 0],
              [0, 0, 0],
              [0, 0, 0]]),
    np.array([[0, 0, 0.5],
              [0, 0, 0],
              [0, 0, 0]])
]

dev = qml.device("default.qutrit.mixed")

@qml.qnode(dev)
def f():
    qml.QutritChannel(kraus, 0)
    return qml.expval(qml.GellMann(wires=0, index=1))

f()
```

This can be done by following these instructions:

1. Implement a `qml.QutritChannel` operation in `pennylane/ops/channel.py`, which defines the qutrit equivalent of the `qml.QubitChannel` operation. This operation should:
    * Take a list containing three $3 \times 3$ Kraus matrices, along with `wires` and `id` (optional string representation of the operator).
    * Define the `compute_kraus_matrices`, which is a `staticmethod` that returns the Kraus matrices in a list, and `_flatten` functions (just like in `qml.QubitChannel`). 

2. Add unit and integration tests:
    * Unit tests must verify the correctness of the operation, its attributes, and its methods.
    * Integration tests must show that the output of a `QNode` that applies the `qml.QutritChannel` operation is correct.
    * Integration tests must also verify that `QNode`s that apply the `qml.QutritChannel` operation can be differentiated using the `parameter-shift`, `finite-diff`, and `backprop` differentiation methods and that the computed gradients/jacobians are correct.
    * Integration tests must be implemented for both the `default.qutrit` and `default.qutrit.mixed` devices.

A successful PR must include:

- The requested feature
- Comprehensive tests
- Relevant updates to the documentation
- An entry in the [release notes](https://github.com/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md)
- Passing all CI tests

Check out our [development guide](https://docs.pennylane.ai/en/stable/development/guide.html) for more information, and don't hesitate to reach out to us if you need assistance.

### How important would you say this feature is?

1: Not important. Would be nice to have.

### Additional information

Existing work is underway to support specific qutrit channels in PRs https://github.com/PennyLaneAI/pennylane/pull/5502 and https://github.com/PennyLaneAI/pennylane/pull/5503.


Discussion:
Hi @trbromley , I want to tackle this issue as part of unitaryhack.
Could you please assign it to me?

Resolving patch:
diff --git a/pennylane/ops/qutrit/__init__.py b/pennylane/ops/qutrit/__init__.py
--- a/pennylane/ops/qutrit/__init__.py
+++ b/pennylane/ops/qutrit/__init__.py
@@ -48,6 +48,6 @@
     "THermitian",
     "GellMann",
 }
-__channels__ = {"QutritDepolarizingChannel", "QutritAmplitudeDamping", "TritFlip"}
+__channels__ = {"QutritDepolarizingChannel", "QutritAmplitudeDamping", "TritFlip", "QutritChannel"}
 
 __all__ = list(__ops__ | __obs__ | __channels__)
diff --git a/pennylane/ops/qutrit/channel.py b/pennylane/ops/qutrit/channel.py
--- a/pennylane/ops/qutrit/channel.py
+++ b/pennylane/ops/qutrit/channel.py
@@ -19,7 +19,7 @@
 import numpy as np
 
 from pennylane import math
-from pennylane.operation import Channel
+from pennylane.operation import AnyWires, Channel
 
 QUDIT_DIM = 3
 
@@ -457,3 +457,73 @@ def compute_kraus_matrices(p_01, p_02, p_12):  # pylint:disable=arguments-differ
             math.cast_like(math.array([[1, 0, 0], [0, 0, 1], [0, 1, 0]]), p_12), p_12
         )
         return [K0, K1, K2, K3]
+
+
+class QutritChannel(Channel):
+    r"""
+    Apply an arbitrary fixed qutrit channel.
+
+    Kraus matrices that represent the fixed channel are provided
+    as a list of NumPy arrays.
+
+    **Details:**
+
+    * Number of wires: Any (the operation can act on any number of wires)
+    * Number of parameters: 1
+    * Gradient recipe: None
+
+    Args:
+        K_list (list[array[complex]]): list of Kraus matrices
+        wires (Union[Wires, Sequence[int], or int]): the wire(s) the operation acts on
+        id (str or None): String representing the operation (optional)
+    """
+
+    num_wires = AnyWires
+    grad_method = None
+
+    def __init__(self, K_list, wires=None, id=None):
+        super().__init__(*K_list, wires=wires, id=id)
+
+        # check all Kraus matrices are square matrices
+        if any(K.shape[0] != K.shape[1] for K in K_list):
+            raise ValueError(
+                "Only channels with the same input and output Hilbert space dimensions can be applied."
+            )
+
+        # check all Kraus matrices have the same shape
+        if any(K.shape != K_list[0].shape for K in K_list):
+            raise ValueError("All Kraus matrices must have the same shape.")
+
+        # check the dimension of all Kraus matrices are valid
+        kraus_dim = QUDIT_DIM ** len(self.wires)
+        if any(K.shape[0] != kraus_dim for K in K_list):
+            raise ValueError(f"Shape of all Kraus matrices must be ({kraus_dim},{kraus_dim}).")
+
+        # check that the channel represents a trace-preserving map
+        if not any(math.is_abstract(K) for K in K_list):
+            K_arr = math.array(K_list)
+            Kraus_sum = math.einsum("ajk,ajl->kl", K_arr.conj(), K_arr)
+            if not math.allclose(Kraus_sum, math.eye(K_list[0].shape[0])):
+                raise ValueError("Only trace preserving channels can be applied.")
+
+    def _flatten(self):
+        return (self.data,), (self.wires, ())
+
+    @staticmethod
+    def compute_kraus_matrices(*kraus_matrices):  # pylint:disable=arguments-differ
+        """Kraus matrices representing the QutritChannel channel.
+
+        Args:
+            *K_list (list[array[complex]]): list of Kraus matrices
+
+        Returns:
+            list (array): list of Kraus matrices
+
+        **Example**
+
+        >>> K_list = qml.QutritDepolarizingChannel(0.75, wires=0).kraus_matrices()
+        >>> res = qml.QutritChannel.compute_kraus_matrices(K_list)
+        >>> all(np.allclose(r, k) for r, k  in zip(res, K_list))
+        True
+        """
+        return list(kraus_matrices)


Test patch:
diff --git a/tests/ops/qutrit/test_qutrit_channel_ops.py b/tests/ops/qutrit/test_qutrit_channel_ops.py
--- a/tests/ops/qutrit/test_qutrit_channel_ops.py
+++ b/tests/ops/qutrit/test_qutrit_channel_ops.py
@@ -391,3 +391,100 @@ def test_kraus_jac_jax(self):
         p_12 = jax.numpy.array(0.23)
         jac = jax.jacobian(self.kraus_fn, argnums=[0, 1, 2])(p_01, p_02, p_12)
         assert qml.math.allclose(jac, self.expected_jac_fn(p_01, p_02, p_12))
+
+
+class TestQutritChannel:
+    """Tests for the quantum channel QubitChannel"""
+
+    def test_input_correctly_handled(self, tol):
+        """Test that Kraus matrices are correctly processed"""
+        K_list = qml.QutritDepolarizingChannel(0.75, wires=0).kraus_matrices()
+        out = qml.QutritChannel(K_list, wires=0).kraus_matrices()
+
+        assert np.allclose(out, K_list, atol=tol, rtol=0)
+
+    def test_kraus_matrices_are_square(self):
+        """Tests that the given Kraus matrices are square"""
+        K_list = [np.zeros((3, 3)), np.zeros((2, 3))]
+        with pytest.raises(
+            ValueError, match="Only channels with the same input and output Hilbert space"
+        ):
+            qml.QutritChannel(K_list, wires=0)
+
+    def test_kraus_matrices_are_of_same_shape(self):
+        """Tests that the given Kraus matrices are of same shape"""
+        K_list = [np.eye(3), np.eye(4)]
+        with pytest.raises(ValueError, match="All Kraus matrices must have the same shape."):
+            qml.QutritChannel(K_list, wires=0)
+
+    def test_kraus_matrices_are_dimensions(self):
+        """Tests that the given Kraus matrices are of right dimension i.e (9,9)"""
+        K_list = [np.eye(3), np.eye(3)]
+        with pytest.raises(ValueError, match=r"Shape of all Kraus matrices must be \(9,9\)."):
+            qml.QutritChannel(K_list, wires=[0, 1])
+
+    def test_kraus_matrices_are_trace_preserved(self):
+        """Tests that the channel represents a trace-preserving map"""
+        K_list = [0.75 * np.eye(3), 0.35j * np.eye(3)]
+        with pytest.raises(ValueError, match="Only trace preserving channels can be applied."):
+            qml.QutritChannel(K_list, wires=0)
+
+    @pytest.mark.parametrize("diff_method", ["parameter-shift", "finite-diff", "backprop"])
+    def test_integrations(self, diff_method):
+        """Test integration"""
+        kraus = [
+            np.array([[1, 0, 0], [0, 0.70710678, 0], [0, 0, 0.8660254]]),
+            np.array([[0, 0.70710678, 0], [0, 0, 0], [0, 0, 0]]),
+            np.array([[0, 0, 0.5], [0, 0, 0], [0, 0, 0]]),
+        ]
+
+        dev = qml.device("default.qutrit.mixed", wires=1)
+
+        @qml.qnode(dev, diff_method=diff_method)
+        def func():
+            qml.QutritChannel(kraus, 0)
+            return qml.expval(qml.GellMann(wires=0, index=1))
+
+        func()
+
+    @pytest.mark.parametrize("diff_method", ["parameter-shift", "finite-diff", "backprop"])
+    def test_integration_grad(self, diff_method):
+        """Test integration with grad"""
+        dev = qml.device("default.qutrit.mixed", wires=1)
+
+        @qml.qnode(dev, diff_method=diff_method)
+        def func(p):
+            kraus = qml.QutritDepolarizingChannel.compute_kraus_matrices(p)
+            qml.QutritChannel(kraus, 0)
+            return qml.expval(qml.GellMann(wires=0, index=1))
+
+        qml.grad(func)(0.5)
+
+    @pytest.mark.parametrize("diff_method", ["parameter-shift", "finite-diff", "backprop"])
+    def test_integration_jacobian(self, diff_method):
+        """Test integration with grad"""
+        dev = qml.device("default.qutrit.mixed", wires=1)
+
+        @qml.qnode(dev, diff_method=diff_method)
+        def func(p):
+            kraus = qml.QutritDepolarizingChannel.compute_kraus_matrices(p)
+            qml.QutritChannel(kraus, 0)
+            return qml.expval(qml.GellMann(wires=0, index=1))
+
+        qml.jacobian(func)(0.5)
+
+    def test_flatten(self):
+        """Test flatten method returns kraus matrices and wires"""
+        kraus = [
+            np.array([[1, 0, 0], [0, 0.70710678, 0], [0, 0, 0.8660254]]),
+            np.array([[0, 0.70710678, 0], [0, 0, 0], [0, 0, 0]]),
+            np.array([[0, 0, 0.5], [0, 0, 0], [0, 0, 0]]),
+        ]
+
+        qutrit_channel = qml.QutritChannel(kraus, 1, id="test")
+        data, metadata = qutrit_channel._flatten()  # pylint: disable=protected-access
+        new_op = qml.QutritChannel._unflatten(data, metadata)  # pylint: disable=protected-access
+        assert qml.equal(qutrit_channel, new_op)
+
+        assert np.allclose(kraus, data)
+        assert metadata == (qml.wires.Wires(1), ())