From cc228f984b0daf970c3034a55e99149ab968bc22 Mon Sep 17 00:00:00 2001 From: huaweil Date: Thu, 9 Jul 2026 13:56:36 +0800 Subject: [PATCH] Fix SPAM-window idle priors across rounds Assign explicit noise-model data-idle priors at the final time step to the SPAM family in every non-final stabilizer round, independent of the data preparation map. Add real-circuit metadata regression coverage for both memory bases, final-round zeroing, and bulk idle rates. Signed-off-by: huaweil --- code/qec/precompute_dem.py | 11 +- .../test_precompute_dem_probabilities.py | 162 ++++++++++++++++++ 2 files changed, 169 insertions(+), 4 deletions(-) create mode 100644 code/tests/test_precompute_dem_probabilities.py diff --git a/code/qec/precompute_dem.py b/code/qec/precompute_dem.py index e5f7bce..5f5e7ca 100644 --- a/code/qec/precompute_dem.py +++ b/code/qec/precompute_dem.py @@ -622,6 +622,11 @@ def build_single_p_marginal( is_ancilla_meas = is_meas and is_meas_qubit is_data_prep = is_prep and is_data is_data_meas = (tt == int(t_total) - 1) and is_data and ((r, q) in prep_basis_map) + # In the explicit noise model, every non-final stabilizer round has a + # data-idle SPAM window after the ancilla measurement/reset step. Do + # not gate this location on prep_basis_map: data qubits are prepared + # only in the boundary rounds, but they idle in this window every round. + is_data_spam_idle = (tt == int(t_total) - 1) and is_data if use_nm: if is_final_round and not (tt == 0 and is_data): @@ -673,17 +678,15 @@ def build_single_p_marginal( p_err[eidx] = float( p_prep_X if prep_basis == 0 else p_prep_Z ) if allowed else 0.0 - elif is_data_meas: + elif is_data_spam_idle: if is_final_round: p_err[eidx] = 0.0 else: p_err[eidx] = float(_nm_single.get(et, {}).get("idle_spam", 0.0)) else: - # Bulk idle: use idle_cnot or idle_spam depending on time step + # Remaining single-qubit locations are bulk/CNOT-layer idles. if is_final_round: p_err[eidx] = 0.0 - elif is_prep or is_data_meas: - p_err[eidx] = float(_nm_single.get(et, {}).get("idle_spam", 0.0)) else: p_err[eidx] = float(_nm_single.get(et, {}).get("idle_cnot", 0.0)) else: diff --git a/code/tests/test_precompute_dem_probabilities.py b/code/tests/test_precompute_dem_probabilities.py new file mode 100644 index 0000000..1e5d6ec --- /dev/null +++ b/code/tests/test_precompute_dem_probabilities.py @@ -0,0 +1,162 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import sys +import unittest +from pathlib import Path + +import numpy as np + +_repo_code = Path(__file__).resolve().parent.parent +if str(_repo_code) not in sys.path: + sys.path.insert(0, str(_repo_code)) + +from qec.noise_model import CNOT_ERROR_TYPES, NoiseModel +from qec.precompute_dem import ( + build_single_p_marginal, + extract_cnot_structure_from_stim_text, + generate_all_errors_local, + replicate_metadata_across_rounds, +) +from qec.surface_code.memory_circuit import MemoryCircuit + + +class TestBuildSinglePMarginalSpamWindows(unittest.TestCase): + + @staticmethod + def _build_probability_vector(basis): + distance = 3 + n_rounds = 4 + p_scalar = 0.004 + noise_model = NoiseModel( + p_prep_X=0.001, + p_prep_Z=0.002, + p_meas_X=0.003, + p_meas_Z=0.004, + p_idle_cnot_X=0.0011, + p_idle_cnot_Y=0.0012, + p_idle_cnot_Z=0.0013, + p_idle_spam_X=0.0021, + p_idle_spam_Y=0.0022, + p_idle_spam_Z=0.0023, + **{f"p_cnot_{error_type}": 0.0001 for error_type in CNOT_ERROR_TYPES}, + ) + + circuit = MemoryCircuit( + distance=distance, + idle_error=p_scalar, + sqgate_error=p_scalar, + tqgate_error=p_scalar, + spam_error=2.0 * p_scalar / 3.0, + n_rounds=n_rounds, + basis=basis, + code_rotation="XV", + noise_model=noise_model, + ) + cnot_circuit, cx_times = extract_cnot_structure_from_stim_text(circuit.circuit) + t_total = int(len(cx_times) + 2) + nq = int(2 * distance * distance - 1) + _, metadata_local = generate_all_errors_local( + t_total=t_total, + nq=nq, + controls_by_layer=cnot_circuit, + cx_times=cx_times, + ) + metadata_global = replicate_metadata_across_rounds( + metadata_local=metadata_local, + n_rounds=n_rounds, + ) + + data_qubits = np.asarray(circuit.code.data_qubits, dtype=np.int32) + xcheck_qubits = np.asarray(circuit.code.xcheck_qubits, dtype=np.int32) + zcheck_qubits = np.asarray(circuit.code.zcheck_qubits, dtype=np.int32) + meas_qubits = np.concatenate([xcheck_qubits, zcheck_qubits]).astype(np.int32) + meas_bases = np.concatenate( + [ + np.zeros(len(xcheck_qubits), dtype=np.int32), + np.ones(len(zcheck_qubits), dtype=np.int32), + ] + ) + probabilities = build_single_p_marginal( + error_metadata_global=metadata_global, + t_total=t_total, + n_rounds=n_rounds, + data_qubits=data_qubits, + xcheck_qubits=xcheck_qubits, + zcheck_qubits=zcheck_qubits, + meas_qubits=meas_qubits, + meas_bases=meas_bases, + basis=basis, + p_scalar=p_scalar, + noise_model=noise_model, + ) + return ( + probabilities, + metadata_global, + noise_model, + set(int(q) for q in data_qubits.tolist()), + t_total, + n_rounds, + ) + + def test_real_metadata_uses_spam_idle_in_every_non_final_round(self): + for basis in ("X", "Z"): + with self.subTest(basis=basis): + probabilities, metadata, noise_model, data_qubits, t_total, n_rounds = ( + self._build_probability_vector(basis) + ) + seen = set() + count = 0 + for error_index, round_index, time_index, qubit, error_type, _ in metadata: + if ( + 0 <= round_index < n_rounds - 1 and time_index == t_total - 1 and + qubit in data_qubits and len(error_type) == 1 + ): + expected = getattr(noise_model, f"p_idle_spam_{error_type}") + self.assertAlmostEqual( + float(probabilities[error_index]), + float(expected), + places=7, + ) + seen.add((round_index, error_type)) + count += 1 + + self.assertEqual( + seen, + { + (round_index, error_type) + for round_index in range(n_rounds - 1) + for error_type in ("X", "Y", "Z") + }, + ) + self.assertEqual(count, len(data_qubits) * 3 * (n_rounds - 1)) + + def test_real_metadata_keeps_final_round_quiet_and_bulk_idles_on_cnot_rates(self): + for basis in ("X", "Z"): + with self.subTest(basis=basis): + probabilities, metadata, noise_model, data_qubits, t_total, n_rounds = ( + self._build_probability_vector(basis) + ) + final_count = 0 + bulk_types = set() + for error_index, round_index, time_index, qubit, error_type, _ in metadata: + if qubit not in data_qubits or len(error_type) != 1: + continue + if round_index == n_rounds - 1 and time_index == t_total - 1: + self.assertEqual(float(probabilities[error_index]), 0.0) + final_count += 1 + elif 0 < time_index < t_total - 1 and round_index < n_rounds - 1: + expected = getattr(noise_model, f"p_idle_cnot_{error_type}") + self.assertAlmostEqual( + float(probabilities[error_index]), + float(expected), + places=7, + ) + bulk_types.add(error_type) + + self.assertEqual(final_count, len(data_qubits) * 3) + self.assertEqual(bulk_types, {"X", "Y", "Z"}) + + +if __name__ == "__main__": + unittest.main()