Repository: PennyLaneAI/pennylane
Pull request #5503 (2024-04-11 23:35:43)

Qutrit channel amplitude damping
**Context:**
`default.qutrit.mixed` device has been added, but only one channel has been added, [qml.QutritDepolarizingChannel](https://github.com/PennyLaneAI/pennylane/pull/5502/), this adds the second channel to the device so the device can simulate the qutrit equivalent of amplitude damping noise.

**Description of the Change:**
Adds new channel module to `qml.ops.qutrit` package. Adding the second qutrit channel `QutritAmplitudeDamping`.

**Benefits:**
Allows for `defualt.qutrit.mixed` to simulate amplitude damping noise. Making `default.qutrit.mixed`, much more useful.

**Possible Drawbacks:**
N/A

**Related GitHub Issues:**
N/A


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5503?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.67%. Comparing base [(`332d090`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/332d09007f885f1e8ae06fbe99b7112d45f840d9?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`5289ca8`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5503?dropdown=coverage&src=pr&el=desc&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    #5503      +/-   ##
==========================================
- Coverage   99.68%   99.67%   -0.01%     
==========================================
  Files         415      415              
  Lines       38878    38602     -276     
==========================================
- Hits        38754    38477     -277     
- Misses        124      125       +1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5503?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).

mudit2812: The jax failures should fix themselves once you update the branch @Gabriel-Bottrill 
Gabriel-Bottrill: Thank you @mudit2812!

Patch:
diff --git a/doc/introduction/operations.rst b/doc/introduction/operations.rst
index 970e7b4ec..b4b9b0d3c 100644
--- a/doc/introduction/operations.rst
+++ b/doc/introduction/operations.rst
@@ -507,6 +507,7 @@ Qutrit noisy channels
     :nosignatures:
 
     ~pennylane.QutritDepolarizingChannel
+    ~pennylane.QutritAmplitudeDamping
 
 :html:`</div>`
 
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index dce3788a2..052a8975c 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -122,6 +122,9 @@
   returning a list of `QuantumTape`s and a post-processing function instead of simply the transformed circuit.
   [(#5693)](https://github.com/PennyLaneAI/pennylane/pull/5693)
 
+* `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)
+
 <h3>Deprecations 👋</h3>
 
 * The `simplify` argument in `qml.Hamiltonian` and `qml.ops.LinearCombination` is deprecated. 
diff --git a/pennylane/ops/qutrit/__init__.py b/pennylane/ops/qutrit/__init__.py
index 670f950ba..da3ee27ca 100644
--- a/pennylane/ops/qutrit/__init__.py
+++ b/pennylane/ops/qutrit/__init__.py
@@ -26,7 +26,7 @@ from .non_parametric_ops import *
 from .observables import *
 from .parametric_ops import *
 from .state_preparation import *
-from .channel import QutritDepolarizingChannel
+from .channel import *
 
 # TODO: Change `qml.Identity` for qutrit support or add `qml.TIdentity` for qutrits
 __ops__ = {
@@ -50,6 +50,7 @@ __obs__ = {
 }
 __channels__ = {
     "QutritDepolarizingChannel",
+    "QutritAmplitudeDamping",
 }
 
 __all__ = list(__ops__ | __obs__ | __channels__)
diff --git a/pennylane/ops/qutrit/channel.py b/pennylane/ops/qutrit/channel.py
index 8443fd2de..c50572cf2 100644
--- a/pennylane/ops/qutrit/channel.py
+++ b/pennylane/ops/qutrit/channel.py
@@ -222,3 +222,94 @@ class QutritDepolarizingChannel(Channel):
         )
 
         return [identity] + Ks
+
+
+class QutritAmplitudeDamping(Channel):
+    r"""
+    Single-qutrit amplitude damping error channel.
+
+    Interaction with the environment can lead to changes in the state populations of a qutrit.
+    This can be modelled by the qutrit amplitude damping channel with the following Kraus matrices:
+
+    .. math::
+        K_0 = \begin{bmatrix}
+                1 & 0 & 0\\
+                0 & \sqrt{1-\gamma_1} & 0 \\
+                0 & 0 & \sqrt{1-\gamma_2}
+                \end{bmatrix}, \quad
+        K_1 = \begin{bmatrix}
+                0 & \sqrt{\gamma_1} & 0 \\
+                0 & 0 & 0 \\
+                0 & 0 & 0
+                \end{bmatrix}, \quad
+        K_2 = \begin{bmatrix}
+                0 & 0 & \sqrt{\gamma_2} \\
+                0 & 0 & 0 \\
+                0 & 0 & 0
+                \end{bmatrix}
+
+    where :math:`\gamma_1 \in [0, 1]` and :math:`\gamma_2 \in [0, 1]` are the amplitude damping
+    probabilities for subspaces (0,1) and (0,2) respectively.
+
+    .. note::
+
+        The Kraus operators :math:`\{K_0, K_1, K_2\}` are adapted from [`1 <https://doi.org/10.48550/arXiv.1905.10481>`_] (Eq. 8).
+
+    **Details:**
+
+    * Number of wires: 1
+    * Number of parameters: 2
+
+    Args:
+        gamma_1 (float): :math:`|1 \rangle \rightarrow |0 \rangle` amplitude damping probability.
+        gamma_2 (float): :math:`|2 \rangle \rightarrow |0 \rangle` amplitude damping probability.
+        wires (Sequence[int] or int): the wire the channel acts on
+        id (str or None): String representing the operation (optional)
+    """
+
+    num_params = 2
+    num_wires = 1
+    grad_method = "F"
+
+    def __init__(self, gamma_1, gamma_2, wires, id=None):
+        # Verify gamma_1 and gamma_2
+        for gamma in (gamma_1, gamma_2):
+            if not (math.is_abstract(gamma_1) or math.is_abstract(gamma_2)):
+                if not 0.0 <= gamma <= 1.0:
+                    raise ValueError("Each probability must be in the interval [0,1]")
+        super().__init__(gamma_1, gamma_2, wires=wires, id=id)
+
+    @staticmethod
+    def compute_kraus_matrices(gamma_1, gamma_2):  # pylint:disable=arguments-differ
+        r"""Kraus matrices representing the ``QutritAmplitudeDamping`` channel.
+
+        Args:
+            gamma_1 (float): :math:`|1\rangle \rightarrow |0\rangle` amplitude damping probability.
+            gamma_2 (float): :math:`|2\rangle \rightarrow |0\rangle` amplitude damping probability.
+
+        Returns:
+            list(array): list of Kraus matrices
+
+        **Example**
+
+        >>> qml.QutritAmplitudeDamping.compute_kraus_matrices(0.5, 0.25)
+        [
+        array([ [1.        , 0.        , 0.        ],
+                [0.        , 0.70710678, 0.        ],
+                [0.        , 0.        , 0.8660254 ]]),
+        array([ [0.        , 0.70710678, 0.        ],
+                [0.        , 0.        , 0.        ],
+                [0.        , 0.        , 0.        ]]),
+        array([ [0.        , 0.        , 0.5       ],
+                [0.        , 0.        , 0.        ],
+                [0.        , 0.        , 0.        ]])
+        ]
+        """
+        K0 = math.diag([1, math.sqrt(1 - gamma_1 + math.eps), math.sqrt(1 - gamma_2 + math.eps)])
+        K1 = math.sqrt(gamma_1 + math.eps) * math.convert_like(
+            math.cast_like(math.array([[0, 1, 0], [0, 0, 0], [0, 0, 0]]), gamma_1), gamma_1
+        )
+        K2 = math.sqrt(gamma_2 + math.eps) * math.convert_like(
+            math.cast_like(math.array([[0, 0, 1], [0, 0, 0], [0, 0, 0]]), gamma_2), gamma_2
+        )
+        return [K0, K1, K2]


Test patch:
diff --git a/tests/ops/qutrit/test_qutrit_channel_ops.py b/tests/ops/qutrit/test_qutrit_channel_ops.py
index be1025d0c..585f623f4 100644
--- a/tests/ops/qutrit/test_qutrit_channel_ops.py
+++ b/tests/ops/qutrit/test_qutrit_channel_ops.py
@@ -19,6 +19,7 @@ import pytest
 from numpy.linalg import matrix_power
 
 import pennylane as qml
+from pennylane import math
 from pennylane import numpy as pnp
 from pennylane.ops.qutrit import channel
 
@@ -74,7 +75,7 @@ class TestQutritDepolarizingChannel:
         dev = qml.device("default.qutrit.mixed")
         prob = pnp.array(0.5, requires_grad=True)
 
-        @qml.qnode(dev)
+        @qml.qnode(dev, diff_method="parameter-shift")
         def circuit(p):
             qml.TRX(angle, wires=0, subspace=(0, 1))
             qml.TRX(angle, wires=0, subspace=(1, 2))
@@ -113,28 +114,24 @@ class TestQutritDepolarizingChannel:
     @staticmethod
     def kraus_fn(p):
         """Gets a matrix of the Kraus matrices to be tested."""
-        return qml.math.stack(channel.QutritDepolarizingChannel(p, wires=0).kraus_matrices())
+        return math.stack(channel.QutritDepolarizingChannel(p, wires=0).kraus_matrices())
 
     @staticmethod
     def kraus_fn_real(p):
         """Gets a matrix of the real part of the Kraus matrices to be tested."""
-        return qml.math.real(
-            qml.math.stack(channel.QutritDepolarizingChannel(p, wires=0).kraus_matrices())
-        )
+        return math.real(math.stack(channel.QutritDepolarizingChannel(p, wires=0).kraus_matrices()))
 
     @staticmethod
     def kraus_fn_imag(p):
         """Gets a matrix of the imaginary part of the Kraus matrices to be tested."""
-        return qml.math.imag(
-            qml.math.stack(channel.QutritDepolarizingChannel(p, wires=0).kraus_matrices())
-        )
+        return math.imag(math.stack(channel.QutritDepolarizingChannel(p, wires=0).kraus_matrices()))
 
     @pytest.mark.autograd
     def test_kraus_jac_autograd(self):
         """Tests Jacobian of Kraus matrices using autograd."""
         p = pnp.array(0.43, requires_grad=True)
         jac = qml.jacobian(self.kraus_fn_real)(p) + 1j * qml.jacobian(self.kraus_fn_imag)(p)
-        assert qml.math.allclose(jac, self.expected_jac_fn(p))
+        assert math.allclose(jac, self.expected_jac_fn(p))
 
     @pytest.mark.torch
     def test_kraus_jac_torch(self):
@@ -144,7 +141,7 @@ class TestQutritDepolarizingChannel:
         p = torch.tensor(0.43, requires_grad=True)
         jacobian = torch.autograd.functional.jacobian
         jac = jacobian(self.kraus_fn_real, p) + 1j * jacobian(self.kraus_fn_imag, p)
-        assert qml.math.allclose(jac, self.expected_jac_fn(p.detach().numpy()))
+        assert math.allclose(jac, self.expected_jac_fn(p.detach().numpy()))
 
     @pytest.mark.tf
     def test_kraus_jac_tf(self):
@@ -157,10 +154,10 @@ class TestQutritDepolarizingChannel:
         with tf.GradientTape() as imag_tape:
             imag_out = self.kraus_fn_imag(p)
 
-        real_jac = qml.math.cast(real_tape.jacobian(real_out, p), complex)
-        imag_jac = qml.math.cast(imag_tape.jacobian(imag_out, p), complex)
+        real_jac = math.cast(real_tape.jacobian(real_out, p), complex)
+        imag_jac = math.cast(imag_tape.jacobian(imag_out, p), complex)
         jac = real_jac + 1j * imag_jac
-        assert qml.math.allclose(jac, self.expected_jac_fn(0.43))
+        assert math.allclose(jac, self.expected_jac_fn(0.43))
 
     @pytest.mark.jax
     def test_kraus_jac_jax(self):
@@ -171,4 +168,98 @@ class TestQutritDepolarizingChannel:
 
         p = jax.numpy.array(0.43, dtype=jax.numpy.complex128)
         jac = jax.jacobian(self.kraus_fn, holomorphic=True)(p)
-        assert qml.math.allclose(jac, self.expected_jac_fn(p))
+        assert math.allclose(jac, self.expected_jac_fn(p))
+
+
+class TestQutritAmplitudeDamping:
+    """Tests for the qutrit quantum channel QutritAmplitudeDamping"""
+
+    def test_gamma_zero(self, tol):
+        """Test gamma_1=gamma_2=0 gives correct Kraus matrices"""
+        kraus_mats = qml.QutritAmplitudeDamping(0, 0, wires=0).kraus_matrices()
+        assert np.allclose(kraus_mats[0], np.eye(3), atol=tol, rtol=0)
+        assert np.allclose(kraus_mats[1], np.zeros((3, 3)), atol=tol, rtol=0)
+        assert np.allclose(kraus_mats[2], np.zeros((3, 3)), atol=tol, rtol=0)
+
+    @pytest.mark.parametrize("gamma1,gamma2", ((0.1, 0.2), (0.75, 0.75)))
+    def test_gamma_arbitrary(self, gamma1, gamma2, tol):
+        """Test the correct Kraus matrices are returned, also ensures that the sum of gammas can be over 1."""
+        K_0 = np.diag((1, np.sqrt(1 - gamma1), np.sqrt(1 - gamma2)))
+
+        K_1 = np.zeros((3, 3))
+        K_1[0, 1] = np.sqrt(gamma1)
+
+        K_2 = np.zeros((3, 3))
+        K_2[0, 2] = np.sqrt(gamma2)
+
+        expected = [K_0, K_1, K_2]
+        damping_channel = qml.QutritAmplitudeDamping(gamma1, gamma2, wires=0)
+        assert np.allclose(damping_channel.kraus_matrices(), expected, atol=tol, rtol=0)
+
+    @pytest.mark.parametrize("gamma1,gamma2", ((1.5, 0.0), (0.0, 1.0 + math.eps)))
+    def test_gamma_invalid_parameter(self, gamma1, gamma2):
+        """Ensures that error is thrown when gamma_1 or gamma_2 are outside [0,1]"""
+        with pytest.raises(ValueError, match="Each probability must be in the interval"):
+            channel.QutritAmplitudeDamping(gamma1, gamma2, wires=0).kraus_matrices()
+
+    @staticmethod
+    def expected_jac_fn(gamma_1, gamma_2):
+        """Gets the expected Jacobian of Kraus matrices"""
+        partial_1 = [math.zeros((3, 3)) for _ in range(3)]
+        partial_1[0][1, 1] = -1 / (2 * math.sqrt(1 - gamma_1))
+        partial_1[1][0, 1] = 1 / (2 * math.sqrt(gamma_1))
+
+        partial_2 = [math.zeros((3, 3)) for _ in range(3)]
+        partial_2[0][2, 2] = -1 / (2 * math.sqrt(1 - gamma_2))
+        partial_2[2][0, 2] = 1 / (2 * math.sqrt(gamma_2))
+
+        return [partial_1, partial_2]
+
+    @staticmethod
+    def kraus_fn(gamma_1, gamma_2):
+        """Gets the Kraus matrices of QutritAmplitudeDamping channel, used for differentiation."""
+        damping_channel = qml.QutritAmplitudeDamping(gamma_1, gamma_2, wires=0)
+        return math.stack(damping_channel.kraus_matrices())
+
+    @pytest.mark.autograd
+    def test_kraus_jac_autograd(self):
+        """Tests Jacobian of Kraus matrices using autograd."""
+        gamma_1 = pnp.array(0.43, requires_grad=True)
+        gamma_2 = pnp.array(0.12, requires_grad=True)
+        jac = qml.jacobian(self.kraus_fn)(gamma_1, gamma_2)
+        assert math.allclose(jac, self.expected_jac_fn(gamma_1, gamma_2))
+
+    @pytest.mark.torch
+    def test_kraus_jac_torch(self):
+        """Tests Jacobian of Kraus matrices using PyTorch."""
+        import torch
+
+        gamma_1 = torch.tensor(0.43, requires_grad=True)
+        gamma_2 = torch.tensor(0.12, requires_grad=True)
+
+        jac = torch.autograd.functional.jacobian(self.kraus_fn, (gamma_1, gamma_2))
+        expected = self.expected_jac_fn(gamma_1.detach().numpy(), gamma_2.detach().numpy())
+        assert math.allclose(jac[0].detach().numpy(), expected[0])
+        assert math.allclose(jac[1].detach().numpy(), expected[1])
+
+    @pytest.mark.tf
+    def test_kraus_jac_tf(self):
+        """Tests Jacobian of Kraus matrices using TensorFlow."""
+        import tensorflow as tf
+
+        gamma_1 = tf.Variable(0.43)
+        gamma_2 = tf.Variable(0.12)
+        with tf.GradientTape() as tape:
+            out = self.kraus_fn(gamma_1, gamma_2)
+        jac = tape.jacobian(out, (gamma_1, gamma_2))
+        assert math.allclose(jac, self.expected_jac_fn(gamma_1, gamma_2))
+
+    @pytest.mark.jax
+    def test_kraus_jac_jax(self):
+        """Tests Jacobian of Kraus matrices using JAX."""
+        import jax
+
+        gamma_1 = jax.numpy.array(0.43)
+        gamma_2 = jax.numpy.array(0.12)
+        jac = jax.jacobian(self.kraus_fn, argnums=[0, 1])(gamma_1, gamma_2)
+        assert math.allclose(jac, self.expected_jac_fn(gamma_1, gamma_2))