feat: add amplitude damping noise model

graphix/channels.py

@@ -200,6 +200,41 @@ def depolarising_channel(prob: float) -> KrausChannel:
     )
 
 
+def amplitude_damping_channel(prob: float) -> KrausChannel:
+    r"""Single-qubit amplitude damping channel.
+
+    Kraus operators:
+
+    .. math::
+        K_0 =
+        \begin{pmatrix}
+            1 & 0 \\
+            0 & \sqrt{1-\gamma}
+        \end{pmatrix},
+        K_1 =
+        \begin{pmatrix}
+            0 & \sqrt{\gamma} \\
+            0 & 0
+        \end{pmatrix}
+
+    Parameters
+    ----------
+    prob : float
+        The damping probability :math:`\gamma` associated to the channel.
+
+    Returns
+    -------
+    :class:`graphix.channels.KrausChannel` object
+        Channel containing the corresponding Kraus operators.
+    """
+    return KrausChannel(
+        [
+            KrausData(1.0, np.array([[1.0, 0.0], [0.0, np.sqrt(1 - prob)]])),
+            KrausData(1.0, np.array([[0.0, np.sqrt(prob)], [0.0, 0.0]])),
+        ]
+    )
+
+
 def pauli_channel(px: float, py: float, pz: float) -> KrausChannel:
     r"""Single-qubit Pauli channel.
 
@@ -220,6 +255,29 @@ def pauli_channel(px: float, py: float, pz: float) -> KrausChannel:
     )
 
 
+def two_qubit_amplitude_damping_channel(prob: float) -> KrausChannel:
+    r"""Two-qubit tensor product of single-qubit amplitude damping channels.
+
+    Parameters
+    ----------
+    prob : float
+        The damping probability :math:`\gamma` associated to each qubit channel.
+
+    Returns
+    -------
+    :class:`graphix.channels.KrausChannel` object
+        Channel containing the tensor-product Kraus operators.
+    """
+    single_qubit_channel = amplitude_damping_channel(prob)
+    return KrausChannel(
+        [
+            KrausData(ki.coef * kj.coef, np.kron(ki.operator, kj.operator))
+            for ki in single_qubit_channel
+            for kj in single_qubit_channel
+        ]
+    )
+
+
 def two_qubit_depolarising_channel(prob: float) -> KrausChannel:
     r"""Two-qubit depolarising channel.
 

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,147 @@
+"""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 probability ``prob``."""
+
+    prob = Probability()
+
+    def __init__(self, prob: float) -> None:
+        """Initialize one-qubit amplitude damping noise.
+
+        Parameters
+        ----------
+        prob : float
+            Probability parameter 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 probability ``prob``."""
+
+    prob = Probability()
+
+    def __init__(self, prob: float) -> None:
+        """Initialize two-qubit amplitude damping noise.
+
+        Parameters
+        ----------
+        prob : float
+            Probability parameter 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):
+    """Amplitude damping noise model."""
+
+    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 _:
+                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 cmd = "M"."""
+        rng = ensure_rng(rng, stacklevel=stacklevel + 1)
+        if rng.uniform() < self.measure_error_prob:
+            return toggle_outcome(result)
+        return result

tests/test_kraus.py

@@ -9,8 +9,10 @@
 from graphix.channels import (
     KrausChannel,
     KrausData,
+    amplitude_damping_channel,
     dephasing_channel,
     depolarising_channel,
+    two_qubit_amplitude_damping_channel,
     two_qubit_depolarising_channel,
     two_qubit_depolarising_tensor_channel,
 )
@@ -119,6 +121,45 @@ def test_depolarising_channel(self, fx_rng: Generator) -> None:
             assert np.allclose(depol_channel[i].coef, data[i].coef)
             assert np.allclose(depol_channel[i].operator, data[i].operator)
 
