Fix SPAM-window idle priors in intermediate surface-code rounds

[huaweil-nv]

Summary

Fix the explicit 25-parameter noise prior assigned to data-qubit idle faults during the ancilla SPAM window.

MemoryCircuit applies idle_kind="spam" to data qubits at the end of every non-final stabilizer round. However, build_single_p_marginal() previously recognized this location only when the data qubit was present in prep_basis_map.

Since data qubits appear in that map only at circuit boundaries, intermediate-round SPAM-window faults were incorrectly assigned p_idle_cnot_X/Y/Z instead of p_idle_spam_X/Y/Z.

Previous behavior

For a seven-round circuit:

Round	Expected prior	Previous prior
0	p_idle_spam_*	p_idle_spam_*
1–5	p_idle_spam_*	p_idle_cnot_*
6, final perfect round	0	0

The previous predicate was:

is_data_meas = (
    tt == t_total - 1
    and is_data
    and (r, q) in prep_basis_map
)

This incorrectly tied a repeated data-idle location to data preparation.

Fix

Identify the SPAM-window data-idle location independently of prep_basis_map:

is_data_spam_idle = (tt == t_total - 1) and is_data

The existing final-round handling keeps the final perfect round at zero.

The change also removes an unreachable fallback condition. After preparation, measurement, and SPAM-window locations have been handled, the remaining single-qubit locations are bulk/CNOT-layer idles and use p_idle_cnot_X/Y/Z.

impact

The public d=7/r=7 configuration uses different values for the two idle families:

p_idle_cnot_X/Y/Z = 0.001
p_idle_spam_X/Y/Z = 0.001996

The previous implementation therefore misclassified:

49 data qubits
× 3 Pauli fault types
× 5 intermediate rounds
= 735 probability-vector entries per basis

Independent Stim validation

I compared the precomputed H/p/A sampler against direct Stim sampling of the same MemoryCircuit.

To isolate this behavior, only the SPAM-idle channels were enabled:

p_idle_spam_X/Y/Z = 0.001996
all other noise parameters = 0

For d=7/r=7, 500,000 shots were sampled for each path and basis:

Basis	Prior	RMSE vs Stim	Maximum statistical deviation
X	Fixed	0.0000540	1.58 standard errors
X	Previous	0.0119817	305.52 standard errors
Z	Fixed	0.0000857	2.32 standard errors
Z	Previous	0.0119722	305.38 standard errors

The corrected sampler agrees with Stim within Monte Carlo variation.

The previous prior produces syndrome activity only from the initial round's SPAM-idle faults. Its later-round syndrome densities collapse to zero, while both Stim and the corrected sampler retain the expected activity in every applicable round. An additional d=3/r=4 comparison with 2,000,000 shots per path and basis gave the same result.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[ivanbasov]

/ok to test cc228f9

[ivanbasov]

Thank you for this contribution — this is an exemplary bug report and fix. The diagnosis is exactly right, the write-up made it easy to verify, and the independent Stim validation with isolated noise channels is precisely the right methodology for this class of bug. Much appreciated!

Verification

I reviewed the predicate change against the circuit generator and re-validated independently:

• MemoryCircuit emits the data-qubit SPAM idle after ancilla measurement in every stabilizer round (code/qec/surface_code/memory_circuit.py:1199-1207), so gating on prep_basis_map (populated for data qubits only at r == 0 and r == n_rounds - 1) was indeed wrong for intermediate rounds. The new is_data_spam_idle predicate matches the circuit.
• The removed fallback (elif is_prep or is_data_meas inside the final else) is confirmed unreachable: every qubit is either a data or a check qubit, so all tt == 0 prep locations are consumed by is_ancilla_prep/is_data_prep, and the SPAM-window locations by the new branch.
• The two new tests pass on this branch, and both fail on main with exactly the misclassification signature (idle_cnot value observed where idle_spam is expected), so they are genuine regression coverage. The completeness assertions (every round × Pauli seen, exact location counts) are a nice touch.
• I reproduced your Stim cross-check with an independent CPU harness (Bernoulli sampling from p pushed through H, per-round syndrome densities vs stim sampling of the same circuit, d=3/r=4, 400k shots): with only p_idle_spam_* enabled, this branch agrees with Stim within 2.7 standard errors across all (round, ancilla) cells.

LGTM.

Two adjacent pre-existing mismatches (not introduced here, non-blocking)

While auditing, I extended the same isolated-channel comparison to the channels your validation left at zero, and found two more circuit-vs-sampler discrepancies in build_single_p_marginal. Both predate this PR and are untouched by it — noting them here because they are the same class of bug and your harness is well positioned to catch them.

1. Reset-window data idles are overcounted with p_idle_cnot_*.
In NoiseModel mode the circuit deliberately applies no idle noise to data qubits during the ancilla prep/reset window — it is folded into the measurement-window SPAM idle (memory_circuit.py:1156-1162). But the marginal builder lets intermediate-round data locations at tt == 0 fall through to the bulk-idle branch and assigns p_idle_cnot_*. With only p_idle_cnot_X/Y/Z = 0.002 enabled (d=3, r=4, X basis, 400k shots):

Round	Sampler	Stim
0	0.01087	0.01079
1	0.03291	0.02165
2	0.03294	0.02177
3	0.01092	0.01102

