dense-evolution 8.3.0__py3-none-win_amd64.whl

dense_evolution/noise/coherent_attack.py

@@ -0,0 +1,264 @@
+"""Coherent, differentiable adversarial noise for stabilizer codes.
+
+Promoted from Dense-Evolution-Discovery's Steane [[7,1,3]] investigation
+(scripts/steane_code_block2_adversarial_noise.py and
+steane_code_block3_linf_constrained_attack.py), generalized here from a
+hard-coded 7-qubit Steane table to any stabilizer list, matching this
+module's own `dense_evolution.qec.compute_syndrome` conventions (equal-
+length Pauli strings over IXYZ, qubit q = character q).
+
+Unlike `dense_evolution.registry.NoiseModel` (stochastic Kraus channels,
+sampled via a hard threshold that is not usefully differentiable), this
+module attacks a genuinely continuous channel: independent per-qubit
+coherent over-rotations rz(delta_q) applied to all qubits at once --
+the natural differentiable analogue of physical miscalibration. Because
+it is JAX-differentiable end-to-end, `craft_adversarial_delta` can
+gradient-search for a worst-case coherent-error direction instead of
+only sampling random ones.
+
+Honest negative result, kept because it is the actual finding, not
+smoothed over: unconstrained (`craft_adversarial_delta`, L2-ball only)
+PGD on `x_stabilizer_leakage` degenerately concentrates the whole error
+budget on whichever single qubit is shared by every X-stabilizer
+generator. A coherent rz error on one qubit alone always collapses to
+weight-0 or weight-1, which a distance-3 code's decoder corrects
+exactly every time -- so the "adversarial" direction found this way has
+ZERO real decoder-failure rate, while random multi-qubit directions of
+the same L2 budget fail readily. `craft_adversarial_delta_constrained`
+fixes this by also capping the per-qubit angle (L-infinity), forcing
+the search to spread the budget across multiple qubits -- the regime
+that actually causes failures. Use the constrained version if the goal
+is a genuine worst-case test; the unconstrained version is kept because
+it is a real, useful example of a differentiable proxy objective being
+anti-correlated with the thing it was meant to approximate.
+"""
+import numpy as np
+import jax
+import jax.numpy as jnp
+
+from dense_evolution.physics.qec import compute_syndrome
+
+__all__ = [
+    "apply_rz_all",
+    "x_stabilizer_leakage",
+    "craft_adversarial_delta",
+    "project_l2_linf",
+    "craft_adversarial_delta_constrained",
+    "decoder_failure_rate",
+    "random_delta_failure_stats",
+]
+
+
+def _flip_mask(pauli_str: str) -> int:
+    """Bitmask of qubits carrying an X or Y term, qubit q = bit (n-1-q) --
+    dense_evolution's own MSB-first statevector-index convention."""
+    n = len(pauli_str)
+    mask = 0
+    for q, p in enumerate(pauli_str):
+        if p in ('X', 'Y'):
+            mask |= 1 << (n - 1 - q)
+    return mask
+
+
+def apply_rz_all(sv0: jnp.ndarray, delta: jnp.ndarray) -> jnp.ndarray:
+    """Coherent per-qubit rz(delta_q) applied to every qubit of `sv0` at
+    once. rz gates are diagonal and all commute, so this is exact
+    elementwise phase multiplication, not a per-gate circuit simulation
+    -- and fully JAX-differentiable in `delta`.
+
+    Parameters
+    ----------
+    sv0 : statevector, length 2**n_qubits
+    delta : real array, length n_qubits -- rz angle for each qubit
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> import jax.numpy as jnp
+    >>> sv0 = jnp.array([1.0, 0.0], dtype=jnp.complex128)
+    >>> sv1 = apply_rz_all(sv0, jnp.array([np.pi]))
+    >>> round(float(jnp.abs(sv1[0]) ** 2), 6)
+    1.0
+    """
+    n = delta.shape[0]
+    dim = sv0.shape[0]
+    idx = jnp.arange(dim, dtype=jnp.int32)
+    bit_pos = jnp.arange(n - 1, -1, -1, dtype=jnp.int32)
+    bits = (idx[:, None] >> bit_pos[None, :]) & 1
+    s = (1 - 2 * bits).astype(jnp.float64)
+    phase_arg = -0.5 * jnp.sum(delta[None, :] * s, axis=1)
+    return sv0 * jnp.exp(1j * phase_arg)
+
+
+def x_stabilizer_leakage(delta: jnp.ndarray, sv0: jnp.ndarray, stabilizers) -> jnp.ndarray:
+    """Total leakage of `sv0`, after a coherent `apply_rz_all(sv0, delta)`
+    perturbation, out of the +1 joint eigenspace of `stabilizers` --
+    sum_i (1 - <stabilizer_i>) / 2, one term per generator. A smooth,
+    bounded ([0, len(stabilizers)]) proxy for how much the coherent
+    error disturbs the syndrome; NOT the actual decoder failure
+    probability (see module docstring for the documented gap between
+    the two). `stabilizers` should be X-type generators to pair with an
+    rz (Z-type-diagonal) coherent error -- the same reasoning
+    `dense_evolution.qec.compute_syndrome` uses generically for any
+    Pauli-string stabilizer list.
+    """
+    sv = apply_rz_all(sv0, delta)
+    dim = sv0.shape[0]
+    idx = jnp.arange(dim, dtype=jnp.int32)
+    total = 0.0
+    for g in stabilizers:
+        mask = _flip_mask(g)
+        src = idx ^ mask
+        expectation = jnp.real(jnp.sum(jnp.conj(sv) * sv[src]))
+        total = total + (1.0 - expectation) / 2.0
+    return total
+
+
+def craft_adversarial_delta(sv0: jnp.ndarray, stabilizers, epsilon: float,
+                             n_steps: int = 150, step_size: float = 0.05, seed: int = 0):
+    """Gradient-ascent PGD on `x_stabilizer_leakage`, projected into the L2
+    epsilon-ball around delta=0 after every step. Returns (best_delta as
+    numpy, best_leakage, leakage_history).
+
+    See the module docstring: on a distance-3 code this unconstrained
+    search finds a direction with real decoder-failure rate 0 -- use
+    `craft_adversarial_delta_constrained` for a genuine worst-case test.
+    """
+    n_qubits = len(stabilizers[0])
+    leakage_grad = jax.grad(lambda d, sv: x_stabilizer_leakage(d, sv, stabilizers), argnums=0)
+
+    rng = np.random.default_rng(seed)
+    init_dir = rng.normal(size=n_qubits)
+    init_dir /= np.linalg.norm(init_dir)
+    init_norm = min(epsilon, 1e-2)
+    delta = jnp.array(init_dir * init_norm)
+
+    best_delta = delta
+    best_leakage = float(x_stabilizer_leakage(delta, sv0, stabilizers))
+    history = [best_leakage]
+
+    for _ in range(n_steps):
+        grad = leakage_grad(delta, sv0)
+        grad_norm = jnp.linalg.norm(grad)
+        step = jnp.where(grad_norm > 1e-12, grad / grad_norm, jnp.zeros_like(grad))
+        delta = delta + step_size * step
+        delta_norm = jnp.linalg.norm(delta)
+        delta = jnp.where(delta_norm > epsilon, delta / delta_norm * epsilon, delta)
+
+        current = float(x_stabilizer_leakage(delta, sv0, stabilizers))
+        history.append(current)
+        if current > best_leakage:
+            best_leakage = current
+            best_delta = delta
+
+    return np.asarray(best_delta), best_leakage, history
+
+
+def project_l2_linf(y: np.ndarray, epsilon: float, linf_cap: float, n_bisect: int = 60) -> np.ndarray:
+    """Exact projection of `y` onto the intersection of an L2 ball of
+    radius `epsilon` and an L-infinity ball (box) of radius `linf_cap` --
+    not the same as clip-then-rescale, which can push coordinates back
+    outside the box. Box-clips first; if that's already inside the L2
+    ball it's the exact answer, otherwise bisects the Lagrange
+    multiplier on the L2 constraint until the clipped, rescaled point
+    lands exactly on the L2 boundary.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> project_l2_linf(np.array([10.0, 0.0]), epsilon=1.0, linf_cap=5.0).round(4)
+    array([1., 0.])
+    >>> project_l2_linf(np.array([0.1, 0.1]), epsilon=1.0, linf_cap=0.05).round(4)
+    array([0.05, 0.05])
+    """
+    y = np.asarray(y, dtype=np.float64)
+    box = np.clip(y, -linf_cap, linf_cap)
+    if np.linalg.norm(box) <= epsilon + 1e-12:
+        return box
+    lo, hi = 0.0, 1e8
+    for _ in range(n_bisect):
+        mid = 0.5 * (lo + hi)
+        z = np.clip(y / (1.0 + mid), -linf_cap, linf_cap)
+        if np.linalg.norm(z) > epsilon:
+            lo = mid
+        else:
+            hi = mid
+    return np.clip(y / (1.0 + hi), -linf_cap, linf_cap)
+
+
+def craft_adversarial_delta_constrained(sv0: jnp.ndarray, stabilizers, epsilon: float, linf_cap: float,
+                                         n_steps: int = 150, step_size: float = 0.05, seed: int = 0):
+    """Same PGD search as `craft_adversarial_delta`, but each step is
+    projected into the L2-epsilon-ball INTERSECTED with an L-infinity
+    box of radius `linf_cap` (via `project_l2_linf`) instead of the L2
+    ball alone. Capping the per-qubit angle forbids the degenerate
+    one-qubit-takes-everything solution and forces the search to spread
+    the budget across multiple qubits -- the regime that actually causes
+    real decoder failures (see module docstring)."""
+    n_qubits = len(stabilizers[0])
+    leakage = lambda d, sv: x_stabilizer_leakage(d, sv, stabilizers)
+    leakage_grad = jax.grad(leakage, argnums=0)
+
+    rng = np.random.default_rng(seed)
+    init_dir = rng.normal(size=n_qubits)
+    init_dir /= np.linalg.norm(init_dir)
+    init_norm = min(epsilon, 1e-2)
+    delta_np = project_l2_linf(init_dir * init_norm, epsilon, linf_cap)
+    delta = jnp.array(delta_np)
+
+    best_delta = delta
+    best_leakage = float(leakage(delta, sv0))
+    history = [best_leakage]
+
+    for _ in range(n_steps):
+        grad = leakage_grad(delta, sv0)
+        grad_norm = jnp.linalg.norm(grad)
+        step = jnp.where(grad_norm > 1e-12, grad / grad_norm, jnp.zeros_like(grad))
+        delta_raw = np.asarray(delta) + step_size * np.asarray(step)
+        delta_np = project_l2_linf(delta_raw, epsilon, linf_cap)
+        delta = jnp.array(delta_np)
+
+        current = float(leakage(delta, sv0))
+        history.append(current)
+        if current > best_leakage:
+            best_leakage = current
+            best_delta = delta
+
+    return np.asarray(best_delta), best_leakage, history
+
+
+def decoder_failure_rate(delta_np: np.ndarray, sv0_np: np.ndarray, decode_fn,
+                          n_trials: int, rng: np.random.Generator) -> float:
+    """Real, discrete evaluation of how often a coherent error `delta_np`
+    actually fools a decoder -- as opposed to `x_stabilizer_leakage`'s
+    smooth proxy. `decode_fn(sv_noisy, rng) -> sv_corrected` is any
+    projective-measurement-based decoder (e.g. built around
+    `dense_evolution.qec.compute_syndrome`); this function applies the
+    coherent error once, then calls `decode_fn` `n_trials` times (the
+    syndrome measurement that collapses the coherently-perturbed state
+    is itself stochastic, so repeated trials are meaningful even though
+    `delta_np` is fixed), and reports the fraction that fail to recover
+    `sv0_np` exactly."""
+    sv_delta = np.asarray(apply_rz_all(jnp.array(sv0_np), jnp.array(delta_np)))
+    n_fail = 0
+    for _ in range(n_trials):
+        sv_corrected = decode_fn(sv_delta.copy(), rng)
+        fidelity = np.abs(np.vdot(sv_corrected, sv0_np)) ** 2
+        if fidelity < 1.0 - 1e-6:
+            n_fail += 1
+    return n_fail / n_trials
+
+
+def random_delta_failure_stats(sv0_np: np.ndarray, decode_fn, epsilon: float, n_qubits: int,
+                                n_random: int, n_trials_each: int, rng: np.random.Generator) -> np.ndarray:
+    """Same evaluation as `decoder_failure_rate`, but over `n_random`
+    random directions of L2 norm `epsilon` instead of one crafted
+    `delta` -- the baseline `craft_adversarial_delta`'s result should be
+    compared against (see module docstring: the unconstrained crafted
+    direction can score WORSE than this random baseline)."""
+    rates = np.zeros(n_random)
+    for i in range(n_random):
+        d = rng.normal(size=n_qubits)
+        d = d / np.linalg.norm(d) * epsilon
+        rates[i] = decoder_failure_rate(d, sv0_np, decode_fn, n_trials_each, rng)
+    return rates

