From cc49e5843911924c3b6b720144ba7b1f60ada224 Mon Sep 17 00:00:00 2001 From: Pritesh Date: Sun, 28 Jun 2026 18:15:34 -0500 Subject: [PATCH] Support custom syndrome-extraction schedules in detector annotation annotate_detectors_automatically silently returned an incomplete set of detectors for circuits whose schedule interleaves collapsing operations with computation (e.g. a depth-optimized rotated surface code whose X-ancilla RX shares a moment with two-qubit gates, or whose Z/X ancillas are measured in separate moments). Such circuits pass is_valid_input_circuit but Fragment only collects leading resets / trailing measurements, and matching is consecutive- only, so interleaved collapsing ops are dropped and cross-round detectors are never found. Add tqecd.schedule, which detects when a fragment would drop a collapsing op, rebuilds a logically-equivalent circuit with the canonical per-round shape, annotates it, and transplants the detectors back onto the original circuit (valid because the measurement record order is preserved). Circuits already handled correctly are left on the unchanged code path. --- src/tqecd/construction.py | 19 ++ src/tqecd/schedule.py | 261 ++++++++++++++++++ src/tqecd/schedule_test.py | 98 +++++++ ...memory_z_distance_3_interleaved_reset.stim | 70 +++++ 4 files changed, 448 insertions(+) create mode 100644 src/tqecd/schedule.py create mode 100644 src/tqecd/schedule_test.py create mode 100644 src/tqecd/test_files/valid/surface_code_rotated_memory_z_distance_3_interleaved_reset.stim diff --git a/src/tqecd/construction.py b/src/tqecd/construction.py index 427faf1..95b3e32 100644 --- a/src/tqecd/construction.py +++ b/src/tqecd/construction.py @@ -9,6 +9,11 @@ from tqecd.fragment import Fragment, FragmentLoop, split_stim_circuit_into_fragments from tqecd.match import MatchedDetector, match_detectors_from_flows_shallow from tqecd.predicates import is_valid_input_circuit +from tqecd.schedule import ( + canonicalize_collapsing_schedule, + fragment_schedule_needs_normalization, + transplant_detectors, +) from tqecd.utils import remove_duplicate_detectors @@ -63,6 +68,20 @@ def annotate_detectors_automatically(circuit: stim.Circuit) -> stim.Circuit: if potential_error_reason is not None: raise TQECDException(potential_error_reason) + # Circuits whose syndrome-extraction schedule interleaves collapsing + # operations with computation (so that fragment splitting would drop resets + # or measurements) are first rescheduled into a logically-equivalent + # canonical form, annotated, and the resulting detectors are transplanted + # back onto the original circuit. See :mod:`tqecd.schedule`. + if fragment_schedule_needs_normalization(circuit): + normalized_circuit = canonicalize_collapsing_schedule(circuit) + annotated_circuit = _annotate_clean_circuit(normalized_circuit) + return transplant_detectors(circuit, annotated_circuit) + + return _annotate_clean_circuit(circuit) + + +def _annotate_clean_circuit(circuit: stim.Circuit) -> stim.Circuit: fragments = split_stim_circuit_into_fragments(circuit) qubit_coords_map: dict[int, tuple[float, ...]] = { q: tuple(coords) for q, coords in circuit.get_final_qubit_coordinates().items() diff --git a/src/tqecd/schedule.py b/src/tqecd/schedule.py new file mode 100644 index 0000000..5bb7ae6 --- /dev/null +++ b/src/tqecd/schedule.py @@ -0,0 +1,261 @@ +"""Normalise the *scheduling* of collapsing operations so that circuits using a +custom syndrome-extraction schedule can be annotated with detectors. + +:func:`~tqecd.construction.annotate_detectors_automatically` relies on splitting +the circuit into :class:`~tqecd.fragment.Fragment` instances, each having the +shape ``[leading resets] [computation] [trailing measurements]``. A +:class:`Fragment` only collects resets from its *leading* contiguous moments and +measurements from its *trailing* contiguous moments (see +:meth:`tqecd.fragment.Fragment.__init__`), and detectors are only matched +between *consecutive* fragments (see +:func:`tqecd.match.match_detectors_from_flows_shallow`). + +Circuits that interleave collapsing operations with computation -- for example a +rotated surface code whose X-ancilla ``RX`` shares a moment with the two-qubit +gates, or whose Z- and X-ancilla measurements live in different moments -- break +those assumptions: the interleaved resets/measurements are silently dropped and +whole rounds are split across two fragments, so many detectors are never found. + +For such circuits the custom schedule is only a *depth* optimisation: the +ancillas that are reset late / measured early are idle up to the round boundary, +so the collapsing operations commute to the boundary. This module rebuilds a +*logically equivalent* circuit in which every round has the canonical shape, and +provides the tooling to transplant the resulting detectors back onto the +original circuit (valid because the measurement *record order* is preserved). +""" + +from __future__ import annotations + +from collections import defaultdict + +import stim + +from tqecd.exceptions import TQECDException +from tqecd.fragment import Fragment, split_stim_circuit_into_fragments +from tqecd.utils import ( + is_measurement, + is_reset, + iter_stim_circuit_by_moments, +) + +_RESET_NAMES = frozenset({"R", "RX", "RY", "RZ"}) +_MEASUREMENT_NAMES = frozenset({"M", "MX", "MY", "MZ"}) +_REGENERATED_ANNOTATIONS = frozenset({"DETECTOR", "SHIFT_COORDS"}) +_PRESERVED_ANNOTATIONS = frozenset({"QUBIT_COORDS", "OBSERVABLE_INCLUDE", "MPAD"}) + + +def _qubit_targets(instruction: stim.CircuitInstruction) -> set[int]: + return { + t.qubit_value + for t in instruction.targets_copy() + if t.is_qubit_target and t.qubit_value is not None + } + + +def _circuit_has_repeat_block(circuit: stim.Circuit) -> bool: + return any(isinstance(inst, stim.CircuitRepeatBlock) for inst in circuit) + + +def fragment_schedule_needs_normalization(circuit: stim.Circuit) -> bool: + """Return whether ``circuit`` has a schedule that fragment splitting cannot + handle directly. + + This is the case if and only if at least one (non-loop) :class:`Fragment` + that would be created from ``circuit`` *drops* a reset or a measurement that + is physically present in the fragment, i.e. a collapsing operation that is + neither in a leading reset moment nor in a trailing measurement moment. + + The check is intentionally conservative: it returns ``False`` for every + circuit that the existing pipeline already handles correctly (so the + behaviour of such circuits is left untouched), and ``False`` for circuits + containing ``stim.CircuitRepeatBlock`` instructions, which this module does + not rewrite. + """ + if _circuit_has_repeat_block(circuit): + return False + fragments = split_stim_circuit_into_fragments(circuit) + for fragment in fragments: + if not isinstance(fragment, Fragment): + return False + collected_reset_qubits = {reset.qubit for reset in fragment.resets} + collected_measurements = fragment.num_measurements + present_reset_qubits: set[int] = set() + present_measurements = 0 + for inst in fragment.circuit: + if isinstance(inst, stim.CircuitRepeatBlock): + continue + if is_reset(inst): + present_reset_qubits |= _qubit_targets(inst) + elif is_measurement(inst): + present_measurements += len(inst.targets_copy()) + if not present_reset_qubits <= collected_reset_qubits: + return True + if present_measurements > collected_measurements: + return True + return False + + +def _split_into_rounds( + moments: list[list[stim.CircuitInstruction]], +) -> list[list[list[stim.CircuitInstruction]]]: + """Group moments into rounds. + + A new round begins at a reset-containing moment that follows a moment in + which a measurement has already been seen for the current round. + """ + rounds: list[list[list[stim.CircuitInstruction]]] = [] + current: list[list[stim.CircuitInstruction]] = [] + seen_measurement = False + for moment in moments: + has_reset = any(inst.name in _RESET_NAMES for inst in moment) + has_measurement = any(inst.name in _MEASUREMENT_NAMES for inst in moment) + if has_reset and seen_measurement: + rounds.append(current) + current = [] + seen_measurement = False + current.append(moment) + if has_measurement: + seen_measurement = True + if current: + rounds.append(current) + return rounds + + +def _moments_without_ticks( + circuit: stim.Circuit, +) -> list[list[stim.CircuitInstruction]]: + moments: list[list[stim.CircuitInstruction]] = [] + for moment in iter_stim_circuit_by_moments(circuit): + if isinstance(moment, stim.CircuitRepeatBlock): + raise TQECDException( + "canonicalize_collapsing_schedule does not support circuits " + "containing stim.CircuitRepeatBlock instructions." + ) + moments.append( + [inst for inst in moment if inst.name != "TICK"] # type: ignore[union-attr] + ) + return moments + + +def canonicalize_collapsing_schedule(circuit: stim.Circuit) -> stim.Circuit: + """Return a logically-equivalent circuit in which every round has the shape + ``[resets] [computation moments...] [single merged measurement moment]``. + + Resets are hoisted to the front of their round and measurements are merged + into a single moment at the end of their round, preserving the original + measurement record order. The two-qubit-gate backbone is left untouched. + + Raises: + TQECDException: if ``circuit`` contains a ``stim.CircuitRepeatBlock``. + """ + moments = _moments_without_ticks(circuit) + qubit_coords = [ + inst for moment in moments for inst in moment if inst.name == "QUBIT_COORDS" + ] + + out = stim.Circuit() + for coord in qubit_coords: + out.append(coord) + out.append(stim.CircuitInstruction("TICK", [])) + + for round_moments in _split_into_rounds(moments): + resets: list[stim.CircuitInstruction] = [] + computation_moments: list[list[stim.CircuitInstruction]] = [] + measurements: list[stim.CircuitInstruction] = [] + for moment in round_moments: + computation: list[stim.CircuitInstruction] = [] + for inst in moment: + if inst.name in _RESET_NAMES: + resets.append(inst) + elif inst.name in _MEASUREMENT_NAMES: + measurements.append(inst) + elif ( + inst.name in _PRESERVED_ANNOTATIONS + or inst.name in _REGENERATED_ANNOTATIONS + ): + continue + else: + computation.append(inst) + if computation: + computation_moments.append(computation) + + for inst in resets: + out.append(inst) + if resets and (computation_moments or measurements): + out.append(stim.CircuitInstruction("TICK", [])) + for computation in computation_moments: + for inst in computation: + out.append(inst) + out.append(stim.CircuitInstruction("TICK", [])) + for inst in measurements: # original record order preserved + out.append(inst) + out.append(stim.CircuitInstruction("TICK", [])) + + while len(out) and out[-1].name == "TICK": + out = out[:-1] + return out + + +def _measurement_record_order(circuit: stim.Circuit) -> list[int]: + return [ + t.qubit_value + for inst in circuit + if not isinstance(inst, stim.CircuitRepeatBlock) and is_measurement(inst) + for t in inst.targets_copy() + if t.qubit_value is not None + ] + + +def transplant_detectors( + original: stim.Circuit, annotated: stim.Circuit +) -> stim.Circuit: + """Copy the DETECTOR annotations computed on ``annotated`` onto ``original``. + + ``original`` and ``annotated`` must have the same measurement record order + (this is guaranteed when ``annotated`` was produced by annotating + :func:`canonicalize_collapsing_schedule(original) `). + Existing ``OBSERVABLE_INCLUDE`` annotations in ``original`` are preserved; + stale ``DETECTOR`` / ``SHIFT_COORDS`` are dropped and replaced. Detectors are + re-emitted just after the measurement moment that brings the cumulative + measurement count to match their anchor position in ``annotated``, so their + ``rec[...]`` offsets stay valid. + + Raises: + TQECDException: if the measurement record orders differ. + """ + if _measurement_record_order(original) != _measurement_record_order(annotated): + raise TQECDException( + "Cannot transplant detectors: the measurement record order of the " + "rescheduled circuit differs from the original." + ) + + detectors_by_anchor: dict[int, list[tuple[tuple[int, ...], list[float]]]] = ( + defaultdict(list) + ) + seen = 0 + for inst in annotated: + if isinstance(inst, stim.CircuitRepeatBlock): + continue + if is_measurement(inst): + seen += len(inst.targets_copy()) + elif inst.name == "DETECTOR": + absolute_targets = tuple( + sorted(seen + t.value for t in inst.targets_copy()) + ) + detectors_by_anchor[seen].append((absolute_targets, inst.gate_args_copy())) + + out = stim.Circuit() + seen = 0 + for inst in original: + if ( + not isinstance(inst, stim.CircuitRepeatBlock) + and inst.name in _REGENERATED_ANNOTATIONS + ): + continue + out.append(inst) + if not isinstance(inst, stim.CircuitRepeatBlock) and is_measurement(inst): + seen += len(inst.targets_copy()) + for absolute_targets, coords in detectors_by_anchor.get(seen, []): + targets = [stim.target_rec(a - seen) for a in absolute_targets] + out.append(stim.CircuitInstruction("DETECTOR", targets, coords)) + return out diff --git a/src/tqecd/schedule_test.py b/src/tqecd/schedule_test.py new file mode 100644 index 0000000..a6151f9 --- /dev/null +++ b/src/tqecd/schedule_test.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +from pathlib import Path + +import stim + +from tqecd.construction import _annotate_clean_circuit, annotate_detectors_automatically +from tqecd.schedule import ( + canonicalize_collapsing_schedule, + fragment_schedule_needs_normalization, + transplant_detectors, +) +from tqecd.utils import remove_annotations + +_HERE = Path(__file__).parent +_VALID = _HERE / "test_files" / "valid" +_INTERLEAVED = ( + _VALID / "surface_code_rotated_memory_z_distance_3_interleaved_reset.stim" +) +_CLEAN = _VALID / "surface_code_rotated_memory_z_distance_3_rounds_2.stim" + + +def _detector_targets(circuit: stim.Circuit) -> set[frozenset[int]]: + """Detectors as sets of absolute measurement indices (placement-independent).""" + detectors: set[frozenset[int]] = set() + seen = 0 + for inst in circuit: + if inst.name in ("M", "MX", "MY", "MZ"): + seen += len(inst.targets_copy()) + elif inst.name == "DETECTOR": + detectors.add(frozenset(seen + t.value for t in inst.targets_copy())) + return detectors + + +def _measurement_record(circuit: stim.Circuit) -> list[int]: + return [ + t.qubit_value + for inst in circuit + if inst.name in ("M", "MX", "MY", "MZ") + for t in inst.targets_copy() + ] + + +def _without_detectors(circuit: stim.Circuit) -> stim.Circuit: + return remove_annotations(circuit, frozenset(["DETECTOR", "SHIFT_COORDS"])) + + +def test_clean_circuit_is_not_rescheduled() -> None: + """A circuit already in canonical form must take the unchanged code path.""" + clean = stim.Circuit.from_file(str(_CLEAN)) + assert not fragment_schedule_needs_normalization(clean) + + +def test_interleaved_reset_circuit_needs_normalization() -> None: + circuit = _without_detectors(stim.Circuit.from_file(str(_INTERLEAVED))) + assert fragment_schedule_needs_normalization(circuit) + + +def test_canonicalize_preserves_measurement_record_order() -> None: + circuit = _without_detectors(stim.Circuit.from_file(str(_INTERLEAVED))) + canonical = canonicalize_collapsing_schedule(circuit) + assert _measurement_record(circuit) == _measurement_record(canonical) + # the rescheduled circuit no longer needs normalization + assert not fragment_schedule_needs_normalization(canonical) + + +def test_interleaved_schedule_recovers_dropped_detectors() -> None: + circuit = _without_detectors(stim.Circuit.from_file(str(_INTERLEAVED))) + + # The unchanged fragment-based path silently drops detectors here ... + naive = _annotate_clean_circuit(circuit) + # ... while the schedule-aware entry point recovers them all. + fixed = annotate_detectors_automatically(circuit) + assert fixed.num_detectors > naive.num_detectors + + # And the recovered detectors match those of the equivalent clean circuit. + clean = _without_detectors(stim.Circuit.from_file(str(_CLEAN))) + ground_truth = annotate_detectors_automatically(clean) + assert _detector_targets(ground_truth) == _detector_targets(fixed) + + +def test_recovered_detectors_are_deterministic() -> None: + circuit = _without_detectors(stim.Circuit.from_file(str(_INTERLEAVED))) + fixed = annotate_detectors_automatically(circuit) + # Raises if any detector is not a deterministic function of the resets. + fixed.detector_error_model(decompose_errors=True) + + +def test_transplant_rejects_mismatched_record_order() -> None: + circuit = _without_detectors(stim.Circuit.from_file(str(_INTERLEAVED))) + annotated = annotate_detectors_automatically(circuit) + reordered = stim.Circuit() + reordered.append("M", [0]) # different measurement record + try: + transplant_detectors(reordered, annotated) + except Exception: # noqa: BLE001 - we only assert that it refuses + return + raise AssertionError("transplant_detectors should reject a mismatched record order") diff --git a/src/tqecd/test_files/valid/surface_code_rotated_memory_z_distance_3_interleaved_reset.stim b/src/tqecd/test_files/valid/surface_code_rotated_memory_z_distance_3_interleaved_reset.stim new file mode 100644 index 0000000..f666777 --- /dev/null +++ b/src/tqecd/test_files/valid/surface_code_rotated_memory_z_distance_3_interleaved_reset.stim @@ -0,0 +1,70 @@ +QUBIT_COORDS(1, 1) 1 +QUBIT_COORDS(2, 0) 2 +QUBIT_COORDS(3, 1) 3 +QUBIT_COORDS(5, 1) 5 +QUBIT_COORDS(1, 3) 8 +QUBIT_COORDS(2, 2) 9 +QUBIT_COORDS(3, 3) 10 +QUBIT_COORDS(4, 2) 11 +QUBIT_COORDS(5, 3) 12 +QUBIT_COORDS(6, 2) 13 +QUBIT_COORDS(0, 4) 14 +QUBIT_COORDS(1, 5) 15 +QUBIT_COORDS(2, 4) 16 +QUBIT_COORDS(3, 5) 17 +QUBIT_COORDS(4, 4) 18 +QUBIT_COORDS(5, 5) 19 +QUBIT_COORDS(4, 6) 25 +R 1 3 5 8 10 12 15 17 19 2 9 11 14 16 18 25 +TICK +H 2 11 16 25 +TICK +CX 2 3 16 17 11 12 15 14 10 9 19 18 +R 13 +TICK +CX 2 1 16 15 11 10 8 14 3 9 12 18 +TICK +CX 16 10 11 5 25 19 8 9 17 18 12 13 +TICK +CX 16 8 11 3 25 17 1 9 10 18 5 13 +TICK +H 2 11 16 25 +TICK +M 2 9 11 13 14 16 18 25 +DETECTOR(4, 4, 0) rec[-2] +DETECTOR(0, 4, 0) rec[-4] +DETECTOR(6, 2, 0) rec[-5] +DETECTOR(2, 2, 0) rec[-7] +TICK +R 2 9 11 14 16 18 25 +TICK +H 2 11 16 25 +TICK +CX 2 3 16 17 11 12 15 14 10 9 19 18 +R 13 +TICK +CX 2 1 16 15 11 10 8 14 3 9 12 18 +TICK +CX 16 10 11 5 25 19 8 9 17 18 12 13 +TICK +CX 16 8 11 3 25 17 1 9 10 18 5 13 +TICK +H 2 11 16 25 +TICK +M 2 9 11 13 14 16 18 25 1 3 5 8 10 12 15 17 19 +DETECTOR(4, 4, 1) rec[-11] rec[-5] rec[-4] rec[-2] rec[-1] +DETECTOR(0, 4, 1) rec[-13] rec[-6] rec[-3] +DETECTOR(6, 2, 1) rec[-14] rec[-7] rec[-4] +DETECTOR(2, 2, 1) rec[-16] rec[-9] rec[-8] rec[-6] rec[-5] +DETECTOR(2, 0, 1) rec[-25] rec[-17] +DETECTOR(2, 2, 1) rec[-24] rec[-16] +DETECTOR(4, 2, 1) rec[-23] rec[-15] +DETECTOR(6, 2, 1) rec[-22] rec[-14] +DETECTOR(0, 4, 1) rec[-21] rec[-13] +DETECTOR(2, 4, 1) rec[-20] rec[-12] +DETECTOR(4, 4, 1) rec[-19] rec[-11] +DETECTOR(4, 6, 1) rec[-18] rec[-10] +DETECTOR(3, 1, 1) rec[-9] rec[-8] rec[-7] +DETECTOR(3, 3, 1) rec[-6] rec[-5] rec[-4] +DETECTOR(3, 5, 1) rec[-3] rec[-2] rec[-1] +OBSERVABLE_INCLUDE(0) rec[-7] rec[-8] rec[-9]