From e5cffcc3de03243d9f82c987aa4ceb3f7e27be31 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 27 Oct 2026 21:00:23 -0700 Subject: [PATCH 1/7] Test the GF(2) relation elimination against brute force _Relations turns a component's classical spiders into output parities and was only covered through whole-circuit distributions. The test sums the sign over the eliminated variables of random small systems and checks the reduced rows, and that the reduction raises exactly when no such rows exist. Co-Authored-By: Claude Opus 5.5 --- test/unit/statevector/test_pbc.py | 78 +++++++++++++++++++++++++++++++ 1 file changed, 78 insertions(+) diff --git a/test/unit/statevector/test_pbc.py b/test/unit/statevector/test_pbc.py index f9f457dc..5b2252f5 100644 --- a/test/unit/statevector/test_pbc.py +++ b/test/unit/statevector/test_pbc.py @@ -266,3 +266,81 @@ def test_a_dependent_postselection_must_hold_with_certainty(rhs): else: _, targets = _frame(reading, [0], 0) assert targets.shape == (0, 0) + + +def _affine(form: int, f: int) -> int: + """Return the affine form ``form`` of ``(1, f_0, f_1, …)`` at the f-bits ``f``.""" + return (form & (1 | f << 1)).bit_count() & 1 + + +@pytest.mark.parametrize("seed", range(300)) +def test_relations_eliminate_and_reduce_like_brute_force(seed): + """Summing the sign over eliminated variables leaves the reduced rows' indicator. + + The weight of an assignment ``v`` is ``(-1)^Φ`` times whether every equality + holds. Summed over the eliminated variables, it must be, for each f-bit + pattern, a constant times the indicator of the reduced rows; the reduction + raises exactly when some pattern breaks that form. + """ + rng = np.random.default_rng(seed) + num_vars, num_f = int(rng.integers(1, 7)), int(rng.integers(0, 3)) + forms = 1 << (1 + num_f) + relations = _Relations(num_vars) + for _ in range(int(rng.integers(0, 2 * num_vars))): + i, j = (int(v) for v in rng.integers(0, num_vars, 2)) + relations.add_sign(i, j) + for i in range(num_vars): + if rng.random() < 0.3: + relations.lin[i] ^= int(rng.integers(0, forms)) + for _ in range(int(rng.integers(0, num_vars + 1))): + mask = int(rng.integers(1, 1 << num_vars)) + relations.equalities.append((mask, int(rng.integers(0, forms)))) + pairs, lin = list(relations.pairs), list(relations.lin) + equalities = list(relations.equalities) + eliminated = [v for v in range(num_vars) if rng.random() < 0.5] + kept = [int(v) for v in rng.permutation(num_vars) if v not in eliminated] + + def summed(f: int) -> dict[int, int]: + """Return the weight of each kept assignment, summed over the others.""" + total: dict[int, int] = {} + for v in range(1 << num_vars): + if not all( + (mask & v).bit_count() & 1 == _affine(rhs, f) + for mask, rhs in equalities + ): + continue + # NOTE: pairs are symmetric, so the pairs below i count each once; a + # diagonal pair is in lin, where add_sign keeps it. + phase = sum( + (v >> i & 1) + * (((pairs[i] & v) & ((1 << i) - 1)).bit_count() + _affine(lin[i], f)) + for i in range(num_vars) + ) + key = sum((v >> k & 1) << k for k in kept) + total[key] = total.get(key, 0) + (-1) ** (phase % 2) + return total + + def kept_assignments(): + for bits in range(1 << len(kept)): + yield sum((bits >> n & 1) << k for n, k in enumerate(kept)) + + weights = [summed(f) for f in range(1 << num_f)] + + def representable(total: dict[int, int]) -> bool: + values = {w for w in total.values() if w} + return len(values) == 1 + + try: + for v in eliminated: + relations.eliminate(v) + rows = relations.reduce(kept) + except ValueError: + assert not all(representable(total) for total in weights) + return + for f, total in enumerate(weights): + assert representable(total) + for v in kept_assignments(): + holds = all( + (mask & v).bit_count() & 1 == _affine(rhs, f) for _, mask, rhs in rows + ) + assert bool(total.get(v, 0)) == holds From 13b77a99ceb7a6c22d85ec7aa4375bd2b6b25cef Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 27 Oct 2026 21:02:12 -0700 Subject: [PATCH 2/7] Keep a component's symbols as boolean rows until the lane program The plan's symbol sets were packed into uint32 words on the tape, which the lane program unpacked again before packing them into uint64 lane words. The tape now holds a column per symbol, so lane_program packs once and unpack_symbols and the uint32 packer are gone. Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/program.py | 28 +++++++++------------------ src/tsim/statevector/shared.py | 14 +++----------- src/tsim/statevector/shots.py | 3 +-- test/helpers/statevector.py | 2 +- test/unit/statevector/test_program.py | 13 +++++++------ test/unit/statevector/test_shots.py | 4 ++-- 6 files changed, 23 insertions(+), 41 deletions(-) diff --git a/src/tsim/statevector/program.py b/src/tsim/statevector/program.py index 020d812c..9eac40c4 100644 --- a/src/tsim/statevector/program.py +++ b/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, diff --git a/src/tsim/statevector/shared.py b/src/tsim/statevector/shared.py index 85a9fac1..32b9832a 100644 --- a/src/tsim/statevector/shared.py +++ b/src/tsim/statevector/shared.py @@ -146,13 +146,6 @@ class KeyBasis(NamedTuple): sign_coords: np.ndarray -def unpack_symbols(words: np.ndarray, num_words: int) -> np.ndarray: - """Return uint32 symbol words as booleans: symbol s is bit s % 32 of word s // 32.""" - shifts = np.arange(32, dtype=np.uint32) - bits = (words[..., None] >> shifts) & 1 - return bits.reshape(*words.shape[:-1], 32 * num_words).astype(bool) - - def pack_bits(bits: np.ndarray, num_words: int) -> np.ndarray: """Return booleans as uint64 words: bit i is bit i % 64 of word i // 64.""" padded = np.zeros((*bits.shape[:-1], 64 * num_words), dtype=bool) @@ -164,7 +157,7 @@ def key_basis(component: StatevectorComponent) -> KeyBasis: """Return a basis of the f-parts of the component's rotation signs.""" tape = component.tape num_f = len(tape.f_selection) - signs = unpack_symbols(tape.sign_row, component.num_words)[:, :num_f] + signs = tape.sign_row[:, :num_f] rotation = np.isin(tape.kind, (ROTATE_DIAG, ROTATE)) signs = signs & rotation[:, None] nonzero = signs.any(axis=1) @@ -197,9 +190,8 @@ def lane_program( num_actions = len(tape.kind) num_outputs = component.num_outputs kind, branch = tape.kind, tape.branch - symbol_bits = 32 * component.num_words - sign_bits = unpack_symbols(tape.sign_row, component.num_words) - output_bits = unpack_symbols(tape.output_row, component.num_words) + symbol_bits = component.num_symbols + sign_bits, output_bits = tape.sign_row, tape.output_row postselected = np.asarray(postselected, dtype=bool) setter = np.full(symbol_bits, -1, dtype=np.int64) diff --git a/src/tsim/statevector/shots.py b/src/tsim/statevector/shots.py index 6cf18dad..8832b70b 100644 --- a/src/tsim/statevector/shots.py +++ b/src/tsim/statevector/shots.py @@ -93,8 +93,7 @@ def shot_layout(program: CompiledProgram, channels: ChannelSampler) -> ShotLayou assert isinstance(component, StatevectorComponent) tape = component.tape selection = np.asarray(tape.f_selection, dtype=np.int64) - output_bits = shared.unpack_symbols(tape.output_row, component.num_words) - slot, f = np.nonzero(output_bits[:, : len(selection)]) + slot, f = np.nonzero(tape.output_row[:, : len(selection)]) f_index.append(selection[f]) bit.append(np.asarray(component.output_indices, dtype=np.int64)[slot]) constant[list(component.output_indices)] = tape.output_const diff --git a/test/helpers/statevector.py b/test/helpers/statevector.py index 74ed9193..6b442dee 100644 --- a/test/helpers/statevector.py +++ b/test/helpers/statevector.py @@ -460,7 +460,7 @@ def sample_lanes( num_f = len(tape.f_selection) num_outputs = compiled.num_outputs f = np.asarray(f_params)[:, list(tape.f_selection)].astype(np.uint8) - output_bits = shared.unpack_symbols(tape.output_row, compiled.num_words) + output_bits = tape.output_row f_parts = (f @ output_bits[:, :num_f].T.astype(np.uint8)) % 2 == 1 constant = np.asarray(tape.output_const, dtype=bool) if flipped is not None: diff --git a/test/unit/statevector/test_program.py b/test/unit/statevector/test_program.py index 77b51cce..673134d8 100644 --- a/test/unit/statevector/test_program.py +++ b/test/unit/statevector/test_program.py @@ -85,17 +85,18 @@ def test_lower_encodes_every_action(): assert tape.pivot.tolist()[2:4] == [1, 0] assert tape.branch.tolist() == [-1, -1, 1, 2, 3] assert tape.sign_const.tolist()[:2] == [True, False] - assert tape.sign_row[:2, 0].tolist() == [1, 0] + assert tape.sign_row[:2].tolist() == [[True, False, False, False], [False] * 4] assert tape.output_const.tolist() == [False, True] - assert tape.output_row[:, 0].tolist() == [0b1010, 0b0101] - assert compiled.num_words == 1 + assert np.flatnonzero(tape.output_row[0]).tolist() == [1, 3] + assert np.flatnonzero(tape.output_row[1]).tolist() == [0, 2] + assert compiled.num_symbols == 4 assert tape.f_selection.tolist() == [4] assert compiled.output_indices == (2, 5) assert (compiled.width, compiled.num_outputs) == (2, 2) assert compiled.num_actions == 5 -def test_lower_packs_symbols_beyond_32_into_more_words(): +def test_lower_keeps_a_column_per_symbol(): planned = Plan( width=0, num_f=40, @@ -103,8 +104,8 @@ def test_lower_packs_symbols_beyond_32_into_more_words(): outputs=(AffineRow(False, frozenset({0, 33, 39})),), ) compiled = lower(planned, ComponentLayout(tuple(range(40)), {}, (0,))) - assert compiled.num_words == 2 - assert compiled.tape.output_row.tolist() == [[1, (1 << 1) | (1 << 7)]] + assert compiled.num_symbols == 40 + assert np.flatnonzero(compiled.tape.output_row[0]).tolist() == [0, 33, 39] def test_too_wide_components_are_rejected(monkeypatch): diff --git a/test/unit/statevector/test_shots.py b/test/unit/statevector/test_shots.py index c27be1b6..abbd4fdf 100644 --- a/test/unit/statevector/test_shots.py +++ b/test/unit/statevector/test_shots.py @@ -69,8 +69,8 @@ def test_noise_writes_each_outputs_f_part_and_each_key(name, rows): ): tape = component.tape f_bits = f[:, list(tape.f_selection)] - output_bits = shared.unpack_symbols(tape.output_row, component.num_words) - f_parts = (f_bits @ output_bits[:, : len(tape.f_selection)].T) % 2 == 1 + output_bits = tape.output_row[:, : len(tape.f_selection)].astype(np.uint8) + f_parts = (f_bits @ output_bits.T) % 2 == 1 expected[:, list(component.output_indices)] = f_parts keys = (f_bits @ key.rows.T.astype(np.uint8)) % 2 == 1 np.testing.assert_array_equal( From 2d5740850ab5bacb090fdde93a630adb1ff68c44 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 27 Oct 2026 21:06:09 -0700 Subject: [PATCH 3/7] Finish pruned shots when their lane dies run_lanes kept dying lanes in a grave buffer and found their shots at the next random measurement or at the end, with codes that pointed into the graves, per-shot grave lists and a hidden mask per entry. A pruning now writes the dying lanes' shots' outputs right away and drops them from the live list, so the graves and their bookkeeping are gone. The samples and lane counts are unchanged. Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/shared.py | 261 +++++++++++++++++----------- test/unit/statevector/test_shots.py | 2 +- 2 files changed, 162 insertions(+), 101 deletions(-) diff --git a/src/tsim/statevector/shared.py b/src/tsim/statevector/shared.py index 32b9832a..47eb116c 100644 --- a/src/tsim/statevector/shared.py +++ b/src/tsim/statevector/shared.py @@ -20,8 +20,8 @@ Pruning. After the action that sets a postselected output's last branch bit (or before the first action, for an output of f-bits only) the lanes in which it reads -1 die. Their shots keep the outcomes drawn so far, and their outputs that need a -later branch bit read False. +1 die. Their shots get their outputs there, from the outcomes drawn so far, and +those that need a later branch bit read False. Outputs. Each shot XORs the branch part of its outputs, from its last lane, and its own fair coins: the dormant measurements that no sign and no postselected @@ -314,11 +314,11 @@ def lane_bytes(component: StatevectorComponent, key: KeyBasis) -> int: """Return the peak memory of one lane of ``run_lanes``. That is three rows of amplitudes (the old and the new array while lanes grow), - four rows of lane symbols (its own, and a grave's while the graves grow) and - 64 bytes of per-lane scratch. + two rows of lane symbols (its own, and its output pattern while it is + pruned) and 64 bytes of per-lane scratch. """ amplitudes = 3 * 2**component.width * np.dtype(np.complex128).itemsize - return amplitudes + 4 * 8 * max_lane_words(component, key) + 64 + return amplitudes + 2 * 8 * max_lane_words(component, key) + 64 @intrinsic @@ -589,54 +589,129 @@ def _group(words, live, n, program, share): return symbols, row, count +@jit +def _lane_pattern(symbols, r, branch_row, output_word, output_bit, num_words): + """Return the branch parts of the outputs of lane ``r``, on the output words.""" + pattern = np.zeros(num_words, dtype=np.uint64) + for o in range(len(output_bit)): + if _row_parity(symbols, r, branch_row[o]): + pattern[output_word[o]] ^= np.uint64(1) << np.uint64(output_bit[o]) + return pattern + + +# NOTE: _finish_shot runs once per shot, so, like _lane_key, it takes only the +# arrays it reads and not the LaneProgram. +@jit +def _finish_shot( + words, shot, pattern, hidden, out_words, coin_pattern, seed, index, slots, first +): + """Add a shot's branch parts and coins to its outputs, then clear ``hidden``. + + The coins come from the shot's draw slots ``first`` on, 64 per slot. + """ + for w in range(len(out_words)): + words[shot, out_words[w]] ^= pattern[w] + coins = np.uint64(0) + for j in range(coin_pattern.shape[0]): + if j % 64 == 0: + coins = _draw(seed, index, slots, first + j // 64) + if (coins >> np.uint64(j % 64)) & np.uint64(1): + for w in range(len(out_words)): + words[shot, out_words[w]] ^= coin_pattern[j, w] + for w in range(len(out_words)): + words[shot, out_words[w]] &= ~hidden[w] + + @jit def _prune( amps, symbols, size, span, - tests, - prune_row, - graves, - grave_at, - dead, + live, + row, + n, code_row, + words, + alive, + first_shot, + seed, + slots, + program, + tests, action, ): - """Bury the lanes in which a tested output reads 1. + """Finish the shots of the lanes in which a tested output reads 1. - The other lanes move down in order, and ``code_row`` maps codes to their new - rows, or to the graves. Returns ``(size, graves, grave_at, dead)``. + Those shots get their outputs as they stand after ``action``, with the + outputs that need a later branch bit cleared, and leave ``live`` and + ``alive``. The other lanes move down in order, and ``code_row`` maps codes + to their new rows. Returns ``(size, n, pruned lanes)``. """ - new_row = np.empty(size, dtype=np.int64) - kept = 0 + dies = np.zeros(size, dtype=np.bool_) + pruned = 0 for r in range(size): - dies = False for o in range(len(tests)): - if tests[o] and _row_parity(symbols, r, prune_row[o]): - dies = True + if tests[o] and _row_parity(symbols, r, program.prune_row[o]): + dies[r] = True + pruned += 1 break - if dies: - if dead == len(grave_at): - graves = _grown(graves, 2 * dead, dead) - at = np.zeros(2 * dead, dtype=np.int64) - at[:dead] = grave_at - grave_at = at - graves[dead] = symbols[r] - grave_at[dead] = action - new_row[r] = -1 - dead - dead += 1 + if pruned == 0: + return size, n, 0 + out_words = program.out_words + hidden = np.zeros(len(out_words), dtype=np.uint64) + for o in range(len(program.output_bit)): + if program.ready[o] > action: + w = program.output_word[o] + hidden[w] |= np.uint64(1) << np.uint64(program.output_bit[o]) + pattern = np.zeros((size, len(out_words)), dtype=np.uint64) + for r in range(size): + if dies[r]: + pattern[r] = _lane_pattern( + symbols, + r, + program.branch_row, + program.output_word, + program.output_bit, + len(out_words), + ) + coins = program.coin_pattern + first_coin = len(program.kind) + m = 0 + for i in range(n): + shot, r = live[i], code_row[row[i]] + if dies[r]: + _finish_shot( + words, + shot, + pattern[r], + hidden, + out_words, + coins, + seed, + first_shot + shot, + slots, + first_coin, + ) + alive[shot] = False + continue + live[m], row[m] = shot, row[i] + m += 1 + new_row = np.empty(size, dtype=np.int64) + kept = 0 + for r in range(size): + if dies[r]: + new_row[r] = -1 continue if kept != r: _copy_lane(amps, kept, r, span) symbols[kept] = symbols[r] new_row[r] = kept kept += 1 - if kept < size: - for code in range(len(code_row)): - if code_row[code] >= 0: - code_row[code] = new_row[code_row[code]] - return kept, graves, grave_at, dead + for code in range(len(code_row)): + if code_row[code] >= 0: + code_row[code] = new_row[code_row[code]] + return kept, m, pruned @kernel @@ -649,8 +724,8 @@ def run_lanes( and ``program`` is the component's ``LaneProgram``. The shots' output bits get the component's branch parts and coins, and those that a pruning hides are cleared. A pruned shot leaves ``alive``, so later components skip it. The - amplitude loops run in up to ``4 * threads`` parallel chunks once - their work reaches ``parallel_work``. + amplitude loops run in up to ``4 * threads`` parallel chunks once their work + reaches ``parallel_work``. Returns: ``(deviation, groups, peak, size, dead)``: the largest normalization @@ -678,30 +753,29 @@ def run_lanes( amps = np.zeros((max(1, size), 1 << program.width), dtype=np.complex128) for r in range(size): amps[r, 0] = 1 - # NOTE: each live shot holds a code, and ``code_row`` maps codes to rows, or to - # graves as ``-1 - grave``: splits and prunes rewrite the map in a step per - # lane, and the next pass over the shots drops those in graves. + # NOTE: each live shot holds a code, and ``code_row`` maps codes to rows: a + # split rewrites the map in a step per lane, and a pruning in a step per code. code_row = np.arange(max(size, 1), dtype=np.int64) - graves = np.zeros((16, program.lane_words), dtype=np.uint64) - grave_at = np.zeros(16, dtype=np.int64) dead = 0 - dead_shot = np.empty(n, dtype=np.int64) - dead_grave = np.empty(n, dtype=np.int64) - num_dead_shots = 0 region = 0 deviation = 0.0 if size and program.prune_first.any(): - size, graves, grave_at, dead = _prune( + size, n, dead = _prune( amps, symbols, size, 1, - program.prune_first, - program.prune_row, - graves, - grave_at, - dead, + live, + row, + n, code_row, + words, + alive, + first_shot, + seed, + slots, + program, + program.prune_first, -1, ) @@ -775,22 +849,14 @@ def run_lanes( new = size if any_random: seen = np.zeros(2 * size, dtype=np.bool_) - m = 0 for i in range(n): shot, r = live[i], code_row[row[i]] - if r < 0: - dead_shot[num_dead_shots] = shot - dead_grave[num_dead_shots] = -1 - r - num_dead_shots += 1 - continue c = certain[r] if c < 0: u = _draw(seed, first_shot + shot, slots, a) c = np.int64(u < threshold[r]) - live[m], row[m] = shot, 2 * r + c + row[i] = 2 * r + c seen[2 * r + c] = True - m += 1 - n = m for r in range(size): if seen[2 * r] and seen[2 * r + 1]: new += 1 @@ -836,56 +902,51 @@ def run_lanes( if outcome[r]: symbols[r, bit >> 6] |= np.uint64(1) << np.uint64(bit & 63) if program.pruned[a]: - size, graves, grave_at, dead = _prune( + size, n, pruned = _prune( amps, symbols, size, 1 << region, - program.prune_outputs[a], - program.prune_row, - graves, - grave_at, - dead, + live, + row, + n, code_row, + words, + alive, + first_shot, + seed, + slots, + program, + program.prune_outputs[a], a, ) + dead += pruned - # NOTE: each entry (a live lane or a grave) has a branch pattern on the outputs - # and hides the outputs that its pruning makes unknown; every shot takes its - # entry's. - entries = size + dead - pattern = np.zeros((entries, len(out_words)), dtype=np.uint64) - hidden = np.zeros((entries, len(out_words)), dtype=np.uint64) - for e in range(entries): - lane_symbols = symbols if e < size else graves - r = e if e < size else e - size - for o in range(len(program.output_bit)): - w = program.output_word[o] - one = np.uint64(1) << np.uint64(program.output_bit[o]) - if _row_parity(lane_symbols, r, program.branch_row[o]): - pattern[e, w] ^= one - if e >= size and program.ready[o] > grave_at[r]: - hidden[e, w] |= one - for i in range(n + num_dead_shots): - if i < n: - shot, e = live[i], code_row[row[i]] - if e < 0: - e = size - 1 - e - else: - shot, e = dead_shot[i - n], size + dead_grave[i - n] - if e >= size: - alive[shot] = False - for w in range(len(out_words)): - words[shot, out_words[w]] ^= pattern[e, w] - coins = np.uint64(0) - for j in range(num_coins): - if j % 64 == 0: - coins = _draw(seed, first_shot + shot, slots, num_actions + j // 64) - if (coins >> np.uint64(j % 64)) & np.uint64(1): - for w in range(len(out_words)): - words[shot, out_words[w]] ^= program.coin_pattern[j, w] - for w in range(len(out_words)): - words[shot, out_words[w]] &= ~hidden[e, w] + pattern = np.zeros((size, len(out_words)), dtype=np.uint64) + for r in range(size): + pattern[r] = _lane_pattern( + symbols, + r, + program.branch_row, + program.output_word, + program.output_bit, + len(out_words), + ) + hidden = np.zeros(len(out_words), dtype=np.uint64) + for i in range(n): + shot = live[i] + _finish_shot( + words, + shot, + pattern[code_row[row[i]]], + hidden, + out_words, + program.coin_pattern, + seed, + first_shot + shot, + slots, + num_actions, + ) return deviation, groups, peak, size, dead diff --git a/test/unit/statevector/test_shots.py b/test/unit/statevector/test_shots.py index abbd4fdf..5f0b0166 100644 --- a/test/unit/statevector/test_shots.py +++ b/test/unit/statevector/test_shots.py @@ -150,7 +150,7 @@ def test_the_memory_estimate_holds_every_lane_symbol(): program = max(sampler.plan(None).programs, key=lambda p: p.lane_words) assert program.lane_words >= 4 amplitudes = 3 * 16 * 2**program.width - assert sampler.bytes_per_shot() >= amplitudes + 4 * 8 * program.lane_words + assert sampler.bytes_per_shot() >= amplitudes + 2 * 8 * program.lane_words @pytest.mark.parametrize("postselect", [False, True]) From 734be13b05cfb4f13097a4073baf64221ab9db90 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 27 Oct 2026 21:11:53 -0700 Subject: [PATCH 4/7] Sample every call through one method of the sampler base CompiledMeasurementSampler.sample and CompiledDetectorSampler.sample each branched on the shot-word and stabilizer-rank backends, and the detector sampler prepared masks and reference samples for both inline. _sample_outputs now takes the columns, the mask and the referenced outputs, and picks the backend once; the stabilizer-rank path keeps its own mask screening and reference handling. The samples are unchanged. Co-Authored-By: Claude Opus 5.5 --- src/tsim/sampler.py | 166 ++++++++++++++++++++++++-------------------- 1 file changed, 91 insertions(+), 75 deletions(-) diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index 28770fcb..d7b8fba7 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -371,7 +371,7 @@ 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, @@ -379,11 +379,38 @@ def _sample_shot_outputs( *, 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 ``columns`` of ``shots`` shots, one array per list. + + 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,12 +980,8 @@ 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,): @@ -931,24 +989,14 @@ def sample( 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, @@ -956,46 +1004,14 @@ def sample( 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] From 32f9915686a21a5bbb10eb44d81ca37e8101955a Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 27 Oct 2026 21:12:24 -0700 Subject: [PATCH 5/7] Add the refactor to the changelog Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index d1e4eba9..be0cd4e3 100644 --- a/CHANGELOG.md +++ b/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`. From df9370773e39a2edaa7000099f418546f50ce93d Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Mon, 2 Nov 2026 16:24:44 -0500 Subject: [PATCH 6/7] Inline the final shot loop again and check the echelon form of reduced rows run_lanes finished every live shot through an out-of-line _finish_shot call with an all-zero hidden mask, which made the shared, unpruned path about 1.4x slower on one thread. The end-of-run loop is inline again, and _finish_shot serves only pruned shots. The _Relations brute-force test now checks that the rows are in reduced echelon form with pivots in kept order, which catches dropping the back-substitution or reversing the pivot priority. Co-Authored-By: Claude Opus 5.5 --- src/tsim/sampler.py | 2 +- src/tsim/statevector/shared.py | 30 ++++++++++++++---------------- test/unit/statevector/test_pbc.py | 5 +++++ 3 files changed, 20 insertions(+), 17 deletions(-) diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index d7b8fba7..38964500 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -381,7 +381,7 @@ def _sample_outputs( mask: np.ndarray | None = None, referenced: np.ndarray | None = None, ) -> list[np.ndarray]: - """Return the output bits ``columns`` of ``shots`` shots, one array per list. + """Return the output bits of ``shots`` shots, an array per entry of ``columns``. Args: shots: Number of shots. diff --git a/src/tsim/statevector/shared.py b/src/tsim/statevector/shared.py index 47eb116c..2d81b6fb 100644 --- a/src/tsim/statevector/shared.py +++ b/src/tsim/statevector/shared.py @@ -599,8 +599,8 @@ def _lane_pattern(symbols, r, branch_row, output_word, output_bit, num_words): return pattern -# NOTE: _finish_shot runs once per shot, so, like _lane_key, it takes only the -# arrays it reads and not the LaneProgram. +# NOTE: _finish_shot runs once per pruned shot, so, like _lane_key, it takes only +# the arrays it reads and not the LaneProgram. @jit def _finish_shot( words, shot, pattern, hidden, out_words, coin_pattern, seed, index, slots, first @@ -932,21 +932,19 @@ def run_lanes( program.output_bit, len(out_words), ) - hidden = np.zeros(len(out_words), dtype=np.uint64) + # NOTE: inline rather than _finish_shot: this loop runs once per live shot, + # and the call costs about 1.7x on one thread when lanes are shared. for i in range(n): - shot = live[i] - _finish_shot( - words, - shot, - pattern[code_row[row[i]]], - hidden, - out_words, - program.coin_pattern, - seed, - first_shot + shot, - slots, - num_actions, - ) + shot, e = live[i], code_row[row[i]] + for w in range(len(out_words)): + words[shot, out_words[w]] ^= pattern[e, w] + coins = np.uint64(0) + for j in range(num_coins): + if j % 64 == 0: + coins = _draw(seed, first_shot + shot, slots, num_actions + j // 64) + if (coins >> np.uint64(j % 64)) & np.uint64(1): + for w in range(len(out_words)): + words[shot, out_words[w]] ^= program.coin_pattern[j, w] return deviation, groups, peak, size, dead diff --git a/test/unit/statevector/test_pbc.py b/test/unit/statevector/test_pbc.py index 5b2252f5..56199ff5 100644 --- a/test/unit/statevector/test_pbc.py +++ b/test/unit/statevector/test_pbc.py @@ -337,6 +337,11 @@ def representable(total: dict[int, int]) -> bool: except ValueError: assert not all(representable(total) for total in weights) return + # NOTE: reduced echelon form, pivots taken in ``kept`` order. + for p, mask, _ in rows: + assert mask >> p & 1 + assert min((v for v in kept if mask >> v & 1), key=kept.index) == p + assert not any(m >> p & 1 for q, m, _ in rows if q != p) for f, total in enumerate(weights): assert representable(total) for v in kept_assignments(): From 2912950cc846e5beed3853a7797c2b86546cb591 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Mon, 2 Nov 2026 16:25:22 -0500 Subject: [PATCH 7/7] Check the pivots' order and the row terms in the _Relations test Co-Authored-By: Claude Opus 5.5 --- test/unit/statevector/test_pbc.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/test/unit/statevector/test_pbc.py b/test/unit/statevector/test_pbc.py index 56199ff5..959105b4 100644 --- a/test/unit/statevector/test_pbc.py +++ b/test/unit/statevector/test_pbc.py @@ -338,9 +338,11 @@ def representable(total: dict[int, int]) -> bool: assert not all(representable(total) for total in weights) return # NOTE: reduced echelon form, pivots taken in ``kept`` order. + pivots = [p for p, _, _ in rows] + assert pivots == sorted(set(pivots), key=kept.index) for p, mask, _ in rows: - assert mask >> p & 1 - assert min((v for v in kept if mask >> v & 1), key=kept.index) == p + terms = [v for v in range(num_vars) if mask >> v & 1] + assert set(terms) <= set(kept) and min(terms, key=kept.index) == p assert not any(m >> p & 1 for q, m, _ in rows if q != p) for f, total in enumerate(weights): assert representable(total)