| 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, 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(): |
| |
| |
| |
| |
| |
| 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(): |
| |
| |
| |
| |
| |
| 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(): |
| |
| |
| |
| 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(): |
| |
| |
| |
| 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) |
| lambdas = np.array([1.0, 2.0, 3.0, 4.0, 5.0]) |
| 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(): |
| |
| |
| |
| 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) |
| |
| |
| 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(): |
| |
| |
| |
| |
| |
| 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(): |
| |
| |
| |
| |
| 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(): |
| |
| |
| |
| |
| |
| 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) |
|
|
|
|
| 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) |
|
|
|
|
| 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(): |
| |
| |
| |
| |
| |
| |
| |
| |
| 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) |
| 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(): |
| |
| |
| |
| 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 |
|
|