Add amplitude damping noise model (#497)

CHANGELOG.md

@@ -11,6 +11,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - #476: Introduced new methods `OpenGraph.extract_circuit`, `CliffordMap.to_tableau` and new function `graphix.circ_ext.compilation.cm_berg_pass`. Circuit extraction can be done natively in Graphix.
 
+- #545: Added an amplitude damping noise model. Introduces `amplitude_damping_channel` / `two_qubit_amplitude_damping_channel`, the `AmplitudeDampingNoise` / `TwoQubitAmplitudeDampingNoise` noise elements, and `AmplitudeDampingNoiseModel`.
+
 - #490: Introduced new `Instruction` and `Command` namespace classes for instruction and command instantiation.
 
 - #505

graphix/__init__.py

@@ -13,7 +13,7 @@
 from graphix.graphsim import GraphState
 from graphix.instruction import Instruction
 from graphix.measurements import BlochMeasurement, Measurement, PauliMeasurement
-from graphix.noise_models import DepolarisingNoiseModel, NoiseModel
+from graphix.noise_models import AmplitudeDampingNoiseModel, DepolarisingNoiseModel, NoiseModel
 from graphix.opengraph import OpenGraph
 from graphix.optimization import StandardizedPattern
 from graphix.parameter import Placeholder
@@ -27,6 +27,7 @@
 
 __all__ = [
     "ANGLE_PI",
+    "AmplitudeDampingNoiseModel",
     "Axis",
     "BasicStates",
     "BlochMeasurement",

graphix/channels.py

@@ -296,3 +296,57 @@ def two_qubit_depolarising_tensor_channel(prob: float) -> KrausChannel:
             KrausData(prob / 3.0, np.kron(Ops.Z, Ops.Y)),
         ]
     )
+
+
+def _amplitude_damping_kraus(prob: float) -> tuple[npt.NDArray[np.complex128], npt.NDArray[np.complex128]]:
+    r"""Return the single-qubit amplitude damping Kraus operators."""
+    k1 = np.array([[1.0, 0.0], [0.0, np.sqrt(1 - prob)]], dtype=np.complex128)
+    k2 = np.array([[0.0, np.sqrt(prob)], [0.0, 0.0]], dtype=np.complex128)
+    return k1, k2
+
+
+def amplitude_damping_channel(prob: float) -> KrausChannel:
+    r"""Single-qubit amplitude damping channel.
+
+    .. math::
+        K_1 = \begin{pmatrix} 1 & 0 \\ 0 & \sqrt{1-\gamma} \end{pmatrix}, \quad
+        K_2 = \begin{pmatrix} 0 & \sqrt{\gamma} \\ 0 & 0 \end{pmatrix}
+
+    Parameters
+    ----------
+    prob : float
+        The damping parameter :math:`\gamma` associated to the channel.
+
+    Returns
+    -------
+    :class:`graphix.channels.KrausChannel` object
+        containing the corresponding Kraus operators
+    """
+    k1, k2 = _amplitude_damping_kraus(prob)
+    return KrausChannel(
+        [
+            KrausData(1.0, k1),
+            KrausData(1.0, k2),
+        ]
+    )
+
+
+def two_qubit_amplitude_damping_channel(prob: float) -> KrausChannel:
+    r"""Two-qubit amplitude damping channel.
+
+    Tensor product of two independent single-qubit amplitude damping channels
+    with the same damping parameter :math:`\gamma`, giving the four Kraus
+    operators :math:`\{K_i \otimes K_j\}` for :math:`i, j \in \{1, 2\}`.
+
+    Parameters
+    ----------
+    prob : float
+        The damping parameter :math:`\gamma` associated to the channel.
+
+    Returns
+    -------
+    :class:`graphix.channels.KrausChannel` object
+        containing the corresponding Kraus operators
+    """
+    operators = _amplitude_damping_kraus(prob)
+    return KrausChannel([KrausData(1.0, np.kron(left, right)) for left in operators for right in operators])

graphix/noise_models/__init__.py

@@ -4,6 +4,11 @@
 
 from typing import TYPE_CHECKING
 
