Sample constant components without compiling them

CHANGELOG.md

@@ -16,6 +16,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   `stim-timeline-text` package. (#171)
 
 ### Changed
+- 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.

docs/index.md

@@ -115,7 +115,9 @@ survivors = samples[keep]
 ```
 
 `sample` always returns exactly `shots` rows. Shots discarded by a **direct**
-postselected detector skip the expensive JAX autoregressive loop; their direct
+postselected detector (one that is a single error bit or a constant, such as a
+deterministic detector of a noiseless circuit) 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

src/tsim/circuit.py

@@ -892,9 +892,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: Stabilizer rank decomposition strategy.

src/tsim/compile/pipeline.py

@@ -37,10 +37,11 @@ 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; every other
-    component has its outputs plugged (all prefixes in ``sequential`` mode, none
-    and all in ``joint`` mode), each plugged graph is decomposed into stabilizer
-    terms, and the terms are compiled into JAX arrays.
+    output equals one error-basis bit takes the direct path, and one whose single
+    output is a constant bit takes the constant path; every other component has
+    its outputs plugged (all prefixes in ``sequential`` mode, none and all in
+    ``joint`` mode), each plugged graph is decomposed into stabilizer terms, and
+    the terms are compiled into JAX arrays.
 
     Args:
         prepared: The prepared graph from ``prepare_graph``.
@@ -55,6 +56,7 @@ def compile_program(
     layout = prepared.layout
 
     direct_entries: list[tuple[int, int, bool]] = []
+    constant_entries: list[tuple[int, bool]] = []
     compiled_components: list[CompiledComponent] = []
     compiled_output_order: list[int] = []
 
@@ -66,6 +68,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, mode, strategy)
         )
@@ -75,8 +81,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)))
@@ -85,6 +95,7 @@ def compile_program(
         components=tuple(compiled_components),
         direct_f_indices=jnp.array([e[1] for e in direct_entries], dtype=jnp.int32),
         direct_flips=jnp.array([e[2] for e in direct_entries], dtype=jnp.bool_),
+        constant_bits=jnp.array([e[1] for e in constant_entries], dtype=jnp.bool_),
         output_order=jnp.asarray(output_order),
         output_reindex=None if is_identity else jnp.asarray(reindex),
         num_outputs=prepared.num_outputs,
@@ -99,6 +110,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
@@ -113,13 +154,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:

src/tsim/core/types.py

@@ -92,9 +92,12 @@ class CompiledProgram:
         components: The compiled components, sorted by number of outputs.
         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.
@@ -105,6 +108,7 @@ class CompiledProgram:
     components: tuple[CompiledComponent, ...]
     direct_f_indices: Array
     direct_flips: Array
+    constant_bits: Array
     output_order: Array
     output_reindex: Array | None
     num_outputs: int

src/tsim/sampler.py

@@ -132,9 +132,9 @@ def sample_program(
 
     """
     results: list[jax.Array] = []
+    batch_size = f_params.shape[0]
 
     if program.num_outputs == 0:
-        batch_size = f_params.shape[0]
         return jnp.zeros((batch_size, 0), dtype=jnp.bool_)
 
     if len(program.direct_f_indices) > 0:
@@ -144,6 +144,12 @@ def sample_program(
         )
         results.append(direct_bits)
 
+    n_constant = len(program.constant_bits)
+    if n_constant > 0:
+        results.append(
+            jnp.broadcast_to(program.constant_bits, (batch_size, n_constant))
+        )
+
     for component in program.components:
         samples, key, max_norm_deviation = sample_component(component, f_params, key)
         if np.isclose(max_norm_deviation, 1):
@@ -228,36 +234,43 @@ def __init__(
         self._direct_global_indices = np.asarray(
             prog.output_order[:n_direct], dtype=np.int32
         )
+        self._constant_bits = np.asarray(prog.constant_bits, dtype=np.bool_)
+        self._constant_global_indices = np.asarray(
+            prog.output_order[n_direct : n_direct + len(self._constant_bits)],
+            dtype=np.int32,
+        )
+        # Outputs known without JAX: direct and constant ones. Postselection
+        # treats both as "direct".
         self._direct_output_mask = np.zeros(prog.num_outputs, dtype=np.bool_)
-        if n_direct > 0:
-            self._direct_output_mask[self._direct_global_indices] = True
+        self._direct_output_mask[self._direct_global_indices] = True
+        self._direct_output_mask[self._constant_global_indices] = True
         self._direct_detector_mask = self._direct_output_mask[
             : self._num_detectors
         ].copy()
 
     def _compute_direct_outputs(self, f_params_np: np.ndarray) -> np.ndarray:
-        """Scatter direct output bits into a full (batch, num_outputs) bool array.
+        """Scatter direct and constant output bits into a (batch, num_outputs) array.
 
-        Non-direct columns are zero.  The zero-copy fast path applies when
-        direct indices are 0..n-1, there are no flips, and no reindex —
+        Compiled-component columns are zero.  The zero-copy fast path applies
+        when direct indices are 0..n-1, there are no flips, and no reindex —
         i.e. the common surface-code case.
         """
         batch = f_params_np.shape[0]
         num_outputs = self._program.num_outputs
         n_direct = len(self._direct_f_indices)
-        if n_direct == 0:
-            return np.zeros((batch, num_outputs), dtype=np.bool_)
         if self._direct_zero_copy and n_direct == num_outputs:
             return f_params_np[:, :n_direct].view(np.bool_).copy()
-        raw = (
-            f_params_np[:, :n_direct].view(np.bool_)
-            if self._direct_zero_copy
-            else (f_params_np[:, self._direct_f_indices] ^ self._direct_flips).view(
-                np.bool_
-            )
-        )
         out = np.zeros((batch, num_outputs), dtype=np.bool_)
-        out[:, self._direct_global_indices] = raw
+        if n_direct > 0:
+            raw = (
+                f_params_np[:, :n_direct].view(np.bool_)
+                if self._direct_zero_copy
+                else (f_params_np[:, self._direct_f_indices] ^ self._direct_flips).view(
+                    np.bool_
+                )
+            )
+            out[:, self._direct_global_indices] = raw
+        out[:, self._constant_global_indices] = self._constant_bits
         return out
 
     def _compute_reference_sample(self) -> np.ndarray:
@@ -546,8 +559,10 @@ def _flush(*, final: bool = False) -> None:
         return result, None, was_discarded
 
     def _sample_direct(self, shots: int) -> np.ndarray:
-        """Fast path when all components are direct (pure numpy, no JAX)."""
+        """Fast path when all components are direct or constant (pure numpy, no JAX)."""
         f_params = self._channel_sampler.sample(shots)
+        if len(self._constant_bits) > 0:
+            return self._compute_direct_outputs(f_params)
         if self._direct_zero_copy:
             return f_params[:, : len(self._direct_f_indices)].view(np.bool_)
         result = f_params[:, self._direct_f_indices] ^ self._direct_flips
@@ -558,6 +573,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 = []
@@ -607,7 +623,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"{np.sum(c_graphs)} graphs, "
             f"{error_channel_bits} error channel bits, "
             f"{np.max(num_outputs) if num_outputs else 0} outputs for largest cc, "
@@ -784,11 +800,12 @@ 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 skip the JAX
-                sampling loop. All ``shots`` rows are still returned: survivors contain
-                the full sample, while discarded rows retain direct detector columns
-                and fill component columns with ``False``. Re-apply the mask to the
-                returned detector columns to recover surviving shots.
+                When set, shots where any masked direct detector (a single error bit
+                or a constant) fires skip the JAX sampling loop. All ``shots`` rows
+                are still returned: survivors contain the full sample, while
+                discarded rows retain direct detector columns and fill component
+                columns with ``False``. Re-apply the mask to the returned detector
+                columns to recover surviving shots.
 
         Returns:
             A numpy array or tuple of numpy arrays containing the samples.
@@ -944,6 +961,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 len(component.compiled_scalar_graphs) == 2
 

test/integration/test_sampler.py

@@ -122,6 +122,7 @@ def test_sample_program_raises_on_component_norm_deviation(monkeypatch):
         components=components,
         direct_f_indices=jnp.array([], dtype=jnp.int32),
         direct_flips=jnp.array([], dtype=jnp.bool_),
+        constant_bits=jnp.array([], dtype=jnp.bool_),
         output_order=jnp.array([0, 1]),
         output_reindex=None,
         num_outputs=2,