dense_evolution/noise/cosmic_ray.py

@@ -0,0 +1,61 @@
+"""A cosmic-ray/gamma-ray-induced quasiparticle burst, as a real
+time-dependent noise-strength profile instead of a constant rate."""
+import jax.numpy as jnp
+
+__all__ = ["cosmic_ray_burst_profile"]
+
+
+def cosmic_ray_burst_profile(time_us, baseline_gamma: float, ratio_intermediate: float = 2.5,
+                              ratio_peak: float = 3.75, tau1_us: float = 3.0,
+                              tau2_us: float = 300.0, tau_decay_ms: float = 25.0) -> jnp.ndarray:
+    """Time-dependent decay-probability profile for a cosmic-ray/gamma-ray-
+    induced quasiparticle burst: a two-stage rise (fast to
+    `ratio_intermediate`x baseline, slower to `ratio_peak`x baseline) times
+    a single-exponential recovery, generalized out of a fixed, paper-number
+    validation (Dense-Evolution-Discovery Experiment 34, reproducing
+    arXiv:2104.05219's real measured event on a 26-qubit chip).
+
+    Feed the result to `continuous_dissipative_evolve` alongside
+    `amplitude_damping_channel` (or any other single-time-varying-parameter
+    channel) to inject a realistic burst into any circuit or QEC study,
+    without re-deriving this shape by hand each time.
+
+    The default ratios/timescales are the paper's own real numbers -- see
+    Experiment 34's docstring for exactly which are paper-fitted (the 25ms
+    decay) versus chosen to match the paper's two described rise points
+    (tau1/tau2). All are overridable for a different event severity or
+    device generation; `baseline_gamma` is never derived here -- pass
+    whatever per-slice decay probability corresponds to your own dt/T1
+    convention (see Experiment 34 for one worked example of that
+    conversion).
+
+    Parameters
+    ----------
+    time_us : array_like
+        Time since impact, in microseconds (t=0 is the impact instant).
+    baseline_gamma : float
+        Undisturbed per-slice decay probability; this profile scales it
+        up, it does not derive it.
+    ratio_intermediate, ratio_peak : float
+        Multiplier on `baseline_gamma` at the two described checkpoints
+        (paper defaults 2.5=10/4, 3.75=15/4, from Fig. 3's ~10us/~1ms
+        readings).
+    tau1_us, tau2_us : float
+        Rise timescales for the two saturating-exponential stages (paper
+        defaults 3, 300 -- chosen to match its ~10us/~1ms descriptions,
+        not fitted by the paper itself).
+    tau_decay_ms : float
+        Recovery time constant (paper default 25 -- its own fitted central
+        value, real range 25-30ms across 415 events).
+
+    Returns
+    -------
+    jnp.ndarray
+        Per-slice decay probability at each entry of `time_us`.
+    """
+    time_us = jnp.asarray(time_us)
+    stage1 = (ratio_intermediate - 1.0) * (1.0 - jnp.exp(-time_us / tau1_us))
+    stage2 = (ratio_peak - ratio_intermediate) * (1.0 - jnp.exp(-time_us / tau2_us))
+    decay = jnp.exp(-time_us / (tau_decay_ms * 1000.0))
+    scaling = 1.0 + (stage1 + stage2) * decay
+    return baseline_gamma * scaling

