Simplify the statevector executor and the sampler's backend dispatch

CHANGELOG.md

@@ -90,6 +90,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - The ZX diagram types are now `"zx"`, `"zx-dets"` and `"zx-meas"`. They return zxparam's HTML drawing object instead of a graph.
 - Internals: `GraphRepresentation` (`tsim.core.instructions`) is now `GraphBuilder` in the new `tsim.core.builder`, and `parse_stim_circuit` returns it. Its bookkeeping functions became methods: `detector` is now `add_detector` and `observable_include` is now `include_observable`; `tick`, `correlated_error` and `finalize_correlated_error` keep their names; `ensure_lane`, `add_lane`, `add_quantum_boundary`, `last_col` and `last_edge` are now private. Its fields are private behind read-only properties (`rec` is now `records`). `ANNOTATION_ROW`, `HELPER_LANE` and `_error_phase` (now `error_phase`) moved to `tsim.core.builder`. `build_sampling_graph` (`tsim.core.graph`) is replaced by `GraphBuilder.build(closed=..., outputs=...)`.
 - Internals moved or removed: `classify_direct` moved from `tsim.core.graph` to `tsim.compile.pipeline`; `tsim.utils.diagram.render_pyzx_d3`, `squash_graph` and `scale_horizontally` are gone, since ZX diagrams use zxparam's `draw` layout and scale options.
