From 78855836e704b2a18d03262bcb3fc673da4cb47c Mon Sep 17 00:00:00 2001 From: Vinny010 Date: Thu, 4 Jun 2026 02:03:51 +0400 Subject: [PATCH] Implement amplitude damping noise model (closes #497) Add amplitude damping noise to the density-matrix simulation stack: - `amplitude_damping_channel` and `two_qubit_amplitude_damping_channel` in `graphix.channels`, returning the corresponding `KrausChannel`. The two-qubit channel is the tensor product of two independent single-qubit channels. - `AmplitudeDampingNoise` and `TwoQubitAmplitudeDampingNoise` (deriving from `Noise`), plus `AmplitudeDampingNoiseModel` (deriving from `NoiseModel`), mirroring the structure of `DepolarisingNoiseModel`. Channel parameters use a `_gamma` suffix to reflect the damping rate. - Tests covering the analytic channel action on basis and superposition states, two-qubit tensor structure, per-step noise insertion in a pattern, zero-damping/noiseless equivalence, and an end-to-end preparation-error simulation. Passes ruff, mypy --strict, pyright, and pytest locally. Co-Authored-By: Claude Opus 4.8 --- graphix/channels.py | 50 ++++++ graphix/noise_models/__init__.py | 8 + graphix/noise_models/amplitude_damping.py | 168 +++++++++++++++++++ tests/test_amplitude_damping.py | 187 ++++++++++++++++++++++ 4 files changed, 413 insertions(+) create mode 100644 graphix/noise_models/amplitude_damping.py create mode 100644 tests/test_amplitude_damping.py diff --git a/graphix/channels.py b/graphix/channels.py index 8bb0ebc1b..6fb5bb773 100644 --- a/graphix/channels.py +++ b/graphix/channels.py @@ -296,3 +296,53 @@ def two_qubit_depolarising_tensor_channel(prob: float) -> KrausChannel: KrausData(prob / 3.0, np.kron(Ops.Z, Ops.Y)), ] ) + + +def amplitude_damping_channel(gamma: float) -> KrausChannel: + r"""Single-qubit amplitude damping channel. + + The channel acts on a density matrix as + :math:`\Phi(\rho) = K_1 \rho K_1^\dagger + K_2 \rho K_2^\dagger` with + + .. 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}. + + It models the decay of the excited state :math:`\lvert 1 \rangle` towards the + ground state :math:`\lvert 0 \rangle` with probability ``gamma``. + + Parameters + ---------- + gamma : float + Damping parameter, between 0 and 1. + + Returns + ------- + :class:`graphix.channels.KrausChannel` + Channel containing the corresponding Kraus operators. + """ + k1 = np.array([[1.0, 0.0], [0.0, np.sqrt(1.0 - gamma)]], dtype=np.complex128) + k2 = np.array([[0.0, np.sqrt(gamma)], [0.0, 0.0]], dtype=np.complex128) + return KrausChannel([KrausData(1.0, k1), KrausData(1.0, k2)]) + + +def two_qubit_amplitude_damping_channel(gamma: float) -> KrausChannel: + r"""Two-qubit amplitude damping channel. + + Tensor product of two independent single-qubit amplitude damping channels + sharing the same parameter ``gamma``. Its Kraus operators are the four + products :math:`K_i \otimes K_j` of the single-qubit Kraus operators returned + by :func:`amplitude_damping_channel`. + + Parameters + ---------- + gamma : float + Damping parameter, between 0 and 1. + + Returns + ------- + :class:`graphix.channels.KrausChannel` + Channel containing the corresponding Kraus operators. + """ + single = [data.coef * data.operator for data in amplitude_damping_channel(gamma)] + return KrausChannel([KrausData(1.0, np.kron(a, b)) for a in single for b in single]) diff --git a/graphix/noise_models/__init__.py b/graphix/noise_models/__init__.py index 1d74beaac..fb7a375e3 100644 --- a/graphix/noise_models/__init__.py +++ b/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", ] diff --git a/graphix/noise_models/amplitude_damping.py b/graphix/noise_models/amplitude_damping.py new file mode 100644 index 000000000..8a60acb06 --- /dev/null +++ b/graphix/noise_models/amplitude_damping.py @@ -0,0 +1,168 @@ +"""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 ``gamma``.""" + + gamma = Probability() + + def __init__(self, gamma: float) -> None: + """Initialize one-qubit amplitude damping noise. + + Parameters + ---------- + gamma : float + Damping parameter of the noise, between 0 and 1. + """ + self.gamma = gamma + + @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.gamma) + + +class TwoQubitAmplitudeDampingNoise(Noise): + """Two-qubit amplitude damping noise with damping parameter ``gamma``.""" + + gamma = Probability() + + def __init__(self, gamma: float) -> None: + """Initialize two-qubit amplitude damping noise. + + Parameters + ---------- + gamma : float + Damping parameter of the noise, between 0 and 1. + """ + self.gamma = gamma + + @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.gamma) + + +class AmplitudeDampingNoiseModel(NoiseModel): + """Amplitude damping noise model. + + Mirrors the structure of + :class:`graphix.noise_models.depolarising.DepolarisingNoiseModel`, applying an + amplitude damping channel at each step of a pattern. Channel parameters are + named with a ``_gamma`` suffix to reflect that they are damping parameters. + + Parameters + ---------- + prepare_error_gamma : float + Damping applied to each freshly prepared node. + x_error_gamma : float + Damping applied after an ``X`` correction (conditioned on its domain). + z_error_gamma : float + Damping applied after a ``Z`` correction (conditioned on its domain). + entanglement_error_gamma : float + Two-qubit damping applied after an entangling ``E`` command. + measure_channel_gamma : float + Damping applied to a node immediately before it is measured. + measure_error_prob : float + Probability of reporting a flipped (classical) measurement outcome. + """ + + def __init__( + self, + prepare_error_gamma: float = 0.0, + x_error_gamma: float = 0.0, + z_error_gamma: float = 0.0, + entanglement_error_gamma: float = 0.0, + measure_channel_gamma: float = 0.0, + measure_error_prob: float = 0.0, + ) -> None: + self.prepare_error_gamma = prepare_error_gamma + self.x_error_gamma = x_error_gamma + self.z_error_gamma = z_error_gamma + self.entanglement_error_gamma = entanglement_error_gamma + self.measure_channel_gamma = measure_channel_gamma + 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_gamma), 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_gamma), nodes=[cmd.node])] + case CommandKind.E: + return [ + cmd, + ApplyNoise( + noise=TwoQubitAmplitudeDampingNoise(self.entanglement_error_gamma), nodes=list(cmd.nodes) + ), + ] + case CommandKind.M: + return [ApplyNoise(noise=AmplitudeDampingNoise(self.measure_channel_gamma), nodes=[cmd.node]), cmd] + case CommandKind.X: + return [ + cmd, + ApplyNoise(noise=AmplitudeDampingNoise(self.x_error_gamma), nodes=[cmd.node], domain=cmd.domain), + ] + case CommandKind.Z: + return [ + cmd, + ApplyNoise(noise=AmplitudeDampingNoise(self.z_error_gamma), 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 a possibly flipped measurement result for cmd = "M".""" + rng = ensure_rng(rng, stacklevel=stacklevel + 1) + if rng.uniform() < self.measure_error_prob: + return toggle_outcome(result) + return result diff --git a/tests/test_amplitude_damping.py b/tests/test_amplitude_damping.py new file mode 100644 index 000000000..f0a21caed --- /dev/null +++ b/tests/test_amplitude_damping.py @@ -0,0 +1,187 @@ +r"""Tests for the amplitude damping channel and noise model. + +Soundness of these tests +------------------------ +The single-qubit amplitude damping channel has the analytic action +:math:`\\Phi_\\gamma(\\rho) = K_1 \\rho K_1^\\dagger + K_2 \\rho K_2^\\dagger`, which on +the computational basis gives :math:`\\Phi_\\gamma(|0\\rangle\\langle 0|) = |0\\rangle\\langle 0|` +(the ground state is unaffected) and +:math:`\\Phi_\\gamma(|1\\rangle\\langle 1|) = \\gamma |0\\rangle\\langle 0| + (1-\\gamma)|1\\rangle\\langle 1|` +(the excited state decays with probability ``gamma``). The two-qubit channel is the +tensor product of two independent single-qubit channels, so its action on a product +state is the Kronecker product of the single-qubit results. These closed forms are used +as references below. The noise-model tests check that the expected ``ApplyNoise`` commands +carrying amplitude damping are inserted at each step of a pattern, and that a zero-damping +model reproduces the noiseless result. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import pytest + +from graphix import command +from graphix.channels import amplitude_damping_channel, two_qubit_amplitude_damping_channel +from graphix.command import CommandKind +from graphix.noise_models import ( + AmplitudeDampingNoise, + AmplitudeDampingNoiseModel, + ApplyNoise, + TwoQubitAmplitudeDampingNoise, +) +from graphix.pattern import Pattern +from graphix.sim.density_matrix import DensityMatrix +from graphix.states import BasicStates +from graphix.transpiler import Circuit + +if TYPE_CHECKING: + from numpy.random import Generator + +GAMMAS = [0.0, 0.2, 0.5, 0.9, 1.0] + + +def _as_apply_noise(cmd: object) -> ApplyNoise: + assert isinstance(cmd, ApplyNoise) + return cmd + + +@pytest.mark.parametrize("gamma", GAMMAS) +def test_amplitude_damping_channel_basis_states(gamma: float) -> None: + # Ground state is unaffected. + ground = DensityMatrix(data=[BasicStates.ZERO]) + ground.apply_channel(amplitude_damping_channel(gamma), [0]) + assert np.allclose(ground.rho, np.array([[1.0, 0.0], [0.0, 0.0]])) + # Excited state decays to the ground state with probability gamma. + excited = DensityMatrix(data=[BasicStates.ONE]) + excited.apply_channel(amplitude_damping_channel(gamma), [0]) + assert np.allclose(excited.rho, np.array([[gamma, 0.0], [0.0, 1.0 - gamma]])) + + +@pytest.mark.parametrize("gamma", GAMMAS) +def test_amplitude_damping_channel_superposition(gamma: float) -> None: + # Analytic action on |+><+|. + plus = DensityMatrix(data=[BasicStates.PLUS]) + plus.apply_channel(amplitude_damping_channel(gamma), [0]) + root = float(np.sqrt(1.0 - gamma)) + expected = 0.5 * np.array([[1.0 + gamma, root], [root, 1.0 - gamma]]) + assert np.allclose(plus.rho, expected) + # Trace is preserved (CPTP). + assert np.isclose(plus.rho.trace(), 1.0) + + +@pytest.mark.parametrize("gamma", GAMMAS) +def test_two_qubit_amplitude_damping_channel(gamma: float) -> None: + # On a product state, the two-qubit channel is the Kronecker product of the + # single-qubit actions. + dm = DensityMatrix(data=[BasicStates.ONE, BasicStates.ONE]) + dm.apply_channel(two_qubit_amplitude_damping_channel(gamma), [0, 1]) + single = np.array([[gamma, 0.0], [0.0, 1.0 - gamma]]) + assert np.allclose(dm.rho, np.kron(single, single)) + + +def test_channel_nqubits() -> None: + assert amplitude_damping_channel(0.3).nqubit == 1 + assert two_qubit_amplitude_damping_channel(0.3).nqubit == 2 + assert len(amplitude_damping_channel(0.3)) == 2 + assert len(two_qubit_amplitude_damping_channel(0.3)) == 4 + + +def test_noise_elements() -> None: + one = AmplitudeDampingNoise(0.4) + assert one.nqubits == 1 + assert one.to_kraus_channel().nqubit == 1 + two = TwoQubitAmplitudeDampingNoise(0.4) + assert two.nqubits == 2 + assert two.to_kraus_channel().nqubit == 2 + + +@pytest.mark.parametrize("bad", [-0.1, 1.1, 2.0]) +def test_noise_gamma_validation(bad: float) -> None: + with pytest.raises(ValueError): + AmplitudeDampingNoise(bad) + + +def test_noise_model_inserts_noise_at_each_step() -> None: + model = AmplitudeDampingNoiseModel( + prepare_error_gamma=0.1, + x_error_gamma=0.2, + z_error_gamma=0.3, + entanglement_error_gamma=0.4, + measure_channel_gamma=0.5, + ) + + # N: prepare error appended after the command. + out = model.command(command.N(0)) + assert [c.kind for c in out] == [CommandKind.N, CommandKind.ApplyNoise] + noise = _as_apply_noise(out[1]) + assert isinstance(noise.noise, AmplitudeDampingNoise) + assert noise.noise.gamma == 0.1 + assert noise.nodes == [0] + + # E: two-qubit entanglement error appended after the command. + out = model.command(command.E((0, 1))) + assert [c.kind for c in out] == [CommandKind.E, CommandKind.ApplyNoise] + noise = _as_apply_noise(out[1]) + assert isinstance(noise.noise, TwoQubitAmplitudeDampingNoise) + assert noise.noise.gamma == 0.4 + assert noise.nodes == [0, 1] + + # M: measurement channel applied *before* the measurement. + out = model.command(command.M(0)) + assert [c.kind for c in out] == [CommandKind.ApplyNoise, CommandKind.M] + noise = _as_apply_noise(out[0]) + assert isinstance(noise.noise, AmplitudeDampingNoise) + assert noise.noise.gamma == 0.5 + + # X / Z: correction errors carry the command's domain. + out = model.command(command.X(0, domain={1})) + assert [c.kind for c in out] == [CommandKind.X, CommandKind.ApplyNoise] + noise = _as_apply_noise(out[1]) + assert isinstance(noise.noise, AmplitudeDampingNoise) + assert noise.noise.gamma == 0.2 + assert noise.domain == {1} + + out = model.command(command.Z(0, domain={1})) + assert [c.kind for c in out] == [CommandKind.Z, CommandKind.ApplyNoise] + noise = _as_apply_noise(out[1]) + assert isinstance(noise.noise, AmplitudeDampingNoise) + assert noise.noise.gamma == 0.3 + assert noise.domain == {1} + + +def test_noise_model_input_nodes() -> None: + model = AmplitudeDampingNoiseModel(prepare_error_gamma=0.25) + out = model.input_nodes([0, 2]) + noises = [_as_apply_noise(c) for c in out] + assert all(isinstance(n.noise, AmplitudeDampingNoise) and n.noise.gamma == 0.25 for n in noises) + assert [n.nodes for n in noises] == [[0], [2]] + + +def test_zero_damping_is_noiseless(fx_rng: Generator) -> None: + # A model with gamma = 0 everywhere leaves the noiseless result unchanged. + pattern = _hadamard_pattern() + res = pattern.simulate_pattern(backend="densitymatrix", noise_model=AmplitudeDampingNoiseModel(), rng=fx_rng) + assert isinstance(res, DensityMatrix) + assert np.allclose(res.rho, np.array([[1.0, 0.0], [0.0, 0.0]])) + + +@pytest.mark.parametrize("gamma", [0.2, 0.6]) +def test_preparation_amplitude_damping(gamma: float, fx_rng: Generator) -> None: + # A single preparation followed by amplitude damping yields the analytic + # action of the channel on the prepared |+> state. + pattern = Pattern(cmds=[command.N(0)]) + res = pattern.simulate_pattern( + backend="densitymatrix", noise_model=AmplitudeDampingNoiseModel(prepare_error_gamma=gamma), rng=fx_rng + ) + assert isinstance(res, DensityMatrix) + root = float(np.sqrt(1.0 - gamma)) + expected = 0.5 * np.array([[1.0 + gamma, root], [root, 1.0 - gamma]]) + assert np.allclose(res.rho, expected) + + +def _hadamard_pattern() -> Pattern: + circ = Circuit(1) + circ.h(0) + return circ.transpile().pattern