dense_evolution/noise/density_matrix_channels.py

@@ -0,0 +1,78 @@
+"""Noise applied directly to a density matrix, instead of a statevector --
+what density-matrix ZNE (`dense_evolution.mitigation.zne_density_matrix`)
+needs its noise ensemble built from."""
+import jax.numpy as jnp
+
+__all__ = ["global_depolarizing_channel", "amplitude_damping_channel", "phaseflip_channel_exact"]
+
+
+def global_depolarizing_channel(rho: jnp.ndarray, p: float) -> jnp.ndarray:
+    """Global n-qubit depolarizing channel, D_p(rho) = (1-p)*rho + (p/dim)*I.
+
+    Distinct from `NoiseModel`'s `'depolarizing'` model, which applies an
+    independent PER-QUBIT local Kraus channel -- a different physical map
+    from this GLOBAL channel, which mixes the whole `dim`-dimensional state
+    toward the fully mixed state as one unit. Use this one when modeling
+    e.g. state-prep/measurement (SPAM) error reported as a single joint
+    depolarizing parameter over the whole register, not per-qubit gate
+    noise (promoted from a real reproduction of arXiv:2608.16716's own
+    SPAM model, Dense-Evolution-Discovery Experiment 33).
+    """
+    rho = jnp.asarray(rho, dtype=jnp.complex128)
+    dim = rho.shape[0]
+    identity = jnp.eye(dim, dtype=jnp.complex128)
+    return (1.0 - p) * rho + (p / dim) * identity
+
+
+def amplitude_damping_channel(rho: jnp.ndarray, gamma: float) -> jnp.ndarray:
+    """Single-qubit amplitude-damping channel: E0 @ rho @ E0.conj().T +
+    E1 @ rho @ E1.conj().T, with E0=diag(1, sqrt(1-gamma)) and
+    E1=[[0,sqrt(gamma)],[0,0]] -- population only ever moves |1>->|0>,
+    never the reverse.
+
+    Distinct from `global_depolarizing_channel` (symmetric, mixes toward
+    the fully-mixed state regardless of which state is |1> or |0>) -- this
+    one is asymmetric by construction, the real signature of energy-relaxation
+    (T1) processes and of quasiparticle poisoning (promoted from a real
+    reproduction of arXiv:2104.05219's measured cosmic-ray-induced error
+    bursts, Dense-Evolution-Discovery Experiment 34, where this asymmetry is
+    exactly the mechanism's own reported signature: decay errors only, no
+    excess excitation errors).
+
+    Single-qubit only (rho must be 2x2) -- unlike `global_depolarizing_channel`,
+    this is not dimension-generic, since amplitude damping is inherently a
+    per-qubit process, not a joint-register one.
+    """
+    rho = jnp.asarray(rho, dtype=jnp.complex128)
+    e0 = jnp.array([[1.0, 0.0], [0.0, jnp.sqrt(1.0 - gamma)]], dtype=jnp.complex128)
+    e1 = jnp.array([[0.0, jnp.sqrt(gamma)], [0.0, 0.0]], dtype=jnp.complex128)
+    return e0 @ rho @ e0.conj().T + e1 @ rho @ e1.conj().T
+
+
+def phaseflip_channel_exact(rho: jnp.ndarray, p: float) -> jnp.ndarray:
+    """Exact (zero-variance) multi-qubit phaseflip channel: K0=sqrt(1-p)*I,
+    K1=sqrt(p)*Z, applied independently per qubit across the whole
+    register, matching `NoiseModel`'s `'phaseflip'` statevector model
+    (`noise/kraus/phaseflip.py`) in the limit of infinitely many Monte
+    Carlo trials -- this function computes that limit directly instead of
+    averaging finite trials, for use as a zero-sampling-variance
+    "classical node" in Classically Augmented ZNE (Scheiber et al., "CA-ZNE:
+    high-noise Richardson nodes replaced by classically simulated
+    estimates", arXiv:2607.25746).
+
+    Since Z is diagonal in the computational basis, conjugating rho by Z_q
+    (Pauli-Z on qubit q, identity elsewhere) only flips the sign of
+    off-diagonal entries whose row and column disagree on bit q --
+    (Z_q rho Z_q)_ij = rho_ij * s_i * s_j, s_k = -1 if bit q of k else +1
+    -- so each qubit's channel (1-p)*rho + p*(Z_q rho Z_q) is one
+    elementwise multiply, no explicit operator ever built.
+    """
+    rho = jnp.asarray(rho, dtype=jnp.complex128)
+    dim = rho.shape[0]
+    n_qubits = dim.bit_length() - 1
+    idx = jnp.arange(dim)
+    for q in range(n_qubits):
+        sign = 1.0 - 2.0 * ((idx >> q) & 1)
+        flip_sign = jnp.outer(sign, sign)
+        rho = (1.0 - p) * rho + p * (flip_sign * rho)
+    return rho

