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