+from graphix.noise_models.amplitude_damping import (
+    AmplitudeDampingNoise,
+    AmplitudeDampingNoiseModel,
+    TwoQubitAmplitudeDampingNoise,
+)
 from graphix.noise_models.depolarising import DepolarisingNoise, DepolarisingNoiseModel, TwoQubitDepolarisingNoise
 from graphix.noise_models.noise_model import (
     ApplyNoise,
@@ -16,11 +21,14 @@
     from graphix.noise_models.noise_model import CommandOrNoise as CommandOrNoise
 
 __all__ = [
+    "AmplitudeDampingNoise",
+    "AmplitudeDampingNoiseModel",
     "ApplyNoise",
     "ComposeNoiseModel",
     "DepolarisingNoise",
     "DepolarisingNoiseModel",
     "Noise",
     "NoiseModel",
+    "TwoQubitAmplitudeDampingNoise",
     "TwoQubitDepolarisingNoise",
 ]

graphix/noise_models/amplitude_damping.py

@@ -0,0 +1,155 @@
+"""Amplitude damping noise model."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import typing_extensions
+
+from graphix.channels import (
+    KrausChannel,
+    amplitude_damping_channel,
+    two_qubit_amplitude_damping_channel,
+)
+from graphix.command import BaseM, CommandKind
+from graphix.measurements import toggle_outcome
+from graphix.noise_models.noise_model import ApplyNoise, Noise, NoiseModel
+from graphix.rng import ensure_rng
+from graphix.utils import Probability
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
+    from numpy.random import Generator
+
+    from graphix.measurements import Outcome
+    from graphix.noise_models.noise_model import CommandOrNoise
+
+
+class AmplitudeDampingNoise(Noise):
+    """One-qubit amplitude damping noise with damping parameter ``prob``."""
+
+    prob = Probability()
+
+    def __init__(self, prob: float) -> None:
+        r"""Initialize one-qubit amplitude damping noise.
+
+        Parameters
+        ----------
+        prob : float
+            Damping parameter :math:`\\gamma` of the noise, between 0 and 1.
+        """
+        self.prob = prob
+
+    @property
+    @typing_extensions.override
+    def nqubits(self) -> int:
+        """Return the number of qubits targetted by the noise element."""
+        return 1
+
+    @typing_extensions.override
+    def to_kraus_channel(self) -> KrausChannel:
+        """Return the Kraus channel describing the noise element."""
+        return amplitude_damping_channel(self.prob)
+
+
+class TwoQubitAmplitudeDampingNoise(Noise):
+    """Two-qubit amplitude damping noise with damping parameter ``prob``."""
+
+    prob = Probability()
+
+    def __init__(self, prob: float) -> None:
+        r"""Initialize two-qubit amplitude damping noise.
+
+        Parameters
+        ----------
+        prob : float
+            Damping parameter :math:`\\gamma` of the noise, between 0 and 1.
+        """
+        self.prob = prob
+
+    @property
+    @typing_extensions.override
+    def nqubits(self) -> int:
+        """Return the number of qubits targetted by the noise element."""
+        return 2
+
+    @typing_extensions.override
+    def to_kraus_channel(self) -> KrausChannel:
+        """Return the Kraus channel describing the noise element."""
+        return two_qubit_amplitude_damping_channel(self.prob)
+
+
+class AmplitudeDampingNoiseModel(NoiseModel):
+    r"""Amplitude damping noise model.
+
+    :param NoiseModel: Parent abstract class class:`NoiseModel`
+    :type NoiseModel: class
+    """
+
+    def __init__(
+        self,
+        prepare_error_prob: float = 0.0,
+        x_error_prob: float = 0.0,
+        z_error_prob: float = 0.0,
+        entanglement_error_prob: float = 0.0,
+        measure_channel_prob: float = 0.0,
+        measure_error_prob: float = 0.0,
+    ) -> None:
+        self.prepare_error_prob = prepare_error_prob
+        self.x_error_prob = x_error_prob
+        self.z_error_prob = z_error_prob
+        self.entanglement_error_prob = entanglement_error_prob
+        self.measure_channel_prob = measure_channel_prob
+        self.measure_error_prob = measure_error_prob
+
+    @typing_extensions.override
+    def input_nodes(
+        self, nodes: Iterable[int], rng: Generator | None = None, *, stacklevel: int = 1
+    ) -> list[CommandOrNoise]:
+        """Return the noise to apply to input nodes."""
+        return [ApplyNoise(noise=AmplitudeDampingNoise(self.prepare_error_prob), nodes=[node]) for node in nodes]
+
+    @typing_extensions.override
+    def command(
+        self, cmd: CommandOrNoise, rng: Generator | None = None, *, stacklevel: int = 1
+    ) -> list[CommandOrNoise]:
+        """Return the noise to apply to the command ``cmd``."""
+        match cmd.kind:
+            case CommandKind.N:
+                return [cmd, ApplyNoise(noise=AmplitudeDampingNoise(self.prepare_error_prob), nodes=[cmd.node])]
+            case CommandKind.E:
+                return [
+                    cmd,
+                    ApplyNoise(
+                        noise=TwoQubitAmplitudeDampingNoise(self.entanglement_error_prob), nodes=list(cmd.nodes)
+                    ),
+                ]
+            case CommandKind.M:
+                return [ApplyNoise(noise=AmplitudeDampingNoise(self.measure_channel_prob), nodes=[cmd.node]), cmd]
+            case CommandKind.X:
+                return [
+                    cmd,
+                    ApplyNoise(noise=AmplitudeDampingNoise(self.x_error_prob), nodes=[cmd.node], domain=cmd.domain),
+                ]
+            case CommandKind.Z:
+                return [
+                    cmd,
+                    ApplyNoise(noise=AmplitudeDampingNoise(self.z_error_prob), nodes=[cmd.node], domain=cmd.domain),
+                ]
+            case CommandKind.C | CommandKind.T | CommandKind.ApplyNoise:
+                return [cmd]
+            case CommandKind.S:
+                raise ValueError("Unexpected signal!")
+            case _:  # pragma: no cover
+                typing_extensions.assert_never(cmd.kind)
+
+    @typing_extensions.override
+    def confuse_result(
+        self, cmd: BaseM, result: Outcome, rng: Generator | None = None, *, stacklevel: int = 1
+    ) -> Outcome:
+        """Assign wrong measurement result with probability ``measure_error_prob``."""
+        rng = ensure_rng(rng, stacklevel=stacklevel + 1)
+        if rng.uniform() < self.measure_error_prob:
+            return toggle_outcome(result)
+        return result

tests/test_density_matrix.py

@@ -12,7 +12,13 @@
 import graphix.random_objects as randobj
 from graphix import command
 from graphix.branch_selector import ConstBranchSelector
-from graphix.channels import KrausChannel, dephasing_channel, depolarising_channel
+from graphix.channels import (
+    KrausChannel,
+    amplitude_damping_channel,
+    dephasing_channel,
+    depolarising_channel,
+    two_qubit_amplitude_damping_channel,
+)
 from graphix.fundamentals import ANGLE_PI, Plane
 from graphix.ops import Ops
 from graphix.sim.density_matrix import DensityMatrix, DensityMatrixBackend
@@ -569,7 +575,7 @@ def test_apply_dephasing_channel(self, fx_rng: Generator) -> None:
         dm = DensityMatrix(randobj.rand_dm(2, fx_rng))
 
         # copy of initial dm
-        rho_test = dm.rho
+        rho_test = np.asarray(dm.rho, dtype=np.complex128)
 
         # create dephasing channel
         prob = fx_rng.uniform()
@@ -647,7 +653,7 @@ def test_apply_depolarising_channel(self, fx_rng: Generator) -> None:
         dm = DensityMatrix(randobj.rand_dm(2, fx_rng))
 
         # copy of initial dm
-        rho_test = dm.rho
+        rho_test = np.asarray(dm.rho, dtype=np.complex128)
 
         # create dephasing channel
         prob = fx_rng.uniform()
@@ -736,6 +742,95 @@ def test_apply_depolarising_channel(self, fx_rng: Generator) -> None:
         assert np.allclose(expected_dm.trace(), 1.0)
         assert np.allclose(dm.rho, expected_dm)
 
+    def test_apply_amplitude_damping_channel(self, fx_rng: Generator) -> None:
+        # check on single qubit first, against the by-hand Kraus sum
+        dm = DensityMatrix(randobj.rand_dm(2, fx_rng))
+        rho_test = np.asarray(dm.rho, dtype=np.complex128)
+
+        gamma = fx_rng.uniform()
+        ad_channel = amplitude_damping_channel(gamma)
+
+        assert isinstance(ad_channel, KrausChannel)
+
+        dm.apply_channel(ad_channel, [0])
+
+        k1 = np.array([[1.0, 0.0], [0.0, np.sqrt(1 - gamma)]], dtype=np.complex128)
+        k2 = np.array([[0.0, np.sqrt(gamma)], [0.0, 0.0]], dtype=np.complex128)
+        expected_dm = k1 @ rho_test @ k1.conj().T + k2 @ rho_test @ k2.conj().T
+
+        assert np.allclose(expected_dm.trace(), 1.0)
+        assert np.allclose(dm.rho, expected_dm)
+
+        # check embedded in a larger random register
+        nqubits = int(fx_rng.integers(2, 5))
+        i = int(fx_rng.integers(0, nqubits))
+
+        psi = _randstate_raw(nqubits, fx_rng)
+        psi /= np.sqrt(np.sum(np.abs(psi) ** 2))
+        dm = DensityMatrix(data=np.outer(psi, psi.conj()))
+
+        gamma = fx_rng.uniform()
+        ad_channel = amplitude_damping_channel(gamma)
+        dm.apply_channel(ad_channel, [i])
+
+        expected_dm = np.zeros((2**nqubits, 2**nqubits), dtype=np.complex128)
+        for elem in ad_channel:
+            psi_evolved = np.tensordot(elem.operator, psi.reshape((2,) * nqubits), (1, i))
+            psi_evolved = np.moveaxis(psi_evolved, 0, i).reshape(2**nqubits)
+            expected_dm += elem.coef * np.conj(elem.coef) * np.outer(psi_evolved, psi_evolved.conj())
+
+        assert np.allclose(expected_dm.trace(), 1.0)
+        assert np.allclose(dm.rho, expected_dm)
+
+    @pytest.mark.parametrize("gamma", [0.0, 0.2, 0.5, 0.9, 1.0])
+    def test_amplitude_damping_ground_state_fixed(self, gamma: float) -> None:
+        # ``|0><0|`` is a fixed point of amplitude damping for any gamma.
+        ket0 = np.array([[1.0, 0.0], [0.0, 0.0]], dtype=np.complex128)
+        dm = DensityMatrix(data=BasicStates.ZERO)
+        dm.apply_channel(amplitude_damping_channel(gamma), [0])
+        assert np.allclose(dm.rho, ket0)
+
+    @pytest.mark.parametrize("gamma", [0.0, 0.2, 0.5, 0.9, 1.0])
+    def test_amplitude_damping_excited_state_decays(self, gamma: float) -> None:
+        # ``|1><1| -> (1 - gamma)|1><1| + gamma|0><0|`` (directional T1 decay).
+        ket0 = np.array([[1.0, 0.0], [0.0, 0.0]], dtype=np.complex128)
+        ket1 = np.array([[0.0, 0.0], [0.0, 1.0]], dtype=np.complex128)
+        dm = DensityMatrix(data=BasicStates.ONE)
+        dm.apply_channel(amplitude_damping_channel(gamma), [0])
+        assert np.allclose(dm.rho, (1 - gamma) * ket1 + gamma * ket0)
+
+    @pytest.mark.parametrize("gamma", [0.1, 0.4, 0.8])
+    def test_amplitude_damping_coherence_decay(self, gamma: float) -> None:
+        # Off-diagonal coherences scale by ``sqrt(1 - gamma)`` (distinct from dephasing).
+        dm = DensityMatrix(data=BasicStates.PLUS)
+        dm.apply_channel(amplitude_damping_channel(gamma), [0])
+        assert np.isclose(dm.rho[0, 1], np.sqrt(1 - gamma) / 2)
+        assert np.isclose(dm.rho[1, 0], np.sqrt(1 - gamma) / 2)
+
+    @pytest.mark.parametrize("gamma", [0.0, 0.25, 0.6, 1.0])
+    def test_apply_two_qubit_amplitude_damping_channel(self, gamma: float, fx_rng: Generator) -> None:
+        # The two-qubit channel equals independent damping on each factor.
+        a = _randstate_raw(1, fx_rng)
+        a /= np.sqrt(np.sum(np.abs(a) ** 2))
+        b = _randstate_raw(1, fx_rng)
+        b /= np.sqrt(np.sum(np.abs(b) ** 2))
+        rho_a = np.outer(a, a.conj())
+        rho_b = np.outer(b, b.conj())
+
+        dm = DensityMatrix(data=np.kron(rho_a, rho_b))
+        dm.apply_channel(two_qubit_amplitude_damping_channel(gamma), [0, 1])
+
+        k1 = np.array([[1.0, 0.0], [0.0, np.sqrt(1 - gamma)]], dtype=np.complex128)
+        k2 = np.array([[0.0, np.sqrt(gamma)], [0.0, 0.0]], dtype=np.complex128)
+
+        def single(rho: npt.NDArray[np.complex128]) -> npt.NDArray[np.complex128]:
+            return k1 @ rho @ k1.conj().T + k2 @ rho @ k2.conj().T
+
+        expected = np.kron(single(rho_a), single(rho_b))
+
+        assert np.allclose(expected.trace(), 1.0)
+        assert np.allclose(dm.rho, expected)
+
     def test_apply_random_channel_one_qubit(self, fx_rng: Generator) -> None:
         """Test using complex parameters."""
         # check against statevector backend by hand for now.