dense_evolution/noise/differentiable.py

@@ -0,0 +1,66 @@
+from typing import Optional, List
+
+import jax
+
+__all__ = ["NoiseSpec"]
+
+
+class NoiseSpec:
+    """Native JAX-differentiable representation of a noise configuration --
+    a real JAX PyTree, so noise parameters thread through jax.jit/jax.grad/
+    jax.vmap natively -- e.g. as the `noise=` argument to
+    circuit_to_energy_fn's energy_fn -- instead of being applied as an
+    external, Python-side step around the already-traced circuit (the old
+    way: build sv, exit the trace, call apply_to_sv separately).
+
+    `model`/`qubits` are static (aux_data): they select which code path
+    runs, not values to differentiate or batch over -- the same role
+    `static_argnames` plays for a plain jax.jit function, but automatic
+    here because it's part of the pytree structure. `p`/`jax_key` are
+    pytree leaves (children): `p` can be a traced/differentiable value
+    (e.g. optimizing noise strength itself), and `jax_key` flows through
+    jit/vmap/scan the way any other JAX array does -- no external
+    Python-level key management, no OS-entropy fallback (unlike
+    apply_to_sv called standalone with jax_key=None), so a NoiseSpec's
+    result is always reproducible from the key it was built with.
+
+    `jax_key` is required (not Optional) -- the whole point of wiring
+    noise into the traced computation this way is to remove the need for
+    an external, ad-hoc key-management workaround; a caller who wants a
+    fresh key per call should split one themselves (`jax.random.split`)
+    and build a fresh NoiseSpec, the same as any other JAX-idiomatic
+    stateless-key pattern.
+
+    Examples
+    --------
+    >>> import jax
+    >>> from dense_evolution.noise import NoiseSpec
+    >>> key = jax.random.PRNGKey(0)
+    >>> spec = NoiseSpec(model="depolarizing", p=0.05, jax_key=key, qubits=[0, 1])
+    >>> spec
+    NoiseSpec(model='depolarizing', p=0.05, qubits=(0, 1))
+    """
+    def __init__(self, model: str, p, jax_key, qubits: Optional[List[int]] = None):
+        self.model = model
+        self.p = p
+        self.jax_key = jax_key
+        self.qubits = tuple(qubits) if qubits is not None else None
+
+    def __repr__(self):
+        return f"NoiseSpec(model={self.model!r}, p={self.p!r}, qubits={self.qubits!r})"
+
+    def tree_flatten(self):
+        children = (self.p, self.jax_key)
+        aux_data = (self.model, self.qubits)
+        return children, aux_data
+
+    @classmethod
+    def tree_unflatten(cls, aux_data, children):
+        model, qubits = aux_data
+        p, jax_key = children
+        return cls(model, p, jax_key, qubits)
+
+
+jax.tree_util.register_pytree_node(
+    NoiseSpec, NoiseSpec.tree_flatten, NoiseSpec.tree_unflatten
+)

