dense-Evolution / test_mitigation.py
Tatopenn's picture
Sync v8.1.40: full jax.jit coverage for dense_evolution.mitigation
7ce2fb5 verified
Raw
History Blame Contribute Delete
21.1 kB
import numpy as np
import jax.numpy as jnp
import pytest
import dense_evolution as de
from dense_evolution.mitigation import (
richardson_extrapolate, zero_noise_extrapolation, polynomial_extrapolate,
project_to_physical, uhlmann_fidelity, zne_density_matrix, zne_density_matrix_jit,
richardson_extrapolate_jit, zero_noise_extrapolation_jit, uhlmann_fidelity_jit,
polynomial_extrapolate_jit,
)
def test_richardson_extrapolate_matches_known_3point_coefficients():
rng = np.random.default_rng(0)
for _ in range(20):
e1, e2, e3 = rng.normal(size=3)
expected = 3.0 * e1 - 3.0 * e2 + 1.0 * e3
got = float(richardson_extrapolate([e1, e2, e3], [1.0, 2.0, 3.0]))
assert got == pytest.approx(expected, abs=1e-9)
def test_richardson_extrapolate_exact_on_linear_data():
# a Richardson/Lagrange extrapolation is exact for any polynomial of
# degree < n_points; for 3 points a linear signal must extrapolate
# to exactly the intercept.
a, b = 5.3, -2.1
lambdas = [1.0, 2.0, 3.0]
values = [a + b * l for l in lambdas]
got = float(richardson_extrapolate(values, lambdas))
assert got == pytest.approx(a, abs=1e-9)
def test_richardson_extrapolate_supports_vector_valued_expectation_values():
# expectation_values[i] doesn't have to be a scalar -- e.g. a full
# probability distribution sampled at noise scale i. Found via a real
# broadcasting bug: coeffs (shape (n,)) times a stacked (n, d) array
# relies on jnp's default trailing-axis alignment, which pairs (n,)
# against d, not n -- fails outright unless d happens to equal n.
rng = np.random.default_rng(1)
v1, v2, v3 = rng.normal(size=4), rng.normal(size=4), rng.normal(size=4)
got = np.asarray(richardson_extrapolate([v1, v2, v3], [1.0, 2.0, 3.0]))
expected = 3.0 * v1 - 3.0 * v2 + 1.0 * v3
np.testing.assert_allclose(got, expected, atol=1e-9)
assert got.shape == (4,)
def test_richardson_extrapolate_preserves_complex_input():
# Regression test: expectation_values used to be forced to jnp.float64
# unconditionally, silently discarding the imaginary part of complex
# input (e.g. density-matrix entries) with only a low-signal
# ComplexWarning. Hand-computed via the textbook (3, -3, 1) Lagrange
# coefficients on real and imaginary parts separately.
a, b, c = 1 + 2j, 3 + 4j, 3 + 4j
d, e, f = 2 + 1j, 6 + 8j, 6 + 8j
got = np.asarray(richardson_extrapolate([[a, d], [b, e], [c, f]], [1.0, 2.0, 3.0]))
expected = np.array([
3.0 * a - 3.0 * b + 1.0 * c,
3.0 * d - 3.0 * e + 1.0 * f,
])
np.testing.assert_allclose(got, expected, atol=1e-9)
assert np.iscomplexobj(got)
assert np.any(np.imag(got) != 0.0) or np.any(np.imag(expected) != 0.0)
def test_richardson_extrapolate_real_input_stays_real():
got = richardson_extrapolate([1.0, 2.0, 3.0], [1.0, 2.0, 3.0])
assert not np.iscomplexobj(np.asarray(got))
def test_zero_noise_extrapolation_healing_branch_preserves_complex_input():
e1, e2, e3 = 1 + 2j, 2 + 1j, 3 + 0.5j
delta_p = abs(9.0 - 10.0) / 10.0
c1, c2, c3 = 3.0 - 0.01 * delta_p, -3.0 + 0.02 * delta_p, 1.0 - 0.01 * delta_p
expected = (c1 * e1 + c2 * e2 + c3 * e3) / (c1 + c2 + c3)
got = complex(zero_noise_extrapolation([e1, e2, e3], [1.0, 2.0, 3.0],
sigma_at_base_noise=9.0,
target_sigma_ideal=10.0))
assert got == pytest.approx(expected, abs=1e-9)
assert got.imag != 0.0
def test_zero_noise_extrapolation_without_sigma_matches_richardson_extrapolate():
values, lambdas = [1.234, 0.876, 0.611], [1.0, 2.0, 3.0]
plain = float(richardson_extrapolate(values, lambdas))
orchestrated = float(zero_noise_extrapolation(values, lambdas))
assert orchestrated == pytest.approx(plain, abs=1e-12)
def test_zero_noise_extrapolation_with_sigma_matches_reference_healing_formula():
# reference formula, promoted verbatim from
# Dense-Evolution-Ising-Tests/tests/test_zne_predictive_healing.py
# (_adaptive_healing_richardson), independently re-derived here.
def reference(e_l1, e_l2, e_l3, delta_p):
c1, c2, c3 = 3.0 - 0.01 * delta_p, -3.0 + 0.02 * delta_p, 1.0 - 0.01 * delta_p
return (c1 * e_l1 + c2 * e_l2 + c3 * e_l3) / (c1 + c2 + c3)
target_sigma = 10.0
for sigma, (e1, e2, e3) in [
(9.5, (1.0, 0.8, 0.6)),
(7.0, (2.3, 1.9, 1.5)),
(10.0, (-0.4, -0.5, -0.6)),
]:
delta_p = abs(sigma - target_sigma) / target_sigma
expected = reference(e1, e2, e3, delta_p)
got = float(zero_noise_extrapolation([e1, e2, e3], [1.0, 2.0, 3.0],
sigma_at_base_noise=sigma,
target_sigma_ideal=target_sigma))
assert got == pytest.approx(expected, abs=1e-9)
def test_zero_noise_extrapolation_rejects_non_3point_healing_request():
with pytest.raises(NotImplementedError):
zero_noise_extrapolation([1.0, 2.0, 3.0, 4.0], [1.0, 2.0, 3.0, 4.0],
sigma_at_base_noise=9.0)
def test_exported_from_package_root():
assert de.richardson_extrapolate is richardson_extrapolate
assert de.zero_noise_extrapolation is zero_noise_extrapolation
assert de.polynomial_extrapolate is polynomial_extrapolate
assert de.project_to_physical is project_to_physical
assert de.uhlmann_fidelity is uhlmann_fidelity
assert de.zne_density_matrix is zne_density_matrix
assert de.zne_density_matrix_jit is zne_density_matrix_jit
assert de.richardson_extrapolate_jit is richardson_extrapolate_jit
assert de.zero_noise_extrapolation_jit is zero_noise_extrapolation_jit
assert de.uhlmann_fidelity_jit is uhlmann_fidelity_jit
def test_polynomial_extrapolate_matches_richardson_at_exact_point_count():
# degree = n_points - 1 means the least-squares fit is exactly
# determined -- must equal the unique interpolating polynomial, i.e.
# richardson_extrapolate, to numerical precision.
rng = np.random.default_rng(10)
for n in (3, 4, 5):
values = rng.normal(size=n) + 1j * rng.normal(size=n)
lambdas = np.arange(1, n + 1, dtype=float)
expected = np.asarray(richardson_extrapolate(values.tolist(), lambdas.tolist()))
got = np.asarray(polynomial_extrapolate(values.tolist(), lambdas.tolist(), degree=n - 1))
np.testing.assert_allclose(got, expected, atol=1e-8)
def test_polynomial_extrapolate_exact_on_matching_degree_polynomial():
rng = np.random.default_rng(11)
coeffs = rng.normal(size=3) # degree-2 polynomial
lambdas = np.array([1.0, 2.0, 3.0, 4.0, 5.0]) # 5 points, overdetermined
values = sum(c * lambdas ** k for k, c in enumerate(coeffs))
expected_intercept = coeffs[0]
got = float(polynomial_extrapolate(values.tolist(), lambdas.tolist(), degree=2))
assert got == pytest.approx(expected_intercept, abs=1e-6)
def test_polynomial_extrapolate_rejects_underdetermined_fit():
with pytest.raises(ValueError):
polynomial_extrapolate([1.0, 2.0], [1.0, 2.0], degree=2)
def test_polynomial_extrapolate_preserves_complex_input():
values = [1 + 2j, 2 + 1j, 3 + 0.5j, 4 - 1j, 5 - 2j]
got = polynomial_extrapolate(values, [1.0, 2.0, 3.0, 4.0, 5.0], degree=2)
assert np.iscomplexobj(np.asarray(got))
assert complex(got).imag != 0.0
def _random_density_matrix(rng, d):
a = rng.normal(size=(d, d)) + 1j * rng.normal(size=(d, d))
rho = a @ a.conj().T
return rho / np.trace(rho)
def test_project_to_physical_matches_paper_worked_example():
# Smolin, Gambetta & Smith (2012), arXiv:1106.5458, Fig. 1: starting
# eigenvalues 3/5, 1/2, 7/20, 1/10, -11/20 (trace 1, one negative)
# project to 9/20, 7/20, 1/5, 0, 0 exactly.
eigvals_in = np.array([3 / 5, 1 / 2, 7 / 20, 1 / 10, -11 / 20])
assert eigvals_in.sum() == pytest.approx(1.0, abs=1e-12)
rho_raw = jnp.asarray(np.diag(eigvals_in), dtype=jnp.complex128)
got = np.asarray(project_to_physical(rho_raw))
got_eigvals = np.sort(np.linalg.eigvalsh(got))[::-1]
expected_eigvals = np.array([9 / 20, 7 / 20, 1 / 5, 0.0, 0.0])
np.testing.assert_allclose(got_eigvals, expected_eigvals, atol=1e-9)
def _random_traceless_hermitian(rng, d):
a = rng.normal(size=(d, d)) + 1j * rng.normal(size=(d, d))
h = (a + a.conj().T) / 2
return h - (np.trace(h).real / d) * np.eye(d)
def test_project_to_physical_output_is_a_valid_density_matrix():
rng = np.random.default_rng(2)
for _ in range(10):
d = rng.integers(2, 6)
rho = _random_density_matrix(rng, d)
# perturb with a traceless Hermitian matrix so trace stays exactly
# 1 while pushing eigenvalues negative (likely, for this scale).
rho_raw = jnp.asarray(rho + 0.5 * _random_traceless_hermitian(rng, d), dtype=jnp.complex128)
got = np.asarray(project_to_physical(rho_raw))
np.testing.assert_allclose(got, got.conj().T, atol=1e-9)
assert np.trace(got).real == pytest.approx(1.0, abs=1e-9)
assert np.trace(got).imag == pytest.approx(0.0, abs=1e-9)
assert np.linalg.eigvalsh(got).min() >= -1e-9
def test_project_to_physical_is_a_near_no_op_on_already_physical_input():
rng = np.random.default_rng(3)
rho = jnp.asarray(_random_density_matrix(rng, 4), dtype=jnp.complex128)
got = project_to_physical(rho)
np.testing.assert_allclose(np.asarray(got), np.asarray(rho), atol=1e-9)
def test_uhlmann_fidelity_self_is_one():
rng = np.random.default_rng(4)
rho = jnp.asarray(_random_density_matrix(rng, 3), dtype=jnp.complex128)
assert uhlmann_fidelity(rho, rho) == pytest.approx(1.0, abs=1e-9)
def test_uhlmann_fidelity_matches_pure_state_overlap():
rng = np.random.default_rng(5)
for _ in range(5):
psi_a = rng.normal(size=4) + 1j * rng.normal(size=4)
psi_a /= np.linalg.norm(psi_a)
psi_b = rng.normal(size=4) + 1j * rng.normal(size=4)
psi_b /= np.linalg.norm(psi_b)
rho_a = jnp.asarray(np.outer(psi_a, psi_a.conj()), dtype=jnp.complex128)
rho_b = jnp.asarray(np.outer(psi_b, psi_b.conj()), dtype=jnp.complex128)
expected = abs(np.vdot(psi_a, psi_b)) ** 2
got = uhlmann_fidelity(rho_a, rho_b)
assert got == pytest.approx(expected, abs=1e-6)
def test_uhlmann_fidelity_is_symmetric():
rng = np.random.default_rng(6)
rho_a = jnp.asarray(_random_density_matrix(rng, 3), dtype=jnp.complex128)
rho_b = jnp.asarray(_random_density_matrix(rng, 3), dtype=jnp.complex128)
assert uhlmann_fidelity(rho_a, rho_b) == pytest.approx(uhlmann_fidelity(rho_b, rho_a), abs=1e-9)
def test_zne_density_matrix_output_is_a_valid_density_matrix():
rng = np.random.default_rng(7)
d = 4
rho_at_scales = jnp.stack([
jnp.asarray(_random_density_matrix(rng, d), dtype=jnp.complex128) for _ in range(3)
])
got = np.asarray(zne_density_matrix(rho_at_scales, [1.0, 2.0, 3.0]))
np.testing.assert_allclose(got, got.conj().T, atol=1e-9)
assert np.trace(got).real == pytest.approx(1.0, abs=1e-9)
assert np.linalg.eigvalsh(got).min() >= -1e-9
def test_zne_density_matrix_is_exactly_richardson_then_projection():
rng = np.random.default_rng(8)
d = 3
rho_at_scales = jnp.stack([
jnp.asarray(_random_density_matrix(rng, d), dtype=jnp.complex128) for _ in range(3)
])
noise_factors = [1.0, 2.0, 3.0]
expected = project_to_physical(richardson_extrapolate(rho_at_scales, noise_factors))
got = zne_density_matrix(rho_at_scales, noise_factors)
np.testing.assert_allclose(np.asarray(got), np.asarray(expected), atol=1e-12)
def test_zne_density_matrix_accepts_more_than_three_scales():
# Unlike richardson_extrapolate's exact interpolation (which gets
# measurably worse with more points, see polynomial_extrapolate's
# docstring), zne_density_matrix's degree=2 least-squares default
# should stay well-behaved -- still a valid density matrix -- when
# given more than 3 noise-scale points.
rng = np.random.default_rng(9)
d = 3
rho_at_scales = jnp.stack([
jnp.asarray(_random_density_matrix(rng, d), dtype=jnp.complex128) for _ in range(6)
])
got = np.asarray(zne_density_matrix(rho_at_scales, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]))
np.testing.assert_allclose(got, got.conj().T, atol=1e-9)
assert np.trace(got).real == pytest.approx(1.0, abs=1e-9)
assert np.linalg.eigvalsh(got).min() >= -1e-9
def test_zne_density_matrix_preserves_complex_off_diagonal_entries():
# A minimal, hand-built 3-scale complex input where the correct
# zero-noise extrapolation has a nonzero imaginary part -- guards
# against the exact bug fixed in richardson_extrapolate resurfacing
# silently through this composed entry point.
base = np.array([[0.6, 0.1 + 0.2j], [0.1 - 0.2j, 0.4]])
rho_at_scales = jnp.stack([
jnp.asarray(base * (1.0 + 0.1 * s), dtype=jnp.complex128) for s in (0, 1, 2)
])
got = np.asarray(zne_density_matrix(rho_at_scales, [1.0, 2.0, 3.0]))
assert np.any(np.abs(got.imag) > 1e-9)
def test_zne_density_matrix_jit_matches_eager_exactly():
rng = np.random.default_rng(20)
d = 4
mats = []
for _ in range(3):
a = rng.normal(size=(d, d)) + 1j * rng.normal(size=(d, d))
m = a @ a.conj().T
mats.append(m / np.trace(m))
rho_at_scales = jnp.asarray(np.stack(mats), dtype=jnp.complex128)
noise_factors = jnp.asarray([1.0, 2.0, 3.0], dtype=jnp.float64)
eager = zne_density_matrix(rho_at_scales, [1.0, 2.0, 3.0], degree=2)
jitted = zne_density_matrix_jit(rho_at_scales, noise_factors, degree=2)
np.testing.assert_array_equal(np.asarray(eager), np.asarray(jitted))
def test_zne_density_matrix_jit_output_is_a_valid_density_matrix():
rng = np.random.default_rng(21)
d = 5
mats = []
for _ in range(5):
a = rng.normal(size=(d, d)) + 1j * rng.normal(size=(d, d))
m = a @ a.conj().T
mats.append(m / np.trace(m))
rho_at_scales = jnp.asarray(np.stack(mats), dtype=jnp.complex128)
noise_factors = jnp.asarray([1.0, 2.0, 3.0, 4.0, 5.0], dtype=jnp.float64)
got = np.asarray(zne_density_matrix_jit(rho_at_scales, noise_factors, degree=2))
np.testing.assert_allclose(got, got.conj().T, atol=1e-9)
assert np.trace(got).real == pytest.approx(1.0, abs=1e-9)
assert np.linalg.eigvalsh(got).min() >= -1e-9
def test_zne_density_matrix_jit_actually_compiles_under_jit():
# zne_density_matrix_jit is already jax.jit-wrapped; this specifically
# checks it doesn't raise a tracing error (e.g. from a stray
# np.iscomplexobj/np.asarray call on a traced value) when called
# through an *additional* outer jax.jit, the realistic case of
# embedding it inside a larger jitted pipeline (e.g. jax.lax.scan).
import jax
rng = np.random.default_rng(22)
d = 3
mats = []
for _ in range(3):
a = rng.normal(size=(d, d)) + 1j * rng.normal(size=(d, d))
m = a @ a.conj().T
mats.append(m / np.trace(m))
rho_at_scales = jnp.asarray(np.stack(mats), dtype=jnp.complex128)
noise_factors = jnp.asarray([1.0, 2.0, 3.0], dtype=jnp.float64)
from dense_evolution.mitigation import _zne_density_matrix_core
import functools
outer_jit = jax.jit(functools.partial(_zne_density_matrix_core, degree=2))
result = outer_jit(rho_at_scales, noise_factors)
result.block_until_ready()
assert np.trace(np.asarray(result)).real == pytest.approx(1.0, abs=1e-9)
def test_polynomial_extrapolate_jit_matches_eager():
lambdas = jnp.asarray([1.0, 2.0, 3.0, 4.0, 5.0], dtype=jnp.float64)
values = jnp.asarray([1 + 2j, 3 + 4j, 3 + 4j, 2 - 1j, 5 + 0.5j], dtype=jnp.complex128)
got = polynomial_extrapolate_jit(values, lambdas, degree=2)
expected = polynomial_extrapolate(
[1 + 2j, 3 + 4j, 3 + 4j, 2 - 1j, 5 + 0.5j], [1.0, 2.0, 3.0, 4.0, 5.0], degree=2)
np.testing.assert_allclose(np.asarray(got), np.asarray(expected))
def test_polynomial_extrapolate_jit_compiles_under_outer_jit():
import jax
import functools
lambdas = jnp.asarray([1.0, 2.0, 3.0], dtype=jnp.float64)
values = jnp.asarray([1.0, 2.0, 3.0], dtype=jnp.float64)
outer_jit = jax.jit(functools.partial(polynomial_extrapolate_jit, degree=2))
result = outer_jit(values, lambdas)
result.block_until_ready()
assert float(result) == pytest.approx(0.0, abs=1e-9) # linear data -> exact intercept 0
def test_richardson_extrapolate_jit_matches_eager():
lambdas = jnp.asarray([1.0, 2.0, 3.0], dtype=jnp.float64)
values = jnp.asarray([1 + 2j, 3 + 4j, 3 + 4j], dtype=jnp.complex128)
got = richardson_extrapolate_jit(values, lambdas)
expected = richardson_extrapolate([1 + 2j, 3 + 4j, 3 + 4j], [1.0, 2.0, 3.0])
np.testing.assert_allclose(np.asarray(got), np.asarray(expected))
def test_richardson_extrapolate_jit_compiles_under_outer_jit():
import jax
lambdas = jnp.asarray([1.0, 2.0, 3.0], dtype=jnp.float64)
values = jnp.asarray([1.0, 2.0, 3.0], dtype=jnp.float64)
outer_jit = jax.jit(richardson_extrapolate_jit)
result = outer_jit(values, lambdas)
result.block_until_ready()
assert float(result) == pytest.approx(0.0, abs=1e-9) # linear data -> exact intercept 0
def test_zero_noise_extrapolation_jit_matches_eager():
values = jnp.asarray([1 + 2j, 3 + 4j, 3 + 4j], dtype=jnp.complex128)
sigma = jnp.asarray(7.5, dtype=jnp.float64)
got = zero_noise_extrapolation_jit(values, sigma, 10.0)
expected = zero_noise_extrapolation([1 + 2j, 3 + 4j, 3 + 4j], [1.0, 2.0, 3.0],
sigma_at_base_noise=7.5, target_sigma_ideal=10.0)
np.testing.assert_allclose(np.asarray(got), np.asarray(expected))
def test_zero_noise_extrapolation_jit_target_sigma_ideal_stays_dynamic():
# target_sigma_ideal must NOT need to be static -- calculate_delta_preemp
# uses jnp.where internally, not a Python if, so it's trace-safe as a
# plain traced float. Regression guard: if this ever needs
# static_argnames, this test starts raising a tracing error (the
# primary thing being checked -- no exception on a second call with a
# different target, no recompilation required). A wide target gap and
# non-degenerate complex values also confirm the coefficients (and so
# the result) actually do move, not just "didn't crash".
import jax
values = jnp.asarray([1 + 2j, 3 + 4j, 3 + 4j], dtype=jnp.complex128)
sigma = jnp.asarray(0.5, dtype=jnp.float64)
@jax.jit
def f(values, sigma, target):
return zero_noise_extrapolation_jit(values, sigma, target)
r1 = f(values, sigma, 10.0)
r2 = f(values, sigma, 1.0) # different target, same compiled function
assert not jnp.allclose(r1, r2)
def test_uhlmann_fidelity_jit_matches_eager():
rng = np.random.default_rng(23)
d = 3
a = rng.normal(size=(d, d)) + 1j * rng.normal(size=(d, d))
rho_a = jnp.asarray(a @ a.conj().T / np.trace(a @ a.conj().T), dtype=jnp.complex128)
b = rng.normal(size=(d, d)) + 1j * rng.normal(size=(d, d))
rho_b = jnp.asarray(b @ b.conj().T / np.trace(b @ b.conj().T), dtype=jnp.complex128)
got = uhlmann_fidelity_jit(rho_a, rho_b)
expected = uhlmann_fidelity(rho_a, rho_b)
assert float(got) == pytest.approx(expected, abs=1e-9)
def test_full_pipeline_composes_under_a_single_outer_jax_jit():
# The realistic case this whole _jit family exists for: several of
# these functions called together inside ONE outer jax.jit (e.g. a
# step function passed to jax.lax.scan), not each jitted in isolation.
import jax
from dense_evolution.mitigation import (
_zero_noise_extrapolation_healing_core, _zne_density_matrix_core, _uhlmann_fidelity_core,
)
rng = np.random.default_rng(24)
d = 3
mats = []
for _ in range(3):
a = rng.normal(size=(d, d)) + 1j * rng.normal(size=(d, d))
m = a @ a.conj().T
mats.append(m / np.trace(m))
rho_at_scales = jnp.asarray(np.stack(mats), dtype=jnp.complex128)
noise_factors = jnp.asarray([1.0, 2.0, 3.0], dtype=jnp.float64)
sigma = jnp.asarray(7.5, dtype=jnp.float64)
rho_target = rho_at_scales[0]
@jax.jit
def full_pipeline(rho_at_scales, noise_factors, sigma, target_sigma, rho_target):
healed = _zero_noise_extrapolation_healing_core(rho_at_scales, sigma, target_sigma)
corrected = _zne_density_matrix_core(rho_at_scales, noise_factors, 2)
fidelity = _uhlmann_fidelity_core(corrected, rho_target)
return healed, corrected, fidelity
healed, corrected, fidelity = full_pipeline(rho_at_scales, noise_factors, sigma, 10.0, rho_target)
jax.block_until_ready((healed, corrected, fidelity))
np.testing.assert_allclose(np.asarray(corrected), np.asarray(corrected).conj().T, atol=1e-9)
assert np.trace(np.asarray(corrected)).real == pytest.approx(1.0, abs=1e-9)
assert 0.0 <= float(fidelity) <= 1.0 + 1e-9