+- Internals of `strategy="active_statevector"`: a component's tape keeps a
+  boolean column per symbol instead of uint32 words, `run_lanes` writes a pruned
+  shot's outputs when its lane dies instead of keeping graves, and both samplers
+  pick their backend in one `_sample_outputs` method. Samples are unchanged.
+  (#223)
 
 ### Deprecated
 - The `"pyzx"`, `"pyzx-dets"` and `"pyzx-meas"` diagram types. They still work as aliases for `"zx"`, `"zx-dets"` and `"zx-meas"` but emit a `DeprecationWarning`.

src/tsim/sampler.py

@@ -371,19 +371,46 @@ def __init__(
         self._direct_detector_mask = np.zeros(self._num_detectors, dtype=np.bool_)
         self._direct_detector_mask[columns[is_detector]] = True
 
-    def _sample_shot_outputs(
+    def _sample_outputs(
         self,
         shots: int,
         batch_size: int | None,
         columns: list[np.ndarray],
         *,
         bit_packed: bool,
         mask: np.ndarray | None = None,
-        flip: np.ndarray | None = None,
+        referenced: np.ndarray | None = None,
     ) -> list[np.ndarray]:
-        """Return the output bits ``columns`` of ``shots`` shots (``ShotSampler.sample``)."""
-        assert self._shots is not None
+        """Return the output bits of ``shots`` shots, an array per entry of ``columns``.
+
+        Args:
+            shots: Number of shots.
+            batch_size: As in ``CompiledDetectorSampler.sample``.
+            columns: The output columns of each array.
+            bit_packed: Whether the arrays are bit-packed.
+            mask: The postselected detectors, as in
+                ``CompiledDetectorSampler.sample``, if any.
+            referenced: The outputs that a noiseless reference sample is XOR-ed
+                into, before the postselection tests, if any.
+
+        """
         _check_sizes(shots, batch_size)
+        if self._shots is None:
+            return self._sample_stabilizer_rank_outputs(
+                shots,
+                batch_size,
+                columns,
+                bit_packed=bit_packed,
+                mask=mask,
+                referenced=referenced,
+            )
+        flip = None if referenced is None else self._reference_sample() & referenced
+        # NOTE: a discarded shot reads False past its direct detectors, which only
+        # the shot words write; the noise sampler writes direct outputs straight
+        # into boolean rows otherwise.
+        if not (self._program.components or bit_packed or mask is not None):
+            samples = self._shots.direct_bools(shots, flip)
+            return [_take_columns(samples, c) for c in columns]
         arrays, deviation = self._shots.sample(
             shots,
             columns,
@@ -395,6 +422,48 @@ def _sample_shot_outputs(
         _check_norm(deviation, StatevectorComponent.norm_warning_threshold)
         return arrays
 
+    def _sample_stabilizer_rank_outputs(
+        self,
+        shots: int,
+        batch_size: int | None,
+        columns: list[np.ndarray],
+        *,
+        bit_packed: bool,
+        mask: np.ndarray | None,
+        referenced: np.ndarray | None,
+    ) -> list[np.ndarray]:
+        """Return the output bits ``columns`` of ``shots`` shots, sampled with JAX.
+
+        The arguments are those of ``_sample_outputs``. A mask only changes the
+        samples through direct detectors that can fire, so other masks are
+        sampled without postselection.
+        """
+        if mask is not None:
+            skips = mask & self._direct_detector_mask
+            indices, flips, direct_columns = self._direct_detectors
+            constant = indices < 0
+            fires = flips[constant]
+            if referenced is not None:
+                fires = fires & ~referenced[direct_columns[constant]]
+            skips[direct_columns[constant]] &= fires
+            if not skips.any():
+                mask = None
+        if mask is not None:
+            samples = self._sample_batches_with_postselection(
+                shots, batch_size, postselection_mask=mask, referenced=referenced
+            )
+        elif referenced is not None:
+            samples, reference = self._sample_batches(
+                shots, batch_size, compute_reference=True
+            )
+            samples ^= reference & referenced
+        else:
+            samples = self._sample_batches(shots, batch_size)
+        return [
+            _maybe_bit_pack(_take_columns(samples, c), bit_packed=bit_packed)
+            for c in columns
+        ]
+
     def _reference_sample(self) -> np.ndarray:
         """Return a noiseless sample of every output (all f-params 0).
 
@@ -737,15 +806,8 @@ def sample(self, shots: int, *, batch_size: int | None = None) -> np.ndarray:
             A numpy array containing the measurement samples.
 
         """
-        if self._shots is not None:
-            columns = [np.arange(self._program.num_outputs)]
-            if not self._program.components:
-                _check_sizes(shots, batch_size)
-                return self._shots.direct_bools(shots)
-            return self._sample_shot_outputs(
-                shots, batch_size, columns, bit_packed=False
-            )[0]
-        return self._sample_batches(shots, batch_size)
+        columns = [np.arange(self._program.num_outputs)]
+        return self._sample_outputs(shots, batch_size, columns, bit_packed=False)[0]
 
 
 def _maybe_bit_pack(array: np.ndarray, *, bit_packed: bool) -> np.ndarray:
@@ -918,84 +980,38 @@ def sample(
                 "append_observables=True or prepend_observables=True"
             )
 
-        compute_reference = (
-            use_detector_reference_sample or use_observable_reference_sample
-        )
         num_detectors = self._num_detectors
-        mask_given = postselection_mask is not None
-
+        mask = None
         if postselection_mask is not None:
             mask = np.asarray(postselection_mask, dtype=np.bool_)
             if mask.shape != (num_detectors,):
                 raise ValueError(
                     f"postselection_mask must have shape ({num_detectors},), "
                     f"got {mask.shape}"
                 )
-            # NOTE: shot words discard and prune with any mask; with a
-            # stabilizer-rank strategy only direct detectors that can fire skip work.
-            if self._shots is not None:
-                postselection_mask = mask
-            else:
-                skips = mask & self._direct_detector_mask
-                indices, flips, columns = self._direct_detectors
-                constant = indices < 0
-                skips[columns[constant]] &= flips[constant] & (
-                    not use_detector_reference_sample
-                )
-                postselection_mask = mask if skips.any() else None
-
         num_outputs = self._program.num_outputs
-        is_detector = np.arange(num_outputs) < num_detectors
-        referenced = np.where(
-            is_detector, use_detector_reference_sample, use_observable_reference_sample
-        )
+        referenced = None
+        if use_detector_reference_sample or use_observable_reference_sample:
+            referenced = np.where(
+                np.arange(num_outputs) < num_detectors,
+                use_detector_reference_sample,
+                use_observable_reference_sample,
+            )
         columns = _output_columns(
             num_detectors,
             num_outputs,
             prepend=prepend_observables,
             append=append_observables,
             separate=separate_observables,
         )
-        if self._shots is not None:
-            # NOTE: the reference bits are XOR-ed in before the postselection tests.
-            flip = self._reference_sample() & referenced if compute_reference else None
-            # NOTE: a discarded shot reads False past its direct detectors, which
-            # only the shot words write; direct_bools serves the other calls.
-            if self._program.components or bit_packed or mask_given:
-                arrays = self._sample_shot_outputs(
-                    shots,
-                    batch_size,
-                    columns,
-                    bit_packed=bit_packed,
-                    mask=postselection_mask,
-                    flip=flip,
-                )
-                return (arrays[0], arrays[1]) if separate_observables else arrays[0]
-            _check_sizes(shots, batch_size)
-            samples = self._shots.direct_bools(shots, flip)
-        elif postselection_mask is not None:
-            if not compute_reference:
-                referenced = None
-            samples = self._sample_batches_with_postselection(
-                shots,
-                batch_size,
-                postselection_mask=postselection_mask,
-                referenced=referenced,
-            )
-        elif compute_reference:
-            samples, reference = self._sample_batches(
-                shots, batch_size, compute_reference=True
-            )
-            if use_detector_reference_sample:
-                samples[:, :num_detectors] ^= reference[:num_detectors]
-            if use_observable_reference_sample:
-                samples[:, num_detectors:] ^= reference[num_detectors:]
-        else:
-            samples = self._sample_batches(shots, batch_size)
-        arrays = [
-            _maybe_bit_pack(_take_columns(samples, c), bit_packed=bit_packed)
-            for c in columns
-        ]
+        arrays = self._sample_outputs(
+            shots,
+            batch_size,
+            columns,
+            bit_packed=bit_packed,
+            mask=mask,
+            referenced=referenced,
+        )
         return (arrays[0], arrays[1]) if separate_observables else arrays[0]
 
 

src/tsim/statevector/program.py

@@ -8,9 +8,8 @@
 
 from __future__ import annotations
 
-from collections.abc import Iterable, Sequence
+from collections.abc import Sequence
 from dataclasses import dataclass
-from math import ceil
 from typing import ClassVar, NamedTuple
 
 import numpy as np
@@ -49,9 +48,9 @@ class Tape(NamedTuple):
             ``2^extent`` and above are zero before and after it.
         branch: Symbol a measurement sets, -1 for a rotation.
         sign_const: Constant of each rotation's sign.
-        sign_row: Symbol bits of each rotation's sign, shape ``(A, num_words)``.
+        sign_row: Symbols of each rotation's sign, shape ``(A, num_symbols)``.
         output_const: Constant of each output, shape ``(num_outputs,)``.
-        output_row: Symbol bits of each output, shape ``(num_outputs, num_words)``.
+        output_row: Symbols of each output, shape ``(num_outputs, num_symbols)``.
 
     """
 
@@ -98,8 +97,8 @@ def num_actions(self) -> int:
         return int(self.tape.kind.shape[0])
 
     @property
-    def num_words(self) -> int:
-        """Return the number of uint32 words that hold the tape's symbols."""
+    def num_symbols(self) -> int:
+        """Return the number of f-bit and branch symbols the tape reads."""
         return int(self.tape.sign_row.shape[1])
 
     @property
@@ -137,18 +136,9 @@ def plan_component(
     return plan(ops, scheduled), component_layout
 
 
-def _pack(terms: Iterable[int], num_words: int) -> np.ndarray:
-    """Return symbol bits as uint32 words: symbol s is bit s % 32 of word s // 32."""
-    row = np.zeros(num_words, dtype=np.uint32)
-    for term in terms:
-        row[term // 32] |= np.uint32(1) << np.uint32(term % 32)
-    return row
-
-
 def lower(plan: Plan, layout: ComponentLayout) -> StatevectorComponent:
     """Return the action tape of ``plan``."""
     num_actions = len(plan.actions)
-    num_words = max(1, ceil(plan.num_symbols / 32))
     kind = np.zeros(num_actions, dtype=np.int32)
     x_mask = np.zeros(num_actions, dtype=np.uint32)
     z_mask = np.zeros(num_actions, dtype=np.uint32)
@@ -157,14 +147,14 @@ def lower(plan: Plan, layout: ComponentLayout) -> StatevectorComponent:
     extent = np.zeros(num_actions, dtype=np.int32)
     branch = np.full(num_actions, -1, dtype=np.int32)
     sign_const = np.zeros(num_actions, dtype=bool)
-    sign_row = np.zeros((num_actions, num_words), dtype=np.uint32)
+    sign_row = np.zeros((num_actions, plan.num_symbols), dtype=bool)
     active = 0
     for a, action in enumerate(plan.actions):
         if isinstance(action, Rotate):
             kind[a] = ROTATE if action.x else ROTATE_DIAG
             x_mask[a], z_mask[a], angle[a] = action.x, action.z, action.angle
             sign_const[a] = action.sign.constant
-            sign_row[a] = _pack(action.sign.terms, num_words)
+            sign_row[a, list(action.sign.terms)] = True
             active |= action.x | action.z
         elif isinstance(action, MeasureActive):
             kind[a] = MEASURE if action.x else MEASURE_DIAG
@@ -176,9 +166,9 @@ def lower(plan: Plan, layout: ComponentLayout) -> StatevectorComponent:
         extent[a] = active.bit_length()
         if isinstance(action, MeasureActive):
             active &= ~(1 << action.pivot)
-    output_row = np.zeros((len(plan.outputs), num_words), dtype=np.uint32)
+    output_row = np.zeros((len(plan.outputs), plan.num_symbols), dtype=bool)
     for slot, output in enumerate(plan.outputs):
-        output_row[slot] = _pack(output.terms, num_words)
+        output_row[slot, list(output.terms)] = True
     tape = Tape(
         f_selection=np.array(layout.f_selection, dtype=np.int32),
         kind=kind,