diff --git a/CHANGELOG.md b/CHANGELOG.md index 41824cd6..cd92d6b5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -121,6 +121,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 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) +- With `strategy="active_statevector"` and a postselection mask, the shots + that pass are gathered by numba kernels in the threads that sampled them, + instead of by NumPy on one thread. Discarded shots are no longer finished: + a pruned shot draws no coins, and its outputs are neither completed nor + cleared. A program without components screens each noise block while it is + in cache, and reuses its noise words across calls: a sampler keeps up to + 16 MiB of them. Samples are unchanged. ### 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/noise/channels.py b/src/tsim/noise/channels.py index 4cd23516..59b3645d 100644 --- a/src/tsim/noise/channels.py +++ b/src/tsim/noise/channels.py @@ -774,6 +774,48 @@ def sample_words( *flips, ) + def sample_screened_words( + self, + out: np.ndarray, + flips: WordTables, + seeds: np.ndarray, + constant: np.ndarray, + screen: np.ndarray, + ) -> tuple[np.ndarray, np.ndarray]: + """``sample_words``, then XOR ``constant`` into each row and screen it. + + Returns whether each row has no ``screen`` bit, and the number of such + rows in each block of ``_BLOCK_SHOTS``. ``constant`` and ``screen`` cover + every word of a row. + """ + num_samples = len(out) + num_blocks = -(-num_samples // _BLOCK_SHOTS) + if len(seeds) != num_blocks: + raise ValueError( + f"{num_samples} samples need {num_blocks} block seeds, " + f"got {len(seeds)}" + ) + if not (len(constant) == len(screen) == out.shape[1]): + raise ValueError("constant and screen must cover every word of a row") + t = self._tables + keep = np.empty(num_samples, dtype=np.bool_) + kept = np.empty(num_blocks, dtype=np.int64) + _sample_screened_blocks( + num_samples * len(t.p_fire), + out, + seeds, + t.p_fire, + t.outcome_start, + t.alias_prob, + t.alias_index, + *flips, + constant, + screen, + keep, + kept, + ) + return keep, kept + class WordTables(NamedTuple): """The words each fault outcome XORs into a row. @@ -949,3 +991,66 @@ def _sample_blocks( word, mask, ) + + +@jit +def _screen_rows(rows, constant, screen, keep): + """XOR ``constant`` into each row and flag in ``keep`` the rows without a ``screen`` bit. + + ``constant`` and ``screen`` cover every word of a row. Returns the number of + rows kept. + """ + num_rows, width = rows.shape + kept = 0 + if width == 1: + # NOTE: a loop over one word would keep LLVM from vectorizing. + c, m = constant[0], screen[0] + for s in range(num_rows): + value = rows[s, 0] ^ c + rows[s, 0] = value + keep[s] = (value & m) == 0 + kept += keep[s] + return kept + for s in range(num_rows): + hit = np.uint64(0) + for w in range(width): + rows[s, w] ^= constant[w] + hit |= rows[s, w] & screen[w] + keep[s] = hit == 0 + kept += keep[s] + return kept + + +@kernel +def _sample_screened_blocks( + out, + seeds, + p_fire, + outcome_start, + alias_prob, + alias_index, + flip_start, + word, + mask, + constant, + screen, + keep, + kept, +): + """``_sample_blocks``, then each block's rows are screened by ``_screen_rows``.""" + for b in numba.prange(len(seeds)): + rows = out[b * _BLOCK_SHOTS : (b + 1) * _BLOCK_SHOTS] + _sample_block( + rows, + seeds[b], + p_fire, + outcome_start, + alias_prob, + alias_index, + flip_start, + word, + mask, + ) + kept[b] = _screen_rows( + rows, constant, screen, keep[b * _BLOCK_SHOTS : (b + 1) * _BLOCK_SHOTS] + ) diff --git a/src/tsim/statevector/shared.py b/src/tsim/statevector/shared.py index 2d81b6fb..f7f67444 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 get their outputs there, from the outcomes drawn so far, and -those that need a later branch bit read False. +1 die. Their shots leave ``alive`` and are not sampled further: they are +discarded, so their outputs are left unfinished. 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 @@ -599,29 +599,6 @@ def _lane_pattern(symbols, r, branch_row, output_word, output_bit, num_words): 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, @@ -632,67 +609,30 @@ def _prune( row, n, code_row, - words, alive, - first_shot, - seed, - slots, - program, + prune_row, tests, - action, ): - """Finish the shots of the lanes in which a tested output reads 1. + """Drop the lanes in which a tested output reads 1, and their shots. - 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)``. + Those shots 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)``. """ dies = np.zeros(size, dtype=np.bool_) pruned = 0 for r in range(size): for o in range(len(tests)): - if tests[o] and _row_parity(symbols, r, program.prune_row[o]): + if tests[o] and _row_parity(symbols, r, prune_row[o]): dies[r] = True pruned += 1 break 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] @@ -722,8 +662,8 @@ def run_lanes( ``words`` holds the chunk's shot words (row ``s`` is shot ``first_shot + s``) 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 + the component's branch parts and coins. A pruned shot leaves ``alive``, so + later components skip it, and its outputs are left unfinished. The amplitude loops run in up to ``4 * threads`` parallel chunks once their work reaches ``parallel_work``. @@ -769,14 +709,9 @@ def run_lanes( row, n, code_row, - words, alive, - first_shot, - seed, - slots, - program, + program.prune_row, program.prune_first, - -1, ) for a in range(num_actions): @@ -911,14 +846,9 @@ def run_lanes( row, n, code_row, - words, alive, - first_shot, - seed, - slots, - program, + program.prune_row, program.prune_outputs[a], - a, ) dead += pruned @@ -932,8 +862,8 @@ def run_lanes( 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. + # NOTE: inline rather than a call: this loop runs once per live shot, and a + # 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)): diff --git a/src/tsim/statevector/shots.py b/src/tsim/statevector/shots.py index 080821dd..e0b6db62 100644 --- a/src/tsim/statevector/shots.py +++ b/src/tsim/statevector/shots.py @@ -17,6 +17,7 @@ from __future__ import annotations import sys +import threading from collections import deque from dataclasses import dataclass @@ -35,6 +36,8 @@ # 2^8 its fixed costs dominate. MAX_CHUNK_SHOTS = 2**16 MIN_CHUNK_SHOTS = 2**8 +# NOTE: the most scratch memory a sampler keeps between calls. +MAX_SCRATCH_BYTES = 2**24 def memory_budget() -> int: @@ -126,22 +129,15 @@ class ShotPlan: Attributes: screen: The postselected direct outputs, as words: a shot with one of them set is discarded before the components. - kept: The outputs a discarded shot keeps (its direct detectors), as words. - pruned_kept: The outputs a shot that component ``k`` prunes keeps, as - words: ``kept`` and component ``k``'s outputs, which it clears - itself past the pruning. None for a component without postselected - outputs, which prunes no shot. programs: Each component's lane program. - rejected: The postselected detectors, as words: a shot with one of them - set was discarded, and the others are returned. + rejected: Whether any detector is postselected: if so, only the shots + that pass the screen and every component are returned. """ screen: np.ndarray - kept: np.ndarray - pruned_kept: tuple[np.ndarray | None, ...] programs: tuple[LaneProgram, ...] - rejected: np.ndarray + rejected: bool def shot_plan( @@ -152,7 +148,6 @@ def shot_plan( postselected = np.zeros(num_outputs, dtype=bool) if mask is not None: postselected[:num_detectors] = mask - is_detector = np.arange(num_outputs) < num_detectors programs = tuple( lane_program( component, key, key_start, postselected[list(component.output_indices)] @@ -161,38 +156,22 @@ def shot_plan( layout.components, layout.keys, layout.key_starts, strict=True ) ) - kept = shared.pack_bits(layout.direct & is_detector, words) - pruned_kept = [] - for component in layout.components: - outputs = list(component.output_indices) - if not postselected[outputs].any(): - pruned_kept.append(None) - continue - own = kept.copy() - for o in outputs: - own[o >> 6] |= np.uint64(1) << np.uint64(o & 63) - pruned_kept.append(own) return ShotPlan( screen=shared.pack_bits(postselected & layout.direct, words), - kept=kept, - pruned_kept=tuple(pruned_kept), programs=programs, - rejected=shared.pack_bits(postselected, words), + rejected=bool(postselected.any()), ) @jit -def _screen(words, constant, screen, kept, alive): - """XOR the constants into the outputs; a screened shot keeps only ``kept``.""" +def _screen(words, constant, screen, alive): + """XOR the constants into the outputs; a shot with a ``screen`` bit is not alive.""" for s in range(words.shape[0]): hit = np.uint64(0) for w in range(len(constant)): words[s, w] ^= constant[w] hit |= words[s, w] & screen[w] alive[s] = hit == 0 - if hit: - for w in range(len(kept)): - words[s, w] &= kept[w] @jit @@ -241,6 +220,87 @@ def _unpack(words, runs, width, out): _spread(bits, out, target) +@jit +def _compact(rows, alive, width): + """Move the first ``width`` words of each alive shot to the front, in order. + + Returns the number of alive shots. + """ + target = 0 + for s in range(rows.shape[0]): + if alive[s]: + if target != s: + for w in range(width): + rows[target, w] = rows[s, w] + target += 1 + return target + + +@kernel +def _copy_kept(words, keep, size, kept, starts, out): + """Copy the rows ``keep`` flags in each run of ``size`` rows to ``out``, in order. + + Run ``i`` keeps ``kept[i]`` rows, from row ``starts[i]`` of ``out`` on; runs + copy in parallel. + """ + width = out.shape[1] + for i in numba.prange(starts.shape[0]): + first = i * size + last = min(first + size, words.shape[0]) - 1 + while last >= first and not keep[last]: + last -= 1 + target = starts[i] + if 2 * kept[i] < last + 1 - first: + # NOTE: few kept rows: copy only them. + for s in range(first, last + 1): + if keep[s]: + for w in range(width): + out[target, w] = words[s, w] + target += 1 + elif width == 1: + # NOTE: every row up to the run's last kept one is written to + # ``target``, which only a kept row moves on, so the loop has no + # branch and never writes past the run's rows of ``out``. + for s in range(first, last + 1): + out[target, 0] = words[s, 0] + target += keep[s] + else: + for s in range(first, last + 1): + for w in range(width): + out[target, w] = words[s, w] + target += keep[s] + + +@jit +def _join(words, kept, size, width): + """Move the ``kept[i]`` rows at the front of each run of ``size`` rows together. + + Only the first ``width`` words move. Returns the number of rows kept. + """ + filled = 0 + for i in range(kept.shape[0]): + source = i * size + if source != filled: + for r in range(kept[i]): + for w in range(width): + words[filled + r, w] = words[source + r, w] + filled += kept[i] + return filled + + +def _shrink(words: np.ndarray, rows: int) -> np.ndarray: + """Return the first ``rows`` rows of ``words``, an array no view shares. + + The memory past them is given back without a copy, so that a view of the + survivors does not keep a whole call's rows alive. + + Warning: the resize skips numpy's reference check, so ``words`` must have no + other view left, or that view would read freed memory. + """ + words.resize((rows, words.shape[1]), refcheck=False) + return words + + def bool_outputs(words: np.ndarray, columns: np.ndarray) -> np.ndarray: """Return the output bits ``columns`` of every shot as booleans.""" columns = np.asarray(columns, dtype=np.int64) @@ -318,6 +378,10 @@ def __init__( self._plans: dict[bytes, ShotPlan] = {} self._widest = max((c.width for c in self._layout.components), default=0) self._direct_tables: WordTables | None = None + # NOTE: noise words of a program without components, reused by the calls + # that screen them, since only their survivors are returned. + self._scratch: np.ndarray | None = None + self._scratch_lock = threading.Lock() if not self._layout.components: f = self._layout.direct_f sets = np.flatnonzero(f >= 0) @@ -400,15 +464,7 @@ def sample( noise_seeds = self._channels.block_seeds(shots) seeds = self._next_seeds(self._seeds) - rejected = plan.rejected if plan.rejected.any() else None - def read(words: np.ndarray) -> list[np.ndarray]: - if rejected is not None: - hits = (words[:, : len(rejected)] & rejected).any(axis=1) - if hits.any(): - # NOTE: the output words alone, so that bit-packed outputs - # can be views of them. - words = words[~hits, : self._layout.output_words] return [ read_outputs(words, c, num_outputs, bit_packed=bit_packed) for c in columns @@ -464,14 +520,19 @@ def _sample_part( ) -> tuple[np.ndarray, float]: """Return the shot words of a call's shots ``start … stop-1`` and their deviation. - ``noise_seeds`` seed the call's noise blocks, so ``start`` is a whole - number of blocks, and ``seeds`` its components. + With a postselected detector, only the shots that pass postselection + are returned, and only their output words. ``noise_seeds`` seed the + call's noise blocks, so ``start`` is a whole number of blocks, and + ``seeds`` its components. """ shots = stop - start + width = self._layout.output_words + blocks = noise_seeds[start // _BLOCK_SHOTS : -(-stop // _BLOCK_SHOTS)] + if shots and not self._layout.components and plan.rejected: + return self._sample_screened(shots, blocks, plan, constant), 0.0 words = np.empty((shots, self._layout.num_words), dtype=np.uint64) if shots == 0: return words, 0.0 - blocks = noise_seeds[start // _BLOCK_SHOTS : -(-stop // _BLOCK_SHOTS)] self._channels.sample_words(words, self._layout.noise, blocks) if not self._layout.components and not (constant.any() or plan.screen.any()): # NOTE: direct outputs without constants are their noise bits. @@ -479,12 +540,70 @@ def _sample_part( threads = numba.get_num_threads() chunk, workers = self._schedule(shots, batch_size, threads) - def run(offset: int) -> list[tuple]: + def run(offset: int) -> tuple[list[tuple], int]: rows = words[offset : offset + chunk] - return self._run_chunk(rows, plan, constant, start + offset, seeds, threads) - - results = run_tasks(run, range(0, shots, chunk), workers) - return words, self._record([stats for result in results for stats in result]) + alive = np.empty(len(rows), dtype=np.bool_) + stats = self._run_chunk( + rows, alive, plan, constant, start + offset, seeds, threads + ) + # NOTE: each chunk moves its survivors to its front on its own thread. + kept = _compact(rows, alive, width) if plan.rejected else len(rows) + return stats, kept + + offsets = range(0, shots, chunk) + results = run_tasks(run, offsets, workers) + deviation = self._record([stats for result, _ in results for stats in result]) + if not plan.rejected: + return words, deviation + kept = np.array([k for _, k in results], dtype=np.int64) + filled = _join(words, kept, chunk, width) + # NOTE: the output words alone, so that bit-packed outputs can be views. + if width == words.shape[1]: + return _shrink(words, filled), deviation + # NOTE: a copy even of one row, which a view would not release. + return words[:filled, :width].copy(), deviation + + def __getstate__(self) -> dict: + """Pickle without the scratch words and their lock.""" + state = self.__dict__.copy() + state["_scratch"] = None + del state["_scratch_lock"] + return state + + def __setstate__(self, state: dict) -> None: + """Unpickle, with a new lock for the scratch words.""" + self.__dict__.update(state) + self._scratch_lock = threading.Lock() + + def _sample_screened( + self, shots: int, blocks: np.ndarray, plan: ShotPlan, constant: np.ndarray + ) -> np.ndarray: + """Return the output words of the shots that pass postselection, of ``shots``. + + For a program without components, which postselects direct detectors + only, all of them in the screen. Each noise block is screened while it is + in cache, and the survivors are copied out of words that later calls + reuse: fresh memory would cost a page fault per 4 kB. + """ + width = self._layout.num_words + assert width == self._layout.output_words + # NOTE: a call that finds the scratch taken (another thread's) allocates. + with self._scratch_lock: + scratch, self._scratch = self._scratch, None + if scratch is None or len(scratch) < shots: + scratch = np.empty((shots, width), dtype=np.uint64) + words = scratch[:shots] + keep, kept = self._channels.sample_screened_words( + words, self._layout.noise, blocks, constant, plan.screen + ) + starts = np.zeros(len(kept), dtype=np.int64) + np.cumsum(kept[:-1], out=starts[1:]) + out = np.empty((int(kept.sum()), width), dtype=np.uint64) + _copy_kept(len(out) * width, words, keep, _BLOCK_SHOTS, kept, starts, out) + if scratch.nbytes <= MAX_SCRATCH_BYTES: + with self._scratch_lock: + self._scratch = scratch + return out def direct_bools(self, shots: int, flip: np.ndarray | None = None) -> np.ndarray: """Return every output bit of a program without components, as booleans. @@ -507,6 +626,7 @@ def reference(self) -> tuple[np.ndarray, float]: words = np.zeros((1, self._layout.num_words), dtype=np.uint64) stats = self._run_chunk( words, + np.empty(1, dtype=np.bool_), self.plan(None), self._constant(None), 0, @@ -556,6 +676,7 @@ def _schedule( def _run_chunk( self, rows: np.ndarray, + alive: np.ndarray, plan: ShotPlan, constant: np.ndarray, first_shot: int, @@ -564,24 +685,19 @@ def _run_chunk( ) -> list[tuple]: """Screen the rows' shots and run the components on them. - A shot that a component prunes is skipped by the later components and - keeps only the outputs of ``plan.pruned_kept``, whatever the order of - the components. Alone on its thread, a chunk runs each component's - amplitude loops on ``threads`` threads. + ``alive`` receives whether each shot passed the screen and every + component's postselected outputs: a shot that one prunes is skipped by + the later ones, and its outputs are left as they are, since it is not + returned. Alone on its thread, a chunk runs each component's amplitude + loops on ``threads`` threads. """ - alive = np.empty(len(rows), dtype=np.bool_) - _screen(rows, constant, plan.screen, plan.kept, alive) + _screen(rows, constant, plan.screen, alive) stats = [] - for component, program, seed, kept in zip( - self._layout.components, - plan.programs, - seeds, - plan.pruned_kept, - strict=True, + for component, program, seed in zip( + self._layout.components, plan.programs, seeds, strict=True ): if not alive.any(): break - before = None if kept is None else alive.copy() stats.append( shared.run_lanes( len(rows) << component.width, @@ -596,11 +712,6 @@ def _run_chunk( threads, ) ) - if kept is not None: - assert before is not None - pruned = before & ~alive - if pruned.any(): - rows[pruned, : len(kept)] &= kept return stats def _record(self, stats: list[tuple]) -> float: diff --git a/test/helpers/statevector.py b/test/helpers/statevector.py index 9878bc71..7dbc3217 100644 --- a/test/helpers/statevector.py +++ b/test/helpers/statevector.py @@ -459,7 +459,9 @@ def sample_lanes( The shot words hold the component's outputs (their f-parts, constants and ``flipped`` bits) and its key, as ``tsim.statevector.shots`` lays them out, and - ``run_lanes`` adds the branch parts. Returns ``(samples, deviation, stats)``. + ``run_lanes`` adds the branch parts. Returns ``(samples, deviation, stats)``; + ``stats["alive"]`` tells which shots no pruning discarded, whose samples are + the only finished ones. """ key, program = component_program(compiled, postselected) tape = compiled.tape @@ -497,5 +499,11 @@ def sample_lanes( bits = np.unpackbits( words[:, :output_words].view(np.uint8), axis=1, bitorder="little" ) - stats = {"groups": groups, "peak lanes": peak, "lanes": size, "dead": dead} + stats = { + "groups": groups, + "peak lanes": peak, + "lanes": size, + "dead": dead, + "alive": alive, + } return bits[:, :num_outputs].astype(bool), float(deviation), stats diff --git a/test/unit/statevector/test_shared.py b/test/unit/statevector/test_shared.py index 792bda99..6f455dbb 100644 --- a/test/unit/statevector/test_shared.py +++ b/test/unit/statevector/test_shared.py @@ -123,34 +123,34 @@ def test_pruning_discards_exactly_the_postselected_failures(name, flip, share): flipped[0] = flip patterns = f_patterns(planned.num_f, limit=3) f_params = _f_params(component_layout, patterns, shots_each) - samples, _, _ = sample_lanes( + samples, _, stats = sample_lanes( compiled, f_params, 5, share, postselected, flipped ) for i, f in enumerate(patterns): rows = samples[i * shots_each : (i + 1) * shots_each] - discarded = np.asarray((rows & postselected).any(axis=1)) + discarded = ~stats["alive"][i * shots_each : (i + 1) * shots_each] expected = _postselected_distribution( plan_distribution(planned, f), postselected, flipped ) _assert_counts(rows ^ flipped, discarded, expected) +@pytest.mark.parametrize("share", [True, False]) @pytest.mark.parametrize("flip", [False, True]) @pytest.mark.parametrize("name", NAMES) -def test_pruned_shots_report_false_after_their_pruning_action(name, flip): +def test_pruning_keeps_exactly_the_shots_an_unpruned_run_accepts(name, flip, share): for planned, component_layout, compiled in _compiled(name): postselected = np.ones(compiled.num_outputs, dtype=bool) + postselected[-1] = False flipped = np.full(compiled.num_outputs, flip) f_params = _f_params(component_layout, f_patterns(planned.num_f, 3), 4_000) - rows, _, _ = sample_lanes( - compiled, f_params, 9, postselected=postselected, flipped=flipped + pruned, _, stats = sample_lanes( + compiled, f_params, 9, share, postselected, flipped ) - lanes = component_program(compiled, postselected)[1] - ready, tested = lanes.ready, lanes.prune_outputs - prune_at = np.where(tested.any(axis=0), tested.argmax(axis=0), ready) - pruning = np.min(np.where(rows, prune_at[None, :], np.inf), axis=1) - late = rows & (ready[None, :] > pruning[:, None]) - assert not np.any(late), "an output set after the pruning action is 1" + unpruned, _, _ = sample_lanes(compiled, f_params, 9, share, None, flipped) + accepted = ~(unpruned & postselected).any(axis=1) + np.testing.assert_array_equal(stats["alive"], accepted) + np.testing.assert_array_equal(pruned[accepted], unpruned[accepted]) @pytest.mark.parametrize("name", ["t_ghz", "midcircuit_feedforward", "injection_check"]) diff --git a/test/unit/statevector/test_shots.py b/test/unit/statevector/test_shots.py index 5f0b0166..15cd46d3 100644 --- a/test/unit/statevector/test_shots.py +++ b/test/unit/statevector/test_shots.py @@ -248,3 +248,102 @@ def test_a_big_endian_host_is_rejected(monkeypatch): monkeypatch.setattr(sys, "byteorder", "big") with pytest.raises(NotImplementedError, match="little-endian"): Circuit(DIRECT_FLIPS).compile_detector_sampler(seed=0) + + +def _kept_rows(words, keep, size): + """The rows ``keep`` flags, and each run of ``size`` rows' count of them.""" + counts = np.add.reduceat(keep.astype(np.int64), np.arange(0, len(keep), size)) + return words[keep], counts + + +@pytest.mark.parametrize("width", [1, 3]) +@pytest.mark.parametrize("acceptance", [0.0, 0.1, 0.9, 1.0]) +def test_copy_kept_copies_the_kept_rows_in_order(width, acceptance): + rng = np.random.default_rng(int(10 * acceptance) + width) + size = 64 + words = rng.integers(0, 2**63, size=(1000, width), dtype=np.uint64) + keep = rng.random(len(words)) < acceptance + expected, kept = _kept_rows(words, keep, size) + starts = np.concatenate([[0], np.cumsum(kept)[:-1]]).astype(np.int64) + out = np.full((len(expected), width), 7, dtype=np.uint64) + for work in (0, 1 << 30): + shots_module._copy_kept(work, words, keep, size, kept, starts, out) + np.testing.assert_array_equal(out, expected) + + +def test_join_moves_each_runs_front_rows_together(): + rng = np.random.default_rng(3) + size, width = 50, 4 + words = rng.integers(0, 2**63, size=(170, width), dtype=np.uint64) + kept = np.array([0, 50, 7, 13], dtype=np.int64) + runs = [words[i * size : i * size + k].copy() for i, k in enumerate(kept)] + filled = shots_module._join(words, kept, size, width) + assert filled == kept.sum() + np.testing.assert_array_equal(words[:filled], np.vstack(runs)) + + +def test_compact_keeps_the_alive_rows_in_order(): + rng = np.random.default_rng(4) + rows = rng.integers(0, 2**63, size=(300, 3), dtype=np.uint64) + alive = rng.random(len(rows)) < 0.4 + expected = rows[alive][:, :2].copy() + kept = shots_module._compact(rows, alive, 2) + assert kept == alive.sum() + np.testing.assert_array_equal(rows[:kept, :2], expected) + + +REPETITION = """ +R 0 1 2 3 4 +X_ERROR(0.05) 0 1 2 3 4 +CNOT 0 1 1 2 2 3 3 4 +M 0 1 2 3 4 +DETECTOR rec[-1] rec[-2] +DETECTOR rec[-2] rec[-3] +DETECTOR rec[-3] rec[-4] +X 4 +M 4 +DETECTOR rec[-1] +OBSERVABLE_INCLUDE(0) rec[-2] +""" + + +@pytest.mark.parametrize("bit_packed", [False, True]) +def test_screened_calls_match_whatever_the_scratch_held(monkeypatch, bit_packed): + """Scratch reuse across shrinking and growing calls leaves the samples alone.""" + circuit = Circuit(REPETITION) + mask = np.array([True, True, False, False]) + calls = [5000, 1, 9000, 4096, 0, 12345, 3] + + def run(): + sampler = circuit.compile_detector_sampler(seed=5) + return [ + sampler.sample(n, postselection_mask=mask, bit_packed=bit_packed) + for n in calls + ] + + reused = run() + assert 0 < len(reused[0]) < calls[0] + monkeypatch.setattr(shots_module, "MAX_SCRATCH_BYTES", 0) + fresh = run() + for a, b in zip(reused, fresh, strict=True): + np.testing.assert_array_equal(a, b) + rows = ( + reused[0] + if not bit_packed + else np.unpackbits(reused[0], axis=1, bitorder="little") + ) + assert not rows[:, :2].any() + + +def test_a_sampler_pickles_after_a_screened_call(): + import pickle + + circuit = Circuit(REPETITION) + mask = np.array([True, True, True, False]) + sampler = circuit.compile_detector_sampler(seed=6) + sampler.sample(100, postselection_mask=mask) + assert shot_sampler(sampler)._scratch is not None + copy = pickle.loads(pickle.dumps(sampler)) # noqa: S301 (a trusted round trip) + expected = sampler.sample(500, postselection_mask=mask) + assert 0 < len(expected) < 500 + np.testing.assert_array_equal(copy.sample(500, postselection_mask=mask), expected)