Registry & Noise Models

registry

NoiseSpec

NoiseSpec(
    model: str,
    p,
    jax_key,
    qubits: Optional[List[int]] = None,
)

JAX PyTree wrapper around a NoiseModel configuration, so noise parameters can be threaded 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.

Source code in dense_evolution/registry.py

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

NoiseModel

Stochastic single-qubit Kraus channels applied directly to a statevector.

All channels are mathematically correct Kraus maps: - trace is preserved (normalisation enforced at the end) - phaseflip applies Z with probability p per qubit (non-deterministic) - amplitude_damping applies the correct K0/K1 Kraus operators - combined is a true worst-case NISQ mixture of all three Pauli errors plus amplitude damping

Supported models

'ideal' identity — no modification 'depolarizing' {√(1-p)I, √(p/3)X, √(p/3)Y, √(p/3)Z} 'bitflip' {√(1-p)I, √p·X} 'phaseflip' {√(1-p)I, √p·Z} ← was broken, now fixed 'amplitude_damping'{K0=diag(1,√(1-γ)), K1=[[0,√γ],[0,0]]} 'combined' depolarizing(p/2) + amplitude_damping(p/3), renormalised

apply_to_sv staticmethod

apply_to_sv(
    sv: ndarray,
    n: int,
    model: str,
    p: float,
    rng: Optional[Generator] = None,
    qubits: Optional[List[int]] = None,
    jax_key: Optional[Any] = None,
) -> np.ndarray

Apply a stochastic Kraus channel to statevector sv in-place (numpy path) or via functional updates (JAX path).

Parameters

