|
1 | 1 | """ |
2 | 2 | Tests for ZNE Mitigator |
3 | 3 | ======================== |
4 | | -Validates ZNE correctness on noiseless simulator: |
5 | | - - Gate folding produces correct circuit structure |
6 | | - - Richardson extrapolation converges to exact value |
7 | | - - Overhead matches expected noise level count |
| 4 | +Validates ZNE correctness on noiseless and noisy simulators: |
| 5 | + - Gate folding produces correct circuit structure (gate count) |
| 6 | + - Richardson extrapolation converges to ideal value on noiseless sim |
| 7 | + - ZNE reduces error vs raw noisy on a simulated noisy backend |
| 8 | + - Overhead and name metadata are correct |
8 | 9 | """ |
9 | 10 |
|
10 | 11 | import pytest |
11 | | -from mitigation.zne import ZNEMitigator |
| 12 | +from qiskit_aer import AerSimulator |
12 | 13 |
|
| 14 | +from benchmarks.circuits import build_noise_model, create_test_circuit |
| 15 | +from mitigation.zne import ( |
| 16 | + ZNEMitigator, |
| 17 | + fold_gates, |
| 18 | + measure_expectation, |
| 19 | + richardson_extrapolate, |
| 20 | +) |
13 | 21 |
|
14 | | -class TestZNEMitigator: |
15 | | - """Test suite for Zero-Noise Extrapolation.""" |
16 | 22 |
|
| 23 | +# --------------------------------------------------------------------------- |
| 24 | +# Fixtures |
| 25 | +# --------------------------------------------------------------------------- |
| 26 | + |
| 27 | + |
| 28 | +@pytest.fixture(scope="module") |
| 29 | +def test_circuit(): |
| 30 | + """2-qubit VQE-like ansatz at depth 2 (same as notebook demo).""" |
| 31 | + return create_test_circuit(depth=2) |
| 32 | + |
| 33 | + |
| 34 | +@pytest.fixture(scope="module") |
| 35 | +def ideal_sim(): |
| 36 | + return AerSimulator() |
| 37 | + |
| 38 | + |
| 39 | +@pytest.fixture(scope="module") |
| 40 | +def noisy_sim(): |
| 41 | + return AerSimulator(noise_model=build_noise_model()) |
| 42 | + |
| 43 | + |
| 44 | +# --------------------------------------------------------------------------- |
| 45 | +# Metadata tests |
| 46 | +# --------------------------------------------------------------------------- |
| 47 | + |
| 48 | + |
| 49 | +class TestZNEMetadata: |
17 | 50 | def test_init_default_noise_levels(self): |
18 | 51 | """Default noise levels should be [1, 3, 5].""" |
19 | 52 | zne = ZNEMitigator() |
20 | 53 | assert zne.noise_levels == [1, 3, 5] |
21 | 54 |
|
22 | 55 | def test_init_custom_noise_levels(self): |
23 | 56 | """Custom noise levels should be stored correctly.""" |
24 | | - zne = ZNEMitigator(noise_levels=[1, 2, 3, 4]) |
25 | | - assert zne.noise_levels == [1, 2, 3, 4] |
| 57 | + zne = ZNEMitigator(noise_levels=[1, 3]) |
| 58 | + assert zne.noise_levels == [1, 3] |
26 | 59 |
|
27 | 60 | def test_overhead(self): |
28 | 61 | """Overhead should equal number of noise levels.""" |
29 | | - zne = ZNEMitigator(noise_levels=[1, 3, 5]) |
30 | | - assert zne.overhead() == 3.0 |
| 62 | + assert ZNEMitigator(noise_levels=[1, 3, 5]).overhead() == 3.0 |
| 63 | + assert ZNEMitigator(noise_levels=[1, 3]).overhead() == 2.0 |
31 | 64 |
|
32 | 65 | def test_name(self): |
33 | 66 | """Name should identify the method.""" |
34 | | - zne = ZNEMitigator() |
35 | | - assert "ZNE" in zne.name() |
| 67 | + assert "ZNE" in ZNEMitigator().name() |
36 | 68 |
|
37 | | - def test_mitigate_not_implemented(self): |
38 | | - """mitigate() should raise NotImplementedError until implemented.""" |
39 | | - zne = ZNEMitigator() |
40 | | - with pytest.raises(NotImplementedError): |
41 | | - zne.mitigate(None, None) |
42 | 69 |
|
43 | | - # TODO: Add after Day 4 implementation |
44 | | - # def test_gate_folding_circuit_structure(self): |
45 | | - # def test_richardson_extrapolation_exact_on_noiseless(self): |
46 | | - # def test_zne_reduces_error_on_noisy_simulator(self): |
| 70 | +# --------------------------------------------------------------------------- |
| 71 | +# Gate folding tests |
| 72 | +# --------------------------------------------------------------------------- |
| 73 | + |
| 74 | + |
| 75 | +class TestGateFolding: |
| 76 | + def test_scale_1_returns_unmeasured_circuit(self, test_circuit): |
| 77 | + """Scale factor 1 should return original circuit stripped of measurements.""" |
| 78 | + folded = fold_gates(test_circuit, 1) |
| 79 | + assert folded.num_clbits == 0, "Scale-1 circuit should have no classical bits" |
| 80 | + |
| 81 | + def test_scale_3_triples_gate_count(self, test_circuit): |
| 82 | + """Scale factor 3 should triple each gate (U → U†UU = 3 gates).""" |
| 83 | + base = fold_gates(test_circuit, 1) |
| 84 | + folded = fold_gates(test_circuit, 3) |
| 85 | + assert folded.size() == base.size() * 3 |
| 86 | + |
| 87 | + def test_scale_5_quintuples_gate_count(self, test_circuit): |
| 88 | + """Scale factor 5 should quintuple each gate (U → U†UU†UU = 5 gates).""" |
| 89 | + base = fold_gates(test_circuit, 1) |
| 90 | + folded = fold_gates(test_circuit, 5) |
| 91 | + assert folded.size() == base.size() * 5 |
| 92 | + |
| 93 | + def test_even_scale_raises(self, test_circuit): |
| 94 | + """Even scale factors are invalid and should raise ValueError.""" |
| 95 | + with pytest.raises(ValueError): |
| 96 | + fold_gates(test_circuit, 2) |
| 97 | + |
| 98 | + |
| 99 | +# --------------------------------------------------------------------------- |
| 100 | +# Richardson extrapolation tests |
| 101 | +# --------------------------------------------------------------------------- |
| 102 | + |
| 103 | + |
| 104 | +class TestRichardsonExtrapolation: |
| 105 | + def test_exact_linear_recovery(self): |
| 106 | + """On a perfectly linear signal the extrapolation should be exact.""" |
| 107 | + # <O>(λ) = -0.96 + 0.01 * λ → zero-noise = -0.96 |
| 108 | + scales = [1, 3, 5] |
| 109 | + values = [-0.96 + 0.01 * s for s in scales] |
| 110 | + mitigated, _, _ = richardson_extrapolate(scales, values, poly_degree=1) |
| 111 | + assert abs(mitigated - (-0.96)) < 1e-10 |
| 112 | + |
| 113 | + def test_returns_three_values(self): |
| 114 | + """Should return (mitigated_value, coefficients, poly_func).""" |
| 115 | + result = richardson_extrapolate([1, 3, 5], [-0.9, -0.85, -0.80]) |
| 116 | + assert len(result) == 3 |
| 117 | + |
| 118 | + |
| 119 | +# --------------------------------------------------------------------------- |
| 120 | +# Full ZNE pipeline tests |
| 121 | +# --------------------------------------------------------------------------- |
| 122 | + |
| 123 | + |
| 124 | +class TestZNEPipeline: |
| 125 | + def test_ideal_expectation_nonzero(self, test_circuit, ideal_sim): |
| 126 | + """Ideal <Z0> on this circuit should be strongly negative (~-0.96).""" |
| 127 | + val = measure_expectation(fold_gates(test_circuit, 1), ideal_sim, num_shots=4096) |
| 128 | + assert val < -0.80, f"Expected <Z0> < -0.80, got {val:.4f}" |
| 129 | + |
| 130 | + def test_mitigate_returns_dict(self, test_circuit, noisy_sim): |
| 131 | + """mitigate() should return a dict with the required keys.""" |
| 132 | + result = ZNEMitigator(num_shots=1024).mitigate(test_circuit, backend=noisy_sim) |
| 133 | + assert set(result.keys()) >= {"mitigated", "noisy", "scales", "poly_func"} |
| 134 | + |
| 135 | + def test_mitigate_noisy_values_degrade_with_scale(self, test_circuit, noisy_sim): |
| 136 | + """Noisy expectation values should become less negative as λ increases.""" |
| 137 | + result = ZNEMitigator(noise_levels=[1, 3, 5], num_shots=2048).mitigate( |
| 138 | + test_circuit, backend=noisy_sim |
| 139 | + ) |
| 140 | + noisy = result["noisy"] |
| 141 | + # Each successive scale should push the value closer to 0 (more noise) |
| 142 | + assert noisy[0] < noisy[1] or abs(noisy[0] - noisy[1]) < 0.05, ( |
| 143 | + "λ=3 value should be >= λ=1 (less negative) within shot noise tolerance" |
| 144 | + ) |
| 145 | + |
| 146 | + def test_mitigate_reduces_error_vs_raw(self, test_circuit, ideal_sim, noisy_sim): |
| 147 | + """ZNE mitigated value should be at least as close to ideal as raw noisy.""" |
| 148 | + ideal = measure_expectation(fold_gates(test_circuit, 1), ideal_sim, num_shots=4096) |
| 149 | + result = ZNEMitigator(noise_levels=[1, 3, 5], num_shots=2048).mitigate( |
| 150 | + test_circuit, backend=noisy_sim |
| 151 | + ) |
| 152 | + raw_error = abs(ideal - result["noisy"][0]) |
| 153 | + zne_error = abs(ideal - result["mitigated"]) |
| 154 | + # Allow up to 50% worse due to shot noise variance — what matters is direction |
| 155 | + assert zne_error <= raw_error * 1.5, ( |
| 156 | + f"ZNE error ({zne_error:.4f}) should not be much worse than " |
| 157 | + f"raw noisy error ({raw_error:.4f})" |
| 158 | + ) |
0 commit comments