dense_evolution/noise/kraus/__init__.py

@@ -0,0 +1,6 @@
+"""One file per Kraus channel: `ideal`, `depolarizing`, `bitflip`,
+`phaseflip`, `amplitude_damping`, `combined`. Each module exposes a single
+`apply(sv_out, idx_0, idx_1, p, rng, key, is_jax) -> (sv_out, key)`
+function; `dense_evolution.noise.kraus_channels.NoiseModel` is the shared
+dispatcher (RNG/key setup, per-qubit loop, normalisation) that calls them.
+"""

dense_evolution/noise/kraus/amplitude_damping.py

@@ -0,0 +1,47 @@
+"""Amplitude-damping channel: K0=diag(1, sqrt(1-gamma)), K1=[[0,
+sqrt(gamma)], [0, 0]]. T1 energy relaxation |1> -> |0> with rate gamma.
+
+Single-trajectory ("quantum jump") unraveling: ONE decay/no-decay
+decision per qubit per shot, using the Born-rule probability aggregated
+over the WHOLE statevector, P(K1) = <psi|K1^dagger K1|psi> =
+gamma * sum_i |v1[i]|^2 (summed over every branch of the other n-1
+qubits, not per branch). If it fires, K1 collapses the entire qubit-q=1
+branch onto q=0, preserving the other qubits' relative amplitudes
+(v1[i]/norm, not flattened to unit phase per branch). Renormalizing by
+the GLOBAL P(K1) (or 1-P(K1) for the no-decay branch), not a per-branch
+norm, is what a correct sequential multi-qubit trajectory needs.
+"""
+import numpy as np
+import jax
+import jax.numpy as jnp
+
+__all__ = ["apply"]
+
+
+def apply(sv_out, idx_0, idx_1, p, rng, key, is_jax):
+    gamma = float(np.clip(p, 0.0, 1.0))
+    if is_jax:
+        key, subkey = jax.random.split(key)
+        r = jax.random.uniform(subkey, shape=(), minval=0.0, maxval=1.0)
+        v0, v1 = sv_out[idx_0], sv_out[idx_1]
+        p1 = jnp.clip(gamma * jnp.sum(jnp.abs(v1) ** 2), 0.0, 1.0)
+        decay = r < p1
+        norm_decay    = jnp.sqrt(jnp.maximum(p1, 1e-15))
+        norm_no_decay = jnp.sqrt(jnp.maximum(1.0 - p1, 1e-15))
+        new_v0 = jnp.where(decay, v1 * jnp.sqrt(gamma) / norm_decay, v0 / norm_no_decay)
+        new_v1 = jnp.where(decay, 0.0 + 0j, v1 * jnp.sqrt(1.0 - gamma) / norm_no_decay)
+        sv_out = sv_out.at[idx_0].set(new_v0)
+        sv_out = sv_out.at[idx_1].set(new_v1)
+    else:
+        r = rng.random()
+        v0, v1 = sv_out[idx_0].copy(), sv_out[idx_1].copy()
+        p1 = float(np.clip(gamma * np.sum(np.abs(v1) ** 2), 0.0, 1.0))
+        if r < p1:
+            norm_decay = np.sqrt(max(p1, 1e-15))
+            sv_out[idx_0] = v1 * np.sqrt(gamma) / norm_decay
+            sv_out[idx_1] = 0.0
+        else:
+            norm_no_decay = np.sqrt(max(1.0 - p1, 1e-15))
+            sv_out[idx_0] = v0 / norm_no_decay
+            sv_out[idx_1] = v1 * np.sqrt(1.0 - gamma) / norm_no_decay
+    return sv_out, key