With the full public 25-parameter values (r13 config), intermediate-round syndrome densities come out ~10% hotter than Stim (0.1020 vs 0.0924) even after this fix. Suggested fix: an explicit elif tt == 0 and is_data: p_err[eidx] = 0.0 for non-boundary rounds (boundary rounds are already consumed by is_data_prep).

2. The final-round data injection uses the wrong rate family.
The circuit represents noisy data readout as a fake data-measurement error at the start of the final perfect round, at p_meas_X/p_meas_Z (memory_circuit.py:1128-1155). The sampler's is_data_prep branch assigns p_prep_X/p_prep_Z instead (its is_final_round split currently has identical arms — it looks like it was meant to carry exactly this distinction). Isolating the channels makes the mirror image visible in round 3: with only p_meas_* = 0.004, Stim shows 0.0099 vs the sampler's 0.0040; with only p_prep_* = 0.004, the sampler shows 0.0099 vs Stim's 0.0040. All shipped configs use p_prep == p_meas, so this cancels numerically today, but it breaks as soon as the two families are calibrated separately.

Since both are pre-existing and independent of this fix, they'll be addressed in a follow-up PR so this one can merge as-is.

Repro harness (CPU, no cuStabilizer needed)

import numpy as np
import stim
import torch

from qec.noise_model import CNOT_ERROR_TYPES, NoiseModel
from qec.precompute_dem import precompute_dem_bundle_surface_code
from qec.dem_sampling import measure_from_stacked_frames
from qec.surface_code.memory_circuit import MemoryCircuit

D, R, BASIS, ROT, P_SCALAR, B, CHUNK = 3, 4, "X", "XV", 0.004, 400_000, 20_000
torch.manual_seed(0)


def make_nm(**kw):
    base = {k: 0.0 for k 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")}
    base.update({f"p_cnot_{t}": 0.0 for t in CNOT_ERROR_TYPES})
    base.update(kw)
    return NoiseModel(**base)


def build_circuit(nm):
    return MemoryCircuit(distance=D, idle_error=P_SCALAR, sqgate_error=P_SCALAR,
                         tqgate_error=P_SCALAR, spam_error=2 / 3 * P_SCALAR,
                         n_rounds=R, basis=BASIS, code_rotation=ROT, noise_model=nm)


def model_syndrome_means(nm):
    art = precompute_dem_bundle_surface_code(
        distance=D, n_rounds=R, basis=BASIS, code_rotation=ROT, p_scalar=P_SCALAR,
        dem_output_dir=None, device=torch.device("cpu"), export=False,
        return_artifacts=True, noise_model=nm)
    H, p, nq = art["H"], art["p"], art["nq"]
    circ = build_circuit(nm)
    xq = np.asarray(circ.code.xcheck_qubits, dtype=np.int64)
    zq = np.asarray(circ.code.zcheck_qubits, dtype=np.int64)
    mq = torch.from_numpy(np.concatenate([xq, zq]))
    mb = torch.from_numpy(np.concatenate([np.zeros(len(xq), np.int64),
                                          np.ones(len(zq), np.int64)]))
    Ht, acc, n = H.t().float(), torch.zeros((R, len(mq)), dtype=torch.float64), 0
    while n < B:
        b = min(CHUNK, B - n)
        err = (torch.rand(b, H.shape[1]) < p[None, :]).float()
        frames = (err @ Ht).remainder(2).to(torch.uint8)
        meas = measure_from_stacked_frames(frames, mq, mb, nq)
        synd = meas.clone()
        synd[:, 1:] ^= meas[:, :-1]
        acc += synd.to(torch.float64).sum(dim=0)
        n += b
    return (acc / B).numpy(), len(xq)


def stim_syndrome_means(nm):
    circ = build_circuit(nm)
    sampler = stim.Circuit(circ.circuit).compile_sampler()
    n_anc = len(circ.code.xcheck_qubits) + len(circ.code.zcheck_qubits)
    m = sampler.sample(B)[:, : R * n_anc].reshape(B, R, n_anc).astype(np.uint8)
    synd = m.copy()
    synd[:, 1:] ^= m[:, :-1]
    return synd.mean(axis=0)


def report(name, nm):
    mod, n_x = model_syndrome_means(nm)
    st = stim_syndrome_means(nm)
    mask = np.ones_like(st, dtype=bool)
    mask[0, n_x:] = False  # round-0 Z checks are projection-random in X basis
    dev = np.abs(mod - st) / np.sqrt(np.maximum(st * (1 - st), 1e-12) / B)
    print(f"\n=== {name} ===")
    for r in range(R):
        print(f"  r={r}: model {mod[r][mask[r]].mean():.6f} | stim {st[r][mask[r]].mean():.6f}")
    print(f"  max deviation: {dev[mask].max():.2f} standard errors")


report("spam-idle only", make_nm(p_idle_spam_X=0.001996, p_idle_spam_Y=0.001996,
                                 p_idle_spam_Z=0.001996))
report("cnot-idle only", make_nm(p_idle_cnot_X=0.002, p_idle_cnot_Y=0.002,
                                 p_idle_cnot_Z=0.002))
report("meas only", make_nm(p_meas_X=0.004, p_meas_Z=0.004))
report("prep only", make_nm(p_prep_X=0.004, p_prep_Z=0.004))

[ivanbasov]

Approving — verification details in my review comment above. Thanks again for the excellent find and fix!