From 2010f3afdde2bd8b5bfe94c943214705c6876889 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Thu, 15 Oct 2026 18:31:01 -0400 Subject: [PATCH 1/3] Sample constant components without compiling them classify_direct only accepted a single-output component whose classical Z spider carries exactly one error-basis parameter. Deterministic detectors of noiseless circuits reduce to the same two-vertex shape with no parameters, so each went through extract/decompose/compile and was sampled as its own JAX component (40 of them for the noiseless distillation circuit). Add classify_constant (sharing the structural check with classify_direct via _output_bit_phase) and a CompiledProgram.constant_bits field. Constants sit between the direct and the compiled outputs in output_order; the sampler broadcasts them in sample_program, scatters them in _compute_direct_outputs (so all-direct-or-constant programs still skip JAX), counts them as direct for postselection, and checks them in CompiledStateProbs.probability_of. Noiseless distillation.stim detector sampler: 41 -> 1 compiled components, 10k shots 2.7 s -> 1.3 s. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 1 + docs/index.md | 17 ++-- src/tsim/circuit.py | 6 +- src/tsim/compile/pipeline.py | 53 ++++++++-- src/tsim/core/types.py | 6 +- src/tsim/sampler.py | 91 +++++++++++++---- test/integration/test_sampler.py | 1 + test/unit/compile/test_pipeline.py | 153 ++++++++++++++++++++++++++++- test/unit/test_postselection.py | 92 +++++++++++++++++ 9 files changed, 378 insertions(+), 42 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8f2ac2b9..8a156179 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -43,6 +43,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 keeps samples independent of the thread count. (#199) - Noise sampling runs in numba kernels, so the samples for a given seed differ from earlier versions; `numba>=0.61` is now a dependency. (#199) +- Single-output components whose output is a constant (for example the deterministic detectors of a noiseless circuit) are no longer compiled: like single error bits, they are read off directly, without JAX. Postselection treats them as direct detectors, and the sampler `repr` reports their count. Samples for a fixed seed can differ from earlier versions. - Tsim now uses zxparam instead of pyzx-param. Circuits are built directly as quantum/classical ZX diagrams (no more composing with the adjoint), and reduction and stabilizer-rank decomposition run in zxparam. - `Circuit.get_graph()` and `Circuit.get_sampling_graph()` return `zxparam.Graph`. - In `Circuit.get_graph()` (and the `"zx"` diagram), every detector and observable has a classical output boundary, with ports after the qubit outputs, so the graph's tensor is the circuit's channel together with the detector and observable bits. diff --git a/docs/index.md b/docs/index.md index e58c94d4..c78138b0 100644 --- a/docs/index.md +++ b/docs/index.md @@ -115,11 +115,14 @@ survivors = samples[keep] ``` `sample` always returns exactly `shots` rows. Shots discarded by a **direct** -postselected detector skip the expensive JAX autoregressive loop; their direct -detector columns are still correct, and all other columns are filled with -`False`. Re-apply the mask to the detector columns (as above) to recover the -surviving shots. Detectors that live inside a JAX component cannot be evaluated -without running JAX, so those shots are always computed in full. +postselected detector (one that is a single error bit or a constant, such as a +deterministic detector of a noiseless circuit) skip the remaining component +sampling; their direct detector columns are still correct, and all other columns +are filled with `False`. With `strategy="active_statevector"`, a shot discarded +by a component detector also stops being simulated, and outputs set after that +detector are filled with `False`. Re-apply the mask to the detector columns (as +above) to recover the surviving shots. Stabilizer-rank components are always +computed in full for shots that pass direct postselection. This is independent of `prepend_observables`, `append_observables`, `separate_observables`, and `bit_packed`. When combined with @@ -127,8 +130,8 @@ This is independent of `prepend_observables`, `append_observables`, postselection discard check. On surviving rows it is applied to every detector column; on direct-discarded partial rows it is applied only to direct detector columns (component columns stay `False`). When combined with -`use_observable_reference_sample`, the reference XOR is applied to every row -that ran JAX; direct-discarded partial rows are left unchanged. +`use_observable_reference_sample`, returned observable values include the +reference XOR; outputs skipped by postselection remain `False`. ## Benchmarks diff --git a/src/tsim/circuit.py b/src/tsim/circuit.py index 98436cae..22629a25 100644 --- a/src/tsim/circuit.py +++ b/src/tsim/circuit.py @@ -904,9 +904,9 @@ def compile_detector_sampler( """Compile circuit into a detector sampler. Connected components whose single output is deterministically given by - one f-variable are handled via a fast direct path (no compilation or - autoregressive sampling). Remaining components go through the full - compilation pipeline. + one f-variable, or is a constant, are handled via a fast direct path (no + compilation or autoregressive sampling). Remaining components go + through the full compilation pipeline. Args: strategy: "active_statevector" (the default) to simulate each diff --git a/src/tsim/compile/pipeline.py b/src/tsim/compile/pipeline.py index bb99ee37..8924b90e 100644 --- a/src/tsim/compile/pipeline.py +++ b/src/tsim/compile/pipeline.py @@ -48,7 +48,8 @@ def compile_program( """Compile a prepared graph into an executable sampling program. The graph is split into connected components. A component whose single - output equals one error-basis bit takes the direct path. With + output equals one error-basis bit takes the direct path, and one whose output + is a constant bit takes the constant path. With ``strategy="active_statevector"`` every other component becomes an action tape for the statevector executor. Otherwise its outputs are plugged (all prefixes in ``sequential`` mode, none and all in ``joint`` mode), each plugged graph @@ -81,6 +82,7 @@ def compile_program( layout = prepared.layout direct_entries: list[tuple[int, int, bool]] = [] + constant_entries: list[tuple[int, bool]] = [] compiled_components: list[CompiledComponent | StatevectorComponent] = [] compiled_output_order: list[int] = [] @@ -92,6 +94,10 @@ def compile_program( basis_index, flip = direct direct_entries.append((output_indices[0], basis_index, flip)) continue + constant = classify_constant(component) + if constant is not None: + constant_entries.append((output_indices[0], constant)) + continue compiled_components.append(compile_component(component, output_indices, layout)) compiled_output_order.extend(output_indices) @@ -99,8 +105,12 @@ def compile_program( # prioritisation in transform_error_basis — the concatenated layout often # matches the original output order, sparing a reindex at sample time. direct_entries.sort() + constant_entries.sort() output_order = np.array( - [e[0] for e in direct_entries] + compiled_output_order, dtype=np.int32 + [e[0] for e in direct_entries] + + [e[0] for e in constant_entries] + + compiled_output_order, + dtype=np.int32, ) reindex = np.argsort(output_order) is_identity = np.array_equal(reindex, np.arange(len(output_order))) @@ -109,6 +119,7 @@ def compile_program( components=tuple(compiled_components), direct_f_indices=np.array([e[1] for e in direct_entries], dtype=np.int32), direct_flips=np.array([e[2] for e in direct_entries], dtype=np.bool_), + constant_bits=np.array([e[1] for e in constant_entries], dtype=np.bool_), output_order=output_order, output_reindex=None if is_identity else reindex, num_outputs=prepared.num_outputs, @@ -123,6 +134,36 @@ def classify_direct(component: Graph, layout: ParamLayout) -> tuple[int, bool] | a Hadamard edge to a classical Z spider whose phase is 0 or π plus one error-basis parameter. """ + phase = _output_bit_phase(component) + if phase is None or len(phase.parameters) != 1: + return None + (param,) = phase.parameters + kind, basis_index = layout.kind(param) + if kind != "f": + return None + return basis_index, phase.value == 1 + + +def classify_constant(component: Graph) -> bool | None: + """Return the output bit if the component's only output is a constant bit. + + That is the case when the component is a classical output boundary joined by + a Hadamard edge to a classical Z spider whose phase is 0 or π, with no + parameters (e.g. a deterministic detector of a noiseless circuit). + """ + phase = _output_bit_phase(component) + if phase is None or phase.parameters: + return None + return phase.value == 1 + + +def _output_bit_phase(component: Graph) -> Phase | None: + """Return the phase that sets the bit of a single-output, single-spider component. + + The component must be a classical output boundary joined by a Hadamard edge + to a classical Z spider whose phase is 0 or π plus any parameters; the output + bit is then the parity of the phase's value (in units of π) and parameters. + """ outputs = component.outputs() if len(outputs) != 1 or component.num_vertices() != 2: return None @@ -137,13 +178,9 @@ def classify_direct(component: Graph, layout: ParamLayout) -> tuple[int, bool] | if component.edge_type(out, spider) != EdgeType.HADAMARD: return None phase = component.phase(spider) - if len(phase.parameters) != 1 or phase.value not in (0, 1): + if phase.value not in (0, 1): return None - (param,) = phase.parameters - kind, basis_index = layout.kind(param) - if kind != "f": - return None - return basis_index, phase.value == 1 + return phase def _port(g: Graph, v: Vertex) -> int: diff --git a/src/tsim/core/types.py b/src/tsim/core/types.py index 95d81d99..dd115116 100644 --- a/src/tsim/core/types.py +++ b/src/tsim/core/types.py @@ -100,9 +100,12 @@ class CompiledProgram: stabilizer-rank ``CompiledComponent``s or ``StatevectorComponent``s. direct_f_indices: Precomputed f-parameter indices for direct components. direct_flips: Precomputed flip flags for direct components. + constant_bits: Output bits of constant components (single-output + components without parameters, e.g. deterministic detectors). output_order: Maps concatenated position to original output index. The first ``len(direct_f_indices)`` entries correspond to direct - components; the remainder to compiled components. + components, the next ``len(constant_bits)`` to constant components, + and the remainder to compiled components. output_reindex: Precomputed ``argsort(output_order)`` permutation, or ``None`` when the outputs are already in order. num_outputs: Total number of outputs across all components. @@ -113,6 +116,7 @@ class CompiledProgram: components: tuple[CompiledComponent | StatevectorComponent, ...] direct_f_indices: np.ndarray direct_flips: np.ndarray + constant_bits: np.ndarray output_order: np.ndarray output_reindex: np.ndarray | None num_outputs: int diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index 9cbfbae0..fa5748c3 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -181,6 +181,12 @@ def sample_program( ^ program.direct_flips ) results.append(direct_bits) + if len(program.constant_bits) > 0: + results.append( + jnp.broadcast_to( + program.constant_bits, (f_params.shape[0], len(program.constant_bits)) + ) + ) for component in program.components: assert isinstance(component, CompiledComponent) samples, key, max_norm_deviation = sample_component(component, f_params, key) @@ -198,6 +204,7 @@ def sample_statevector_program( seed: int, layouts: tuple[SharingLayout, ...], direct_flips: np.ndarray, + constant_bits: np.ndarray, *, share: bool, sharing_stats: deque[dict] | None = None, @@ -210,6 +217,7 @@ def sample_statevector_program( seed: Seed of the batch's random numbers, below 2^64. layouts: The sharing layout of each component. direct_flips: The flip of each direct output, with its reference bit. + constant_bits: The constant outputs, with their reference bits. share: Whether shots with equal histories share a lane. sharing_stats: If given, the executor appends its lane counts. @@ -240,6 +248,7 @@ def sample_statevector_program( f_params, program.direct_f_indices, direct_flips, + constant_bits, component_bits, program.output_order, out, @@ -251,14 +260,16 @@ def sample_statevector_program( def _write_direct_detectors(f, f_indices, flips, columns, postselect, out): """Clear every row of ``out`` and write its direct detector ``j`` at ``columns[j]``. - Detector ``j`` is ``f[row, f_indices[j]] ^ flips[j]``. Returns which rows - have a postselected detector of 1. + Detector ``j`` is ``f[row, f_indices[j]] ^ flips[j]``, or just ``flips[j]`` + when its index is -1. Returns which rows have a postselected detector of 1. """ discarded = np.zeros(f.shape[0], dtype=np.bool_) for row in numba.prange(f.shape[0]): out[row] = False for j in range(len(f_indices)): - bit = (f[row, f_indices[j]] != 0) ^ flips[j] + bit = flips[j] + if f_indices[j] >= 0: + bit ^= f[row, f_indices[j]] != 0 out[row, columns[j]] = bit if postselect[j] and bit: discarded[row] = True @@ -266,18 +277,22 @@ def _write_direct_detectors(f, f_indices, flips, columns, postselect, out): @kernel -def _gather_outputs(f, f_indices, flips, component_bits, order, out): +def _gather_outputs(f, f_indices, flips, constant_bits, component_bits, order, out): """Write every row's outputs into ``out`` in their original order. Output ``k`` of the concatenation goes to column ``order[k]``: first the - direct outputs ``f[row, f_indices[j]] ^ flips[j]``, then ``component_bits``. + direct outputs ``f[row, f_indices[j]] ^ flips[j]``, then the constant bits, + then ``component_bits``. """ num_direct = len(f_indices) + num_constant = len(constant_bits) for row in numba.prange(f.shape[0]): for j in range(num_direct): out[row, order[j]] = (f[row, f_indices[j]] != 0) ^ flips[j] + for j in range(num_constant): + out[row, order[num_direct + j]] = constant_bits[j] for j in range(component_bits.shape[1]): - out[row, order[num_direct + j]] = component_bits[row, j] + out[row, order[num_direct + num_constant + j]] = component_bits[row, j] def _nbytes(tree: object) -> int: @@ -423,7 +438,9 @@ def __init__( self._num_detectors = num_detectors self._share = share self._seeds = np.random.default_rng([seed, 1]) - self._layouts: dict[bytes, tuple[tuple[SharingLayout, ...], np.ndarray]] = {} + self._layouts: dict[ + bytes, tuple[tuple[SharingLayout, ...], np.ndarray, np.ndarray] + ] = {} self.sharing_stats: deque[dict] = deque(maxlen=10_000) def max_threads(self) -> int: @@ -457,21 +474,22 @@ def sample( pad_to: int | None = None, ) -> np.ndarray: """Sample the batch; shots that a ``mask`` detector discards are pruned.""" - layouts, direct_flips = self.layouts(mask, flip) + layouts, direct_flips, constant_bits = self.layouts(mask, flip) return sample_statevector_program( self._program, f_params, seed, layouts, direct_flips, + constant_bits, share=self._share, sharing_stats=self.sharing_stats, ) def layouts( self, mask: np.ndarray | None, flip: np.ndarray | None - ) -> tuple[tuple[SharingLayout, ...], np.ndarray]: - """Return the components' sharing layouts and the direct outputs' flips.""" + ) -> tuple[tuple[SharingLayout, ...], np.ndarray, np.ndarray]: + """Return sharing layouts and the referenced direct and constant bits.""" program = self._program postselected = np.zeros(program.num_outputs, dtype=np.bool_) if mask is not None: @@ -487,8 +505,16 @@ def layouts( layouts.append( sharing_layout(component, postselected[outputs], flip[outputs]) ) - direct = program.output_order[: len(program.direct_f_indices)] - self._layouts[key] = (tuple(layouts), program.direct_flips ^ flip[direct]) + num_direct = len(program.direct_f_indices) + direct = program.output_order[:num_direct] + constant = program.output_order[ + num_direct : num_direct + len(program.constant_bits) + ] + self._layouts[key] = ( + tuple(layouts), + program.direct_flips ^ flip[direct], + program.constant_bits ^ flip[constant], + ) return self._layouts[key] def concatenate(self, batches: list) -> np.ndarray: @@ -572,11 +598,21 @@ def __init__( assert np.array_equal( np.sort(program.output_order), np.arange(program.num_outputs) ) - columns = program.output_order[: len(program.direct_f_indices)] + num_direct = len(program.direct_f_indices) + self._constant_bits = program.constant_bits + self._constant_global_indices = program.output_order[ + num_direct : num_direct + len(program.constant_bits) + ] + columns = program.output_order[: num_direct + len(program.constant_bits)] is_detector = columns < self._num_detectors self._direct_detectors = ( - program.direct_f_indices[is_detector], - program.direct_flips[is_detector], + np.concatenate( + [ + program.direct_f_indices, + np.full(len(program.constant_bits), -1, dtype=np.int32), + ] + )[is_detector], + np.concatenate([program.direct_flips, program.constant_bits])[is_detector], columns[is_detector], ) self._direct_detector_mask = np.zeros(self._num_detectors, dtype=np.bool_) @@ -607,6 +643,7 @@ def _compute_reference_sample(self) -> np.ndarray: reference[program.output_order[: len(program.direct_f_indices)]] = ( program.direct_flips ) + reference[self._constant_global_indices] = self._constant_bits return reference num_f = self._channel_sampler.signature_matrix.shape[1] f_ref = np.zeros((1, num_f), dtype=np.uint8) @@ -881,13 +918,14 @@ def dispatch(task: tuple) -> None: return result def _sample_direct(self, shots: int) -> np.ndarray: - """Sample a program without components, whose outputs are all direct.""" + """Sample a program whose outputs are all direct or constant.""" program = self._program f_params = self._channel_sampler.sample(shots) num_direct = len(program.direct_f_indices) # NOTE: typical surface-code detector circuits need no copy. if ( program.output_reindex is None + and len(program.constant_bits) == 0 and not program.direct_flips.any() and np.array_equal(program.direct_f_indices, np.arange(num_direct)) ): @@ -898,6 +936,7 @@ def _sample_direct(self, shots: int) -> np.ndarray: f_params, program.direct_f_indices, program.direct_flips, + program.constant_bits, np.zeros((shots, 0), dtype=np.bool_), program.output_order, out, @@ -907,6 +946,7 @@ def _sample_direct(self, shots: int) -> np.ndarray: def __repr__(self) -> str: """Return a string representation with compilation statistics.""" n_direct = len(self._program.direct_f_indices) + n_constant = len(self._program.constant_bits) c_graphs = [] c_params = [] @@ -968,7 +1008,7 @@ def _format_bytes(n: int) -> str: ) return ( - f"{type(self).__name__}({n_direct} direct, " + f"{type(self).__name__}({n_direct} direct, {n_constant} constant, " f"{statevector}" f"{np.sum(c_graphs)} graphs, " f"{error_channel_bits} error channel bits, " @@ -1162,8 +1202,9 @@ def sample( only be used when observables are deterministic. Otherwise, it can unpredictably change the results. postselection_mask: Optional boolean array of length ``num_detectors``. - When set, shots where any masked direct detector fires are not - sampled further. All ``shots`` rows are still returned: survivors + When set, shots where any masked direct detector (an error bit + or a constant) fires are not sampled further. All ``shots`` rows + are still returned: survivors contain the full sample, while discarded rows retain direct detector columns and fill component columns with ``False``; with ``strategy="active_statevector"``, a shot discarded by a masked @@ -1201,7 +1242,13 @@ def sample( if not self._backend.prunes: skips = mask & self._direct_detector_mask else: - skips = mask + skips = mask.copy() + constant_detectors = self._constant_global_indices < num_detectors + skips[ + self._constant_global_indices[constant_detectors] + ] &= self._constant_bits[constant_detectors] & ( + not use_detector_reference_sample + ) postselection_mask = ( mask if skips.any() and self._program.components else None ) @@ -1318,6 +1365,10 @@ def probability_of(self, state: np.ndarray, *, batch_size: int) -> np.ndarray: targets = state[self._program.output_order[:n_direct]] p_joint = p_joint * (direct_bits == targets).all(axis=1) + if len(self._constant_bits) > 0: + targets = state[self._constant_global_indices] + p_joint = p_joint * np.array_equal(self._constant_bits, targets) + for component in self._program.components: assert isinstance(component, CompiledComponent) assert len(component.compiled_scalar_graphs) == 2 diff --git a/test/integration/test_sampler.py b/test/integration/test_sampler.py index c7e3e64d..82dca8ef 100644 --- a/test/integration/test_sampler.py +++ b/test/integration/test_sampler.py @@ -119,6 +119,7 @@ def test_sample_program_raises_on_component_norm_deviation(monkeypatch): components=components, direct_f_indices=np.array([], dtype=np.int32), direct_flips=np.array([], dtype=np.bool_), + constant_bits=np.array([], dtype=np.bool_), output_order=np.array([0, 1]), output_reindex=None, num_outputs=2, diff --git a/test/unit/compile/test_pipeline.py b/test/unit/compile/test_pipeline.py index aa4571d7..44734f5d 100644 --- a/test/unit/compile/test_pipeline.py +++ b/test/unit/compile/test_pipeline.py @@ -1,4 +1,4 @@ -"""compile_program on zxparam: plugging normalization and the direct path.""" +"""compile_program on zxparam: plugging normalization, the direct and constant paths.""" import itertools from fractions import Fraction @@ -9,9 +9,15 @@ import pytest from zxparam import Domain, EdgeType, Graph, Phase, VertexType, decompose +import tsim.sampler as sampler_module from tsim.circuit import Circuit from tsim.compile.evaluate import evaluate -from tsim.compile.pipeline import _plugged, classify_direct, compile_program +from tsim.compile.pipeline import ( + _plugged, + classify_constant, + classify_direct, + compile_program, +) from tsim.core.graph import prepare_graph from tsim.core.params import ParamLayout from tsim.core.types import CompiledComponent @@ -107,11 +113,123 @@ def test_single_error_detectors_take_the_direct_path(): assert not np.asarray(program.direct_flips).any() +# Detector 0 is random (H T H |0>); detectors 1-3 are deterministic with values +# 0, 1 and 1, the last one through a T · T_DAG pair that cancels. +NOISELESS_DETERMINISTIC_DETECTORS = """ +R 0 1 2 3 +H 0 +T 0 +H 0 +X 2 +H 3 +T 3 +T_DAG 3 +H 3 +X 3 +M 0 1 2 3 +DETECTOR rec[-4] +DETECTOR rec[-3] +DETECTOR rec[-2] +DETECTOR rec[-1] +""" + + +@pytest.mark.parametrize("strategy", ["cat5", "active_statevector"]) +def test_deterministic_detectors_of_a_noiseless_circuit_take_the_constant_path( + strategy, +): + circuit = Circuit(NOISELESS_DETERMINISTIC_DETECTORS) + program = compile_program( + prepare_graph(circuit, sample_detectors=True), + mode="sequential", + strategy=strategy, + ) + assert len(program.components) == 1 + assert program.components[0].output_indices == (0,) + assert len(program.direct_f_indices) == 0 + assert np.asarray(program.constant_bits).tolist() == [False, True, True] + assert np.asarray(program.output_order).tolist() == [1, 2, 3, 0] + assert program.output_reindex is not None + + samples = circuit.compile_detector_sampler(strategy=strategy, seed=0).sample( + 200, batch_size=50 + ) + assert samples.shape == (200, 4) + assert (samples[:, 1:] == [False, True, True]).all() + assert samples[:, 0].any() and not samples[:, 0].all() + + +@pytest.mark.parametrize("strategy", ["cat5", "active_statevector"]) +def test_all_constant_detectors_skip_jax(monkeypatch, strategy): + circuit = Circuit( + "R 0 1\nX 1\nM 0 1\nDETECTOR rec[-2]\nDETECTOR rec[-1]\n" + "OBSERVABLE_INCLUDE(0) rec[-1]" + ) + sampler = circuit.compile_detector_sampler(strategy=strategy, seed=0) + assert sampler._program.components == () + + def fail(*args, **kwargs): + raise AssertionError("sample_program must not be called") + + monkeypatch.setattr(sampler_module, "sample_program", fail) + monkeypatch.setattr(sampler_module, "sample_statevector_program", fail) + dets, obs = sampler.sample(10, separate_observables=True) + assert (dets == [False, True]).all() + assert obs.all() + dets, obs = sampler.sample( + 10, + separate_observables=True, + use_detector_reference_sample=True, + use_observable_reference_sample=True, + ) + assert not dets.any() and not obs.any() + + +@pytest.mark.parametrize("strategy", ["cat5", "active_statevector"]) +@pytest.mark.parametrize("with_component", [False, True]) +def test_direct_and_constant_detectors_are_reassembled_in_order( + strategy, with_component +): + # Detector 0 is constant 1, detector 1 a direct error bit and detector 2 + # constant 0, so the direct output comes first in the concatenated layout; + # detector 3 (if present) is a compiled component. + text = ( + "R 0 1 2 3\nX 0\nX_ERROR(0.5) 1\nM 0 1 2\n" + "DETECTOR rec[-3]\nDETECTOR rec[-2]\nDETECTOR rec[-1]" + ) + if with_component: + text += "\nH 3\nM 3\nDETECTOR rec[-1]" + sampler = Circuit(text).compile_detector_sampler(strategy=strategy, seed=0) + program = sampler._program + assert len(program.components) == int(with_component) + assert len(program.direct_f_indices) == 1 + assert np.asarray(program.constant_bits).tolist() == [True, False] + assert program.output_reindex is not None + + samples = sampler.sample(200, batch_size=50) + assert samples[:, 0].all() and not samples[:, 2].any() + for column in range(1, samples.shape[1], 2): + assert samples[:, column].any() and not samples[:, column].all() + + +def test_constant_outputs_condition_the_joint_probability(): + # Measurement 0 is random (H T H |0>); measurements 1 and 2 are constant. + circuit = Circuit("R 0 1 2\nH 0\nT 0\nH 0\nX 2\nM 0 1 2") + probs = CompiledStateProbs(circuit, seed=0) + assert len(probs._program.constant_bits) == 2 + p_one = np.sin(np.pi / 8) ** 2 + for m0, p in ((False, 1 - p_one), (True, p_one)): + got = probs.probability_of(np.array([m0, False, True]), batch_size=1)[0] + np.testing.assert_allclose(got, p, atol=1e-5) + for wrong in ([m0, True, True], [m0, False, False]): + assert probs.probability_of(np.array(wrong), batch_size=1)[0] == 0 + + LAYOUT = ParamLayout(num_errors=3) def _direct_candidate( - value: int = 0, + value: int | Fraction = 0, params: tuple[int, ...] = (LAYOUT.basis(1),), edge_type: EdgeType = EdgeType.HADAMARD, output_domain: Domain = Domain.CLASSICAL, @@ -156,6 +274,35 @@ def test_classify_direct_rejects_other_components(candidate): assert classify_direct(candidate, LAYOUT) is None +@pytest.mark.parametrize("value, bit", [(0, False), (1, True)]) +def test_classify_constant_reads_the_bit(value, bit): + assert classify_constant(_direct_candidate(value=value, params=())) is bit + + +@pytest.mark.parametrize( + "candidate", + [ + pytest.param(_direct_candidate(), id="basis-param"), + pytest.param(_direct_candidate(params=(LAYOUT.error(0),)), id="error-param"), + pytest.param( + _direct_candidate(params=(), edge_type=EdgeType.SIMPLE), id="simple-edge" + ), + pytest.param( + _direct_candidate(params=(), spider_type=VertexType.X), id="x-spider" + ), + pytest.param( + _direct_candidate(params=(), output_domain=Domain.QUANTUM), + id="quantum-output", + ), + pytest.param( + _direct_candidate(value=Fraction(1, 2), params=()), id="non-pauli-phase" + ), + ], +) +def test_classify_constant_rejects_other_components(candidate): + assert classify_constant(candidate) is None + + @pytest.mark.parametrize("output_domain", [Domain.CLASSICAL, Domain.QUANTUM]) def test_plugged_outputs_give_outcome_probabilities(output_domain): g = Graph() diff --git a/test/unit/test_postselection.py b/test/unit/test_postselection.py index 148f8e7e..fc43b18e 100644 --- a/test/unit/test_postselection.py +++ b/test/unit/test_postselection.py @@ -52,6 +52,16 @@ OBSERVABLE_INCLUDE(0) rec[-1] """ +# Detector 0 is constant 1 (no noise); detector 1 is a random compiled component. +CONSTANT_AND_COMPONENT_CIRCUIT = """ +R 0 1 +X 0 +H 1 +M 0 1 +DETECTOR rec[-2] +DETECTOR rec[-1] +""" + # Circuit with detectors and an observable for output-layout tests. DET_OBS_CIRCUIT = """ R 0 1 2 @@ -287,6 +297,88 @@ def test_postselection_jax_never_called_for_all_direct_discards(monkeypatch): assert calls == [] +@pytest.mark.parametrize("use_reference", [False, True]) +@pytest.mark.parametrize("strategy", ["cat5", "active_statevector"]) +def test_postselection_on_a_constant_detector(monkeypatch, use_reference, strategy): + sampler = Circuit(CONSTANT_AND_COMPONENT_CIRCUIT).compile_detector_sampler( + strategy=strategy, seed=0 + ) + assert sampler._direct_detector_mask.tolist() == [True, False] + mask = np.array([True, False]) + calls: list[int] = [] + + _record_rows(monkeypatch, calls) + samples = sampler.sample( + 16, + batch_size=4, + postselection_mask=mask, + use_detector_reference_sample=use_reference, + ) + if use_reference: + assert sum(calls) >= 16 + assert not samples[:, 0].any() + assert samples[:, 1].any() + else: + assert calls == [] + assert samples[:, 0].all() + assert not samples[:, 1].any() + + +@pytest.mark.parametrize("strategy", ["cat5", "active_statevector"]) +@pytest.mark.parametrize("constant,use_reference", [(0, False), (0, True), (1, True)]) +def test_nonrejecting_constant_postselection_uses_normal_batches( + monkeypatch, strategy, constant, use_reference +): + circuit = CONSTANT_AND_COMPONENT_CIRCUIT + if not constant: + circuit = circuit.replace("X 0", "") + sampler = Circuit(circuit).compile_detector_sampler(strategy=strategy, seed=0) + + def fail(*args, **kwargs): + raise AssertionError("A nonrejecting constant must not buffer survivors") + + monkeypatch.setattr(sampler, "_sample_batches_with_postselection", fail) + mask = np.array([True, False]) + samples = sampler.sample( + 19, + batch_size=5, + postselection_mask=mask, + use_detector_reference_sample=use_reference, + ) + assert not samples[:, 0].any() + assert mask.tolist() == [True, False] + + +@pytest.mark.parametrize("strategy", ["cat5", "active_statevector"]) +@pytest.mark.parametrize("use_reference", [False, True]) +def test_constant_outputs_with_component_postselection(strategy, use_reference): + circuit = Circuit(""" + X 0 + H 1 + T 1 + H 1 + X_ERROR(0.2) 2 + M 0 1 2 + DETECTOR rec[-1] + DETECTOR rec[-2] + DETECTOR rec[-3] + OBSERVABLE_INCLUDE(0) rec[-3] + """) + sampler = circuit.compile_detector_sampler(strategy=strategy, seed=0) + mask = np.array([True, True, False]) + detectors, observables = sampler.sample( + 128, + batch_size=17, + postselection_mask=mask, + separate_observables=True, + use_observable_reference_sample=use_reference, + ) + survivors = ~detectors[:, :2].any(axis=1) + assert survivors.any() + assert detectors[survivors, 2].all() + assert (observables[survivors, 0] == (not use_reference)).all() + + def test_postselection_jax_rows_less_than_shots(): """Total JAX rows < shots when some shots are direct-discarded.""" sampler = _make(MIXED_DIRECT_CIRCUIT, seed=0) From d04afb720ab983da0eed1eed205ba27c21bdd41c Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 20 Oct 2026 11:20:41 -0400 Subject: [PATCH 2/3] Skip constant components in the statevector plan tests test_ops_match_stabilizer_rank and test_plans_match_stabilizer_rank walk every connected component, skip the direct ones and look the rest up among the compiled components. Constant components are no longer compiled, so the classical_wiring circuit failed with KeyError: (0,). Skip them too, and keep quantum_components in step with what the pipeline compiles. Co-Authored-By: Claude Opus 5.5 --- test/helpers/statevector.py | 6 ++++-- test/unit/statevector/test_pbc.py | 7 +++++-- test/unit/statevector/test_program.py | 7 +++++-- 3 files changed, 14 insertions(+), 6 deletions(-) diff --git a/test/helpers/statevector.py b/test/helpers/statevector.py index 98574dc1..1d8aeb00 100644 --- a/test/helpers/statevector.py +++ b/test/helpers/statevector.py @@ -16,7 +16,7 @@ from tsim.circuit import Circuit from tsim.compile.evaluate import evaluate -from tsim.compile.pipeline import classify_direct +from tsim.compile.pipeline import classify_constant, classify_direct from tsim.core.graph import prepare_graph from tsim.core.params import ParamLayout from tsim.core.types import CompiledComponent @@ -173,12 +173,14 @@ def quantum_components( text: str, *, sample_detectors: bool ) -> tuple[list[Graph], ParamLayout]: - """Return the non-direct components of a circuit's prepared sampling graph.""" + """Return the components of a circuit's prepared sampling graph that are + compiled: neither direct nor constant.""" prepared = prepare_graph(Circuit(text), sample_detectors=sample_detectors) components = [ component for component in prepared.graph.connected_components() if classify_direct(component, prepared.layout) is None + and classify_constant(component) is None ] return components, prepared.layout diff --git a/test/unit/statevector/test_pbc.py b/test/unit/statevector/test_pbc.py index 00986d90..38bf868f 100644 --- a/test/unit/statevector/test_pbc.py +++ b/test/unit/statevector/test_pbc.py @@ -16,7 +16,7 @@ from zxparam import Domain, Phase, VertexType from tsim.circuit import Circuit -from tsim.compile.pipeline import classify_direct, compile_program +from tsim.compile.pipeline import classify_constant, classify_direct, compile_program from tsim.core.graph import prepare_graph from tsim.core.types import CompiledComponent from tsim.statevector.pbc import ( @@ -37,7 +37,10 @@ def test_ops_match_stabilizer_rank(name, sample_detectors): joint = compile_program(prepared, mode="joint", strategy="cat5").components by_outputs = {c.output_indices: c for c in joint} for component in prepared.graph.connected_components(): - if classify_direct(component, prepared.layout) is not None: + if ( + classify_direct(component, prepared.layout) is not None + or classify_constant(component) is not None + ): continue indices = output_indices(component) layout = ComponentLayout.of(component, indices, prepared.layout) diff --git a/test/unit/statevector/test_program.py b/test/unit/statevector/test_program.py index 77b51cce..40847071 100644 --- a/test/unit/statevector/test_program.py +++ b/test/unit/statevector/test_program.py @@ -13,7 +13,7 @@ import tsim.statevector.program as program from tsim.circuit import Circuit -from tsim.compile.pipeline import classify_direct, compile_program +from tsim.compile.pipeline import classify_constant, classify_direct, compile_program from tsim.core.graph import prepare_graph from tsim.core.types import CompiledComponent from tsim.statevector.pbc import ComponentLayout @@ -36,7 +36,10 @@ def test_plans_match_stabilizer_rank(name, sample_detectors): joint = compile_program(prepared, mode="joint", strategy="cat5").components by_outputs = {c.output_indices: c for c in joint} for component in prepared.graph.connected_components(): - if classify_direct(component, prepared.layout) is not None: + if ( + classify_direct(component, prepared.layout) is not None + or classify_constant(component) is not None + ): continue indices = output_indices(component) planned, component_layout = plan_component(component, indices, prepared.layout) From 6a8eb0454a46507e15c31565ef639dcf66811fee Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 20 Oct 2026 11:24:27 -0400 Subject: [PATCH 3/3] Reuse direct outputs for constant components --- CHANGELOG.md | 4 +- src/tsim/compile/pipeline.py | 61 +++--------- src/tsim/core/types.py | 9 +- src/tsim/sampler.py | 104 +++++++------------- test/helpers/statevector.py | 6 +- test/integration/test_sampler.py | 1 - test/integration/test_trajectory_sharing.py | 4 +- test/unit/compile/test_pipeline.py | 59 +++-------- test/unit/statevector/test_pbc.py | 7 +- test/unit/statevector/test_program.py | 7 +- test/unit/test_postselection.py | 6 +- 11 files changed, 79 insertions(+), 189 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8a156179..5d7b5622 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -43,7 +43,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 keeps samples independent of the thread count. (#199) - Noise sampling runs in numba kernels, so the samples for a given seed differ from earlier versions; `numba>=0.61` is now a dependency. (#199) -- Single-output components whose output is a constant (for example the deterministic detectors of a noiseless circuit) are no longer compiled: like single error bits, they are read off directly, without JAX. Postselection treats them as direct detectors, and the sampler `repr` reports their count. Samples for a fixed seed can differ from earlier versions. +- Constant outputs now use the existing direct-output path, skipping component + compilation and sampling. Postselection treats them as direct detectors. + Samples for a fixed seed can differ from earlier versions. (#201) - Tsim now uses zxparam instead of pyzx-param. Circuits are built directly as quantum/classical ZX diagrams (no more composing with the adjoint), and reduction and stabilizer-rank decomposition run in zxparam. - `Circuit.get_graph()` and `Circuit.get_sampling_graph()` return `zxparam.Graph`. - In `Circuit.get_graph()` (and the `"zx"` diagram), every detector and observable has a classical output boundary, with ports after the qubit outputs, so the graph's tensor is the circuit's channel together with the detector and observable bits. diff --git a/src/tsim/compile/pipeline.py b/src/tsim/compile/pipeline.py index 8924b90e..3754f617 100644 --- a/src/tsim/compile/pipeline.py +++ b/src/tsim/compile/pipeline.py @@ -48,8 +48,7 @@ def compile_program( """Compile a prepared graph into an executable sampling program. The graph is split into connected components. A component whose single - output equals one error-basis bit takes the direct path, and one whose output - is a constant bit takes the constant path. With + output equals one error-basis bit or a constant takes the direct path. With ``strategy="active_statevector"`` every other component becomes an action tape for the statevector executor. Otherwise its outputs are plugged (all prefixes in ``sequential`` mode, none and all in ``joint`` mode), each plugged graph @@ -82,7 +81,6 @@ def compile_program( layout = prepared.layout direct_entries: list[tuple[int, int, bool]] = [] - constant_entries: list[tuple[int, bool]] = [] compiled_components: list[CompiledComponent | StatevectorComponent] = [] compiled_output_order: list[int] = [] @@ -94,10 +92,6 @@ def compile_program( basis_index, flip = direct direct_entries.append((output_indices[0], basis_index, flip)) continue - constant = classify_constant(component) - if constant is not None: - constant_entries.append((output_indices[0], constant)) - continue compiled_components.append(compile_component(component, output_indices, layout)) compiled_output_order.extend(output_indices) @@ -105,12 +99,8 @@ def compile_program( # prioritisation in transform_error_basis — the concatenated layout often # matches the original output order, sparing a reindex at sample time. direct_entries.sort() - constant_entries.sort() output_order = np.array( - [e[0] for e in direct_entries] - + [e[0] for e in constant_entries] - + compiled_output_order, - dtype=np.int32, + [e[0] for e in direct_entries] + compiled_output_order, dtype=np.int32 ) reindex = np.argsort(output_order) is_identity = np.array_equal(reindex, np.arange(len(output_order))) @@ -119,7 +109,6 @@ def compile_program( components=tuple(compiled_components), direct_f_indices=np.array([e[1] for e in direct_entries], dtype=np.int32), direct_flips=np.array([e[2] for e in direct_entries], dtype=np.bool_), - constant_bits=np.array([e[1] for e in constant_entries], dtype=np.bool_), output_order=output_order, output_reindex=None if is_identity else reindex, num_outputs=prepared.num_outputs, @@ -128,41 +117,11 @@ def compile_program( def classify_direct(component: Graph, layout: ParamLayout) -> tuple[int, bool] | None: - """Return ``(basis index, flip)`` if the component's only output is one basis bit. - - That is the case when the component is a classical output boundary joined by - a Hadamard edge to a classical Z spider whose phase is 0 or π plus one - error-basis parameter. - """ - phase = _output_bit_phase(component) - if phase is None or len(phase.parameters) != 1: - return None - (param,) = phase.parameters - kind, basis_index = layout.kind(param) - if kind != "f": - return None - return basis_index, phase.value == 1 - - -def classify_constant(component: Graph) -> bool | None: - """Return the output bit if the component's only output is a constant bit. + """Return ``(basis index, flip)`` for a direct output; index -1 means constant. That is the case when the component is a classical output boundary joined by - a Hadamard edge to a classical Z spider whose phase is 0 or π, with no - parameters (e.g. a deterministic detector of a noiseless circuit). - """ - phase = _output_bit_phase(component) - if phase is None or phase.parameters: - return None - return phase.value == 1 - - -def _output_bit_phase(component: Graph) -> Phase | None: - """Return the phase that sets the bit of a single-output, single-spider component. - - The component must be a classical output boundary joined by a Hadamard edge - to a classical Z spider whose phase is 0 or π plus any parameters; the output - bit is then the parity of the phase's value (in units of π) and parameters. + a Hadamard edge to a classical Z spider whose phase is 0 or π plus at most + one error-basis parameter. """ outputs = component.outputs() if len(outputs) != 1 or component.num_vertices() != 2: @@ -178,9 +137,15 @@ def _output_bit_phase(component: Graph) -> Phase | None: if component.edge_type(out, spider) != EdgeType.HADAMARD: return None phase = component.phase(spider) - if phase.value not in (0, 1): + if len(phase.parameters) > 1 or phase.value not in (0, 1): return None - return phase + if not phase.parameters: + return -1, phase.value == 1 + (param,) = phase.parameters + kind, basis_index = layout.kind(param) + if kind != "f": + return None + return basis_index, phase.value == 1 def _port(g: Graph, v: Vertex) -> int: diff --git a/src/tsim/core/types.py b/src/tsim/core/types.py index dd115116..0d351047 100644 --- a/src/tsim/core/types.py +++ b/src/tsim/core/types.py @@ -98,14 +98,12 @@ class CompiledProgram: Attributes: components: The compiled components, sorted by number of outputs: stabilizer-rank ``CompiledComponent``s or ``StatevectorComponent``s. - direct_f_indices: Precomputed f-parameter indices for direct components. + direct_f_indices: Precomputed f-parameter indices for direct components; + -1 denotes a constant output whose value is its direct flip. direct_flips: Precomputed flip flags for direct components. - constant_bits: Output bits of constant components (single-output - components without parameters, e.g. deterministic detectors). output_order: Maps concatenated position to original output index. The first ``len(direct_f_indices)`` entries correspond to direct - components, the next ``len(constant_bits)`` to constant components, - and the remainder to compiled components. + components; the remainder to compiled components. output_reindex: Precomputed ``argsort(output_order)`` permutation, or ``None`` when the outputs are already in order. num_outputs: Total number of outputs across all components. @@ -116,7 +114,6 @@ class CompiledProgram: components: tuple[CompiledComponent | StatevectorComponent, ...] direct_f_indices: np.ndarray direct_flips: np.ndarray - constant_bits: np.ndarray output_order: np.ndarray output_reindex: np.ndarray | None num_outputs: int diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index fa5748c3..6103a555 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -157,6 +157,16 @@ def _check_norm(max_norm_deviation: float | jax.Array, threshold: float) -> None ) +def _direct_bits(program: CompiledProgram, f_params: jax.Array) -> jax.Array: + """Evaluate direct outputs, including constants with f-index -1.""" + indices = program.direct_f_indices + has_f = indices >= 0 + bits = jnp.broadcast_to(program.direct_flips, (f_params.shape[0], len(indices))) + return bits.at[:, has_f].set( + bits[:, has_f] ^ f_params[:, indices[has_f]].astype(jnp.bool_) + ) + + def sample_program( program: CompiledProgram, f_params: jax.Array, key: PRNGKey ) -> jax.Array: @@ -176,17 +186,7 @@ def sample_program( return jnp.zeros((f_params.shape[0], 0), dtype=jnp.bool_) results = [] if len(program.direct_f_indices) > 0: - direct_bits = ( - f_params[:, program.direct_f_indices].astype(jnp.bool_) - ^ program.direct_flips - ) - results.append(direct_bits) - if len(program.constant_bits) > 0: - results.append( - jnp.broadcast_to( - program.constant_bits, (f_params.shape[0], len(program.constant_bits)) - ) - ) + results.append(_direct_bits(program, f_params)) for component in program.components: assert isinstance(component, CompiledComponent) samples, key, max_norm_deviation = sample_component(component, f_params, key) @@ -204,7 +204,6 @@ def sample_statevector_program( seed: int, layouts: tuple[SharingLayout, ...], direct_flips: np.ndarray, - constant_bits: np.ndarray, *, share: bool, sharing_stats: deque[dict] | None = None, @@ -217,7 +216,6 @@ def sample_statevector_program( seed: Seed of the batch's random numbers, below 2^64. layouts: The sharing layout of each component. direct_flips: The flip of each direct output, with its reference bit. - constant_bits: The constant outputs, with their reference bits. share: Whether shots with equal histories share a lane. sharing_stats: If given, the executor appends its lane counts. @@ -248,7 +246,6 @@ def sample_statevector_program( f_params, program.direct_f_indices, direct_flips, - constant_bits, component_bits, program.output_order, out, @@ -277,22 +274,22 @@ def _write_direct_detectors(f, f_indices, flips, columns, postselect, out): @kernel -def _gather_outputs(f, f_indices, flips, constant_bits, component_bits, order, out): +def _gather_outputs(f, f_indices, flips, component_bits, order, out): """Write every row's outputs into ``out`` in their original order. Output ``k`` of the concatenation goes to column ``order[k]``: first the - direct outputs ``f[row, f_indices[j]] ^ flips[j]``, then the constant bits, - then ``component_bits``. + direct outputs (just ``flips[j]`` when ``f_indices[j]`` is -1), then + ``component_bits``. """ num_direct = len(f_indices) - num_constant = len(constant_bits) for row in numba.prange(f.shape[0]): for j in range(num_direct): - out[row, order[j]] = (f[row, f_indices[j]] != 0) ^ flips[j] - for j in range(num_constant): - out[row, order[num_direct + j]] = constant_bits[j] + bit = flips[j] + if f_indices[j] >= 0: + bit ^= f[row, f_indices[j]] != 0 + out[row, order[j]] = bit for j in range(component_bits.shape[1]): - out[row, order[num_direct + num_constant + j]] = component_bits[row, j] + out[row, order[num_direct + j]] = component_bits[row, j] def _nbytes(tree: object) -> int: @@ -438,9 +435,7 @@ def __init__( self._num_detectors = num_detectors self._share = share self._seeds = np.random.default_rng([seed, 1]) - self._layouts: dict[ - bytes, tuple[tuple[SharingLayout, ...], np.ndarray, np.ndarray] - ] = {} + self._layouts: dict[bytes, tuple[tuple[SharingLayout, ...], np.ndarray]] = {} self.sharing_stats: deque[dict] = deque(maxlen=10_000) def max_threads(self) -> int: @@ -474,22 +469,21 @@ def sample( pad_to: int | None = None, ) -> np.ndarray: """Sample the batch; shots that a ``mask`` detector discards are pruned.""" - layouts, direct_flips, constant_bits = self.layouts(mask, flip) + layouts, direct_flips = self.layouts(mask, flip) return sample_statevector_program( self._program, f_params, seed, layouts, direct_flips, - constant_bits, share=self._share, sharing_stats=self.sharing_stats, ) def layouts( self, mask: np.ndarray | None, flip: np.ndarray | None - ) -> tuple[tuple[SharingLayout, ...], np.ndarray, np.ndarray]: - """Return sharing layouts and the referenced direct and constant bits.""" + ) -> tuple[tuple[SharingLayout, ...], np.ndarray]: + """Return the components' sharing layouts and the direct outputs' flips.""" program = self._program postselected = np.zeros(program.num_outputs, dtype=np.bool_) if mask is not None: @@ -505,16 +499,8 @@ def layouts( layouts.append( sharing_layout(component, postselected[outputs], flip[outputs]) ) - num_direct = len(program.direct_f_indices) - direct = program.output_order[:num_direct] - constant = program.output_order[ - num_direct : num_direct + len(program.constant_bits) - ] - self._layouts[key] = ( - tuple(layouts), - program.direct_flips ^ flip[direct], - program.constant_bits ^ flip[constant], - ) + direct = program.output_order[: len(program.direct_f_indices)] + self._layouts[key] = (tuple(layouts), program.direct_flips ^ flip[direct]) return self._layouts[key] def concatenate(self, batches: list) -> np.ndarray: @@ -598,21 +584,11 @@ def __init__( assert np.array_equal( np.sort(program.output_order), np.arange(program.num_outputs) ) - num_direct = len(program.direct_f_indices) - self._constant_bits = program.constant_bits - self._constant_global_indices = program.output_order[ - num_direct : num_direct + len(program.constant_bits) - ] - columns = program.output_order[: num_direct + len(program.constant_bits)] + columns = program.output_order[: len(program.direct_f_indices)] is_detector = columns < self._num_detectors self._direct_detectors = ( - np.concatenate( - [ - program.direct_f_indices, - np.full(len(program.constant_bits), -1, dtype=np.int32), - ] - )[is_detector], - np.concatenate([program.direct_flips, program.constant_bits])[is_detector], + program.direct_f_indices[is_detector], + program.direct_flips[is_detector], columns[is_detector], ) self._direct_detector_mask = np.zeros(self._num_detectors, dtype=np.bool_) @@ -643,7 +619,6 @@ def _compute_reference_sample(self) -> np.ndarray: reference[program.output_order[: len(program.direct_f_indices)]] = ( program.direct_flips ) - reference[self._constant_global_indices] = self._constant_bits return reference num_f = self._channel_sampler.signature_matrix.shape[1] f_ref = np.zeros((1, num_f), dtype=np.uint8) @@ -918,14 +893,13 @@ def dispatch(task: tuple) -> None: return result def _sample_direct(self, shots: int) -> np.ndarray: - """Sample a program whose outputs are all direct or constant.""" + """Sample a program without components, whose outputs are all direct.""" program = self._program f_params = self._channel_sampler.sample(shots) num_direct = len(program.direct_f_indices) # NOTE: typical surface-code detector circuits need no copy. if ( program.output_reindex is None - and len(program.constant_bits) == 0 and not program.direct_flips.any() and np.array_equal(program.direct_f_indices, np.arange(num_direct)) ): @@ -936,7 +910,6 @@ def _sample_direct(self, shots: int) -> np.ndarray: f_params, program.direct_f_indices, program.direct_flips, - program.constant_bits, np.zeros((shots, 0), dtype=np.bool_), program.output_order, out, @@ -946,7 +919,6 @@ def _sample_direct(self, shots: int) -> np.ndarray: def __repr__(self) -> str: """Return a string representation with compilation statistics.""" n_direct = len(self._program.direct_f_indices) - n_constant = len(self._program.constant_bits) c_graphs = [] c_params = [] @@ -1008,7 +980,7 @@ def _format_bytes(n: int) -> str: ) return ( - f"{type(self).__name__}({n_direct} direct, {n_constant} constant, " + f"{type(self).__name__}({n_direct} direct, " f"{statevector}" f"{np.sum(c_graphs)} graphs, " f"{error_channel_bits} error channel bits, " @@ -1243,10 +1215,9 @@ def sample( skips = mask & self._direct_detector_mask else: skips = mask.copy() - constant_detectors = self._constant_global_indices < num_detectors - skips[ - self._constant_global_indices[constant_detectors] - ] &= self._constant_bits[constant_detectors] & ( + indices, flips, columns = self._direct_detectors + constant = indices < 0 + skips[columns[constant]] &= flips[constant] & ( not use_detector_reference_sample ) postselection_mask = ( @@ -1357,18 +1328,11 @@ def probability_of(self, state: np.ndarray, *, batch_size: int) -> np.ndarray: p_joint = jnp.ones(batch_size) if len(self._program.direct_f_indices) > 0: - direct_bits = ( - f_samples[:, self._program.direct_f_indices].astype(jnp.bool_) - ^ self._program.direct_flips - ) + direct_bits = _direct_bits(self._program, f_samples) n_direct = len(self._program.direct_f_indices) targets = state[self._program.output_order[:n_direct]] p_joint = p_joint * (direct_bits == targets).all(axis=1) - if len(self._constant_bits) > 0: - targets = state[self._constant_global_indices] - p_joint = p_joint * np.array_equal(self._constant_bits, targets) - for component in self._program.components: assert isinstance(component, CompiledComponent) assert len(component.compiled_scalar_graphs) == 2 diff --git a/test/helpers/statevector.py b/test/helpers/statevector.py index 1d8aeb00..98574dc1 100644 --- a/test/helpers/statevector.py +++ b/test/helpers/statevector.py @@ -16,7 +16,7 @@ from tsim.circuit import Circuit from tsim.compile.evaluate import evaluate -from tsim.compile.pipeline import classify_constant, classify_direct +from tsim.compile.pipeline import classify_direct from tsim.core.graph import prepare_graph from tsim.core.params import ParamLayout from tsim.core.types import CompiledComponent @@ -173,14 +173,12 @@ def quantum_components( text: str, *, sample_detectors: bool ) -> tuple[list[Graph], ParamLayout]: - """Return the components of a circuit's prepared sampling graph that are - compiled: neither direct nor constant.""" + """Return the non-direct components of a circuit's prepared sampling graph.""" prepared = prepare_graph(Circuit(text), sample_detectors=sample_detectors) components = [ component for component in prepared.graph.connected_components() if classify_direct(component, prepared.layout) is None - and classify_constant(component) is None ] return components, prepared.layout diff --git a/test/integration/test_sampler.py b/test/integration/test_sampler.py index 82dca8ef..c7e3e64d 100644 --- a/test/integration/test_sampler.py +++ b/test/integration/test_sampler.py @@ -119,7 +119,6 @@ def test_sample_program_raises_on_component_norm_deviation(monkeypatch): components=components, direct_f_indices=np.array([], dtype=np.int32), direct_flips=np.array([], dtype=np.bool_), - constant_bits=np.array([], dtype=np.bool_), output_order=np.array([0, 1]), output_reindex=None, num_outputs=2, diff --git a/test/integration/test_trajectory_sharing.py b/test/integration/test_trajectory_sharing.py index 99b45257..b0cb21b3 100644 --- a/test/integration/test_trajectory_sharing.py +++ b/test/integration/test_trajectory_sharing.py @@ -360,7 +360,9 @@ def test_direct_outputs_land_in_their_columns(name, rows): f_params = sampler._channel_sampler.sample(rows) samples = sampler._backend.sample(f_params, 2) columns = program.output_order[: len(program.direct_f_indices)] - expected = (f_params[:, program.direct_f_indices] != 0) ^ program.direct_flips + expected = np.broadcast_to(program.direct_flips, (rows, len(columns))).copy() + has_f = program.direct_f_indices >= 0 + expected[:, has_f] ^= f_params[:, program.direct_f_indices[has_f]] != 0 np.testing.assert_array_equal(samples[:, columns], expected) diff --git a/test/unit/compile/test_pipeline.py b/test/unit/compile/test_pipeline.py index 44734f5d..01eba319 100644 --- a/test/unit/compile/test_pipeline.py +++ b/test/unit/compile/test_pipeline.py @@ -1,4 +1,4 @@ -"""compile_program on zxparam: plugging normalization, the direct and constant paths.""" +"""compile_program on zxparam: plugging normalization and direct outputs.""" import itertools from fractions import Fraction @@ -14,7 +14,6 @@ from tsim.compile.evaluate import evaluate from tsim.compile.pipeline import ( _plugged, - classify_constant, classify_direct, compile_program, ) @@ -135,7 +134,7 @@ def test_single_error_detectors_take_the_direct_path(): @pytest.mark.parametrize("strategy", ["cat5", "active_statevector"]) -def test_deterministic_detectors_of_a_noiseless_circuit_take_the_constant_path( +def test_deterministic_detectors_of_a_noiseless_circuit_take_the_direct_path( strategy, ): circuit = Circuit(NOISELESS_DETERMINISTIC_DETECTORS) @@ -146,8 +145,8 @@ def test_deterministic_detectors_of_a_noiseless_circuit_take_the_constant_path( ) assert len(program.components) == 1 assert program.components[0].output_indices == (0,) - assert len(program.direct_f_indices) == 0 - assert np.asarray(program.constant_bits).tolist() == [False, True, True] + assert program.direct_f_indices.tolist() == [-1, -1, -1] + assert program.direct_flips.tolist() == [False, True, True] assert np.asarray(program.output_order).tolist() == [1, 2, 3, 0] assert program.output_reindex is not None @@ -190,9 +189,6 @@ def fail(*args, **kwargs): def test_direct_and_constant_detectors_are_reassembled_in_order( strategy, with_component ): - # Detector 0 is constant 1, detector 1 a direct error bit and detector 2 - # constant 0, so the direct output comes first in the concatenated layout; - # detector 3 (if present) is a compiled component. text = ( "R 0 1 2 3\nX 0\nX_ERROR(0.5) 1\nM 0 1 2\n" "DETECTOR rec[-3]\nDETECTOR rec[-2]\nDETECTOR rec[-1]" @@ -202,9 +198,9 @@ def test_direct_and_constant_detectors_are_reassembled_in_order( sampler = Circuit(text).compile_detector_sampler(strategy=strategy, seed=0) program = sampler._program assert len(program.components) == int(with_component) - assert len(program.direct_f_indices) == 1 - assert np.asarray(program.constant_bits).tolist() == [True, False] - assert program.output_reindex is not None + assert program.direct_f_indices.tolist() == [-1, 0, -1] + assert program.direct_flips.tolist() == [True, False, False] + assert program.output_reindex is None samples = sampler.sample(200, batch_size=50) assert samples[:, 0].all() and not samples[:, 2].any() @@ -216,7 +212,7 @@ def test_constant_outputs_condition_the_joint_probability(): # Measurement 0 is random (H T H |0>); measurements 1 and 2 are constant. circuit = Circuit("R 0 1 2\nH 0\nT 0\nH 0\nX 2\nM 0 1 2") probs = CompiledStateProbs(circuit, seed=0) - assert len(probs._program.constant_bits) == 2 + assert probs._program.direct_f_indices.tolist() == [-1, -1] p_one = np.sin(np.pi / 8) ** 2 for m0, p in ((False, 1 - p_one), (True, p_one)): got = probs.probability_of(np.array([m0, False, True]), batch_size=1)[0] @@ -245,9 +241,10 @@ def _direct_candidate( @pytest.mark.parametrize("value, flip", [(0, False), (1, True)]) -@pytest.mark.parametrize("basis_index", [1, 5]) +@pytest.mark.parametrize("basis_index", [-1, 1, 5]) def test_classify_direct_reads_the_basis_index_and_flip(value, flip, basis_index): - component = _direct_candidate(value=value, params=(LAYOUT.basis(basis_index),)) + params = () if basis_index == -1 else (LAYOUT.basis(basis_index),) + component = _direct_candidate(value=value, params=params) assert classify_direct(component, LAYOUT) == (basis_index, flip) @@ -260,8 +257,9 @@ def test_classify_direct_reads_the_basis_index_and_flip(value, flip, basis_index _direct_candidate(params=(LAYOUT.basis(0), LAYOUT.basis(1))), id="two-params", ), - pytest.param(_direct_candidate(params=()), id="no-param"), - pytest.param(_direct_candidate(value=1, params=()), id="no-param-flip"), + pytest.param( + _direct_candidate(value=Fraction(1, 2), params=()), id="non-pauli-phase" + ), pytest.param(_direct_candidate(edge_type=EdgeType.SIMPLE), id="simple-edge"), pytest.param(_direct_candidate(spider_type=VertexType.X), id="x-spider"), pytest.param( @@ -274,35 +272,6 @@ def test_classify_direct_rejects_other_components(candidate): assert classify_direct(candidate, LAYOUT) is None -@pytest.mark.parametrize("value, bit", [(0, False), (1, True)]) -def test_classify_constant_reads_the_bit(value, bit): - assert classify_constant(_direct_candidate(value=value, params=())) is bit - - -@pytest.mark.parametrize( - "candidate", - [ - pytest.param(_direct_candidate(), id="basis-param"), - pytest.param(_direct_candidate(params=(LAYOUT.error(0),)), id="error-param"), - pytest.param( - _direct_candidate(params=(), edge_type=EdgeType.SIMPLE), id="simple-edge" - ), - pytest.param( - _direct_candidate(params=(), spider_type=VertexType.X), id="x-spider" - ), - pytest.param( - _direct_candidate(params=(), output_domain=Domain.QUANTUM), - id="quantum-output", - ), - pytest.param( - _direct_candidate(value=Fraction(1, 2), params=()), id="non-pauli-phase" - ), - ], -) -def test_classify_constant_rejects_other_components(candidate): - assert classify_constant(candidate) is None - - @pytest.mark.parametrize("output_domain", [Domain.CLASSICAL, Domain.QUANTUM]) def test_plugged_outputs_give_outcome_probabilities(output_domain): g = Graph() diff --git a/test/unit/statevector/test_pbc.py b/test/unit/statevector/test_pbc.py index 38bf868f..00986d90 100644 --- a/test/unit/statevector/test_pbc.py +++ b/test/unit/statevector/test_pbc.py @@ -16,7 +16,7 @@ from zxparam import Domain, Phase, VertexType from tsim.circuit import Circuit -from tsim.compile.pipeline import classify_constant, classify_direct, compile_program +from tsim.compile.pipeline import classify_direct, compile_program from tsim.core.graph import prepare_graph from tsim.core.types import CompiledComponent from tsim.statevector.pbc import ( @@ -37,10 +37,7 @@ def test_ops_match_stabilizer_rank(name, sample_detectors): joint = compile_program(prepared, mode="joint", strategy="cat5").components by_outputs = {c.output_indices: c for c in joint} for component in prepared.graph.connected_components(): - if ( - classify_direct(component, prepared.layout) is not None - or classify_constant(component) is not None - ): + if classify_direct(component, prepared.layout) is not None: continue indices = output_indices(component) layout = ComponentLayout.of(component, indices, prepared.layout) diff --git a/test/unit/statevector/test_program.py b/test/unit/statevector/test_program.py index 40847071..77b51cce 100644 --- a/test/unit/statevector/test_program.py +++ b/test/unit/statevector/test_program.py @@ -13,7 +13,7 @@ import tsim.statevector.program as program from tsim.circuit import Circuit -from tsim.compile.pipeline import classify_constant, classify_direct, compile_program +from tsim.compile.pipeline import classify_direct, compile_program from tsim.core.graph import prepare_graph from tsim.core.types import CompiledComponent from tsim.statevector.pbc import ComponentLayout @@ -36,10 +36,7 @@ def test_plans_match_stabilizer_rank(name, sample_detectors): joint = compile_program(prepared, mode="joint", strategy="cat5").components by_outputs = {c.output_indices: c for c in joint} for component in prepared.graph.connected_components(): - if ( - classify_direct(component, prepared.layout) is not None - or classify_constant(component) is not None - ): + if classify_direct(component, prepared.layout) is not None: continue indices = output_indices(component) planned, component_layout = plan_component(component, indices, prepared.layout) diff --git a/test/unit/test_postselection.py b/test/unit/test_postselection.py index fc43b18e..1e8c9979 100644 --- a/test/unit/test_postselection.py +++ b/test/unit/test_postselection.py @@ -82,9 +82,9 @@ def _direct_outputs(sampler, f_params: np.ndarray) -> np.ndarray: program = sampler._program columns = program.output_order[: len(program.direct_f_indices)] out = np.zeros((len(f_params), program.num_outputs), dtype=np.bool_) - out[:, columns] = ( - f_params[:, program.direct_f_indices] != 0 - ) ^ program.direct_flips + out[:, columns] = program.direct_flips + has_f = program.direct_f_indices >= 0 + out[:, columns[has_f]] ^= f_params[:, program.direct_f_indices[has_f]] != 0 return out