dense_evolution/noise/kraus/bitflip.py

@@ -0,0 +1,22 @@
+"""Bit-flip channel: K0=sqrt(1-p)I, K1=sqrt(p)X. X applied with
+probability p, once per qubit per shot, uniformly across the whole
+statevector."""
+import jax
+import jax.numpy as jnp
+
+__all__ = ["apply"]
+
+
+def apply(sv_out, idx_0, idx_1, p, rng, key, is_jax):
+    if is_jax:
+        key, subkey = jax.random.split(key)
+        r    = jax.random.uniform(subkey, shape=(), minval=0.0, maxval=1.0)
+        fire = r < p
+        v0, v1 = sv_out[idx_0], sv_out[idx_1]
+        sv_out = sv_out.at[idx_0].set(jnp.where(fire, v1, v0))
+        sv_out = sv_out.at[idx_1].set(jnp.where(fire, v0, v1))
+    else:
+        r = rng.random()
+        if r < p:
+            sv_out[idx_0], sv_out[idx_1] = sv_out[idx_1].copy(), sv_out[idx_0].copy()
+    return sv_out, key

dense_evolution/noise/kraus/combined.py

@@ -0,0 +1,16 @@
+"""Combined channel: depolarizing(p/2) then amplitude_damping(p/3),
+applied sequentially on the same qubit -- a worst-case NISQ mixture
+(dephasing + relaxation). Reuses `depolarizing.apply` and
+`amplitude_damping.apply` directly rather than duplicating their logic,
+each called with its own scaled probability."""
+from . import depolarizing, amplitude_damping
+
+__all__ = ["apply"]
+
+
+def apply(sv_out, idx_0, idx_1, p, rng, key, is_jax):
+    p_dep = p * 0.5
+    p_damp = p * 0.333333
+    sv_out, key = depolarizing.apply(sv_out, idx_0, idx_1, p_dep, rng, key, is_jax)
+    sv_out, key = amplitude_damping.apply(sv_out, idx_0, idx_1, p_damp, rng, key, is_jax)
+    return sv_out, key

