diff --git a/CHANGELOG.md b/CHANGELOG.md index 8f2ac2b9..5d7b5622 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -43,6 +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) +- 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/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..3754f617 100644 --- a/src/tsim/compile/pipeline.py +++ b/src/tsim/compile/pipeline.py @@ -48,7 +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. 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 @@ -117,11 +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. + """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 π plus one - error-basis parameter. + 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: @@ -137,8 +137,10 @@ 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 len(phase.parameters) > 1 or phase.value not in (0, 1): return None + if not phase.parameters: + return -1, phase.value == 1 (param,) = phase.parameters kind, basis_index = layout.kind(param) if kind != "f": diff --git a/src/tsim/core/types.py b/src/tsim/core/types.py index 95d81d99..0d351047 100644 --- a/src/tsim/core/types.py +++ b/src/tsim/core/types.py @@ -98,7 +98,8 @@ 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. output_order: Maps concatenated position to original output index. The first ``len(direct_f_indices)`` entries correspond to direct diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index 9cbfbae0..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,11 +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) + 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) @@ -251,14 +257,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 @@ -270,12 +278,16 @@ 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 ``component_bits``. + direct outputs (just ``flips[j]`` when ``f_indices[j]`` is -1), then + ``component_bits``. """ num_direct = len(f_indices) 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] + 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 + j]] = component_bits[row, j] @@ -1162,8 +1174,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 +1214,12 @@ def sample( if not self._backend.prunes: skips = mask & self._direct_detector_mask else: - skips = mask + skips = mask.copy() + indices, flips, columns = self._direct_detectors + constant = indices < 0 + skips[columns[constant]] &= flips[constant] & ( + not use_detector_reference_sample + ) postselection_mask = ( mask if skips.any() and self._program.components else None ) @@ -1310,10 +1328,7 @@ 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) 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 aa4571d7..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 and the direct path.""" +"""compile_program on zxparam: plugging normalization and direct outputs.""" import itertools from fractions import Fraction @@ -9,9 +9,14 @@ 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_direct, + compile_program, +) from tsim.core.graph import prepare_graph from tsim.core.params import ParamLayout from tsim.core.types import CompiledComponent @@ -107,11 +112,120 @@ 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_direct_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 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 + + 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 +): + 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 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() + 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 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] + 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, @@ -127,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) @@ -142,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( diff --git a/test/unit/test_postselection.py b/test/unit/test_postselection.py index 148f8e7e..1e8c9979 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 @@ -72,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 @@ -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)