diff --git a/README.md b/README.md index 1f34366..57b62e8 100644 --- a/README.md +++ b/README.md @@ -205,6 +205,11 @@ If you are not training locally, you can run inference using pre-trained models. - `models/Ising-Decoder-SurfaceCode-1-Fast.pt` (receptive field R=9) - `models/Ising-Decoder-SurfaceCode-1-Accurate.pt` (receptive field R=13) + These checkpoints target the uniform circuit-level depolarizing setting + encoded by the public configs. Custom, non-uniform 25-parameter noise models + are supported for training by the pipeline below; they are a training-time + customization rather than a property of the shipped checkpoints. + Clones get the files via `git lfs pull`. Optionally, set `PREDECODER_MODEL_URL` to the LFS/raw URL to fetch files when not in the working tree (e.g. in a minimal checkout or CI). 3. Set: @@ -559,28 +564,39 @@ LOGICAL Z (lz): #### Noise model (public default) - `data.noise_model`: a **25-parameter circuit-level** noise model (SPAM, idles, and CNOT Pauli channels). +- The shipped configs use a **uniform circuit-level depolarizing** mapping, where all 25 values are derived from a single physical error rate `p` (for example `p_prep_{X,Z}=2*p/3`, `p_idle_cnot_{X,Y,Z}=p/3`, and `p_cnot_*=p/15`). +- You may edit `data.noise_model` to train on a non-uniform/custom 25-parameter model. In that case the Torch training generator refreshes the sampling probability vector from the active 25p model instead of collapsing back to the scalar uniform-depolarizing path. #### Training noise upscaling (surface code) -When training a surface-code pre-decoder the noise parameters you specify may be very small (e.g. `p = 1e-4`), which produces extremely sparse syndromes and slow convergence. To address this, the training pipeline **automatically upscales** all 25 noise-model parameters so that the largest grouped total `max(P_prep, P_meas, P_idle_cnot, P_idle_spam, P_cnot)` equals a fixed target of **6 × 10⁻³** (just below the surface-code threshold of ~7.5 × 10⁻³). +When training a surface-code pre-decoder the noise parameters you specify may be very small (e.g. `p = 1e-4`), which produces extremely sparse syndromes and slow convergence. To address this, the training pipeline **automatically upscales** all 25 noise-model parameters so that the largest *effective* fault-channel probability equals a fixed target of **6 × 10⁻³** (just below the surface-code threshold of ~7.5 × 10⁻³). -The five grouped totals are: +The seven channels considered (the "capital P's") are: -| Group | Sum of | -|-------|--------| -| P_prep | `p_prep_X + p_prep_Z` | -| P_meas | `p_meas_X + p_meas_Z` | +| Channel | Value | +|---------|-------| +| P_prep_X | `p_prep_X` | +| P_prep_Z | `p_prep_Z` | +| P_meas_X | `p_meas_X` | +| P_meas_Z | `p_meas_Z` | | P_idle_cnot | `p_idle_cnot_X + p_idle_cnot_Y + p_idle_cnot_Z` | -| P_idle_spam | `p_idle_spam_X + p_idle_spam_Y + p_idle_spam_Z` | +| P_idle_spam (effective) | `0.5 × (p_idle_spam_X + p_idle_spam_Y + p_idle_spam_Z)` | | P_cnot | sum of all 15 `p_cnot_*` | +`max_group = max(P_prep_X, P_prep_Z, P_meas_X, P_meas_Z, P_idle_cnot, P_idle_spam_effective, P_cnot)`. + +Two design notes: + +- **X / Z prep and measurement are kept separate.** They are independent one-Pauli fault channels — summing `p_prep_X + p_prep_Z` (or `p_meas_X + p_meas_Z`) double-counts the effective channel probability and would inflate `max_group` for an otherwise on-target depolarising noise model. +- **`p_idle_spam_*` is halved before the comparison.** The SPAM-window idle is built from a two-step model (one per state-prep and one per ancilla-reset half), so the raw configured total represents two depolarising steps. The scaling decision uses the per-step effective value `0.5 × p_idle_spam_raw`; the raw value is still reported in logs as `idle_spam_raw`. + **Upscaling rules:** - If `max_group < 6e-3`: all 25 p's are multiplied by `6e-3 / max_group` for training data generation only. Evaluation always uses the original user-specified noise model as-is. - If `max_group >= 6e-3`: parameters are **not** modified (the training log emits a warning in case this indicates a configuration error). - Non-surface-code types (`code_type != "surface_code"`) are never upscaled. -**Algorithm in brief:** The pipeline stores `p_max = max(P_prep, P_meas, P_idle_cnot, P_idle_spam, P_cnot)` from the full 25-parameter noise vector and rescales the entire vector by `0.006 / p_max` so that `p_max` is raised to **0.6%** (6 × 10⁻³). The original noise model is preserved unchanged for evaluation. +**Algorithm in brief:** The pipeline computes the seven channels above, takes `p_max = max(...)`, and rescales the entire 25-parameter vector by `0.006 / p_max` so that `p_max` is raised to **0.6%** (6 × 10⁻³). The original noise model is preserved unchanged for evaluation. We have found that training on denser syndromes and then evaluating on sparser data produces better results than training directly on sparse data. @@ -622,6 +638,14 @@ If frames are missing, the code can fall back to on-the-fly generation, but it i python3 code/data/precompute_frames.py --distance 13 --n_rounds 13 --basis X Z --rotation O1 ``` +Precomputed DEM/frame artifacts are structural: they encode which detector +responses each possible error column can produce for a given distance, number of +rounds, basis, and rotation. The active scalar or 25-parameter noise model +controls the per-column sampling probabilities. Therefore cached structural +artifacts can be reused when only the probabilities change; the training +generator refreshes the probability vector from the active noise model at load +time. + ### Resuming training and running inference on a trained model - **Inference uses the trained model from `outputs//models/`**, so keep the same `EXPERIMENT_NAME` when you switch from training to inference. diff --git a/code/data/generator_torch.py b/code/data/generator_torch.py index 0131748..41db2f4 100644 --- a/code/data/generator_torch.py +++ b/code/data/generator_torch.py @@ -14,6 +14,7 @@ # limitations under the License. import torch +from pathlib import Path class QCDataGeneratorTorch: @@ -76,17 +77,15 @@ def __init__( raise ValueError( "decompose_y is not supported in the Torch-only generator (set decompose_y=false)." ) - if noise_model is not None: - # TODO: wire noise_model through to precompute_dem_bundle_surface_code() - # so train/eval use the same 25-parameter noise distribution. - # build_single_p_marginal() already supports noise_model; only this - # generator-level plumbing is missing. - raise ValueError( - "noise_model is not supported in the Torch-only generator (simple single-p only)." - ) + self.noise_model = noise_model from qec.surface_code.memory_circuit_torch import MemoryCircuitTorch - from qec.precompute_dem import precompute_dem_bundle_surface_code + from qec.precompute_dem import ( + build_probability_vector_surface_code, + dem_artifact_metadata_matches, + load_dem_artifact_metadata, + precompute_dem_bundle_surface_code, + ) import threading self._early_compile_threads: list[threading.Thread] = [] @@ -110,17 +109,73 @@ def __init__( self._early_compile_threads.append(t) dem_cache = {} - if precomputed_frames_dir is None: - # Pick a nominal p for building the single-p marginal vector. - # (This matches existing behavior when using a precomputed directory.) - p_nom = float( - p_error if p_error is not None else (p_max if p_max is not None else 0.004) + p_overrides = {} + bases_needed = ["X", "Z"] if self._mixed else [self._single_basis] + p_nom = float(p_error if p_error is not None else (p_max if p_max is not None else 0.004)) + effective_precomputed_frames_dir = precomputed_frames_dir + metadata_reason = None + if precomputed_frames_dir is not None: + precomputed_dir = Path(precomputed_frames_dir) + for b in bases_needed: + p_path = ( + precomputed_dir / + f"surface_d{self.distance}_r{self.n_rounds}_{b}_frame_predecoder.p.npz" + ) + if not p_path.exists(): + # Asymmetric handling on purpose: + # - Scalar mode preserves the legacy behaviour of letting + # MemoryCircuitTorch raise its own clear FileNotFoundError + # downstream (so `continue` here keeps the dir intact). + # - 25p mode self-heals by falling back to an in-memory + # build, since the p artifact is needed to determine the + # cached error-column count. + if noise_model is None: + continue + effective_precomputed_frames_dir = None + metadata_reason = f"missing p artifact for basis {b}" + break + metadata = load_dem_artifact_metadata(p_path) + ok, reason = dem_artifact_metadata_matches( + metadata, + distance=self.distance, + n_rounds=self.n_rounds, + basis=b, + code_rotation=self.code_rotation, + p_scalar=p_nom, + noise_model=noise_model, + ) + if not ok: + effective_precomputed_frames_dir = None + metadata_reason = f"basis {b}: {reason}" + break + p_overrides[b] = torch.from_numpy( + build_probability_vector_surface_code( + distance=self.distance, + n_rounds=self.n_rounds, + basis=b, + code_rotation=self.code_rotation, + p_scalar=p_nom, + noise_model=noise_model, + ) + ) + + if effective_precomputed_frames_dir is None: + nm_tag = ", noise_model=25p" if noise_model is not None else "" + if precomputed_frames_dir is None: + source = "precomputed_frames_dir=None" + else: + source = f"precomputed DEM metadata mismatch ({metadata_reason})" + # Always announce a metadata-driven rebuild on rank 0 so silent + # rebuilds are visible in non-verbose distributed runs. The + # precomputed_frames_dir=None path is only logged when verbose. + should_log = self.verbose or ( + precomputed_frames_dir is not None and int(self.rank) == 0 ) - if self.verbose: + if should_log: print( - f"[QCDataGeneratorTorch] precomputed_frames_dir=None -> building in-memory DEM bundle at p={p_nom}" + f"[QCDataGeneratorTorch] {source} -> building in-memory DEM bundle " + f"at p={p_nom}{nm_tag}" ) - bases_needed = ["X", "Z"] if self._mixed else [self._single_basis] for b in bases_needed: dem_cache[b] = precompute_dem_bundle_surface_code( distance=self.distance, @@ -132,8 +187,17 @@ def __init__( device=self.device, export=False, return_artifacts=True, - # TODO: pass noise_model=noise_model here for circuit-level noise support + noise_model=noise_model, ) + elif self.verbose: + nm_tag = ( + f", refreshed_p_from_noise_model={noise_model.sha256()}" + if noise_model is not None else f", refreshed_p_from_scalar={p_nom:g}" + ) + print( + f"[QCDataGeneratorTorch] using disk DEM structure from " + f"{effective_precomputed_frames_dir}{nm_tag}" + ) _he_kwargs = dict( timelike_he=timelike_he, @@ -154,24 +218,26 @@ def __init__( distance=self.distance, n_rounds=self.n_rounds, basis="X", - precomputed_frames_dir=precomputed_frames_dir, + precomputed_frames_dir=effective_precomputed_frames_dir, code_rotation=self.code_rotation, device=self.device, H=(dem_cache.get("X", {}).get("H") if dem_cache else None), p=(dem_cache.get("X", {}).get("p") if dem_cache else None), A=(dem_cache.get("X", {}).get("A") if dem_cache else None), + p_override=p_overrides.get("X"), **_he_kwargs, ) self.sim_Z = MemoryCircuitTorch( distance=self.distance, n_rounds=self.n_rounds, basis="Z", - precomputed_frames_dir=precomputed_frames_dir, + precomputed_frames_dir=effective_precomputed_frames_dir, code_rotation=self.code_rotation, device=self.device, H=(dem_cache.get("Z", {}).get("H") if dem_cache else None), p=(dem_cache.get("Z", {}).get("p") if dem_cache else None), A=(dem_cache.get("Z", {}).get("A") if dem_cache else None), + p_override=p_overrides.get("Z"), **_he_kwargs, ) else: @@ -179,12 +245,13 @@ def __init__( distance=self.distance, n_rounds=self.n_rounds, basis=self._single_basis, - precomputed_frames_dir=precomputed_frames_dir, + precomputed_frames_dir=effective_precomputed_frames_dir, code_rotation=self.code_rotation, device=self.device, H=(dem_cache.get(self._single_basis, {}).get("H") if dem_cache else None), p=(dem_cache.get(self._single_basis, {}).get("p") if dem_cache else None), A=(dem_cache.get(self._single_basis, {}).get("A") if dem_cache else None), + p_override=p_overrides.get(self._single_basis), **_he_kwargs, ) diff --git a/code/data/precompute_frames.py b/code/data/precompute_frames.py index add8547..5f9d031 100644 --- a/code/data/precompute_frames.py +++ b/code/data/precompute_frames.py @@ -34,6 +34,7 @@ """ import argparse +import json from pathlib import Path import sys import time @@ -44,6 +45,7 @@ sys.path.insert(0, str(Path(__file__).parent.parent)) from qec.precompute_dem import precompute_dem_bundle_surface_code +from qec.noise_model import NoiseModel def _normalize_rotation(rotation: str) -> str: @@ -56,6 +58,42 @@ def _normalize_rotation(rotation: str) -> str: return internal_rot +def _load_config_mapping(path: str) -> dict: + path_obj = Path(path) + if path_obj.suffix.lower() == ".json": + with path_obj.open("r", encoding="utf-8") as f: + return json.load(f) + + try: + from omegaconf import OmegaConf + cfg = OmegaConf.load(path_obj) + return OmegaConf.to_container(cfg, resolve=True) + except Exception: + try: + import yaml + except ImportError as exc: + raise RuntimeError( + "YAML noise model configs require omegaconf or PyYAML to be installed" + ) from exc + with path_obj.open("r", encoding="utf-8") as f: + return yaml.safe_load(f) + + +def _load_noise_model(path: str) -> NoiseModel: + cfg = _load_config_mapping(path) + if not isinstance(cfg, dict): + raise ValueError(f"Noise model config must be a mapping, got {type(cfg).__name__}") + + if isinstance(cfg.get("data"), dict) and isinstance(cfg["data"].get("noise_model"), dict): + noise_model_cfg = cfg["data"]["noise_model"] + elif isinstance(cfg.get("noise_model"), dict): + noise_model_cfg = cfg["noise_model"] + else: + noise_model_cfg = cfg + + return NoiseModel.from_config_dict(noise_model_cfg) + + def main() -> None: parser = argparse.ArgumentParser( description="Precompute DEM bundles for MemoryCircuitTorch", @@ -101,6 +139,18 @@ def main() -> None: default=0.01, help="Scalar p for exporting single-p marginals", ) + parser.add_argument( + "--noise_model_config", + type=str, + default=None, + help="YAML/JSON config containing data.noise_model, noise_model, or a direct 25p mapping", + ) + parser.add_argument( + "--noise_model_json", + type=str, + default=None, + help="JSON file containing data.noise_model, noise_model, or a direct 25p mapping", + ) parser.add_argument( "--dem_output_dir", type=str, @@ -132,6 +182,13 @@ def main() -> None: args.dem_output_dir = args.output_dir if args.dem_output_dir is None: args.dem_output_dir = str(Path(__file__).parent.parent / "frames_data") + if args.noise_model_config is not None and args.noise_model_json is not None: + print("Error: Use only one of --noise_model_config or --noise_model_json") + sys.exit(1) + noise_model = None + noise_model_path = args.noise_model_config or args.noise_model_json + if noise_model_path is not None: + noise_model = _load_noise_model(noise_model_path) if args.n_rounds is None: args.n_rounds = args.distance @@ -158,6 +215,10 @@ def main() -> None: print(f"# Bases: {args.basis}") print(f"# Rotation: {args.rotation} (internal={internal_rot})") print(f"# Output: {args.dem_output_dir}") + if noise_model is None: + print(f"# Noise: scalar p={float(args.p)}") + else: + print(f"# Noise: 25p noise_model sha256={noise_model.sha256()}") print(f"{'#' * 60}") total_t0 = time.time() @@ -175,6 +236,7 @@ def main() -> None: dem_output_dir=str(args.dem_output_dir), device=device, export=True, + noise_model=noise_model, ) output_dirs.add(str(dem_dir)) except Exception as e: diff --git a/code/export/generate_test_data.py b/code/export/generate_test_data.py index 96d2ac0..47c773c 100644 --- a/code/export/generate_test_data.py +++ b/code/export/generate_test_data.py @@ -106,9 +106,9 @@ "p_idle_cnot_X": 0.001, "p_idle_cnot_Y": 0.001, "p_idle_cnot_Z": 0.001, - "p_idle_spam_X": 0.001998, - "p_idle_spam_Y": 0.001998, - "p_idle_spam_Z": 0.001998, + "p_idle_spam_X": 0.001996, + "p_idle_spam_Y": 0.001996, + "p_idle_spam_Z": 0.001996, "p_cnot_IX": 0.0002, "p_cnot_IY": 0.0002, "p_cnot_IZ": 0.0002, diff --git a/code/qec/noise_model.py b/code/qec/noise_model.py index 7d0593d..4805a95 100644 --- a/code/qec/noise_model.py +++ b/code/qec/noise_model.py @@ -43,6 +43,9 @@ from dataclasses import dataclass, field, asdict from typing import Dict, List, Optional, Tuple, Any +import hashlib +import json +import math import numpy as np # Surface-code training upscale target (below threshold ~7.5e-3). Used when sampling training data. @@ -194,6 +197,23 @@ def to_config_dict(self) -> Dict[str, float]: """ return {k: v for k, v in asdict(self).items() if not k.startswith("_")} + def canonical_parameters(self) -> Dict[str, float]: + """Stable public 25p parameter mapping for metadata and hashing.""" + return {k: float(v) for k, v in sorted(self.to_config_dict().items())} + + def canonical_json(self) -> str: + """Stable JSON representation of public parameters only.""" + return json.dumps( + self.canonical_parameters(), + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ) + + def sha256(self) -> str: + """SHA-256 of the canonical public 25p parameter JSON.""" + return hashlib.sha256(self.canonical_json().encode("utf-8")).hexdigest() + def copy(self) -> "NoiseModel": """Deep-ish copy preserving reference parameters.""" nm = NoiseModel.from_config_dict(self.to_config_dict()) @@ -400,22 +420,44 @@ def __repr__(self) -> str: def get_grouped_totals(nm: NoiseModel) -> Dict[str, float]: """ - Compute the sum of p's per fault channel (capital P's) for the 25-p noise model. + Compute effective fault-channel totals (capital P's) for 25-p training scaling. Returns: - Dict with keys: p_prep, p_meas, p_idle_cnot, p_idle_spam, p_cnot, max_group. + Dict with separate prep/meas channels, idle/cnot totals, and max_group. + + Notes: + - X/Z prep and measurement are separate one-Pauli fault channels; summing + them would double-count the effective channel probability. + - p_idle_spam_* models a two-step SPAM window, so the scaling decision uses + half the raw total while still reporting the raw configured total. """ - p_prep = float(nm.p_prep_X + nm.p_prep_Z) - p_meas = float(nm.p_meas_X + nm.p_meas_Z) - p_idle_cnot = float(nm.get_total_idle_cnot_probability()) - p_idle_spam = float(nm.get_total_idle_spam_probability()) - p_cnot = float(nm.get_total_cnot_probability()) - max_group = max(p_prep, p_meas, p_idle_cnot, p_idle_spam, p_cnot) + p_prep_X = float(nm.p_prep_X) + p_prep_Z = float(nm.p_prep_Z) + p_meas_X = float(nm.p_meas_X) + p_meas_Z = float(nm.p_meas_Z) + p_idle_cnot = math.fsum(float(p) for p in nm.get_idle_cnot_probabilities()) + p_idle_spam_raw = math.fsum(float(p) for p in nm.get_idle_spam_probabilities()) + p_idle_spam_effective = 0.5 * p_idle_spam_raw + p_cnot = math.fsum(float(p) for p in nm.get_cnot_probabilities()) + max_group = max( + p_prep_X, + p_prep_Z, + p_meas_X, + p_meas_Z, + p_idle_cnot, + p_idle_spam_effective, + p_cnot, + ) return { - "p_prep": p_prep, - "p_meas": p_meas, + "p_prep_X": p_prep_X, + "p_prep_Z": p_prep_Z, + "p_meas_X": p_meas_X, + "p_meas_Z": p_meas_Z, + "p_prep_total": p_prep_X + p_prep_Z, + "p_meas_total": p_meas_X + p_meas_Z, "p_idle_cnot": p_idle_cnot, - "p_idle_spam": p_idle_spam, + "p_idle_spam_raw": p_idle_spam_raw, + "p_idle_spam_effective": p_idle_spam_effective, "p_cnot": p_cnot, "max_group": max_group, } @@ -448,7 +490,7 @@ def get_training_upscaled_noise_model( - applied_upscale: bool - scale_factor: float (only if upscaling applied) - max_group: float - - group_totals: dict (p_prep, p_meas, ...) + - group_totals: dict (p_prep_X, p_prep_Z, p_meas_X, p_meas_Z, ...) - above_target_warning: bool (max_group > UPSCALE_TARGET) - downscale_skipped: bool (max_group > target, params not modified) - skipped_by_user: bool (skip_upscale was True) @@ -481,8 +523,11 @@ def get_training_upscaled_noise_model( if max_group <= 0.0: raise ValueError( "Invalid noise_model: all grouped totals are <= 0 " - f"(prep={totals['p_prep']}, meas={totals['p_meas']}, " - f"idle_cnot={totals['p_idle_cnot']}, idle_spam={totals['p_idle_spam']}, cnot={totals['p_cnot']})." + f"(prep_X={totals['p_prep_X']}, prep_Z={totals['p_prep_Z']}, " + f"meas_X={totals['p_meas_X']}, meas_Z={totals['p_meas_Z']}, " + f"idle_cnot={totals['p_idle_cnot']}, " + f"idle_spam_effective={totals['p_idle_spam_effective']}, " + f"cnot={totals['p_cnot']})." ) scale_factor = target / max_group diff --git a/code/qec/precompute_dem.py b/code/qec/precompute_dem.py index d3c3332..e5f7bce 100644 --- a/code/qec/precompute_dem.py +++ b/code/qec/precompute_dem.py @@ -30,8 +30,9 @@ from __future__ import annotations import argparse +import json from pathlib import Path -from typing import Iterable, List, Tuple +from typing import Any, Iterable, List, Tuple import sys @@ -47,6 +48,99 @@ # Stim parsing helpers (pure python) # ============================================================================= +DEM_ARTIFACT_METADATA_VERSION = 1 +DEM_ARTIFACT_METADATA_KEY = "metadata_json" + + +def build_dem_artifact_metadata( + *, + distance: int, + n_rounds: int, + basis: str, + code_rotation: str, + p_scalar: float, + noise_model=None, +) -> dict[str, Any]: + """Build metadata for a DEM artifact. + + The structural keys identify whether cached H/A frame artifacts can be + reused. The probability keys record provenance for the stored p vector, but + probability changes alone do not invalidate the structural DEM. + """ + metadata: dict[str, Any] = { + "schema_version": DEM_ARTIFACT_METADATA_VERSION, + "distance": int(distance), + "n_rounds": int(n_rounds), + "basis": str(basis).upper(), + "code_rotation": str(code_rotation).upper(), + } + if noise_model is None: + metadata.update({ + "noise_mode": "scalar", + "p_scalar": float(p_scalar), + }) + else: + metadata.update( + { + "noise_mode": "noise_model", + "p_scalar_placeholder": float(p_scalar), + "noise_model_sha256": noise_model.sha256(), + "noise_model": noise_model.canonical_parameters(), + } + ) + return metadata + + +def encode_dem_artifact_metadata(metadata: dict[str, Any]) -> str: + """Serialize metadata in a deterministic JSON form.""" + return json.dumps(metadata, sort_keys=True, separators=(",", ":"), allow_nan=False) + + +def decode_dem_artifact_metadata(value) -> dict[str, Any]: + """Decode metadata loaded from an npz scalar/string array.""" + if isinstance(value, np.ndarray): + value = value.item() if value.shape == () else value.reshape(-1)[0] + if isinstance(value, bytes): + value = value.decode("utf-8") + return json.loads(str(value)) + + +def load_dem_artifact_metadata(p_path: Path) -> dict[str, Any] | None: + """Return metadata from a .p.npz file, or None for legacy artifacts.""" + with np.load(p_path, allow_pickle=False) as z: + if DEM_ARTIFACT_METADATA_KEY not in z.files: + return None + return decode_dem_artifact_metadata(z[DEM_ARTIFACT_METADATA_KEY]) + + +def dem_artifact_metadata_matches( + metadata: dict[str, Any] | None, + *, + distance: int, + n_rounds: int, + basis: str, + code_rotation: str, + p_scalar: float, + noise_model=None, +) -> tuple[bool, str]: + """Check whether on-disk DEM metadata matches the requested structure.""" + if metadata is None: + return True, "legacy artifact without structural metadata" + + expected = build_dem_artifact_metadata( + distance=distance, + n_rounds=n_rounds, + basis=basis, + code_rotation=code_rotation, + p_scalar=p_scalar, + noise_model=noise_model, + ) + for key in ("schema_version", "distance", "n_rounds", "basis", "code_rotation"): + if metadata.get(key) != expected.get(key): + return False, f"metadata {key}={metadata.get(key)!r} != expected {expected.get(key)!r}" + + return True, "structural metadata matches" + def extract_cnot_structure_from_stim_text(circuit_string: str) -> tuple[np.ndarray, np.ndarray]: """ @@ -542,7 +636,20 @@ def build_single_p_marginal( p_err[eidx] = float(_nm_cnot.get(et, 0.0)) elif len(et) == 1: if is_ancilla_prep: - p_err[eidx] = 0.0 + # Stim emits Z_ERROR(p_prep_X) on X-basis-reset ancillas and + # X_ERROR(p_prep_Z) on Z-basis-reset ancillas (see MemoryCircuit. + # add_reset / add_single_error). Treat prep as its own one-Pauli + # fault channel, consistent with get_grouped_totals' "X/Z prep + # and measurement are separate one-Pauli fault channels" rule + # and with linear behaviour under uniform noise upscaling. + prep_basis = int(prep_basis_map[(r, q)]) + if prep_basis == 0: + allowed = (et == "Z") + else: + allowed = (et == "X") + p_err[eidx] = float( + p_prep_X if prep_basis == 0 else p_prep_Z + ) if allowed else 0.0 elif is_ancilla_meas: meas_basis = int(meas_basis_map[(r, q)]) if meas_basis == 0: @@ -616,6 +723,75 @@ def build_single_p_marginal( return p_err +def build_probability_vector_surface_code( + *, + distance: int, + n_rounds: int, + basis: str, + code_rotation: str, + p_scalar: float, + noise_model=None, +) -> np.ndarray: + """Build only the per-error probability vector for a surface-code DEM. + + Cached DEM frame artifacts encode the possible detector responses. The + active scalar/noise model only determines how likely each error column is, + so callers can reuse cached H/A artifacts and refresh p with this helper. + """ + from qec.surface_code.memory_circuit import MemoryCircuit + + distance = int(distance) + n_rounds = int(n_rounds) + basis = str(basis).upper() + code_rotation = str(code_rotation).upper() + p_scalar = float(p_scalar) + + circ = MemoryCircuit( + distance=distance, + idle_error=p_scalar, + sqgate_error=p_scalar, + tqgate_error=p_scalar, + spam_error=2.0 / 3.0 * p_scalar, + n_rounds=n_rounds, + basis=basis, + code_rotation=code_rotation, + noise_model=noise_model, + ) + circ.set_error_rates() + cnot_circuit, cx_times = extract_cnot_structure_from_stim_text(circ.circuit) + t_total = int(len(cx_times) + 2) + nq = int(2 * distance * distance - 1) + + data_qubits = np.array(circ.code.data_qubits, dtype=np.int32) + xcheck_qubits = np.array(circ.code.xcheck_qubits, dtype=np.int32) + zcheck_qubits = np.array(circ.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), np.int32), + np.ones(len(zcheck_qubits), np.int32)] + ).astype(np.int32) + + _, 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 + ) + return 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, + ).astype(np.float32) + + # ============================================================================= # End-to-end entrypoint # ============================================================================= @@ -784,8 +960,19 @@ def precompute_dem_bundle_surface_code( np.savez_compressed(dem_dir / f"{prefix}.X.npz", HX=HX.cpu().numpy().astype(np.uint8)) np.savez_compressed(dem_dir / f"{prefix}.Z.npz", HZ=HZ.cpu().numpy().astype(np.uint8)) + metadata = build_dem_artifact_metadata( + distance=distance, + n_rounds=n_rounds, + basis=basis, + code_rotation=code_rotation, + p_scalar=p_scalar, + noise_model=noise_model, + ) np.savez_compressed( - dem_dir / f"{prefix}.p.npz", p=p_err, p_nominal=np.array(p_scalar, dtype=np.float32) + dem_dir / f"{prefix}.p.npz", + p=p_err, + p_nominal=np.array(p_scalar, dtype=np.float32), + **{DEM_ARTIFACT_METADATA_KEY: np.array(encode_dem_artifact_metadata(metadata))}, ) np.savez_compressed(dem_dir / f"{prefix}.A.npz", A=A.astype(np.uint8)) return dem_dir @@ -803,6 +990,18 @@ def main() -> None: ap.add_argument( "--p", type=float, default=0.01, help="Scalar p for exporting single-p marginals" ) + ap.add_argument( + "--noise_model_config", + type=str, + default=None, + help="YAML/JSON config containing data.noise_model, noise_model, or a direct 25p mapping", + ) + ap.add_argument( + "--noise_model_json", + type=str, + default=None, + help="JSON file containing data.noise_model, noise_model, or a direct 25p mapping", + ) ap.add_argument("--dem_output_dir", type=str, default=None) ap.add_argument( "--no_save", action="store_true", help="Run precompute but do not write any files" @@ -812,6 +1011,16 @@ def main() -> None: ) args = ap.parse_args() + if args.noise_model_config is not None and args.noise_model_json is not None: + ap.error("Use only one of --noise_model_config or --noise_model_json") + noise_model = None + noise_model_path = args.noise_model_config or args.noise_model_json + if noise_model_path is not None: + # Defer the import to avoid a circular dependency with data.precompute_frames + # at module-import time (precompute_frames.py imports from this module). + from data.precompute_frames import _load_noise_model + noise_model = _load_noise_model(noise_model_path) + d = int(args.distance) r = int(args.n_rounds) if args.n_rounds is not None else d dev = ( @@ -827,6 +1036,7 @@ def main() -> None: dem_output_dir=(str(args.dem_output_dir) if args.dem_output_dir is not None else None), device=dev, export=(not bool(args.no_save)), + noise_model=noise_model, ) diff --git a/code/qec/surface_code/memory_circuit_torch.py b/code/qec/surface_code/memory_circuit_torch.py index 1fe0a14..d688ab2 100644 --- a/code/qec/surface_code/memory_circuit_torch.py +++ b/code/qec/surface_code/memory_circuit_torch.py @@ -88,6 +88,7 @@ def __init__( H: torch.Tensor | None = None, # (2*num_detectors, num_errors) uint8 p: torch.Tensor | None = None, # (num_errors,) float32 A: torch.Tensor | None = None, # (n_rounds*num_meas, 2*num_detectors) uint8 + p_override: torch.Tensor | np.ndarray | None = None, ): self.distance = int(distance) self.n_rounds = int(n_rounds) @@ -184,6 +185,17 @@ def __init__( hx = np.asarray(hx, dtype=np.uint8) hz = np.asarray(hz, dtype=np.uint8) p_arr = np.asarray(p_arr).reshape(-1) + if p_override is not None: + if isinstance(p_override, torch.Tensor): + override_src = p_override.detach().cpu().numpy() + else: + override_src = p_override + override = np.asarray(override_src).reshape(-1) + if int(override.shape[0]) != errors: + raise ValueError( + f"p_override length {override.shape[0]} != DEM artifact error count {errors}" + ) + p_arr = override hx = hx if hx.shape[1] == errors else hx.T hz = hz if hz.shape[1] == errors else hz.T self.H = torch.from_numpy(np.concatenate([hx, hz], diff --git a/code/tests/test_noise_model.py b/code/tests/test_noise_model.py index 05e4b49..ce5982e 100644 --- a/code/tests/test_noise_model.py +++ b/code/tests/test_noise_model.py @@ -27,6 +27,7 @@ import os import sys +import tempfile import unittest from pathlib import Path @@ -280,6 +281,19 @@ def test_noise_model_roundtrip_and_invariants(self): with self.assertRaises(ValueError): NoiseModel(p_prep_X=1.5) + def test_canonical_noise_model_hash_uses_public_parameters_only(self): + nm = _noise_model_from_p(0.006) + nm_copy = nm.copy() + self.assertEqual(nm.sha256(), nm_copy.sha256()) + self.assertNotIn("_reference", nm.canonical_parameters()) + + nm_copy._reference = {k: float(v) * 1.7 for k, v in nm_copy._reference.items()} + self.assertEqual(nm.sha256(), nm_copy.sha256()) + + changed = nm.copy() + changed.p_prep_X += 1e-9 + self.assertNotEqual(nm.sha256(), changed.sha256()) + def test_stim_circuit_audit_no_cnot_noise_in_logical_measurement_section(self): # Non-trivial noise model: ensure PAULI_CHANNEL_2 appears in repeat block but NOT after it. D = 5 @@ -431,24 +445,180 @@ class TestNoiseModelUpscaling(unittest.TestCase): """Tests for surface-code training noise model upscaling (get_training_upscaled_noise_model).""" def test_get_grouped_totals(self): - """get_grouped_totals returns correct P_prep, P_meas, P_idle_cnot, P_idle_spam, P_cnot and max_group.""" + """get_grouped_totals returns effective channels used for training scaling.""" nm = _noise_model_from_p(0.01) tot = get_grouped_totals(nm) - self.assertAlmostEqual( - tot["p_prep"], 2.0 * 0.01 / 3.0 * 2, places=12 - ) # p_prep_X + p_prep_Z - self.assertAlmostEqual(tot["p_meas"], 2.0 * 0.01 / 3.0 * 2, places=12) + self.assertAlmostEqual(tot["p_prep_X"], 2.0 * 0.01 / 3.0, places=12) + self.assertAlmostEqual(tot["p_prep_Z"], 2.0 * 0.01 / 3.0, places=12) + self.assertAlmostEqual(tot["p_meas_X"], 2.0 * 0.01 / 3.0, places=12) + self.assertAlmostEqual(tot["p_meas_Z"], 2.0 * 0.01 / 3.0, places=12) + self.assertAlmostEqual(tot["p_prep_total"], 2.0 * 0.01 / 3.0 * 2, places=12) + self.assertAlmostEqual(tot["p_meas_total"], 2.0 * 0.01 / 3.0 * 2, places=12) self.assertAlmostEqual(tot["p_idle_cnot"], 0.01, places=12) - self.assertAlmostEqual(tot["p_idle_spam"], nm.get_total_idle_spam_probability(), places=12) + self.assertAlmostEqual( + tot["p_idle_spam_raw"], nm.get_total_idle_spam_probability(), places=12 + ) + self.assertAlmostEqual( + tot["p_idle_spam_effective"], nm.get_total_idle_spam_probability() / 2.0, places=12 + ) self.assertAlmostEqual(tot["p_cnot"], 0.01, places=12) self.assertGreater(tot["max_group"], 0) self.assertEqual( tot["max_group"], max( - tot["p_prep"], tot["p_meas"], tot["p_idle_cnot"], tot["p_idle_spam"], tot["p_cnot"] + tot["p_prep_X"], + tot["p_prep_Z"], + tot["p_meas_X"], + tot["p_meas_Z"], + tot["p_idle_cnot"], + tot["p_idle_spam_effective"], + tot["p_cnot"], + ) + ) + + def test_depolarizing_p006_has_target_effective_max_group(self): + """The p=6e-3 config should not look above target due to channel double-counting.""" + nm = NoiseModel( + 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_{k}": 0.0004 for k in CNOT_ERROR_TYPES} + ) + tot = get_grouped_totals(nm) + self.assertAlmostEqual(tot["p_prep_X"], 0.004, places=12) + self.assertAlmostEqual(tot["p_prep_Z"], 0.004, places=12) + self.assertAlmostEqual(tot["p_meas_X"], 0.004, places=12) + self.assertAlmostEqual(tot["p_meas_Z"], 0.004, places=12) + self.assertAlmostEqual(tot["p_idle_cnot"], 0.006, places=12) + self.assertAlmostEqual(tot["p_idle_spam_raw"], 0.011952, places=12) + self.assertAlmostEqual(tot["p_idle_spam_effective"], 0.005976, places=12) + self.assertAlmostEqual(tot["p_cnot"], 0.006, places=12) + self.assertAlmostEqual(tot["max_group"], SURFACE_CODE_TRAINING_UPSCALE_TARGET, places=12) + + training_nm, info = get_training_upscaled_noise_model(nm, code_type="surface_code") + self.assertTrue(info["applied_upscale"]) + self.assertFalse(info["downscale_skipped"]) + self.assertFalse(info["above_target_warning"]) + self.assertAlmostEqual(info["scale_factor"], 1.0, places=12) + self.assertEqual(training_nm.to_config_dict(), nm.to_config_dict()) + + def test_precomputed_dem_probability_vector_uses_25p_values(self): + """DEM precompute should build p from the explicit 25p model, not scalar p/3 or p/15.""" + import torch + from qec.precompute_dem import precompute_dem_bundle_surface_code + + cnot_probs = {f"p_cnot_{k}": 0.00011 + i * 0.00001 for i, k in enumerate(CNOT_ERROR_TYPES)} + nm = NoiseModel( + p_prep_X=0.0011, + p_prep_Z=0.0022, + p_meas_X=0.0033, + p_meas_Z=0.0044, + p_idle_cnot_X=0.0051, + p_idle_cnot_Y=0.0052, + p_idle_cnot_Z=0.0053, + p_idle_spam_X=0.0061, + p_idle_spam_Y=0.0062, + p_idle_spam_Z=0.0063, + **cnot_probs + ) + + observed = [] + for basis in ("X", "Z"): + artifacts = precompute_dem_bundle_surface_code( + distance=3, + n_rounds=3, + basis=basis, + code_rotation="XV", + p_scalar=0.1234, + dem_output_dir=None, + device=torch.device("cpu"), + export=False, + return_artifacts=True, + noise_model=nm, ) + observed.append(artifacts["p"].cpu().numpy()) + p_values = np.concatenate(observed) + + expected_values = [ + nm.p_prep_X, + nm.p_prep_Z, + nm.p_meas_X, + nm.p_meas_Z, + nm.p_idle_cnot_X, + nm.p_idle_cnot_Y, + nm.p_idle_cnot_Z, + nm.p_idle_spam_X, + nm.p_idle_spam_Y, + nm.p_idle_spam_Z, + nm.p_cnot_IX, + nm.p_cnot_ZZ, + ] + for expected in expected_values: + self.assertTrue( + np.any(np.isclose(p_values, expected, rtol=0.0, atol=1e-9)), + f"Expected 25p probability {expected} in DEM p vector", + ) + + scalar_derived_values = [0.1234 / 3.0, 0.1234 / 15.0, 2.0 * 0.1234 / 3.0] + for scalar_value in scalar_derived_values: + self.assertFalse( + np.any(np.isclose(p_values, scalar_value, rtol=0.0, atol=1e-9)), + f"Unexpected scalar-derived probability {scalar_value} in 25p DEM p vector", + ) + + def test_precompute_dem_export_writes_noise_metadata(self): + import torch + from qec.precompute_dem import ( + load_dem_artifact_metadata, + precompute_dem_bundle_surface_code, ) + with tempfile.TemporaryDirectory() as tmp: + precompute_dem_bundle_surface_code( + distance=3, + n_rounds=3, + basis="X", + code_rotation="XV", + p_scalar=0.004, + dem_output_dir=tmp, + device=torch.device("cpu"), + export=True, + ) + scalar_meta = load_dem_artifact_metadata( + Path(tmp) / "surface_d3_r3_X_frame_predecoder.p.npz" + ) + self.assertEqual(scalar_meta["noise_mode"], "scalar") + self.assertEqual(scalar_meta["distance"], 3) + self.assertEqual(scalar_meta["basis"], "X") + self.assertAlmostEqual(float(scalar_meta["p_scalar"]), 0.004, places=12) + + nm = _noise_model_from_p(0.005) + with tempfile.TemporaryDirectory() as tmp: + precompute_dem_bundle_surface_code( + distance=3, + n_rounds=3, + basis="Z", + code_rotation="XV", + p_scalar=0.123, + dem_output_dir=tmp, + device=torch.device("cpu"), + export=True, + noise_model=nm, + ) + nm_meta = load_dem_artifact_metadata( + Path(tmp) / "surface_d3_r3_Z_frame_predecoder.p.npz" + ) + self.assertEqual(nm_meta["noise_mode"], "noise_model") + self.assertEqual(nm_meta["noise_model_sha256"], nm.sha256()) + self.assertEqual(nm_meta["noise_model"], nm.canonical_parameters()) + def test_upscale_small_noise(self): """When max_group < target, all 25 p's are scaled so that new max_group = target.""" # Single-p 1e-4 -> max_group is around 1e-4 (order of magnitude) diff --git a/code/tests/test_torch_setup.py b/code/tests/test_torch_setup.py index 971753e..e1b5b7f 100644 --- a/code/tests/test_torch_setup.py +++ b/code/tests/test_torch_setup.py @@ -5,6 +5,8 @@ # Runs in CI with code/tests (PYTHONPATH=code). import sys +import json +import tempfile import unittest from pathlib import Path @@ -14,6 +16,7 @@ sys.path.insert(0, str(_repo_code)) import torch +import numpy as np import qec.dem_sampling as _dem_mod @@ -89,3 +92,285 @@ def test_generator_init_and_batch(self): trainX, trainY = gen.generate_batch(step=0, batch_size=2) self.assertEqual(trainX.dim(), 5) self.assertEqual(trainY.dim(), 5) + + def test_generator_uses_noise_model_for_in_memory_dem_precompute(self): + from data.generator_torch import QCDataGeneratorTorch + from qec.noise_model import NoiseModel, CNOT_ERROR_TYPES + + nm = 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_{k}": 0.0001 for k in CNOT_ERROR_TYPES} + ) + + gen = QCDataGeneratorTorch( + distance=3, + n_rounds=3, + p_error=0.004, + measure_basis="both", + rank=0, + mode="train", + verbose=False, + timelike_he=False, + decompose_y=False, + precomputed_frames_dir=None, + code_rotation="XV", + noise_model=nm, + device=torch.device("cpu"), + ) + + self.assertIs(gen.noise_model, nm) + p_values = torch.cat([gen.sim_X.p.cpu(), gen.sim_Z.p.cpu()]) + for expected in (nm.p_prep_X, nm.p_prep_Z, nm.p_meas_X, nm.p_meas_Z, nm.p_idle_spam_Z): + self.assertTrue( + torch.isclose(p_values, torch.tensor(expected, dtype=p_values.dtype)).any(), + f"Expected 25p probability {expected} in generated DEM p vector", + ) + + def test_disk_dem_refreshes_probabilities_from_active_noise_model(self): + from data.generator_torch import QCDataGeneratorTorch + from qec.noise_model import NoiseModel, CNOT_ERROR_TYPES + from qec.precompute_dem import ( + DEM_ARTIFACT_METADATA_KEY, + precompute_dem_bundle_surface_code, + ) + + nm = 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_{k}": 0.0001 for k in CNOT_ERROR_TYPES} + ) + + with tempfile.TemporaryDirectory() as tmp: + precompute_dem_bundle_surface_code( + distance=3, + n_rounds=3, + basis="X", + code_rotation="XV", + p_scalar=0.004, + dem_output_dir=tmp, + device=torch.device("cpu"), + export=True, + noise_model=nm, + ) + p_path = Path(tmp) / "surface_d3_r3_X_frame_predecoder.p.npz" + with np.load(p_path, allow_pickle=False) as z: + p_arr = z["p"].copy() + p_nominal = z["p_nominal"].copy() + metadata_json = z[DEM_ARTIFACT_METADATA_KEY].copy() + sentinel = np.float32(0.987654) + p_arr[0] = sentinel + np.savez_compressed( + p_path, + p=p_arr, + p_nominal=p_nominal, + **{DEM_ARTIFACT_METADATA_KEY: metadata_json}, + ) + + gen = QCDataGeneratorTorch( + distance=3, + n_rounds=3, + p_error=0.004, + measure_basis="X", + rank=0, + mode="train", + verbose=False, + timelike_he=False, + decompose_y=False, + precomputed_frames_dir=tmp, + code_rotation="XV", + noise_model=nm, + device=torch.device("cpu"), + ) + + self.assertIs(gen.noise_model, nm) + self.assertFalse(torch.isclose(gen.sim.p.cpu(), torch.tensor(float(sentinel))).any()) + self.assertTrue( + torch.isclose(gen.sim.p.cpu(), torch.tensor(nm.p_prep_X)).any(), + "Expected cached DEM structure to use probabilities refreshed from the active noise_model", + ) + + def test_mismatched_noise_model_metadata_reuses_structure_with_refreshed_p(self): + from data.generator_torch import QCDataGeneratorTorch + from qec.noise_model import NoiseModel, CNOT_ERROR_TYPES + from qec.precompute_dem import ( + DEM_ARTIFACT_METADATA_KEY, + precompute_dem_bundle_surface_code, + ) + from unittest import mock + + nm_disk = 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_{k}": 0.0001 for k in CNOT_ERROR_TYPES} + ) + nm_active = nm_disk.copy() + nm_active.p_prep_X += 1e-6 + + with tempfile.TemporaryDirectory() as tmp: + precompute_dem_bundle_surface_code( + distance=3, + n_rounds=3, + basis="X", + code_rotation="XV", + p_scalar=0.004, + dem_output_dir=tmp, + device=torch.device("cpu"), + export=True, + noise_model=nm_disk, + ) + p_path = Path(tmp) / "surface_d3_r3_X_frame_predecoder.p.npz" + with np.load(p_path, allow_pickle=False) as z: + p_arr = z["p"].copy() + p_nominal = z["p_nominal"].copy() + metadata_json = z[DEM_ARTIFACT_METADATA_KEY].copy() + sentinel = np.float32(0.987654) + p_arr[0] = sentinel + np.savez_compressed( + p_path, + p=p_arr, + p_nominal=p_nominal, + **{DEM_ARTIFACT_METADATA_KEY: metadata_json}, + ) + + with mock.patch("qec.precompute_dem.precompute_dem_bundle_surface_code") as rebuild: + rebuild.side_effect = AssertionError("should reuse cached structural DEM") + gen = QCDataGeneratorTorch( + distance=3, + n_rounds=3, + p_error=0.004, + measure_basis="X", + rank=0, + mode="train", + verbose=False, + timelike_he=False, + decompose_y=False, + precomputed_frames_dir=tmp, + code_rotation="XV", + noise_model=nm_active, + device=torch.device("cpu"), + ) + rebuild.assert_not_called() + + self.assertIs(gen.noise_model, nm_active) + self.assertFalse(torch.isclose(gen.sim.p.cpu(), torch.tensor(float(sentinel))).any()) + self.assertTrue( + torch.isclose(gen.sim.p.cpu(), torch.tensor(nm_active.p_prep_X)).any(), + "Expected cached DEM structure to use probabilities refreshed from the active noise_model", + ) + + def test_legacy_scalar_metadata_free_frames_still_load(self): + from data.generator_torch import QCDataGeneratorTorch + from qec.precompute_dem import ( + build_probability_vector_surface_code, + precompute_dem_bundle_surface_code, + ) + + with tempfile.TemporaryDirectory() as tmp: + precompute_dem_bundle_surface_code( + distance=3, + n_rounds=3, + basis="X", + code_rotation="XV", + p_scalar=0.004, + dem_output_dir=tmp, + device=torch.device("cpu"), + export=True, + ) + p_path = Path(tmp) / "surface_d3_r3_X_frame_predecoder.p.npz" + with np.load(p_path, allow_pickle=False) as z: + p_arr = z["p"].copy() + p_nominal = z["p_nominal"].copy() + sentinel = np.float32(0.876543) + p_arr[0] = sentinel + np.savez_compressed(p_path, p=p_arr, p_nominal=p_nominal) + + gen = QCDataGeneratorTorch( + distance=3, + n_rounds=3, + p_error=0.004, + measure_basis="X", + rank=0, + mode="train", + verbose=False, + timelike_he=False, + decompose_y=False, + precomputed_frames_dir=tmp, + code_rotation="XV", + device=torch.device("cpu"), + ) + + self.assertFalse(torch.isclose(gen.sim.p.cpu(), torch.tensor(float(sentinel))).any()) + # The legacy (metadata-free) artifact must be refreshed from the active + # scalar p_error rather than reused as-is. The freshly rebuilt scalar p + # vector is the ground truth here; in particular its max is the ancilla + # syndrome-readout probability 2*spam*(1-spam) (with spam = 2/3 * p), + # which is ~5.32e-3 for p=0.004 and therefore strictly larger than p. + expected_p = build_probability_vector_surface_code( + distance=3, + n_rounds=3, + basis="X", + code_rotation="XV", + p_scalar=0.004, + noise_model=None, + ) + self.assertTrue( + torch.allclose( + gen.sim.p.cpu(), + torch.from_numpy(expected_p).to(gen.sim.p.cpu().dtype), + ) + ) + + def test_precompute_frames_loads_nested_noise_model_config(self): + from data.precompute_frames import _load_noise_model + from qec.noise_model import NoiseModel, CNOT_ERROR_TYPES + + nm = 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_{k}": 0.0001 for k in CNOT_ERROR_TYPES} + ) + + with tempfile.TemporaryDirectory() as tmp: + config_path = Path(tmp) / "noise_model.json" + config_path.write_text( + json.dumps({"data": { + "noise_model": nm.to_config_dict() + }}), + encoding="utf-8", + ) + loaded = _load_noise_model(str(config_path)) + + self.assertEqual(loaded.sha256(), nm.sha256()) diff --git a/code/training/train.py b/code/training/train.py index 11e8895..294e367 100644 --- a/code/training/train.py +++ b/code/training/train.py @@ -781,26 +781,25 @@ def init_process_group_with_timeout(*args, **kwargs): if nm_dict is not None: noise_model_user_obj = NoiseModel.from_config_dict(dict(nm_dict)) - # Compute grouped noise totals for logging and scaling decisions. - p_prep = float(noise_model_user_obj.p_prep_X + noise_model_user_obj.p_prep_Z) - p_meas = float(noise_model_user_obj.p_meas_X + noise_model_user_obj.p_meas_Z) - p_idle_cnot = float(noise_model_user_obj.get_total_idle_cnot_probability()) - p_idle_spam = float(noise_model_user_obj.get_total_idle_spam_probability()) - p_cnot = float(noise_model_user_obj.get_total_cnot_probability()) - max_group = max(p_prep, p_meas, p_idle_cnot, p_idle_spam, p_cnot) - if max_group <= 0.0: - raise ValueError( - "Invalid noise_model: all grouped totals are <= 0 " - f"(prep={p_prep}, meas={p_meas}, idle_cnot={p_idle_cnot}, idle_spam={p_idle_spam}, cnot={p_cnot})." - ) - # Surface-code training upscaling: bring max(P's) to target for training # data only. Evaluation uses the user-specified noise model as-is. from qec.noise_model import ( + get_grouped_totals, get_training_upscaled_noise_model, SURFACE_CODE_TRAINING_UPSCALE_TARGET, SURFACE_CODE_THRESHOLD_APPROX, ) + group_totals = get_grouped_totals(noise_model_user_obj) + max_group = float(group_totals["max_group"]) + if max_group <= 0.0: + raise ValueError( + "Invalid noise_model: all grouped totals are <= 0 " + f"(prep_X={group_totals['p_prep_X']}, prep_Z={group_totals['p_prep_Z']}, " + f"meas_X={group_totals['p_meas_X']}, meas_Z={group_totals['p_meas_Z']}, " + f"idle_cnot={group_totals['p_idle_cnot']}, " + f"idle_spam_effective={group_totals['p_idle_spam_effective']}, " + f"cnot={group_totals['p_cnot']})." + ) code_type = getattr(cfg.data, "code_type", "surface_code") skip_upscale = bool(getattr(cfg.data, "skip_noise_upscaling", False)) if os.environ.get("PREDECODER_SKIP_NOISE_UPSCALING", "0") == "1": @@ -818,11 +817,19 @@ def init_process_group_with_timeout(*args, **kwargs): p_min_value = p_error_value p_max_value = p_error_value if dist.rank == 0: + group_totals = upscale_info["group_totals"] + max_group = float(upscale_info["max_group"]) # Always print the grouped totals + decision to make verification easy from logs. print( "[Train] noise_model grouped totals: " - f"prep={p_prep:.6g}, meas={p_meas:.6g}, " - f"idle_cnot={p_idle_cnot:.6g}, idle_spam={p_idle_spam:.6g}, cnot={p_cnot:.6g}; " + f"prep_X={group_totals['p_prep_X']:.6g}, " + f"prep_Z={group_totals['p_prep_Z']:.6g}, " + f"meas_X={group_totals['p_meas_X']:.6g}, " + f"meas_Z={group_totals['p_meas_Z']:.6g}, " + f"idle_cnot={group_totals['p_idle_cnot']:.6g}, " + f"idle_spam_raw={group_totals['p_idle_spam_raw']:.6g}, " + f"idle_spam_effective={group_totals['p_idle_spam_effective']:.6g}, " + f"cnot={group_totals['p_cnot']:.6g}; " f"max_group={max_group:.6g}" ) print(f"[Train] {upscale_info['message']}") @@ -876,10 +883,12 @@ def init_process_group_with_timeout(*args, **kwargs): ) print( "[Train] noise_model totals: " - f"prep_total={p_prep:.6g}, meas_total={p_meas:.6g}, " - f"idle_cnot_total={noise_model_user_obj.get_total_idle_cnot_probability():.6g}, " - f"idle_spam_total={noise_model_user_obj.get_total_idle_spam_probability():.6g}, " - f"cnot_total={noise_model_user_obj.get_total_cnot_probability():.6g}" + f"prep_total={group_totals['p_prep_total']:.6g}, " + f"meas_total={group_totals['p_meas_total']:.6g}, " + f"idle_cnot_total={group_totals['p_idle_cnot']:.6g}, " + f"idle_spam_total={group_totals['p_idle_spam_raw']:.6g}, " + f"idle_spam_effective={group_totals['p_idle_spam_effective']:.6g}, " + f"cnot_total={group_totals['p_cnot']:.6g}" ) elif dist.rank == 0: print("[Train] noise_model: null (using legacy single-p / p-range sampling)") @@ -909,6 +918,16 @@ def is_list_like(obj): precomputed_frames_dir = resolve_precomputed_frames_dir( precomputed_frames_dir, cfg.distance, cfg.n_rounds, cfg.meas_basis, dist.rank ) + if dist.rank == 0: + if noise_model_train_obj is not None: + print(f"[Train] active noise_model_sha256={noise_model_train_obj.sha256()}") + if precomputed_frames_dir is None: + print("[Train] DEM artifacts: building in memory") + else: + print( + "[Train] DEM artifacts: using disk request " + f"{precomputed_frames_dir} with noise metadata validation" + ) _compute_dtype_raw = getattr(cfg.data, 'compute_dtype', None) _compute_dtype = None @@ -958,6 +977,7 @@ def is_list_like(obj): decompose_y=False, precomputed_frames_dir=precomputed_frames_dir, code_rotation=code_rotation, + noise_model=noise_model_train_obj, **_he_accel_kwargs, ) val_generator = QCDataGeneratorTorch( @@ -978,6 +998,7 @@ def is_list_like(obj): decompose_y=False, precomputed_frames_dir=precomputed_frames_dir, code_rotation=code_rotation, + noise_model=noise_model_train_obj, **_he_accel_kwargs, ) diff --git a/conf/config_local_test.yaml b/conf/config_local_test.yaml index ff7fa10..68eedb5 100644 --- a/conf/config_local_test.yaml +++ b/conf/config_local_test.yaml @@ -37,9 +37,9 @@ data: p_idle_cnot_X: 0.001 p_idle_cnot_Y: 0.001 p_idle_cnot_Z: 0.001 - p_idle_spam_X: 0.001998 - p_idle_spam_Y: 0.001998 - p_idle_spam_Z: 0.001998 + p_idle_spam_X: 0.001996 + p_idle_spam_Y: 0.001996 + p_idle_spam_Z: 0.001996 p_cnot_IX: 0.0002 p_cnot_IY: 0.0002 p_cnot_IZ: 0.0002 diff --git a/conf/config_public.yaml b/conf/config_public.yaml index d7e4c64..1170e64 100644 --- a/conf/config_public.yaml +++ b/conf/config_public.yaml @@ -49,9 +49,10 @@ data: p_idle_cnot_Y: 0.001 # p/3 p_idle_cnot_Z: 0.001 # p/3 # Idle during SPAM window (ancilla prep+reset) on data qubits only (3) - p_idle_spam_X: 0.001998 # 2*p/3 - 2*p^2/9 - p_idle_spam_Y: 0.001998 # 2*p/3 - 2*p^2/9 - p_idle_spam_Z: 0.001998 # 2*p/3 - 2*p^2/9 + # Two-step depolarising composition (XOR of two p/3-per-axis steps): 2*p/3 - 4*p^2/9. + p_idle_spam_X: 0.001996 # 2*p/3 - 4*p^2/9 + p_idle_spam_Y: 0.001996 # 2*p/3 - 4*p^2/9 + p_idle_spam_Z: 0.001996 # 2*p/3 - 4*p^2/9 # CNOT two-qubit errors (15) - keys are p_cnot_{Pauli}{Pauli} excluding II, p/15 p_cnot_IX: 0.0002 p_cnot_IY: 0.0002