dense_evolution/noise/kraus/depolarizing.py

@@ -0,0 +1,47 @@
+"""Depolarizing channel: K0=sqrt(1-p)I, K1=sqrt(p/3)X, K2=sqrt(p/3)Y,
+K3=sqrt(p/3)Z. Isotropic Pauli error -- equiprobable X, Y, Z given that
+the channel fires at all.
+
+One fire/no-fire decision per qubit per shot (probability p), then --
+only if it fired -- one equiprobable Pauli choice among X/Y/Z. The
+1-in-3 choice thresholds are fixed at 1/3 and 2/3 regardless of p: p
+only gates whether an error happens at all, not which Pauli it is.
+"""
+import jax
+import jax.numpy as jnp
+
+__all__ = ["apply"]
+
+_THIRD = 1.0 / 3.0
+
+
+def apply(sv_out, idx_0, idx_1, p, rng, key, is_jax):
+    if is_jax:
+        key, sk1, sk2 = jax.random.split(key, 3)
+        r  = jax.random.uniform(sk1, shape=(), minval=0.0, maxval=1.0)
+        ch = jax.random.uniform(sk2, shape=(), minval=0.0, maxval=1.0)
+        fire   = r < p
+        x_gate = fire & (ch < _THIRD)
+        y_gate = fire & (ch >= _THIRD) & (ch < 2.0 * _THIRD)
+        z_gate = fire & (ch >= 2.0 * _THIRD)
+        v0, v1 = sv_out[idx_0], sv_out[idx_1]
+        new_v0 = jnp.where(x_gate,  v1,
+                 jnp.where(y_gate, -1j * v1, v0))
+        new_v1 = jnp.where(x_gate,  v0,
+                 jnp.where(y_gate,  1j * v0,
+                 jnp.where(z_gate, -v1, v1)))
+        sv_out = sv_out.at[idx_0].set(new_v0)
+        sv_out = sv_out.at[idx_1].set(new_v1)
+    else:
+        r  = rng.random()
+        ch = rng.random()
+        if r < p:
+            if ch < _THIRD:
+                sv_out[idx_0], sv_out[idx_1] = sv_out[idx_1].copy(), sv_out[idx_0].copy()
+            elif ch < 2.0 * _THIRD:
+                v0, v1 = sv_out[idx_0].copy(), sv_out[idx_1].copy()
+                sv_out[idx_0] = -1j * v1
+                sv_out[idx_1] =  1j * v0
+            else:
+                sv_out[idx_1] = -sv_out[idx_1]
+    return sv_out, key