sv : complex statevector of length 2n n : number of qubits model : one of NoiseModel.MODELS p : error probability (or damping rate γ for amplitude_damping) rng : optional pre-seeded numpy Generator. Used directly when sv is a NumPy array. When sv is a JAX array, rng used to be silently ignored in favor of jax_key (or a non-reproducible OS-entropy key if that was also None -- issue #7); it is now used to derive a reproducible jax_key (rng.integers(...) seeds jax.random.PRNGKey) whenever jax_key isn't given explicitly, so seeding rng has the effect a caller expects on both array types instead of only on one of them. qubits : subset of qubits to apply the channel to; defaults to all jax_key : optional JAX PRNGKey, only meaningful when sv is a JAX array. Takes precedence over rng when both are given (explicit key beats a derived one). Created from OS entropy if neither jax_key nor rng is given.

Returns

Normalised statevector (same array type as input).

Source code in dense_evolution/registry.py

@staticmethod
def apply_to_sv(
    sv:       np.ndarray,
    n:        int,
    model:    str,
    p:        float,
    rng:      Optional[np.random.Generator] = None,
    qubits:   Optional[List[int]] = None,
    jax_key:  Optional[Any] = None,
) -> np.ndarray:
    """
    Apply a stochastic Kraus channel to statevector *sv* in-place
    (numpy path) or via functional updates (JAX path).

    Parameters
    ----------
    sv      : complex statevector of length 2**n
    n       : number of qubits
    model   : one of NoiseModel.MODELS
    p       : error probability (or damping rate γ for amplitude_damping)
    rng     : optional pre-seeded numpy Generator. Used directly when
              *sv* is a NumPy array. When *sv* is a JAX array, `rng`
              used to be silently ignored in favor of `jax_key` (or a
              non-reproducible OS-entropy key if that was also None
              -- issue #7); it is now used to *derive* a reproducible
              jax_key (`rng.integers(...)` seeds `jax.random.PRNGKey`)
              whenever `jax_key` isn't given explicitly, so seeding
              `rng` has the effect a caller expects on both array
              types instead of only on one of them.
    qubits  : subset of qubits to apply the channel to; defaults to all
    jax_key : optional JAX PRNGKey, only meaningful when *sv* is a JAX
              array. Takes precedence over `rng` when both are given
              (explicit key beats a derived one). Created from OS
              entropy if neither `jax_key` nor `rng` is given.

    Returns
    -------
    Normalised statevector (same array type as input).
    """
    if model == 'ideal':
        return sv
    try:
        if p <= 0.0:
            return sv
    except jax.errors.TracerBoolConversionError:
        # p is a traced value (e.g. flowing through jax.jit/vmap/grad
        # as a NoiseSpec pytree leaf) -- can't early-exit on a Python
        # bool of it. Falling through is still correct: every channel
        # below already reduces to a no-op at p=0 (`fire = r < p` is
        # always False), this skips only the eager-mode optimization,
        # not correctness.
        pass

    is_jax = HAS_JAX and isinstance(sv, jnp.ndarray)
    dim    = len(sv)

    # ── RNG initialisation ────────────────────────────────────────
    if is_jax:
        if jax_key is not None:
            key = jax_key
        elif rng is not None:
            # Derive a reproducible JAX key from the caller's seeded
            # NumPy generator instead of silently ignoring it -- each
            # call advances `rng`'s state, so a fresh, identically-
            # seeded `rng` reproduces the exact same sequence of keys
            # across separate runs (same guarantee the NumPy path
            # already gives).
            key = jax.random.PRNGKey(int(rng.integers(0, 2**32 - 1)))
        else:
            seed_bytes = os.urandom(4)
            jax_seed   = int.from_bytes(seed_bytes, byteorder='big')
            jax_seed  ^= time.perf_counter_ns() & 0xFFFF_FFFF
            key = jax.random.PRNGKey(jax_seed)
    else:
        if rng is None:
            rng = _fresh_rng()

    target_qubits = qubits if qubits is not None else list(range(n))
    sv_out = sv  # JAX: functional; NumPy: will be modified in-place copy

    if not is_jax:
        sv_out = sv.copy()  # never mutate the caller's array

    for q in target_qubits:
        # ── correct index pair construction ───────────────────────
        idx_0, idx_1 = _qubit_index_pairs(dim, q)
        half = len(idx_0)   # == dim // 2

        # ── draw random numbers ───────────────────────────────────
        if is_jax:
            key, subkey = jax.random.split(key)
            r = jax.random.uniform(subkey, shape=(half,), minval=0.0, maxval=1.0)
        else:
            r = rng.random(half)            # uniform [0, 1)

        # ── channel application ───────────────────────────────────
        if model == 'depolarizing':
            # Three equiprobable Pauli errors GIVEN that the channel
            # fired (fire = r < p already gates the overall p rate) —
            # the choice AMONG X/Y/Z must be a uniform 1-in-3 pick,
            # independent of p. ch is drawn uniform on [0,1), so the
            # thresholds here are fixed at 1/3 and 2/3, not p/3 and
            # 2p/3 (that was comparing a full-range [0,1) draw against
            # thresholds scaled for a [0,p) draw, skewing the outcome
            # heavily toward Z for any p < 1 — verified directly:
            # p=0.3 gave P(X|fire)=P(Y|fire)=10%, P(Z|fire)=80% instead
            # of the correct 33.3% each).
            THIRD = 1.0 / 3.0
            if is_jax:
                key, subkey2 = jax.random.split(key)
                ch = jax.random.uniform(subkey2, shape=(half,), minval=0.0, maxval=1.0)
                v0, v1      = sv_out[idx_0], sv_out[idx_1]
                fire        = r < p
                x_gate      = fire & (ch < THIRD)
                y_gate      = fire & (ch >= THIRD) & (ch < 2.0 * THIRD)
                z_gate      = fire & (ch >= 2.0 * THIRD)
                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:
                ch     = rng.random(half)
                v0, v1 = sv_out[idx_0].copy(), sv_out[idx_1].copy()
                fire   = r < p
                x_gate = fire & (ch < THIRD)
                y_gate = fire & (ch >= THIRD) & (ch < 2.0 * THIRD)
                z_gate = fire & (ch >= 2.0 * THIRD)
                sv_out[idx_0] = np.where(x_gate,  v1,
                                np.where(y_gate, -1j * v1, v0))
                sv_out[idx_1] = np.where(x_gate,  v0,
                                np.where(y_gate,  1j * v0,
                                np.where(z_gate, -v1, v1)))

        elif model == 'bitflip':
            # X gate applied with probability p
            fire = r < p
            if is_jax:
                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:
                v0, v1 = sv_out[idx_0].copy(), sv_out[idx_1].copy()
                sv_out[idx_0] = np.where(fire, v1, v0)
                sv_out[idx_1] = np.where(fire, v0, v1)

        elif model == 'phaseflip':
            # Z gate applied with probability p:
            # Z|0⟩ = |0⟩  →  no change to idx_0 amplitudes
            # Z|1⟩ = -|1⟩ →  negate idx_1 amplitudes when fired
            fire = r < p
            if is_jax:
                v1     = sv_out[idx_1]
                sv_out = sv_out.at[idx_1].set(jnp.where(fire, -v1, v1))
            else:
                v1            = sv_out[idx_1].copy()
                sv_out[idx_1] = np.where(fire, -v1, v1)

        elif model == 'amplitude_damping':
            # K0 = [[1, 0], [0, √(1-γ)]]  — no decay
            # K1 = [[0, √γ], [0, 0]]       — decay |1⟩ → |0⟩
            # Applied stochastically: with probability γ the qubit
            # decays (K1 path), otherwise K0 is applied.
            gamma = float(np.clip(p, 0.0, 1.0))
            decay = r < gamma
            if is_jax:
                v0, v1 = sv_out[idx_0], sv_out[idx_1]
                # decay path:    v0 += v1 * √γ,  v1 = 0
                # no-decay path: v0 unchanged,   v1 *= √(1-γ)
                sq_gamma    = jnp.sqrt(gamma)
                sq_1m_gamma = jnp.sqrt(1.0 - gamma)
                new_v0 = jnp.where(decay, v0 + v1 * sq_gamma, v0)
                new_v1 = jnp.where(decay, 0.0 + 0j,           v1 * sq_1m_gamma)
                sv_out = sv_out.at[idx_0].set(new_v0)
                sv_out = sv_out.at[idx_1].set(new_v1)
            else:
                v0, v1      = sv_out[idx_0].copy(), sv_out[idx_1].copy()
                sq_gamma    = np.sqrt(gamma)
                sq_1m_gamma = np.sqrt(1.0 - gamma)
                sv_out[idx_0] = np.where(decay, v0 + v1 * sq_gamma, v0)
                sv_out[idx_1] = np.where(decay, 0.0 + 0j,           v1 * sq_1m_gamma)

        elif model == 'combined':
            # Worst-case NISQ: depolarizing(p/2) + amplitude_damping(p/3)
            # applied sequentially on the same qubit.
            p_dep   = p * 0.5
            p_damp  = p * 0.333333
            THIRD   = 1.0 / 3.0  # see the 'depolarizing' branch above for why
                                 # this must be fixed at 1/3, not p_dep/3

            # — depolarizing sub-channel —
            if is_jax:
                key, sk1, sk2 = jax.random.split(key, 3)
                r_dep = jax.random.uniform(sk1, shape=(half,), minval=0.0, maxval=1.0)
                ch    = jax.random.uniform(sk2, shape=(half,), minval=0.0, maxval=1.0)
                v0, v1 = sv_out[idx_0], sv_out[idx_1]
                fire   = r_dep < p_dep
                x_gate = fire & (ch < THIRD)
                y_gate = fire & (ch >= THIRD) & (ch < 2.0 * THIRD)
                z_gate = fire & (ch >= 2.0 * THIRD)
                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)
                # — amplitude_damping sub-channel —
                key, sk3   = jax.random.split(key)
                r_damp     = jax.random.uniform(sk3, shape=(half,), minval=0.0, maxval=1.0)
                decay      = r_damp < p_damp
                v0, v1     = sv_out[idx_0], sv_out[idx_1]
                sq_g       = jnp.sqrt(p_damp)
                sq_1mg     = jnp.sqrt(1.0 - p_damp)
                sv_out     = sv_out.at[idx_0].set(jnp.where(decay, v0 + v1 * sq_g, v0))
                sv_out     = sv_out.at[idx_1].set(jnp.where(decay, 0.0 + 0j,       v1 * sq_1mg))
            else:
                r_dep  = rng.random(half)
                ch     = rng.random(half)
                v0, v1 = sv_out[idx_0].copy(), sv_out[idx_1].copy()
                fire   = r_dep < p_dep
                x_gate = fire & (ch < THIRD)
                y_gate = fire & (ch >= THIRD) & (ch < 2.0 * THIRD)
                z_gate = fire & (ch >= 2.0 * THIRD)
                sv_out[idx_0] = np.where(x_gate,  v1,
                                np.where(y_gate, -1j * v1, v0))
                sv_out[idx_1] = np.where(x_gate,  v0,
                                np.where(y_gate,  1j * v0,
                                np.where(z_gate, -v1, v1)))
                r_damp        = rng.random(half)
                decay         = r_damp < p_damp
                v0, v1        = sv_out[idx_0].copy(), sv_out[idx_1].copy()
                sq_g          = np.sqrt(p_damp)
                sq_1mg        = np.sqrt(1.0 - p_damp)
                sv_out[idx_0] = np.where(decay, v0 + v1 * sq_g, v0)
                sv_out[idx_1] = np.where(decay, 0.0 + 0j,       v1 * sq_1mg)

    # ── normalise ─────────────────────────────────────────────────
    if is_jax:
        norm = jnp.linalg.norm(sv_out)
        return sv_out / (norm + 1e-15)
    else:
        norm = np.linalg.norm(sv_out)
        return sv_out / (norm + 1e-15)