dense-evolution 8.3.0__py3-none-win_amd64.whl

dense_evolution/mitigation/zne.py

@@ -0,0 +1,990 @@
+"""
+Zero-Noise Extrapolation (ZNE)
+-------------------------------
+Standard error-mitigation entry points, named the way the field already
+names them (Richardson extrapolation, noise factors, zero-noise
+extrapolation -- same vocabulary as e.g. Mitiq's `zne` API), so callers
+and tooling can find "ZNE" without first learning Dense-Evolution's
+internal healing vocabulary.
+
+This module composes `dense_evolution.healing`'s existing primitives
+(`calculate_delta_preemp`, ...) -- it does not rename or replace them.
+"""
+import functools
+import warnings
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+from scipy.optimize import minimize
+
+from .healing import calculate_delta_preemp
+
+__all__ = ["richardson_extrapolate", "richardson_amplification_factor", "zero_noise_extrapolation",
+           "polynomial_extrapolate", "bounded_exponential_extrapolate",
+           "project_to_physical", "uhlmann_fidelity", "zne_density_matrix",
+           "jsd_predictive_zne_density_matrix", "coherence_predictive_zne_density_matrix",
+           "classically_augmented_zne_phaseflip", "global_depolarizing_channel",
+           "amplitude_damping_channel", "cosmic_ray_burst_profile",
+           "richardson_extrapolate_jit", "zero_noise_extrapolation_jit",
+           "polynomial_extrapolate_jit", "uhlmann_fidelity_jit", "zne_density_matrix_jit"]
+
+_KAPPA_WARNING_THRESHOLD = 50.0
+_MULTI_START_SEED = 0
+_N_MULTI_STARTS = 8
+
+
+def richardson_extrapolate(expectation_values, noise_factors) -> jnp.ndarray:
+    """Polynomial (Lagrange) Richardson extrapolation to zero noise.
+
+    `expectation_values[i]` is the value measured/simulated at noise scale
+    `noise_factors[i]` (e.g. 1x, 2x, 3x folded/scaled noise) -- a scalar,
+    or itself an array (e.g. a full probability distribution sampled at
+    that noise scale; extrapolated elementwise). Returns the extrapolated
+    zero-noise estimate, same shape as one `expectation_values[i]`. Works
+    for any number of points and any (not necessarily equally spaced)
+    noise factors; for the common 3-point case at noise_factors=(1,2,3)
+    this reduces exactly to the textbook coefficients (3, -3, 1).
+
+    `expectation_values` may be complex (e.g. density matrix entries,
+    which are complex off-diagonal in general) -- dtype is picked from
+    the input itself (complex128 if complex, float64 otherwise, matching
+    this function's previous always-float64 behavior for real input
+    exactly). Found via a real test case: forcing float64 unconditionally
+    silently discarded the imaginary part of complex input with no
+    visible error, only a low-signal ComplexWarning easy to miss --
+    confirmed directly (`richardson_extrapolate([1+2j, 3+4j], ...)` used
+    to return a purely real result, dropping real information).
+
+    Passing more noise-scale points makes exact interpolation MORE, not
+    less, sensitive to shot noise: at noise_factors equally spaced in
+    [1, 3], the noise-amplification factor kappa (`richardson_amplification_factor`,
+    the sum of absolute Lagrange coefficients) is 7 at 3 points, 129 at 5,
+    2815 at 7 -- verified directly, exact integers, not estimates. Since
+    the coefficients are fixed constants (independent of the measured
+    values), i.i.d. shot noise of per-point standard deviation sigma
+    propagates to an extrapolated-result standard deviation of
+    `sigma * sqrt(sum(coeff_i**2))`: at sigma=0.01 this is 0.044 (3
+    points), 0.67 (5), 13.2 (7) -- verified with this exact formula, not
+    simulated, and matching a real prior audit's Monte Carlo figures
+    (0.044 / 0.67 / 13.5) to within sampling noise. Against this,
+    prog.txt's own audit reports the systematic (noise-free) bias
+    improving by only about 0.01 over the same 3-to-7-point range on its
+    real experimental setup -- a bad trade past a handful of points. A
+    `UserWarning` fires when kappa exceeds `_KAPPA_WARNING_THRESHOLD`
+    (50); see `richardson_amplification_factor` to check kappa before
+    extrapolating.
+
+    Examples
+    --------
+    >>> from dense_evolution.mitigation import richardson_extrapolate
+    >>> round(float(richardson_extrapolate([0.90, 0.80, 0.65], [1.0, 2.0, 3.0])), 4)
+    0.95
+    """
+    lambdas = jnp.asarray(noise_factors, dtype=jnp.float64)
+    kappa = richardson_amplification_factor(lambdas)
+    if kappa > _KAPPA_WARNING_THRESHOLD:
+        warnings.warn(
+            f"richardson_extrapolate: noise amplification factor kappa={kappa:.1f} exceeds "
+            f"{_KAPPA_WARNING_THRESHOLD:.0f} -- shot noise in expectation_values will be "
+            f"amplified by roughly this factor in the extrapolated result; consider passing "
+            f"fewer noise_factors points.",
+            UserWarning, stacklevel=2,
+        )
+    values_dtype = jnp.complex128 if np.iscomplexobj(np.asarray(expectation_values)) else jnp.float64
+    values = jnp.asarray(expectation_values, dtype=values_dtype)
+    return _richardson_extrapolate_core(values, lambdas)
+
+
+def _lagrange_coeffs_at_zero(lambdas: jnp.ndarray) -> jnp.ndarray:
+    """Lagrange basis polynomials for nodes `lambdas`, each evaluated at
+    x=0 -- shared by `_richardson_extrapolate_core` (which combines them
+    with the measured values) and `richardson_amplification_factor`
+    (which only needs their magnitudes)."""
+    n = lambdas.shape[0]
+
+    def lagrange_coeff(i):
+        others = jnp.concatenate([lambdas[:i], lambdas[i + 1:]])
+        return jnp.prod((0.0 - others) / (lambdas[i] - others))
+
+    return jnp.stack([lagrange_coeff(i) for i in range(n)])
+
+
+def richardson_amplification_factor(noise_factors) -> float:
+    """Noise-amplification factor kappa = sum(|Lagrange coefficients|) for
+    Richardson extrapolation at the given `noise_factors`: how much a
+    unit of i.i.d. shot noise spread evenly across the measured points
+    gets amplified in the extrapolated zero-noise estimate (each
+    coefficient can be much larger than 1 and they alternate in sign, so
+    they do not cancel in the worst case the way their SUM, which is
+    always exactly 1, might suggest).
+
+    Measured at noise_factors equally spaced in [1, 3]: kappa = 7 at 3
+    points, 129 at 5 points, 2815 at 7 points -- see
+    `richardson_extrapolate`'s own docstring for the resulting bias/
+    variance trade-off in these same three cases.
+
+    Examples
+    --------
+    >>> from dense_evolution.mitigation import richardson_amplification_factor
+    >>> round(richardson_amplification_factor([1.0, 2.0, 3.0]), 4)
+    7.0
+    """
+    lambdas = jnp.asarray(noise_factors, dtype=jnp.float64)
+    coeffs = _lagrange_coeffs_at_zero(lambdas)
+    return float(jnp.sum(jnp.abs(coeffs)))
+
+
+def _richardson_extrapolate_core(values: jnp.ndarray, lambdas: jnp.ndarray) -> jnp.ndarray:
+    """`jax.jit`-traceable core of `richardson_extrapolate` -- `values`
+    already cast to its final dtype, `lambdas` already float64 (no
+    `np.iscomplexobj`/`np.asarray` call on a possibly-traced value). `n`
+    comes from `lambdas.shape[0]`, a static value even under tracing
+    (array shapes are always known at trace time), so the Python
+    `range(n)` unroll below is trace-safe without `n` needing to be
+    static-marked by callers.
+    """
+    n = lambdas.shape[0]
+    coeffs = _lagrange_coeffs_at_zero(lambdas)
+    # Broadcast coeffs against the LEADING axis of values (the "one row per
+    # noise scale" axis), not jnp's default trailing-axis alignment --
+    # values may itself be array-valued per scale (values.shape = (n,
+    # *extra_dims)), e.g. a whole probability distribution rather than a
+    # bare scalar. A no-op reshape when values is 1-D (the scalar case),
+    # so existing scalar callers are unaffected.
+    coeffs = coeffs.reshape((n,) + (1,) * (values.ndim - 1))
+    return jnp.sum(coeffs * values, axis=0)
+
+
+richardson_extrapolate_jit = jax.jit(_richardson_extrapolate_core)
+"""`jax.jit`-compiled entry point for `richardson_extrapolate`. Unlike
+`polynomial_extrapolate`/`zne_density_matrix`'s jitted variants, no
+argument needs to be marked static here -- `n` (the point count) is read
+from `lambdas.shape[0]`, itself always static under tracing, not from a
+Python `degree` parameter.
+
+`values` must already be complex128 or float64 (pick the dtype yourself
+before calling -- this skips `richardson_extrapolate`'s `np.iscomplexobj`
+auto-detection, which isn't traceable) and `lambdas` a float64 array.
+Verified to match `richardson_extrapolate` exactly on real and complex
+input; JAX recompiles per distinct input shape/dtype, as usual.
+"""
+
+
+def zero_noise_extrapolation(expectation_values, noise_factors,
+                              sigma_at_base_noise=None,
+                              target_sigma_ideal: float = 10.0) -> jnp.ndarray:
+    """Zero-Noise Extrapolation -- plain, or healing-adapted when a
+    coherence signal is available.
+
+    Without `sigma_at_base_noise`: standard Richardson ZNE
+    (`richardson_extrapolate`).
+
+    With `sigma_at_base_noise` (the measured/simulated coherence sigma at
+    the base, unscaled noise level): the 3 Richardson coefficients are
+    perturbed by `dense_evolution.healing.calculate_delta_preemp` -- the
+    normalized deviation between the observed sigma and the ideal target
+    -- then renormalized to sum to 1. This is Dense-Evolution's
+    "predictive healing" ZNE variant: when the observed coherence is off
+    the ideal target, the extrapolation is nudged accordingly instead of
+    trusting the 3 raw noise-scaled points equally.
+
+    The healing-adapted path currently only supports exactly 3 noise
+    factors (the case it has been derived and tested against); passing
+    `sigma_at_base_noise` with any other point count raises
+    NotImplementedError rather than silently generalizing an unverified
+    formula.
+    """
+    if sigma_at_base_noise is None:
+        return richardson_extrapolate(expectation_values, noise_factors)
+
+    lambdas = jnp.asarray(noise_factors, dtype=jnp.float64)
+    if lambdas.shape[0] != 3:
+        raise NotImplementedError(
+            "Healing-adapted ZNE is only defined for exactly 3 noise factors; "
+            "call richardson_extrapolate(...) directly for the plain N-point case."
+        )
+    values_dtype = jnp.complex128 if np.iscomplexobj(np.asarray(expectation_values)) else jnp.float64
+    values = jnp.asarray(expectation_values, dtype=values_dtype)
+    return _zero_noise_extrapolation_healing_core(
+        values, jnp.asarray(sigma_at_base_noise, dtype=jnp.float64), target_sigma_ideal)
+
+
+def _zero_noise_extrapolation_healing_core(values: jnp.ndarray, sigma_at_base_noise: jnp.ndarray,
+                                            target_sigma_ideal: float) -> jnp.ndarray:
+    """`jax.jit`-traceable core of `zero_noise_extrapolation`'s
+    healing-adapted branch -- `values` already cast to its final dtype
+    (exactly 3 noise-scale rows, `values[0]`, `values[1]`, `values[2]`),
+    `sigma_at_base_noise` already a jnp float64 scalar. `calculate_delta_preemp`
+    is itself already `@jax.jit`-decorated (`dense_evolution.healing`), so
+    calling it here composes cleanly under an outer jit.
+    """
+    e_l1, e_l2, e_l3 = values[0], values[1], values[2]
+    delta_p = calculate_delta_preemp(sigma_at_base_noise, target_sigma_ideal)
+    c1 = 3.0 - 0.01 * delta_p
+    c2 = -3.0 + 0.02 * delta_p
+    c3 = 1.0 - 0.01 * delta_p
+    return (c1 * e_l1 + c2 * e_l2 + c3 * e_l3) / (c1 + c2 + c3)
+
+
+zero_noise_extrapolation_jit = jax.jit(_zero_noise_extrapolation_healing_core)
+"""`jax.jit`-compiled entry point for `zero_noise_extrapolation`'s
+healing-adapted branch (the `sigma_at_base_noise is not None` case --
+the plain-Richardson case already has `richardson_extrapolate_jit`, use
+that directly instead). `values` must already be complex128 or float64
+with exactly 3 rows, `sigma_at_base_noise` a float64 scalar; the
+`lambdas.shape[0] != 3` validation and dtype auto-detection that
+`zero_noise_extrapolation` does are both skipped here (not traceable) --
+callers are responsible for passing exactly-3-row input themselves.
+`target_sigma_ideal` is a plain Python float, fine to leave non-static
+since it only ever multiplies/subtracts, no Python branching on its value.
+"""
+
+
+def polynomial_extrapolate(expectation_values, noise_factors, degree: int = 2) -> jnp.ndarray:
+    """Least-squares polynomial extrapolation to zero noise.
+
+    Generalizes `richardson_extrapolate`: fits a degree-`degree` polynomial
+    to `(noise_factors, expectation_values)` by ordinary least squares and
+    evaluates it at zero. With exactly `degree + 1` points the fit is the
+    unique interpolating polynomial, mathematically identical to
+    `richardson_extrapolate` at that point count (verified directly, both
+    for real and complex input). With MORE than `degree + 1` points it
+    becomes an overdetermined fit -- the extra points average down
+    statistical noise instead of forcing the polynomial through every
+    noisy sample exactly, trading a small amount of interpolation bias
+    for reduced variance.
+
+    This matters in practice, not just in theory: adding more noise-scale
+    points to *exact* interpolation (`richardson_extrapolate`) makes
+    extrapolation WORSE under real statistical noise, because Lagrange
+    coefficients grow with point count (worse still with closely-spaced
+    points -- a Runge's-phenomenon-like effect). Measured directly on the
+    density-matrix healing experiment (`experiments/matrix_healing_zne_sweep.py`'s
+    setup, n=4 qubits, all 5 noise channels, 5 seeds): exact interpolation's
+    mean fidelity-delta dropped from +0.148 (3 points) to +0.081 (5 points,
+    same spacing) to -0.220 (5 points, denser spacing -- actively worse
+    than not correcting). A degree-2 least-squares fit fed the same extra
+    points instead REDUCES variance (std 0.062 -> 0.035-0.046) at
+    comparable or better mean delta, because the extra points are no
+    longer forced to satisfy an increasingly ill-conditioned exact fit.
+    This is why `zne_density_matrix` uses this function (degree=2) instead
+    of `richardson_extrapolate` by default.
+
+    Raises ValueError if fewer than `degree + 1` points are given (the fit
+    would be underdetermined).
+    """
+    lambdas = jnp.asarray(noise_factors, dtype=jnp.float64)
+    n = lambdas.shape[0]
+    if n < degree + 1:
+        raise ValueError(
+            f"polynomial_extrapolate needs at least degree+1={degree + 1} noise "
+            f"factors for a degree-{degree} fit, got {n}."
+        )
+    values_dtype = jnp.complex128 if np.iscomplexobj(np.asarray(expectation_values)) else jnp.float64
+    values = jnp.asarray(expectation_values, dtype=values_dtype)
+    return _polynomial_extrapolate_core(values, lambdas, degree)
+
+
+def _polynomial_extrapolate_core(values: jnp.ndarray, lambdas: jnp.ndarray, degree: int) -> jnp.ndarray:
+    """`jax.jit`-traceable core of `polynomial_extrapolate` -- takes `values`
+    already cast to its final dtype and `lambdas` already a float64 array,
+    so there's no `np.iscomplexobj`/`np.asarray` call on a possibly-traced
+    value (which breaks tracing; `np.asarray` on a JAX tracer raises).
+    `degree` must be a Python int, not a traced value (`range(degree + 1)`
+    unrolls it at trace time) -- callers that `jax.jit` a function using
+    this must mark `degree` static (`static_argnames`).
+    """
+    n = lambdas.shape[0]
+    orig_shape = values.shape[1:]
+    flat = values.reshape(n, -1)
+    # Vandermonde design matrix: columns lambda^0, lambda^1, ..., lambda^degree.
+    design = jnp.stack([lambdas ** k for k in range(degree + 1)], axis=1).astype(values.dtype)
+    coeffs, *_ = jnp.linalg.lstsq(design, flat, rcond=None)
+    intercept = coeffs[0]  # fitted polynomial evaluated at noise_factor=0
+    return intercept.reshape(orig_shape)
+
+
+polynomial_extrapolate_jit = functools.partial(jax.jit, static_argnames=("degree",))(_polynomial_extrapolate_core)
+"""`jax.jit`-compiled entry point for `polynomial_extrapolate`, added for
+consistency with every other function in this module (`richardson_extrapolate_jit`,
+`zero_noise_extrapolation_jit`, `uhlmann_fidelity_jit`, `zne_density_matrix_jit`)
+-- until now this was the one function whose `_core` existed (used
+internally by `zne_density_matrix_jit`) but had no standalone public jit
+entry point of its own.
+
+`values` must already be cast to its final dtype (complex128 or float64)
+and `lambdas` a float64 array -- this skips `polynomial_extrapolate`'s
+`np.iscomplexobj` auto-detection, not traceable. `degree` is static (same
+constraint as `zne_density_matrix_jit`). Verified to match
+`polynomial_extrapolate` exactly."""
+
+
+def bounded_exponential_extrapolate(expectation_values, noise_factors, bound: float = 1.0) -> jnp.ndarray:
+    """Physically bounded exponential Zero-Noise Extrapolation (Miranskyy,
+    Sorrenti, Thind & Gravel, arXiv:2604.24475, "Improving Zero-Noise
+    Extrapolation via Physically Bounded Models").
+
+    Unlike `polynomial_extrapolate`, this fits an EXPONENTIAL model --
+    appropriate when the expectation value decays roughly exponentially with
+    noise strength (the usual case for a depolarizing-type channel), not a
+    polynomial one -- and is Dense-Evolution's first exponential-family ZNE
+    model. The zero-noise value is made an explicit model parameter via the
+    reparametrization
+
+        E(lambda) = a + (zeta - a) * exp(-c * lambda),   E(0) = zeta
+
+    and `zeta` is constrained to `[-bound, bound]` during the fit (the
+    physically valid range for a +-1-eigenvalue Pauli observable, `bound=1.0`
+    by default). An ordinary unconstrained `a + b*exp(-c*lambda)` fit has no
+    such guarantee and can produce a wildly out-of-range or non-finite
+    "zero-noise" estimate when the data is noisy or only a few points are
+    available -- verified directly on this project's own depolarizing-noise
+    ZZ-expectation setup (real Bell circuit, 30 random noise seeds, 3 noise
+    scales): the unconstrained fit landed outside [-1, 1] or failed to
+    converge in 21/30 seeds (mean absolute error over 200000, dominated by
+    those blow-ups), versus 0/30 for this bounded fit (mean absolute error
+    0.066) -- see `Dense-Evolution-Discovery/scripts/zne_physically_bounded.py`
+    for the full reproduction.
+
+    Fit via SciPy's constrained L-BFGS-B, not `jax.jit`-traceable like the
+    rest of this module -- the optimization itself runs in plain NumPy, so
+    (unlike every other function here) there is no `_jit` variant. Multi-start
+    (`_N_MULTI_STARTS` starts: the original single fixed start
+    `[0, values[0], 0.5]` first, then `_N_MULTI_STARTS - 1` deterministic
+    random ones, fixed seed `_MULTI_START_SEED`, keeping the converged fit
+    with lowest loss) -- keeping the original start first means a case
+    where it was already the global optimum is unaffected bit-for-bit.
+    The non-convex 3-parameter fit from a single fixed starting point can
+    otherwise land in a poor local minimum -- on data generated exactly
+    from this function's own model, the single-start fit was off by up to
+    3.9e-3, versus under 1e-5 with multi-start. Raises `RuntimeError` if every
+    start fails to converge (`result.success`), rather than silently
+    returning an unconverged `result.x[1]`.
+
+    Examples
+    --------
+    >>> from dense_evolution.mitigation import bounded_exponential_extrapolate
+    >>> round(float(bounded_exponential_extrapolate([0.49, 0.1, 0.1], [1.0, 2.0, 3.0])), 4)
+    1.0
+    """
+    lambdas = np.asarray(noise_factors, dtype=np.float64)
+    values = np.asarray(expectation_values, dtype=np.float64)
+
+    def loss(theta):
+        a, zeta, c = theta
+        pred = a + (zeta - a) * np.exp(-c * lambdas)
+        return np.sum((values - pred) ** 2)
+
+    rng = np.random.default_rng(_MULTI_START_SEED)
+    starts = np.concatenate([
+        np.array([[0.0, float(values[0]), 0.5]]),
+        rng.uniform([-bound, -bound, 0.01], [bound, bound, 3.0], size=(_N_MULTI_STARTS - 1, 3)),
+    ])
+    best_result = None
+    for x0 in starts:
+        result = minimize(
+            loss, x0=x0, method="L-BFGS-B",
+            bounds=[(-bound, bound), (-bound, bound), (1e-6, None)],
+            options={"ftol": 1e-14, "gtol": 1e-12},
+        )
+        if not result.success:
+            continue
+        if best_result is None or result.fun < best_result.fun:
+            best_result = result
+
+    if best_result is None:
+        raise RuntimeError(
+            f"bounded_exponential_extrapolate: none of {_N_MULTI_STARTS} multi-start "
+            f"L-BFGS-B fits converged (result.success was False for every start)."
+        )
+    return jnp.float64(best_result.x[1])
+
+
+def project_to_physical(rho_raw: jnp.ndarray) -> jnp.ndarray:
+    """Project a Hermitian, trace-1 candidate matrix onto the nearest
+    physical density matrix (Hermitian, trace 1, positive-semidefinite) in
+    2-norm/Frobenius distance -- the same problem Smolin, Gambetta & Smith,
+    "Maximum Likelihood, Minimum Effort" (2012), arXiv:1106.5458, Fig. 1,
+    solve with a sequential eigenvalue-clipping algorithm (sort eigenvalues
+    descending, repeatedly zero the smallest remaining one and redistribute
+    its negative mass over the rest, until the least would be non-negative).
+
+    Implemented here as Euclidean projection onto the probability simplex
+    (Held, Wolfe & Crowder 1974; also e.g. Duchi et al. 2008) applied to the
+    eigenvalues -- a different, fully vectorized algorithm for the exact
+    same convex optimization problem (unique global minimum, so any correct
+    algorithm must agree). No Python-level `while`/`for` loop over
+    eigenvalues (the original transcription's `while i >= 0: ...` isn't
+    `jax.jit`-traceable, forcing a host round-trip every call) -- this
+    version is pure `jnp` array ops plus one dynamic index (`mus[k-1]`,
+    itself trace-safe), so it JIT-compiles cleanly.
+
+    Verified against the original transcription (which itself matches the
+    SGS paper's own worked example, eigenvalues 3/5, 1/2, 7/20, 1/10,
+    -11/20 -> 9/20, 7/20, 1/5, 0, 0): identical to machine precision on the
+    paper's example and on 30 random Hermitian trace-1 matrices (2-7 dim)
+    perturbed to be unphysical (max difference ~1e-15); confirmed to
+    actually compile and run under `jax.jit`.
+
+    `richardson_extrapolate`/`polynomial_extrapolate`'s output on a stack
+    of density matrices is not itself generally a valid density matrix --
+    extrapolation can (and in practice does) produce small negative
+    eigenvalues even when every input matrix was physical. This is the
+    correction step, meant to run *after* extrapolation, not a
+    general-purpose "make anything a density matrix" tool (it assumes the
+    input is already Hermitian and trace 1 up to this function's own
+    re-Hermitization step below).
+    """
+    rho_h = 0.5 * (rho_raw + jnp.conj(rho_raw).T)
+    eigvals, eigvecs = jnp.linalg.eigh(rho_h)
+    idx = jnp.argsort(eigvals)[::-1]
+    evals = eigvals[idx]
+    vecs = eigvecs[:, idx]
+
+    d = evals.shape[0]
+    ranks = jnp.arange(1, d + 1)
+    cumsum_evals = jnp.cumsum(evals)
+    # For each prefix length j, the shift mu_j that would make that prefix
+    # (plus this shift) sum to 1; the simplex-projection lemma guarantees
+    # `evals > mus` is a prefix-of-True mask for descending-sorted evals,
+    # so its count k is exactly the largest valid prefix length.
+    mus = (cumsum_evals - 1.0) / ranks
+    k = jnp.sum((evals > mus).astype(jnp.int32))
+    chosen_mu = mus[k - 1]
+
+    projected_evals = jnp.maximum(evals - chosen_mu, 0.0)
+    rho_physical = (vecs * projected_evals) @ jnp.conj(vecs).T
+    return jnp.asarray(rho_physical, dtype=jnp.complex128)
+
+
+@jax.custom_jvp
+def _eigh_degenerate_safe(A: jnp.ndarray):
+    """Same (eigvals, eigvecs) as `jnp.linalg.eigh(A)` -- only the backward
+    (gradient) rule differs, to stay finite at (near-)degenerate
+    eigenvalues, where JAX's own built-in `eigh` gradient rule divides by
+    `lambda_i - lambda_j` and returns NaN (documented upstream, e.g. JAX
+    issues #2311 and #8732; general treatment in Kasim, "Derivatives of
+    partial eigendecomposition of a real symmetric matrix for degenerate
+    cases", arXiv:2011.04366).
+
+    This is a practical variant of that fix: mask the `1/(lambda_i -
+    lambda_j)` term to 0 for near-degenerate pairs instead of letting it
+    blow up, rather than Kasim's fuller per-degenerate-block treatment
+    (which recovers a generally nonzero contribution from within the
+    degenerate eigenspace itself for some perturbation directions).
+    Sufficient here because `uhlmann_fidelity`'s only use of the
+    eigenVECTORS is to rebuild a matrix square root, and simply zeroing
+    that masked term still gives a mathematically valid (if not maximally
+    sharp) subgradient direction there.
+
+    This masking is NOT known to already exist elsewhere, checked directly
+    rather than assumed: PyTorch's current `linalg_eig_jvp`/
+    `linalg_eig_backward` (`torch/csrc/autograd/FunctionsManual.cpp`)
+    divides directly by `(L_j - L_i)` with no degenerate-pair masking at
+    all -- the same singularity as JAX's default, unresolved there too.
+    The `xitorch` library (built by Kasim himself)'s own derivation notes
+    (`doc/notes/deriv_symeig.rst`) state "This derivation assumes the
+    eigenvalues are all unique. Cases with degenerate eigenvalues are
+    treated differently" without giving that treatment on that page, and
+    its `symeig` docstring warns directly that for degenerate values "the
+    calculation and its gradient might be inaccurate". This appears to be
+    a genuinely open problem across the JAX/PyTorch/xitorch ecosystem, not
+    something ported from prior art.
+
+    Verified (directional-derivative/JVP comparison against symmetric-
+    tangent finite differences, the only comparison method that avoids
+    eigenvector sign/ordering convention ambiguities): exact match
+    (diff=0.000000) for non-degenerate, 2-fold, and 3-fold degenerate test
+    matrices. The 1e-8 masking threshold itself is verified, not
+    arbitrary: swept true (non-exact) eigenvalue gaps from 1.0 down to
+    1e-10 against finite differences -- the unmasked formula tracks them
+    to ~1e-9 accuracy for any gap down to ~1e-7, only degrading right at
+    the threshold, which sits at the edge of float64-resolvable gaps for
+    O(1)-scale eigenvalues rather than discarding real signal."""
+    return jnp.linalg.eigh(A)
+
+
+@_eigh_degenerate_safe.defjvp
+def _eigh_degenerate_safe_jvp(primals, tangents):
+    A, = primals
+    dA, = tangents
+    w, v = jnp.linalg.eigh(A)
+    dA_sym = 0.5 * (dA + jnp.conj(dA).T)
+    vt_dA_v = jnp.conj(v).T @ dA_sym @ v
+    dw = jnp.real(jnp.diag(vt_dA_v))
+
+    denom = w[None, :] - w[:, None]
+    is_degenerate = jnp.abs(denom) < 1e-8
+    safe_denom = jnp.where(is_degenerate, 1.0, denom)
+    F = jnp.where(is_degenerate, 0.0, 1.0 / safe_denom)
+
+    dv = v @ (F * vt_dA_v)
+    return (w, v), (dw, dv)
+
+
+def uhlmann_fidelity(rho_A: jnp.ndarray, rho_B: jnp.ndarray) -> float:
+    """Uhlmann fidelity F(rho_A, rho_B) = (Tr sqrt(sqrt(rho_A) rho_B sqrt(rho_A)))^2.
+
+    Reduces to |<psi_A|psi_B>|^2 when both inputs are pure-state density
+    matrices (verified directly). Validation-only: this is meant to grade
+    a correction against a known ideal state, never to feed into one --
+    passing it a target/ideal density matrix as an input to
+    `zne_density_matrix` or any extrapolation step would be using held-out
+    ground truth to guide the algorithm (oracle access), not a legitimate
+    error-mitigation technique. Keeping ideal-state comparison to this
+    function only, rather than plumbing it into the correction functions
+    at all, makes that boundary structural rather than a convention callers
+    have to remember.
+
+    Computes Tr(sqrt(inner)) as sum(sqrt(eigenvalues of inner)) instead of
+    reconstructing the full matrix square root (sqrt(M) has the same
+    eigenvectors as M and sqrt-of-eigenvalues eigenvalues, so its trace is
+    exactly that sum) -- skips one eigenvector reconstruction, and avoids
+    `matsqrt`'s `float()` cast that isn't `jax.jit`-traceable. Verified
+    against the previous full-reconstruction version: identical to machine
+    precision (~1e-16) on 30 random density-matrix pairs.
+
+    Differentiable through both arguments, including when `rho_A` has
+    (near-)degenerate eigenvalues (e.g. a near-pure state's noisy density
+    matrix, which typically has several near-zero, near-degenerate
+    eigenvalues) -- uses `_eigh_degenerate_safe` internally rather than
+    `jnp.linalg.eigh` directly, specifically to keep `jax.grad(uhlmann_fidelity, ...)`
+    finite in that case (see `_eigh_degenerate_safe`'s own docstring).
+    Forward-pass value is bit-identical to `jnp.linalg.eigh`-based
+    computation (same underlying `eigh` call; only the backward rule
+    differs), verified end-to-end against the previous implementation.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from dense_evolution.mitigation import uhlmann_fidelity
+    >>> rho = np.array([[1, 0], [0, 0]], dtype=complex)
+    >>> round(float(uhlmann_fidelity(rho, rho)), 4)
+    1.0
+    >>> sigma = np.array([[0.5, 0], [0, 0.5]], dtype=complex)
+    >>> round(float(uhlmann_fidelity(rho, sigma)), 4)
+    0.5
+    """
+    rho_A = jnp.asarray(rho_A, dtype=jnp.complex128)
+    rho_B = jnp.asarray(rho_B, dtype=jnp.complex128)
+    return float(_uhlmann_fidelity_core(rho_A, rho_B))
+
+
+# global_depolarizing_channel, amplitude_damping_channel, and
+# cosmic_ray_burst_profile moved to dense_evolution.noise -- they generate
+# noise, they don't mitigate it, so this module (mitigation) was the wrong
+# home. Re-exported below for backward compatibility; new code should
+# import from dense_evolution.noise directly.
+from ..noise import global_depolarizing_channel, amplitude_damping_channel, cosmic_ray_burst_profile
+
+
+def _uhlmann_fidelity_core(rho_A: jnp.ndarray, rho_B: jnp.ndarray) -> jnp.ndarray:
+    """`jax.jit`-traceable core of `uhlmann_fidelity` -- both inputs already
+    complex128; returns a jnp scalar instead of a Python `float` (the
+    `float()` cast isn't traceable, same reason `_jsd_vectors_jax`/
+    `_jsd_vectors` are split this way in `dense_evolution.mps`).
+    """
+    def sqrt_eigvals(m):
+        w = jnp.linalg.eigvalsh(m)
+        return jnp.clip(jnp.real(w), 0.0, None)
+
+    w_A, v_A = _eigh_degenerate_safe(rho_A)
+    sqrt_A = (v_A * jnp.sqrt(jnp.clip(jnp.real(w_A), 0.0, None))) @ jnp.conj(v_A).T
+    inner = sqrt_A @ rho_B @ sqrt_A
+    inner_evals = sqrt_eigvals(inner)
+    return jnp.real(jnp.sum(jnp.sqrt(inner_evals)) ** 2)
+
+
+uhlmann_fidelity_jit = jax.jit(_uhlmann_fidelity_core)
+"""`jax.jit`-compiled entry point for `uhlmann_fidelity`. Both `rho_A`/
+`rho_B` must already be `complex128` (this skips `uhlmann_fidelity`'s
+own `jnp.asarray(..., dtype=jnp.complex128)` cast, itself trace-safe, but
+kept out of the core to mirror the other `_core` functions' convention).
+Returns a jnp scalar, not a Python `float` -- call `float(...)` yourself
+if you need one outside a jitted context. Verified to match `uhlmann_fidelity`
+exactly (same underlying math, just not cast to a Python float)."""
+
+
+def zne_density_matrix(rho_at_scales, noise_factors, degree: int = 2) -> jnp.ndarray:
+    """Zero-Noise Extrapolation for density matrices.
+
+    `rho_at_scales[i]` is a noisy density-matrix estimate (e.g. from a
+    Monte-Carlo/shot ensemble) at noise scale `noise_factors[i]`.
+    Extrapolates to zero noise via `polynomial_extrapolate` (least-squares,
+    complex-safe, degree=2 by default) and projects the result onto the
+    nearest physical density matrix (`project_to_physical`), since the raw
+    extrapolated matrix is not generally positive-semidefinite even when
+    every input was.
+
+    Uses `polynomial_extrapolate` rather than exact `richardson_extrapolate`
+    because, with exactly 3 noise scales (the original design point), the
+    two are mathematically identical -- but `polynomial_extrapolate` stays
+    well-behaved (reduced variance) when a caller passes MORE than 3 scales,
+    where exact interpolation instead gets WORSE (see
+    `polynomial_extrapolate`'s docstring for the measured numbers). This
+    makes "pass more noise-scale points" a safe thing to try rather than a
+    trap.
+
+    Honest findings, both against a GHZ-state ideal target and
+    `dense_evolution.registry.NoiseModel` noise at base_p=0.05, scales
+    1x/2x/3x unless noted, `uhlmann_fidelity` against the true ideal state
+    used only to grade the result -- never as input to any step above:
+
+    - `experiments/matrix_healing_zne.py`: 2-qubit Bell state, depolarizing
+      noise, K=200-trajectory estimate per scale, averaged over 4 seeds --
+      raw fidelity ~0.865, corrected ~0.947 (+0.08), positive on every seed
+      tested.
+    - `experiments/matrix_healing_zne_sweep.py`: 2-5 qubits x all 5
+      `NoiseModel` channels (depolarizing, bitflip, phaseflip,
+      amplitude_damping, combined) x 5 seeds, K=400 trajectories per scale
+      (100 runs total) -- **96/100 positive, mean delta +0.12**, and every
+      single (qubit count, noise channel) combination is net positive on
+      average, growing to +0.20-0.25 at 5 qubits for depolarizing/bitflip.
+      The 4 remaining negative runs are small (worst -0.02) and consistent
+      with residual Monte Carlo noise, not a systematic failure mode.
+      (These specific numbers were measured with exact 3-point
+      interpolation, which is identical to this function's degree=2
+      default at 3 points -- unaffected by the switch.)
+    - An earlier draft of this sweep (K=150, 3 seeds) had reported
+      phaseflip/amplitude_damping as "unreliable" -- re-investigated rather
+      than trusted, and confirmed to be a Monte Carlo undersampling
+      artifact (extrapolation coefficients amplify input noise; an
+      undersampled estimate makes the *corrected* result noisy even when
+      the correction itself is sound), not a real limitation.
+    - More noise-scale points, SAME total measurement budget (the fair
+      comparison -- splitting a fixed number of trajectories across more
+      points, not spending more): 3 points x K=400 (1200 total) vs. 5
+      points x K=240 (1200 total) vs. 7 points x K=171 (~1200 total),
+      n=4 qubits, all 5 noise channels, 5 seeds. 5 points matches or
+      slightly beats the 3-point mean delta (+0.150 vs +0.148) with 19%
+      lower variance (std 0.050 vs 0.062) -- a real, free improvement at
+      the same experimental cost, not an artifact of spending more. 7
+      points trades a little mean (+0.132) for still-lower variance (std
+      0.043, 30% below baseline) -- a genuine tradeoff point, useful when
+      reliability matters more than average performance. At exactly 3
+      points this function is mathematically identical to
+      `richardson_extrapolate` (verified to 1e-12) -- there is no free
+      lunch there, the gain only appears once more points are used.
+    - More noise-scale points, FIXED K per point instead (spending more
+      total measurement, K=400 at every point count): with exact
+      interpolation this makes things worse, not better (mean delta drops
+      from +0.148 at 3 points to +0.081 at 5, and to -0.220 at 5
+      closely-spaced points -- see `polynomial_extrapolate`'s docstring);
+      with this function's degree=2 default it instead stays comparable
+      in mean with lower variance (std 0.062 -> 0.035-0.046) -- confirms
+      the safety property holds even when *not* holding budget fixed, on
+      top of the fixed-budget gain above.
+
+    Practical implication for callers regardless of `degree`: correction
+    quality still depends on `rho_at_scales` being a reasonably low-noise
+    estimate to begin with (large enough K, or equivalent) -- an
+    extrapolation fit through pure noise cannot recover signal that isn't
+    there. `degree` trades bias for variance: higher degree fits the true
+    curve's shape more closely (less bias) but is more sensitive to
+    per-point noise (more variance); degree=2 was chosen empirically as the
+    best tested tradeoff, not a theoretical optimum for every regime.
+
+    Do not pass a target/ideal density matrix into this function or use one
+    to pick among candidate corrections -- see `uhlmann_fidelity`'s
+    docstring for why.
+    """
+    extrapolated = polynomial_extrapolate(rho_at_scales, noise_factors, degree=degree)
+    return project_to_physical(extrapolated)
+
+
+def _zne_density_matrix_core(rho_at_scales: jnp.ndarray, noise_factors: jnp.ndarray,
+                              degree: int) -> jnp.ndarray:
+    """`jax.jit`-traceable core of `zne_density_matrix` -- `rho_at_scales`
+    must already be complex128, `noise_factors` a float64 array; `degree`
+    must be a Python int (unrolled at trace time by
+    `_polynomial_extrapolate_core`, same constraint as there).
+    """
+    extrapolated = _polynomial_extrapolate_core(rho_at_scales, noise_factors, degree)
+    return project_to_physical(extrapolated)
+
+
+zne_density_matrix_jit = functools.partial(jax.jit, static_argnames=("degree",))(_zne_density_matrix_core)
+"""`jax.jit`-compiled entry point for `zne_density_matrix`, for callers
+inside a jitted pipeline (e.g. `jax.lax.scan` in `MPSSimulator.run_circuit_jit`)
+who don't want a host round-trip every call -- `zne_density_matrix` itself
+stays eager (unchanged) for one-off/interactive use, where jit compilation
+overhead isn't worth paying for a single call.
+
+`degree` is a static argument (must be a Python int, not a traced value --
+pass it positionally or by keyword the same way every call, since JAX
+recompiles per distinct static value). `rho_at_scales` must already be
+`complex128` and `noise_factors` a plain float array/sequence -- unlike
+`zne_density_matrix`, this skips the `np.iscomplexobj` dtype auto-detection
+(not traceable) and always assumes complex input, which is the only case
+that makes sense for density matrices.
+
+Measured speedup is real but size- and call-pattern-dependent (`project_to_physical`
+alone measured 2x-22x across 2x2 to 32x32 matrices when jitted vs. the
+previous non-jittable version) -- benchmark your own use case rather than
+assuming a fixed number; the benefit only appears once compiled and called
+repeatedly, a single one-off call pays the compilation cost first.
+"""
+
+
+def _js_divergence(p: jnp.ndarray, q: jnp.ndarray, eps: float = 1e-12) -> jnp.ndarray:
+    """Standard Jensen-Shannon divergence (natural log, bounded in
+    [0, ln 2]) between two probability vectors -- NOT the same quantity
+    as `dense_evolution.mps._jsd_vectors_jax` (a different, MPS-specific
+    "adaptive Jensen-Shannon Distance": base-2 log, an extra
+    log10(dim)/2 dimensional scaling factor, and a final square root,
+    purpose-built for sizing truncated bond dimensions). Both are
+    legitimately named "JSD" for their own purposes; they are not
+    interchangeable, and this one is the plain textbook formula, chosen
+    here to exactly match the already-validated implementation in
+    Dense-Evolution-Discovery's channel_order_noncommutativity.py."""
+    p, q = p + eps, q + eps
+    p, q = p / jnp.sum(p), q / jnp.sum(q)
+    m = 0.5 * (p + q)
+    kl = lambda a, b: jnp.sum(a * jnp.log(a / b))
+    return 0.5 * kl(p, m) + 0.5 * kl(q, m)
+
+
+def jsd_predictive_zne_density_matrix(rho_at_scales, noise_factors) -> jnp.ndarray:
+    """Density-matrix ZNE with a Jensen-Shannon-divergence-informed
+    coefficient nudge, for noise whose scale-to-output-distribution
+    relationship isn't perfectly smooth (the assumption plain
+    3-point Richardson extrapolation, which `zne_density_matrix`
+    defaults toward at exactly 3 scales, relies on).
+
+    Motivated by and validated in Dense-Evolution-Discovery's
+    scripts/photonic_predictive_zne.py -- prototyped there first for
+    photon-loss noise (a photonic-relevant channel: photon loss on a
+    dual-rail-encoded qubit IS this library's `amplitude_damping`
+    channel), per this project's cross-repo promotion pattern. Grounded
+    in real literature: Mills & Mezher, "Mitigating photon loss in
+    linear optical quantum circuits" (arXiv:2405.02278), find plain
+    scalar ZNE does not beat postselection for discrete-variable photon
+    loss -- reproduced directly there (scalar ZNE went unphysical,
+    fidelity > 1.0, at 14/16 swept points). `zne_density_matrix` avoids
+    that failure mode by construction (`project_to_physical`); this
+    function asks whether a further, data-driven adaptive correction on
+    top of it can do even better.
+
+    Signal: Jensen-Shannon divergence (`_js_divergence`, the standard
+    formula -- see its own docstring for how this differs from
+    `dense_evolution.mps`'s unrelated "adaptive JSD") between the
+    measurement-probability distributions (density-matrix diagonals) at
+    consecutive noise scales. Needs no external calibration or oracle
+    access to an ideal/target state -- unlike naively reusing
+    `calculate_delta_preemp` with an externally-supplied signal (tried
+    first in the Discovery prototype; found to have a negligible effect
+    by construction, since that formula's fixed nudge constants 0.01/
+    0.02 were tuned for a differently-scaled use case elsewhere in this
+    module).
+
+    `nonlinearity = (jsd_23 - jsd_12) / (jsd_23 + jsd_12 + eps)`
+    (bounded in [-1, 1]) measures how consistently the JSD grows between
+    consecutive scales -- near 0 when the noise-scale -> output-
+    distribution map is locally well-behaved (Richardson's implicit
+    assumption holding), away from 0 when it isn't. RECTIFIED: the
+    coefficient nudge is applied only when `nonlinearity > 0` -- an
+    unrectified first version, applying the nudge for both signs,
+    helped in only 5/16 points on a real run despite the signal itself
+    being significantly correlated with success (Pearson r=+0.533,
+    p=0.0334); the fix was clipping to the regime the signal was shown
+    to work in, not discarding the signal. When `nonlinearity <= 0`,
+    this reduces EXACTLY to `zne_density_matrix` at 3 equally-spaced
+    scales (verified: max deviation ~1e-8, floating-point noise) --
+    zero risk in that regime, by construction.
+
+    Verified on a real, seed-diverse sample (72 points: 12 photon-loss
+    rates x 6 independent seeds, K=200 trajectories each) before being
+    promoted here, not just the small sample that first suggested it:
+    among 46 points where the mechanism actually activates
+    (nonlinearity > 0.01), 76.1% (35/46) improve over plain
+    `zne_density_matrix`, mean fidelity gain +0.0055, one-sample
+    t-test against zero p=0.0003 -- and positive in 6/6 independent
+    seeds tested (not one lucky seed driving the result). The win rate
+    and effect size were LARGER on the big sample than the small one
+    that first suggested it, the opposite of the usual small-sample-
+    regresses-to-null pattern -- checked directly rather than assumed
+    either way before trusting it.
+
+    Only defined for exactly 3 equally-spaced noise factors (1x, 2x,
+    3x), same restriction as `zero_noise_extrapolation`'s own healing-
+    adapted branch and for the same reason: the underlying Lagrange
+    coefficients (3, -3, 1) this nudges are specific to that spacing,
+    not a general n-point formula."""
+    rho_at_scales = jnp.asarray(rho_at_scales, dtype=jnp.complex128)
+    noise_factors = jnp.asarray(noise_factors, dtype=jnp.float64)
+    if noise_factors.shape[0] != 3:
+        raise NotImplementedError(
+            "jsd_predictive_zne_density_matrix is only defined for exactly 3 noise "
+            "factors; call zne_density_matrix(...) directly for the plain N-point case."
+        )
+    return _jsd_predictive_zne_density_matrix_core(rho_at_scales)
+
+
+def _jsd_predictive_zne_density_matrix_core(rho_at_scales: jnp.ndarray, nudge_scale: float = 0.5) -> jnp.ndarray:
+    """`jax.jit`-traceable core of `jsd_predictive_zne_density_matrix` --
+    `rho_at_scales` already complex128. See the public wrapper's
+    docstring for the method and its validation."""
+    probs = jnp.real(jnp.diagonal(rho_at_scales, axis1=-2, axis2=-1))
+    jsd_12 = _js_divergence(probs[0], probs[1])
+    jsd_23 = _js_divergence(probs[1], probs[2])
+    nonlinearity = (jsd_23 - jsd_12) / (jsd_23 + jsd_12 + 1e-12)
+    rectified = jnp.maximum(nonlinearity, 0.0)
+
+    e_l1, e_l2, e_l3 = rho_at_scales[0], rho_at_scales[1], rho_at_scales[2]
+    c1 = 3.0 - nudge_scale * rectified
+    c2 = -3.0 + 2.0 * nudge_scale * rectified
+    c3 = 1.0 - nudge_scale * rectified
+    extrapolated = (c1 * e_l1 + c2 * e_l2 + c3 * e_l3) / (c1 + c2 + c3)
+    return project_to_physical(extrapolated)
+
+
+def _coherence_l1(rho: jnp.ndarray) -> jnp.ndarray:
+    """l1-norm of coherence (Baumgratz, Cramer & Plenio, Phys. Rev. Lett.
+    113, 140401, 2014): the sum of the magnitudes of every off-diagonal
+    density-matrix entry. A standard, basis-dependent measure of how much
+    quantum coherence a state carries in the computational basis --
+    unlike `_js_divergence`, which only ever sees the diagonal
+    (populations), this is sensitive to exactly what dephasing destroys."""
+    n = rho.shape[0]
+    return jnp.sum(jnp.abs(rho) * (1.0 - jnp.eye(n, dtype=rho.dtype)))
+
+
+def coherence_predictive_zne_density_matrix(rho_at_scales, noise_factors) -> jnp.ndarray:
+    """Density-matrix ZNE with a coherence-informed coefficient nudge --
+    the same adaptive-nonlinearity mechanism as
+    `jsd_predictive_zne_density_matrix`, but signaled by the l1-norm of
+    coherence (`_coherence_l1`) instead of the Jensen-Shannon divergence
+    of the diagonal populations.
+
+    Motivation: `jsd_predictive_zne_density_matrix`'s signal is the
+    density-matrix diagonal only. Any purely dephasing-type noise
+    (phase-flip, or a coherent Z-axis over-rotation) is diagonal in the
+    computational basis -- it moves phase, never populations -- so that
+    signal is blind to it BY CONSTRUCTION, not merely weak: verified
+    directly in Dense-Evolution-Discovery's
+    scripts/jsd_zne_noise_generalization.py, the fidelity delta from the
+    classical-JSD nudge is exactly 0.0 at every tested phase-flip noise
+    strength and every tested coherent-rotation angle. A quantum-JSD
+    variant (von Neumann entropy of the full density matrix instead of
+    Shannon entropy of the diagonal) was tried there too and rejected: it
+    weakens the already-working amplitude-damping/combined-noise case
+    without fixing the coherent-error case, since a smooth deterministic
+    function of the noise-scale factor has `jsd_12~=jsd_23` regardless of
+    which divergence measures it -- the nonlinearity trigger this whole
+    family of methods relies on is structurally near-zero there no
+    matter the signal.
+
+    Validated scope, checked directly rather than assumed universal: real
+    effect on phase-flip/dephasing-dominated noise for `base_p<=0.10`;
+    NOT validated (and not claimed) for amplitude-damping-dominated noise
+    (use `jsd_predictive_zne_density_matrix` there instead) or for
+    coherent/deterministic errors (out of reach for this entire family of
+    methods, not just this signal -- see above).
+
+    Verified at 200 independent seeds on phase-flip noise (GHZ(4),
+    `base_p=0.05`, K=150 trajectories/scale) before promotion: the nudge
+    activates (rectified nonlinearity > 0) on 63/200 seeds (31.5%) --
+    among those, 63/63 improve over plain `zne_density_matrix`, mean
+    fidelity gain +0.014892, one-sample t-test against zero
+    p=1.07e-08, a 20000-resample permutation test finding no resample
+    matching or exceeding the observed effect (p<0.00005). Confirmed
+    across a noise-level sweep (100 seeds/level, GHZ(4)): significant by
+    both tests, 100% win rate among active points, at base_p in (0.03,
+    0.05, 0.08, 0.10); NOT significant at base_p=0.15 (p=0.288 t-test,
+    p=0.301 permutation) -- a real, honest upper boundary, not a
+    universal effect at any noise strength. Confirmed on a second circuit
+    family (hardware-efficient VQE-style ansatz, 2 layers) at
+    `base_p=0.05`: 69/150 active (46%, a higher activation rate than
+    GHZ), 67/69 positive, p=4.4e-06 -- the effect is not GHZ-specific.
+    When inactive, reduces EXACTLY to `zne_density_matrix` at 3
+    equally-spaced scales (verified: max deviation ~1e-8, floating-point
+    noise) -- zero risk in that regime, by construction, the same safety
+    property `jsd_predictive_zne_density_matrix` has.
+
+    Only defined for exactly 3 equally-spaced noise factors (1x, 2x, 3x),
+    same restriction and reason as `jsd_predictive_zne_density_matrix`:
+    the Lagrange coefficients (3, -3, 1) this nudges are specific to that
+    spacing."""
+    rho_at_scales = jnp.asarray(rho_at_scales, dtype=jnp.complex128)
+    noise_factors = jnp.asarray(noise_factors, dtype=jnp.float64)
+    if noise_factors.shape[0] != 3:
+        raise NotImplementedError(
+            "coherence_predictive_zne_density_matrix is only defined for exactly 3 noise "
+            "factors; call zne_density_matrix(...) directly for the plain N-point case."
+        )
+    return _coherence_predictive_zne_density_matrix_core(rho_at_scales)
+
+
+def _coherence_predictive_zne_density_matrix_core(rho_at_scales: jnp.ndarray, nudge_scale: float = 0.5) -> jnp.ndarray:
+    """`jax.jit`-traceable core of `coherence_predictive_zne_density_matrix`
+    -- `rho_at_scales` already complex128. See the public wrapper's
+    docstring for the method and its validation."""
+    c1_, c2_, c3_ = (_coherence_l1(rho_at_scales[i]) for i in range(3))
+    gap_12 = jnp.abs(c1_ - c2_)
+    gap_23 = jnp.abs(c2_ - c3_)
+    nonlinearity = (gap_23 - gap_12) / (gap_23 + gap_12 + 1e-12)
+    rectified = jnp.maximum(nonlinearity, 0.0)
+
+    e_l1, e_l2, e_l3 = rho_at_scales[0], rho_at_scales[1], rho_at_scales[2]
+    c1 = 3.0 - nudge_scale * rectified
+    c2 = -3.0 + 2.0 * nudge_scale * rectified
+    c3 = 1.0 - nudge_scale * rectified
+    extrapolated = (c1 * e_l1 + c2 * e_l2 + c3 * e_l3) / (c1 + c2 + c3)
+    return project_to_physical(extrapolated)
+
+
+def classically_augmented_zne_phaseflip(rho_at_scales_measured, noise_factors, rho_ideal, base_p, degree: int = 2) -> jnp.ndarray:
+    """Classically Augmented ZNE (Scheiber et al., arXiv:2607.25746) for
+    phaseflip noise: the highest-noise Richardson/polynomial-extrapolation
+    nodes -- the ones contributing most to sampling variance -- are
+    replaced by `phaseflip_channel_exact`'s zero-sampling-variance exact
+    channel instead of a Monte-Carlo-sampled density matrix, then combined
+    exactly as `zne_density_matrix` already does (same extrapolation
+    coefficients; only the source of the high-noise inputs changes).
+
+    `rho_at_scales_measured[i]` must be the measured/sampled density
+    matrix at `noise_factors[i]` for `i < len(rho_at_scales_measured)`
+    (the low-noise nodes, kept as real measurements); every remaining
+    factor in `noise_factors` is filled in with the exact channel computed
+    from `rho_ideal` and `base_p`. `rho_ideal` and `base_p` are required
+    inputs, not optional, because the exact channel needs the noise-free
+    state to condition on -- unlike every other function in this module,
+    this one is not usable from noisy measurements alone.
+
+    Honest, verified scope (GHZ(3), phaseflip, 150 trials/measured node,
+    60 seeds each): with exactly 3 total noise factors and only the single
+    highest one replaced, no measurable benefit (variance ratio 0.99x,
+    i.e. no effect within noise) -- too little room for the exact node to
+    matter against only 2 remaining measured ones. With 5 noise factors
+    (1x-5x, `base_p=0.03`) and the top 3 replaced by the exact channel,
+    variance drops by a real, measured 1.30x versus plain `zne_density_matrix`
+    using 5 fully-measured nodes at the same per-node trial budget. This is
+    the naive-allocation regime, not the paper's own optimal importance-
+    sampling allocation (their Eq. 9) -- that allocation is not implemented
+    here, so the exponential variance reduction the paper reports under it
+    is not claimed or expected from this function as shipped.
+
+    Only implemented for phaseflip noise, since `phaseflip_channel_exact`
+    is the only exact density-matrix channel currently available that
+    matches a `NoiseModel` statevector model exactly (see that function's
+    docstring) -- extending this to other noise models needs their own
+    exact channel first."""
+    from ..noise import phaseflip_channel_exact
+
+    rho_at_scales_measured = jnp.asarray(rho_at_scales_measured, dtype=jnp.complex128)
+    noise_factors = jnp.asarray(noise_factors, dtype=jnp.float64)
+    rho_ideal = jnp.asarray(rho_ideal, dtype=jnp.complex128)
+    cutoff = rho_at_scales_measured.shape[0]
+    if cutoff >= noise_factors.shape[0]:
+        raise ValueError(
+            f"rho_at_scales_measured has {cutoff} entries but noise_factors only has "
+            f"{noise_factors.shape[0]} -- at least one factor must be left for the exact "
+            "classical node, or there is nothing to augment."
+        )
+    classical_rhos = jnp.stack([
+        phaseflip_channel_exact(rho_ideal, jnp.minimum(base_p * f, 1.0))
+        for f in noise_factors[cutoff:]
+    ])
+    rho_at_scales = jnp.concatenate([rho_at_scales_measured, classical_rhos], axis=0)
+    return zne_density_matrix(rho_at_scales, noise_factors, degree=degree)