dense_evolution/noise/kraus/ideal.py

@@ -0,0 +1,10 @@
+"""Identity channel -- no noise. K0 = I."""
+
+__all__ = ["apply"]
+
+
+def apply(sv_out, idx_0, idx_1, p, rng, key, is_jax):
+    """No-op: returns `(sv_out, key)` unchanged. Kept as its own channel
+    (rather than a special case elsewhere) so `NoiseModel.MODELS` has one
+    real implementation per entry, `'ideal'` included."""
+    return sv_out, key

dense_evolution/noise/kraus/phaseflip.py

@@ -0,0 +1,21 @@
+"""Phase-flip channel: K0=sqrt(1-p)I, K1=sqrt(p)Z. Z applied with
+probability p, once per qubit per shot: Z|0> = |0> (idx_0 amplitudes
+untouched), Z|1> = -|1> (idx_1 amplitudes negated when fired)."""
+import jax
+import jax.numpy as jnp
+
+__all__ = ["apply"]
+
+
+def apply(sv_out, idx_0, idx_1, p, rng, key, is_jax):
+    if is_jax:
+        key, subkey = jax.random.split(key)
+        r    = jax.random.uniform(subkey, shape=(), minval=0.0, maxval=1.0)
+        fire = r < p
+        v1     = sv_out[idx_1]
+        sv_out = sv_out.at[idx_1].set(jnp.where(fire, -v1, v1))
+    else:
+        r = rng.random()
+        if r < p:
+            sv_out[idx_1] = -sv_out[idx_1]
+    return sv_out, key