Added |2> -> |1> amplitude damping to amplitude damping. (#5757)

**Context:**
I recently added `qml.QutritAmplitudeDamping` based on specifications
from this [paper](https://doi.org/10.48550/arXiv.1905.10481). This paper
only considers the transition from |2> to |0>, in reality on a super
conducting device it is more likely to have a |2> to |1> transition.
This PR adds a variable and a new Kraus matrix to
`qml.QutritAmplitudeDamping` to allow for the simulation of the |2> to
|1> transition.

**Description of the Change:**
Adds a new variable, \gamma_3, to `qml.QutritAmplitudeDamping` that
defines the new Kraus matrix K_3. This new Kraus operator allows for the
simulation of |2> to |1> relaxations.

**Benefits:**
Allows for simulation of |2> to |1> relaxation. This will make for more
accurate simulations and is useful for measurement error for qutrits.

**Possible Drawbacks:**
This changes `qml.QutritAmplitudeDamping` and adds an extra parameter,
the naming may be a bit more confusing.

**Related GitHub Issues:**
N/A

---------

Co-authored-by: Olivia Di Matteo <2068515+glassnotes@users.noreply.github.com>

doc/releases/changelog-dev.md

@@ -152,6 +152,10 @@
 * ``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.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)
+  [(#5757)](https://github.com/PennyLaneAI/pennylane/pull/5757)
+
 <h3>Breaking changes 💔</h3>
 
 * A custom decomposition can no longer be provided to `QDrift`. Instead, apply the operations in your custom
@@ -177,9 +181,6 @@
 * `Controlled.wires` does not include `self.work_wires` anymore. That can be accessed separately through `Controlled.work_wires`.
   Consequently, `Controlled.active_wires` has been removed in favour of the more common `Controlled.wires`.
   [(#5728)](https://github.com/PennyLaneAI/pennylane/pull/5728)
-
-* `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>
 

pennylane/__init__.py

@@ -218,6 +218,9 @@ def device(name, *args, **kwargs):
     * :mod:`'default.qutrit' <pennylane.devices.default_qutrit>`: a simple
       state simulator of qutrit-based quantum circuit architectures.
 
+    * :mod:`'default.qutrit.mixed' <pennylane.devices.default_qutrit_mixed>`: a
+      mixed-state simulator of qutrit-based quantum circuit architectures.
+
     * :mod:`'default.gaussian' <pennylane.devices.default_gaussian>`: a simple simulator
       of Gaussian states and operations on continuous-variable circuit architectures.
 

pennylane/ops/qutrit/channel.py

@@ -235,8 +235,10 @@ class QutritAmplitudeDamping(Channel):
         K_0 = \begin{bmatrix}
                 1 & 0 & 0\\
                 0 & \sqrt{1-\gamma_1} & 0 \\
-                0 & 0 & \sqrt{1-\gamma_2}
-                \end{bmatrix}, \quad
+                0 & 0 & \sqrt{1-(\gamma_2+\gamma_3)}
+                \end{bmatrix}
+
+    .. math::
         K_1 = \begin{bmatrix}
                 0 & \sqrt{\gamma_1} & 0 \\
                 0 & 0 & 0 \\
@@ -246,70 +248,95 @@ class QutritAmplitudeDamping(Channel):
                 0 & 0 & \sqrt{\gamma_2} \\
                 0 & 0 & 0 \\
                 0 & 0 & 0
+                \end{bmatrix}, \quad
+        K_3 = \begin{bmatrix}
+                0 & 0 & 0 \\
+                0 & 0 & \sqrt{\gamma_3} \\
+                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.
+    where :math:`\gamma_1, \gamma_2, \gamma_3 \in [0, 1]` are the amplitude damping
+    probabilities for subspaces (0,1), (0,2), and (1,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).
+        When :math:`\gamma_3=0` then Kraus operators :math:`\{K_0, K_1, K_2\}` are adapted from
+        [`1 <https://doi.org/10.48550/arXiv.1905.10481>`_] (Eq. 8).
+
+        The Kraus operator :math:`K_3` represents the :math:`|2 \rangle \rightarrow |1 \rangle` transition which is more
+        likely on some devices [`2 <https://arxiv.org/abs/2003.03307>`_] (Sec II.A).
+
+        To maintain normalization :math:`\gamma_2 + \gamma_3 \leq 1`.
+
 
     **Details:**
 
     * Number of wires: 1
-    * Number of parameters: 2
+    * Number of parameters: 3
 
     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.
+        gamma_3 (float): :math:`|2 \rangle \rightarrow |1 \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_params = 3
     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)):
+    def __init__(self, gamma_1, gamma_2, gamma_3, wires, id=None):
+        # Verify input
+        for gamma in (gamma_1, gamma_2, gamma_3):
+            if not math.is_abstract(gamma):
                 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)
+        if not (math.is_abstract(gamma_2) or math.is_abstract(gamma_3)):
+            if not 0.0 <= gamma_2 + gamma_3 <= 1.0:
+                raise ValueError(r"\gamma_2+\gamma_3 must be in the interval [0,1]")
+        super().__init__(gamma_1, gamma_2, gamma_3, wires=wires, id=id)
 
     @staticmethod
-    def compute_kraus_matrices(gamma_1, gamma_2):  # pylint:disable=arguments-differ
+    def compute_kraus_matrices(gamma_1, gamma_2, gamma_3):  # 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.
+            gamma_3 (float): :math:`|2\rangle \rightarrow |1\rangle` amplitude damping probability.
 
         Returns:
             list(array): list of Kraus matrices
 
         **Example**
 
-        >>> qml.QutritAmplitudeDamping.compute_kraus_matrices(0.5, 0.25)
+        >>> qml.QutritAmplitudeDamping.compute_kraus_matrices(0.5, 0.25, 0.36)
         [
         array([ [1.        , 0.        , 0.        ],
                 [0.        , 0.70710678, 0.        ],
-                [0.        , 0.        , 0.8660254 ]]),
+                [0.        , 0.        , 0.6244998 ]]),
         array([ [0.        , 0.70710678, 0.        ],
                 [0.        , 0.        , 0.        ],
                 [0.        , 0.        , 0.        ]]),
         array([ [0.        , 0.        , 0.5       ],
                 [0.        , 0.        , 0.        ],
                 [0.        , 0.        , 0.        ]])
+        array([ [0.        , 0.        , 0.        ],
+                [0.        , 0.        , 0.6       ],
+                [0.        , 0.        , 0.        ]])
         ]
         """
-        K0 = math.diag([1, math.sqrt(1 - gamma_1 + math.eps), math.sqrt(1 - gamma_2 + math.eps)])
+        K0 = math.diag(
+            [1, math.sqrt(1 - gamma_1 + math.eps), math.sqrt(1 - gamma_2 - gamma_3 + 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]
+        K3 = math.sqrt(gamma_3 + math.eps) * math.convert_like(
+            math.cast_like(math.array([[0, 0, 0], [0, 0, 1], [0, 0, 0]]), gamma_3), gamma_3
+        )
+        return [K0, K1, K2, K3]

tests/devices/qutrit_mixed/test_qutrit_mixed_preprocessing.py

@@ -126,6 +126,7 @@ def test_measurement_is_swapped_out(self, mp_fn, mp_cls, shots):
             (qml.Snapshot(), True),
             (qml.TRX(1.1, 0), True),
             (qml.QutritDepolarizingChannel(0.4, 0), True),
+            (qml.QutritAmplitudeDamping(0.1, 0.2, 0.12, 0), True),
         ],
     )
     def test_accepted_operator(self, op, expected):

tests/ops/qutrit/test_qutrit_channel_ops.py

@@ -176,58 +176,75 @@ class TestQutritAmplitudeDamping:
 
     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()
+        kraus_mats = qml.QutritAmplitudeDamping(0, 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)
+        for kraus_mat in kraus_mats[1:]:
+            assert np.allclose(kraus_mat, 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)))
+    @pytest.mark.parametrize("gamma1,gamma2,gamma3", ((0.1, 0.2, 0.3), (0.75, 0.75, 0.25)))
+    def test_gamma_arbitrary(self, gamma1, gamma2, gamma3, tol):
+        """Test the correct Kraus matrices are returned."""
+        K_0 = np.diag((1, np.sqrt(1 - gamma1), np.sqrt(1 - gamma2 - gamma3)))
 
         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)
+        K_3 = np.zeros((3, 3))
+        K_3[1, 2] = np.sqrt(gamma3)
+
+        expected = [K_0, K_1, K_2, K_3]
+        damping_channel = qml.QutritAmplitudeDamping(gamma1, gamma2, gamma3, 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()
+    @pytest.mark.parametrize(
+        "gamma1,gamma2,gamma3",
+        (
+            (1.5, 0.0, 0.0),
+            (0.0, 1.0 + math.eps, 0.0),
+            (0.0, 0.0, 1.1),
+            (0.0, 0.33, 0.67 + math.eps),
+        ),
+    )
+    def test_gamma_invalid_parameter(self, gamma1, gamma2, gamma3):
+        """Ensures that error is thrown when gamma_1, gamma_2, gamma_3, or (gamma_2 + gamma_3) are outside [0,1]"""
+        with pytest.raises(ValueError, match="must be in the interval"):
+            channel.QutritAmplitudeDamping(gamma1, gamma2, gamma3, wires=0).kraus_matrices()
 
     @staticmethod
-    def expected_jac_fn(gamma_1, gamma_2):
+    def expected_jac_fn(gamma_1, gamma_2, gamma_3):
         """Gets the expected Jacobian of Kraus matrices"""
-        partial_1 = [math.zeros((3, 3)) for _ in range(3)]
+        partial_1 = [math.zeros((3, 3)) for _ in range(4)]
         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 = [math.zeros((3, 3)) for _ in range(4)]
+        partial_2[0][2, 2] = -1 / (2 * math.sqrt(1 - gamma_2 - gamma_3))
         partial_2[2][0, 2] = 1 / (2 * math.sqrt(gamma_2))
 
-        return [partial_1, partial_2]
+        partial_3 = [math.zeros((3, 3)) for _ in range(4)]
+        partial_3[0][2, 2] = -1 / (2 * math.sqrt(1 - gamma_2 - gamma_3))
+        partial_3[3][1, 2] = 1 / (2 * math.sqrt(gamma_3))
+
+        return [partial_1, partial_2, partial_3]
 
     @staticmethod
-    def kraus_fn(gamma_1, gamma_2):
+    def kraus_fn(gamma_1, gamma_2, gamma_3):
         """Gets the Kraus matrices of QutritAmplitudeDamping channel, used for differentiation."""
-        damping_channel = qml.QutritAmplitudeDamping(gamma_1, gamma_2, wires=0)
+        damping_channel = qml.QutritAmplitudeDamping(gamma_1, gamma_2, gamma_3, 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))
+        gamma_3 = pnp.array(0.35, requires_grad=True)
+
+        jac = qml.jacobian(self.kraus_fn)(gamma_1, gamma_2, gamma_3)
+        assert math.allclose(jac, self.expected_jac_fn(gamma_1, gamma_2, gamma_3))
 
     @pytest.mark.torch
     def test_kraus_jac_torch(self):
@@ -236,11 +253,15 @@ def test_kraus_jac_torch(self):
 
         gamma_1 = torch.tensor(0.43, requires_grad=True)
         gamma_2 = torch.tensor(0.12, requires_grad=True)
+        gamma_3 = torch.tensor(0.35, 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])
+        jac = torch.autograd.functional.jacobian(self.kraus_fn, (gamma_1, gamma_2, gamma_3))
+        expected = self.expected_jac_fn(
+            gamma_1.detach().numpy(), gamma_2.detach().numpy(), gamma_3.detach().numpy()
+        )
+
+        for res_partial, exp_partial in zip(jac, expected):
+            assert math.allclose(res_partial.detach().numpy(), exp_partial)
 
     @pytest.mark.tf
     def test_kraus_jac_tf(self):
@@ -249,10 +270,12 @@ def test_kraus_jac_tf(self):
 
         gamma_1 = tf.Variable(0.43)
         gamma_2 = tf.Variable(0.12)
+        gamma_3 = tf.Variable(0.35)
+
         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))
+            out = self.kraus_fn(gamma_1, gamma_2, gamma_3)
+        jac = tape.jacobian(out, (gamma_1, gamma_2, gamma_3))
+        assert math.allclose(jac, self.expected_jac_fn(gamma_1, gamma_2, gamma_3))
 
     @pytest.mark.jax
     def test_kraus_jac_jax(self):
@@ -261,5 +284,7 @@ def test_kraus_jac_jax(self):
 
         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))
+        gamma_3 = jax.numpy.array(0.35)
+
+        jac = jax.jacobian(self.kraus_fn, argnums=[0, 1, 2])(gamma_1, gamma_2, gamma_3)
+        assert math.allclose(jac, self.expected_jac_fn(gamma_1, gamma_2, gamma_3))