From 5685fda82725a3dc6dac91682ce43961313def89 Mon Sep 17 00:00:00 2001 From: Jaipreeth Tiruvaipati Date: Mon, 15 Jun 2026 01:10:19 +0530 Subject: [PATCH] Reimplement PauliString with a symplectic bit representation Replace the dict-based PauliString with two integer bitmasks (the symplectic X/Z representation), turning multiplication, (anti)commutation, collapsing and cover-vector encoding into pure integer bit operations and removing the repeated stim.PauliString <-> dict conversions. - pauli.py: symplectic masks + __slots__ + precomputed hash; O(1) exact and commuting cover vectors; bit-packed stim conversions; a CollapsingOperators value type that stores a boundary's single-qubit collapse as combined masks. - boundary.py: cache after_collapse; represent collapsing operations by their CollapsingOperators instead of a frozenset of PauliString objects. - cover.py: encode the GF(2) cover system over whole strings. - flow.py: share per-fragment collapse masks; find the measurements/resets a stabilizer overlaps with a bitwise AND instead of scanning every operator. - match.py: replace the loop-sanity deepcopy with a shallow flow clone. The annotated detectors are unchanged; on tqec surface-code memory circuits this is ~7x faster than the dict baseline and the gap grows with code size. Fixes #47 --- src/tqecd/boundary.py | 64 ++++++-- src/tqecd/boundary_test.py | 13 ++ src/tqecd/cover.py | 100 +++++------- src/tqecd/flow.py | 78 ++++++--- src/tqecd/match.py | 36 ++++- src/tqecd/match_test.py | 42 +++++ src/tqecd/pauli.py | 314 +++++++++++++++++++++++++++---------- src/tqecd/pauli_test.py | 98 +++++++++++- 8 files changed, 546 insertions(+), 199 deletions(-) create mode 100644 src/tqecd/match_test.py diff --git a/src/tqecd/boundary.py b/src/tqecd/boundary.py index b48d57d..32f1d5a 100644 --- a/src/tqecd/boundary.py +++ b/src/tqecd/boundary.py @@ -6,14 +6,14 @@ from tqecd.exceptions import TQECDException from tqecd.measurement import RelativeMeasurementLocation -from tqecd.pauli import PauliString +from tqecd.pauli import CollapsingOperators, PauliString class BoundaryStabilizer: def __init__( self, stabilizer: PauliString, - collapsing_operations: Iterable[PauliString], + collapsing_operations: Iterable[PauliString] | CollapsingOperators, measurements: list[RelativeMeasurementLocation], reset_qubits: frozenset[int], forward: bool, @@ -29,6 +29,10 @@ def __init__( operation has been applied. collapsing_operations: The collapsing operations the stabilizer will have to go through to exit the :class:`~tqecd.fragment.Fragment`. + Either the individual single-qubit operators, or an already-built + :class:`~tqecd.pauli.CollapsingOperators` (used by + :meth:`merge`, :meth:`with_measurement_offset` and fragment-level + construction to share the precomputed masks). measurements: measurement offsets relative to the **end** of the fragment (even if the created :class:`BoundaryStabilizer` instance represents a stabilizer on the beginning boundary) of @@ -44,12 +48,24 @@ def __init__( """ self._stabilizer = stabilizer self._measurements = measurements - self._collapsing_operations = frozenset(collapsing_operations) - self._has_anticommuting_collapsing_operations = any( - co.anticommutes(stabilizer) for co in collapsing_operations + self._collapse = ( + collapsing_operations + if isinstance(collapsing_operations, CollapsingOperators) + else CollapsingOperators.from_paulis(collapsing_operations) + ) + # The individual operators are reconstructed from the masks only on + # demand (repr, external access); see ``collapsing_operations``. + self._collapsing_operations_cache: frozenset[PauliString] | None = None + self._has_anticommuting_collapsing_operations = ( + self._collapse.anticommutes_with(stabilizer) ) self._reset_qubits: frozenset[int] = reset_qubits self._is_forward = forward + # Lazily-cached return value of ``after_collapse``. ``BoundaryStabilizer`` + # is immutable after construction, so the collapsed stabilizer is + # computed at most once and reused (it is accessed repeatedly while + # matching detectors). + self._after_collapse: PauliString | None = None @property def has_anticommuting_operations(self) -> bool: @@ -67,6 +83,10 @@ def after_collapse(self) -> PauliString: """Compute the stabilizer obtained after applying the collapsing operations. + The result is computed once on first access and cached for subsequent + accesses, which is safe because ``BoundaryStabilizer`` is immutable + after construction. + Raises: TQECDException: If any of the collapsing operation anti-commutes with the stored stabilizer. @@ -80,7 +100,9 @@ def after_collapse(self) -> PauliString: "Cannot collapse a BoundaryStabilizer if it has " "anticommuting operations." ) - return self._stabilizer.collapse_by(self._collapsing_operations) + if self._after_collapse is None: + self._after_collapse = self._collapse.collapse(self._stabilizer) + return self._after_collapse @property def before_collapse(self) -> PauliString: @@ -96,9 +118,21 @@ def before_collapse(self) -> PauliString: @property def collapsing_operations(self) -> Iterable[PauliString]: - """Iterator on all the collapsing operations defining the boundary this - stabilizer is applied to.""" - return self._collapsing_operations + """All the collapsing operations defining the boundary this stabilizer is + applied to. + + The single-qubit operations are reconstructed (and cached) from the + stored ``(X, Z)`` masks on first access. + """ + if self._collapsing_operations_cache is None: + self._collapsing_operations_cache = self._collapse.to_paulis() + return self._collapsing_operations_cache + + @property + def collapse(self) -> CollapsingOperators: + """The collapsing operators defining the boundary this stabilizer is + applied to.""" + return self._collapse @property def measurements(self) -> list[RelativeMeasurementLocation]: @@ -127,14 +161,12 @@ def merge(self, other: BoundaryStabilizer) -> BoundaryStabilizer: operations (i.e., the same boundary), but with the two pre-collapsing stabilizers multiplied together. """ - self_collapsing_operations = set(self.collapsing_operations) - other_collapsing_operations = set(other.collapsing_operations) - if self_collapsing_operations != other_collapsing_operations: + if self.collapse != other.collapse: raise TQECDException( "Breaking pre-condition: trying to merge two BoundaryStabilizer " "instances that are not defined on the same boundary.\n" - f"Collapsing operations for left-hand side: {self_collapsing_operations}.\n" - f"Collapsing operations for right-hand side: {other_collapsing_operations}.\n" + f"Collapsing operations for left-hand side: {set(self.collapsing_operations)}.\n" + f"Collapsing operations for right-hand side: {set(other.collapsing_operations)}.\n" ) if self._is_forward != other._is_forward: raise TQECDException( @@ -179,7 +211,7 @@ def merge(self, other: BoundaryStabilizer) -> BoundaryStabilizer: ) return BoundaryStabilizer( stabilizer, - self_collapsing_operations, + self.collapse, measurements, reset_qubits, is_forward_merge, @@ -227,7 +259,7 @@ def coordinates( def with_measurement_offset(self, offset: int) -> BoundaryStabilizer: return BoundaryStabilizer( self._stabilizer, - self.collapsing_operations, + self.collapse, [m.offset_by(offset) for m in self.measurements], self._reset_qubits, self._is_forward, diff --git a/src/tqecd/boundary_test.py b/src/tqecd/boundary_test.py index 0fbe3a2..92e7110 100644 --- a/src/tqecd/boundary_test.py +++ b/src/tqecd/boundary_test.py @@ -75,6 +75,19 @@ def test_after_and_before_collapse() -> None: assert stab.after_collapse == X0Z1 +def test_after_collapse_is_cached() -> None: + X0Z1 = PauliString({0: "X", 1: "Z"}) + X0 = PauliString({0: "X"}) + Z1 = PauliString({1: "Z"}) + + stab = BoundaryStabilizer(X0Z1, [X0, Z1], [], frozenset([0]), True) + first = stab.after_collapse + second = stab.after_collapse + assert first == PauliString({}) + # The collapsed stabilizer is computed once and reused on later accesses. + assert first is second + + def test_merge() -> None: X0Z1 = PauliString({0: "X", 1: "Z"}) Z0Z1 = PauliString({0: "Z", 1: "Z"}) diff --git a/src/tqecd/cover.py b/src/tqecd/cover.py index 0f13e34..2a5936b 100644 --- a/src/tqecd/cover.py +++ b/src/tqecd/cover.py @@ -5,45 +5,6 @@ from tqecd.pauli import PauliString -def _find_cover( - target: PauliString, - sources: list[PauliString], - on_qubits: frozenset[int], - commute_with: PauliString | None = None, -) -> list[int] | None: - """Try to find a set of boundary stabilizers from ``sources`` that generate - ``target`` (or commute with ``target``) on qubits ``on_qubits`` (a "cover"). - - If multiple valid covers exist, only one will be returned. This choice is - deterministic, but not necessarily with the lowest weight. - - Args: - target: the stabilizers to cover with stabilizers from ``sources``. - sources: stabilizers that can be used to cover ``target``. - on_qubits: qubits to consider when trying to cover ``target`` with - ``sources``. - commute_with: if provided, find a commuting-cover to the provided Pauli - string; otherwise, find an exact cover. - - Returns: - A list of source indices forming the exact or commuting cover, or - ``None`` if no such cover could be found. - """ - basis: dict[int, tuple[int, int]] = {} - for i, source in enumerate(sources): - result = _solve_linear_system( - basis, source.to_int(on_qubits, commute_with), 1 << i - ) - if result is not None and commute_with is not None: - return _int_to_bit_indices(result) - if commute_with is not None: - return None - result = _solve_linear_system( - basis, target.to_int(on_qubits, commute_with), update_basis=False - ) - return None if result is None else _int_to_bit_indices(result) - - def find_exact_cover( target: PauliString, sources: list[PauliString] ) -> list[int] | None: @@ -64,6 +25,9 @@ def find_exact_cover( resulting_pauli_string = resulting_pauli_string * sources[i] assert resulting_pauli_string == target, "Should hold" + If multiple valid covers exist, only one will be returned. This choice is + deterministic, but not necessarily with the lowest weight. + Args: target: the stabilizers to cover with stabilizers from ``sources``. sources: stabilizers that can be used to cover ``target``. @@ -82,38 +46,34 @@ def find_exact_cover( if not sources: return None - # We want an exact (i.e., equality) cover on all qubits, to be sure that - # the post-condition in the docstring holds. For that, it is sufficient to - # only consider all the qubits where either `target` or at least one of the - # items of `sources` acts non-trivially (i.e., something else than the identity). - involved_qubits = frozenset(target.qubits) + # Each Pauli string is encoded as a GF(2) bit-vector concatenating its X and + # Z masks (see ``PauliString.exact_cover_vector``). ``shift`` keeps the two + # halves disjoint and must exceed the highest involved qubit index, so it is + # derived from the union of all the supports. + support = target.support for source in sources: - involved_qubits |= frozenset(source.qubits) + support |= source.support + shift = support.bit_length() - return _find_cover(target, sources, involved_qubits) + basis: dict[int, tuple[int, int]] = {} + for i, source in enumerate(sources): + _solve_linear_system(basis, source.exact_cover_vector(shift), 1 << i) + result = _solve_linear_system( + basis, target.exact_cover_vector(shift), update_basis=False + ) + return None if result is None else _int_to_bit_indices(result) def find_commuting_cover_on_target_qubits( target: PauliString, sources: list[PauliString] ) -> list[int] | None: - """Try to find a set of boundary stabilizers from ``sources`` that generate a - superset of ``target``. - - This function tries to find a set of Pauli strings from ``sources`` that - includes ``target`` (i.e., on every qubit where ``target`` is non-trivial, - the product of each of the returned Pauli strings should commute with - ``target``). - - The differences with :func:`find_cover` are: + """Try to find a set of boundary stabilizers from ``sources`` whose product + commutes with ``target`` on every qubit where ``target`` is non-trivial. - 1. this function does not restrict the output of the product of each of - the returned Pauli string on qubits where ``target`` acts trivially (i.e. - "I"). So in practice, on qubits where ``target[qubit] == "I"``, the value - of the returned Pauli string can be anything. - 2. this function does not restrict the output of the product of each of - the returned Pauli string to exactly match ``target`` on qubits where - it acts non-trivially, but rather requires the output to commute with - ``target`` on those qubits. + Unlike :func:`find_exact_cover`, this function does not constrain the product + of the returned Pauli strings on qubits where ``target`` acts trivially, and + only requires the product to commute with ``target`` (rather than to equal + it) on qubits where ``target`` is non-trivial. Args: target: the stabilizers to cover with stabilizers from ``sources``. @@ -125,7 +85,19 @@ def find_commuting_cover_on_target_qubits( """ if not sources: return None - return _find_cover(target, sources, frozenset(target.qubits), commute_with=target) + + # Each source is encoded as the bit-vector of its per-qubit anti-commutation + # with ``target`` (see ``PauliString.commuting_cover_vector``). A subset of + # sources commutes with ``target`` exactly when their vectors XOR to zero, + # i.e. when a source is linearly dependent on the previous ones. + basis: dict[int, tuple[int, int]] = {} + for i, source in enumerate(sources): + result = _solve_linear_system( + basis, source.commuting_cover_vector(target), 1 << i + ) + if result is not None: + return _int_to_bit_indices(result) + return None def _solve_linear_system( diff --git a/src/tqecd/flow.py b/src/tqecd/flow.py index e8c5a46..6870e93 100644 --- a/src/tqecd/flow.py +++ b/src/tqecd/flow.py @@ -7,8 +7,8 @@ from tqecd.cover import find_commuting_cover_on_target_qubits from tqecd.exceptions import TQECDException from tqecd.fragment import Fragment, FragmentLoop -from tqecd.measurement import get_relative_measurement_index -from tqecd.pauli import PauliString, pauli_product +from tqecd.measurement import RelativeMeasurementLocation +from tqecd.pauli import CollapsingOperators, PauliString def _anti_commuting_stabilizers_indices(flows: list[BoundaryStabilizer]) -> list[int]: @@ -43,29 +43,38 @@ def _try_merge_anticommuting_flows_inplace(flows: list[BoundaryStabilizer]) -> N if not anti_commuting_index_to_flows_index: return - collapsing_operations: list[set[PauliString]] = [ - set(flows[fi].collapsing_operations) - for fi in anti_commuting_index_to_flows_index - ] - # Checking that all the provided flows are defined on the same boundary. - # This is checked by comparing the collapsing operations for each - # anti-commuting stabilizer and asserting that they are all equal. - for i in range(1, len(collapsing_operations)): - if collapsing_operations[0] != collapsing_operations[i]: + # Two anti-commuting flows are on the same boundary iff their + # CollapsingOperators compare equal (value equality on the combined masks), + # which is cheaper than rebuilding and comparing sets of PauliString objects. + collapse = [flows[fi].collapse for fi in anti_commuting_index_to_flows_index] + # Checking that all the provided flows are defined on the same boundary, by + # asserting that their collapsing operations are all equal. + for i in range(1, len(collapse)): + if collapse[0] != collapse[i]: raise TQECDException( "Cannot merge anti-commuting flows defined on different collapsing " "operations. Found the following difference:\nFlow 0 has the " "collapsing operations:\n\t" - + "\n\t".join(f"- {c}" for c in collapsing_operations[0]) + + "\n\t".join( + f"- {c}" + for c in flows[ + anti_commuting_index_to_flows_index[0] + ].collapsing_operations + ) + f"\nFlow {i} has the collapsing operations:\n\t" - + "\n\t".join(f"- {c}" for c in collapsing_operations[i]) + + "\n\t".join( + f"- {c}" + for c in flows[ + anti_commuting_index_to_flows_index[i] + ].collapsing_operations + ) + "\n" ) - # Computation of the Pauli string representing all the collapsing operations. + # Pauli string representing the product of all the collapsing operations. # The goal of this method will be to find a cover from the provided flows such # as the resulting propagated Pauli string commutes with this one. - collapsing_pauli = pauli_product(collapsing_operations[0]) + collapsing_pauli = flows[anti_commuting_index_to_flows_index[0]].collapse.pauli # Now, we want to find flows in anticommuting_stabilizers that, when taken # into account together, commute with collapsing_pauli. @@ -251,23 +260,39 @@ def _build_flows_from_fragment(fragment: Fragment) -> FragmentFlows: targets = list(range(len(tableau))) sorted_qubit_involved_in_measurements = fragment.measurements_qubits + # The collapsing operations of every creation (resp. destruction) flow of the + # fragment are exactly its measurements (resp. resets), so their combined + # masks are computed once and shared by all the boundary stabilizers below. + measurement_collapse = CollapsingOperators.from_paulis(fragment.measurements) + reset_collapse = CollapsingOperators.from_paulis(fragment.resets) + # Support masks (one bit per touched qubit) and a qubit -> relative + # measurement location map, precomputed once so that finding the measurements + # (resp. resets) a propagated stabilizer overlaps reduces to a bitwise AND + # plus a lookup, instead of testing the stabilizer against every operation. + measurement_support = measurement_collapse.support + reset_support = reset_collapse.support + number_of_measured_qubits = len(sorted_qubit_involved_in_measurements) + measurement_location_by_qubit = { + qubit: RelativeMeasurementLocation(index - number_of_measured_qubits, qubit) + for index, qubit in enumerate(sorted_qubit_involved_in_measurements) + } + # First compute the flows created within the Fragment (i.e., originating from # reset instructions). creation_flows: list[BoundaryStabilizer] = [] for reset_stabilizer in fragment.resets: final_stabilizer = reset_stabilizer.after(tableau, targets) + measurement_overlap = final_stabilizer.support & measurement_support involved_measurements_offsets = [ - get_relative_measurement_index( - sorted_qubit_involved_in_measurements, m.qubit - ) - for m in fragment.measurements - if final_stabilizer.overlaps(m) + measurement_location_by_qubit[qubit] + for qubit in range(measurement_overlap.bit_length()) + if (measurement_overlap >> qubit) & 1 ] involved_resets_qubits = [reset_stabilizer.qubit] creation_flows.append( BoundaryStabilizer( final_stabilizer, - fragment.measurements, + measurement_collapse, involved_measurements_offsets, frozenset(involved_resets_qubits), forward=True, @@ -287,17 +312,18 @@ def _build_flows_from_fragment(fragment: Fragment) -> FragmentFlows: ) initial_stabilizer = measurement.after(tableau_inv, targets) involved_measurements_offsets = [ - get_relative_measurement_index( - sorted_qubit_involved_in_measurements, measurement.qubit - ) + measurement_location_by_qubit[measurement.qubit] ] + reset_overlap = initial_stabilizer.support & reset_support involved_resets_qubits = [ - m.qubit for m in fragment.resets if initial_stabilizer.overlaps(m) + qubit + for qubit in range(reset_overlap.bit_length()) + if (reset_overlap >> qubit) & 1 ] destruction_flows.append( BoundaryStabilizer( initial_stabilizer, - fragment.resets, + reset_collapse, involved_measurements_offsets, frozenset(involved_resets_qubits), forward=False, diff --git a/src/tqecd/match.py b/src/tqecd/match.py index 753cf41..c0b619f 100644 --- a/src/tqecd/match.py +++ b/src/tqecd/match.py @@ -1,8 +1,7 @@ from __future__ import annotations -from copy import deepcopy from dataclasses import dataclass -from typing import Final, Iterator, Mapping +from typing import Final, Iterator, Mapping, cast import numpy import stim @@ -246,6 +245,29 @@ def _find_non_propagating_non_trivial_flows( yield i +def _shallow_copy_flows( + flows: FragmentFlows | FragmentLoopFlows, +) -> FragmentFlows | FragmentLoopFlows: + """Duplicate ``flows`` for a throwaway, mutating match. + + Only the mutable ``creation``/``destruction`` lists (recursively, for a + :class:`~tqecd.flow.FragmentLoopFlows`) are copied; the + :class:`~tqecd.boundary.BoundaryStabilizer` entries are shared, as the + matching functions never mutate a boundary stabilizer in place (they only + add/remove entries from the lists). This is the cheap equivalent of a + ``deepcopy`` for the sole purpose of the sanity check below. + """ + if isinstance(flows, FragmentLoopFlows): + return FragmentLoopFlows( + [_shallow_copy_flows(f) for f in flows.fragment_flows], flows.repeat + ) + return FragmentFlows( + list(flows.creation), + list(flows.destruction), + flows.total_number_of_measurements, + ) + + def match_boundary_stabilizers( left_flows: FragmentFlows | FragmentLoopFlows, right_flows: FragmentFlows | FragmentLoopFlows, @@ -295,12 +317,12 @@ def match_boundary_stabilizers( isinstance(right_flows, FragmentLoopFlows) and right_flows.repeat > 1 ) if should_sanity_check: + # right_flows is guaranteed to be a FragmentLoopFlows here (per the + # value of should_sanity_check), so the cast below is safe. + loop_flows = cast(FragmentLoopFlows, right_flows) matched_detectors_within_loop = match_boundary_stabilizers( - # Type checking is disabled below. right_flows is guaranteed to be of type - # FragmentLoopFlows (per the value of should_sanity_check), and so have a - # "fragment_flows" attribute. - deepcopy(right_flows.fragment_flows[-1]), # type: ignore - deepcopy(right_flows.fragment_flows[0]), # type: ignore + _shallow_copy_flows(loop_flows.fragment_flows[-1]), + _shallow_copy_flows(loop_flows.fragment_flows[0]), qubit_coordinates, ) diff --git a/src/tqecd/match_test.py b/src/tqecd/match_test.py new file mode 100644 index 0000000..225a288 --- /dev/null +++ b/src/tqecd/match_test.py @@ -0,0 +1,42 @@ +from tqecd.boundary import BoundaryStabilizer +from tqecd.flow import FragmentFlows, FragmentLoopFlows +from tqecd.match import _shallow_copy_flows +from tqecd.pauli import PauliString + + +def _flows() -> FragmentFlows: + X0 = PauliString({0: "X"}) + Z1 = PauliString({1: "Z"}) + creation = [BoundaryStabilizer(X0, [], [], frozenset([0]), True)] + destruction = [BoundaryStabilizer(Z1, [], [], frozenset([1]), False)] + return FragmentFlows(creation, destruction, total_number_of_measurements=2) + + +def test_shallow_copy_flows_duplicates_lists_but_shares_stabilizers() -> None: + flows = _flows() + copied = _shallow_copy_flows(flows) + assert isinstance(copied, FragmentFlows) + + # The lists are independent: mutating the copy leaves the original intact. + assert copied.creation is not flows.creation + assert copied.destruction is not flows.destruction + copied.creation.pop() + assert len(flows.creation) == 1 + assert copied.total_number_of_measurements == flows.total_number_of_measurements + + # The boundary stabilizers themselves are shared (never mutated in place). + other = _shallow_copy_flows(flows) + assert isinstance(other, FragmentFlows) + assert other.creation[0] is flows.creation[0] + + +def test_shallow_copy_flows_recurses_into_loops() -> None: + inner = _flows() + loop = FragmentLoopFlows([inner], repeat=3) + copied = _shallow_copy_flows(loop) + assert isinstance(copied, FragmentLoopFlows) + assert copied.repeat == 3 + assert copied.fragment_flows is not loop.fragment_flows + assert isinstance(copied.fragment_flows[0], FragmentFlows) + assert copied.fragment_flows[0].creation is not inner.creation + assert copied.fragment_flows[0].creation[0] is inner.creation[0] diff --git a/src/tqecd/pauli.py b/src/tqecd/pauli.py index 9772ca7..0f740b9 100644 --- a/src/tqecd/pauli.py +++ b/src/tqecd/pauli.py @@ -4,15 +4,23 @@ that are used across the package. This class can easily be converted from and to `stim.PauliString` and implement a subset of the `stim.PauliString` API. + +Internally a `PauliString` is stored in the symplectic representation: two +integers ``_xs`` and ``_zs`` whose bit ``q`` holds, respectively, the ``X`` +and ``Z`` component of the Pauli on qubit ``q`` (``I=00``, ``X=10``, ``Z=01``, +``Y=11``). The operations on the detector-search hot path -- multiplication, +(anti)commutation, collapsing and cover-vector encoding -- are therefore pure +integer bit operations rather than per-qubit dictionary lookups. """ from __future__ import annotations import operator +from dataclasses import dataclass from functools import reduce -from itertools import chain from typing import Iterable, Literal +import numpy import stim from tqecd.exceptions import TQECDException @@ -20,44 +28,97 @@ PAULI_STRING_TYPE = Literal["I", "X", "Y", "Z"] _IXYZ: list[PAULI_STRING_TYPE] = ["I", "X", "Y", "Z"] _IXZY: list[PAULI_STRING_TYPE] = ["I", "X", "Z", "Y"] +# (x, z) symplectic bits for each Pauli literal, used when building a +# ``PauliString`` from a ``{qubit: literal}`` mapping. +_XZ_BY_LITERAL: dict[PAULI_STRING_TYPE, tuple[int, int]] = { + "I": (0, 0), + "X": (1, 0), + "Y": (1, 1), + "Z": (0, 1), +} class PauliString: """A mapping from qubits to Pauli operators that represent a Pauli string. + The string is stored in the symplectic representation as two integer + bit-masks ``_xs``/``_zs``; bit ``q`` of each holds the ``X``/``Z`` + component of the Pauli on qubit ``q``. + Invariant: - This class never stores identity Pauli terms. Any missing Pauli term is - considered to be an identity. - As such, it is illegal to initialise this class with an identity term. + This class never stores identity Pauli terms: qubit ``q`` belongs to + the string if and only if bit ``q`` of ``_xs`` or ``_zs`` is set. + Initialising with an identity term is a no-op for that qubit. + + Note: + Instances are immutable -- ``_xs``, ``_zs`` and the cached ``_hash`` + are set once at construction and every operation returns a new + instance. Performance optimisations across this package rely on this: + the precomputed ``__hash__`` here and the cached ``after_collapse`` of + :class:`~tqecd.boundary.BoundaryStabilizer`. """ + # ``__slots__`` drops the per-instance ``__dict__``: a detector search + # allocates millions of small ``PauliString`` objects, so this measurably + # reduces memory and attribute-access cost. The trade-off is that the set + # of instance attributes is fixed. See + # https://docs.python.org/3/reference/datamodel.html#slots + __slots__ = ("_xs", "_zs", "_hash") + def __init__(self, pauli_by_qubit: dict[int, PAULI_STRING_TYPE]) -> None: + xs = 0 + zs = 0 for qubit, pauli in pauli_by_qubit.items(): - if pauli not in _IXYZ: + xz = _XZ_BY_LITERAL.get(pauli) + if xz is None: raise TQECDException( - f"Invalid Pauli operator {pauli} for qubit {qubit}, expected I, X, Y, or Z." + f"Invalid Pauli operator {pauli} for qubit {qubit}, " + "expected I, X, Y, or Z." ) - self._pauli_by_qubit: dict[int, PAULI_STRING_TYPE] = { - q: pauli for q, pauli in sorted(pauli_by_qubit.items()) if pauli != "I" - } - self._hash = hash(tuple(self._pauli_by_qubit.items())) + x, z = xz + xs |= x << qubit + zs |= z << qubit + self._xs = xs + self._zs = zs + self._hash = hash((xs, zs)) + + @classmethod + def _from_masks(cls, xs: int, zs: int) -> PauliString: + """Build a ``PauliString`` straight from its symplectic masks. + + Skips the ``{qubit: literal}`` validation and iteration of + ``__init__``; used internally by the operations that already produce + valid masks. + """ + ret = cls.__new__(cls) + ret._xs = xs + ret._zs = zs + ret._hash = hash((xs, zs)) + return ret @property def non_trivial_pauli_count(self) -> int: - return len(self._pauli_by_qubit) + return (self._xs | self._zs).bit_count() @property def qubits(self) -> Iterable[int]: - return self._pauli_by_qubit.keys() + support = self._xs | self._zs + return [q for q in range(support.bit_length()) if (support >> q) & 1] + + @property + def support(self) -> int: + """Bit-mask of the qubits on which this Pauli string is non-trivial.""" + return self._xs | self._zs @property def qubit(self) -> int: - if len(self._pauli_by_qubit) != 1: + support = self._xs | self._zs + if support.bit_count() != 1: raise TQECDException( "Cannot retrieve only one qubit from a Pauli string with " - f"{len(self._pauli_by_qubit)} qubits." + f"{support.bit_count()} qubits." ) - return next(iter(self.qubits)) + return support.bit_length() - 1 @staticmethod def from_stim_pauli_string( @@ -65,12 +126,13 @@ def from_stim_pauli_string( ) -> PauliString: """Convert a `stim.PauliString` to a `PauliString` instance, ignoring the sign.""" - return PauliString( - { - q: _IXYZ[stim_pauli_string[q]] - for q in range(len(stim_pauli_string)) - if stim_pauli_string[q] != 0 - } + # ``to_numpy(bit_packed=True)`` returns the X and Z components as + # little-endian bit-packed ``uint8`` arrays, which map directly onto the + # integer masks (bit ``q`` of the bytes is qubit ``q``). + xs, zs = stim_pauli_string.to_numpy(bit_packed=True) + return PauliString._from_masks( + int.from_bytes(xs.tobytes(), "little"), + int.from_bytes(zs.tobytes(), "little"), ) def to_stim_pauli_string(self, length: int | None) -> stim.PauliString: @@ -80,46 +142,35 @@ def to_stim_pauli_string(self, length: int | None) -> stim.PauliString: length: The length of the `stim.PauliString`. If `None`, the length is set to the maximum qubit index in the `PauliString` plus one. """ - max_qubit_index = max(self._pauli_by_qubit.keys()) + max_qubit_index = (self._xs | self._zs).bit_length() - 1 length = length or max_qubit_index + 1 if length <= max_qubit_index: raise TQECDException( f"The length specified {length} <= the maximum qubit index {max_qubit_index} in the pauli string." ) - stim_pauli_string = stim.PauliString(length) - for q, p in self._pauli_by_qubit.items(): - stim_pauli_string[q] = p - return stim_pauli_string + # Feed stim the X and Z masks as little-endian bit-packed ``uint8`` + # arrays, the inverse of the conversion in ``from_stim_pauli_string``. + num_bytes = (length + 7) // 8 + xs = numpy.frombuffer(self._xs.to_bytes(num_bytes, "little"), dtype=numpy.uint8) + zs = numpy.frombuffer(self._zs.to_bytes(num_bytes, "little"), dtype=numpy.uint8) + return stim.PauliString.from_numpy(xs=xs, zs=zs, num_qubits=length) def __bool__(self) -> bool: - return bool(self._pauli_by_qubit) + return bool(self._xs | self._zs) def __mul__(self, other: PauliString) -> PauliString: - result: dict[int, PAULI_STRING_TYPE] = {} - for q in self._pauli_by_qubit.keys() | other._pauli_by_qubit.keys(): - a = self._pauli_by_qubit.get(q, "I") - b = other._pauli_by_qubit.get(q, "I") - ax = a in "XY" - az = a in "YZ" - bx = b in "XY" - bz = b in "YZ" - cx = ax ^ bx - cz = az ^ bz - c = _IXZY[cx + cz * 2] - if c != "I": - result[q] = c - return PauliString(result) + # Pauli product ignoring the phase is the symplectic XOR of the masks. + return PauliString._from_masks(self._xs ^ other._xs, self._zs ^ other._zs) def __repr__(self) -> str: - return f"PauliString(qubits={self._pauli_by_qubit!r})" + inner = ", ".join(f"{q}: {self[q]!r}" for q in self.qubits) + return f"PauliString(qubits={{{inner}}})" def __str__(self) -> str: - return "*".join( - f"{self._pauli_by_qubit[q]}{q}" for q in sorted(self._pauli_by_qubit.keys()) - ) + return "*".join(f"{self[q]}{q}" for q in self.qubits) def __len__(self) -> int: - return len(self._pauli_by_qubit) + return (self._xs | self._zs).bit_count() def commutes(self, other: PauliString) -> bool: """Check if this Pauli string commutes with another Pauli string.""" @@ -128,10 +179,8 @@ def commutes(self, other: PauliString) -> bool: def anticommutes(self, other: PauliString) -> bool: """Check if this Pauli string anticommutes with another Pauli string.""" - t = 0 - for q in self._pauli_by_qubit.keys() & other._pauli_by_qubit.keys(): - t += self._pauli_by_qubit[q] != other._pauli_by_qubit[q] - return t % 2 == 1 + # Parity of the symplectic inner product: (X_self & Z_other) ^ (Z_self & X_other). + return ((self._xs & other._zs) ^ (self._zs & other._xs)).bit_count() & 1 == 1 def collapse_by(self, collapse_operators: Iterable[PauliString]) -> PauliString: """Collapse the provided Pauli string by the provided operators. @@ -154,29 +203,46 @@ def collapse_by(self, collapse_operators: Iterable[PauliString]) -> PauliString: Returns: a copy of self, collapsed by the provided operators. """ - ret = PauliString(self._pauli_by_qubit.copy()) + xs, zs = self._xs, self._zs for op in collapse_operators: - if not ret.commutes(op): + # Commute-check against the *current* partially-collapsed result, not + # the original ``self``: collapsing is performed sequentially. + if ((xs & op._zs) ^ (zs & op._xs)).bit_count() & 1: raise TQECDException( - f"Cannot collapse {ret} by a non-commuting operator {op}." + f"Cannot collapse {PauliString._from_masks(xs, zs)} by a " + f"non-commuting operator {op}." ) - for qubit in op.qubits: - if qubit in ret._pauli_by_qubit: - del ret._pauli_by_qubit[qubit] - return ret + support = op._xs | op._zs + xs &= ~support + zs &= ~support + return PauliString._from_masks(xs, zs) def after(self, tableau: stim.Tableau, targets: Iterable[int]) -> PauliString: - stim_pauli_string = self.to_stim_pauli_string( - length=max(list(targets) + list(self._pauli_by_qubit.keys())) + 1 - ) + targets = tuple(targets) + # Every target indexes into ``tableau`` (so is ``< len(tableau)``); the + # propagated string therefore fits in ``len(tableau)`` qubits, or in this + # string's own length if it is wider. Using ``len(tableau)`` avoids an + # O(len(targets)) scan for the largest target. + max_qubit_index = (self._xs | self._zs).bit_length() - 1 + length = max(len(tableau), max_qubit_index + 1) + stim_pauli_string = self.to_stim_pauli_string(length=length) stim_pauli_string_after = stim_pauli_string.after(tableau, targets=targets) return PauliString.from_stim_pauli_string(stim_pauli_string_after) def contains(self, other: PauliString) -> bool: - return self._pauli_by_qubit.items() >= other._pauli_by_qubit.items() + """Check whether ``other`` is a sub-string of ``self``. + + ``self`` must act exactly like ``other`` on every qubit where ``other`` + is non-trivial. Qubits outside ``other``'s support are intentionally + ignored, so e.g. ``X0*X1*X2`` contains ``X2``. + """ + other_support = other._xs | other._zs + return (self._xs & other_support) == other._xs and ( + self._zs & other_support + ) == other._zs def overlaps(self, other: PauliString) -> bool: - return bool(self._pauli_by_qubit.keys() & other._pauli_by_qubit.keys()) + return bool((self._xs | self._zs) & (other._xs | other._zs)) def __eq__(self, other: object) -> bool: """Check if two PauliString are equal. @@ -189,39 +255,45 @@ def __eq__(self, other: object) -> bool: """ return ( isinstance(other, PauliString) - and self._pauli_by_qubit == other._pauli_by_qubit + and self._xs == other._xs + and self._zs == other._zs ) def __hash__(self) -> int: return self._hash def __getitem__(self, index: int) -> PAULI_STRING_TYPE: - return self._pauli_by_qubit.get(index, "I") + return _IXZY[((self._xs >> index) & 1) + 2 * ((self._zs >> index) & 1)] - def to_int( - self, qubits: Iterable[int], reference: PauliString | None = None - ) -> int: - """Convert the Pauli string to an integer representation on the provided qubits. + def exact_cover_vector(self, shift: int) -> int: + """Encode this Pauli string as a GF(2) bit-vector for exact-cover solving. + + The ``X`` and ``Z`` masks are concatenated into a single integer, with + the ``Z`` mask shifted up by ``shift`` bits so the two halves never + overlap. ``shift`` must be strictly greater than the highest qubit index + involved in the cover (and shared by every vector of the same system). + The XOR of two such vectors is the vector of the Pauli product of the + two strings, so GF(2) linear combinations correspond to Pauli products. Args: - qubits: the qubits over which to encode. - reference: if ``None``, each qubit contributes 2 bits encoding the - Pauli at that qubit. If a reference Pauli string is provided, - each qubit contributes 1 bit indicating whether the Pauli at - that qubit anti-commutes with the reference's Pauli at the same qubit. + shift: bit offset applied to the ``Z`` mask; must exceed the highest + involved qubit index. """ - if reference is None: - return reduce( - lambda acc, bit: acc << 1 | bit, - chain.from_iterable(pauli_literal_to_bools(self[q]) for q in qubits), - 0, - ) - result = 0 - for q in qubits: - sxt, szt = pauli_literal_to_bools(self[q]) - rxt, rzt = pauli_literal_to_bools(reference[q]) - result = (result << 1) | int((sxt and rzt) ^ (szt and rxt)) - return result + return self._xs | (self._zs << shift) + + def commuting_cover_vector(self, reference: PauliString) -> int: + """Encode this Pauli string as a GF(2) bit-vector of per-qubit + anti-commutation with ``reference``. + + Bit ``q`` is set when this Pauli string anti-commutes with ``reference`` + on qubit ``q``. Bits where ``reference`` is trivial are always zero, so + the encoding is naturally restricted to ``reference``'s support. + + Args: + reference: the Pauli string to measure per-qubit anti-commutation + against. + """ + return (self._xs & reference._zs) ^ (self._zs & reference._xs) def pauli_literal_to_bools( @@ -239,3 +311,75 @@ def pauli_literal_to_bools( def pauli_product(paulis: Iterable[PauliString]) -> PauliString: return reduce(operator.mul, paulis, PauliString({})) + + +@dataclass(frozen=True) +class CollapsingOperators: + """The collapsing operators (measurements or resets) acting at one boundary + of a fragment, stored by their combined symplectic masks. + + Collapsing operators are single-qubit Pauli operators acting on distinct + qubits, so the whole set is captured losslessly by the bitwise OR of the + operators' ``X`` and ``Z`` masks. Keeping only the masks turns the boundary + computations -- anti-commutation with a stabilizer, the collapse itself, and + the product of the operators -- into single integer bit operations, and + avoids materialising a set of one-qubit :class:`PauliString` objects on the + hot path. + + Attributes: + xs: combined ``X`` mask (bit ``q`` set iff some operator has an ``X`` or + ``Y`` component on qubit ``q``). + zs: combined ``Z`` mask (bit ``q`` set iff some operator has a ``Z`` or + ``Y`` component on qubit ``q``). + """ + + xs: int + zs: int + + @staticmethod + def from_paulis(operators: Iterable[PauliString]) -> CollapsingOperators: + """Build the combined masks from the individual single-qubit operators.""" + xs = 0 + zs = 0 + for pauli in operators: + xs |= pauli._xs + zs |= pauli._zs + return CollapsingOperators(xs, zs) + + @property + def support(self) -> int: + """Bit-mask of the qubits touched by the collapsing operators.""" + return self.xs | self.zs + + def anticommutes_with(self, stabilizer: PauliString) -> bool: + """Whether ``stabilizer`` anti-commutes with at least one operator. + + Because the operators act on distinct qubits, the per-qubit symplectic + product has a non-zero bit exactly when a single operator anti-commutes, + so one ``!= 0`` test replaces iterating over the operators. + """ + return ((stabilizer._xs & self.zs) ^ (stabilizer._zs & self.xs)) != 0 + + def collapse(self, stabilizer: PauliString) -> PauliString: + """Return ``stabilizer`` with every collapsed qubit reset to identity. + + Assumes ``stabilizer`` commutes with the operators (checked by the caller + via :meth:`anticommutes_with`); collapsing then simply removes the + measured/reset qubits from the propagated stabilizer. + """ + keep = ~self.support + return PauliString._from_masks(stabilizer._xs & keep, stabilizer._zs & keep) + + @property + def pauli(self) -> PauliString: + """Product of all the operators as a single :class:`PauliString`.""" + return PauliString._from_masks(self.xs, self.zs) + + def to_paulis(self) -> frozenset[PauliString]: + """Reconstruct the individual single-qubit operators.""" + support = self.support + return frozenset( + PauliString._from_masks(self.xs & (1 << q), self.zs & (1 << q)) + for q in range(support.bit_length()) + if (support >> q) & 1 + ) diff --git a/src/tqecd/pauli_test.py b/src/tqecd/pauli_test.py index 5617558..a4a202c 100644 --- a/src/tqecd/pauli_test.py +++ b/src/tqecd/pauli_test.py @@ -2,7 +2,12 @@ import stim from tqecd.exceptions import TQECDException -from tqecd.pauli import PauliString, pauli_literal_to_bools, pauli_product +from tqecd.pauli import ( + CollapsingOperators, + PauliString, + pauli_literal_to_bools, + pauli_product, +) def test_pauli_string_construction() -> None: @@ -116,3 +121,94 @@ def test_pauli_product() -> None: assert pauli_product([X0Z1 for _ in range(20)]) == I0to20 assert pauli_product([X0Z1 for _ in range(21)]) == X0Z1 assert pauli_product([X0Z1, Z0]) == PauliString({0: "Y", 1: "Z"}) + + +# Sizes spanning empty, single-qubit and the 8/16/64-bit word boundaries to +# exercise the symplectic integer representation around byte/word edges. +@pytest.mark.parametrize("num_qubits", [0, 1, 7, 8, 9, 16, 23, 64, 65]) +def test_pauli_string_matches_stim(num_qubits: int) -> None: + lhs = stim.PauliString.random(num_qubits=num_qubits) + rhs = stim.PauliString.random(num_qubits=num_qubits) + pauli_lhs = PauliString.from_stim_pauli_string(lhs) + pauli_rhs = PauliString.from_stim_pauli_string(rhs) + + # Multiplication agrees with stim (both ignore the sign). + assert pauli_lhs * pauli_rhs == PauliString.from_stim_pauli_string(lhs * rhs) + # (Anti)commutation agrees with stim. + assert pauli_lhs.anticommutes(pauli_rhs) == (not lhs.commutes(rhs)) + assert pauli_lhs.commutes(pauli_rhs) == lhs.commutes(rhs) + # Weight (number of non-identity terms) agrees with stim. + assert pauli_lhs.non_trivial_pauli_count == lhs.weight + + +def test_pauli_string_contains() -> None: + XXX = PauliString({0: "X", 1: "X", 2: "X"}) + # ``contains`` ignores qubits outside ``other``'s support: a Pauli string + # contains any of its single-qubit sub-strings. + assert XXX.contains(PauliString({2: "X"})) + assert XXX.contains(PauliString({0: "X", 2: "X"})) + assert XXX.contains(PauliString({})) + # Same support but different Pauli is not contained. + assert not PauliString({0: "X"}).contains(PauliString({0: "Z"})) + assert not XXX.contains(PauliString({0: "X", 3: "X"})) + + +def test_pauli_string_collapse_by_is_sequential() -> None: + X0X1 = PauliString({0: "X", 1: "X"}) + X0 = PauliString({0: "X"}) + Z0Z1 = PauliString({0: "Z", 1: "Z"}) + + # ``X0*X1`` commutes with ``Z0*Z1`` as a whole... + assert X0X1.commutes(Z0Z1) + # ...but once ``X0`` has been collapsed away, the remaining ``X1`` + # anti-commutes with ``Z0*Z1``, so the sequential collapse must raise. + with pytest.raises(TQECDException, match=r"^Cannot collapse .* non-commuting"): + X0X1.collapse_by([X0, Z0Z1]) + + +def test_pauli_string_exact_cover_vector() -> None: + # X and Z masks concatenated, Z half shifted up by ``shift``. + assert PauliString({0: "X"}).exact_cover_vector(4) == 0b0001 + assert PauliString({0: "Z"}).exact_cover_vector(4) == 1 << 4 + assert PauliString({0: "Y"}).exact_cover_vector(4) == (1 << 4) | 1 + # XOR of two exact-cover vectors equals the vector of their Pauli product. + a = PauliString({0: "X", 1: "Z"}) + b = PauliString({0: "Z", 1: "Z"}) + assert a.exact_cover_vector(8) ^ b.exact_cover_vector(8) == ( + a * b + ).exact_cover_vector(8) + + +def test_pauli_string_commuting_cover_vector() -> None: + reference = PauliString({0: "Z", 1: "Z"}) + # Bit set where self anti-commutes with the reference. + assert PauliString({0: "X"}).commuting_cover_vector(reference) == 0b01 + assert PauliString({0: "X", 1: "X"}).commuting_cover_vector(reference) == 0b11 + assert PauliString({0: "Z"}).commuting_cover_vector(reference) == 0b00 + # Qubits outside the reference's support never contribute. + assert PauliString({2: "X"}).commuting_cover_vector(reference) == 0b00 + + +def test_collapsing_operators() -> None: + X0 = PauliString({0: "X"}) + Z1 = PauliString({1: "Z"}) + Y2 = PauliString({2: "Y"}) + collapse = CollapsingOperators.from_paulis([X0, Z1, Y2]) + + # Combined masks capture the qubits touched and the product of the operators. + assert collapse.support == 0b111 + assert collapse.pauli == PauliString({0: "X", 1: "Z", 2: "Y"}) + # The individual single-qubit operators round-trip out of the masks. + assert collapse.to_paulis() == frozenset({X0, Z1, Y2}) + # Value equality (frozen dataclass), order-independent. + assert collapse == CollapsingOperators.from_paulis([Y2, X0, Z1]) + + # anticommutes_with == "any operator anti-commutes". + assert collapse.anticommutes_with(PauliString({0: "Z"})) # Z0 vs X0 + assert not collapse.anticommutes_with(PauliString({0: "X"})) # commutes with all + # collapse removes exactly the touched qubits, leaving the rest untouched. + assert collapse.collapse(PauliString({0: "X", 3: "Z"})) == PauliString({3: "Z"}) + + empty = CollapsingOperators.from_paulis([]) + assert empty.support == 0 + assert not empty.to_paulis()