From d65f671326b71bd9e755a157ec913af31888934b Mon Sep 17 00:00:00 2001 From: Tarun Kumar Allamsetty <70093909+Tarun-Kumar07@users.noreply.github.com> Date: Tue, 11 Jun 2024 04:32:40 +0530 Subject: [PATCH] Qutrit channel operation (#5793) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit **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 Co-authored-by: Thomas R. Bromley <49409390+trbromley@users.noreply.github.com> Co-authored-by: Christina Lee --- doc/introduction/operations.rst | 1 + doc/releases/changelog-dev.md | 4 + pennylane/ops/qutrit/__init__.py | 2 +- pennylane/ops/qutrit/channel.py | 72 ++++++++++++++- tests/ops/qutrit/test_qutrit_channel_ops.py | 97 +++++++++++++++++++++ 5 files changed, 174 insertions(+), 2 deletions(-) diff --git a/doc/introduction/operations.rst b/doc/introduction/operations.rst index ac7e23efcf3..bb978640dbe 100644 --- a/doc/introduction/operations.rst +++ b/doc/introduction/operations.rst @@ -509,6 +509,7 @@ Qutrit noisy channels ~pennylane.QutritDepolarizingChannel ~pennylane.QutritAmplitudeDamping ~pennylane.TritFlip + ~pennylane.QutritChannel :html:`` diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 11fc89c7c1f..12959fa6688 100644 --- a/doc/releases/changelog-dev.md +++ b/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, diff --git a/pennylane/ops/qutrit/__init__.py b/pennylane/ops/qutrit/__init__.py index 7dbc601a62d..c91466c5db0 100644 --- 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 index e9f632761d7..2f667b81edf 100644 --- 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) diff --git a/tests/ops/qutrit/test_qutrit_channel_ops.py b/tests/ops/qutrit/test_qutrit_channel_ops.py index 3497c650b50..049044a74a4 100644 --- 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), ())