Qutrit channel operation (#5793)

**Context:**
The goal of this PR is to create a qutrit counterpart to the
QubitChannel, enabling noise to be defined using a set of
3×3 Kraus operators.

**Description of the Change:**
- Implement `QutritChannel` similar to `QubitChannel`

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**
Fixes #5649

---------

Co-authored-by: Mudit Pandey <mudit.pandey@xanadu.ai>
Co-authored-by: Thomas R. Bromley <49409390+trbromley@users.noreply.github.com>
Co-authored-by: Christina Lee <chrissie.c.l@gmail.com>

doc/introduction/operations.rst

@@ -509,6 +509,7 @@ Qutrit noisy channels
  ~pennylane.QutritDepolarizingChannel
  ~pennylane.QutritAmplitudeDamping
  ~pennylane.TritFlip
+ ~pennylane.QutritChannel
 
 :html:`</div>`
 

doc/releases/changelog-dev.md

@@ -225,6 +225,9 @@
 
 * `qml.QutritDepolarizingChannel` has been added, allowing for depolarizing noise to be simulated on the `default.qutrit.mixed` device.
  [(#5502)](https://github.com/PennyLaneAI/pennylane/pull/5502)
+
+* `qml.QutritChannel` has been added, enabling the specification of noise using a collection of (3x3) Kraus matrices on the `default.qutrit.mixed` device.
+ [(#5793)](https://github.com/PennyLaneAI/pennylane/issues/5793)
 
 * `qml.QutritAmplitudeDamping` channel has been added, allowing for noise processes modelled by amplitude damping to be simulated on the `default.qutrit.mixed` device.
  [(#5503)](https://github.com/PennyLaneAI/pennylane/pull/5503)
@@ -370,6 +373,7 @@
 
 This release contains contributions from (in alphabetical order):
 
+Tarun Kumar Allamsetty,
 Guillermo Alonso-Linaje,
 Lillian M. A. Frederiksen,
 Gabriel Bottrill,

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__)

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)

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), ())