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`. diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index 28770fcb..38964500 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 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,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] 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..2d81b6fb 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 @@ -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) @@ -322,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 @@ -597,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 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 +): + """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 @@ -657,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 @@ -686,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, ) @@ -783,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 @@ -844,45 +902,40 @@ 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 + 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), + ) + # 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, 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) @@ -892,8 +945,6 @@ def run_lanes( 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] return deviation, groups, peak, size, dead 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_pbc.py b/test/unit/statevector/test_pbc.py index f9f457dc..959105b4 100644 --- a/test/unit/statevector/test_pbc.py +++ b/test/unit/statevector/test_pbc.py @@ -266,3 +266,88 @@ 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 + # 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: + 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) + 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 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..5f0b0166 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( @@ -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])