From b757cfcaa14488140ca9a5aba6fd1808203ad031 Mon Sep 17 00:00:00 2001 From: Ivan Basov Date: Thu, 9 Jul 2026 20:21:35 +0000 Subject: [PATCH 1/2] Fix reset-window and final-round data-qubit priors in precompute_dem Two circuit-vs-sampler mismatches in build_single_p_marginal, both in the explicit 25-parameter noise model path: 1. Data qubits idling during the ancilla prep/reset window (tt == 0, intermediate rounds) were assigned p_idle_cnot_*, but the circuit emits no noise there: that idle is folded into the measurement-window SPAM idle. Assign 0 instead. 2. The final-round data injection was assigned p_prep_*, but the circuit represents noisy data readout as a fake data-measurement flip at p_meas_* (Z_ERROR(p_meas_X) for X-basis readout, X_ERROR(p_meas_Z) for Z-basis). Use the measurement rates. Adds unit tests for both locations plus a seeded end-to-end guard that compares per-round syndrome densities from Bernoulli(p) errors pushed through H against direct Stim sampling of the same MemoryCircuit. Co-Authored-By: Claude Fable 5 --- code/qec/precompute_dem.py | 15 +- .../test_precompute_dem_probabilities.py | 234 ++++++++++++++++++ 2 files changed, 248 insertions(+), 1 deletion(-) diff --git a/code/qec/precompute_dem.py b/code/qec/precompute_dem.py index 5f5e7ca..822ddbc 100644 --- a/code/qec/precompute_dem.py +++ b/code/qec/precompute_dem.py @@ -671,8 +671,13 @@ def build_single_p_marginal( else: allowed = (et == "X") if is_final_round: + # The circuit represents noisy data readout as a fake + # data-measurement flip injected at the start of the final + # perfect round, at the *measurement* rates: Z_ERROR(p_meas_X) + # for X-basis readout, X_ERROR(p_meas_Z) for Z-basis (see + # MemoryCircuit's logical_measurement injection). p_err[eidx] = float( - p_prep_X if prep_basis == 0 else p_prep_Z + p_meas_X if prep_basis == 0 else p_meas_Z ) if allowed else 0.0 else: p_err[eidx] = float( @@ -683,6 +688,14 @@ def build_single_p_marginal( p_err[eidx] = 0.0 else: p_err[eidx] = float(_nm_single.get(et, {}).get("idle_spam", 0.0)) + elif tt == 0 and is_data: + # Data qubits idling during the ancilla prep/reset window carry + # no separate noise location: the circuit folds that idle into + # the measurement-window SPAM idle above and emits nothing here + # (see MemoryCircuit: "IGNORE data-idle during ancilla + # prep/reset"). Boundary rounds never reach this branch: their + # tt == 0 data locations are consumed by is_data_prep. + p_err[eidx] = 0.0 else: # Remaining single-qubit locations are bulk/CNOT-layer idles. if is_final_round: diff --git a/code/tests/test_precompute_dem_probabilities.py b/code/tests/test_precompute_dem_probabilities.py index 1e5d6ec..6f38e76 100644 --- a/code/tests/test_precompute_dem_probabilities.py +++ b/code/tests/test_precompute_dem_probabilities.py @@ -158,5 +158,239 @@ def test_real_metadata_keeps_final_round_quiet_and_bulk_idles_on_cnot_rates(self self.assertEqual(bulk_types, {"X", "Y", "Z"}) +class TestBuildSinglePMarginalResetWindowAndFinalRound(unittest.TestCase): + + def test_reset_window_data_idles_carry_no_prior_in_intermediate_rounds(self): + # The circuit folds the data-qubit idle during the ancilla prep/reset + # window into the measurement-window SPAM idle and emits no noise at + # tt == 0 for data qubits outside the boundary rounds. + for basis in ("X", "Z"): + with self.subTest(basis=basis): + probabilities, metadata, _, data_qubits, _, n_rounds = ( + TestBuildSinglePMarginalSpamWindows._build_probability_vector(basis) + ) + count = 0 + for error_index, round_index, time_index, qubit, error_type, _ in metadata: + if ( + 0 < round_index < n_rounds - 1 and time_index == 0 and + qubit in data_qubits and len(error_type) == 1 + ): + self.assertEqual(float(probabilities[error_index]), 0.0) + count += 1 + + self.assertEqual(count, len(data_qubits) * 3 * (n_rounds - 2)) + + def test_final_round_data_injection_uses_measurement_rate(self): + # The final perfect round represents noisy data readout as a fake + # data-measurement flip at the start of the round: Z_ERROR(p_meas_X) + # for X-basis readout, X_ERROR(p_meas_Z) for Z-basis readout. + for basis, allowed_type, rate_name in ( + ("X", "Z", "p_meas_X"), + ("Z", "X", "p_meas_Z"), + ): + with self.subTest(basis=basis): + probabilities, metadata, noise_model, data_qubits, _, n_rounds = ( + TestBuildSinglePMarginalSpamWindows._build_probability_vector(basis) + ) + count = 0 + for error_index, round_index, time_index, qubit, error_type, _ in metadata: + if ( + round_index == n_rounds - 1 and time_index == 0 and + qubit in data_qubits and len(error_type) == 1 + ): + expected = ( + getattr(noise_model, rate_name) + if error_type == allowed_type else 0.0 + ) + self.assertAlmostEqual( + float(probabilities[error_index]), + float(expected), + places=7, + ) + count += 1 + + self.assertEqual(count, len(data_qubits) * 3) + + def test_round_zero_data_prep_keeps_prep_rate(self): + for basis, allowed_type, rate_name in ( + ("X", "Z", "p_prep_X"), + ("Z", "X", "p_prep_Z"), + ): + with self.subTest(basis=basis): + probabilities, metadata, noise_model, data_qubits, _, _ = ( + TestBuildSinglePMarginalSpamWindows._build_probability_vector(basis) + ) + count = 0 + for error_index, round_index, time_index, qubit, error_type, _ in metadata: + if ( + round_index == 0 and time_index == 0 and + qubit in data_qubits and len(error_type) == 1 + ): + expected = ( + getattr(noise_model, rate_name) + if error_type == allowed_type else 0.0 + ) + self.assertAlmostEqual( + float(probabilities[error_index]), + float(expected), + places=7, + ) + count += 1 + + self.assertEqual(count, len(data_qubits) * 3) + + +class TestSamplerMatchesStimSyndromeDensities(unittest.TestCase): + """Seeded end-to-end guard against circuit-vs-sampler prior drift. + + Per-round syndrome densities obtained by pushing Bernoulli(p) errors + through the precomputed H must match direct Stim sampling of the same + MemoryCircuit. Both paths are seeded, so the comparison is deterministic. + """ + + DISTANCE = 3 + N_ROUNDS = 4 + BASIS = "X" + SHOTS = 50_000 + CHUNK = 10_000 + + @staticmethod + def _make_noise_model(**overrides): + params = { + key: 0.0 + for key in ( + "p_prep_X", "p_prep_Z", "p_meas_X", "p_meas_Z", + "p_idle_cnot_X", "p_idle_cnot_Y", "p_idle_cnot_Z", + "p_idle_spam_X", "p_idle_spam_Y", "p_idle_spam_Z", + ) + } + params.update({f"p_cnot_{error_type}": 0.0 for error_type in CNOT_ERROR_TYPES}) + params.update(overrides) + return NoiseModel(**params) + + def _build_circuit(self, noise_model): + return MemoryCircuit( + distance=self.DISTANCE, + idle_error=0.004, + sqgate_error=0.004, + tqgate_error=0.004, + spam_error=2.0 * 0.004 / 3.0, + n_rounds=self.N_ROUNDS, + basis=self.BASIS, + code_rotation="XV", + noise_model=noise_model, + ) + + def _model_round_densities(self, noise_model): + import torch + + from qec.dem_sampling import measure_from_stacked_frames + from qec.precompute_dem import precompute_dem_bundle_surface_code + + artifacts = precompute_dem_bundle_surface_code( + distance=self.DISTANCE, + n_rounds=self.N_ROUNDS, + basis=self.BASIS, + code_rotation="XV", + p_scalar=0.004, + dem_output_dir=None, + device=torch.device("cpu"), + export=False, + return_artifacts=True, + noise_model=noise_model, + ) + H, p, nq = artifacts["H"], artifacts["p"], artifacts["nq"] + circuit = self._build_circuit(noise_model) + xcheck_qubits = np.asarray(circuit.code.xcheck_qubits, dtype=np.int64) + zcheck_qubits = np.asarray(circuit.code.zcheck_qubits, dtype=np.int64) + meas_qubits = torch.from_numpy(np.concatenate([xcheck_qubits, zcheck_qubits])) + meas_bases = torch.from_numpy( + np.concatenate([ + np.zeros(len(xcheck_qubits), np.int64), + np.ones(len(zcheck_qubits), np.int64), + ]) + ) + + generator = torch.Generator().manual_seed(20260709) + H_t = H.t().float() + totals = torch.zeros((self.N_ROUNDS, len(meas_qubits)), dtype=torch.float64) + remaining = self.SHOTS + while remaining > 0: + batch = min(self.CHUNK, remaining) + errors = ( + torch.rand((batch, H.shape[1]), generator=generator) < p[None, :] + ).float() + frames = (errors @ H_t).remainder(2).to(torch.uint8) + measurements = measure_from_stacked_frames(frames, meas_qubits, meas_bases, nq) + syndromes = measurements.clone() + syndromes[:, 1:] ^= measurements[:, :-1] + totals += syndromes.to(torch.float64).sum(dim=0) + remaining -= batch + return (totals / self.SHOTS).numpy(), len(xcheck_qubits) + + def _stim_round_densities(self, noise_model): + import stim + + circuit = self._build_circuit(noise_model) + sampler = stim.Circuit(circuit.circuit).compile_sampler(seed=20260709) + n_ancillas = len(circuit.code.xcheck_qubits) + len(circuit.code.zcheck_qubits) + measurements = sampler.sample(self.SHOTS) + measurements = ( + measurements[:, : self.N_ROUNDS * n_ancillas] + .reshape(self.SHOTS, self.N_ROUNDS, n_ancillas) + .astype(np.uint8) + ) + syndromes = measurements.copy() + syndromes[:, 1:] ^= measurements[:, :-1] + return syndromes.mean(axis=0) + + def _assert_densities_match(self, noise_model): + model, n_xchecks = self._model_round_densities(noise_model) + reference = self._stim_round_densities(noise_model) + # Round-0 Z-check outcomes are projection-random in the X basis; the + # frame-based model measures deviations from the noise-free reference, + # so exclude those cells. + valid = np.ones_like(reference, dtype=bool) + valid[0, n_xchecks:] = False + for round_index in range(self.N_ROUNDS): + model_mean = model[round_index][valid[round_index]].mean() + reference_mean = reference[round_index][valid[round_index]].mean() + n_cells = int(valid[round_index].sum()) * self.SHOTS + stderr = float( + np.sqrt(max(reference_mean * (1.0 - reference_mean), 1e-12) / n_cells) + ) + self.assertLess( + abs(model_mean - reference_mean), + 5.0 * stderr + 1e-4, + msg=( + f"round {round_index}: model syndrome density {model_mean:.6f} " + f"vs stim {reference_mean:.6f}" + ), + ) + + def test_cnot_idle_only_matches_stim(self): + self._assert_densities_match( + self._make_noise_model( + p_idle_cnot_X=0.002, p_idle_cnot_Y=0.002, p_idle_cnot_Z=0.002 + ) + ) + + def test_measurement_only_matches_stim(self): + self._assert_densities_match( + self._make_noise_model(p_meas_X=0.004, p_meas_Z=0.004) + ) + + def test_full_noise_model_matches_stim(self): + self._assert_densities_match( + self._make_noise_model( + p_prep_X=0.004, p_prep_Z=0.004, + p_meas_X=0.004, p_meas_Z=0.004, + p_idle_cnot_X=0.002, p_idle_cnot_Y=0.002, p_idle_cnot_Z=0.002, + p_idle_spam_X=0.003984, p_idle_spam_Y=0.003984, p_idle_spam_Z=0.003984, + **{f"p_cnot_{error_type}": 0.0004 for error_type in CNOT_ERROR_TYPES}, + ) + ) + + if __name__ == "__main__": unittest.main() From bde141d936a596b88cf7ec2fc1ff5cc03b01b944 Mon Sep 17 00:00:00 2001 From: Ivan Basov Date: Thu, 9 Jul 2026 21:12:33 +0000 Subject: [PATCH 2/2] Apply yapf formatting Co-Authored-By: Claude Fable 5 --- .../test_precompute_dem_probabilities.py | 76 ++++++++++--------- 1 file changed, 40 insertions(+), 36 deletions(-) diff --git a/code/tests/test_precompute_dem_probabilities.py b/code/tests/test_precompute_dem_probabilities.py index 6f38e76..425ad77 100644 --- a/code/tests/test_precompute_dem_probabilities.py +++ b/code/tests/test_precompute_dem_probabilities.py @@ -195,12 +195,11 @@ def test_final_round_data_injection_uses_measurement_rate(self): count = 0 for error_index, round_index, time_index, qubit, error_type, _ in metadata: if ( - round_index == n_rounds - 1 and time_index == 0 and - qubit in data_qubits and len(error_type) == 1 + round_index == n_rounds - 1 and time_index == 0 and qubit in data_qubits and + len(error_type) == 1 ): expected = ( - getattr(noise_model, rate_name) - if error_type == allowed_type else 0.0 + getattr(noise_model, rate_name) if error_type == allowed_type else 0.0 ) self.assertAlmostEqual( float(probabilities[error_index]), @@ -223,12 +222,11 @@ def test_round_zero_data_prep_keeps_prep_rate(self): count = 0 for error_index, round_index, time_index, qubit, error_type, _ in metadata: if ( - round_index == 0 and time_index == 0 and - qubit in data_qubits and len(error_type) == 1 + round_index == 0 and time_index == 0 and qubit in data_qubits and + len(error_type) == 1 ): expected = ( - getattr(noise_model, rate_name) - if error_type == allowed_type else 0.0 + getattr(noise_model, rate_name) if error_type == allowed_type else 0.0 ) self.assertAlmostEqual( float(probabilities[error_index]), @@ -257,11 +255,17 @@ class TestSamplerMatchesStimSyndromeDensities(unittest.TestCase): @staticmethod def _make_noise_model(**overrides): params = { - key: 0.0 - for key in ( - "p_prep_X", "p_prep_Z", "p_meas_X", "p_meas_Z", - "p_idle_cnot_X", "p_idle_cnot_Y", "p_idle_cnot_Z", - "p_idle_spam_X", "p_idle_spam_Y", "p_idle_spam_Z", + key: 0.0 for key in ( + "p_prep_X", + "p_prep_Z", + "p_meas_X", + "p_meas_Z", + "p_idle_cnot_X", + "p_idle_cnot_Y", + "p_idle_cnot_Z", + "p_idle_spam_X", + "p_idle_spam_Y", + "p_idle_spam_Z", ) } params.update({f"p_cnot_{error_type}": 0.0 for error_type in CNOT_ERROR_TYPES}) @@ -305,10 +309,12 @@ def _model_round_densities(self, noise_model): zcheck_qubits = np.asarray(circuit.code.zcheck_qubits, dtype=np.int64) meas_qubits = torch.from_numpy(np.concatenate([xcheck_qubits, zcheck_qubits])) meas_bases = torch.from_numpy( - np.concatenate([ - np.zeros(len(xcheck_qubits), np.int64), - np.ones(len(zcheck_qubits), np.int64), - ]) + np.concatenate( + [ + np.zeros(len(xcheck_qubits), np.int64), + np.ones(len(zcheck_qubits), np.int64), + ] + ) ) generator = torch.Generator().manual_seed(20260709) @@ -317,9 +323,7 @@ def _model_round_densities(self, noise_model): remaining = self.SHOTS while remaining > 0: batch = min(self.CHUNK, remaining) - errors = ( - torch.rand((batch, H.shape[1]), generator=generator) < p[None, :] - ).float() + errors = (torch.rand((batch, H.shape[1]), generator=generator) < p[None, :]).float() frames = (errors @ H_t).remainder(2).to(torch.uint8) measurements = measure_from_stacked_frames(frames, meas_qubits, meas_bases, nq) syndromes = measurements.clone() @@ -336,9 +340,9 @@ def _stim_round_densities(self, noise_model): n_ancillas = len(circuit.code.xcheck_qubits) + len(circuit.code.zcheck_qubits) measurements = sampler.sample(self.SHOTS) measurements = ( - measurements[:, : self.N_ROUNDS * n_ancillas] - .reshape(self.SHOTS, self.N_ROUNDS, n_ancillas) - .astype(np.uint8) + measurements[:, :self.N_ROUNDS * + n_ancillas].reshape(self.SHOTS, self.N_ROUNDS, + n_ancillas).astype(np.uint8) ) syndromes = measurements.copy() syndromes[:, 1:] ^= measurements[:, :-1] @@ -356,9 +360,7 @@ def _assert_densities_match(self, noise_model): model_mean = model[round_index][valid[round_index]].mean() reference_mean = reference[round_index][valid[round_index]].mean() n_cells = int(valid[round_index].sum()) * self.SHOTS - stderr = float( - np.sqrt(max(reference_mean * (1.0 - reference_mean), 1e-12) / n_cells) - ) + stderr = float(np.sqrt(max(reference_mean * (1.0 - reference_mean), 1e-12) / n_cells)) self.assertLess( abs(model_mean - reference_mean), 5.0 * stderr + 1e-4, @@ -370,23 +372,25 @@ def _assert_densities_match(self, noise_model): def test_cnot_idle_only_matches_stim(self): self._assert_densities_match( - self._make_noise_model( - p_idle_cnot_X=0.002, p_idle_cnot_Y=0.002, p_idle_cnot_Z=0.002 - ) + self._make_noise_model(p_idle_cnot_X=0.002, p_idle_cnot_Y=0.002, p_idle_cnot_Z=0.002) ) def test_measurement_only_matches_stim(self): - self._assert_densities_match( - self._make_noise_model(p_meas_X=0.004, p_meas_Z=0.004) - ) + self._assert_densities_match(self._make_noise_model(p_meas_X=0.004, p_meas_Z=0.004)) def test_full_noise_model_matches_stim(self): self._assert_densities_match( self._make_noise_model( - p_prep_X=0.004, p_prep_Z=0.004, - p_meas_X=0.004, p_meas_Z=0.004, - p_idle_cnot_X=0.002, p_idle_cnot_Y=0.002, p_idle_cnot_Z=0.002, - p_idle_spam_X=0.003984, p_idle_spam_Y=0.003984, p_idle_spam_Z=0.003984, + p_prep_X=0.004, + p_prep_Z=0.004, + p_meas_X=0.004, + p_meas_Z=0.004, + p_idle_cnot_X=0.002, + p_idle_cnot_Y=0.002, + p_idle_cnot_Z=0.002, + p_idle_spam_X=0.003984, + p_idle_spam_Y=0.003984, + p_idle_spam_Z=0.003984, **{f"p_cnot_{error_type}": 0.0004 for error_type in CNOT_ERROR_TYPES}, ) )