+    def test_amplitude_damping_channel(self, fx_rng: Generator) -> None:
+        prob = fx_rng.uniform()
+        data = [
+            KrausData(1.0, np.array([[1.0, 0.0], [0.0, np.sqrt(1 - prob)]])),
+            KrausData(1.0, np.array([[0.0, np.sqrt(prob)], [0.0, 0.0]])),
+        ]
+
+        damping_channel = amplitude_damping_channel(prob)
+
+        assert isinstance(damping_channel, KrausChannel)
+        assert damping_channel.nqubit == 1
+        assert len(damping_channel) == 2
+
+        for i in range(len(damping_channel)):
+            assert np.allclose(damping_channel[i].coef, data[i].coef)
+            assert np.allclose(damping_channel[i].operator, data[i].operator)
+
+    def test_2_qubit_amplitude_damping_channel(self, fx_rng: Generator) -> None:
+        prob = fx_rng.uniform()
+        single_qubit_data = [
+            KrausData(1.0, np.array([[1.0, 0.0], [0.0, np.sqrt(1 - prob)]])),
+            KrausData(1.0, np.array([[0.0, np.sqrt(prob)], [0.0, 0.0]])),
+        ]
+        data = [
+            KrausData(ki.coef * kj.coef, np.kron(ki.operator, kj.operator))
+            for ki in single_qubit_data
+            for kj in single_qubit_data
+        ]
+
+        damping_channel_2_qubit = two_qubit_amplitude_damping_channel(prob)
+
+        assert isinstance(damping_channel_2_qubit, KrausChannel)
+        assert damping_channel_2_qubit.nqubit == 2
+        assert len(damping_channel_2_qubit) == 4
+
+        for i in range(len(damping_channel_2_qubit)):
+            assert np.allclose(damping_channel_2_qubit[i].coef, data[i].coef)
+            assert np.allclose(damping_channel_2_qubit[i].operator, data[i].operator)
+
     def test_2_qubit_depolarising_channel(self, fx_rng: Generator) -> None:
         prob = fx_rng.uniform()
         data = [

tests/test_noise_model.py

@@ -8,10 +8,13 @@
 from graphix import Pattern
 from graphix.command import CommandKind, M, N
 from graphix.noise_models import (
+    AmplitudeDampingNoise,
+    AmplitudeDampingNoiseModel,
     ApplyNoise,
     ComposeNoiseModel,
     DepolarisingNoise,
     DepolarisingNoiseModel,
+    TwoQubitAmplitudeDampingNoise,
     TwoQubitDepolarisingNoise,
 )
 from graphix.noise_models.noise_model import NoiselessNoiseModel
@@ -86,6 +89,44 @@ def check_noise_command(cmd: CommandOrNoise, prob: float, two_qubits: bool) -> N
                 check_noise_command(next(iterator), 0, False)
 
 
+def test_amplitude_damping_noise_model_transpile(fx_rng: Generator) -> None:
+    nqubits = 5
+    depth = 5
+    circuit = rand_circuit(nqubits, depth, rng=fx_rng)
+    pattern = circuit.transpile().pattern
+    noise_model = AmplitudeDampingNoiseModel(
+        prepare_error_prob=0.1,
+        x_error_prob=0.2,
+        z_error_prob=0.3,
+        entanglement_error_prob=0.4,
+        measure_channel_prob=0.5,
+    )
+    noisy_pattern = noise_model.transpile(pattern, rng=fx_rng)
+    iterator = iter(noisy_pattern)
+
+    def check_noise_command(cmd: CommandOrNoise, prob: float, two_qubits: bool) -> None:
+        assert isinstance(cmd, ApplyNoise)
+        if two_qubits:
+            assert isinstance(cmd.noise, TwoQubitAmplitudeDampingNoise)
+        else:
+            assert isinstance(cmd.noise, AmplitudeDampingNoise)
+        assert cmd.noise.prob == prob
+
+    for cmd in pattern:
+        if cmd.kind == CommandKind.M:
+            check_noise_command(next(iterator), 0.5, False)
+        assert next(iterator) == cmd
+        match cmd.kind:
+            case CommandKind.N:
+                check_noise_command(next(iterator), 0.1, False)
+            case CommandKind.E:
+                check_noise_command(next(iterator), 0.4, True)
+            case CommandKind.X:
+                check_noise_command(next(iterator), 0.2, False)
+            case CommandKind.Z:
+                check_noise_command(next(iterator), 0.3, False)
+
+
 def test_compose_noise_model_simulation(fx_rng: Generator) -> None:
     nqubits = 5
     depth = 5