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]