Blame for dense_evolution/mitigation/zne.py

"""

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 numpy as np
import jax
import jax.numpy as jnp

from .healing import calculate_delta_preemp

__all__ = ["richardson_extrapolate", "zero_noise_extrapolation", "polynomial_extrapolate",
           "project_to_physical", "uhlmann_fidelity", "zne_density_matrix",
           "jsd_predictive_zne_density_matrix",
           "richardson_extrapolate_jit", "zero_noise_extrapolation_jit",
           "polynomial_extrapolate_jit", "uhlmann_fidelity_jit", "zne_density_matrix_jit"]


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).

    """
    lambdas = jnp.asarray(noise_factors, dtype=jnp.float64)
    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 _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]

    def lagrange_coeff(i):
        others = jnp.concatenate([lambdas[:i], lambdas[i + 1:]])
        return jnp.prod((0.0 - others) / (lambdas[i] - others))

    coeffs = jnp.stack([lagrange_coeff(i) for i in range(n)])
    # 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 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.

    """
    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))


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)