From 9776785b4b50b0d6c06b64a06d89856c180b0720 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Sat, 10 Oct 2026 16:28:08 -0400 Subject: [PATCH 01/18] Add Heisenberg-frame ops for statevector components Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/__init__.py | 8 + src/tsim/statevector/heisenberg.py | 276 +++++++++++++++++ test/helpers/statevector.py | 365 +++++++++++++++++++++++ test/unit/statevector/__init__.py | 0 test/unit/statevector/test_heisenberg.py | 114 +++++++ 5 files changed, 763 insertions(+) create mode 100644 src/tsim/statevector/__init__.py create mode 100644 src/tsim/statevector/heisenberg.py create mode 100644 test/helpers/statevector.py create mode 100644 test/unit/statevector/__init__.py create mode 100644 test/unit/statevector/test_heisenberg.py diff --git a/src/tsim/statevector/__init__.py b/src/tsim/statevector/__init__.py new file mode 100644 index 00000000..6cac2946 --- /dev/null +++ b/src/tsim/statevector/__init__.py @@ -0,0 +1,8 @@ +"""Statevector sampling of Tsim's quantum components (``strategy="statevector"``). + +A non-direct component of the reduced sampling graph is extracted as a circuit +(``zx_param.extract``), rewritten as rotations and measurements in the frame of +the initial ``|0…0⟩`` state (``heisenberg``), reordered to keep the dense array +small (``schedule``), planned on a fixed number of array axes (``plan``), +lowered to JAX arrays (``program``) and interpreted lane by lane (``execute``). +""" diff --git a/src/tsim/statevector/heisenberg.py b/src/tsim/statevector/heisenberg.py new file mode 100644 index 00000000..d25ed60d --- /dev/null +++ b/src/tsim/statevector/heisenberg.py @@ -0,0 +1,276 @@ +"""Heisenberg-frame operations of an extracted circuit. + +``heisenberg_ops`` rewrites an ``ExtractedCircuit`` as rotations and +measurements whose Paulis are written in the frame of the initial ``|0…0⟩`` +state. The parameter part of a Z phase is noise: a Z applied when a parity of +f-bits is 1. Moving it past every later operation to the end of the circuit, +where it changes no outcome, flips the sign of each operation it anticommutes +with. So every operation carries a fixed GF(2) sign row, and operations that +commute can be reordered without touching their signs. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from fractions import Fraction +from math import pi + +import numpy as np +import stim +from zx_param import Graph, Vertex +from zx_param.extract import ExtractedCircuit, Gate + +from tsim.core.params import ParamLayout + +ROTATION = 0 +MEASUREMENT = 1 + +_CLIFFORD_QUARTERS = {1: "s", 2: "z", 3: "s_dag"} + + +@dataclass(frozen=True) +class ComponentLayout: + """Where a component's parameters and outputs go. + + Attributes: + f_selection: Global f-bit index of each component f-bit, ascending. + f_position: Position in ``f_selection`` of each phase parameter ID. + output_indices: Global output index of each output slot, in port order + (the order of ``Graph.outputs()``, as for stabilizer-rank + components). + + """ + + f_selection: tuple[int, ...] + f_position: Mapping[int, int] + output_indices: tuple[int, ...] + + @classmethod + def of(cls, component: Graph, layout: ParamLayout) -> ComponentLayout: + """Return the layout of ``component``'s f-bits and output ports.""" + f_selection = tuple( + sorted( + { + layout.basis_index(param) + for v in component.vertices() + for param in component.phase(v).parameters + } + ) + ) + f_position = {layout.basis(i): pos for pos, i in enumerate(f_selection)} + outputs = tuple(_port(component, v) for v in component.outputs()) + return cls(f_selection, f_position, outputs) + + +def _port(g: Graph, v: Vertex) -> int: + port = g.port(v) + assert port is not None + return port + + +@dataclass(frozen=True) +class HeisenbergOps: + """Rotations and measurements in the frame of the initial ``|0…0⟩`` state. + + Op ``j`` is the rotation ``exp(-iθ (-1)^s P)`` or the measurement of + ``(-1)^s P``, where ``P`` has the bit rows ``x[j]`` and ``z[j]`` (``Y`` + where both are set) and ``s`` is the parity of ``sign[j]`` against + ``[1, f_0, f_1, …]``. A measurement records its eigenvalue bit. + + Attributes: + num_lines: Number of qubit lines. + num_f: Number of component f-bits the signs depend on. + kind: ``ROTATION`` or ``MEASUREMENT``, shape ``(N,)``. + x: X part of each Pauli, shape ``(N, num_lines)``. + z: Z part of each Pauli, shape ``(N, num_lines)``. + angle: θ of each rotation, 0 for measurements, shape ``(N,)``. + sign: Constant, then one column per f-bit, shape ``(N, 1 + num_f)``. + slot: Output slot of a recorded measurement, -1 for a forgotten + measurement or a rotation, shape ``(N,)``. + + """ + + num_lines: int + num_f: int + kind: np.ndarray + x: np.ndarray + z: np.ndarray + angle: np.ndarray + sign: np.ndarray + slot: np.ndarray + + @property + def num_ops(self) -> int: + """Return the number of operations.""" + return len(self.kind) + + @property + def num_outputs(self) -> int: + """Return the number of recorded measurements.""" + return int(np.count_nonzero(self.slot >= 0)) + + def take(self, order: np.ndarray) -> HeisenbergOps: + """Return the operations in the given order. + + The result describes the same distribution when ``order`` keeps every + pair of anticommuting operations in their original order. + """ + return HeisenbergOps( + num_lines=self.num_lines, + num_f=self.num_f, + kind=self.kind[order], + x=self.x[order], + z=self.z[order], + angle=self.angle[order], + sign=self.sign[order], + slot=self.slot[order], + ) + + +def heisenberg_ops(circuit: ExtractedCircuit, layout: ComponentLayout) -> HeisenbergOps: + """Return the operations of ``circuit`` in the ``|0…0⟩`` frame, in time order. + + Args: + circuit: A circuit extracted from a closed sampling graph: every line + starts with a ``|+⟩`` preparation and every output is classical. + layout: The component's f-bits and output slots. + + Raises: + ValueError: If a line starts at an input port, ends at a quantum + output, or a gate kind is unknown. + + """ + _check_lines(circuit) + builder = _Builder(circuit.n_lines, len(layout.f_selection)) + for gate in circuit.gates: + builder.gate(gate, layout.f_position) + slot_of_port = {port: slot for slot, port in enumerate(layout.output_indices)} + for line in range(circuit.n_lines): + port = circuit.line_end[line] + if port is not None: + builder.measure(line, slot_of_port[port]) + return builder.finish() + + +def _check_lines(circuit: ExtractedCircuit) -> None: + for line in range(circuit.n_lines): + start = circuit.line_start[line] + if start is not None: + raise ValueError( + f"line {line} starts at input port {start}; " + "sampling graphs have no inputs" + ) + port = circuit.line_end[line] + if port is not None and port not in circuit.classical_ports: + raise ValueError(f"line {line} ends at the quantum output port {port}") + + +class _Builder: + """Walks the gates with a stim simulator that holds the Clifford prefix U.""" + + def __init__(self, num_lines: int, num_f: int) -> None: + """Start every line in ``|+⟩`` (``H`` on ``|0⟩``).""" + self.num_lines = num_lines + self.num_f = num_f + self.sim = stim.TableauSimulator() + self.sim.set_num_qubits(num_lines) + for line in range(num_lines): + self.sim.h(line) + self.inverse: stim.Tableau | None = None + self.time = 0 + self.ops: list[tuple[int, np.ndarray, np.ndarray, bool, float, int, int]] = [] + self.noise: list[tuple[np.ndarray, np.ndarray, frozenset[int], int]] = [] + + def _mapped(self, line: int) -> tuple[np.ndarray, np.ndarray, bool]: + """Return U†·Z_line·U as bit rows and a sign flag.""" + if self.inverse is None: + self.inverse = self.sim.current_inverse_tableau() + image = self.inverse.z_output(line) + x, z = image.to_numpy() + return x, z, bool(image.sign == -1) + + def _clifford(self, name: str, *lines: int) -> None: + getattr(self.sim, name)(*lines) + self.inverse = None + + def gate(self, gate: Gate, f_position: Mapping[int, int]) -> None: + """Absorb a Clifford gate, or record a rotation, noise or measurement.""" + if gate.kind == "H": + self._clifford("h", *gate.lines) + elif gate.kind == "CNOT": + self._clifford("cx", *gate.lines) + elif gate.kind == "CZ": + self._clifford("cz", *gate.lines) + elif gate.kind == "Z": + self._phase(gate, f_position) + elif gate.kind in ("MEASURE", "DISCARD"): + self.measure(gate.lines[0], -1) + else: + raise ValueError(f"unknown gate kind {gate.kind!r}") + + def _phase(self, gate: Gate, f_position: Mapping[int, int]) -> None: + (line,) = gate.lines + phase = gate.phase + assert phase is not None + quarters = round(2 * phase.value) + remainder = phase.value - Fraction(quarters, 2) + clifford = _CLIFFORD_QUARTERS.get(quarters % 4) + if clifford is not None: + self._clifford(clifford, line) + if remainder: + x, z, negative = self._mapped(line) + self._op(ROTATION, x, z, negative, pi * float(remainder) / 2, -1) + if phase.parameters: + x, z, _ = self._mapped(line) + positions = frozenset(f_position[p] for p in phase.parameters) + self.noise.append((x, z, positions, self.time)) + self.time += 1 + + def measure(self, line: int, slot: int) -> None: + """Record a Z measurement of ``line`` (``slot`` -1 if forgotten).""" + x, z, negative = self._mapped(line) + self._op(MEASUREMENT, x, z, negative, 0.0, slot) + + def _op( + self, + kind: int, + x: np.ndarray, + z: np.ndarray, + negative: bool, + angle: float, + slot: int, + ) -> None: + self.ops.append((kind, x, z, negative, angle, slot, self.time)) + self.time += 1 + + def finish(self) -> HeisenbergOps: + """Return the ops, with every noise event folded into their sign rows.""" + num_ops, num_lines = len(self.ops), self.num_lines + x = np.array([op[1] for op in self.ops], dtype=bool).reshape(num_ops, num_lines) + z = np.array([op[2] for op in self.ops], dtype=bool).reshape(num_ops, num_lines) + sign = np.zeros((num_ops, 1 + self.num_f), dtype=bool) + sign[:, 0] = [op[3] for op in self.ops] + if self.noise and num_ops: + noise_x = np.array([e[0] for e in self.noise], dtype=np.float32) + noise_z = np.array([e[1] for e in self.noise], dtype=np.float32) + anti = ( + noise_x @ z.T.astype(np.float32) + noise_z @ x.T.astype(np.float32) + ) % 2 == 1 + noise_time = np.array([e[3] for e in self.noise]) + op_time = np.array([op[6] for op in self.ops]) + flips = anti & (noise_time[:, None] < op_time[None, :]) + condition = np.zeros((len(self.noise), self.num_f), dtype=np.float32) + for e, (_, _, positions, _) in enumerate(self.noise): + condition[e, sorted(positions)] = 1 + sign[:, 1:] = (flips.T.astype(np.float32) @ condition) % 2 == 1 + return HeisenbergOps( + num_lines=num_lines, + num_f=self.num_f, + kind=np.array([op[0] for op in self.ops], dtype=np.uint8), + x=x, + z=z, + angle=np.array([op[4] for op in self.ops], dtype=np.float64), + sign=sign, + slot=np.array([op[5] for op in self.ops], dtype=np.int32), + ) diff --git a/test/helpers/statevector.py b/test/helpers/statevector.py new file mode 100644 index 00000000..d821b315 --- /dev/null +++ b/test/helpers/statevector.py @@ -0,0 +1,365 @@ +"""Dense reference simulations for testing ``strategy="statevector"``. + +Bit ``q`` of a basis-state index is qubit ``q`` (a line, or an array axis). +""" + +from __future__ import annotations + +import itertools +from collections.abc import Iterator, Mapping, Sequence +from fractions import Fraction +from math import cos, pi, sin + +import numpy as np +from zx_param import Graph, Phase, Scalar +from zx_param.extract import ExtractedCircuit, Gate, extract_circuit + +from tsim.circuit import Circuit +from tsim.compile.pipeline import classify_direct +from tsim.core.graph import prepare_graph +from tsim.core.params import ParamLayout +from tsim.statevector.heisenberg import ( + ROTATION, + ComponentLayout, + HeisenbergOps, + heisenberg_ops, +) + +Distribution = dict[tuple[int, ...], float] + +CIRCUITS: dict[str, str] = { + "t_interference": """ + R 0 + H 0 + T 0 + H 0 + T 0 + H 0 + X_ERROR(0.1) 0 + T 0 + H 0 + M 0 + DETECTOR rec[-1] + """, + "t_ghz": """ + R 0 1 2 + H 0 + CNOT 0 1 + CNOT 1 2 + T 0 1 2 + H 0 1 2 + DEPOLARIZE1(0.05) 0 1 2 + T 0 1 + CZ 0 1 + H 0 1 + M 0 1 2 + DETECTOR rec[-3] rec[-2] + DETECTOR rec[-2] rec[-1] + OBSERVABLE_INCLUDE(0) rec[-1] + """, + "rotations": """ + R 0 1 + U3(0.3, 0.24, 0.49) 0 + CNOT 0 1 + R_Z(0.13) 1 + R_X(0.21) 0 + Z_ERROR(0.2) 0 + R_Y(0.4) 1 + CZ 0 1 + H 1 + M 0 1 + DETECTOR rec[-1] + DETECTOR rec[-2] + """, + "injection_check": """ + R 0 1 2 3 + H 0 + T 0 + CNOT 0 1 + CNOT 0 2 + DEPOLARIZE1(0.02) 0 1 2 + CNOT 0 3 + CNOT 1 3 + M 3 + DETECTOR rec[-1] + R 3 + CNOT 1 3 + CNOT 2 3 + X_ERROR(0.05) 3 + M 3 + DETECTOR rec[-1] + CNOT 0 2 + CNOT 0 1 + T_DAG 0 + H 0 + M 0 1 2 + DETECTOR rec[-2] + DETECTOR rec[-1] + OBSERVABLE_INCLUDE(0) rec[-3] + """, + "midcircuit_feedforward": """ + R 0 1 2 + H 0 1 + T 0 1 + CNOT 0 2 + CNOT 1 2 + X_ERROR(0.1) 2 + M 2 + CX rec[-1] 0 + T 0 + CNOT 0 1 + H 0 1 + Z_ERROR(0.1) 1 + T 1 + H 1 + MR 0 + H 0 + T 0 + H 0 + M 0 1 + DETECTOR rec[-1] rec[-4] + DETECTOR rec[-2] rec[-3] + OBSERVABLE_INCLUDE(0) rec[-2] + """, + "two_components": """ + R 0 1 2 3 4 + X_ERROR(0.2) 4 + H 0 2 + T 0 2 + H 0 2 + T 0 2 + CNOT 0 1 + CNOT 2 3 + X_ERROR(0.1) 1 3 + T 1 3 + H 0 1 2 3 + M 4 0 1 2 3 + DETECTOR rec[-5] + DETECTOR rec[-4] rec[-3] + DETECTOR rec[-2] rec[-1] + OBSERVABLE_INCLUDE(0) rec[-4] + """, + "clifford_component": """ + X_ERROR(0.5) 0 + R 1 + H 1 + M 0 1 + DETECTOR rec[-2] + DETECTOR rec[-1] rec[-2] + """, +} + + +def quantum_components( + text: str, *, sample_detectors: bool +) -> tuple[list[Graph], ParamLayout]: + """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 + ] + return components, prepared.layout + + +def hand_built_circuits() -> list[tuple[str, ExtractedCircuit, ComponentLayout]]: + """Return small extracted circuits with a DISCARD and a reused MEASURE line.""" + t = Phase(Fraction(1, 4)) + noise = Phase(Fraction(0), frozenset({7})) + discard = ExtractedCircuit( + gates=[ + Gate("Z", (0,), t), + Gate("CZ", (0, 1)), + Gate("Z", (1,), Phase(Fraction(1, 4), frozenset({7}))), + Gate("H", (1,)), + Gate("DISCARD", (1,)), + Gate("H", (0,)), + Gate("Z", (0,), Phase(Fraction(3, 4))), + ], + n_lines=2, + line_start={0: None, 1: None}, + line_end={0: 0, 1: None}, + scalar=Scalar.one(), + classical_ports=frozenset({0}), + ) + measure = ExtractedCircuit( + gates=[ + Gate("Z", (0,), t), + Gate("CNOT", (0, 1)), + Gate("Z", (1,), noise), + Gate("H", (1,)), + Gate("Z", (1,), Phase(Fraction(1, 8))), + Gate("MEASURE", (1,)), + Gate("CNOT", (1, 0)), + Gate("H", (0,)), + Gate("Z", (0,), Phase(Fraction(7, 4))), + Gate("H", (0,)), + ], + n_lines=2, + line_start={0: None, 1: None}, + line_end={0: 0, 1: 1}, + scalar=Scalar.one(), + classical_ports=frozenset({0, 1}), + ) + return [ + ("discard", discard, ComponentLayout((0,), {7: 0}, (0,))), + ("measure", measure, ComponentLayout((0,), {7: 0}, (0, 1))), + ] + + +def component_ops() -> Iterator[HeisenbergOps]: + """Yield the ops of every test component (both modes) and hand-built circuit.""" + for name in sorted(CIRCUITS): + for sample_detectors in (True, False): + components, layout = quantum_components( + CIRCUITS[name], sample_detectors=sample_detectors + ) + for component in components: + component_layout = ComponentLayout.of(component, layout) + yield heisenberg_ops(extract_circuit(component), component_layout) + for _, circuit, component_layout in hand_built_circuits(): + yield heisenberg_ops(circuit, component_layout) + + +def f_patterns(num_f: int, limit: int = 32) -> list[tuple[int, ...]]: + """Return every f-bit pattern, or all-zero plus ``limit`` random ones.""" + if 2**num_f <= limit: + return [tuple(p) for p in itertools.product((0, 1), repeat=num_f)] + rng = np.random.default_rng(0) + random = [tuple(int(b) for b in rng.integers(0, 2, num_f)) for _ in range(limit)] + return [(0,) * num_f, *random] + + +def assert_distributions_close( + got: Distribution, expected: Distribution, atol: float = 1e-6 +) -> None: + """Assert that two outcome distributions agree entry by entry.""" + for key in set(got) | set(expected): + a, b = got.get(key, 0.0), expected.get(key, 0.0) + assert abs(a - b) <= atol, (key, a, b) + + +def parity(values: np.ndarray | int) -> np.ndarray: + """Return the parity of the set bits of each non-negative integer.""" + remaining = np.asarray(values, dtype=np.uint64).copy() + result = np.zeros(remaining.shape, dtype=np.uint64) + while remaining.any(): + result ^= remaining & np.uint64(1) + remaining >>= np.uint64(1) + return result.astype(bool) + + +def mask_of(bits: Sequence[bool] | np.ndarray) -> int: + """Return the integer whose bit ``q`` is ``bits[q]``.""" + return sum(1 << q for q, bit in enumerate(bits) if bit) + + +def apply_pauli(psi: np.ndarray, x: int, z: int) -> np.ndarray: + """Return P·ψ for P = i^|x∧z| X^x Z^z.""" + index = np.arange(len(psi)) + flipped = index ^ x + i_power = (1, 1j, -1, -1j)[bin(x & z).count("1") % 4] + signs = np.where(parity(flipped & z), -1, 1) + return i_power * signs * psi[flipped] + + +_H = np.array([[1, 1], [1, -1]]) / np.sqrt(2) + + +def _single_qubit(psi: np.ndarray, num_qubits: int, qubit: int, matrix: np.ndarray): + tensor = psi.reshape([2] * num_qubits) + axis = num_qubits - 1 - qubit + tensor = np.moveaxis(np.tensordot(matrix, tensor, axes=([1], [axis])), 0, axis) + return tensor.reshape(-1) + + +def _extracted_gate( + psi: np.ndarray, num_lines: int, gate: Gate, bindings: Mapping[int, int] +) -> list[np.ndarray]: + index = np.arange(len(psi)) + if gate.kind == "H": + return [_single_qubit(psi, num_lines, gate.lines[0], _H)] + if gate.kind == "CNOT": + control, target = gate.lines + return [psi[index ^ (((index >> control) & 1) << target)]] + if gate.kind == "CZ": + a, b = gate.lines + both = ((index >> a) & (index >> b) & 1) == 1 + return [psi * np.where(both, -1, 1)] + (line,) = gate.lines + bit = (index >> line) & 1 + if gate.kind == "Z": + assert gate.phase is not None + flips = sum(bindings[p] for p in gate.phase.parameters) + value = float(gate.phase.value) + flips + return [psi * np.where(bit == 1, np.exp(1j * pi * value), 1)] + assert gate.kind in ("MEASURE", "DISCARD") + return [psi * (bit == 0), psi * (bit == 1)] + + +def extracted_distribution( + circuit: ExtractedCircuit, layout: ComponentLayout, bindings: Mapping[int, int] +) -> Distribution: + """Return the exact output distribution of an extracted circuit. + + Lines start in |+⟩. MEASURE and DISCARD dephase their line in the Z basis + (each outcome becomes a branch). Output lines are read in the Z basis at + the end, in slot order. ``bindings`` gives each phase parameter ID its bit. + """ + num_lines = circuit.n_lines + branches = [np.full(2**num_lines, 2 ** (-num_lines / 2), dtype=complex)] + for gate in circuit.gates: + branches = [ + new + for psi in branches + for new in _extracted_gate(psi, num_lines, gate, bindings) + ] + slot_of_port = {port: slot for slot, port in enumerate(layout.output_indices)} + lines = [0] * len(layout.output_indices) + for line, port in circuit.line_end.items(): + if port is not None: + lines[slot_of_port[port]] = line + result: Distribution = {} + for psi in branches: + probabilities = np.abs(psi) ** 2 + for i in np.flatnonzero(probabilities > 1e-15): + key = tuple(int((i >> line) & 1) for line in lines) + result[key] = result.get(key, 0.0) + float(probabilities[i]) + return result + + +def ops_distribution(ops: HeisenbergOps, f: Sequence[int]) -> Distribution: + """Return the exact distribution of the recorded outcomes of ``ops``. + + The state starts in |0…0⟩. A measurement records the eigenvalue bit of its + signed Pauli; a forgotten one only branches. + """ + f_vector = np.asarray(f, dtype=np.int64) + psi = np.zeros(2**ops.num_lines, dtype=complex) + psi[0] = 1 + branches: list[tuple[np.ndarray, dict[int, int]]] = [(psi, {})] + for j in range(ops.num_ops): + sign = (int(ops.sign[j, 0]) + int(ops.sign[j, 1:] @ f_vector)) & 1 + x, z = mask_of(ops.x[j]), mask_of(ops.z[j]) + slot = int(ops.slot[j]) + updated: list[tuple[np.ndarray, dict[int, int]]] = [] + for state, record in branches: + paulied = (-1) ** sign * apply_pauli(state, x, z) + if ops.kind[j] == ROTATION: + rotated = cos(ops.angle[j]) * state - 1j * sin(ops.angle[j]) * paulied + updated.append((rotated, record)) + continue + for bit in (0, 1): + projected = (state + (-1) ** bit * paulied) / 2 + if np.vdot(projected, projected).real < 1e-15: + continue + updated.append( + (projected, {**record, slot: bit} if slot >= 0 else record) + ) + branches = updated + result: Distribution = {} + for state, record in branches: + key = tuple(record[s] for s in range(ops.num_outputs)) + result[key] = result.get(key, 0.0) + float(np.vdot(state, state).real) + return result diff --git a/test/unit/statevector/__init__.py b/test/unit/statevector/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/test/unit/statevector/test_heisenberg.py b/test/unit/statevector/test_heisenberg.py new file mode 100644 index 00000000..0fb571ad --- /dev/null +++ b/test/unit/statevector/test_heisenberg.py @@ -0,0 +1,114 @@ +from fractions import Fraction +from math import pi +from test.helpers.statevector import ( + CIRCUITS, + assert_distributions_close, + extracted_distribution, + f_patterns, + hand_built_circuits, + ops_distribution, + quantum_components, +) + +import pytest +from zx_param import Phase, Scalar +from zx_param.extract import ExtractedCircuit, Gate, extract_circuit + +from tsim.statevector.heisenberg import ( + MEASUREMENT, + ROTATION, + ComponentLayout, + heisenberg_ops, +) + +LAYOUT = ComponentLayout(f_selection=(0,), f_position={7: 0}, output_indices=(0,)) + + +def _one_line(gates, *, line_start=None): + return ExtractedCircuit( + gates=gates, + n_lines=1, + line_start={0: line_start}, + line_end={0: 0}, + scalar=Scalar.one(), + classical_ports=frozenset({0}), + ) + + +def test_t_on_plus_is_a_rotation_about_x(): + ops = heisenberg_ops(_one_line([Gate("Z", (0,), Phase(Fraction(1, 4)))]), LAYOUT) + assert ops.kind.tolist() == [ROTATION, MEASUREMENT] + assert ops.x.tolist() == [[True], [True]] + assert ops.z.tolist() == [[False], [False]] + assert ops.angle[0] == pytest.approx(pi / 8) + assert ops.sign.tolist() == [[False, False], [False, False]] + assert ops.slot.tolist() == [-1, 0] + assert ops.num_outputs == 1 + + +def test_clifford_part_of_a_phase_goes_into_the_frame(): + # 3π/4 = π (a Z, into the frame) − π/4 (a rotation by −π/8). + ops = heisenberg_ops(_one_line([Gate("Z", (0,), Phase(Fraction(3, 4)))]), LAYOUT) + assert ops.kind.tolist() == [ROTATION, MEASUREMENT] + assert ops.angle[0] == pytest.approx(-pi / 8) + assert ops.sign[:, 0].tolist() == [False, False] + + +def test_noise_flips_the_ops_it_anticommutes_with(): + # |+⟩, Z^f, H, T, measure Z: in the |0⟩ frame the noise is X while the T + # and the measurement are about Z, so both signs depend on f. + gates = [ + Gate("Z", (0,), Phase(Fraction(0), frozenset({7}))), + Gate("H", (0,)), + Gate("Z", (0,), Phase(Fraction(1, 4))), + ] + ops = heisenberg_ops(_one_line(gates), LAYOUT) + assert ops.kind.tolist() == [ROTATION, MEASUREMENT] + assert ops.x.tolist() == [[False], [False]] + assert ops.z.tolist() == [[True], [True]] + assert ops.sign.tolist() == [[False, True], [False, True]] + + +def test_a_line_from_an_input_is_rejected(): + with pytest.raises(ValueError, match="input port"): + heisenberg_ops(_one_line([], line_start=3), LAYOUT) + + +def test_component_layout_reads_f_bits_and_ports(): + (component,), layout = quantum_components(CIRCUITS["t_ghz"], sample_detectors=True) + component_layout = ComponentLayout.of(component, layout) + assert list(component_layout.f_selection) == sorted(component_layout.f_selection) + assert len(component_layout.f_selection) == 3 + assert component_layout.output_indices == (0, 1, 2) + assert sorted(component_layout.f_position.values()) == [0, 1, 2] + + +@pytest.mark.parametrize("name", sorted(CIRCUITS)) +@pytest.mark.parametrize("sample_detectors", [True, False]) +def test_ops_reproduce_the_extracted_circuit(name, sample_detectors): + components, layout = quantum_components( + CIRCUITS[name], sample_detectors=sample_detectors + ) + assert components + for component in components: + component_layout = ComponentLayout.of(component, layout) + circuit = extract_circuit(component) + ops = heisenberg_ops(circuit, component_layout) + for f in f_patterns(ops.num_f): + bindings = {p: f[i] for p, i in component_layout.f_position.items()} + assert_distributions_close( + ops_distribution(ops, f), + extracted_distribution(circuit, component_layout, bindings), + ) + + +@pytest.mark.parametrize("index", [0, 1]) +def test_ops_reproduce_hand_built_circuits(index): + _, circuit, component_layout = hand_built_circuits()[index] + ops = heisenberg_ops(circuit, component_layout) + for f in f_patterns(ops.num_f): + bindings = {p: f[i] for p, i in component_layout.f_position.items()} + assert_distributions_close( + ops_distribution(ops, f), + extracted_distribution(circuit, component_layout, bindings), + ) From 77338bd853b1ea20c1bf73c35c76465e76033888 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Sat, 10 Oct 2026 16:38:44 -0400 Subject: [PATCH 02/18] Schedule statevector ops to keep the active width small Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/schedule.py | 168 +++++++++++++++++++++++++ test/unit/statevector/test_schedule.py | 88 +++++++++++++ 2 files changed, 256 insertions(+) create mode 100644 src/tsim/statevector/schedule.py create mode 100644 test/unit/statevector/test_schedule.py diff --git a/src/tsim/statevector/schedule.py b/src/tsim/statevector/schedule.py new file mode 100644 index 00000000..a43ca4cd --- /dev/null +++ b/src/tsim/statevector/schedule.py @@ -0,0 +1,168 @@ +"""Reordering of Heisenberg-frame operations to keep the dense array small. + +Two operations must keep their relative order exactly when they anticommute; +every order that respects these pairs gives the same distribution. The active +width k = n − dim S is tracked with the Pauli group S that stabilizes every +trajectory up to sign. ``schedule`` runs measurements as early as possible and +widening rotations as late as possible: when nothing else can run, it runs the +pending measurement with the fewest unrun rotations among its ancestors. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + +from tsim.statevector.heisenberg import MEASUREMENT, HeisenbergOps + + +class StabilizerGroup: + """Generators, up to sign, of the Pauli group stabilizing every trajectory.""" + + def __init__(self, num_qubits: int) -> None: + """Start from ``|0…0⟩``: the generators are ``Z_0 … Z_{n-1}``.""" + self.x = np.zeros((num_qubits, num_qubits), dtype=bool) + self.z = np.eye(num_qubits, dtype=bool) + + @property + def width(self) -> int: + """Return the active width: qubits minus independent generators.""" + return self.x.shape[1] - self.x.shape[0] + + def anticommuting(self, x: np.ndarray, z: np.ndarray) -> np.ndarray: + """Return which generators anticommute with the Pauli ``(x, z)``.""" + return np.logical_xor.reduce(self.x & z, axis=1) ^ np.logical_xor.reduce( + self.z & x, axis=1 + ) + + def rotate(self, x: np.ndarray, z: np.ndarray) -> None: + """Keep the generators that a rotation about ``(x, z)`` preserves.""" + anti = np.flatnonzero(self.anticommuting(x, z)) + if len(anti) == 0: + return + first, rest = anti[0], anti[1:] + self.x[rest] ^= self.x[first] + self.z[rest] ^= self.z[first] + self.x = np.delete(self.x, first, axis=0) + self.z = np.delete(self.z, first, axis=0) + + def measure(self, x: np.ndarray, z: np.ndarray) -> None: + """Update the generators after measuring the Pauli ``(x, z)``.""" + anti = np.flatnonzero(self.anticommuting(x, z)) + if len(anti): + first, rest = anti[0], anti[1:] + self.x[rest] ^= self.x[first] + self.z[rest] ^= self.z[first] + self.x[first] = x + self.z[first] = z + elif not self.contains(x, z): + self.x = np.vstack([self.x, x[None]]) + self.z = np.vstack([self.z, z[None]]) + + def contains(self, x: np.ndarray, z: np.ndarray) -> bool: + """Return whether ``(x, z)`` is a product of the generators, up to sign.""" + matrix = np.hstack([self.x, self.z]) + extended = np.vstack([matrix, np.concatenate([x, z])[None]]) + return _rank(extended) == _rank(matrix) + + +def _rank(matrix: np.ndarray) -> int: + """Return the GF(2) rank of a boolean matrix.""" + m = matrix.copy() + rank = 0 + for column in range(m.shape[1]): + if rank == m.shape[0]: + break + rows = np.flatnonzero(m[rank:, column]) + if len(rows) == 0: + continue + pivot = rank + rows[0] + m[[rank, pivot]] = m[[pivot, rank]] + others = np.flatnonzero(m[:, column]) + others = others[others != rank] + m[others] ^= m[rank] + rank += 1 + return rank + + +def anticommutation(ops: HeisenbergOps) -> np.ndarray: + """Return the symmetric matrix of which pairs of operations anticommute.""" + x = ops.x.astype(np.float32) + z = ops.z.astype(np.float32) + return (x @ z.T + z @ x.T) % 2 == 1 + + +def _ancestors(before: np.ndarray) -> np.ndarray: + """Return ``a[j, i]``: op ``i`` must run before op ``j`` (transitively).""" + num_ops = before.shape[0] + ancestors = np.zeros((num_ops, num_ops), dtype=bool) + for j in range(num_ops): + preds = np.flatnonzero(before[:, j]) + if len(preds): + ancestors[j] = ancestors[preds].any(axis=0) + ancestors[j, preds] = True + return ancestors + + +@dataclass(frozen=True) +class Schedule: + """An order of the operations and the largest active width it reaches. + + Attributes: + order: Indices into the ``HeisenbergOps``, in execution order. + width: The largest structural active width along ``order`` (W). + + """ + + order: np.ndarray + width: int + + +def schedule(ops: HeisenbergOps) -> Schedule: + """Return an order of ``ops`` that keeps the active width small.""" + num_ops = ops.num_ops + before = np.triu(anticommutation(ops), 1) + ancestors = _ancestors(before) + successors = [np.flatnonzero(before[i]) for i in range(num_ops)] + waiting = before.sum(axis=0) + is_measurement = ops.kind == MEASUREMENT + done = np.zeros(num_ops, dtype=bool) + group = StabilizerGroup(ops.num_lines) + order: list[int] = [] + width = 0 + + def run(j: int) -> None: + nonlocal width + if is_measurement[j]: + group.measure(ops.x[j], ops.z[j]) + else: + group.rotate(ops.x[j], ops.z[j]) + done[j] = True + order.append(j) + width = max(width, group.width) + waiting[successors[j]] -= 1 + + while not done.all(): + progressed = True + while progressed: + progressed = False + for j in np.flatnonzero(~done): + if waiting[j] or done[j]: + continue + keeps_width = not group.anticommuting(ops.x[j], ops.z[j]).any() + if is_measurement[j] or keeps_width: + run(int(j)) + progressed = True + pending = np.flatnonzero(~done & is_measurement) + if len(pending) == 0: + for j in np.flatnonzero(~done): + run(int(j)) + break + unrun_rotations = ~done & ~is_measurement + costs = (ancestors[pending] & unrun_rotations).sum(axis=1) + target = int(pending[np.argmin(costs)]) + for j in np.flatnonzero(ancestors[target] & ~done): + run(int(j)) + run(target) + return Schedule(order=np.asarray(order, dtype=np.int32), width=width) diff --git a/test/unit/statevector/test_schedule.py b/test/unit/statevector/test_schedule.py new file mode 100644 index 00000000..22f5044e --- /dev/null +++ b/test/unit/statevector/test_schedule.py @@ -0,0 +1,88 @@ +from test.helpers.statevector import ( + assert_distributions_close, + component_ops, + f_patterns, + ops_distribution, +) + +import numpy as np + +from tsim.statevector.heisenberg import MEASUREMENT, ROTATION, HeisenbergOps +from tsim.statevector.schedule import StabilizerGroup, anticommutation, schedule + + +def _bits(*rows): + return np.array(rows, dtype=bool) + + +def _two_expansions_then_two_measurements(): + # Rotate about X0, rotate about X1, measure Z0, measure Z1. + return HeisenbergOps( + num_lines=2, + num_f=0, + kind=np.array([ROTATION, ROTATION, MEASUREMENT, MEASUREMENT], dtype=np.uint8), + x=_bits([1, 0], [0, 1], [0, 0], [0, 0]), + z=_bits([0, 0], [0, 0], [1, 0], [0, 1]), + angle=np.array([0.3, 0.3, 0.0, 0.0]), + sign=np.zeros((4, 1), dtype=bool), + slot=np.array([-1, -1, 0, 1], dtype=np.int32), + ) + + +def test_stabilizer_group_tracks_the_width(): + group = StabilizerGroup(2) + assert group.width == 0 + group.rotate(np.array([True, False]), np.array([False, False])) + assert group.width == 1 + group.rotate(np.array([False, False]), np.array([True, False])) + assert group.width == 1 + group.measure(np.array([True, False]), np.array([False, False])) + assert group.width == 0 + group.measure(np.array([False, False]), np.array([True, False])) + assert group.width == 0 + assert group.contains(np.array([False, False]), np.array([True, False])) + + +def test_anticommutation_finds_the_pairs(): + anti = anticommutation(_two_expansions_then_two_measurements()) + assert anti[0, 2] and anti[2, 0] and anti[1, 3] and anti[3, 1] + assert not anti[0, 1] and not anti[0, 3] and not anti[2, 3] + + +def test_measurements_run_before_unrelated_expansions(): + scheduled = schedule(_two_expansions_then_two_measurements()) + assert scheduled.order.tolist() == [0, 2, 1, 3] + assert scheduled.width == 1 + + +def test_schedules_keep_every_anticommuting_pair_in_order(): + for ops in component_ops(): + order = schedule(ops).order + assert sorted(order.tolist()) == list(range(ops.num_ops)) + position = np.empty(ops.num_ops, dtype=int) + position[order] = np.arange(ops.num_ops) + first, second = np.nonzero(np.triu(anticommutation(ops), 1)) + assert np.all(position[first] < position[second]) + + +def test_schedules_preserve_the_distribution(): + for ops in component_ops(): + reordered = ops.take(schedule(ops).order) + for f in f_patterns(ops.num_f): + assert_distributions_close( + ops_distribution(reordered, f), ops_distribution(ops, f) + ) + + +def test_width_is_the_largest_structural_width_of_the_order(): + for ops in component_ops(): + scheduled = schedule(ops) + group = StabilizerGroup(ops.num_lines) + widths = [0] + for j in scheduled.order: + if ops.kind[j] == MEASUREMENT: + group.measure(ops.x[j], ops.z[j]) + else: + group.rotate(ops.x[j], ops.z[j]) + widths.append(group.width) + assert scheduled.width == max(widths) From 39dad029ce1055821d71f7623564dfc6d06851be Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Sat, 10 Oct 2026 16:46:46 -0400 Subject: [PATCH 03/18] Plan statevector ops on a fixed number of array axes Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/plan.py | 315 +++++++++++++++++++++++++++++ test/helpers/statevector.py | 85 ++++++++ test/unit/statevector/test_plan.py | 95 +++++++++ 3 files changed, 495 insertions(+) create mode 100644 src/tsim/statevector/plan.py create mode 100644 test/unit/statevector/test_plan.py diff --git a/src/tsim/statevector/plan.py b/src/tsim/statevector/plan.py new file mode 100644 index 00000000..077e1d91 --- /dev/null +++ b/src/tsim/statevector/plan.py @@ -0,0 +1,315 @@ +"""Fixed-width sampling plans. + +``plan`` walks scheduled Heisenberg ops with a Clifford frame ``F``: a +``stim.Tableau`` that maps Paulis of the ``|0…0⟩`` frame to virtual +coordinates. Virtual coordinates ``0 … W-1`` are the axes of a dense array of +``2^W`` amplitudes, each active or free (in ``|0⟩``); coordinates ``W … n-1`` +stay dormant (in ``|0⟩``). Each op becomes at most one action, and every sign an +``AffineRow`` over the symbols: the component f-bits, then the branch bits of +random measurements. A random measurement keeps the reference state on its +``b = 0`` branch; its ``b = 1`` branch differs by a correction Pauli, which +flips the sign of every later op it anticommutes with. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Sequence +from dataclasses import dataclass + +import numpy as np +import stim + +from tsim.statevector.heisenberg import ROTATION, HeisenbergOps +from tsim.statevector.schedule import Schedule + + +@dataclass(frozen=True) +class AffineRow: + """The parity ``constant ⊕ (XOR of the symbols in terms)``.""" + + constant: bool + terms: frozenset[int] + + def __xor__(self, other: AffineRow) -> AffineRow: + """Return the XOR of two parities.""" + return AffineRow(self.constant ^ other.constant, self.terms ^ other.terms) + + def evaluate(self, symbols: Sequence[int]) -> int: + """Return the parity's value for the given symbol bits.""" + return int(self.constant) ^ (sum(symbols[t] for t in self.terms) & 1) + + +@dataclass(frozen=True) +class Rotate: + """Apply ``exp(-iθ (-1)^sign P)`` for the Pauli with W-bit masks ``x``, ``z``.""" + + x: int + z: int + angle: float + sign: AffineRow + + +@dataclass(frozen=True) +class MeasureActive: + """Measure the Pauli ``x``, ``z`` on the array, then reset ``pivot`` to ``|0⟩``. + + The sampled eigenvalue bit becomes symbol ``branch``; ``slot`` (if not -1) + records ``branch ⊕ sign``. + """ + + x: int + z: int + pivot: int + branch: int + sign: AffineRow + slot: int + + +@dataclass(frozen=True) +class MeasureDormant: + """Sample a fair bit into symbol ``branch``; ``slot`` records it ⊕ ``sign``.""" + + branch: int + sign: AffineRow + slot: int + + +@dataclass(frozen=True) +class Record: + """Record the deterministic value ``sign`` in ``slot``.""" + + sign: AffineRow + slot: int + + +Action = Rotate | MeasureActive | MeasureDormant | Record + + +@dataclass(frozen=True) +class Plan: + """Actions on ``2^width`` amplitudes, with signs over the symbols. + + Attributes: + width: Number of array axes (W). + num_f: Number of f-bit symbols (symbols ``0 … num_f-1``). + num_branches: Number of branch symbols (``num_f …``). + num_outputs: Number of output slots. + actions: The actions, in execution order. + + """ + + width: int + num_f: int + num_branches: int + num_outputs: int + actions: tuple[Action, ...] + + @property + def num_symbols(self) -> int: + """Return the number of f-bit and branch symbols.""" + return self.num_f + self.num_branches + + +def plan(ops: HeisenbergOps, schedule: Schedule) -> Plan: + """Return the fixed-width plan of ``ops`` in schedule order. + + Raises: + RuntimeError: If the ops need more than ``schedule.width`` active axes. + + """ + planner = _Planner(ops, schedule) + for position in range(len(schedule.order)): + planner.step(position) + return Plan( + width=schedule.width, + num_f=ops.num_f, + num_branches=planner.num_branches, + num_outputs=ops.num_outputs, + actions=tuple(planner.actions), + ) + + +def _anticommutes( + x: np.ndarray, z: np.ndarray, xs: np.ndarray, zs: np.ndarray +) -> np.ndarray: + """Return which rows ``(xs[i], zs[i])`` anticommute with the Pauli ``(x, z)``.""" + return np.logical_xor.reduce(xs & z, axis=1) ^ np.logical_xor.reduce(zs & x, axis=1) + + +def _mask(bits: np.ndarray, axes: Iterable[int]) -> int: + return sum(1 << int(q) for q in axes if bits[q]) + + +_CONTROLLED_PAULI = {(True, False): "CX", (False, True): "CZ", (True, True): "CY"} + + +class _Planner: + """The frame, the active axes and the branch corrections during planning.""" + + def __init__(self, ops: HeisenbergOps, schedule: Schedule) -> None: + """Start with the identity frame and every coordinate dormant.""" + self.ops = ops + self.order = schedule.order + self.width = schedule.width + self.frame = stim.Tableau(ops.num_lines) + self.active: set[int] = set() + self.num_branches = 0 + self.correction_x = np.zeros((0, ops.num_lines), dtype=bool) + self.correction_z = np.zeros((0, ops.num_lines), dtype=bool) + self.correction_condition: list[AffineRow] = [] + self.actions: list[Action] = [] + + def _apply(self, gate: str, *targets: int) -> None: + """Apply a virtual Clifford: ``F ← gate ∘ F``.""" + self.frame.append(stim.Tableau.from_named_gate(gate), list(targets)) + + def _virtual(self, j: int) -> tuple[np.ndarray, np.ndarray, bool]: + """Return ``F(P_j)`` as bit rows and a sign flag.""" + pauli = stim.PauliString.from_numpy(xs=self.ops.x[j], zs=self.ops.z[j]) + image = self.frame(pauli) + x, z = image.to_numpy() + return x, z, bool(image.sign == -1) + + def _sign(self, j: int, negative: bool) -> AffineRow: + """Return op ``j``'s sign: noise row, frame sign and branch corrections.""" + row = self.ops.sign[j] + terms = frozenset(int(t) for t in np.flatnonzero(row[1:])) + sign = AffineRow(bool(row[0]) ^ negative, terms) + anti = _anticommutes( + self.ops.x[j], self.ops.z[j], self.correction_x, self.correction_z + ) + for m in np.flatnonzero(anti): + sign = sign ^ self.correction_condition[m] + return sign + + def _new_branch(self) -> int: + branch = self.ops.num_f + self.num_branches + self.num_branches += 1 + return branch + + def _correction(self, axis: int) -> tuple[np.ndarray, np.ndarray]: + """Return ``F⁻¹(X_axis)``, the correction Pauli in the ``|0…0⟩`` frame.""" + x, z = self.frame.inverse_x_output(axis).to_numpy() + return x, z + + def _add_correction(self, axis: int, condition: AffineRow) -> None: + x, z = self._correction(axis) + self.correction_x = np.vstack([self.correction_x, x[None]]) + self.correction_z = np.vstack([self.correction_z, z[None]]) + self.correction_condition.append(condition) + + def step(self, position: int) -> None: + """Plan the op at ``position`` of the schedule.""" + j = int(self.order[position]) + x, z, negative = self._virtual(j) + sign = self._sign(j, negative) + dormant_x = [int(q) for q in np.flatnonzero(x) if int(q) not in self.active] + on_active = any(x[q] or z[q] for q in self.active) + if self.ops.kind[j] == ROTATION: + if dormant_x: + self._promote(j, dormant_x) + elif on_active: + axes = sorted(self.active) + angle = float(self.ops.angle[j]) + self.actions.append(Rotate(_mask(x, axes), _mask(z, axes), angle, sign)) + return + slot = int(self.ops.slot[j]) + if dormant_x: + self._measure_dormant(position, j, dormant_x, negative, sign, slot) + elif on_active: + self._measure_active(j, x, z, negative, sign, slot) + elif slot >= 0: + self.actions.append(Record(sign, slot)) + + def _promote(self, j: int, dormant_x: list[int]) -> None: + """Move the rotation's dormant X support onto one free axis and rotate.""" + free = [q for q in range(self.width) if q not in self.active] + if not free: + raise RuntimeError( + f"the plan needs more active axes than the schedule's width {self.width}" + ) + in_array = [q for q in dormant_x if q < self.width] + if in_array: + axis = in_array[0] + else: + axis = free[0] + self._apply("SWAP", axis, dormant_x[0]) + x, _, _ = self._virtual(j) + for q in np.flatnonzero(x): + if int(q) != axis and int(q) not in self.active: + self._apply("CX", axis, int(q)) + x, z, negative = self._virtual(j) + axes = sorted(self.active | {axis}) + angle = float(self.ops.angle[j]) + sign = self._sign(j, negative) + self.actions.append(Rotate(_mask(x, axes), _mask(z, axes), angle, sign)) + self.active.add(axis) + + def _measure_dormant( + self, + position: int, + j: int, + dormant_x: list[int], + negative: bool, + sign: AffineRow, + slot: int, + ) -> None: + """Turn the measured Pauli into ``±X_d`` without touching the state, then H.""" + d = dormant_x[0] + for q in dormant_x[1:]: + self._apply("CX", d, q) + _, z, _ = self._virtual(j) + if z[d]: + self._apply("S_DAG", d) + x, z, _ = self._virtual(j) + for q in range(self.ops.num_lines): + if q != d and (x[q] or z[q]): + self._apply(_CONTROLLED_PAULI[(bool(x[q]), bool(z[q]))], d, q) + x, z, after = self._virtual(j) + assert x[d] and x.sum() == 1 and not z.any() + self._apply("H", d) + later = self.order[position + 1 :] + kx, kz = self._correction(d) + anti = _anticommutes(kx, kz, self.ops.x[later], self.ops.z[later]) + if slot < 0 and not anti.any(): + return + branch = self._new_branch() + self._add_correction(d, AffineRow(after ^ negative, frozenset({branch}))) + self.actions.append(MeasureDormant(branch, sign, slot)) + + def _measure_active( + self, + j: int, + x: np.ndarray, + z: np.ndarray, + negative: bool, + sign: AffineRow, + slot: int, + ) -> None: + """Measure on the array, then turn the Pauli into ``±Z_pivot`` and free it.""" + axes = sorted(self.active) + x_mask, z_mask = _mask(x, axes), _mask(z, axes) + x_support = [q for q in axes if x[q]] + if x_support: + pivot = x_support[0] + for q in x_support[1:]: + self._apply("CX", pivot, q) + _, z_now, _ = self._virtual(j) + for q in axes: + if q != pivot and z_now[q]: + self._apply("CZ", pivot, q) + _, z_now, _ = self._virtual(j) + if z_now[pivot]: + self._apply("S_DAG", pivot) + self._apply("H", pivot) + else: + z_support = [q for q in axes if z[q]] + pivot = z_support[0] + for q in z_support[1:]: + self._apply("CX", q, pivot) + _, _, after = self._virtual(j) + assert x_support or after == negative + branch = self._new_branch() + self.actions.append(MeasureActive(x_mask, z_mask, pivot, branch, sign, slot)) + self._add_correction(pivot, AffineRow(after ^ negative, frozenset({branch}))) + self.active.remove(pivot) diff --git a/test/helpers/statevector.py b/test/helpers/statevector.py index d821b315..80d2f7c5 100644 --- a/test/helpers/statevector.py +++ b/test/helpers/statevector.py @@ -24,6 +24,7 @@ HeisenbergOps, heisenberg_ops, ) +from tsim.statevector.plan import MeasureActive, Plan, Record, Rotate Distribution = dict[tuple[int, ...], float] @@ -363,3 +364,87 @@ def ops_distribution(ops: HeisenbergOps, f: Sequence[int]) -> Distribution: key = tuple(record[s] for s in range(ops.num_outputs)) result[key] = result.get(key, 0.0) + float(np.vdot(state, state).real) return result + + +def reference_rotate( + psi: np.ndarray, x: int, z: int, angle: float, negate: bool +) -> np.ndarray: + """Return exp(-iθ (±P)) ψ, the ``ROTATE`` and ``ROTATE_DIAG`` kernels.""" + sign = -1 if negate else 1 + return cos(angle) * psi - 1j * sin(angle) * sign * apply_pauli(psi, x, z) + + +def reference_measure( + psi: np.ndarray, x: int, z: int, pivot: int, outcome: int +) -> tuple[np.ndarray, float]: + """Return the ``MEASURE`` / ``MEASURE_DIAG`` kernel's state and outcome probability. + + The kernel projects onto the (-1)^outcome eigenspace of P and resets the + pivot axis to |0⟩: it keeps the pivot-0 half and renormalizes. + """ + index = np.arange(len(psi)) + keep = ((index >> pivot) & 1) == 0 + norm = float(np.vdot(psi, psi).real) + if x: + paulied = apply_pauli(psi, x, z) + expectation = float(np.vdot(psi, paulied).real) + p1 = min(max((norm - expectation) / (2 * norm), 0.0), 1.0) + p = p1 if outcome else 1 - p1 + projected = (psi + (-1) ** outcome * paulied) / np.sqrt( + max(2 * norm * p, 1e-300) + ) + return np.where(keep, projected, 0), p + odd = parity(index & z) + p1 = float(np.sum(np.abs(psi[odd]) ** 2)) / norm + p = p1 if outcome else 1 - p1 + source = index ^ ((odd ^ bool(outcome)).astype(np.int64) << pivot) + return np.where(keep, psi[source] / np.sqrt(max(norm * p, 1e-300)), 0), p + + +def plan_distribution(plan: Plan, f: Sequence[int]) -> Distribution: + """Return the exact output distribution of a plan by enumerating its branches.""" + result: Distribution = {} + + def written(record: list[int], slot: int, value: int) -> list[int]: + if slot < 0: + return record + updated = list(record) + updated[slot] = value + return updated + + def walk( + a: int, psi: np.ndarray, symbols: list[int], record: list[int], weight: float + ) -> None: + if weight < 1e-14: + return + if a == len(plan.actions): + key = tuple(record) + result[key] = result.get(key, 0.0) + weight + return + action = plan.actions[a] + sign = action.sign.evaluate(symbols) + if isinstance(action, Rotate): + rotated = reference_rotate( + psi, action.x, action.z, action.angle, bool(sign) + ) + walk(a + 1, rotated, symbols, record, weight) + return + if isinstance(action, Record): + walk(a + 1, psi, symbols, written(record, action.slot, sign), weight) + return + for outcome in (0, 1): + if isinstance(action, MeasureActive): + new, p = reference_measure( + psi, action.x, action.z, action.pivot, outcome + ) + else: + new, p = psi, 0.5 + branched = list(symbols) + branched[action.branch] = outcome + value = outcome ^ sign + walk(a + 1, new, branched, written(record, action.slot, value), weight * p) + + psi = np.zeros(2**plan.width, dtype=complex) + psi[0] = 1 + walk(0, psi, [*f, *[0] * plan.num_branches], [0] * plan.num_outputs, 1.0) + return result diff --git a/test/unit/statevector/test_plan.py b/test/unit/statevector/test_plan.py new file mode 100644 index 00000000..cc119410 --- /dev/null +++ b/test/unit/statevector/test_plan.py @@ -0,0 +1,95 @@ +from fractions import Fraction +from math import pi +from test.helpers.statevector import ( + assert_distributions_close, + component_ops, + f_patterns, + ops_distribution, + plan_distribution, +) + +import numpy as np +import pytest +from zx_param import Phase, Scalar +from zx_param.extract import ExtractedCircuit, Gate + +from tsim.statevector.heisenberg import ComponentLayout, heisenberg_ops +from tsim.statevector.plan import ( + AffineRow, + MeasureActive, + MeasureDormant, + Rotate, + plan, +) +from tsim.statevector.schedule import Schedule, schedule + + +def _t_then_h(): + # |+⟩, T, H, measure Z: in the |0⟩ frame a rotation about X, then a + # measurement of Z, which anticommute. + circuit = ExtractedCircuit( + gates=[Gate("Z", (0,), Phase(Fraction(1, 4))), Gate("H", (0,))], + n_lines=1, + line_start={0: None}, + line_end={0: 0}, + scalar=Scalar.one(), + classical_ports=frozenset({0}), + ) + return heisenberg_ops(circuit, ComponentLayout((0,), {7: 0}, (0,))) + + +def test_t_then_h_promotes_one_axis_and_measures_it(): + ops = _t_then_h() + planned = plan(ops, schedule(ops)) + assert planned.width == 1 + assert planned.num_branches == 1 + rotate, measure = planned.actions + assert isinstance(rotate, Rotate) + assert (rotate.x, rotate.z) == (1, 0) + assert rotate.angle == pytest.approx(pi / 8) + assert rotate.sign == AffineRow(False, frozenset()) + assert measure == MeasureActive( + x=0, z=1, pivot=0, branch=1, sign=AffineRow(False, frozenset()), slot=0 + ) + + +def test_affine_rows_xor_and_evaluate(): + row = AffineRow(True, frozenset({0, 2})) ^ AffineRow(False, frozenset({2, 3})) + assert row == AffineRow(True, frozenset({0, 3})) + assert row.evaluate([1, 0, 0, 0]) == 0 + assert row.evaluate([1, 0, 1, 1]) == 1 + + +def test_a_schedule_that_is_too_narrow_is_rejected(): + ops = _t_then_h() + with pytest.raises(RuntimeError, match="active axes"): + plan(ops, Schedule(order=np.arange(ops.num_ops, dtype=np.int32), width=0)) + + +def test_plans_reproduce_the_scheduled_ops(): + for ops in component_ops(): + scheduled = schedule(ops) + planned = plan(ops, scheduled) + reordered = ops.take(scheduled.order) + for f in f_patterns(ops.num_f): + assert_distributions_close( + plan_distribution(planned, f), ops_distribution(reordered, f), atol=1e-9 + ) + + +def test_plans_are_well_formed(): + for ops in component_ops(): + scheduled = schedule(ops) + planned = plan(ops, scheduled) + assert planned.width == scheduled.width + assert planned.num_outputs == ops.num_outputs + defined = set(range(planned.num_f)) + for action in planned.actions: + assert action.sign.terms <= defined + if isinstance(action, (Rotate, MeasureActive)): + assert 0 <= action.x < 2**planned.width + assert 0 <= action.z < 2**planned.width + if isinstance(action, (MeasureActive, MeasureDormant)): + assert action.branch not in defined + defined.add(action.branch) + assert defined == set(range(planned.num_symbols)) From 1574005e43e74d18e89c739ff78cfa12bf63c713 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Sat, 10 Oct 2026 16:56:50 -0400 Subject: [PATCH 04/18] Compile statevector components into action tapes Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/program.py | 193 ++++++++++++++++++++++++++ test/helpers/statevector.py | 17 +++ test/unit/statevector/test_program.py | 158 +++++++++++++++++++++ 3 files changed, 368 insertions(+) create mode 100644 src/tsim/statevector/program.py create mode 100644 test/unit/statevector/test_program.py diff --git a/src/tsim/statevector/program.py b/src/tsim/statevector/program.py new file mode 100644 index 00000000..e7616d61 --- /dev/null +++ b/src/tsim/statevector/program.py @@ -0,0 +1,193 @@ +"""Compilation of one component for ``strategy="statevector"``. + +``plan_component`` runs extraction, the Heisenberg rewrite, scheduling and +planning; ``lower`` turns the plan into the arrays of a ``StatevectorComponent``, +an action tape that ``execute.run`` interprets. +""" + +from __future__ import annotations + +import equinox as eqx +import jax +import jax.numpy as jnp +import numpy as np +from jax import Array +from zx_param import Graph +from zx_param.extract import extract_circuit + +from tsim.core.params import ParamLayout +from tsim.statevector.heisenberg import ComponentLayout, heisenberg_ops +from tsim.statevector.plan import MeasureActive, MeasureDormant, Plan, Rotate, plan +from tsim.statevector.schedule import schedule + +ROTATE_DIAG = 0 +ROTATE = 1 +MEASURE_DIAG = 2 +MEASURE = 3 +MEASURE_DORMANT = 4 +RECORD = 5 + +MAX_WIDTH = 30 + +_HINT = "Use a stabilizer-rank strategy such as strategy='cat5' for this circuit." + + +class StatevectorComponent(eqx.Module): + """A component compiled for ``strategy="statevector"``: an action tape. + + Attributes: + output_indices: Global output index of each output slot. + width: Number of array axes W; lanes hold ``2^W`` amplitudes. + num_words: Number of uint32 words that hold a lane's symbols. + num_outputs: Number of output slots. + dtype: ``"complex64"`` or ``"complex128"``. + f_selection: Global f-bit index of each component f-bit. + kind: Kernel of each action (``ROTATE_DIAG`` … ``RECORD``). + x_mask: X mask of the action's Pauli on the array axes. + z_mask: Z mask of the action's Pauli on the array axes. + cos: cos θ of each rotation. + sin: sin θ of each rotation. + pivot: Axis an active measurement resets to ``|0⟩``. + branch: Symbol a measurement sets, -1 if none. + slot: Output slot an action writes, -1 if none. + sign_const: Constant of each action's sign. + sign_row: Symbol bits of each action's sign, shape ``(A, num_words)``. + + """ + + output_indices: tuple[int, ...] = eqx.field(static=True) + width: int = eqx.field(static=True) + num_words: int = eqx.field(static=True) + num_outputs: int = eqx.field(static=True) + dtype: str = eqx.field(static=True) + f_selection: Array + kind: Array + x_mask: Array + z_mask: Array + cos: Array + sin: Array + pivot: Array + branch: Array + slot: Array + sign_const: Array + sign_row: Array + + @property + def num_actions(self) -> int: + """Return the number of actions on the tape.""" + return int(self.kind.shape[0]) + + +def plan_component( + component: Graph, layout: ParamLayout +) -> tuple[Plan, ComponentLayout]: + """Extract, rewrite, schedule and plan one non-direct component. + + Raises: + ValueError: If extraction fails or the plan needs more than + ``MAX_WIDTH`` array axes. + + """ + component_layout = ComponentLayout.of(component, layout) + outputs = list(component_layout.output_indices) + try: + circuit = extract_circuit(component) + except (RuntimeError, ValueError) as error: + raise ValueError( + "strategy='statevector' could not extract a circuit from the component " + f"with outputs {outputs}: {error}. {_HINT}" + ) from error + ops = heisenberg_ops(circuit, component_layout) + scheduled = schedule(ops) + if scheduled.width > MAX_WIDTH: + raise ValueError( + f"strategy='statevector' needs {scheduled.width} active qubits for the " + f"component with outputs {outputs}, more than {MAX_WIDTH}. {_HINT}" + ) + return plan(ops, scheduled), component_layout + + +def resolve_dtype(dtype: str | None) -> str: + """Return the amplitude dtype: complex128 if JAX has 64-bit types, else complex64. + + Raises: + ValueError: If ``dtype`` is unknown, or complex128 without + ``jax_enable_x64``. + + """ + x64 = jax.dtypes.canonicalize_dtype(np.complex128) == np.dtype(np.complex128) + if dtype is None: + return "complex128" if x64 else "complex64" + if dtype not in ("complex64", "complex128"): + raise ValueError(f"dtype must be 'complex64' or 'complex128', got {dtype!r}") + if dtype == "complex128" and not x64: + raise ValueError("complex128 amplitudes need jax_enable_x64") + return dtype + + +def lower(plan: Plan, layout: ComponentLayout, dtype: str) -> StatevectorComponent: + """Return the action tape of ``plan``.""" + num_actions = len(plan.actions) + num_words = max(1, -(-plan.num_symbols // 32)) + kind = np.zeros(num_actions, dtype=np.int32) + x_mask = np.zeros(num_actions, dtype=np.uint32) + z_mask = np.zeros(num_actions, dtype=np.uint32) + angle = np.zeros(num_actions, dtype=np.float64) + pivot = np.zeros(num_actions, dtype=np.int32) + branch = np.full(num_actions, -1, dtype=np.int32) + slot = np.full(num_actions, -1, dtype=np.int32) + sign_const = np.zeros(num_actions, dtype=bool) + sign_row = np.zeros((num_actions, num_words), dtype=np.uint32) + for a, action in enumerate(plan.actions): + if isinstance(action, Rotate): + kind[a] = ROTATE if action.x else ROTATE_DIAG + x_mask[a], z_mask[a], angle[a] = action.x, action.z, action.angle + elif isinstance(action, MeasureActive): + kind[a] = MEASURE if action.x else MEASURE_DIAG + x_mask[a], z_mask[a] = action.x, action.z + pivot[a], branch[a], slot[a] = action.pivot, action.branch, action.slot + elif isinstance(action, MeasureDormant): + kind[a] = MEASURE_DORMANT + branch[a], slot[a] = action.branch, action.slot + else: + kind[a] = RECORD + slot[a] = action.slot + sign_const[a] = action.sign.constant + for term in action.sign.terms: + sign_row[a, term // 32] |= np.uint32(1) << np.uint32(term % 32) + real = np.float32 if dtype == "complex64" else np.float64 + return StatevectorComponent( + output_indices=layout.output_indices, + width=plan.width, + num_words=num_words, + num_outputs=plan.num_outputs, + dtype=dtype, + f_selection=jnp.asarray(np.array(layout.f_selection, dtype=np.int32)), + kind=jnp.asarray(kind), + x_mask=jnp.asarray(x_mask), + z_mask=jnp.asarray(z_mask), + cos=jnp.asarray(np.cos(angle).astype(real)), + sin=jnp.asarray(np.sin(angle).astype(real)), + pivot=jnp.asarray(pivot), + branch=jnp.asarray(branch), + slot=jnp.asarray(slot), + sign_const=jnp.asarray(sign_const), + sign_row=jnp.asarray(sign_row), + ) + + +def compile_statevector_component( + component: Graph, layout: ParamLayout, *, dtype: str | None = None +) -> StatevectorComponent: + """Compile one non-direct component for ``strategy="statevector"``. + + Args: + component: A connected component of the prepared sampling graph. + layout: The prepared graph's parameter layout. + dtype: Amplitude dtype; ``None`` picks complex128 when JAX has 64-bit + types enabled, else complex64. + + """ + resolved = resolve_dtype(dtype) + planned, component_layout = plan_component(component, layout) + return lower(planned, component_layout, resolved) diff --git a/test/helpers/statevector.py b/test/helpers/statevector.py index 80d2f7c5..71912702 100644 --- a/test/helpers/statevector.py +++ b/test/helpers/statevector.py @@ -10,14 +10,17 @@ from fractions import Fraction from math import cos, pi, sin +import jax.numpy as jnp import numpy as np from zx_param import Graph, Phase, Scalar from zx_param.extract import ExtractedCircuit, Gate, extract_circuit from tsim.circuit import Circuit +from tsim.compile.evaluate import evaluate 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 from tsim.statevector.heisenberg import ( ROTATION, ComponentLayout, @@ -448,3 +451,17 @@ def walk( psi[0] = 1 walk(0, psi, [*f, *[0] * plan.num_branches], [0] * plan.num_outputs, 1.0) return result + + +def stabilizer_rank_distribution( + compiled: CompiledComponent, f: Sequence[int] +) -> Distribution: + """Return P(outputs | f) of a joint-mode stabilizer-rank component.""" + norm_graphs, joint_graphs = compiled.compiled_scalar_graphs + outcomes = list(itertools.product((0, 1), repeat=len(compiled.output_indices))) + params = np.array([[*f, *o] for o in outcomes], dtype=np.uint8) + params = params.reshape(len(outcomes), len(f) + len(compiled.output_indices)) + joint = np.abs(np.asarray(evaluate(joint_graphs, jnp.asarray(params)))) + f_row = np.array([list(f)], dtype=np.uint8).reshape(1, len(f)) + norm = float(np.abs(np.asarray(evaluate(norm_graphs, jnp.asarray(f_row))))[0]) + return {o: float(p) / norm for o, p in zip(outcomes, joint, strict=True)} diff --git a/test/unit/statevector/test_program.py b/test/unit/statevector/test_program.py new file mode 100644 index 00000000..5cb863a4 --- /dev/null +++ b/test/unit/statevector/test_program.py @@ -0,0 +1,158 @@ +from test.helpers.statevector import ( + CIRCUITS, + assert_distributions_close, + f_patterns, + plan_distribution, + quantum_components, + stabilizer_rank_distribution, +) + +import jax +import numpy as np +import pytest + +import tsim.statevector.program as program +from tsim.circuit import Circuit +from tsim.compile.pipeline import classify_direct, compile_program +from tsim.core.graph import prepare_graph +from tsim.statevector.heisenberg import ComponentLayout +from tsim.statevector.plan import ( + AffineRow, + MeasureActive, + MeasureDormant, + Plan, + Record, + Rotate, +) +from tsim.statevector.program import ( + MAX_WIDTH, + MEASURE, + MEASURE_DIAG, + MEASURE_DORMANT, + RECORD, + ROTATE, + ROTATE_DIAG, + compile_statevector_component, + lower, + plan_component, + resolve_dtype, +) + + +@pytest.mark.parametrize("name", sorted(CIRCUITS)) +@pytest.mark.parametrize("sample_detectors", [True, False]) +def test_plans_match_stabilizer_rank(name, sample_detectors): + prepared = prepare_graph(Circuit(CIRCUITS[name]), sample_detectors=sample_detectors) + joint = compile_program(prepared, mode="joint").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: + continue + planned, component_layout = plan_component(component, prepared.layout) + reference = by_outputs[component_layout.output_indices] + selection = tuple(np.asarray(reference.f_selection).tolist()) + assert selection == component_layout.f_selection + for f in f_patterns(planned.num_f): + assert_distributions_close( + plan_distribution(planned, f), + stabilizer_rank_distribution(reference, f), + atol=2e-5, + ) + + +def test_lower_encodes_every_action(): + planned = Plan( + width=2, + num_f=1, + num_branches=3, + num_outputs=2, + actions=( + Rotate(x=0, z=3, angle=0.25, sign=AffineRow(True, frozenset({0}))), + Rotate(x=1, z=2, angle=-0.5, sign=AffineRow(False, frozenset())), + MeasureActive( + 0, 2, 1, branch=1, sign=AffineRow(False, frozenset({0})), slot=1 + ), + MeasureActive( + 1, 0, 0, branch=2, sign=AffineRow(True, frozenset({1})), slot=-1 + ), + MeasureDormant(branch=3, sign=AffineRow(False, frozenset({1, 2})), slot=0), + Record(sign=AffineRow(True, frozenset({3})), slot=1), + ), + ) + layout = ComponentLayout(f_selection=(4,), f_position={}, output_indices=(2, 5)) + compiled = lower(planned, layout, "complex64") + assert compiled.kind.tolist() == [ + ROTATE_DIAG, + ROTATE, + MEASURE_DIAG, + MEASURE, + MEASURE_DORMANT, + RECORD, + ] + assert compiled.x_mask.tolist() == [0, 1, 0, 1, 0, 0] + assert compiled.z_mask.tolist() == [3, 2, 2, 0, 0, 0] + np.testing.assert_allclose(compiled.cos[:2], np.cos([0.25, -0.5]), rtol=1e-6) + np.testing.assert_allclose(compiled.sin[:2], np.sin([0.25, -0.5]), rtol=1e-6) + assert compiled.pivot.tolist()[2:4] == [1, 0] + assert compiled.branch.tolist() == [-1, -1, 1, 2, 3, -1] + assert compiled.slot.tolist() == [-1, -1, 1, -1, 0, 1] + assert compiled.sign_const.tolist() == [True, False, False, True, False, True] + assert compiled.sign_row[:, 0].tolist() == [1, 0, 1, 2, 6, 8] + assert compiled.num_words == 1 + assert compiled.f_selection.tolist() == [4] + assert compiled.output_indices == (2, 5) + assert (compiled.width, compiled.num_outputs, compiled.dtype) == (2, 2, "complex64") + assert compiled.num_actions == 6 + + +def test_lower_packs_symbols_beyond_32_into_more_words(): + planned = Plan( + width=0, + num_f=40, + num_branches=0, + num_outputs=1, + actions=(Record(sign=AffineRow(False, frozenset({0, 33, 39})), slot=0),), + ) + compiled = lower(planned, ComponentLayout(tuple(range(40)), {}, (0,)), "complex64") + assert compiled.num_words == 2 + assert compiled.sign_row.tolist() == [[1, (1 << 1) | (1 << 7)]] + + +def test_compile_statevector_component(): + (component,), layout = quantum_components(CIRCUITS["t_ghz"], sample_detectors=True) + compiled = compile_statevector_component(component, layout) + planned, component_layout = plan_component(component, layout) + assert compiled.width == planned.width + assert compiled.num_actions == len(planned.actions) + assert compiled.output_indices == component_layout.output_indices + assert compiled.dtype == "complex64" + + +def test_too_wide_components_are_rejected(monkeypatch): + assert MAX_WIDTH == 30 + monkeypatch.setattr(program, "MAX_WIDTH", 0) + (component,), layout = quantum_components(CIRCUITS["t_ghz"], sample_detectors=True) + with pytest.raises(ValueError, match="cat5"): + plan_component(component, layout) + + +def test_extraction_failures_are_value_errors(monkeypatch): + def fail(_): + raise RuntimeError("no gflow") + + monkeypatch.setattr(program, "extract_circuit", fail) + (component,), layout = quantum_components(CIRCUITS["t_ghz"], sample_detectors=True) + with pytest.raises(ValueError, match="no gflow"): + plan_component(component, layout) + + +def test_resolve_dtype(): + assert resolve_dtype(None) == "complex64" + assert resolve_dtype("complex64") == "complex64" + with pytest.raises(ValueError, match="jax_enable_x64"): + resolve_dtype("complex128") + with pytest.raises(ValueError, match="complex32"): + resolve_dtype("complex32") + with jax.enable_x64(True): + assert resolve_dtype(None) == "complex128" + assert resolve_dtype("complex128") == "complex128" From 81336d162470f3e3b864f95f4ee905fc89011199 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Sat, 10 Oct 2026 17:03:49 -0400 Subject: [PATCH 05/18] Run statevector action tapes in JAX Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/execute.py | 261 ++++++++++++++++++++++++++ test/helpers/statevector.py | 36 ++++ test/unit/statevector/test_execute.py | 167 ++++++++++++++++ 3 files changed, 464 insertions(+) create mode 100644 src/tsim/statevector/execute.py create mode 100644 test/unit/statevector/test_execute.py diff --git a/src/tsim/statevector/execute.py b/src/tsim/statevector/execute.py new file mode 100644 index 00000000..98ac013a --- /dev/null +++ b/src/tsim/statevector/execute.py @@ -0,0 +1,261 @@ +"""JAX executor for statevector components. + +``run`` interprets a component's action tape with one ``lax.fori_loop``; each +step dispatches on the action's kind with ``lax.switch``. The loop carry is a +``LaneState``: one lane per shot, each with ``2^W`` amplitudes, its symbols +(f-bits, then branch bits as they are sampled) and its output records. +All randomness goes through ``choose_outcome``, and ``count`` holds the number +of shots a lane stands for (always 1 for now), so that trajectory sharing can +later split lanes without touching the kernels. +""" + +from __future__ import annotations + +from functools import partial +from typing import NamedTuple + +import jax +import jax.numpy as jnp +from jax import Array + +from tsim.statevector.program import StatevectorComponent + +_I_POWERS = (1, 1j, -1, -1j) +_TINY = 1e-30 + + +class LaneState(NamedTuple): + """Per-lane state of the executor; the leading axis is the lane. + + Attributes: + amps: Amplitudes, shape ``(L, 2^W)``. + symbols: f-bits, then branch bits, packed into uint32 words, + shape ``(L, num_words)``. + records: Output bits, plus one scratch column for actions without a + slot, shape ``(L, num_outputs + 1)``. + count: Shots each lane stands for, shape ``(L,)``. + max_norm_deviation: Largest ``|‖ψ‖² - 1|`` seen at a measurement. + + """ + + amps: Array + symbols: Array + records: Array + count: Array # pyright: ignore[reportIncompatibleMethodOverride] + max_norm_deviation: Array + + +def choose_outcome(p1: Array, count: Array, key: Array) -> Array: + """Return each lane's outcome: True with probability ``p1``. + + ``count`` is unused for now; trajectory sharing will split it binomially. + """ + del count + return jax.random.uniform(key, p1.shape, dtype=p1.dtype) < p1 + + +def pack_bits(bits: Array, num_words: int) -> Array: + """Pack boolean columns into uint32 words: column s is bit s % 32 of word s // 32.""" + lanes, num_bits = bits.shape + padded = jnp.zeros((lanes, num_words * 32), dtype=jnp.uint32) + padded = padded.at[:, :num_bits].set(bits.astype(jnp.uint32)) + shifts = jnp.arange(32, dtype=jnp.uint32) + words = padded.reshape(lanes, num_words, 32) << shifts + return jnp.sum(words, axis=-1, dtype=jnp.uint32) + + +def initial_state(component: StatevectorComponent, f_bits: Array) -> LaneState: + """Return lanes in ``|0…0⟩`` with the given f-bits, shape ``(L, num_f)``.""" + lanes = f_bits.shape[0] + dtype = jnp.dtype(component.dtype) + amps = jnp.zeros((lanes, 2**component.width), dtype=dtype).at[:, 0].set(1) + return LaneState( + amps=amps, + symbols=pack_bits(f_bits, component.num_words), + records=jnp.zeros((lanes, component.num_outputs + 1), dtype=jnp.bool_), + count=jnp.ones(lanes, dtype=jnp.int32), + max_norm_deviation=jnp.zeros((), dtype=jnp.finfo(dtype).dtype), + ) + + +def _index(component: StatevectorComponent) -> Array: + return jnp.arange(2**component.width, dtype=jnp.uint32) + + +def _parity(values: Array) -> Array: + return (jax.lax.population_count(values) & 1).astype(jnp.bool_) + + +def _sign(component: StatevectorComponent, symbols: Array, a: Array) -> Array: + """Return each lane's sign bit of action ``a``.""" + ones = jnp.sum(jax.lax.population_count(symbols & component.sign_row[a]), axis=1) + return ((ones & 1) == 1) ^ component.sign_const[a] + + +def _apply_pauli(amps: Array, index: Array, x: Array, z: Array) -> Array: + """Return P·ψ lane by lane, for P = i^|x∧z| X^x Z^z.""" + flipped = index ^ x + quarter_turns = (jax.lax.population_count(x & z) & 3).astype(jnp.int32) + i_power = jnp.asarray(_I_POWERS, dtype=amps.dtype)[quarter_turns] + signs = jnp.where(_parity(flipped & z), -1, 1).astype(amps.dtype) + return i_power * signs * jnp.take(amps, flipped.astype(jnp.int32), axis=1) + + +def _write( + component: StatevectorComponent, state: LaneState, a: Array, outcome: Array +) -> LaneState: + """Record ``outcome ⊕ sign`` in the action's slot and set its branch bit.""" + value = outcome ^ _sign(component, state.symbols, a) + slot = component.slot[a] + column = jnp.where(slot >= 0, slot, component.num_outputs) + records = state.records.at[:, column].set(value) + branch = component.branch[a] + word = jnp.maximum(branch, 0) // 32 + bit = jnp.where( + branch >= 0, + jnp.uint32(1) << (jnp.maximum(branch, 0) % 32).astype(jnp.uint32), + jnp.uint32(0), + ) + flags = jnp.where(outcome, bit, jnp.uint32(0)) + symbols = state.symbols.at[:, word].set(state.symbols[:, word] | flags) + return state._replace(records=records, symbols=symbols) + + +def _rotate_diag( + component: StatevectorComponent, state: LaneState, a: Array, key: Array +) -> LaneState: + del key + odd = _parity(_index(component) & component.z_mask[a]) + flip = _sign(component, state.symbols, a)[:, None] ^ odd[None, :] + factor = component.cos[a] - 1j * component.sin[a] * jnp.where(flip, -1, 1) + return state._replace(amps=(state.amps * factor).astype(state.amps.dtype)) + + +def _rotate( + component: StatevectorComponent, state: LaneState, a: Array, key: Array +) -> LaneState: + del key + paulied = _apply_pauli( + state.amps, _index(component), component.x_mask[a], component.z_mask[a] + ) + negate = jnp.where(_sign(component, state.symbols, a), -1, 1)[:, None] + amps = component.cos[a] * state.amps - 1j * component.sin[a] * negate * paulied + return state._replace(amps=amps.astype(state.amps.dtype)) + + +def _measured(state: LaneState, new: Array, norm: Array) -> LaneState: + deviation = jnp.maximum(state.max_norm_deviation, jnp.max(jnp.abs(norm - 1))) + return state._replace( + amps=new.astype(state.amps.dtype), + max_norm_deviation=deviation.astype(state.max_norm_deviation.dtype), + ) + + +def _measure( + component: StatevectorComponent, state: LaneState, a: Array, key: Array +) -> LaneState: + index = _index(component) + amps = state.amps + paulied = _apply_pauli(amps, index, component.x_mask[a], component.z_mask[a]) + norm = jnp.sum(jnp.abs(amps) ** 2, axis=1) + expectation = jnp.real(jnp.sum(jnp.conj(amps) * paulied, axis=1)) + p1 = jnp.clip((norm - expectation) / (2 * norm), 0, 1) + outcome = choose_outcome(p1, state.count, jax.random.fold_in(key, a)) + p_outcome = jnp.where(outcome, p1, 1 - p1) + scale = jax.lax.rsqrt(jnp.maximum(2 * norm * p_outcome, _TINY))[:, None] + signed = jnp.where(outcome, -1, 1)[:, None] + keep = ((index >> component.pivot[a].astype(jnp.uint32)) & 1) == 0 + new = jnp.where(keep[None, :], (amps + signed * paulied) * scale, 0) + return _write(component, _measured(state, new, norm), a, outcome) + + +def _measure_diag( + component: StatevectorComponent, state: LaneState, a: Array, key: Array +) -> LaneState: + index = _index(component) + amps = state.amps + odd = _parity(index & component.z_mask[a]) + weights = jnp.abs(amps) ** 2 + norm = jnp.sum(weights, axis=1) + p1 = jnp.clip(jnp.sum(jnp.where(odd[None, :], weights, 0), axis=1) / norm, 0, 1) + outcome = choose_outcome(p1, state.count, jax.random.fold_in(key, a)) + p_outcome = jnp.where(outcome, p1, 1 - p1) + pivot = component.pivot[a].astype(jnp.uint32) + bit = jnp.uint32(1) << pivot + source_if_0 = (index ^ (bit * odd.astype(jnp.uint32))).astype(jnp.int32) + source_if_1 = (index ^ (bit * (~odd).astype(jnp.uint32))).astype(jnp.int32) + gathered = jnp.where( + outcome[:, None], + jnp.take(amps, source_if_1, axis=1), + jnp.take(amps, source_if_0, axis=1), + ) + keep = ((index >> pivot) & 1) == 0 + scale = jax.lax.rsqrt(jnp.maximum(norm * p_outcome, _TINY))[:, None] + new = jnp.where(keep[None, :], gathered * scale, 0) + return _write(component, _measured(state, new, norm), a, outcome) + + +def _measure_dormant( + component: StatevectorComponent, state: LaneState, a: Array, key: Array +) -> LaneState: + half = jnp.full(state.count.shape, 0.5, dtype=state.max_norm_deviation.dtype) + outcome = choose_outcome(half, state.count, jax.random.fold_in(key, a)) + return _write(component, state, a, outcome) + + +def _record( + component: StatevectorComponent, state: LaneState, a: Array, key: Array +) -> LaneState: + del key + return _write(component, state, a, jnp.zeros(state.count.shape, dtype=jnp.bool_)) + + +# NOTE: the order must match the kind codes ROTATE_DIAG … RECORD in program.py. +_KERNELS = (_rotate_diag, _rotate, _measure_diag, _measure, _measure_dormant, _record) + + +def run( + component: StatevectorComponent, + state: LaneState, + key: Array, + start: int | Array, + stop: int | Array, +) -> LaneState: + """Run actions ``start … stop-1`` of the tape on every lane.""" + branches = [partial(kernel, component) for kernel in _KERNELS] + + def step(a: Array, carry: LaneState) -> LaneState: + return jax.lax.switch(component.kind[a], branches, carry, a, key) + + return jax.lax.fori_loop(start, stop, step, state) + + +@jax.jit +def _run_lanes( + component: StatevectorComponent, f_bits: Array, key: Array +) -> tuple[Array, Array]: + state = run( + component, initial_state(component, f_bits), key, 0, component.num_actions + ) + return state.records[:, : component.num_outputs], state.max_norm_deviation + + +def sample_statevector_component( + component: StatevectorComponent, f_params: Array, key: Array +) -> tuple[Array, Array, Array]: + """Sample one component's outputs; the contract of ``sampler.sample_component``. + + Args: + component: The compiled component. + f_params: Error-basis bits of every shot, shape ``(B, num_f_global)``. + key: JAX random key. + + Returns: + ``(samples, next_key, max_norm_deviation)`` with ``samples`` of shape + ``(B, num_outputs)``. + + """ + f_bits = f_params[:, component.f_selection].astype(jnp.bool_) + key, subkey = jax.random.split(key) + samples, deviation = _run_lanes(component, f_bits, subkey) + return samples, key, deviation diff --git a/test/helpers/statevector.py b/test/helpers/statevector.py index 71912702..56d46345 100644 --- a/test/helpers/statevector.py +++ b/test/helpers/statevector.py @@ -465,3 +465,39 @@ def stabilizer_rank_distribution( f_row = np.array([list(f)], dtype=np.uint8).reshape(1, len(f)) norm = float(np.abs(np.asarray(evaluate(norm_graphs, jnp.asarray(f_row))))[0]) return {o: float(p) / norm for o, p in zip(outcomes, joint, strict=True)} + + +_PAULI_MATRICES = { + (0, 0): np.eye(2), + (1, 0): np.array([[0, 1], [1, 0]]), + (0, 1): np.diag([1, -1]), + (1, 1): np.array([[0, -1j], [1j, 0]]), +} + + +def pauli_matrix(x: int, z: int, num_qubits: int) -> np.ndarray: + """Return the dense matrix of P = i^|x∧z| X^x Z^z (bit q is qubit q).""" + matrix = np.ones((1, 1), dtype=complex) + for q in reversed(range(num_qubits)): + matrix = np.kron(matrix, _PAULI_MATRICES[((x >> q) & 1, (z >> q) & 1)]) + return matrix + + +def assert_matches_distribution( + samples: np.ndarray, distribution: Distribution +) -> None: + """Assert that sampled rows follow ``distribution`` (5σ + 3 per outcome).""" + shots = samples.shape[0] + counts: dict[tuple[int, ...], int] = {} + for row in samples.astype(int): + key = tuple(row.tolist()) + counts[key] = counts.get(key, 0) + 1 + for key in set(counts) | set(distribution): + p = distribution.get(key, 0.0) + expected = shots * p + sigma = np.sqrt(shots * p * (1 - p)) + assert abs(counts.get(key, 0) - expected) <= 5 * sigma + 3, ( + key, + counts.get(key, 0), + expected, + ) diff --git a/test/unit/statevector/test_execute.py b/test/unit/statevector/test_execute.py new file mode 100644 index 00000000..d4ac93bd --- /dev/null +++ b/test/unit/statevector/test_execute.py @@ -0,0 +1,167 @@ +from test.helpers.statevector import ( + CIRCUITS, + assert_matches_distribution, + f_patterns, + pauli_matrix, + plan_distribution, + quantum_components, + reference_measure, +) + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +import tsim.statevector.execute as execute +from tsim.statevector.execute import ( + initial_state, + pack_bits, + run, + sample_statevector_component, +) +from tsim.statevector.heisenberg import ComponentLayout +from tsim.statevector.plan import ( + AffineRow, + MeasureActive, + MeasureDormant, + Plan, + Record, + Rotate, +) +from tsim.statevector.program import lower, plan_component, resolve_dtype + + +def _component( + actions, *, width, num_f=1, num_branches=0, num_outputs=0, dtype="complex64" +): + planned = Plan( + width=width, + num_f=num_f, + num_branches=num_branches, + num_outputs=num_outputs, + actions=tuple(actions), + ) + layout = ComponentLayout(tuple(range(num_f)), {}, tuple(range(num_outputs))) + return lower(planned, layout, dtype) + + +def _random_states(lanes, width, seed=0): + rng = np.random.default_rng(seed) + psi = rng.normal(size=(lanes, 2**width)) + 1j * rng.normal(size=(lanes, 2**width)) + return psi / np.linalg.norm(psi, axis=1, keepdims=True) + + +def test_pack_bits(): + bits = np.zeros((2, 40), dtype=bool) + bits[0, [0, 31, 32, 39]] = True + bits[1, 5] = True + words = np.asarray(pack_bits(jnp.asarray(bits), 2)) + assert words.dtype == np.uint32 + assert words.tolist() == [[1 | (1 << 31), 1 | (1 << 7)], [1 << 5, 0]] + + +@pytest.mark.parametrize("dtype", ["complex64", "complex128"]) +@pytest.mark.parametrize("x, z", [(0, 0b101), (0b011, 0), (0b110, 0b011), (1, 1)]) +def test_rotations_match_dense_matrices(x, z, dtype): + with jax.enable_x64(dtype == "complex128"): + sign = AffineRow(False, frozenset({0})) + component = _component([Rotate(x, z, 0.37, sign)], width=3, dtype=dtype) + psi = _random_states(2, 3) + f_bits = jnp.array([[False], [True]]) + amps = jnp.asarray(psi, jnp.dtype(dtype)) + state = initial_state(component, f_bits)._replace(amps=amps) + out = run(component, state, jax.random.key(0), 0, 1) + assert out.amps.dtype == jnp.dtype(dtype) + pauli = pauli_matrix(x, z, 3) + atol = 1e-5 if dtype == "complex64" else 1e-12 + for lane, negate in ((0, 1), (1, -1)): + rotation = np.cos(0.37) * np.eye(8) - 1j * np.sin(0.37) * negate * pauli + np.testing.assert_allclose( + np.asarray(out.amps[lane]), rotation @ psi[lane], atol=atol + ) + + +@pytest.mark.parametrize( + "x, z, pivot", [(0b011, 0b110, 0), (0b010, 0b011, 1), (0, 0b101, 2), (0, 0b011, 0)] +) +def test_measurements_match_the_reference_kernel(monkeypatch, x, z, pivot): + action = MeasureActive( + x, z, pivot, branch=1, sign=AffineRow(False, frozenset({0})), slot=0 + ) + component = _component([action], width=3, num_branches=1, num_outputs=1) + outcomes = jnp.array([False, True, True]) + monkeypatch.setattr(execute, "choose_outcome", lambda p1, count, key: outcomes) + psi = _random_states(3, 3) + f_bits = jnp.array([[False], [False], [True]]) + state = initial_state(component, f_bits)._replace( + amps=jnp.asarray(psi, jnp.complex64) + ) + out = run(component, state, jax.random.key(0), 0, 1) + for lane, outcome in enumerate([0, 1, 1]): + expected, _ = reference_measure(psi[lane], x, z, pivot, outcome) + np.testing.assert_allclose(np.asarray(out.amps[lane]), expected, atol=1e-5) + assert np.asarray(out.records[:, 0]).tolist() == [False, True, False] + assert ((np.asarray(out.symbols[:, 0]) >> 1) & 1).tolist() == [0, 1, 1] + + +def test_dormant_measurements_and_records_write_outcome_xor_sign(monkeypatch): + component = _component( + [ + MeasureDormant(branch=1, sign=AffineRow(True, frozenset()), slot=0), + Record(sign=AffineRow(False, frozenset({0, 1})), slot=1), + ], + width=0, + num_branches=1, + num_outputs=2, + ) + monkeypatch.setattr( + execute, "choose_outcome", lambda p1, count, key: jnp.array([False, True]) + ) + state = initial_state(component, jnp.array([[True], [True]])) + out = run(component, state, jax.random.key(0), 0, 2) + assert np.asarray(out.records[:, :2]).tolist() == [[True, True], [False, False]] + + +@pytest.mark.parametrize( + "name", + [ + "t_ghz", + "midcircuit_feedforward", + "two_components", + "clifford_component", + "rotations", + ], +) +def test_sampling_matches_the_plan_distribution(name): + shots = 20_000 + components, layout = quantum_components(CIRCUITS[name], sample_detectors=True) + for component in components: + planned, component_layout = plan_component(component, layout) + compiled = lower(planned, component_layout, "complex64") + num_columns = max(component_layout.f_selection, default=-1) + 1 + for f in f_patterns(planned.num_f, limit=2): + f_params = np.zeros((shots, num_columns), dtype=np.uint8) + if component_layout.f_selection: + f_params[:, list(component_layout.f_selection)] = f + samples, _, deviation = sample_statevector_component( + compiled, jnp.asarray(f_params), jax.random.key(7) + ) + assert float(deviation) < 1e-4 + assert_matches_distribution( + np.asarray(samples), plan_distribution(planned, f) + ) + + +def test_complex128_amplitudes(): + with jax.enable_x64(True): + (component,), layout = quantum_components( + CIRCUITS["t_ghz"], sample_detectors=True + ) + planned, component_layout = plan_component(component, layout) + compiled = lower(planned, component_layout, resolve_dtype(None)) + assert compiled.dtype == "complex128" + state = initial_state(compiled, jnp.zeros((4, planned.num_f), dtype=bool)) + assert state.amps.dtype == jnp.complex128 + out = run(compiled, state, jax.random.key(0), 0, compiled.num_actions) + assert out.amps.dtype == jnp.complex128 From a0804e5223c420fe27c2bf09a0b22a8359fc0b51 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Sat, 10 Oct 2026 17:06:59 -0400 Subject: [PATCH 06/18] Return the state unchanged for statevector tapes without actions Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/execute.py | 7 ++++++- test/unit/statevector/test_execute.py | 14 ++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/src/tsim/statevector/execute.py b/src/tsim/statevector/execute.py index 98ac013a..cac4369a 100644 --- a/src/tsim/statevector/execute.py +++ b/src/tsim/statevector/execute.py @@ -221,7 +221,12 @@ def run( start: int | Array, stop: int | Array, ) -> LaneState: - """Run actions ``start … stop-1`` of the tape on every lane.""" + """Run actions ``start … stop-1`` of the tape on every lane. + + A component without actions returns ``state`` unchanged. + """ + if component.num_actions == 0: + return state branches = [partial(kernel, component) for kernel in _KERNELS] def step(a: Array, carry: LaneState) -> LaneState: diff --git a/test/unit/statevector/test_execute.py b/test/unit/statevector/test_execute.py index d4ac93bd..a7d79c4a 100644 --- a/test/unit/statevector/test_execute.py +++ b/test/unit/statevector/test_execute.py @@ -123,6 +123,20 @@ def test_dormant_measurements_and_records_write_outcome_xor_sign(monkeypatch): assert np.asarray(out.records[:, :2]).tolist() == [[True, True], [False, False]] +def test_a_component_without_actions_samples_no_outputs(): + component = _component([], width=0, num_f=0) + assert component.num_actions == 0 + state = initial_state(component, jnp.zeros((3, 0), dtype=bool)) + out = run(component, state, jax.random.key(0), 0, component.num_actions) + for got, initial in zip(out, state, strict=True): + np.testing.assert_array_equal(np.asarray(got), np.asarray(initial)) + samples, _, deviation = sample_statevector_component( + component, jnp.zeros((3, 0), dtype=jnp.uint8), jax.random.key(0) + ) + assert samples.shape == (3, 0) + assert float(deviation) == 0.0 + + @pytest.mark.parametrize( "name", [ From 3e9bf3a518cf8d2c6b301b0370972414d66e3d63 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Sat, 10 Oct 2026 17:21:51 -0400 Subject: [PATCH 07/18] Add strategy="statevector" to the samplers Co-Authored-By: Claude Opus 5.5 --- src/tsim/__init__.py | 1 + src/tsim/circuit.py | 17 +-- src/tsim/compile/pipeline.py | 44 +++++-- src/tsim/core/types.py | 12 +- src/tsim/sampler.py | 88 ++++++++++--- test/helpers/statevector.py | 15 +++ test/integration/test_sampler.py | 35 ++++-- test/integration/test_statevector_sampler.py | 123 +++++++++++++++++++ test/unit/compile/test_pipeline.py | 2 + test/unit/statevector/test_program.py | 2 + 10 files changed, 287 insertions(+), 52 deletions(-) create mode 100644 test/integration/test_statevector_sampler.py diff --git a/src/tsim/__init__.py b/src/tsim/__init__.py index 912fac71..04820977 100644 --- a/src/tsim/__init__.py +++ b/src/tsim/__init__.py @@ -10,6 +10,7 @@ - `sampler.py`: Orchestrates compilation and evaluation to sample measurements or detectors. - `core/`: Builds quantum/classical ZX graphs with zxparam; parsing and data types. - `compile/`: Compiles zxparam's stabilizer-rank decompositions to JAX. +- `statevector/`: Samples non-direct components as extracted circuits on a dense state vector (`strategy="statevector"`). - `noise/`: Pauli noise channels and detector error models. - `utils/`: Utility functions for visualization and linear algebra. """ diff --git a/src/tsim/circuit.py b/src/tsim/circuit.py index 927271e1..22c63d7f 100644 --- a/src/tsim/circuit.py +++ b/src/tsim/circuit.py @@ -16,7 +16,7 @@ import numpy as np import stim import zxparam as zx -from zxparam import DecompositionStrategy, Graph +from zxparam import Graph from zxparam.phase import Bit if TYPE_CHECKING: @@ -25,6 +25,7 @@ from tsim.core.graph import transform_error_basis from tsim.core.parse import parse_parametric_tag, parse_stim_circuit from tsim.core.tags import encode_t_tag +from tsim.core.types import SamplingStrategy from tsim.noise.dem import get_detector_error_model from tsim.utils.clifford import expand_clifford_rotations, is_clifford from tsim.utils.diagram import draw_zx, render_svg @@ -862,14 +863,15 @@ def get_sampling_graph(self, sample_detectors: bool = False) -> Graph: def compile_sampler( self, *, - strategy: DecompositionStrategy = "cat5", + strategy: SamplingStrategy = "cat5", seed: int | None = None, ) -> CompiledMeasurementSampler: """Compile circuit into a measurement sampler. Args: - strategy: Stabilizer rank decomposition strategy. - Must be one of "cat5", "bss", "cutting". + strategy: Stabilizer rank decomposition strategy ("cat5", "bss", + "cutting"), or "statevector" to simulate each non-direct + component as an extracted circuit on a dense state vector. seed: Random seed for the sampler. IMPORTANT: Currently, the sampler will only produce deterministic samples for fixed batch size. If deterministic samples are needed, the batch size should be set @@ -886,7 +888,7 @@ def compile_sampler( def compile_detector_sampler( self, *, - strategy: DecompositionStrategy = "cat5", + strategy: SamplingStrategy = "cat5", seed: int | None = None, ) -> CompiledDetectorSampler: """Compile circuit into a detector sampler. @@ -897,8 +899,9 @@ def compile_detector_sampler( compilation pipeline. Args: - strategy: Stabilizer rank decomposition strategy. - Must be one of "cat5", "bss", "cutting". + strategy: Stabilizer rank decomposition strategy ("cat5", "bss", + "cutting"), or "statevector" to simulate each non-direct + component as an extracted circuit on a dense state vector. seed: Random seed for the sampler. IMPORTANT: Currently, the sampler will only produce deterministic samples for fixed batch size. If deterministic samples are needed, the batch size should be set diff --git a/src/tsim/compile/pipeline.py b/src/tsim/compile/pipeline.py index 14daf2e2..bb8b39d8 100644 --- a/src/tsim/compile/pipeline.py +++ b/src/tsim/compile/pipeline.py @@ -23,7 +23,16 @@ from tsim.compile.compile import CompiledScalarGraphs, compile_scalars from tsim.core.params import ParamLayout -from tsim.core.types import CompiledComponent, CompiledProgram, SamplingGraph +from tsim.core.types import ( + CompiledComponent, + CompiledProgram, + SamplingGraph, + SamplingStrategy, +) +from tsim.statevector.program import ( + StatevectorComponent, + compile_statevector_component, +) DecompositionMode = Literal["sequential", "joint"] @@ -32,30 +41,40 @@ def compile_program( prepared: SamplingGraph, *, mode: DecompositionMode, - strategy: DecompositionStrategy = "cat5", + strategy: SamplingStrategy = "cat5", ) -> CompiledProgram: """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. With + ``strategy="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 + is decomposed into stabilizer terms, and the terms are compiled into JAX + arrays. Args: prepared: The prepared graph from ``prepare_graph``. mode: ``"sequential"`` for sampling, ``"joint"`` for probabilities. - strategy: Stabilizer-rank decomposition strategy. + strategy: A stabilizer-rank decomposition strategy, or ``"statevector"``. Returns: A CompiledProgram ready for sampling. + Raises: + ValueError: If ``strategy="statevector"`` is combined with ``mode="joint"``. + """ + if strategy == "statevector" and mode == "joint": + raise ValueError( + "strategy='statevector' samples outcomes and cannot compute " + "probabilities; use a stabilizer-rank strategy such as 'cat5'" + ) graph = prepared.graph layout = prepared.layout direct_entries: list[tuple[int, int, bool]] = [] - compiled_components: list[CompiledComponent] = [] + compiled_components: list[CompiledComponent | StatevectorComponent] = [] compiled_output_order: list[int] = [] components = sorted(graph.connected_components(), key=lambda c: len(c.outputs())) @@ -66,9 +85,12 @@ def compile_program( basis_index, flip = direct direct_entries.append((output_indices[0], basis_index, flip)) continue - compiled_components.append( - _compile_component(component, output_indices, layout, mode, strategy) - ) + if strategy == "statevector": + compiled_components.append(compile_statevector_component(component, layout)) + else: + compiled_components.append( + _compile_component(component, output_indices, layout, mode, strategy) + ) compiled_output_order.extend(output_indices) # Sort direct entries by output index so that — together with the output diff --git a/src/tsim/core/types.py b/src/tsim/core/types.py index aa34bcb7..e92877c3 100644 --- a/src/tsim/core/types.py +++ b/src/tsim/core/types.py @@ -11,17 +11,22 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Literal import equinox as eqx import numpy as np from jax import Array +from zxparam import DecompositionStrategy if TYPE_CHECKING: from zxparam import Graph from tsim.compile.compile import CompiledScalarGraphs from tsim.core.params import ParamLayout + from tsim.statevector.program import StatevectorComponent + +SamplingStrategy = DecompositionStrategy | Literal["statevector"] +"""A stabilizer-rank decomposition strategy, or ``"statevector"``.""" @dataclass(frozen=True) @@ -89,7 +94,8 @@ class CompiledProgram: needed to sample from the circuit. Attributes: - components: The compiled components, sorted by number of outputs. + 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_flips: Precomputed flip flags for direct components. output_order: Maps concatenated position to original output index. @@ -102,7 +108,7 @@ class CompiledProgram: """ - components: tuple[CompiledComponent, ...] + components: tuple[CompiledComponent | StatevectorComponent, ...] direct_f_indices: Array direct_flips: Array output_order: Array diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index 514edc96..4a885830 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -15,8 +15,10 @@ from tsim.compile.evaluate import evaluate from tsim.compile.pipeline import compile_program from tsim.core.graph import prepare_graph -from tsim.core.types import CompiledComponent, CompiledProgram +from tsim.core.types import CompiledComponent, CompiledProgram, SamplingStrategy from tsim.noise.channels import ChannelSampler +from tsim.statevector.execute import sample_statevector_component +from tsim.statevector.program import StatevectorComponent from tsim.utils.cuda_helpers import copy_d2h if TYPE_CHECKING: @@ -145,7 +147,14 @@ def sample_program( results.append(direct_bits) for component in program.components: - samples, key, max_norm_deviation = sample_component(component, f_params, key) + if isinstance(component, StatevectorComponent): + samples, key, max_norm_deviation = sample_statevector_component( + component, f_params, key + ) + else: + samples, key, max_norm_deviation = sample_component( + component, f_params, key + ) if np.isclose(max_norm_deviation, 1): raise ValueError( "A vanishing marginal probability distribution was encountered (normalization 0). " @@ -176,7 +185,7 @@ def __init__( *, sample_detectors: bool, mode: Literal["sequential", "joint"], - strategy: DecompositionStrategy = "cat5", + strategy: SamplingStrategy = "cat5", seed: int | None = None, ): """Initialize the sampler by compiling the circuit. @@ -185,8 +194,9 @@ def __init__( circuit: The quantum circuit to compile. sample_detectors: If True, sample detectors/observables instead of measurements. mode: Compilation mode - "sequential" for autoregressive, "joint" for probabilities. - strategy: Stabilizer rank decomposition strategy. - Must be one of "cat5", "bss", "cutting". + strategy: Stabilizer rank decomposition strategy ("cat5", "bss", + "cutting"), or "statevector" to simulate each non-direct + component as an extracted circuit on a dense state vector. seed: Random seed. If None, a random seed is generated. Note that deterministic results are only guaranteed for a fixed batch size and fixed reference sample settings. @@ -235,6 +245,18 @@ def __init__( : self._num_detectors ].copy() + statevector = [ + c for c in self._program.components if isinstance(c, StatevectorComponent) + ] + if statevector and self._peak_bytes_per_sample() > self._available_memory(): + widest = max(statevector, key=lambda c: c.width) + raise ValueError( + f"strategy='statevector' needs {self._peak_bytes_per_sample()} bytes " + f"per shot for the component with outputs {list(widest.output_indices)} " + f"(width {widest.width}), more than the available memory. Use a " + "stabilizer-rank strategy such as strategy='cat5'." + ) + def _compute_direct_outputs(self, f_params_np: np.ndarray) -> np.ndarray: """Scatter direct output bits into a full (batch, num_outputs) bool array. @@ -295,6 +317,12 @@ def _peak_bytes_per_sample(self) -> int: """Estimate peak device memory per sample from compiled program structure.""" peak = 0 for component in self._program.components: + if isinstance(component, StatevectorComponent): + itemsize = 8 if component.dtype == "complex64" else 16 + lane = 3 * 2**component.width * itemsize + lane += 4 * component.num_words + component.num_outputs + 1 + peak = max(peak, lane) + continue for circuit in component.compiled_scalar_graphs: G = circuit.num_graphs max_a = circuit.node_phases.phases.shape[1] @@ -306,18 +334,18 @@ def _peak_bytes_per_sample(self) -> int: peak = max(peak, G * largest * 3) return max(peak, 1) - def _estimate_batch_size(self) -> int: - """Estimate the largest batch size that fits in available device memory.""" + def _available_memory(self) -> int: + """Return the free memory of the first JAX device, or of the host.""" device = jax.devices()[0] if device.platform == "gpu": stats = device.memory_stats() - available = stats.get("bytes_limit", 8 * 1024**3) - stats.get( - "bytes_in_use", 0 - ) - else: - available = psutil.virtual_memory().available + limit = stats.get("bytes_limit", 8 * 1024**3) + return limit - stats.get("bytes_in_use", 0) + return psutil.virtual_memory().available - half_of_available = int(available * 0.5) # conservative estimate + def _estimate_batch_size(self) -> int: + """Estimate the largest batch size that fits in available device memory.""" + half_of_available = int(self._available_memory() * 0.5) # conservative return max(1, half_of_available // self._peak_bytes_per_sample()) @overload @@ -568,8 +596,20 @@ def __repr__(self) -> str: c_e_terms = [] total_memory_bytes = 0 num_outputs = [] + statevector_widths: list[int] = [] + statevector_actions: list[int] = [] for component in self._program.components: + if isinstance(component, StatevectorComponent): + statevector_widths.append(component.width) + statevector_actions.append(component.num_actions) + num_outputs.append(len(component.output_indices)) + total_memory_bytes += sum( + v.nbytes + for v in jax.tree_util.tree_leaves(component) + if isinstance(v, jax.Array) + ) + continue for circuit in component.compiled_scalar_graphs: num_outputs.append(len(component.output_indices)) c_graphs.append(circuit.num_graphs) @@ -606,8 +646,17 @@ def _format_bytes(n: int) -> str: channel.num_bits for channel in self._channel_sampler.channels ) + statevector = ( + f"{len(statevector_widths)} statevector components " + f"(width ≤ {max(statevector_widths)}, " + f"{sum(statevector_actions)} actions), " + if statevector_widths + else "" + ) + return ( f"{type(self).__name__}({n_direct} direct, " + f"{statevector}" 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, " @@ -631,15 +680,15 @@ def __init__( self, circuit: Circuit, *, - strategy: DecompositionStrategy = "cat5", + strategy: SamplingStrategy = "cat5", seed: int | None = None, ): """Create a measurement sampler. Args: circuit: The quantum circuit to compile. - strategy: Stabilizer rank decomposition strategy. - Must be one of "cat5", "bss", "cutting". + strategy: Stabilizer rank decomposition strategy ("cat5", "bss", + "cutting"), or "statevector". seed: Random seed for the sampler. IMPORTANT: Currently, the sampler will only produce deterministic samples for fixed batch size. If deterministic samples are needed, the batch size should be set @@ -686,15 +735,15 @@ def __init__( self, circuit: Circuit, *, - strategy: DecompositionStrategy = "cat5", + strategy: SamplingStrategy = "cat5", seed: int | None = None, ): """Create a detector sampler. Args: circuit: The quantum circuit to compile. - strategy: Stabilizer rank decomposition strategy. - Must be one of "cat5", "bss", "cutting". + strategy: Stabilizer rank decomposition strategy ("cat5", "bss", + "cutting"), or "statevector". seed: Random seed for the sampler. IMPORTANT: Currently, the sampler will only produce deterministic samples for fixed batch size and fixed reference sample settings. If deterministic samples are @@ -945,6 +994,7 @@ def probability_of(self, state: np.ndarray, *, batch_size: int) -> np.ndarray: p_joint = p_joint * (direct_bits == targets).all(axis=1) for component in self._program.components: + assert isinstance(component, CompiledComponent) assert len(component.compiled_scalar_graphs) == 2 f_selected = f_samples[:, component.f_selection] diff --git a/test/helpers/statevector.py b/test/helpers/statevector.py index 56d46345..80fb3366 100644 --- a/test/helpers/statevector.py +++ b/test/helpers/statevector.py @@ -501,3 +501,18 @@ def assert_matches_distribution( counts.get(key, 0), expected, ) + + +def assert_counts_match(samples: np.ndarray, reference: np.ndarray) -> None: + """Assert that two sets of sampled rows have the same row distribution.""" + width = samples.shape[1] + powers = 2 ** np.arange(width) + bins = 2**width + a = np.bincount(samples.astype(np.int64) @ powers, minlength=bins) + b = np.bincount(reference.astype(np.int64) @ powers, minlength=bins) + scale_a, scale_b = 1 / samples.shape[0], 1 / reference.shape[0] + pooled = (a + b) / (samples.shape[0] + reference.shape[0]) + sigma = np.sqrt(pooled * (1 - pooled) * (scale_a + scale_b)) + assert np.all( + np.abs(a * scale_a - b * scale_b) <= 5 * sigma + 3 * (scale_a + scale_b) + ) diff --git a/test/integration/test_sampler.py b/test/integration/test_sampler.py index 13c51a7e..d2199f31 100644 --- a/test/integration/test_sampler.py +++ b/test/integration/test_sampler.py @@ -12,7 +12,7 @@ import tsim.sampler as sampler_module from tsim.circuit import Circuit -from tsim.core.types import CompiledComponent, CompiledProgram +from tsim.core.types import CompiledComponent, CompiledProgram, SamplingStrategy from tsim.external.vec_sim.vec_sampler import VecSampler from tsim.sampler import CompiledStateProbs @@ -63,10 +63,8 @@ def assert_samples_match(samples1: np.ndarray, samples2: np.ndarray): "color_code:memory_xyz", ], ) -@pytest.mark.parametrize("strategy", ["cat5", "bss", "cutting"]) -def test_quantum_memory_codes_without_noise( - code_task: str, strategy: DecompositionStrategy -): +@pytest.mark.parametrize("strategy", ["cat5", "bss", "cutting", "statevector"]) +def test_quantum_memory_codes_without_noise(code_task: str, strategy: SamplingStrategy): circ = stim.Circuit.generated( code_task, @@ -84,7 +82,7 @@ def test_quantum_memory_codes_without_noise( @pytest.mark.parametrize("seed", [1, 2, 42]) -@pytest.mark.parametrize("strategy", ["cat5", "bss", "cutting"]) +@pytest.mark.parametrize("strategy", ["cat5", "bss", "cutting", "statevector"]) def test_sampler(seed, strategy): num_qubits = 3 stim_circuit = gen_stim_circuit( @@ -218,7 +216,10 @@ def test_memory_error_correction_and_compare_to_stim(code_task: str): "before_round_data_depolarization", ], ) -def test_rotated_surface_code_single_noise_channel_matches_stim(channel: str): +@pytest.mark.parametrize("strategy", ["cat5", "statevector"]) +def test_rotated_surface_code_single_noise_channel_matches_stim( + channel: str, strategy: SamplingStrategy +): """Regression test for single-channel detector sampling noise handling.""" noise_kwargs = { "after_clifford_depolarization": 0.0, @@ -241,7 +242,7 @@ def test_rotated_surface_code_single_noise_channel_matches_stim(channel: str): ) tsim_dets = ( Circuit.from_stim_program(circ) - .compile_detector_sampler(seed=42) + .compile_detector_sampler(seed=42, strategy=strategy) .sample( shots=shots, batch_size=shots // 10, @@ -293,10 +294,15 @@ def test_rotated_surface_code_single_noise_channel_matches_stim(channel: str): "two_correlated_errors_with_else", ], ) -def test_correlated_error(stim_program: str): +@pytest.mark.parametrize("strategy", ["cat5", "statevector"]) +def test_correlated_error(stim_program: str, strategy: SamplingStrategy): n_samples = 10_000 stim_samples = stim.Circuit(stim_program).compile_sampler(seed=0).sample(n_samples) - tsim_samples = Circuit(stim_program).compile_sampler(seed=0).sample(n_samples) + tsim_samples = ( + Circuit(stim_program) + .compile_sampler(seed=0, strategy=strategy) + .sample(n_samples) + ) assert_samples_match(stim_samples, tsim_samples) @@ -375,11 +381,16 @@ def test_correlated_error(stim_program: str): "multiple_products", ], ) -def test_mpp_inversion_parity(stim_program: str): +@pytest.mark.parametrize("strategy", ["cat5", "statevector"]) +def test_mpp_inversion_parity(stim_program: str, strategy: SamplingStrategy): """Test that MPP correctly handles inversions on any Pauli (only parity matters).""" n_samples = 10_000 stim_samples = stim.Circuit(stim_program).compile_sampler(seed=0).sample(n_samples) - tsim_samples = Circuit(stim_program).compile_sampler(seed=0).sample(n_samples) + tsim_samples = ( + Circuit(stim_program) + .compile_sampler(seed=0, strategy=strategy) + .sample(n_samples) + ) assert_samples_match(stim_samples, tsim_samples) diff --git a/test/integration/test_statevector_sampler.py b/test/integration/test_statevector_sampler.py new file mode 100644 index 00000000..e88d1bf4 --- /dev/null +++ b/test/integration/test_statevector_sampler.py @@ -0,0 +1,123 @@ +"""strategy="statevector" end to end, against stabilizer rank and VecSampler.""" + +from test.helpers.gen import gen_stim_circuit +from test.helpers.statevector import CIRCUITS, assert_counts_match + +import jax +import numpy as np +import pytest + +import tsim.sampler as sampler_module +from tsim.circuit import Circuit +from tsim.compile.pipeline import compile_program +from tsim.core.graph import prepare_graph +from tsim.core.types import SamplingStrategy +from tsim.external.vec_sim.vec_sampler import VecSampler +from tsim.statevector.program import StatevectorComponent + + +@pytest.mark.parametrize("name", sorted(CIRCUITS)) +def test_detector_sampling_matches_stabilizer_rank(name): + circuit = Circuit(CIRCUITS[name]) + shots = 20_000 + expected = circuit.compile_detector_sampler(strategy="cat5", seed=1).sample( + shots, append_observables=True + ) + got = circuit.compile_detector_sampler(strategy="statevector", seed=2).sample( + shots, batch_size=5_000, append_observables=True + ) + assert_counts_match(got, expected) + + +@pytest.mark.parametrize("name", sorted(CIRCUITS)) +def test_measurement_sampling_matches_stabilizer_rank(name): + circuit = Circuit(CIRCUITS[name]) + shots = 20_000 + expected = circuit.compile_sampler(strategy="cat5", seed=1).sample(shots) + got = circuit.compile_sampler(strategy="statevector", seed=2).sample( + shots, batch_size=5_000 + ) + assert_counts_match(got, expected) + + +@pytest.mark.parametrize("seed", [1, 2, 3]) +def test_measurement_sampling_matches_vec_sampler(seed): + stim_circuit = gen_stim_circuit(3, 60, p_t=3, p_x_error=0.5, p_r_z=0.5, seed=seed) + shots = 512 + expected, _, _ = VecSampler(stim_circuit).sample(shots) + sampler = Circuit.from_stim_program(stim_circuit).compile_sampler( + strategy="statevector", seed=seed + ) + assert_counts_match(sampler.sample(shots, batch_size=shots), expected) + + +def test_statevector_samplers_hold_statevector_components(): + sampler = Circuit(CIRCUITS["t_ghz"]).compile_detector_sampler( + strategy="statevector" + ) + components = sampler._program.components + assert components + assert all(isinstance(c, StatevectorComponent) for c in components) + assert "statevector components" in repr(sampler) + widest = max(c.width for c in components if isinstance(c, StatevectorComponent)) + assert sampler._peak_bytes_per_sample() >= 3 * 2**widest * 8 + + +def test_postselection_and_partial_batches(): + circuit = Circuit(CIRCUITS["two_components"]) + mask = np.ones(circuit.num_detectors, dtype=bool) + shots = 30_000 + rates = [] + cases: tuple[tuple[SamplingStrategy, int], ...] = (("cat5", 1), ("statevector", 2)) + for strategy, seed in cases: + sampler = circuit.compile_detector_sampler(strategy=strategy, seed=seed) + detectors, observables = sampler.sample( + shots, batch_size=7_000, separate_observables=True, postselection_mask=mask + ) + assert detectors.shape == (shots, circuit.num_detectors) + assert observables.shape == (shots, circuit.num_observables) + keep = ~detectors.any(axis=1) + rates.append((keep.mean(), observables[keep, 0].mean())) + (accept_a, error_a), (accept_b, error_b) = rates + accept_sigma = np.sqrt(accept_a * (1 - accept_a) * 2 / shots) + assert abs(accept_a - accept_b) <= 5 * accept_sigma + 1e-3 + error_sigma = np.sqrt(max(error_a * (1 - error_a), 1e-4) * 2 / (accept_a * shots)) + assert abs(error_a - error_b) <= 5 * error_sigma + 1e-3 + + +def test_reference_samples_work_with_statevector(): + circuit = Circuit(CIRCUITS["injection_check"]) + sampler = circuit.compile_detector_sampler(strategy="statevector", seed=0) + samples = sampler.sample( + 1_000, use_detector_reference_sample=True, use_observable_reference_sample=True + ) + assert samples.shape == (1_000, circuit.num_detectors) + + +def test_probabilities_reject_statevector(): + prepared = prepare_graph(Circuit(CIRCUITS["t_ghz"]), sample_detectors=False) + with pytest.raises(ValueError, match="statevector"): + compile_program(prepared, mode="joint", strategy="statevector") + + +def test_complex128_sampling_matches_stabilizer_rank(): + circuit = Circuit(CIRCUITS["t_ghz"]) + expected = circuit.compile_detector_sampler(strategy="cat5", seed=1).sample( + 10_000, append_observables=True + ) + with jax.enable_x64(True): + sampler = circuit.compile_detector_sampler(strategy="statevector", seed=2) + assert all( + isinstance(c, StatevectorComponent) and c.dtype == "complex128" + for c in sampler._program.components + ) + got = sampler.sample(10_000, append_observables=True) + assert_counts_match(got, expected) + + +def test_a_lane_that_does_not_fit_in_memory_is_rejected(monkeypatch): + monkeypatch.setattr( + sampler_module._CompiledSamplerBase, "_available_memory", lambda self: 1 + ) + with pytest.raises(ValueError, match="cat5"): + Circuit(CIRCUITS["t_ghz"]).compile_detector_sampler(strategy="statevector") diff --git a/test/unit/compile/test_pipeline.py b/test/unit/compile/test_pipeline.py index b5b1867a..dd644a70 100644 --- a/test/unit/compile/test_pipeline.py +++ b/test/unit/compile/test_pipeline.py @@ -14,6 +14,7 @@ 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 from tsim.sampler import CompiledStateProbs I2 = np.eye(2) @@ -31,6 +32,7 @@ def test_sequential_plugged_graphs_are_consistent_marginals(): ) assert any(len(c.output_indices) > 1 for c in program.components) for component in program.components: + assert isinstance(component, CompiledComponent) graphs = component.compiled_scalar_graphs for k in range(len(graphs) - 1): for prefix in itertools.product((0, 1), repeat=k): diff --git a/test/unit/statevector/test_program.py b/test/unit/statevector/test_program.py index 5cb863a4..c8386378 100644 --- a/test/unit/statevector/test_program.py +++ b/test/unit/statevector/test_program.py @@ -15,6 +15,7 @@ from tsim.circuit import Circuit 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.heisenberg import ComponentLayout from tsim.statevector.plan import ( AffineRow, @@ -50,6 +51,7 @@ def test_plans_match_stabilizer_rank(name, sample_detectors): continue planned, component_layout = plan_component(component, prepared.layout) reference = by_outputs[component_layout.output_indices] + assert isinstance(reference, CompiledComponent) selection = tuple(np.asarray(reference.f_selection).tolist()) assert selection == component_layout.f_selection for f in f_patterns(planned.num_f): From af3bc5e989056a67c8e60176802ac7ef43cb2eac Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Sat, 10 Oct 2026 17:32:57 -0400 Subject: [PATCH 08/18] Benchmark strategy="statevector" on the clifft-bench circuits Co-Authored-By: Claude Opus 5.5 --- scripts/statevector/benchmark.py | 96 ++++++++++++++++++++++++++++++++ scripts/statevector/results.md | 53 ++++++++++++++++++ 2 files changed, 149 insertions(+) create mode 100644 scripts/statevector/benchmark.py create mode 100644 scripts/statevector/results.md diff --git a/scripts/statevector/benchmark.py b/scripts/statevector/benchmark.py new file mode 100644 index 00000000..beee56c7 --- /dev/null +++ b/scripts/statevector/benchmark.py @@ -0,0 +1,96 @@ +"""Benchmark strategy="statevector" on the clifft-bench circuits. + +Run from the repository root: + + uv run python scripts/statevector/benchmark.py [--seconds S] [NAME ...] + +Circuits are read from the directory in ``CLIFFT_BENCH_CIRCUITS`` (default +``/Users/rafaelhaenel/Documents/quera/clifft-bench/workloads/circuits``). For +every circuit this compiles a detector sampler with ``strategy="statevector"``, +makes one untimed warmup call (JIT), then samples for at least ``--seconds`` +with postselection on every detector, counting as clifft-bench does: attempted +shots, accepted shots (no detector fired) and logical errors (observable 0 +among accepted shots). Acceptance is checked against Clifft 0.10.0's value +(clifft-bench release-v1-20260909-130000) within 5 standard errors. +""" + +import argparse +import math +import os +import time +from pathlib import Path + +import numpy as np + +import tsim + +CIRCUITS = Path( + os.environ.get( + "CLIFFT_BENCH_CIRCUITS", + "/Users/rafaelhaenel/Documents/quera/clifft-bench/workloads/circuits", + ) +) + +# name: (file, Clifft 0.10.0 acceptance, shots per sample call, batch size) +WORKLOADS = { + "distillation": ("distillation.stim", 0.0798, 262_144, 65_536), + "cultivation_d3": ("msc_d3_inject_cultivate_p1e-3.stim", 0.6868, 262_144, 65_536), + "coherent_d3_r3": ("coherent_d3_r3.stim", 0.3732, 262_144, 65_536), + "cultivation_d5": ("msc_d5_inject_cultivate_p1e-3.stim", 0.1440, 65_536, 16_384), + "coherent_d5_r5": ("coherent_d5_r5.stim", 0.0078, 16_384, 4_096), +} + + +def count(sampler, shots: int, batch_size: int, num_detectors: int) -> tuple[int, int]: + """Return (accepted shots, logical errors) of one sample call.""" + detectors, observables = sampler.sample( + shots, + batch_size=batch_size, + separate_observables=True, + postselection_mask=np.ones(num_detectors, dtype=bool), + ) + keep = ~detectors.any(axis=1) + return int(keep.sum()), int(observables[keep, 0].sum()) + + +def main() -> None: + """Benchmark the requested workloads and print one Markdown row each.""" + parser = argparse.ArgumentParser() + parser.add_argument("names", nargs="*", default=list(WORKLOADS)) + parser.add_argument("--seconds", type=float, default=10.0) + args = parser.parse_args() + print( + "| circuit | compile s | warmup s | shots/s | acceptance | Clifft " + "acceptance | check | logical errors / accepted | sampler |" + ) + print("|---|---:|---:|---:|---:|---:|---|---:|---|") + for name in args.names: + file, reference, shots, batch_size = WORKLOADS[name] + circuit = tsim.Circuit.from_file(str(CIRCUITS / file)) + start = time.perf_counter() + sampler = circuit.compile_detector_sampler(strategy="statevector", seed=0) + compile_seconds = time.perf_counter() - start + start = time.perf_counter() + count(sampler, shots, batch_size, circuit.num_detectors) + warmup_seconds = time.perf_counter() - start + attempted = accepted = errors = 0 + start = time.perf_counter() + while time.perf_counter() - start < args.seconds: + kept, wrong = count(sampler, shots, batch_size, circuit.num_detectors) + attempted += shots + accepted += kept + errors += wrong + elapsed = time.perf_counter() - start + acceptance = accepted / attempted + sigma = math.sqrt(reference * (1 - reference) / attempted) + check = "ok" if abs(acceptance - reference) <= 5 * sigma else "MISMATCH" + print( + f"| {name} | {compile_seconds:.2f} | {warmup_seconds:.2f} | " + f"{attempted / elapsed:,.0f} | {acceptance:.4f} | {reference:.4f} | " + f"{check} | {errors} / {accepted} | `{sampler!r}` |", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/statevector/results.md b/scripts/statevector/results.md new file mode 100644 index 00000000..7f37478a --- /dev/null +++ b/scripts/statevector/results.md @@ -0,0 +1,53 @@ +# strategy="statevector" on the clifft-bench circuits + +- **Date:** 2026-10-10 +- **Machine:** Apple M4 Pro (JAX CPU backend) +- **JAX:** 0.9.2 +- **Commit:** `bf80f7e` + +Command, from the repository root: + + uv run python scripts/statevector/benchmark.py --seconds 10 + +See the docstring of `benchmark.py` for what is timed and counted. Amplitudes +are complex64 (the default). + +## Measured + +| circuit | compile s | warmup s | shots/s | acceptance | Clifft acceptance | check | logical errors / accepted | sampler | +|---|---:|---:|---:|---:|---:|---|---:|---| +| distillation | 0.16 | 0.61 | 2,926,016 | 0.0799 | 0.0798 | ok | 990898 / 2344917 | `CompiledDetectorSampler(40 direct, 1 statevector components (width ≤ 3, 15 actions), 0.0 graphs, 473 error channel bits, 5 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 579 B)` | +| cultivation_d3 | 0.11 | 0.74 | 1,189,716 | 0.6866 | 0.6868 | ok | 6 / 8279256 | `CompiledDetectorSampler(16 direct, 1 statevector components (width ≤ 5, 36 actions), 0.0 graphs, 268 error channel bits, 5 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 1.6 kB)` | +| coherent_d3_r3 | 0.07 | 1.02 | 570,139 | 0.3732 | 0.3732 | ok | 0 / 2152025 | `CompiledDetectorSampler(17 direct, 1 statevector components (width ≤ 5, 73 actions), 0.0 graphs, 54 error channel bits, 8 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 3.0 kB)` | +| cultivation_d5 | 0.92 | 0.69 | 594,757 | 0.1442 | 0.1440 | ok | 0 / 859817 | `CompiledDetectorSampler(93 direct, 1 statevector components (width ≤ 6, 109 actions), 0.0 graphs, 1688 error channel bits, 15 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 6.7 kB)` | +| coherent_d5_r5 | 0.74 | 18.72 | 901 | 0.0094 | 0.0078 | ok | 0 / 154 | `CompiledDetectorSampler(73 direct, 1 statevector components (width ≤ 12, 427 actions), 0.0 graphs, 234 error channel bits, 48 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 26.0 kB)` | + +## Reference + +| circuit | Clifft 0.10.0 shots/s (one core, AMD EPYC 9R14) | Tsim stabilizer rank on this machine | +|---|---:|---:| +| distillation | 1.87M | 603k (cat5) | +| cultivation_d3 | 2.82M | 567 (cutting) | +| coherent_d3_r3 | 1.05M | does not compile in 120 s | +| cultivation_d5 | 197k | does not compile in 120 s | +| coherent_d5_r5 | 5.9k | does not compile in 120 s | + +## Summary + +All five acceptances agree with Clifft 0.10.0 within 5 standard errors, and +every circuit compiles in under 1 s, including the three that the stabilizer +rank strategy does not compile in 120 s. Shots/s exceeds Clifft's one-core +figure (measured on a different machine) on distillation and cultivation_d5 +and falls below it on cultivation_d3, coherent_d3_r3 and coherent_d5_r5 +(901 against 5.9k). + +Run notes: + +- coherent_d5_r5: one sample call of 16,384 shots takes about 18 s, so the + timed window holds a single call and its acceptance check rests on 16,384 + attempted shots (154 accepted). +- coherent_d5_r5 printed six `UserWarning`s from `sampler.py` ("A marginal + probability was not normalized correctly"), with deviations from 2.0e-05 to + 4.1e-05 (the warning threshold is 1e-05). With `JAX_ENABLE_X64=1` + (complex128), `--seconds 1 coherent_d5_r5` printed no warning, with + acceptance 0.0079 (129 accepted of 16,384) at 599 shots/s. From b55970c748e5130e302fe9f781cdd11c7bd59b4c Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Sat, 10 Oct 2026 17:39:21 -0400 Subject: [PATCH 09/18] Compare logical-error rates and JIT time in the statevector results Co-Authored-By: Claude Opus 5.5 --- scripts/statevector/results.md | 60 ++++++++++++++++++++++++++++++---- 1 file changed, 53 insertions(+), 7 deletions(-) diff --git a/scripts/statevector/results.md b/scripts/statevector/results.md index 7f37478a..87d81bdf 100644 --- a/scripts/statevector/results.md +++ b/scripts/statevector/results.md @@ -1,7 +1,12 @@ # strategy="statevector" on the clifft-bench circuits - **Date:** 2026-10-10 -- **Machine:** Apple M4 Pro (JAX CPU backend) +- **Machine:** Apple M4 Pro, 14 cores (10 performance, 4 efficiency) +- **Backend:** JAX CPU, default settings (no `XLA_FLAGS` or thread limits), so + XLA may use every core +- **Threads used:** the full run took 219.23 s user + 14.67 s system over + 82.98 s wall (shell `time`), about 2.8 cores busy on average. Tsim's shots/s + below are multi-threaded, not one-core figures. - **JAX:** 0.9.2 - **Commit:** `bf80f7e` @@ -10,7 +15,9 @@ Command, from the repository root: uv run python scripts/statevector/benchmark.py --seconds 10 See the docstring of `benchmark.py` for what is timed and counted. Amplitudes -are complex64 (the default). +are complex64 (the default). "warmup s" is the first sample call, which +includes JIT compilation and one full sample call of the workload's shots per +call (see JIT time below). ## Measured @@ -32,14 +39,53 @@ are complex64 (the default). | cultivation_d5 | 197k | does not compile in 120 s | | coherent_d5_r5 | 5.9k | does not compile in 120 s | +## Logical-error rate + +Tsim's rate is logical errors / accepted from the measured table. Clifft +0.10.0's rates are from the design spec (section 11). "Expected" is accepted +times Clifft's rate. + +| circuit | Tsim errors / accepted | Tsim rate | Clifft 0.10.0 rate | expected errors at Clifft's rate | +|---|---:|---:|---:|---:| +| distillation | 990898 / 2344917 | 0.4226 | 0.422 | 989,555 (binomial sd 756) | +| cultivation_d3 | 6 / 8279256 | 7.2e-7 | 9.6e-7 | 7.95 | +| coherent_d3_r3 | 0 / 2152025 | 0 | 5.1e-8 | 0.11 | +| cultivation_d5 | 0 / 859817 | 0 | 0 (in 3.0e7 shots) | 0 | +| coherent_d5_r5 | 0 / 154 | 0 | 0 (in 8.8e5 shots) | 0 | + +distillation is 1,343 errors (1.8 sd) above the count expected at 0.422. The +reference's third-digit rounding (±0.0005) by itself moves that expected count +by up to ±1,172. For cultivation_d3, a Poisson count with mean 7.95 is 6 or +fewer with probability 0.32. For coherent_d3_r3, a Poisson count with mean +0.11 is 0 with probability 0.90. + +## JIT time + +The script does not time JIT separately. Estimate it as warmup s minus one +call's time, where one call's time is shots per call divided by the measured +shots/s. This assumes the warmup call samples at the timed rate. The estimate +also includes any other first-call overhead, and warmup s is rounded to +0.01 s. + +| circuit | warmup s | shots per call | one call s | JIT estimate s | +|---|---:|---:|---:|---:| +| distillation | 0.61 | 262,144 | 0.090 | 0.52 | +| cultivation_d3 | 0.74 | 262,144 | 0.220 | 0.52 | +| coherent_d3_r3 | 1.02 | 262,144 | 0.460 | 0.56 | +| cultivation_d5 | 0.69 | 65,536 | 0.110 | 0.58 | +| coherent_d5_r5 | 18.72 | 16,384 | 18.184 | 0.54 | + ## Summary All five acceptances agree with Clifft 0.10.0 within 5 standard errors, and -every circuit compiles in under 1 s, including the three that the stabilizer -rank strategy does not compile in 120 s. Shots/s exceeds Clifft's one-core -figure (measured on a different machine) on distillation and cultivation_d5 -and falls below it on cultivation_d3, coherent_d3_r3 and coherent_d5_r5 -(901 against 5.9k). +all five logical-error counts are consistent with Clifft's rates. Every +circuit compiles in under 1 s, including the three that the stabilizer rank +strategy does not compile in 120 s, and JIT takes about 0.5 s. Tsim's +multi-threaded shots/s on this machine (about 2.8 cores busy on average) is +above Clifft's one-core EPYC figure on distillation and cultivation_d5, and +below it on cultivation_d3, coherent_d3_r3 and coherent_d5_r5 (901 against +5.9k). This compares different machines and thread counts, so it is not a +per-core comparison. Run notes: From cff71f141a25aa2e771961f55dd7068f5350a7b4 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Sat, 10 Oct 2026 18:27:43 -0400 Subject: [PATCH 10/18] Sample outcomes with 2^-64 resolution in choose_outcome MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A float32 uniform draw is a multiple of 2^-23, so any p1 below 2^-23 gave outcome 1 with probability 2^-23: deterministic outcomes, whose p1 rounds to ~1e-15 instead of 0, fired at ~1.2e-7 per measurement under complex64. Compare two uint32 draws with p1·2^32 instead, independent of the dtype. Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/execute.py | 15 ++++++++++-- test/unit/statevector/test_execute.py | 33 +++++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 2 deletions(-) diff --git a/src/tsim/statevector/execute.py b/src/tsim/statevector/execute.py index cac4369a..9c0c9cbf 100644 --- a/src/tsim/statevector/execute.py +++ b/src/tsim/statevector/execute.py @@ -22,6 +22,7 @@ _I_POWERS = (1, 1j, -1, -1j) _TINY = 1e-30 +_TWO_TO_32 = 2.0**32 class LaneState(NamedTuple): @@ -48,10 +49,20 @@ class LaneState(NamedTuple): def choose_outcome(p1: Array, count: Array, key: Array) -> Array: """Return each lane's outcome: True with probability ``p1``. - ``count`` is unused for now; trajectory sharing will split it binomially. + Two uint32 draws per lane form a uniform 64-bit fraction ``high + low·2^-32`` + (in units of 2^-32), compared with ``t = p1·2^32`` (an exact scaling), so the + resolution is 2^-64 for any float dtype of ``p1``. ``p1 = 0`` never and + ``p1 = 1`` always gives True. ``count`` is unused for now; trajectory sharing + will split it binomially. """ del count - return jax.random.uniform(key, p1.shape, dtype=p1.dtype) < p1 + high, low = jax.random.bits(key, (2, *p1.shape), dtype=jnp.uint32) + scaled = jnp.clip(p1, 0, 1) * _TWO_TO_32 + whole = jnp.floor(scaled) + high_bound = whole.astype(jnp.uint32) + low_bound = jnp.floor((scaled - whole) * _TWO_TO_32).astype(jnp.uint32) + below = (high < high_bound) | ((high == high_bound) & (low < low_bound)) + return (p1 >= 1) | ((p1 > 0) & below) def pack_bits(bits: Array, num_words: int) -> Array: diff --git a/test/unit/statevector/test_execute.py b/test/unit/statevector/test_execute.py index a7d79c4a..4885e877 100644 --- a/test/unit/statevector/test_execute.py +++ b/test/unit/statevector/test_execute.py @@ -179,3 +179,36 @@ def test_complex128_amplitudes(): assert state.amps.dtype == jnp.complex128 out = run(compiled, state, jax.random.key(0), 0, compiled.num_actions) assert out.amps.dtype == jnp.complex128 + + +def _count_outcomes(p1, dtype, draws, chunk=2**22): + @jax.jit + def ones_in_chunk(key): + probabilities = jnp.full(chunk, p1, dtype=dtype) + count = jnp.ones(chunk, dtype=jnp.int32) + return jnp.sum(execute.choose_outcome(probabilities, count, key)) + + return sum( + int(ones_in_chunk(jax.random.key(seed))) for seed in range(draws // chunk) + ) + + +@pytest.mark.parametrize("p1", [0.0, 1e-15]) +def test_choose_outcome_never_returns_true_below_its_resolution(p1): + assert _count_outcomes(p1, jnp.float32, 2**26) == 0 + + +@pytest.mark.parametrize("dtype", ["float32", "float64"]) +def test_choose_outcome_always_returns_true_for_p1_one(dtype): + with jax.enable_x64(dtype == "float64"): + assert _count_outcomes(1.0, jnp.dtype(dtype), 2**22) == 2**22 + + +@pytest.mark.parametrize("dtype", ["float32", "float64"]) +@pytest.mark.parametrize("p1", [0.5, 0.3, 1e-3, 1e-6]) +def test_choose_outcome_frequency_matches_p1(p1, dtype): + draws = 2**24 + with jax.enable_x64(dtype == "float64"): + ones = _count_outcomes(p1, jnp.dtype(dtype), draws) + expected = draws * p1 + assert abs(ones - expected) <= 5 * np.sqrt(draws * p1 * (1 - p1)) + 1 From 944a1e5e7bec2028c10d2dc8de69cc7110fc33ab Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Sat, 10 Oct 2026 18:31:07 -0400 Subject: [PATCH 11/18] Compute measurement probabilities from the outcome weights p1 = (nu -

)/(2 nu) cancels catastrophically (p1 ~ 1e-7 instead of ~1e-15 on eigenstates under complex64), and renormalizing by 1 - p1 loses the relative precision of a small p0. Sum each outcome's weight directly (|psi +- P psi|^2 / 4, or the even and odd parity halves) and renormalize by the selected weight. Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/execute.py | 39 +++++++++------ test/unit/statevector/test_execute.py | 71 +++++++++++++++++++++++++++ 2 files changed, 96 insertions(+), 14 deletions(-) diff --git a/src/tsim/statevector/execute.py b/src/tsim/statevector/execute.py index 9c0c9cbf..ebd6b85f 100644 --- a/src/tsim/statevector/execute.py +++ b/src/tsim/statevector/execute.py @@ -154,6 +154,15 @@ def _rotate( return state._replace(amps=amps.astype(state.amps.dtype)) +def _squared(values: Array) -> Array: + return jnp.real(values) ** 2 + jnp.imag(values) ** 2 + + +def _probability_of_one(weight_0: Array, weight_1: Array) -> Array: + """Return ``w1 / (w0 + w1)`` from the outcome weights, without cancellation.""" + return jnp.clip(weight_1 / jnp.maximum(weight_0 + weight_1, _TINY), 0, 1) + + def _measured(state: LaneState, new: Array, norm: Array) -> LaneState: deviation = jnp.maximum(state.max_norm_deviation, jnp.max(jnp.abs(norm - 1))) return state._replace( @@ -168,16 +177,17 @@ def _measure( index = _index(component) amps = state.amps paulied = _apply_pauli(amps, index, component.x_mask[a], component.z_mask[a]) - norm = jnp.sum(jnp.abs(amps) ** 2, axis=1) - expectation = jnp.real(jnp.sum(jnp.conj(amps) * paulied, axis=1)) - p1 = jnp.clip((norm - expectation) / (2 * norm), 0, 1) + plus, minus = amps + paulied, amps - paulied + weight_0 = jnp.sum(_squared(plus), axis=1) / 4 + weight_1 = jnp.sum(_squared(minus), axis=1) / 4 + p1 = _probability_of_one(weight_0, weight_1) outcome = choose_outcome(p1, state.count, jax.random.fold_in(key, a)) - p_outcome = jnp.where(outcome, p1, 1 - p1) - scale = jax.lax.rsqrt(jnp.maximum(2 * norm * p_outcome, _TINY))[:, None] - signed = jnp.where(outcome, -1, 1)[:, None] + weight = jnp.where(outcome, weight_1, weight_0) + scale = jax.lax.rsqrt(jnp.maximum(2 * weight, _TINY))[:, None] keep = ((index >> component.pivot[a].astype(jnp.uint32)) & 1) == 0 - new = jnp.where(keep[None, :], (amps + signed * paulied) * scale, 0) - return _write(component, _measured(state, new, norm), a, outcome) + projected = jnp.where(outcome[:, None], minus, plus) + new = jnp.where(keep[None, :], projected * scale, 0) + return _write(component, _measured(state, new, weight_0 + weight_1), a, outcome) def _measure_diag( @@ -186,11 +196,12 @@ def _measure_diag( index = _index(component) amps = state.amps odd = _parity(index & component.z_mask[a]) - weights = jnp.abs(amps) ** 2 - norm = jnp.sum(weights, axis=1) - p1 = jnp.clip(jnp.sum(jnp.where(odd[None, :], weights, 0), axis=1) / norm, 0, 1) + weights = _squared(amps) + weight_0 = jnp.sum(jnp.where(odd[None, :], 0, weights), axis=1) + weight_1 = jnp.sum(jnp.where(odd[None, :], weights, 0), axis=1) + p1 = _probability_of_one(weight_0, weight_1) outcome = choose_outcome(p1, state.count, jax.random.fold_in(key, a)) - p_outcome = jnp.where(outcome, p1, 1 - p1) + weight = jnp.where(outcome, weight_1, weight_0) pivot = component.pivot[a].astype(jnp.uint32) bit = jnp.uint32(1) << pivot source_if_0 = (index ^ (bit * odd.astype(jnp.uint32))).astype(jnp.int32) @@ -201,9 +212,9 @@ def _measure_diag( jnp.take(amps, source_if_0, axis=1), ) keep = ((index >> pivot) & 1) == 0 - scale = jax.lax.rsqrt(jnp.maximum(norm * p_outcome, _TINY))[:, None] + scale = jax.lax.rsqrt(jnp.maximum(weight, _TINY))[:, None] new = jnp.where(keep[None, :], gathered * scale, 0) - return _write(component, _measured(state, new, norm), a, outcome) + return _write(component, _measured(state, new, weight_0 + weight_1), a, outcome) def _measure_dormant( diff --git a/test/unit/statevector/test_execute.py b/test/unit/statevector/test_execute.py index 4885e877..117dfc2c 100644 --- a/test/unit/statevector/test_execute.py +++ b/test/unit/statevector/test_execute.py @@ -1,5 +1,6 @@ from test.helpers.statevector import ( CIRCUITS, + apply_pauli, assert_matches_distribution, f_patterns, pauli_matrix, @@ -212,3 +213,73 @@ def test_choose_outcome_frequency_matches_p1(p1, dtype): ones = _count_outcomes(p1, jnp.dtype(dtype), draws) expected = draws * p1 assert abs(ones - expected) <= 5 * np.sqrt(draws * p1 * (1 - p1)) + 1 + + +def _projected(vectors, x, z, eigenvalue): + projected = np.array([v + eigenvalue * apply_pauli(v, x, z) for v in vectors]) + return projected / np.linalg.norm(projected, axis=1, keepdims=True) + + +def _run_recording_p1(monkeypatch, component, state): + recorded = [] + choose = execute.choose_outcome + + def recording(p1, count, key): + recorded.append(np.asarray(p1, dtype=np.float64)) + return choose(p1, count, key) + + monkeypatch.setattr(execute, "choose_outcome", recording) + with jax.disable_jit(): + out = run(component, state, jax.random.key(3), 0, 1) + (p1,) = recorded + return p1, out + + +_MEASUREMENTS = [(0b011, 0b110, 0), (0b110, 0b010, 2), (0, 0b101, 2), (0, 0b011, 0)] + + +@pytest.mark.parametrize("eigenvalue", [1, -1]) +@pytest.mark.parametrize("x, z, pivot", _MEASUREMENTS) +def test_measuring_an_eigenstate_never_samples_the_impossible_outcome( + monkeypatch, x, z, pivot, eigenvalue +): + lanes = 4096 + action = MeasureActive( + x, z, pivot, branch=1, sign=AffineRow(False, frozenset()), slot=0 + ) + component = _component([action], width=3, num_branches=1, num_outputs=1) + # NOTE: the noise is of the size rotations leave in complex64 amplitudes. + rounding = 1e-7 * _random_states(lanes, 3, seed=1) + psi = _projected(_random_states(lanes, 3), x, z, eigenvalue) + rounding + state = initial_state(component, jnp.zeros((lanes, 1), dtype=bool))._replace( + amps=jnp.asarray(psi, jnp.complex64) + ) + p1, out = _run_recording_p1(monkeypatch, component, state) + impossible = p1 if eigenvalue == 1 else 1 - p1 + assert impossible.max() < 1e-12 + assert np.all(np.asarray(out.records[:, 0]) == (eigenvalue == -1)) + + +@pytest.mark.parametrize("unlikely", [0, 1]) +@pytest.mark.parametrize("x, z, pivot", _MEASUREMENTS) +def test_an_unlikely_outcome_leaves_a_normalized_state( + monkeypatch, x, z, pivot, unlikely +): + lanes = 64 + action = MeasureActive( + x, z, pivot, branch=1, sign=AffineRow(False, frozenset()), slot=0 + ) + component = _component([action], width=3, num_branches=1, num_outputs=1) + vectors = _random_states(lanes, 3) + likely = _projected(vectors, x, z, (-1) ** (1 - unlikely)) + rare = _projected(vectors, x, z, (-1) ** unlikely) + p = np.logspace(-5, -3, lanes)[:, None] + psi = np.sqrt(1 - p) * likely + np.sqrt(p) * rare + outcomes = jnp.full(lanes, bool(unlikely)) + monkeypatch.setattr(execute, "choose_outcome", lambda p1, count, key: outcomes) + state = initial_state(component, jnp.zeros((lanes, 1), dtype=bool))._replace( + amps=jnp.asarray(psi, jnp.complex64) + ) + out = run(component, state, jax.random.key(0), 0, 1) + norms = np.linalg.norm(np.asarray(out.amps, dtype=np.complex128), axis=1) ** 2 + assert np.abs(norms - 1).max() < 1e-5 From 812d94c39e42ce426172708f80093d4878a4f045 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Sat, 10 Oct 2026 18:32:57 -0400 Subject: [PATCH 12/18] Warn about statevector norm drift at a dtype-aware threshold The 1e-5 threshold was built for stabilizer-rank arithmetic. Under complex64 each rotation adds ~1e-7 relative drift, which is harmless because every probability is normalized by the measured weights, so warn above 1e-3 for complex64 and 1e-9 for complex128. Stabilizer-rank components keep 1e-5. Co-Authored-By: Claude Opus 5.5 --- src/tsim/sampler.py | 9 +++- src/tsim/statevector/execute.py | 5 ++ test/integration/test_statevector_sampler.py | 53 ++++++++++++++++++++ 3 files changed, 65 insertions(+), 2 deletions(-) diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index 4a885830..56aad280 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -17,7 +17,10 @@ from tsim.core.graph import prepare_graph from tsim.core.types import CompiledComponent, CompiledProgram, SamplingStrategy from tsim.noise.channels import ChannelSampler -from tsim.statevector.execute import sample_statevector_component +from tsim.statevector.execute import ( + NORM_WARNING_THRESHOLDS, + sample_statevector_component, +) from tsim.statevector.program import StatevectorComponent from tsim.utils.cuda_helpers import copy_d2h @@ -151,17 +154,19 @@ def sample_program( samples, key, max_norm_deviation = sample_statevector_component( component, f_params, key ) + warning_threshold = NORM_WARNING_THRESHOLDS[component.dtype] else: samples, key, max_norm_deviation = sample_component( component, f_params, key ) + warning_threshold = 1e-5 if np.isclose(max_norm_deviation, 1): raise ValueError( "A vanishing marginal probability distribution was encountered (normalization 0). " "This is likely the result of an underflow error. Please report this " "as a bug at https://github.com/QuEraComputing/tsim/issues/new." ) # pragma: no cover - if max_norm_deviation > 1e-5: + if max_norm_deviation > warning_threshold: warnings.warn( "A marginal probability was not normalized correctly " f"(normalization deviated from 1 by {max_norm_deviation:.1e}). " diff --git a/src/tsim/statevector/execute.py b/src/tsim/statevector/execute.py index ebd6b85f..f68c9a79 100644 --- a/src/tsim/statevector/execute.py +++ b/src/tsim/statevector/execute.py @@ -24,6 +24,11 @@ _TINY = 1e-30 _TWO_TO_32 = 2.0**32 +# NOTE: every outcome probability is normalized by the measured weights, so norm +# drift is harmless until it exceeds what rounding can explain: each rotation +# adds ~1e-7 relative drift under complex64. +NORM_WARNING_THRESHOLDS = {"complex64": 1e-3, "complex128": 1e-9} + class LaneState(NamedTuple): """Per-lane state of the executor; the leading axis is the lane. diff --git a/test/integration/test_statevector_sampler.py b/test/integration/test_statevector_sampler.py index e88d1bf4..91f8fc05 100644 --- a/test/integration/test_statevector_sampler.py +++ b/test/integration/test_statevector_sampler.py @@ -1,9 +1,11 @@ """strategy="statevector" end to end, against stabilizer rank and VecSampler.""" +import warnings from test.helpers.gen import gen_stim_circuit from test.helpers.statevector import CIRCUITS, assert_counts_match import jax +import jax.numpy as jnp import numpy as np import pytest @@ -121,3 +123,54 @@ def test_a_lane_that_does_not_fit_in_memory_is_rejected(monkeypatch): ) with pytest.raises(ValueError, match="cat5"): Circuit(CIRCUITS["t_ghz"]).compile_detector_sampler(strategy="statevector") + + +def _norm_warnings(monkeypatch, strategy, deviation, *, x64=False): + name = ( + "sample_statevector_component" + if strategy == "statevector" + else "sample_component" + ) + sample = getattr(sampler_module, name) + + def drifting(component, f_params, key): + samples, key, _ = sample(component, f_params, key) + return samples, key, jnp.asarray(deviation) + + monkeypatch.setattr(sampler_module, name, drifting) + with jax.enable_x64(x64): + sampler = Circuit(CIRCUITS["t_ghz"]).compile_detector_sampler( + strategy=strategy, seed=0 + ) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + sampler.sample(10) + return [w for w in caught if "not normalized" in str(w.message)] + + +@pytest.mark.parametrize( + "x64, deviation, warns", + [ + (False, 5e-4, False), + (False, 2e-3, True), + (True, 1e-10, False), + (True, 1e-8, True), + ], +) +def test_statevector_norm_warnings_depend_on_the_amplitude_dtype( + monkeypatch, x64, deviation, warns +): + caught = _norm_warnings(monkeypatch, "statevector", deviation, x64=x64) + assert bool(caught) == warns + + +@pytest.mark.parametrize("deviation, warns", [(5e-6, False), (2e-5, True)]) +def test_stabilizer_rank_norm_warnings_keep_their_threshold( + monkeypatch, deviation, warns +): + assert bool(_norm_warnings(monkeypatch, "cat5", deviation)) == warns + + +def test_a_vanishing_statevector_norm_raises(monkeypatch): + with pytest.raises(ValueError, match="vanishing"): + _norm_warnings(monkeypatch, "statevector", 1.0) From 8892c5df99ccc1da18346f77383397790aa1c2fc Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Sat, 10 Oct 2026 18:35:27 -0400 Subject: [PATCH 13/18] Cap and recount the automatic batch size of statevector samplers The per-shot estimate counted only the amplitudes; peak memory was up to 2.7x higher, and the memory-derived batches (millions of lanes) were also slower. Count the f-params and records on host and device, cap the automatic batch at 2^16 lanes, and require one lane to fit in the same 50% budget that sizes batches. Co-Authored-By: Claude Opus 5.5 --- src/tsim/sampler.py | 30 ++++++++--- test/integration/test_statevector_sampler.py | 53 ++++++++++++++++++++ 2 files changed, 77 insertions(+), 6 deletions(-) diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index 56aad280..e4b14716 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -29,6 +29,11 @@ from tsim.circuit import Circuit +# NOTE: statevector batches beyond 2^16 lanes cost memory without being faster: +# on CPU, 2^16 lanes sampled cultivation d3 1.7x faster than the memory-derived +# batch, at a fifth of its peak memory. +STATEVECTOR_MAX_BATCH_SIZE = 2**16 + def _sample_component( component: CompiledComponent, @@ -253,13 +258,13 @@ def __init__( statevector = [ c for c in self._program.components if isinstance(c, StatevectorComponent) ] - if statevector and self._peak_bytes_per_sample() > self._available_memory(): + if statevector and self._peak_bytes_per_sample() > self._memory_budget(): widest = max(statevector, key=lambda c: c.width) raise ValueError( f"strategy='statevector' needs {self._peak_bytes_per_sample()} bytes " f"per shot for the component with outputs {list(widest.output_indices)} " - f"(width {widest.width}), more than the available memory. Use a " - "stabilizer-rank strategy such as strategy='cat5'." + f"(width {widest.width}), more than half of the available memory. " + "Use a stabilizer-rank strategy such as strategy='cat5'." ) def _compute_direct_outputs(self, f_params_np: np.ndarray) -> np.ndarray: @@ -320,12 +325,15 @@ def _resolve_batch_size( def _peak_bytes_per_sample(self) -> int: """Estimate peak device memory per sample from compiled program structure.""" + num_f = self._channel_sampler.signature_matrix.shape[1] peak = 0 for component in self._program.components: if isinstance(component, StatevectorComponent): itemsize = 8 if component.dtype == "complex64" else 16 lane = 3 * 2**component.width * itemsize lane += 4 * component.num_words + component.num_outputs + 1 + # NOTE: the f-params and the output records, on host and device. + lane += 2 * num_f + 2 * self._program.num_outputs peak = max(peak, lane) continue for circuit in component.compiled_scalar_graphs: @@ -348,10 +356,20 @@ def _available_memory(self) -> int: return limit - stats.get("bytes_in_use", 0) return psutil.virtual_memory().available + def _memory_budget(self) -> int: + """Return the memory one batch may use: half of the available memory.""" + return int(self._available_memory() * 0.5) # conservative + def _estimate_batch_size(self) -> int: - """Estimate the largest batch size that fits in available device memory.""" - half_of_available = int(self._available_memory() * 0.5) # conservative - return max(1, half_of_available // self._peak_bytes_per_sample()) + """Estimate the largest batch size that fits in available device memory. + + Programs with statevector components use at most + ``STATEVECTOR_MAX_BATCH_SIZE`` lanes per batch. + """ + batch_size = max(1, self._memory_budget() // self._peak_bytes_per_sample()) + if any(isinstance(c, StatevectorComponent) for c in self._program.components): + batch_size = min(batch_size, STATEVECTOR_MAX_BATCH_SIZE) + return batch_size @overload def _sample_batches( diff --git a/test/integration/test_statevector_sampler.py b/test/integration/test_statevector_sampler.py index 91f8fc05..b880ab88 100644 --- a/test/integration/test_statevector_sampler.py +++ b/test/integration/test_statevector_sampler.py @@ -174,3 +174,56 @@ def test_stabilizer_rank_norm_warnings_keep_their_threshold( def test_a_vanishing_statevector_norm_raises(monkeypatch): with pytest.raises(ValueError, match="vanishing"): _norm_warnings(monkeypatch, "statevector", 1.0) + + +def test_the_automatic_batch_size_of_a_statevector_sampler_is_capped(monkeypatch): + monkeypatch.setattr( + sampler_module._CompiledSamplerBase, "_available_memory", lambda self: 2**50 + ) + circuit = Circuit(CIRCUITS["t_ghz"]) + sampler = circuit.compile_detector_sampler(strategy="statevector", seed=0) + assert sampler._estimate_batch_size() == 2**16 + assert sampler._resolve_batch_size(10**7, None, compute_reference=False) <= 2**16 + lanes = [] + sample_program = sampler_module.sample_program + + def recording(program, f_params, key): + lanes.append(f_params.shape[0]) + return sample_program(program, f_params, key) + + monkeypatch.setattr(sampler_module, "sample_program", recording) + sampler.sample(2**16 + 1) + sampler.sample( + 2**16 + 1, postselection_mask=np.ones(circuit.num_detectors, dtype=bool) + ) + assert lanes and max(lanes) <= 2**16 + stabilizer_rank = circuit.compile_detector_sampler(strategy="cat5") + assert stabilizer_rank._estimate_batch_size() > 2**16 + + +def test_the_per_shot_estimate_counts_f_params_and_records(): + sampler = Circuit(CIRCUITS["t_ghz"]).compile_detector_sampler( + strategy="statevector" + ) + widest = max( + c.width + for c in sampler._program.components + if isinstance(c, StatevectorComponent) + ) + num_f = sampler._channel_sampler.signature_matrix.shape[1] + records = sampler._program.num_outputs + amplitudes = 3 * 2**widest * 8 + assert sampler._peak_bytes_per_sample() >= amplitudes + 2 * num_f + 2 * records + + +def test_a_lane_must_fit_in_half_of_the_available_memory(monkeypatch): + circuit = Circuit(CIRCUITS["t_ghz"]) + peak = circuit.compile_detector_sampler( + strategy="statevector" + )._peak_bytes_per_sample() + base = sampler_module._CompiledSamplerBase + monkeypatch.setattr(base, "_available_memory", lambda self: 2 * peak - 1) + with pytest.raises(ValueError, match="cat5"): + circuit.compile_detector_sampler(strategy="statevector") + monkeypatch.setattr(base, "_available_memory", lambda self: 2 * peak) + circuit.compile_detector_sampler(strategy="statevector") From e7fce813eafcf8e9f2c6455fdfaad5e1628661eb Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Sat, 10 Oct 2026 18:39:00 -0400 Subject: [PATCH 14/18] Test the untested planner paths and the run resume hook Hand-built ops reach a dormant measurement of Y (S_DAG on the pivot), a forgotten dormant measurement that no later op depends on, and a diagonal measurement on two active axes. Running a tape as [0, k) then [k, A) reproduces running [0, A) field by field. Co-Authored-By: Claude Opus 5.5 --- test/unit/statevector/test_execute.py | 29 +++++++- test/unit/statevector/test_plan.py | 101 +++++++++++++++++++++++++- 2 files changed, 128 insertions(+), 2 deletions(-) diff --git a/test/unit/statevector/test_execute.py b/test/unit/statevector/test_execute.py index 117dfc2c..3586de75 100644 --- a/test/unit/statevector/test_execute.py +++ b/test/unit/statevector/test_execute.py @@ -30,7 +30,14 @@ Record, Rotate, ) -from tsim.statevector.program import lower, plan_component, resolve_dtype +from tsim.statevector.program import ( + MEASURE, + MEASURE_DIAG, + MEASURE_DORMANT, + lower, + plan_component, + resolve_dtype, +) def _component( @@ -283,3 +290,23 @@ def test_an_unlikely_outcome_leaves_a_normalized_state( out = run(component, state, jax.random.key(0), 0, 1) norms = np.linalg.norm(np.asarray(out.amps, dtype=np.complex128), axis=1) ** 2 assert np.abs(norms - 1).max() < 1e-5 + + +def test_running_a_tape_in_two_parts_equals_running_it_whole(): + (component,), layout = quantum_components( + CIRCUITS["midcircuit_feedforward"], sample_detectors=True + ) + planned, component_layout = plan_component(component, layout) + compiled = lower(planned, component_layout, "complex64") + lanes, split, stop = 512, 4, compiled.num_actions + kinds = np.asarray(compiled.kind) + measurements = np.isin(kinds, [MEASURE_DIAG, MEASURE, MEASURE_DORMANT]) + assert measurements[:split].any() and measurements[split:].any() + rng = np.random.default_rng(0) + f_bits = jnp.asarray(rng.integers(0, 2, (lanes, planned.num_f)).astype(bool)) + state = initial_state(compiled, f_bits) + key = jax.random.key(5) + whole = run(compiled, state, key, 0, stop) + parts = run(compiled, run(compiled, state, key, 0, split), key, split, stop) + for got, expected in zip(parts, whole, strict=True): + np.testing.assert_array_equal(np.asarray(got), np.asarray(expected)) diff --git a/test/unit/statevector/test_plan.py b/test/unit/statevector/test_plan.py index cc119410..ed690189 100644 --- a/test/unit/statevector/test_plan.py +++ b/test/unit/statevector/test_plan.py @@ -10,15 +10,23 @@ import numpy as np import pytest +import stim from zx_param import Phase, Scalar from zx_param.extract import ExtractedCircuit, Gate -from tsim.statevector.heisenberg import ComponentLayout, heisenberg_ops +from tsim.statevector.heisenberg import ( + MEASUREMENT, + ROTATION, + ComponentLayout, + HeisenbergOps, + heisenberg_ops, +) from tsim.statevector.plan import ( AffineRow, MeasureActive, MeasureDormant, Rotate, + _Planner, plan, ) from tsim.statevector.schedule import Schedule, schedule @@ -93,3 +101,94 @@ def test_plans_are_well_formed(): assert action.branch not in defined defined.add(action.branch) assert defined == set(range(planned.num_symbols)) + + +def _ops(entries): + """Return ops from (Pauli, angle or None for a measurement, f-sign, slot).""" + paulis = [stim.PauliString(pauli).to_numpy() for pauli, _, _, _ in entries] + return HeisenbergOps( + num_lines=len(entries[0][0]), + num_f=1, + kind=np.array( + [MEASUREMENT if angle is None else ROTATION for _, angle, _, _ in entries], + dtype=np.uint8, + ), + x=np.array([x for x, _ in paulis], dtype=bool), + z=np.array([z for _, z in paulis], dtype=bool), + angle=np.array([angle or 0.0 for _, angle, _, _ in entries]), + sign=np.array([[False, f_sign] for _, _, f_sign, _ in entries], dtype=bool), + slot=np.array([slot for _, _, _, slot in entries], dtype=np.int32), + ) + + +def _plan_recording_gates(monkeypatch, ops): + """Plan ``ops`` and check it; return the plan and the applied virtual gates.""" + applied = [] + apply = _Planner._apply + + def recording(self, gate, *targets): + applied.append((gate, targets, frozenset(self.active))) + apply(self, gate, *targets) + + monkeypatch.setattr(_Planner, "_apply", recording) + scheduled = schedule(ops) + planned = plan(ops, scheduled) + reordered = ops.take(scheduled.order) + for f in f_patterns(ops.num_f): + assert_distributions_close( + plan_distribution(planned, f), ops_distribution(reordered, f), atol=1e-9 + ) + return planned, applied + + +def test_a_dormant_measurement_of_y_turns_its_pivot_with_s_dag(monkeypatch): + ops = _ops( + [ + ("YX", None, True, 0), + ("X_", 0.3, False, -1), + ("_Y", 0.7, True, -1), + ("Z_", None, False, 1), + ("_Z", None, False, 2), + ] + ) + planned, applied = _plan_recording_gates(monkeypatch, ops) + assert isinstance(planned.actions[0], MeasureDormant) + assert any( + gate == "S_DAG" and targets[0] not in active + for gate, targets, active in applied + ) + + +def test_a_forgotten_dormant_measurement_nothing_depends_on_is_skipped(monkeypatch): + ops = _ops( + [ + ("X_", 0.4, True, -1), + ("_X", None, False, -1), + ("Z_", 0.3, False, -1), + ("Z_", 0.2, True, -1), + ("Y_", None, False, 0), + ] + ) + planned, _ = _plan_recording_gates(monkeypatch, ops) + assert not any(isinstance(action, MeasureDormant) for action in planned.actions) + assert planned.num_branches == 1 + + +def test_a_diagonal_measurement_on_several_active_axes(monkeypatch): + ops = _ops( + [ + ("X_", 0.3, True, -1), + ("_X", 0.5, False, -1), + ("ZZ", None, False, 0), + ("_Y", 0.6, True, -1), + ("_Z", None, False, 1), + ] + ) + planned, applied = _plan_recording_gates(monkeypatch, ops) + assert any( + isinstance(action, MeasureActive) + and action.x == 0 + and bin(action.z).count("1") == 2 + for action in planned.actions + ) + assert ("CX", (1, 0), frozenset({0, 1})) in applied From 2a422ce55aeb7724509bcc4f4049f3ad1b6c8b3b Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Sat, 10 Oct 2026 18:45:26 -0400 Subject: [PATCH 15/18] Rerun the statevector benchmark with the executor fixes All five acceptances still agree with Clifft 0.10.0 and no norm warning is printed. Write zero logical-error counts as upper bounds, and note that the previous logical-error section could not resolve the ~1e-7 sampling floor. Co-Authored-By: Claude Opus 5.5 --- scripts/statevector/results.md | 84 +++++++++++++++++++++------------- 1 file changed, 51 insertions(+), 33 deletions(-) diff --git a/scripts/statevector/results.md b/scripts/statevector/results.md index 87d81bdf..0671b3b6 100644 --- a/scripts/statevector/results.md +++ b/scripts/statevector/results.md @@ -4,11 +4,13 @@ - **Machine:** Apple M4 Pro, 14 cores (10 performance, 4 efficiency) - **Backend:** JAX CPU, default settings (no `XLA_FLAGS` or thread limits), so XLA may use every core -- **Threads used:** the full run took 219.23 s user + 14.67 s system over - 82.98 s wall (shell `time`), about 2.8 cores busy on average. Tsim's shots/s +- **Threads used:** the full run took 227.61 s user + 14.68 s system over + 84.06 s wall (shell `time`), about 2.9 cores busy on average. Tsim's shots/s below are multi-threaded, not one-core figures. - **JAX:** 0.9.2 -- **Commit:** `bf80f7e` +- **Commit:** `0d6b4ab`, which includes the executor fixes of the final review: + outcomes drawn with 2^-64 resolution, outcome probabilities computed from + the outcome weights, and dtype-aware norm warnings Command, from the repository root: @@ -17,17 +19,18 @@ Command, from the repository root: See the docstring of `benchmark.py` for what is timed and counted. Amplitudes are complex64 (the default). "warmup s" is the first sample call, which includes JIT compilation and one full sample call of the workload's shots per -call (see JIT time below). +call (see JIT time below). The script sets every batch size explicitly, so the +automatic batch-size cap (2^16 lanes) does not apply here. ## Measured | circuit | compile s | warmup s | shots/s | acceptance | Clifft acceptance | check | logical errors / accepted | sampler | |---|---:|---:|---:|---:|---:|---|---:|---| -| distillation | 0.16 | 0.61 | 2,926,016 | 0.0799 | 0.0798 | ok | 990898 / 2344917 | `CompiledDetectorSampler(40 direct, 1 statevector components (width ≤ 3, 15 actions), 0.0 graphs, 473 error channel bits, 5 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 579 B)` | -| cultivation_d3 | 0.11 | 0.74 | 1,189,716 | 0.6866 | 0.6868 | ok | 6 / 8279256 | `CompiledDetectorSampler(16 direct, 1 statevector components (width ≤ 5, 36 actions), 0.0 graphs, 268 error channel bits, 5 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 1.6 kB)` | -| coherent_d3_r3 | 0.07 | 1.02 | 570,139 | 0.3732 | 0.3732 | ok | 0 / 2152025 | `CompiledDetectorSampler(17 direct, 1 statevector components (width ≤ 5, 73 actions), 0.0 graphs, 54 error channel bits, 8 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 3.0 kB)` | -| cultivation_d5 | 0.92 | 0.69 | 594,757 | 0.1442 | 0.1440 | ok | 0 / 859817 | `CompiledDetectorSampler(93 direct, 1 statevector components (width ≤ 6, 109 actions), 0.0 graphs, 1688 error channel bits, 15 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 6.7 kB)` | -| coherent_d5_r5 | 0.74 | 18.72 | 901 | 0.0094 | 0.0078 | ok | 0 / 154 | `CompiledDetectorSampler(73 direct, 1 statevector components (width ≤ 12, 427 actions), 0.0 graphs, 234 error channel bits, 48 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 26.0 kB)` | +| distillation | 0.17 | 0.63 | 2,869,787 | 0.0799 | 0.0798 | ok | 973000 / 2302973 | `CompiledDetectorSampler(40 direct, 1 statevector components (width ≤ 3, 15 actions), 0.0 graphs, 473 error channel bits, 5 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 579 B)` | +| cultivation_d3 | 0.11 | 0.78 | 1,096,444 | 0.6866 | 0.6868 | ok | 4 / 7559234 | `CompiledDetectorSampler(16 direct, 1 statevector components (width ≤ 5, 36 actions), 0.0 graphs, 268 error channel bits, 5 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 1.6 kB)` | +| coherent_d3_r3 | 0.07 | 1.07 | 522,613 | 0.3730 | 0.3732 | ok | 0 / 1955689 | `CompiledDetectorSampler(17 direct, 1 statevector components (width ≤ 5, 73 actions), 0.0 graphs, 54 error channel bits, 8 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 3.0 kB)` | +| cultivation_d5 | 0.92 | 0.69 | 572,691 | 0.1441 | 0.1440 | ok | 0 / 831104 | `CompiledDetectorSampler(93 direct, 1 statevector components (width ≤ 6, 109 actions), 0.0 graphs, 1688 error channel bits, 15 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 6.7 kB)` | +| coherent_d5_r5 | 0.73 | 19.22 | 874 | 0.0094 | 0.0078 | ok | 0 / 154 | `CompiledDetectorSampler(73 direct, 1 statevector components (width ≤ 12, 427 actions), 0.0 graphs, 234 error channel bits, 48 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 26.0 kB)` | ## Reference @@ -43,21 +46,32 @@ call (see JIT time below). Tsim's rate is logical errors / accepted from the measured table. Clifft 0.10.0's rates are from the design spec (section 11). "Expected" is accepted -times Clifft's rate. +times Clifft's rate. A count of 0 is written as a 95% upper bound, 3/n for n +shots rounded up (n accepted for Tsim; n as quoted by the spec for Clifft). | circuit | Tsim errors / accepted | Tsim rate | Clifft 0.10.0 rate | expected errors at Clifft's rate | |---|---:|---:|---:|---:| -| distillation | 990898 / 2344917 | 0.4226 | 0.422 | 989,555 (binomial sd 756) | -| cultivation_d3 | 6 / 8279256 | 7.2e-7 | 9.6e-7 | 7.95 | -| coherent_d3_r3 | 0 / 2152025 | 0 | 5.1e-8 | 0.11 | -| cultivation_d5 | 0 / 859817 | 0 | 0 (in 3.0e7 shots) | 0 | -| coherent_d5_r5 | 0 / 154 | 0 | 0 (in 8.8e5 shots) | 0 | +| distillation | 973000 / 2302973 | 0.4225 | 0.422 | 971,855 (binomial sd 750) | +| cultivation_d3 | 4 / 7559234 | 5.3e-7 | 9.6e-7 | 7.26 | +| coherent_d3_r3 | 0 / 1955689 | < 1.6e-6 | 5.1e-8 | 0.10 | +| cultivation_d5 | 0 / 831104 | < 3.7e-6 | < 1.0e-7 (0 in 3.0e7 shots) | < 0.09 | +| coherent_d5_r5 | 0 / 154 | < 2.0e-2 | < 3.5e-6 (0 in 8.8e5 shots) | < 0.001 | -distillation is 1,343 errors (1.8 sd) above the count expected at 0.422. The +distillation is 1,145 errors (1.5 sd) above the count expected at 0.422. The reference's third-digit rounding (±0.0005) by itself moves that expected count -by up to ±1,172. For cultivation_d3, a Poisson count with mean 7.95 is 6 or -fewer with probability 0.32. For coherent_d3_r3, a Poisson count with mean -0.11 is 0 with probability 0.90. +by up to ±1,151. For cultivation_d3, a Poisson count with mean 7.26 is 4 or +fewer with probability 0.15. For coherent_d3_r3, a Poisson count with mean +0.10 is 0 with probability 0.90. + +The previous version of this section (commit `bf80f7e`) could not detect the +outcome-sampling floor that the final review found: complex64 sampling gave +an impossible outcome with probability ~1.2e-7 per measurement, which on +cultivation_d5 flipped observable 0 in about 1.0e-7 of the accepted shots. +Over the 859,817 accepted shots of that run this predicts 0.09 extra logical +errors, and the acceptance shift is far below the 5-sd check, so neither +column could show it. Resolving a floor of that size needs ~3e7 or more +accepted shots. The run above includes the fix; its counts are again +consistent with Clifft's rates, at the same limited resolution. ## JIT time @@ -69,11 +83,11 @@ also includes any other first-call overhead, and warmup s is rounded to | circuit | warmup s | shots per call | one call s | JIT estimate s | |---|---:|---:|---:|---:| -| distillation | 0.61 | 262,144 | 0.090 | 0.52 | -| cultivation_d3 | 0.74 | 262,144 | 0.220 | 0.52 | -| coherent_d3_r3 | 1.02 | 262,144 | 0.460 | 0.56 | -| cultivation_d5 | 0.69 | 65,536 | 0.110 | 0.58 | -| coherent_d5_r5 | 18.72 | 16,384 | 18.184 | 0.54 | +| distillation | 0.63 | 262,144 | 0.091 | 0.54 | +| cultivation_d3 | 0.78 | 262,144 | 0.239 | 0.54 | +| coherent_d3_r3 | 1.07 | 262,144 | 0.502 | 0.57 | +| cultivation_d5 | 0.69 | 65,536 | 0.114 | 0.58 | +| coherent_d5_r5 | 19.22 | 16,384 | 18.746 | 0.47 | ## Summary @@ -81,19 +95,23 @@ All five acceptances agree with Clifft 0.10.0 within 5 standard errors, and all five logical-error counts are consistent with Clifft's rates. Every circuit compiles in under 1 s, including the three that the stabilizer rank strategy does not compile in 120 s, and JIT takes about 0.5 s. Tsim's -multi-threaded shots/s on this machine (about 2.8 cores busy on average) is +multi-threaded shots/s on this machine (about 2.9 cores busy on average) is above Clifft's one-core EPYC figure on distillation and cultivation_d5, and -below it on cultivation_d3, coherent_d3_r3 and coherent_d5_r5 (901 against +below it on cultivation_d3, coherent_d3_r3 and coherent_d5_r5 (874 against 5.9k). This compares different machines and thread counts, so it is not a -per-core comparison. +per-core comparison. Against the previous run (commit `bf80f7e`), shots/s is +2% to 8% lower (distillation 2.93M to 2.87M, cultivation_d3 1.19M to 1.10M, +coherent_d3_r3 570k to 523k, cultivation_d5 595k to 573k, coherent_d5_r5 901 +to 874); the fixes draw two uint32 words per outcome and sum two outcome +weights per measurement. Run notes: -- coherent_d5_r5: one sample call of 16,384 shots takes about 18 s, so the +- coherent_d5_r5: one sample call of 16,384 shots takes about 19 s, so the timed window holds a single call and its acceptance check rests on 16,384 attempted shots (154 accepted). -- coherent_d5_r5 printed six `UserWarning`s from `sampler.py` ("A marginal - probability was not normalized correctly"), with deviations from 2.0e-05 to - 4.1e-05 (the warning threshold is 1e-05). With `JAX_ENABLE_X64=1` - (complex128), `--seconds 1 coherent_d5_r5` printed no warning, with - acceptance 0.0079 (129 accepted of 16,384) at 599 shots/s. +- No norm warning was printed. The previous run printed six on + coherent_d5_r5, with deviations from 2.0e-05 to 4.1e-05 against the old + 1e-05 threshold. With the fixed kernels, four batches of 4,096 lanes of + coherent_d5_r5 reach a largest norm deviation of 1.4e-06 to 1.6e-06, and + complex64 components now warn above 1e-03. From 10f01026c6021a46787ea4dd5e126b806f6f63dd Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Sun, 11 Oct 2026 13:06:10 -0400 Subject: [PATCH 16/18] Benchmark Clifft on the same machine and pin Tsim to one thread Add benchmark_clifft.py, which runs Clifft on the five clifft-bench workloads the way clifft-bench does, and workloads.py, which both benchmark scripts share. benchmark.py now reports the cores kept busy in the timed window and takes --seed; NPROC=1 pins JAX to one thread. results.md compares Tsim with Clifft 0.11.0 on one thread and on all threads of the same machine, lists both libraries' active widths, and checks the logical-error rates against Clifft and SymFT. Co-Authored-By: Claude Opus 5.5 --- scripts/statevector/benchmark.py | 56 +++--- scripts/statevector/benchmark_clifft.py | 97 ++++++++++ scripts/statevector/results.md | 245 ++++++++++++++++-------- scripts/statevector/workloads.py | 31 +++ 4 files changed, 316 insertions(+), 113 deletions(-) create mode 100644 scripts/statevector/benchmark_clifft.py create mode 100644 scripts/statevector/workloads.py diff --git a/scripts/statevector/benchmark.py b/scripts/statevector/benchmark.py index beee56c7..251b1d87 100644 --- a/scripts/statevector/benchmark.py +++ b/scripts/statevector/benchmark.py @@ -2,44 +2,30 @@ Run from the repository root: - uv run python scripts/statevector/benchmark.py [--seconds S] [NAME ...] + uv run python scripts/statevector/benchmark.py [--seconds S] [--seed N] [NAME ...] -Circuits are read from the directory in ``CLIFFT_BENCH_CIRCUITS`` (default -``/Users/rafaelhaenel/Documents/quera/clifft-bench/workloads/circuits``). For -every circuit this compiles a detector sampler with ``strategy="statevector"``, -makes one untimed warmup call (JIT), then samples for at least ``--seconds`` -with postselection on every detector, counting as clifft-bench does: attempted -shots, accepted shots (no detector fired) and logical errors (observable 0 -among accepted shots). Acceptance is checked against Clifft 0.10.0's value -(clifft-bench release-v1-20260909-130000) within 5 standard errors. +Prefix the command with ``NPROC=1`` to run JAX on one thread: XLA sizes its CPU +thread pools from ``NPROC``. + +The workloads are defined in ``workloads.py``. For every circuit this compiles +a detector sampler with ``strategy="statevector"``, makes one untimed warmup +call (JIT), then samples for at least ``--seconds`` with postselection on every +detector, counting as clifft-bench does: attempted shots, accepted shots (no +detector fired) and logical errors (observable 0 among accepted shots). +Acceptance is checked against Clifft 0.10.0's value (clifft-bench +release-v1-20260909-130000) within 5 standard errors. "cores busy" is the +process CPU time divided by wall time over the timed window. """ import argparse import math -import os import time -from pathlib import Path import numpy as np +from workloads import CIRCUITS, WORKLOADS, cpu_seconds import tsim -CIRCUITS = Path( - os.environ.get( - "CLIFFT_BENCH_CIRCUITS", - "/Users/rafaelhaenel/Documents/quera/clifft-bench/workloads/circuits", - ) -) - -# name: (file, Clifft 0.10.0 acceptance, shots per sample call, batch size) -WORKLOADS = { - "distillation": ("distillation.stim", 0.0798, 262_144, 65_536), - "cultivation_d3": ("msc_d3_inject_cultivate_p1e-3.stim", 0.6868, 262_144, 65_536), - "coherent_d3_r3": ("coherent_d3_r3.stim", 0.3732, 262_144, 65_536), - "cultivation_d5": ("msc_d5_inject_cultivate_p1e-3.stim", 0.1440, 65_536, 16_384), - "coherent_d5_r5": ("coherent_d5_r5.stim", 0.0078, 16_384, 4_096), -} - def count(sampler, shots: int, batch_size: int, num_detectors: int) -> tuple[int, int]: """Return (accepted shots, logical errors) of one sample call.""" @@ -58,22 +44,26 @@ def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("names", nargs="*", default=list(WORKLOADS)) parser.add_argument("--seconds", type=float, default=10.0) + parser.add_argument("--seed", type=int, default=0) args = parser.parse_args() print( - "| circuit | compile s | warmup s | shots/s | acceptance | Clifft " - "acceptance | check | logical errors / accepted | sampler |" + "| circuit | compile s | warmup s | shots/s | cores busy | acceptance | " + "Clifft acceptance | check | logical errors / accepted | sampler |" ) - print("|---|---:|---:|---:|---:|---:|---|---:|---|") + print("|---|---:|---:|---:|---:|---:|---:|---|---:|---|") for name in args.names: file, reference, shots, batch_size = WORKLOADS[name] circuit = tsim.Circuit.from_file(str(CIRCUITS / file)) start = time.perf_counter() - sampler = circuit.compile_detector_sampler(strategy="statevector", seed=0) + sampler = circuit.compile_detector_sampler( + strategy="statevector", seed=args.seed + ) compile_seconds = time.perf_counter() - start start = time.perf_counter() count(sampler, shots, batch_size, circuit.num_detectors) warmup_seconds = time.perf_counter() - start attempted = accepted = errors = 0 + cpu_start = cpu_seconds() start = time.perf_counter() while time.perf_counter() - start < args.seconds: kept, wrong = count(sampler, shots, batch_size, circuit.num_detectors) @@ -81,12 +71,14 @@ def main() -> None: accepted += kept errors += wrong elapsed = time.perf_counter() - start + cores = (cpu_seconds() - cpu_start) / elapsed acceptance = accepted / attempted sigma = math.sqrt(reference * (1 - reference) / attempted) check = "ok" if abs(acceptance - reference) <= 5 * sigma else "MISMATCH" print( f"| {name} | {compile_seconds:.2f} | {warmup_seconds:.2f} | " - f"{attempted / elapsed:,.0f} | {acceptance:.4f} | {reference:.4f} | " + f"{attempted / elapsed:,.0f} | {cores:.1f} | {acceptance:.4f} | " + f"{reference:.4f} | " f"{check} | {errors} / {accepted} | `{sampler!r}` |", flush=True, ) diff --git a/scripts/statevector/benchmark_clifft.py b/scripts/statevector/benchmark_clifft.py new file mode 100644 index 00000000..627ba157 --- /dev/null +++ b/scripts/statevector/benchmark_clifft.py @@ -0,0 +1,97 @@ +"""Benchmark Clifft on the clifft-bench circuits, as a same-machine reference. + +Clifft is not a tsim dependency. Install it into a separate environment and run +this script with that interpreter, from the repository root: + + uv venv /tmp/clifft-venv && uv pip install --python /tmp/clifft-venv clifft + /tmp/clifft-venv/bin/python scripts/statevector/benchmark_clifft.py [--seconds S] [NAME ...] + +Every circuit is compiled as clifft-bench does: default HIR passes, every +detector post-selected. Then, for one thread and for all threads +(``threads="auto"``), it samples with ``sample_survivors`` for at least +``--seconds`` after one untimed call, counting attempted shots, accepted shots +and logical errors (observable 0 among accepted shots). Each thread setting is +timed with scalar execution and with packed lanes of 2048 shots (unless Clifft +rejects that batch as too large), and the faster one is reported, as +clifft-bench's batch calibration does. "cores busy" is the process CPU time +divided by wall time over the timed window. +""" + +import argparse +import math +import time + +import clifft +from workloads import CIRCUITS, WORKLOADS, cpu_seconds + +BATCH_SIZES = (1, 2048) + + +def measure(program, shots: int, seconds: float, threads, batch_size: int): + """Return (shots/s, cores busy, attempted, accepted, logical errors).""" + clifft.sample_survivors( + program, shots, seed=0, threads=threads, batch_size=batch_size + ) + attempted = accepted = errors = calls = 0 + cpu_start = cpu_seconds() + start = time.perf_counter() + while time.perf_counter() - start < seconds: + calls += 1 + result = clifft.sample_survivors( + program, shots, seed=calls, threads=threads, batch_size=batch_size + ) + attempted += int(result.total_shots) + accepted += int(result.passed_shots) + errors += int(result.observable_ones[0]) + elapsed = time.perf_counter() - start + cores = (cpu_seconds() - cpu_start) / elapsed + return attempted / elapsed, cores, attempted, accepted, errors + + +def main() -> None: + """Benchmark the requested workloads and print one Markdown row per thread setting.""" + parser = argparse.ArgumentParser() + parser.add_argument("names", nargs="*", default=list(WORKLOADS)) + parser.add_argument("--seconds", type=float, default=10.0) + args = parser.parse_args() + print(f"Clifft {clifft.__version__}") + print( + "| circuit | threads | batch | compile s | k_max | shots/s | cores busy | " + "acceptance | Clifft 0.10.0 acceptance | check | logical errors / accepted |" + ) + print("|---|---|---:|---:|---:|---:|---:|---:|---:|---|---:|") + for name in args.names: + file, reference, shots, _ = WORKLOADS[name] + text = (CIRCUITS / file).read_text() + start = time.perf_counter() + hir = clifft.trace(clifft.parse(text)) + clifft.default_hir_pass_manager().run(hir) + program = clifft.lower(hir, postselection_mask=[1] * int(hir.num_detectors)) + compile_seconds = time.perf_counter() - start + for threads in (1, "auto"): + runs = {} + for batch_size in BATCH_SIZES: + try: + runs[batch_size] = measure( + program, shots, args.seconds, threads, batch_size + ) + except ValueError: + # NOTE: Clifft rejects packed batches above its 64 MiB state limit. + continue + batch_size, (rate, cores, attempted, accepted, errors) = max( + runs.items(), key=lambda item: item[1][0] + ) + acceptance = accepted / attempted + sigma = math.sqrt(reference * (1 - reference) / attempted) + check = "ok" if abs(acceptance - reference) <= 5 * sigma else "MISMATCH" + print( + f"| {name} | {threads} | {batch_size} | {compile_seconds:.2f} | " + f"{program.peak_active_width} | {rate:,.0f} | {cores:.1f} | " + f"{acceptance:.4f} | {reference:.4f} | {check} | " + f"{errors} / {accepted} |", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/statevector/results.md b/scripts/statevector/results.md index 0671b3b6..b83b6916 100644 --- a/scripts/statevector/results.md +++ b/scripts/statevector/results.md @@ -2,35 +2,110 @@ - **Date:** 2026-10-10 - **Machine:** Apple M4 Pro, 14 cores (10 performance, 4 efficiency) -- **Backend:** JAX CPU, default settings (no `XLA_FLAGS` or thread limits), so - XLA may use every core -- **Threads used:** the full run took 227.61 s user + 14.68 s system over - 84.06 s wall (shell `time`), about 2.9 cores busy on average. Tsim's shots/s - below are multi-threaded, not one-core figures. -- **JAX:** 0.9.2 -- **Commit:** `0d6b4ab`, which includes the executor fixes of the final review: - outcomes drawn with 2^-64 resolution, outcome probabilities computed from - the outcome weights, and dtype-aware norm warnings - -Command, from the repository root: +- **Tsim:** commit `883b125` (the sampler code is unchanged since `0d6b4ab`), + JAX 0.9.2 on CPU, complex64 amplitudes (the default) +- **Clifft:** 0.11.0 from PyPI, in its own environment, on the same machine +- **Threads:** "all threads" is each library's default: XLA's CPU thread pools + for Tsim, `threads="auto"` for Clifft. "1 thread" is `NPROC=1` for Tsim (XLA + sizes its CPU thread pools from `NPROC`; macOS has no CPU affinity) and + `threads=1` for Clifft. "cores busy" is the process CPU time divided by wall + time over the timed window; it is 1.0 in every 1-thread row. + +Commands, from the repository root (`benchmark_clifft.py` describes how to +install Clifft): uv run python scripts/statevector/benchmark.py --seconds 10 - -See the docstring of `benchmark.py` for what is timed and counted. Amplitudes -are complex64 (the default). "warmup s" is the first sample call, which -includes JIT compilation and one full sample call of the workload's shots per -call (see JIT time below). The script sets every batch size explicitly, so the -automatic batch-size cap (2^16 lanes) does not apply here. + NPROC=1 uv run python scripts/statevector/benchmark.py --seconds 10 + /tmp/clifft-venv/bin/python scripts/statevector/benchmark_clifft.py --seconds 10 + +See the docstrings of both scripts for what is timed and counted. Both +postselect on every detector, as clifft-bench does. Tsim's "warmup s" is the +first sample call, which includes JIT compilation and one full sample call of +the workload's shots per call (see JIT time below). The scripts set every batch +size explicitly, so Tsim's automatic batch-size cap (2^16 lanes) does not apply. + +## Same-machine comparison + +Shots/s from the measured tables below. Tsim's width is the width of its state +vector; Clifft's k_max is its peak active width under its default passes. + +| circuit | Tsim width | Clifft k_max | Tsim, 1 thread | Clifft, 1 thread | Tsim, all threads | Clifft, all threads | +|---|---:|---:|---:|---:|---:|---:| +| distillation | 3 | 5 | 2,462,770 | 2,755,248 | 2,869,491 | 20,196,109 | +| cultivation_d3 | 5 | 4 | 562,902 | 4,742,046 | 1,108,828 | 30,405,370 | +| coherent_d3_r3 | 5 | 5 | 248,544 | 1,920,603 | 527,961 | 13,664,724 | +| cultivation_d5 | 6 | 10 | 312,294 | 136,869 | 572,464 | 1,207,485 | +| coherent_d5_r5 | 12 | 13 | 272 | 3,692 | 886 | 25,476 | + +Clifft's shots/s divided by Tsim's, and the cores each keeps busy on all +threads: + +| circuit | 1 thread | all threads | cores busy, Tsim / Clifft | +|---|---:|---:|---:| +| distillation | 1.1 | 7.0 | 1.3 / 8.7 | +| cultivation_d3 | 8.4 | 27 | 2.5 / 8.4 | +| coherent_d3_r3 | 7.7 | 26 | 2.9 / 9.1 | +| cultivation_d5 | 0.44 | 2.1 | 2.4 / 10.7 | +| coherent_d5_r5 | 13.6 | 29 | 3.7 / 9.5 | ## Measured -| circuit | compile s | warmup s | shots/s | acceptance | Clifft acceptance | check | logical errors / accepted | sampler | -|---|---:|---:|---:|---:|---:|---|---:|---| -| distillation | 0.17 | 0.63 | 2,869,787 | 0.0799 | 0.0798 | ok | 973000 / 2302973 | `CompiledDetectorSampler(40 direct, 1 statevector components (width ≤ 3, 15 actions), 0.0 graphs, 473 error channel bits, 5 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 579 B)` | -| cultivation_d3 | 0.11 | 0.78 | 1,096,444 | 0.6866 | 0.6868 | ok | 4 / 7559234 | `CompiledDetectorSampler(16 direct, 1 statevector components (width ≤ 5, 36 actions), 0.0 graphs, 268 error channel bits, 5 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 1.6 kB)` | -| coherent_d3_r3 | 0.07 | 1.07 | 522,613 | 0.3730 | 0.3732 | ok | 0 / 1955689 | `CompiledDetectorSampler(17 direct, 1 statevector components (width ≤ 5, 73 actions), 0.0 graphs, 54 error channel bits, 8 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 3.0 kB)` | -| cultivation_d5 | 0.92 | 0.69 | 572,691 | 0.1441 | 0.1440 | ok | 0 / 831104 | `CompiledDetectorSampler(93 direct, 1 statevector components (width ≤ 6, 109 actions), 0.0 graphs, 1688 error channel bits, 15 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 6.7 kB)` | -| coherent_d5_r5 | 0.73 | 19.22 | 874 | 0.0094 | 0.0078 | ok | 0 / 154 | `CompiledDetectorSampler(73 direct, 1 statevector components (width ≤ 12, 427 actions), 0.0 graphs, 234 error channel bits, 48 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 26.0 kB)` | +### Tsim, all threads + +| circuit | compile s | warmup s | shots/s | cores busy | acceptance | Clifft acceptance | check | logical errors / accepted | sampler | +|---|---:|---:|---:|---:|---:|---:|---|---:|---| +| distillation | 0.17 | 0.63 | 2,869,491 | 1.3 | 0.0799 | 0.0798 | ok | 973000 / 2302973 | `CompiledDetectorSampler(40 direct, 1 statevector components (width ≤ 3, 15 actions), 0.0 graphs, 473 error channel bits, 5 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 579 B)` | +| cultivation_d3 | 0.11 | 0.76 | 1,108,828 | 2.5 | 0.6866 | 0.6868 | ok | 4 / 7739165 | `CompiledDetectorSampler(16 direct, 1 statevector components (width ≤ 5, 36 actions), 0.0 graphs, 268 error channel bits, 5 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 1.6 kB)` | +| coherent_d3_r3 | 0.07 | 1.05 | 527,961 | 2.9 | 0.3731 | 0.3732 | ok | 0 / 2053991 | `CompiledDetectorSampler(17 direct, 1 statevector components (width ≤ 5, 73 actions), 0.0 graphs, 54 error channel bits, 8 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 3.0 kB)` | +| cultivation_d5 | 0.92 | 0.66 | 572,464 | 2.4 | 0.1441 | 0.1440 | ok | 0 / 831104 | `CompiledDetectorSampler(93 direct, 1 statevector components (width ≤ 6, 109 actions), 0.0 graphs, 1688 error channel bits, 15 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 6.7 kB)` | +| coherent_d5_r5 | 0.73 | 19.32 | 886 | 3.7 | 0.0094 | 0.0078 | ok | 0 / 154 | `CompiledDetectorSampler(73 direct, 1 statevector components (width ≤ 12, 427 actions), 0.0 graphs, 234 error channel bits, 48 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 26.0 kB)` | + +### Tsim, 1 thread (`NPROC=1`) + +| circuit | compile s | warmup s | shots/s | cores busy | acceptance | Clifft acceptance | check | logical errors / accepted | sampler | +|---|---:|---:|---:|---:|---:|---:|---|---:|---| +| distillation | 0.17 | 0.64 | 2,462,770 | 1.0 | 0.0798 | 0.0798 | ok | 831265 / 1967388 | `CompiledDetectorSampler(40 direct, 1 statevector components (width ≤ 3, 15 actions), 0.0 graphs, 473 error channel bits, 5 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 579 B)` | +| cultivation_d3 | 0.12 | 1.00 | 562,902 | 1.0 | 0.6867 | 0.6868 | ok | 0 / 3960347 | `CompiledDetectorSampler(16 direct, 1 statevector components (width ≤ 5, 36 actions), 0.0 graphs, 268 error channel bits, 5 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 1.6 kB)` | +| coherent_d3_r3 | 0.07 | 1.57 | 248,544 | 1.0 | 0.3733 | 0.3732 | ok | 0 / 978504 | `CompiledDetectorSampler(17 direct, 1 statevector components (width ≤ 5, 73 actions), 0.0 graphs, 54 error channel bits, 8 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 3.0 kB)` | +| cultivation_d5 | 0.92 | 0.75 | 312,294 | 1.0 | 0.1440 | 0.1440 | ok | 0 / 453106 | `CompiledDetectorSampler(93 direct, 1 statevector components (width ≤ 6, 109 actions), 0.0 graphs, 1688 error channel bits, 15 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 6.7 kB)` | +| coherent_d5_r5 | 0.75 | 58.14 | 272 | 1.0 | 0.0094 | 0.0078 | ok | 0 / 154 | `CompiledDetectorSampler(73 direct, 1 statevector components (width ≤ 12, 427 actions), 0.0 graphs, 234 error channel bits, 48 outputs for largest cc, ≤ 0 parameters, 0.0 A terms, 0.0 B terms, 0.0 C terms, 0.0 D terms, 0.0 E terms, 26.0 kB)` | + +### Clifft 0.11.0 + +"batch" is the faster of scalar execution (1) and packed lanes of 2048 shots. +Clifft rejects packed lanes of 2048 on coherent_d5_r5 (they exceed its 64 MiB +packed-state limit at k_max 13), so those rows are scalar. + +| circuit | threads | batch | compile s | k_max | shots/s | cores busy | acceptance | Clifft 0.10.0 acceptance | check | logical errors / accepted | +|---|---|---:|---:|---:|---:|---:|---:|---:|---|---:| +| distillation | 1 | 2048 | 0.00 | 5 | 2,755,248 | 1.0 | 0.0798 | 0.0798 | ok | 935589 / 2217580 | +| distillation | auto | 2048 | 0.00 | 5 | 20,196,109 | 8.7 | 0.0797 | 0.0798 | ok | 6798775 / 16115015 | +| cultivation_d3 | 1 | 2048 | 0.00 | 4 | 4,742,046 | 1.0 | 0.6868 | 0.6868 | ok | 42 / 32589642 | +| cultivation_d3 | auto | 2048 | 0.00 | 4 | 30,405,370 | 8.4 | 0.6868 | 0.6868 | ok | 222 / 208838618 | +| coherent_d3_r3 | 1 | 2048 | 0.00 | 5 | 1,920,603 | 1.0 | 0.3732 | 0.3732 | ok | 0 / 7239813 | +| coherent_d3_r3 | auto | 2048 | 0.00 | 5 | 13,664,724 | 9.1 | 0.3732 | 0.3732 | ok | 1 / 51070460 | +| cultivation_d5 | 1 | 1 | 0.01 | 10 | 136,869 | 1.0 | 0.1444 | 0.1440 | ok | 0 / 198758 | +| cultivation_d5 | auto | 1 | 0.01 | 10 | 1,207,485 | 10.7 | 0.1440 | 0.1440 | ok | 0 / 1746413 | +| coherent_d5_r5 | 1 | 1 | 0.01 | 13 | 3,692 | 1.0 | 0.0075 | 0.0078 | ok | 0 / 369 | +| coherent_d5_r5 | auto | 1 | 0.01 | 13 | 25,476 | 9.5 | 0.0074 | 0.0078 | ok | 0 / 1931 | + +## Active width + +| circuit | Clifft, no passes | Clifft, default passes | Clifft + `ActiveWidthSchedulePass` | Tsim, detectors | Tsim, measurements | +|---|---:|---:|---:|---:|---:| +| distillation | 5 | 5 | 3 | 3 | 3 | +| cultivation_d3 | 4 | 4 | 4 | 5 | 5 | +| coherent_d3_r3 | 8 | 5 | 4 | 5 | 4 | +| cultivation_d5 | 10 | 10 | 10 | 6 | 6 | +| coherent_d5_r5 | 24 | 13 | 13 | 12 | 15 | + +Clifft's columns are `peak_active_width` of `clifft.compile(text, +hir_passes=None)`, of `clifft.compile(text)`, and of `clifft.compile` with +`default_hir_pass_manager()` plus `ActiveWidthSchedulePass()`, which is opt-in +in 0.11.0. The benchmark uses the default passes. The Clifft paper's table +matches the no-pass column. Tsim's columns are the largest `width` over the +statevector components of `compile_detector_sampler` (the benchmark) and of +`compile_sampler`. ## Reference @@ -42,76 +117,84 @@ automatic batch-size cap (2^16 lanes) does not apply here. | cultivation_d5 | 197k | does not compile in 120 s | | coherent_d5_r5 | 5.9k | does not compile in 120 s | +Clifft 0.10.0's figures are from clifft-bench release-v1-20260909-130000, on a +different machine. + ## Logical-error rate -Tsim's rate is logical errors / accepted from the measured table. Clifft -0.10.0's rates are from the design spec (section 11). "Expected" is accepted -times Clifft's rate. A count of 0 is written as a 95% upper bound, 3/n for n -shots rounded up (n accepted for Tsim; n as quoted by the spec for Clifft). +Tsim's counts pool the benchmark's run (seed 0) with independent runs: -| circuit | Tsim errors / accepted | Tsim rate | Clifft 0.10.0 rate | expected errors at Clifft's rate | -|---|---:|---:|---:|---:| -| distillation | 973000 / 2302973 | 0.4225 | 0.422 | 971,855 (binomial sd 750) | -| cultivation_d3 | 4 / 7559234 | 5.3e-7 | 9.6e-7 | 7.26 | -| coherent_d3_r3 | 0 / 1955689 | < 1.6e-6 | 5.1e-8 | 0.10 | -| cultivation_d5 | 0 / 831104 | < 3.7e-6 | < 1.0e-7 (0 in 3.0e7 shots) | < 0.09 | -| coherent_d5_r5 | 0 / 154 | < 2.0e-2 | < 3.5e-6 (0 in 8.8e5 shots) | < 0.001 | - -distillation is 1,145 errors (1.5 sd) above the count expected at 0.422. The -reference's third-digit rounding (±0.0005) by itself moves that expected count -by up to ±1,151. For cultivation_d3, a Poisson count with mean 7.26 is 4 or -fewer with probability 0.15. For coherent_d3_r3, a Poisson count with mean -0.10 is 0 with probability 0.90. - -The previous version of this section (commit `bf80f7e`) could not detect the -outcome-sampling floor that the final review found: complex64 sampling gave -an impossible outcome with probability ~1.2e-7 per measurement, which on -cultivation_d5 flipped observable 0 in about 1.0e-7 of the accepted shots. -Over the 859,817 accepted shots of that run this predicts 0.09 extra logical -errors, and the acceptance shift is far below the 5-sd check, so neither -column could show it. Resolving a floor of that size needs ~3e7 or more -accepted shots. The run above includes the fix; its counts are again -consistent with Clifft's rates, at the same limited resolution. + uv run python scripts/statevector/benchmark.py --seconds 120 --seed 1 cultivation_d3 + uv run python scripts/statevector/benchmark.py --seconds 45 --seed 1 distillation + uv run python scripts/statevector/benchmark.py --seconds 45 --seed 2 distillation + +which gave 78 / 89470166, 4242388 / 10057137 and 4275565 / 10138745 (errors / +accepted). The references are Clifft 0.11.0 on this machine (both thread +settings pooled), and Clifft 0.10.0 and SymFT from the raw data of clifft-bench +release-v1-20260909-130000. z is Tsim's deviation from each reference in +standard deviations, counting the statistical error of both. A count of 0 is +written as a 95% upper bound, 3/n. + +| circuit | Tsim errors / accepted | Tsim rate | Clifft 0.11.0 | Clifft 0.10.0 | SymFT | +|---|---:|---:|---:|---:|---:| +| distillation | 9490953 / 22498855 | 0.42184 | 0.42189 (z −0.3) | 0.42152 (z +2.2) | 0.42158 (z +1.8) | +| cultivation_d3 | 82 / 97209331 | 8.4e-7 | 1.09e-6 (z −2.0) | 9.6e-7 (z −1.0) | 8.5e-7 (z −0.1) | +| coherent_d3_r3 | 0 / 2053991 | < 1.5e-6 | 1 / 58310273 | 3 / 58747248 | 1 / 31579774 | +| cultivation_d5 | 0 / 831104 | < 3.6e-6 | 0 / 1945171 | 0 / 4272772 | 0 / 3685293 | +| coherent_d5_r5 | 0 / 154 | < 2.0e-2 | 0 / 2300 | 0 / 6851 | 0 / 63 | + +The references scatter by as much among themselves: Clifft 0.11.0 is 2.4 sd +from Clifft 0.10.0 on distillation and 2.8 sd from SymFT on cultivation_d3, and +Tsim lies within that scatter. For coherent_d3_r3, the pooled rate of Clifft +0.10.0 and 0.11.0 (4 / 117057521) predicts 0.07 Tsim errors. + +The outcome-sampling floor that the final review found (complex64 drew an +impossible outcome with probability ~1.2e-7 per measurement) was fixed in +`0d6b4ab`, before all of these runs. ## JIT time -The script does not time JIT separately. Estimate it as warmup s minus one -call's time, where one call's time is shots per call divided by the measured -shots/s. This assumes the warmup call samples at the timed rate. The estimate -also includes any other first-call overhead, and warmup s is rounded to -0.01 s. +The script does not time JIT separately. Estimate it from the all-threads run +as warmup s minus one call's time, where one call's time is shots per call +divided by the measured shots/s. This assumes the warmup call samples at the +timed rate. The estimate also includes any other first-call overhead, and +warmup s is rounded to 0.01 s. | circuit | warmup s | shots per call | one call s | JIT estimate s | |---|---:|---:|---:|---:| | distillation | 0.63 | 262,144 | 0.091 | 0.54 | -| cultivation_d3 | 0.78 | 262,144 | 0.239 | 0.54 | -| coherent_d3_r3 | 1.07 | 262,144 | 0.502 | 0.57 | -| cultivation_d5 | 0.69 | 65,536 | 0.114 | 0.58 | -| coherent_d5_r5 | 19.22 | 16,384 | 18.746 | 0.47 | +| cultivation_d3 | 0.76 | 262,144 | 0.236 | 0.52 | +| coherent_d3_r3 | 1.05 | 262,144 | 0.497 | 0.55 | +| cultivation_d5 | 0.66 | 65,536 | 0.114 | 0.55 | +| coherent_d5_r5 | 19.32 | 16,384 | 18.492 | 0.83 | ## Summary -All five acceptances agree with Clifft 0.10.0 within 5 standard errors, and -all five logical-error counts are consistent with Clifft's rates. Every -circuit compiles in under 1 s, including the three that the stabilizer rank -strategy does not compile in 120 s, and JIT takes about 0.5 s. Tsim's -multi-threaded shots/s on this machine (about 2.9 cores busy on average) is -above Clifft's one-core EPYC figure on distillation and cultivation_d5, and -below it on cultivation_d3, coherent_d3_r3 and coherent_d5_r5 (874 against -5.9k). This compares different machines and thread counts, so it is not a -per-core comparison. Against the previous run (commit `bf80f7e`), shots/s is -2% to 8% lower (distillation 2.93M to 2.87M, cultivation_d3 1.19M to 1.10M, -coherent_d3_r3 570k to 523k, cultivation_d5 595k to 573k, coherent_d5_r5 901 -to 874); the fixes draw two uint32 words per outcome and sum two outcome -weights per measurement. +All acceptances, Tsim's on both thread settings and Clifft 0.11.0's, agree with +Clifft 0.10.0's within 5 standard errors, and Tsim's logical-error rates agree +with the exact references within their scatter. On this machine Clifft 0.11.0 +is faster than Tsim on four of the five circuits. On one thread it is 1.1× +faster on distillation, 7.7× and 8.4× on coherent_d3_r3 and cultivation_d3, and +13.6× on coherent_d5_r5. Tsim is faster only on cultivation_d5, by 2.3× on one +thread, where its width is 6 against Clifft's 10. On all threads the gaps grow +(2.1× on cultivation_d5, 7× to 29× on the others), because Clifft keeps 8 to 11 +cores busy and Tsim 1.3 to 3.7. Tsim compiles in 0.07 to 0.92 s, Clifft in +0.01 s or less. + +The previous version of this file compared Tsim on all threads with Clifft +0.10.0's one-core EPYC figures and found Tsim faster on distillation and +cultivation_d5. Measured on the same machine against Clifft 0.11.0, that holds +only for cultivation_d5, and only per thread. Run notes: -- coherent_d5_r5: one sample call of 16,384 shots takes about 19 s, so the - timed window holds a single call and its acceptance check rests on 16,384 - attempted shots (154 accepted). -- No norm warning was printed. The previous run printed six on - coherent_d5_r5, with deviations from 2.0e-05 to 4.1e-05 against the old - 1e-05 threshold. With the fixed kernels, four batches of 4,096 lanes of - coherent_d5_r5 reach a largest norm deviation of 1.4e-06 to 1.6e-06, and - complex64 components now warn above 1e-03. +- coherent_d5_r5: one Tsim sample call of 16,384 shots takes about 18 s on all + threads and 60 s on one, so each timed window holds a single call and the + acceptance check rests on 16,384 attempted shots (154 accepted). +- Clifft's coherent_d5_r5 rows come from a rerun of that workload, after the + script learned to skip a batch size that Clifft rejects; the other rows are + unaffected by that change. +- No norm warning was printed. +- Whole runs (shell `time`): Tsim on all threads took 84.4 s wall, 226.6 s user + and 15.1 s system; Tsim on one thread took 166.2 s wall, 161.4 s user and + 4.5 s system. diff --git a/scripts/statevector/workloads.py b/scripts/statevector/workloads.py new file mode 100644 index 00000000..1ff5b89e --- /dev/null +++ b/scripts/statevector/workloads.py @@ -0,0 +1,31 @@ +"""The clifft-bench workloads shared by ``benchmark.py`` and ``benchmark_clifft.py``. + +Circuits are read from the directory in ``CLIFFT_BENCH_CIRCUITS`` (default +``/Users/rafaelhaenel/Documents/quera/clifft-bench/workloads/circuits``). +""" + +import os +import resource +from pathlib import Path + +CIRCUITS = Path( + os.environ.get( + "CLIFFT_BENCH_CIRCUITS", + "/Users/rafaelhaenel/Documents/quera/clifft-bench/workloads/circuits", + ) +) + +# name: (file, Clifft 0.10.0 acceptance, shots per sample call, batch size) +WORKLOADS = { + "distillation": ("distillation.stim", 0.0798, 262_144, 65_536), + "cultivation_d3": ("msc_d3_inject_cultivate_p1e-3.stim", 0.6868, 262_144, 65_536), + "coherent_d3_r3": ("coherent_d3_r3.stim", 0.3732, 262_144, 65_536), + "cultivation_d5": ("msc_d5_inject_cultivate_p1e-3.stim", 0.1440, 65_536, 16_384), + "coherent_d5_r5": ("coherent_d5_r5.stim", 0.0078, 16_384, 4_096), +} + + +def cpu_seconds() -> float: + """Return the CPU time of this process (user plus system) in seconds.""" + usage = resource.getrusage(resource.RUSAGE_SELF) + return usage.ru_utime + usage.ru_stime From 0c638b1f6d72101bb6b1e2a6acf1fbae5d11b207 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Sun, 11 Oct 2026 23:47:10 -0400 Subject: [PATCH 17/18] Port the statevector sampler to zxparam zx_param is now zxparam: rename the imports in the statevector modules, their tests and the test helper. Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/__init__.py | 2 +- src/tsim/statevector/heisenberg.py | 4 ++-- src/tsim/statevector/program.py | 4 ++-- test/helpers/statevector.py | 4 ++-- test/unit/statevector/test_heisenberg.py | 4 ++-- test/unit/statevector/test_plan.py | 4 ++-- 6 files changed, 11 insertions(+), 11 deletions(-) diff --git a/src/tsim/statevector/__init__.py b/src/tsim/statevector/__init__.py index 6cac2946..5c66c918 100644 --- a/src/tsim/statevector/__init__.py +++ b/src/tsim/statevector/__init__.py @@ -1,7 +1,7 @@ """Statevector sampling of Tsim's quantum components (``strategy="statevector"``). A non-direct component of the reduced sampling graph is extracted as a circuit -(``zx_param.extract``), rewritten as rotations and measurements in the frame of +(``zxparam.extract``), rewritten as rotations and measurements in the frame of the initial ``|0…0⟩`` state (``heisenberg``), reordered to keep the dense array small (``schedule``), planned on a fixed number of array axes (``plan``), lowered to JAX arrays (``program``) and interpreted lane by lane (``execute``). diff --git a/src/tsim/statevector/heisenberg.py b/src/tsim/statevector/heisenberg.py index d25ed60d..b3a60583 100644 --- a/src/tsim/statevector/heisenberg.py +++ b/src/tsim/statevector/heisenberg.py @@ -18,8 +18,8 @@ import numpy as np import stim -from zx_param import Graph, Vertex -from zx_param.extract import ExtractedCircuit, Gate +from zxparam import Graph, Vertex +from zxparam.extract import ExtractedCircuit, Gate from tsim.core.params import ParamLayout diff --git a/src/tsim/statevector/program.py b/src/tsim/statevector/program.py index e7616d61..a6a0305c 100644 --- a/src/tsim/statevector/program.py +++ b/src/tsim/statevector/program.py @@ -12,8 +12,8 @@ import jax.numpy as jnp import numpy as np from jax import Array -from zx_param import Graph -from zx_param.extract import extract_circuit +from zxparam import Graph +from zxparam.extract import extract_circuit from tsim.core.params import ParamLayout from tsim.statevector.heisenberg import ComponentLayout, heisenberg_ops diff --git a/test/helpers/statevector.py b/test/helpers/statevector.py index 80fb3366..602c4ba2 100644 --- a/test/helpers/statevector.py +++ b/test/helpers/statevector.py @@ -12,8 +12,8 @@ import jax.numpy as jnp import numpy as np -from zx_param import Graph, Phase, Scalar -from zx_param.extract import ExtractedCircuit, Gate, extract_circuit +from zxparam import Graph, Phase, Scalar +from zxparam.extract import ExtractedCircuit, Gate, extract_circuit from tsim.circuit import Circuit from tsim.compile.evaluate import evaluate diff --git a/test/unit/statevector/test_heisenberg.py b/test/unit/statevector/test_heisenberg.py index 0fb571ad..7ac2009b 100644 --- a/test/unit/statevector/test_heisenberg.py +++ b/test/unit/statevector/test_heisenberg.py @@ -11,8 +11,8 @@ ) import pytest -from zx_param import Phase, Scalar -from zx_param.extract import ExtractedCircuit, Gate, extract_circuit +from zxparam import Phase, Scalar +from zxparam.extract import ExtractedCircuit, Gate, extract_circuit from tsim.statevector.heisenberg import ( MEASUREMENT, diff --git a/test/unit/statevector/test_plan.py b/test/unit/statevector/test_plan.py index ed690189..667327c6 100644 --- a/test/unit/statevector/test_plan.py +++ b/test/unit/statevector/test_plan.py @@ -11,8 +11,8 @@ import numpy as np import pytest import stim -from zx_param import Phase, Scalar -from zx_param.extract import ExtractedCircuit, Gate +from zxparam import Phase, Scalar +from zxparam.extract import ExtractedCircuit, Gate from tsim.statevector.heisenberg import ( MEASUREMENT, From ec8caf677528f9bcdd916eaaafe96eb53e922773 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Mon, 12 Oct 2026 12:40:31 -0400 Subject: [PATCH 18/18] Simplify the statevector sampler Read every output off the final symbols instead of carrying records through the tape: an output is the measurement's sign, XOR its branch bit when the measurement is random. This removes the Record action, the RECORD kernel, slots, the scratch record column and the sign of measurement actions. - Dispatch on the component type in sample_component and give each component type its own norm-warning threshold, so sample_program is the loop it is on main. - Pick the per-component compiler once in compile_program and reuse its output indices and f-bit selection in ComponentLayout.of. - Share one anticommute helper across heisenberg, schedule and plan, and one elimination step between StabilizerGroup.rotate and measure. - Derive num_lines, num_f, num_words, num_outputs and num_branches instead of storing them; drop the unused dtype option. - Gather once in the diagonal measurement kernel and jit sample_statevector_component as a whole. - Plan on ops in schedule order and emit each Rotate in one place. - Share test fixtures, merge duplicated tests and drop subsumed ones. - Default the benchmark circuits to a clifft-bench checkout next to the repository instead of a personal path. Co-Authored-By: Claude Opus 5.5 --- scripts/statevector/benchmark.py | 10 +- scripts/statevector/benchmark_clifft.py | 10 +- scripts/statevector/workloads.py | 16 +- src/tsim/compile/pipeline.py | 23 +-- src/tsim/core/types.py | 5 +- src/tsim/sampler.py | 92 ++++----- src/tsim/statevector/execute.py | 195 ++++++++++--------- src/tsim/statevector/heisenberg.py | 144 ++++++-------- src/tsim/statevector/plan.py | 167 +++++++--------- src/tsim/statevector/program.py | 126 ++++++------ src/tsim/statevector/schedule.py | 42 ++-- test/helpers/statevector.py | 83 ++++---- test/integration/test_statevector_sampler.py | 33 +--- test/unit/statevector/test_execute.py | 174 +++++++++-------- test/unit/statevector/test_heisenberg.py | 60 ++---- test/unit/statevector/test_plan.py | 66 +++---- test/unit/statevector/test_program.py | 87 +++------ test/unit/statevector/test_schedule.py | 8 +- 18 files changed, 615 insertions(+), 726 deletions(-) diff --git a/scripts/statevector/benchmark.py b/scripts/statevector/benchmark.py index 251b1d87..aa294bbd 100644 --- a/scripts/statevector/benchmark.py +++ b/scripts/statevector/benchmark.py @@ -18,11 +18,10 @@ """ import argparse -import math import time import numpy as np -from workloads import CIRCUITS, WORKLOADS, cpu_seconds +from workloads import CIRCUITS, WORKLOADS, acceptance_check import tsim @@ -63,7 +62,7 @@ def main() -> None: count(sampler, shots, batch_size, circuit.num_detectors) warmup_seconds = time.perf_counter() - start attempted = accepted = errors = 0 - cpu_start = cpu_seconds() + cpu_start = time.process_time() start = time.perf_counter() while time.perf_counter() - start < args.seconds: kept, wrong = count(sampler, shots, batch_size, circuit.num_detectors) @@ -71,10 +70,9 @@ def main() -> None: accepted += kept errors += wrong elapsed = time.perf_counter() - start - cores = (cpu_seconds() - cpu_start) / elapsed + cores = (time.process_time() - cpu_start) / elapsed acceptance = accepted / attempted - sigma = math.sqrt(reference * (1 - reference) / attempted) - check = "ok" if abs(acceptance - reference) <= 5 * sigma else "MISMATCH" + check = acceptance_check(acceptance, reference, attempted) print( f"| {name} | {compile_seconds:.2f} | {warmup_seconds:.2f} | " f"{attempted / elapsed:,.0f} | {cores:.1f} | {acceptance:.4f} | " diff --git a/scripts/statevector/benchmark_clifft.py b/scripts/statevector/benchmark_clifft.py index 627ba157..fd874f02 100644 --- a/scripts/statevector/benchmark_clifft.py +++ b/scripts/statevector/benchmark_clifft.py @@ -18,11 +18,10 @@ """ import argparse -import math import time import clifft -from workloads import CIRCUITS, WORKLOADS, cpu_seconds +from workloads import CIRCUITS, WORKLOADS, acceptance_check BATCH_SIZES = (1, 2048) @@ -33,7 +32,7 @@ def measure(program, shots: int, seconds: float, threads, batch_size: int): program, shots, seed=0, threads=threads, batch_size=batch_size ) attempted = accepted = errors = calls = 0 - cpu_start = cpu_seconds() + cpu_start = time.process_time() start = time.perf_counter() while time.perf_counter() - start < seconds: calls += 1 @@ -44,7 +43,7 @@ def measure(program, shots: int, seconds: float, threads, batch_size: int): accepted += int(result.passed_shots) errors += int(result.observable_ones[0]) elapsed = time.perf_counter() - start - cores = (cpu_seconds() - cpu_start) / elapsed + cores = (time.process_time() - cpu_start) / elapsed return attempted / elapsed, cores, attempted, accepted, errors @@ -82,8 +81,7 @@ def main() -> None: runs.items(), key=lambda item: item[1][0] ) acceptance = accepted / attempted - sigma = math.sqrt(reference * (1 - reference) / attempted) - check = "ok" if abs(acceptance - reference) <= 5 * sigma else "MISMATCH" + check = acceptance_check(acceptance, reference, attempted) print( f"| {name} | {threads} | {batch_size} | {compile_seconds:.2f} | " f"{program.peak_active_width} | {rate:,.0f} | {cores:.1f} | " diff --git a/scripts/statevector/workloads.py b/scripts/statevector/workloads.py index 1ff5b89e..ef470271 100644 --- a/scripts/statevector/workloads.py +++ b/scripts/statevector/workloads.py @@ -1,17 +1,17 @@ """The clifft-bench workloads shared by ``benchmark.py`` and ``benchmark_clifft.py``. Circuits are read from the directory in ``CLIFFT_BENCH_CIRCUITS`` (default -``/Users/rafaelhaenel/Documents/quera/clifft-bench/workloads/circuits``). +``workloads/circuits`` of a clifft-bench checkout next to this repository). """ +import math import os -import resource from pathlib import Path +_SIBLING_CHECKOUT = Path(__file__).resolve().parents[3] / "clifft-bench" CIRCUITS = Path( os.environ.get( - "CLIFFT_BENCH_CIRCUITS", - "/Users/rafaelhaenel/Documents/quera/clifft-bench/workloads/circuits", + "CLIFFT_BENCH_CIRCUITS", _SIBLING_CHECKOUT / "workloads" / "circuits" ) ) @@ -25,7 +25,7 @@ } -def cpu_seconds() -> float: - """Return the CPU time of this process (user plus system) in seconds.""" - usage = resource.getrusage(resource.RUSAGE_SELF) - return usage.ru_utime + usage.ru_stime +def acceptance_check(acceptance: float, reference: float, attempted: int) -> str: + """Return "ok" if ``acceptance`` is within 5σ of the reference, else "MISMATCH".""" + sigma = math.sqrt(reference * (1 - reference) / attempted) + return "ok" if abs(acceptance - reference) <= 5 * sigma else "MISMATCH" diff --git a/src/tsim/compile/pipeline.py b/src/tsim/compile/pipeline.py index bb8b39d8..4c21a876 100644 --- a/src/tsim/compile/pipeline.py +++ b/src/tsim/compile/pipeline.py @@ -3,6 +3,7 @@ from __future__ import annotations from fractions import Fraction +from functools import partial from typing import Literal import jax.numpy as jnp @@ -30,6 +31,7 @@ SamplingStrategy, ) from tsim.statevector.program import ( + STABILIZER_RANK_HINT, StatevectorComponent, compile_statevector_component, ) @@ -65,11 +67,15 @@ def compile_program( ValueError: If ``strategy="statevector"`` is combined with ``mode="joint"``. """ - if strategy == "statevector" and mode == "joint": - raise ValueError( - "strategy='statevector' samples outcomes and cannot compute " - "probabilities; use a stabilizer-rank strategy such as 'cat5'" - ) + if strategy == "statevector": + if mode == "joint": + raise ValueError( + "strategy='statevector' samples outcomes and cannot compute " + f"probabilities. {STABILIZER_RANK_HINT}" + ) + compile_component = compile_statevector_component + else: + compile_component = partial(_compile_component, mode=mode, strategy=strategy) graph = prepared.graph layout = prepared.layout @@ -85,12 +91,7 @@ def compile_program( basis_index, flip = direct direct_entries.append((output_indices[0], basis_index, flip)) continue - if strategy == "statevector": - compiled_components.append(compile_statevector_component(component, layout)) - else: - compiled_components.append( - _compile_component(component, output_indices, layout, mode, strategy) - ) + compiled_components.append(compile_component(component, output_indices, layout)) compiled_output_order.extend(output_indices) # Sort direct entries by output index so that — together with the output diff --git a/src/tsim/core/types.py b/src/tsim/core/types.py index e92877c3..58ada1ae 100644 --- a/src/tsim/core/types.py +++ b/src/tsim/core/types.py @@ -11,7 +11,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING, Literal +from typing import TYPE_CHECKING, ClassVar, Literal import equinox as eqx import numpy as np @@ -85,6 +85,9 @@ class CompiledComponent(eqx.Module): f_selection: Array compiled_scalar_graphs: tuple[CompiledScalarGraphs, ...] + norm_warning_threshold: ClassVar[float] = 1e-5 + """Largest normalization deviation that sampling tolerates without a warning.""" + @dataclass(frozen=True) class CompiledProgram: diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index e4b14716..c1e29f58 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -17,11 +17,8 @@ from tsim.core.graph import prepare_graph from tsim.core.types import CompiledComponent, CompiledProgram, SamplingStrategy from tsim.noise.channels import ChannelSampler -from tsim.statevector.execute import ( - NORM_WARNING_THRESHOLDS, - sample_statevector_component, -) -from tsim.statevector.program import StatevectorComponent +from tsim.statevector.execute import lane_bytes, sample_statevector_component +from tsim.statevector.program import STABILIZER_RANK_HINT, StatevectorComponent from tsim.utils.cuda_helpers import copy_d2h if TYPE_CHECKING: @@ -102,11 +99,11 @@ def _sample_component_jit( def sample_component( - component: CompiledComponent, + component: CompiledComponent | StatevectorComponent, f_params: jax.Array, key: PRNGKey, ) -> tuple[jax.Array, PRNGKey, jax.Array]: - """Sample outputs from a single component using autoregressive sampling. + """Sample outputs from a single component. Args: component: The compiled component to sample from. @@ -118,6 +115,8 @@ def sample_component( (batch_size, num_outputs_for_component). """ + if isinstance(component, StatevectorComponent): + return sample_statevector_component(component, f_params, key) # Skip JIT for small components (overhead not worth it) if len(component.output_indices) <= 1: return _sample_component(component, f_params, key) @@ -155,23 +154,14 @@ def sample_program( results.append(direct_bits) for component in program.components: - if isinstance(component, StatevectorComponent): - samples, key, max_norm_deviation = sample_statevector_component( - component, f_params, key - ) - warning_threshold = NORM_WARNING_THRESHOLDS[component.dtype] - else: - samples, key, max_norm_deviation = sample_component( - component, f_params, key - ) - warning_threshold = 1e-5 + samples, key, max_norm_deviation = sample_component(component, f_params, key) if np.isclose(max_norm_deviation, 1): raise ValueError( "A vanishing marginal probability distribution was encountered (normalization 0). " "This is likely the result of an underflow error. Please report this " "as a bug at https://github.com/QuEraComputing/tsim/issues/new." ) # pragma: no cover - if max_norm_deviation > warning_threshold: + if max_norm_deviation > component.norm_warning_threshold: warnings.warn( "A marginal probability was not normalized correctly " f"(normalization deviated from 1 by {max_norm_deviation:.1e}). " @@ -186,6 +176,12 @@ def sample_program( return combined +def _nbytes(tree: object) -> int: + return sum( + v.nbytes for v in jax.tree_util.tree_leaves(tree) if isinstance(v, jax.Array) + ) + + class _CompiledSamplerBase: """Base class for compiled samplers with common initialization logic.""" @@ -258,14 +254,16 @@ def __init__( statevector = [ c for c in self._program.components if isinstance(c, StatevectorComponent) ] - if statevector and self._peak_bytes_per_sample() > self._memory_budget(): - widest = max(statevector, key=lambda c: c.width) - raise ValueError( - f"strategy='statevector' needs {self._peak_bytes_per_sample()} bytes " - f"per shot for the component with outputs {list(widest.output_indices)} " - f"(width {widest.width}), more than half of the available memory. " - "Use a stabilizer-rank strategy such as strategy='cat5'." - ) + if statevector: + peak = self._peak_bytes_per_sample() + if peak > self._memory_budget(): + widest = max(statevector, key=lambda c: c.width) + raise ValueError( + f"strategy='statevector' needs {peak} bytes per shot for the " + f"component with outputs {list(widest.output_indices)} (width " + f"{widest.width}), more than half of the available memory. " + f"{STABILIZER_RANK_HINT}" + ) def _compute_direct_outputs(self, f_params_np: np.ndarray) -> np.ndarray: """Scatter direct output bits into a full (batch, num_outputs) bool array. @@ -329,11 +327,8 @@ def _peak_bytes_per_sample(self) -> int: peak = 0 for component in self._program.components: if isinstance(component, StatevectorComponent): - itemsize = 8 if component.dtype == "complex64" else 16 - lane = 3 * 2**component.width * itemsize - lane += 4 * component.num_words + component.num_outputs + 1 - # NOTE: the f-params and the output records, on host and device. - lane += 2 * num_f + 2 * self._program.num_outputs + # NOTE: plus the f-params and the output records, on host and device. + lane = lane_bytes(component) + 2 * num_f + 2 * self._program.num_outputs peak = max(peak, lane) continue for circuit in component.compiled_scalar_graphs: @@ -347,18 +342,17 @@ def _peak_bytes_per_sample(self) -> int: peak = max(peak, G * largest * 3) return max(peak, 1) - def _available_memory(self) -> int: - """Return the free memory of the first JAX device, or of the host.""" + def _memory_budget(self) -> int: + """Return the memory one batch may use: half of the free device memory.""" device = jax.devices()[0] if device.platform == "gpu": stats = device.memory_stats() - limit = stats.get("bytes_limit", 8 * 1024**3) - return limit - stats.get("bytes_in_use", 0) - return psutil.virtual_memory().available - - def _memory_budget(self) -> int: - """Return the memory one batch may use: half of the available memory.""" - return int(self._available_memory() * 0.5) # conservative + available = stats.get("bytes_limit", 8 * 1024**3) - stats.get( + "bytes_in_use", 0 + ) + else: + available = psutil.virtual_memory().available + return int(available * 0.5) # conservative def _estimate_batch_size(self) -> int: """Estimate the largest batch size that fits in available device memory. @@ -620,18 +614,14 @@ def __repr__(self) -> str: total_memory_bytes = 0 num_outputs = [] statevector_widths: list[int] = [] - statevector_actions: list[int] = [] + statevector_actions = 0 for component in self._program.components: if isinstance(component, StatevectorComponent): statevector_widths.append(component.width) - statevector_actions.append(component.num_actions) + statevector_actions += component.num_actions num_outputs.append(len(component.output_indices)) - total_memory_bytes += sum( - v.nbytes - for v in jax.tree_util.tree_leaves(component) - if isinstance(v, jax.Array) - ) + total_memory_bytes += _nbytes(component) continue for circuit in component.compiled_scalar_graphs: num_outputs.append(len(component.output_indices)) @@ -651,11 +641,7 @@ def __repr__(self) -> str: * circuit.float_terms.table.shape[1] ) - total_memory_bytes += sum( - v.nbytes - for v in jax.tree_util.tree_leaves(circuit) - if isinstance(v, jax.Array) - ) + total_memory_bytes += _nbytes(circuit) def _format_bytes(n: int) -> str: if n < 1024: @@ -672,7 +658,7 @@ def _format_bytes(n: int) -> str: statevector = ( f"{len(statevector_widths)} statevector components " f"(width ≤ {max(statevector_widths)}, " - f"{sum(statevector_actions)} actions), " + f"{statevector_actions} actions), " if statevector_widths else "" ) diff --git a/src/tsim/statevector/execute.py b/src/tsim/statevector/execute.py index f68c9a79..490c7756 100644 --- a/src/tsim/statevector/execute.py +++ b/src/tsim/statevector/execute.py @@ -2,11 +2,12 @@ ``run`` interprets a component's action tape with one ``lax.fori_loop``; each step dispatches on the action's kind with ``lax.switch``. The loop carry is a -``LaneState``: one lane per shot, each with ``2^W`` amplitudes, its symbols -(f-bits, then branch bits as they are sampled) and its output records. -All randomness goes through ``choose_outcome``, and ``count`` holds the number -of shots a lane stands for (always 1 for now), so that trajectory sharing can -later split lanes without touching the kernels. +``LaneState``: one lane per shot, each with ``2^W`` amplitudes and its symbols +(f-bits, then branch bits as they are sampled), from which ``outputs`` reads +the output bits once the tape has run. All randomness goes through +``choose_outcome``, and ``count`` holds the number of shots a lane stands for +(always 1 for now), so that trajectory sharing can later split lanes without +touching the kernels. """ from __future__ import annotations @@ -18,17 +19,19 @@ import jax.numpy as jnp from jax import Array -from tsim.statevector.program import StatevectorComponent +from tsim.statevector.program import ( + MEASURE, + MEASURE_DIAG, + MEASURE_DORMANT, + ROTATE, + ROTATE_DIAG, + StatevectorComponent, +) _I_POWERS = (1, 1j, -1, -1j) _TINY = 1e-30 _TWO_TO_32 = 2.0**32 -# NOTE: every outcome probability is normalized by the measured weights, so norm -# drift is harmless until it exceeds what rounding can explain: each rotation -# adds ~1e-7 relative drift under complex64. -NORM_WARNING_THRESHOLDS = {"complex64": 1e-3, "complex128": 1e-9} - class LaneState(NamedTuple): """Per-lane state of the executor; the leading axis is the lane. @@ -37,8 +40,6 @@ class LaneState(NamedTuple): amps: Amplitudes, shape ``(L, 2^W)``. symbols: f-bits, then branch bits, packed into uint32 words, shape ``(L, num_words)``. - records: Output bits, plus one scratch column for actions without a - slot, shape ``(L, num_outputs + 1)``. count: Shots each lane stands for, shape ``(L,)``. max_norm_deviation: Largest ``|‖ψ‖² - 1|`` seen at a measurement. @@ -46,11 +47,16 @@ class LaneState(NamedTuple): amps: Array symbols: Array - records: Array count: Array # pyright: ignore[reportIncompatibleMethodOverride] max_norm_deviation: Array +def lane_bytes(component: StatevectorComponent) -> int: + """Return the peak device memory of one lane, temporary amplitudes included.""" + amplitudes = 3 * 2**component.width * jnp.dtype(component.dtype).itemsize + return amplitudes + 4 * component.num_words + component.num_outputs + + def choose_outcome(p1: Array, count: Array, key: Array) -> Array: """Return each lane's outcome: True with probability ``p1``. @@ -88,24 +94,35 @@ def initial_state(component: StatevectorComponent, f_bits: Array) -> LaneState: return LaneState( amps=amps, symbols=pack_bits(f_bits, component.num_words), - records=jnp.zeros((lanes, component.num_outputs + 1), dtype=jnp.bool_), count=jnp.ones(lanes, dtype=jnp.int32), max_norm_deviation=jnp.zeros((), dtype=jnp.finfo(dtype).dtype), ) -def _index(component: StatevectorComponent) -> Array: - return jnp.arange(2**component.width, dtype=jnp.uint32) - - def _parity(values: Array) -> Array: return (jax.lax.population_count(values) & 1).astype(jnp.bool_) -def _sign(component: StatevectorComponent, symbols: Array, a: Array) -> Array: - """Return each lane's sign bit of action ``a``.""" - ones = jnp.sum(jax.lax.population_count(symbols & component.sign_row[a]), axis=1) - return ((ones & 1) == 1) ^ component.sign_const[a] +def _affine(symbols: Array, rows: Array, constant: Array) -> Array: + """Return ``constant ⊕ parity(symbols ∧ rows)`` over the last (word) axis.""" + ones = jnp.sum(jax.lax.population_count(symbols & rows), axis=-1) + return ((ones & 1) == 1) ^ constant + + +def outputs(component: StatevectorComponent, state: LaneState) -> Array: + """Return each lane's output bits, shape ``(L, num_outputs)``.""" + return _affine( + state.symbols[:, None, :], component.output_row, component.output_const + ) + + +def _index(component: StatevectorComponent) -> Array: + return jnp.arange(2**component.width, dtype=jnp.uint32) + + +def _sign(component: StatevectorComponent, state: LaneState, a: Array) -> Array: + """Return each lane's sign bit of rotation ``a``.""" + return _affine(state.symbols, component.sign_row[a], component.sign_const[a]) def _apply_pauli(amps: Array, index: Array, x: Array, z: Array) -> Array: @@ -117,32 +134,12 @@ def _apply_pauli(amps: Array, index: Array, x: Array, z: Array) -> Array: return i_power * signs * jnp.take(amps, flipped.astype(jnp.int32), axis=1) -def _write( - component: StatevectorComponent, state: LaneState, a: Array, outcome: Array -) -> LaneState: - """Record ``outcome ⊕ sign`` in the action's slot and set its branch bit.""" - value = outcome ^ _sign(component, state.symbols, a) - slot = component.slot[a] - column = jnp.where(slot >= 0, slot, component.num_outputs) - records = state.records.at[:, column].set(value) - branch = component.branch[a] - word = jnp.maximum(branch, 0) // 32 - bit = jnp.where( - branch >= 0, - jnp.uint32(1) << (jnp.maximum(branch, 0) % 32).astype(jnp.uint32), - jnp.uint32(0), - ) - flags = jnp.where(outcome, bit, jnp.uint32(0)) - symbols = state.symbols.at[:, word].set(state.symbols[:, word] | flags) - return state._replace(records=records, symbols=symbols) - - def _rotate_diag( component: StatevectorComponent, state: LaneState, a: Array, key: Array ) -> LaneState: del key odd = _parity(_index(component) & component.z_mask[a]) - flip = _sign(component, state.symbols, a)[:, None] ^ odd[None, :] + flip = _sign(component, state, a)[:, None] ^ odd[None, :] factor = component.cos[a] - 1j * component.sin[a] * jnp.where(flip, -1, 1) return state._replace(amps=(state.amps * factor).astype(state.amps.dtype)) @@ -154,7 +151,7 @@ def _rotate( paulied = _apply_pauli( state.amps, _index(component), component.x_mask[a], component.z_mask[a] ) - negate = jnp.where(_sign(component, state.symbols, a), -1, 1)[:, None] + negate = jnp.where(_sign(component, state, a), -1, 1)[:, None] amps = component.cos[a] * state.amps - 1j * component.sin[a] * negate * paulied return state._replace(amps=amps.astype(state.amps.dtype)) @@ -163,36 +160,59 @@ def _squared(values: Array) -> Array: return jnp.real(values) ** 2 + jnp.imag(values) ** 2 -def _probability_of_one(weight_0: Array, weight_1: Array) -> Array: - """Return ``w1 / (w0 + w1)`` from the outcome weights, without cancellation.""" - return jnp.clip(weight_1 / jnp.maximum(weight_0 + weight_1, _TINY), 0, 1) +def _set_branch( + component: StatevectorComponent, state: LaneState, a: Array, outcome: Array +) -> LaneState: + """Set the measurement's branch symbol to ``outcome``.""" + branch = component.branch[a] + word = branch // 32 + bit = jnp.uint32(1) << (branch % 32).astype(jnp.uint32) + flags = jnp.where(outcome, bit, jnp.uint32(0)) + symbols = state.symbols.at[:, word].set(state.symbols[:, word] | flags) + return state._replace(symbols=symbols) + + +def _sample_outcome( + state: LaneState, a: Array, key: Array, weight_0: Array, weight_1: Array +) -> tuple[Array, Array]: + """Return each lane's outcome and its weight, without cancellation in ``p1``.""" + p1 = jnp.clip(weight_1 / jnp.maximum(weight_0 + weight_1, _TINY), 0, 1) + outcome = choose_outcome(p1, state.count, jax.random.fold_in(key, a)) + return outcome, jnp.where(outcome, weight_1, weight_0) -def _measured(state: LaneState, new: Array, norm: Array) -> LaneState: +def _collapse( + component: StatevectorComponent, + state: LaneState, + a: Array, + outcome: Array, + projected: Array, + norm: Array, +) -> LaneState: + """Keep the pivot-0 half of the normalized ``projected``; branch on ``outcome``.""" + keep = ((_index(component) >> component.pivot[a].astype(jnp.uint32)) & 1) == 0 deviation = jnp.maximum(state.max_norm_deviation, jnp.max(jnp.abs(norm - 1))) - return state._replace( - amps=new.astype(state.amps.dtype), + state = state._replace( + amps=jnp.where(keep[None, :], projected, 0).astype(state.amps.dtype), max_norm_deviation=deviation.astype(state.max_norm_deviation.dtype), ) + return _set_branch(component, state, a, outcome) def _measure( component: StatevectorComponent, state: LaneState, a: Array, key: Array ) -> LaneState: - index = _index(component) amps = state.amps - paulied = _apply_pauli(amps, index, component.x_mask[a], component.z_mask[a]) + paulied = _apply_pauli( + amps, _index(component), component.x_mask[a], component.z_mask[a] + ) plus, minus = amps + paulied, amps - paulied weight_0 = jnp.sum(_squared(plus), axis=1) / 4 weight_1 = jnp.sum(_squared(minus), axis=1) / 4 - p1 = _probability_of_one(weight_0, weight_1) - outcome = choose_outcome(p1, state.count, jax.random.fold_in(key, a)) - weight = jnp.where(outcome, weight_1, weight_0) + outcome, weight = _sample_outcome(state, a, key, weight_0, weight_1) scale = jax.lax.rsqrt(jnp.maximum(2 * weight, _TINY))[:, None] - keep = ((index >> component.pivot[a].astype(jnp.uint32)) & 1) == 0 - projected = jnp.where(outcome[:, None], minus, plus) - new = jnp.where(keep[None, :], projected * scale, 0) - return _write(component, _measured(state, new, weight_0 + weight_1), a, outcome) + projected = jnp.where(outcome[:, None], minus, plus) * scale + return _collapse(component, state, a, outcome, projected, weight_0 + weight_1) def _measure_diag( @@ -204,22 +224,13 @@ def _measure_diag( weights = _squared(amps) weight_0 = jnp.sum(jnp.where(odd[None, :], 0, weights), axis=1) weight_1 = jnp.sum(jnp.where(odd[None, :], weights, 0), axis=1) - p1 = _probability_of_one(weight_0, weight_1) - outcome = choose_outcome(p1, state.count, jax.random.fold_in(key, a)) - weight = jnp.where(outcome, weight_1, weight_0) - pivot = component.pivot[a].astype(jnp.uint32) - bit = jnp.uint32(1) << pivot - source_if_0 = (index ^ (bit * odd.astype(jnp.uint32))).astype(jnp.int32) - source_if_1 = (index ^ (bit * (~odd).astype(jnp.uint32))).astype(jnp.int32) - gathered = jnp.where( - outcome[:, None], - jnp.take(amps, source_if_1, axis=1), - jnp.take(amps, source_if_0, axis=1), - ) - keep = ((index >> pivot) & 1) == 0 + outcome, weight = _sample_outcome(state, a, key, weight_0, weight_1) + bit = jnp.uint32(1) << component.pivot[a].astype(jnp.uint32) + partner = jnp.take(amps, (index ^ bit).astype(jnp.int32), axis=1) + gathered = jnp.where(odd[None, :] == outcome[:, None], amps, partner) scale = jax.lax.rsqrt(jnp.maximum(weight, _TINY))[:, None] - new = jnp.where(keep[None, :], gathered * scale, 0) - return _write(component, _measured(state, new, weight_0 + weight_1), a, outcome) + projected = gathered * scale + return _collapse(component, state, a, outcome, projected, weight_0 + weight_1) def _measure_dormant( @@ -227,18 +238,16 @@ def _measure_dormant( ) -> LaneState: half = jnp.full(state.count.shape, 0.5, dtype=state.max_norm_deviation.dtype) outcome = choose_outcome(half, state.count, jax.random.fold_in(key, a)) - return _write(component, state, a, outcome) - - -def _record( - component: StatevectorComponent, state: LaneState, a: Array, key: Array -) -> LaneState: - del key - return _write(component, state, a, jnp.zeros(state.count.shape, dtype=jnp.bool_)) + return _set_branch(component, state, a, outcome) -# NOTE: the order must match the kind codes ROTATE_DIAG … RECORD in program.py. -_KERNELS = (_rotate_diag, _rotate, _measure_diag, _measure, _measure_dormant, _record) +_KERNELS = { + ROTATE_DIAG: _rotate_diag, + ROTATE: _rotate, + MEASURE_DIAG: _measure_diag, + MEASURE: _measure, + MEASURE_DORMANT: _measure_dormant, +} def run( @@ -254,7 +263,7 @@ def run( """ if component.num_actions == 0: return state - branches = [partial(kernel, component) for kernel in _KERNELS] + branches = [partial(_KERNELS[kind], component) for kind in range(len(_KERNELS))] def step(a: Array, carry: LaneState) -> LaneState: return jax.lax.switch(component.kind[a], branches, carry, a, key) @@ -263,15 +272,6 @@ def step(a: Array, carry: LaneState) -> LaneState: @jax.jit -def _run_lanes( - component: StatevectorComponent, f_bits: Array, key: Array -) -> tuple[Array, Array]: - state = run( - component, initial_state(component, f_bits), key, 0, component.num_actions - ) - return state.records[:, : component.num_outputs], state.max_norm_deviation - - def sample_statevector_component( component: StatevectorComponent, f_params: Array, key: Array ) -> tuple[Array, Array, Array]: @@ -289,5 +289,6 @@ def sample_statevector_component( """ f_bits = f_params[:, component.f_selection].astype(jnp.bool_) key, subkey = jax.random.split(key) - samples, deviation = _run_lanes(component, f_bits, subkey) - return samples, key, deviation + state = initial_state(component, f_bits) + state = run(component, state, subkey, 0, component.num_actions) + return outputs(component, state), key, state.max_norm_deviation diff --git a/src/tsim/statevector/heisenberg.py b/src/tsim/statevector/heisenberg.py index b3a60583..1ab0a834 100644 --- a/src/tsim/statevector/heisenberg.py +++ b/src/tsim/statevector/heisenberg.py @@ -11,14 +11,14 @@ from __future__ import annotations -from collections.abc import Mapping -from dataclasses import dataclass +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, fields from fractions import Fraction from math import pi import numpy as np import stim -from zxparam import Graph, Vertex +from zxparam import Graph from zxparam.extract import ExtractedCircuit, Gate from tsim.core.params import ParamLayout @@ -26,6 +26,7 @@ ROTATION = 0 MEASUREMENT = 1 +_CLIFFORD_GATES = {"H": "h", "CNOT": "cx", "CZ": "cz"} _CLIFFORD_QUARTERS = {1: "s", 2: "z", 3: "s_dag"} @@ -36,9 +37,7 @@ class ComponentLayout: Attributes: f_selection: Global f-bit index of each component f-bit, ascending. f_position: Position in ``f_selection`` of each phase parameter ID. - output_indices: Global output index of each output slot, in port order - (the order of ``Graph.outputs()``, as for stabilizer-rank - components). + output_indices: Global output index of each output slot. """ @@ -47,26 +46,27 @@ class ComponentLayout: output_indices: tuple[int, ...] @classmethod - def of(cls, component: Graph, layout: ParamLayout) -> ComponentLayout: - """Return the layout of ``component``'s f-bits and output ports.""" + def of( + cls, component: Graph, output_indices: Sequence[int], layout: ParamLayout + ) -> ComponentLayout: + """Return the layout of ``component``, whose outputs are ``output_indices``.""" f_selection = tuple( - sorted( - { - layout.basis_index(param) - for v in component.vertices() - for param in component.phase(v).parameters - } - ) + sorted(layout.basis_index(p) for p in component.parameters()) ) f_position = {layout.basis(i): pos for pos, i in enumerate(f_selection)} - outputs = tuple(_port(component, v) for v in component.outputs()) - return cls(f_selection, f_position, outputs) + return cls(f_selection, f_position, tuple(output_indices)) + +def anticommute( + x1: np.ndarray, z1: np.ndarray, x2: np.ndarray, z2: np.ndarray +) -> np.ndarray: + """Return which Paulis ``(x1, z1)`` anticommute with which Paulis ``(x2, z2)``. -def _port(g: Graph, v: Vertex) -> int: - port = g.port(v) - assert port is not None - return port + Each row of an argument is a Pauli; a 1-D argument is one Pauli, whose axis + the result drops. + """ + f = np.float32 + return (x1.astype(f) @ z2.T.astype(f) + z1.astype(f) @ x2.T.astype(f)) % 2 == 1 @dataclass(frozen=True) @@ -79,8 +79,6 @@ class HeisenbergOps: ``[1, f_0, f_1, …]``. A measurement records its eigenvalue bit. Attributes: - num_lines: Number of qubit lines. - num_f: Number of component f-bits the signs depend on. kind: ``ROTATION`` or ``MEASUREMENT``, shape ``(N,)``. x: X part of each Pauli, shape ``(N, num_lines)``. z: Z part of each Pauli, shape ``(N, num_lines)``. @@ -91,8 +89,6 @@ class HeisenbergOps: """ - num_lines: int - num_f: int kind: np.ndarray x: np.ndarray z: np.ndarray @@ -105,6 +101,16 @@ def num_ops(self) -> int: """Return the number of operations.""" return len(self.kind) + @property + def num_lines(self) -> int: + """Return the number of qubit lines.""" + return self.x.shape[1] + + @property + def num_f(self) -> int: + """Return the number of component f-bits the signs depend on.""" + return self.sign.shape[1] - 1 + @property def num_outputs(self) -> int: """Return the number of recorded measurements.""" @@ -116,16 +122,7 @@ def take(self, order: np.ndarray) -> HeisenbergOps: The result describes the same distribution when ``order`` keeps every pair of anticommuting operations in their original order. """ - return HeisenbergOps( - num_lines=self.num_lines, - num_f=self.num_f, - kind=self.kind[order], - x=self.x[order], - z=self.z[order], - angle=self.angle[order], - sign=self.sign[order], - slot=self.slot[order], - ) + return HeisenbergOps(*(getattr(self, f.name)[order] for f in fields(self))) def heisenberg_ops(circuit: ExtractedCircuit, layout: ComponentLayout) -> HeisenbergOps: @@ -142,9 +139,9 @@ def heisenberg_ops(circuit: ExtractedCircuit, layout: ComponentLayout) -> Heisen """ _check_lines(circuit) - builder = _Builder(circuit.n_lines, len(layout.f_selection)) + builder = _Builder(circuit.n_lines, layout) for gate in circuit.gates: - builder.gate(gate, layout.f_position) + builder.gate(gate) slot_of_port = {port: slot for slot, port in enumerate(layout.output_indices)} for line in range(circuit.n_lines): port = circuit.line_end[line] @@ -169,18 +166,18 @@ def _check_lines(circuit: ExtractedCircuit) -> None: class _Builder: """Walks the gates with a stim simulator that holds the Clifford prefix U.""" - def __init__(self, num_lines: int, num_f: int) -> None: + def __init__(self, num_lines: int, layout: ComponentLayout) -> None: """Start every line in ``|+⟩`` (``H`` on ``|0⟩``).""" self.num_lines = num_lines - self.num_f = num_f + self.num_f = len(layout.f_selection) + self.f_position = layout.f_position self.sim = stim.TableauSimulator() self.sim.set_num_qubits(num_lines) for line in range(num_lines): self.sim.h(line) self.inverse: stim.Tableau | None = None - self.time = 0 - self.ops: list[tuple[int, np.ndarray, np.ndarray, bool, float, int, int]] = [] - self.noise: list[tuple[np.ndarray, np.ndarray, frozenset[int], int]] = [] + self.ops: list[tuple[int, np.ndarray, np.ndarray, bool, float, int]] = [] + self.noise: list[tuple[np.ndarray, np.ndarray, np.ndarray, int]] = [] def _mapped(self, line: int) -> tuple[np.ndarray, np.ndarray, bool]: """Return U†·Z_line·U as bit rows and a sign flag.""" @@ -194,22 +191,18 @@ def _clifford(self, name: str, *lines: int) -> None: getattr(self.sim, name)(*lines) self.inverse = None - def gate(self, gate: Gate, f_position: Mapping[int, int]) -> None: + def gate(self, gate: Gate) -> None: """Absorb a Clifford gate, or record a rotation, noise or measurement.""" - if gate.kind == "H": - self._clifford("h", *gate.lines) - elif gate.kind == "CNOT": - self._clifford("cx", *gate.lines) - elif gate.kind == "CZ": - self._clifford("cz", *gate.lines) + if gate.kind in _CLIFFORD_GATES: + self._clifford(_CLIFFORD_GATES[gate.kind], *gate.lines) elif gate.kind == "Z": - self._phase(gate, f_position) + self._phase(gate) elif gate.kind in ("MEASURE", "DISCARD"): self.measure(gate.lines[0], -1) else: raise ValueError(f"unknown gate kind {gate.kind!r}") - def _phase(self, gate: Gate, f_position: Mapping[int, int]) -> None: + def _phase(self, gate: Gate) -> None: (line,) = gate.lines phase = gate.phase assert phase is not None @@ -220,53 +213,36 @@ def _phase(self, gate: Gate, f_position: Mapping[int, int]) -> None: self._clifford(clifford, line) if remainder: x, z, negative = self._mapped(line) - self._op(ROTATION, x, z, negative, pi * float(remainder) / 2, -1) + self.ops.append((ROTATION, x, z, negative, pi * float(remainder) / 2, -1)) if phase.parameters: x, z, _ = self._mapped(line) - positions = frozenset(f_position[p] for p in phase.parameters) - self.noise.append((x, z, positions, self.time)) - self.time += 1 + condition = np.zeros(self.num_f, dtype=bool) + condition[[self.f_position[p] for p in phase.parameters]] = True + self.noise.append((x, z, condition, len(self.ops))) def measure(self, line: int, slot: int) -> None: """Record a Z measurement of ``line`` (``slot`` -1 if forgotten).""" x, z, negative = self._mapped(line) - self._op(MEASUREMENT, x, z, negative, 0.0, slot) - - def _op( - self, - kind: int, - x: np.ndarray, - z: np.ndarray, - negative: bool, - angle: float, - slot: int, - ) -> None: - self.ops.append((kind, x, z, negative, angle, slot, self.time)) - self.time += 1 + self.ops.append((MEASUREMENT, x, z, negative, 0.0, slot)) def finish(self) -> HeisenbergOps: - """Return the ops, with every noise event folded into their sign rows.""" + """Return the ops, with every noise event folded into their sign rows. + + A noise event flips the sign of every later op it anticommutes with. + """ num_ops, num_lines = len(self.ops), self.num_lines x = np.array([op[1] for op in self.ops], dtype=bool).reshape(num_ops, num_lines) z = np.array([op[2] for op in self.ops], dtype=bool).reshape(num_ops, num_lines) sign = np.zeros((num_ops, 1 + self.num_f), dtype=bool) sign[:, 0] = [op[3] for op in self.ops] - if self.noise and num_ops: - noise_x = np.array([e[0] for e in self.noise], dtype=np.float32) - noise_z = np.array([e[1] for e in self.noise], dtype=np.float32) - anti = ( - noise_x @ z.T.astype(np.float32) + noise_z @ x.T.astype(np.float32) - ) % 2 == 1 - noise_time = np.array([e[3] for e in self.noise]) - op_time = np.array([op[6] for op in self.ops]) - flips = anti & (noise_time[:, None] < op_time[None, :]) - condition = np.zeros((len(self.noise), self.num_f), dtype=np.float32) - for e, (_, _, positions, _) in enumerate(self.noise): - condition[e, sorted(positions)] = 1 - sign[:, 1:] = (flips.T.astype(np.float32) @ condition) % 2 == 1 + if self.noise: + noise_x, noise_z, condition, start = map( + np.array, zip(*self.noise, strict=True) + ) + later = start[:, None] <= np.arange(num_ops)[None, :] + flips = (anticommute(noise_x, noise_z, x, z) & later).astype(np.float32) + sign[:, 1:] = (flips.T @ condition.astype(np.float32)) % 2 == 1 return HeisenbergOps( - num_lines=num_lines, - num_f=self.num_f, kind=np.array([op[0] for op in self.ops], dtype=np.uint8), x=x, z=z, diff --git a/src/tsim/statevector/plan.py b/src/tsim/statevector/plan.py index 077e1d91..46436558 100644 --- a/src/tsim/statevector/plan.py +++ b/src/tsim/statevector/plan.py @@ -8,7 +8,8 @@ ``AffineRow`` over the symbols: the component f-bits, then the branch bits of random measurements. A random measurement keeps the reference state on its ``b = 0`` branch; its ``b = 1`` branch differs by a correction Pauli, which -flips the sign of every later op it anticommutes with. +flips the sign of every later op it anticommutes with. Each output is an +``AffineRow`` too: the measurement's sign, XOR its branch bit if it is random. """ from __future__ import annotations @@ -19,7 +20,7 @@ import numpy as np import stim -from tsim.statevector.heisenberg import ROTATION, HeisenbergOps +from tsim.statevector.heisenberg import ROTATION, HeisenbergOps, anticommute from tsim.statevector.schedule import Schedule @@ -53,56 +54,47 @@ class Rotate: class MeasureActive: """Measure the Pauli ``x``, ``z`` on the array, then reset ``pivot`` to ``|0⟩``. - The sampled eigenvalue bit becomes symbol ``branch``; ``slot`` (if not -1) - records ``branch ⊕ sign``. + The sampled eigenvalue bit becomes symbol ``branch``. """ x: int z: int pivot: int branch: int - sign: AffineRow - slot: int @dataclass(frozen=True) class MeasureDormant: - """Sample a fair bit into symbol ``branch``; ``slot`` records it ⊕ ``sign``.""" + """Sample a fair bit into symbol ``branch``.""" branch: int - sign: AffineRow - slot: int - -@dataclass(frozen=True) -class Record: - """Record the deterministic value ``sign`` in ``slot``.""" - - sign: AffineRow - slot: int - -Action = Rotate | MeasureActive | MeasureDormant | Record +Action = Rotate | MeasureActive | MeasureDormant @dataclass(frozen=True) class Plan: - """Actions on ``2^width`` amplitudes, with signs over the symbols. + """Actions on ``2^width`` amplitudes, and outputs over the symbols. Attributes: width: Number of array axes (W). - num_f: Number of f-bit symbols (symbols ``0 … num_f-1``). - num_branches: Number of branch symbols (``num_f …``). - num_outputs: Number of output slots. + num_f: Number of f-bit symbols (symbols ``0 … num_f-1``); the branch + symbols follow, one per measurement action. actions: The actions, in execution order. + outputs: The value of each output slot. """ width: int num_f: int - num_branches: int - num_outputs: int actions: tuple[Action, ...] + outputs: tuple[AffineRow, ...] + + @property + def num_branches(self) -> int: + """Return the number of branch symbols.""" + return sum(not isinstance(action, Rotate) for action in self.actions) @property def num_symbols(self) -> int: @@ -117,25 +109,17 @@ def plan(ops: HeisenbergOps, schedule: Schedule) -> Plan: RuntimeError: If the ops need more than ``schedule.width`` active axes. """ - planner = _Planner(ops, schedule) - for position in range(len(schedule.order)): - planner.step(position) + planner = _Planner(ops.take(schedule.order), schedule.width) + for j in range(ops.num_ops): + planner.step(j) return Plan( width=schedule.width, num_f=ops.num_f, - num_branches=planner.num_branches, - num_outputs=ops.num_outputs, actions=tuple(planner.actions), + outputs=tuple(planner.outputs[slot] for slot in range(ops.num_outputs)), ) -def _anticommutes( - x: np.ndarray, z: np.ndarray, xs: np.ndarray, zs: np.ndarray -) -> np.ndarray: - """Return which rows ``(xs[i], zs[i])`` anticommute with the Pauli ``(x, z)``.""" - return np.logical_xor.reduce(xs & z, axis=1) ^ np.logical_xor.reduce(zs & x, axis=1) - - def _mask(bits: np.ndarray, axes: Iterable[int]) -> int: return sum(1 << int(q) for q in axes if bits[q]) @@ -146,18 +130,17 @@ def _mask(bits: np.ndarray, axes: Iterable[int]) -> int: class _Planner: """The frame, the active axes and the branch corrections during planning.""" - def __init__(self, ops: HeisenbergOps, schedule: Schedule) -> None: + def __init__(self, ops: HeisenbergOps, width: int) -> None: """Start with the identity frame and every coordinate dormant.""" self.ops = ops - self.order = schedule.order - self.width = schedule.width + self.width = width self.frame = stim.Tableau(ops.num_lines) self.active: set[int] = set() - self.num_branches = 0 self.correction_x = np.zeros((0, ops.num_lines), dtype=bool) self.correction_z = np.zeros((0, ops.num_lines), dtype=bool) self.correction_condition: list[AffineRow] = [] self.actions: list[Action] = [] + self.outputs: dict[int, AffineRow] = {} def _apply(self, gate: str, *targets: int) -> None: """Apply a virtual Clifford: ``F ← gate ∘ F``.""" @@ -175,54 +158,58 @@ def _sign(self, j: int, negative: bool) -> AffineRow: row = self.ops.sign[j] terms = frozenset(int(t) for t in np.flatnonzero(row[1:])) sign = AffineRow(bool(row[0]) ^ negative, terms) - anti = _anticommutes( - self.ops.x[j], self.ops.z[j], self.correction_x, self.correction_z + anti = anticommute( + self.correction_x, self.correction_z, self.ops.x[j], self.ops.z[j] ) for m in np.flatnonzero(anti): sign = sign ^ self.correction_condition[m] return sign - def _new_branch(self) -> int: - branch = self.ops.num_f + self.num_branches - self.num_branches += 1 - return branch - def _correction(self, axis: int) -> tuple[np.ndarray, np.ndarray]: """Return ``F⁻¹(X_axis)``, the correction Pauli in the ``|0…0⟩`` frame.""" x, z = self.frame.inverse_x_output(axis).to_numpy() return x, z - def _add_correction(self, axis: int, condition: AffineRow) -> None: - x, z = self._correction(axis) + def _branch(self, correction: tuple[np.ndarray, np.ndarray], flip: bool) -> int: + """Return a new branch symbol ``b``; ``correction`` applies when ``b ⊕ flip``.""" + branch = self.ops.num_f + len(self.correction_condition) + x, z = correction self.correction_x = np.vstack([self.correction_x, x[None]]) self.correction_z = np.vstack([self.correction_z, z[None]]) - self.correction_condition.append(condition) + self.correction_condition.append(AffineRow(flip, frozenset({branch}))) + return branch - def step(self, position: int) -> None: - """Plan the op at ``position`` of the schedule.""" - j = int(self.order[position]) + def step(self, j: int) -> None: + """Plan op ``j``.""" x, z, negative = self._virtual(j) - sign = self._sign(j, negative) dormant_x = [int(q) for q in np.flatnonzero(x) if int(q) not in self.active] on_active = any(x[q] or z[q] for q in self.active) if self.ops.kind[j] == ROTATION: if dormant_x: self._promote(j, dormant_x) - elif on_active: - axes = sorted(self.active) - angle = float(self.ops.angle[j]) - self.actions.append(Rotate(_mask(x, axes), _mask(z, axes), angle, sign)) + x, z, negative = self._virtual(j) + elif not on_active: + return + axes = sorted(self.active) + angle = float(self.ops.angle[j]) + sign = self._sign(j, negative) + self.actions.append(Rotate(_mask(x, axes), _mask(z, axes), angle, sign)) return - slot = int(self.ops.slot[j]) + output = self._sign(j, negative) if dormant_x: - self._measure_dormant(position, j, dormant_x, negative, sign, slot) + branch = self._measure_dormant(j, dormant_x, negative) elif on_active: - self._measure_active(j, x, z, negative, sign, slot) - elif slot >= 0: - self.actions.append(Record(sign, slot)) + branch = self._measure_active(j, x, z, negative) + else: + branch = None + if branch is not None: + output ^= AffineRow(False, frozenset({branch})) + slot = int(self.ops.slot[j]) + if slot >= 0: + self.outputs[slot] = output def _promote(self, j: int, dormant_x: list[int]) -> None: - """Move the rotation's dormant X support onto one free axis and rotate.""" + """Move the rotation's dormant X support onto one free axis and activate it.""" free = [q for q in range(self.width) if q not in self.active] if not free: raise RuntimeError( @@ -238,23 +225,15 @@ def _promote(self, j: int, dormant_x: list[int]) -> None: for q in np.flatnonzero(x): if int(q) != axis and int(q) not in self.active: self._apply("CX", axis, int(q)) - x, z, negative = self._virtual(j) - axes = sorted(self.active | {axis}) - angle = float(self.ops.angle[j]) - sign = self._sign(j, negative) - self.actions.append(Rotate(_mask(x, axes), _mask(z, axes), angle, sign)) self.active.add(axis) def _measure_dormant( - self, - position: int, - j: int, - dormant_x: list[int], - negative: bool, - sign: AffineRow, - slot: int, - ) -> None: - """Turn the measured Pauli into ``±X_d`` without touching the state, then H.""" + self, j: int, dormant_x: list[int], negative: bool + ) -> int | None: + """Turn the measured Pauli into ``±X_d`` without touching the state, then H. + + Return the branch symbol, or None when no output or later op needs it. + """ d = dormant_x[0] for q in dormant_x[1:]: self._apply("CX", d, q) @@ -268,24 +247,18 @@ def _measure_dormant( x, z, after = self._virtual(j) assert x[d] and x.sum() == 1 and not z.any() self._apply("H", d) - later = self.order[position + 1 :] - kx, kz = self._correction(d) - anti = _anticommutes(kx, kz, self.ops.x[later], self.ops.z[later]) - if slot < 0 and not anti.any(): - return - branch = self._new_branch() - self._add_correction(d, AffineRow(after ^ negative, frozenset({branch}))) - self.actions.append(MeasureDormant(branch, sign, slot)) + correction = self._correction(d) + later = slice(j + 1, None) + anti = anticommute(self.ops.x[later], self.ops.z[later], *correction) + if self.ops.slot[j] < 0 and not anti.any(): + return None + branch = self._branch(correction, after ^ negative) + self.actions.append(MeasureDormant(branch)) + return branch def _measure_active( - self, - j: int, - x: np.ndarray, - z: np.ndarray, - negative: bool, - sign: AffineRow, - slot: int, - ) -> None: + self, j: int, x: np.ndarray, z: np.ndarray, negative: bool + ) -> int: """Measure on the array, then turn the Pauli into ``±Z_pivot`` and free it.""" axes = sorted(self.active) x_mask, z_mask = _mask(x, axes), _mask(z, axes) @@ -309,7 +282,7 @@ def _measure_active( self._apply("CX", q, pivot) _, _, after = self._virtual(j) assert x_support or after == negative - branch = self._new_branch() - self.actions.append(MeasureActive(x_mask, z_mask, pivot, branch, sign, slot)) - self._add_correction(pivot, AffineRow(after ^ negative, frozenset({branch}))) + branch = self._branch(self._correction(pivot), after ^ negative) + self.actions.append(MeasureActive(x_mask, z_mask, pivot, branch)) self.active.remove(pivot) + return branch diff --git a/src/tsim/statevector/program.py b/src/tsim/statevector/program.py index a6a0305c..2a48e4d7 100644 --- a/src/tsim/statevector/program.py +++ b/src/tsim/statevector/program.py @@ -1,12 +1,16 @@ """Compilation of one component for ``strategy="statevector"``. ``plan_component`` runs extraction, the Heisenberg rewrite, scheduling and -planning; ``lower`` turns the plan into the arrays of a ``StatevectorComponent``, -an action tape that ``execute.run`` interprets. +planning; ``lower`` turns the plan into the arrays of a ``StatevectorComponent``: +an action tape that ``execute.run`` interprets, and the outputs as parities of +the symbols the tape leaves behind. """ from __future__ import annotations +from collections.abc import Iterable, Sequence +from math import ceil + import equinox as eqx import jax import jax.numpy as jnp @@ -17,7 +21,7 @@ from tsim.core.params import ParamLayout from tsim.statevector.heisenberg import ComponentLayout, heisenberg_ops -from tsim.statevector.plan import MeasureActive, MeasureDormant, Plan, Rotate, plan +from tsim.statevector.plan import MeasureActive, Plan, Rotate, plan from tsim.statevector.schedule import schedule ROTATE_DIAG = 0 @@ -25,11 +29,17 @@ MEASURE_DIAG = 2 MEASURE = 3 MEASURE_DORMANT = 4 -RECORD = 5 MAX_WIDTH = 30 -_HINT = "Use a stabilizer-rank strategy such as strategy='cat5' for this circuit." +# NOTE: every outcome probability is normalized by the measured weights, so norm +# drift is harmless until it exceeds what rounding can explain: each rotation +# adds ~1e-7 relative drift under complex64. +_NORM_WARNING_THRESHOLDS = {"complex64": 1e-3, "complex128": 1e-9} + +STABILIZER_RANK_HINT = ( + "Use a stabilizer-rank strategy such as strategy='cat5' for this circuit." +) class StatevectorComponent(eqx.Module): @@ -38,27 +48,24 @@ class StatevectorComponent(eqx.Module): Attributes: output_indices: Global output index of each output slot. width: Number of array axes W; lanes hold ``2^W`` amplitudes. - num_words: Number of uint32 words that hold a lane's symbols. - num_outputs: Number of output slots. dtype: ``"complex64"`` or ``"complex128"``. f_selection: Global f-bit index of each component f-bit. - kind: Kernel of each action (``ROTATE_DIAG`` … ``RECORD``). + kind: Kernel of each action (``ROTATE_DIAG`` … ``MEASURE_DORMANT``). x_mask: X mask of the action's Pauli on the array axes. z_mask: Z mask of the action's Pauli on the array axes. cos: cos θ of each rotation. sin: sin θ of each rotation. pivot: Axis an active measurement resets to ``|0⟩``. - branch: Symbol a measurement sets, -1 if none. - slot: Output slot an action writes, -1 if none. - sign_const: Constant of each action's sign. - sign_row: Symbol bits of each action's sign, shape ``(A, num_words)``. + branch: Symbol a measurement sets, -1 for a rotation. + sign_const: Constant of each rotation's sign. + sign_row: Symbol bits of each rotation's sign, shape ``(A, num_words)``. + output_const: Constant of each output, shape ``(num_outputs,)``. + output_row: Symbol bits of each output, shape ``(num_outputs, num_words)``. """ output_indices: tuple[int, ...] = eqx.field(static=True) width: int = eqx.field(static=True) - num_words: int = eqx.field(static=True) - num_outputs: int = eqx.field(static=True) dtype: str = eqx.field(static=True) f_selection: Array kind: Array @@ -68,18 +75,34 @@ class StatevectorComponent(eqx.Module): sin: Array pivot: Array branch: Array - slot: Array sign_const: Array sign_row: Array + output_const: Array + output_row: Array @property def num_actions(self) -> int: """Return the number of actions on the tape.""" return int(self.kind.shape[0]) + @property + def num_words(self) -> int: + """Return the number of uint32 words that hold a lane's symbols.""" + return int(self.sign_row.shape[1]) + + @property + def num_outputs(self) -> int: + """Return the number of output slots.""" + return len(self.output_indices) + + @property + def norm_warning_threshold(self) -> float: + """Return the largest normalization deviation that passes without a warning.""" + return _NORM_WARNING_THRESHOLDS[self.dtype] + def plan_component( - component: Graph, layout: ParamLayout + component: Graph, output_indices: Sequence[int], layout: ParamLayout ) -> tuple[Plan, ComponentLayout]: """Extract, rewrite, schedule and plan one non-direct component. @@ -88,79 +111,66 @@ def plan_component( ``MAX_WIDTH`` array axes. """ - component_layout = ComponentLayout.of(component, layout) - outputs = list(component_layout.output_indices) + component_layout = ComponentLayout.of(component, output_indices, layout) + outputs = list(output_indices) try: circuit = extract_circuit(component) except (RuntimeError, ValueError) as error: raise ValueError( "strategy='statevector' could not extract a circuit from the component " - f"with outputs {outputs}: {error}. {_HINT}" + f"with outputs {outputs}: {error}. {STABILIZER_RANK_HINT}" ) from error ops = heisenberg_ops(circuit, component_layout) scheduled = schedule(ops) if scheduled.width > MAX_WIDTH: raise ValueError( f"strategy='statevector' needs {scheduled.width} active qubits for the " - f"component with outputs {outputs}, more than {MAX_WIDTH}. {_HINT}" + f"component with outputs {outputs}, more than {MAX_WIDTH}. " + f"{STABILIZER_RANK_HINT}" ) return plan(ops, scheduled), component_layout -def resolve_dtype(dtype: str | None) -> str: - """Return the amplitude dtype: complex128 if JAX has 64-bit types, else complex64. - - Raises: - ValueError: If ``dtype`` is unknown, or complex128 without - ``jax_enable_x64``. - - """ - x64 = jax.dtypes.canonicalize_dtype(np.complex128) == np.dtype(np.complex128) - if dtype is None: - return "complex128" if x64 else "complex64" - if dtype not in ("complex64", "complex128"): - raise ValueError(f"dtype must be 'complex64' or 'complex128', got {dtype!r}") - if dtype == "complex128" and not x64: - raise ValueError("complex128 amplitudes need jax_enable_x64") - return dtype +def _pack(terms: Iterable[int], num_words: int) -> np.ndarray: + """Return symbol bits as uint32 words: symbol s is bit s % 32 of word s // 32.""" + row = np.zeros(num_words, dtype=np.uint32) + for term in terms: + row[term // 32] |= np.uint32(1) << np.uint32(term % 32) + return row def lower(plan: Plan, layout: ComponentLayout, dtype: str) -> StatevectorComponent: """Return the action tape of ``plan``.""" num_actions = len(plan.actions) - num_words = max(1, -(-plan.num_symbols // 32)) + num_words = max(1, ceil(plan.num_symbols / 32)) kind = np.zeros(num_actions, dtype=np.int32) x_mask = np.zeros(num_actions, dtype=np.uint32) z_mask = np.zeros(num_actions, dtype=np.uint32) angle = np.zeros(num_actions, dtype=np.float64) pivot = np.zeros(num_actions, dtype=np.int32) branch = np.full(num_actions, -1, dtype=np.int32) - slot = np.full(num_actions, -1, dtype=np.int32) sign_const = np.zeros(num_actions, dtype=bool) sign_row = np.zeros((num_actions, num_words), dtype=np.uint32) for a, action in enumerate(plan.actions): if isinstance(action, Rotate): kind[a] = ROTATE if action.x else ROTATE_DIAG x_mask[a], z_mask[a], angle[a] = action.x, action.z, action.angle + sign_const[a] = action.sign.constant + sign_row[a] = _pack(action.sign.terms, num_words) elif isinstance(action, MeasureActive): kind[a] = MEASURE if action.x else MEASURE_DIAG x_mask[a], z_mask[a] = action.x, action.z - pivot[a], branch[a], slot[a] = action.pivot, action.branch, action.slot - elif isinstance(action, MeasureDormant): - kind[a] = MEASURE_DORMANT - branch[a], slot[a] = action.branch, action.slot + pivot[a], branch[a] = action.pivot, action.branch else: - kind[a] = RECORD - slot[a] = action.slot - sign_const[a] = action.sign.constant - for term in action.sign.terms: - sign_row[a, term // 32] |= np.uint32(1) << np.uint32(term % 32) - real = np.float32 if dtype == "complex64" else np.float64 + kind[a] = MEASURE_DORMANT + branch[a] = action.branch + output_row = np.zeros((len(plan.outputs), num_words), dtype=np.uint32) + for slot, output in enumerate(plan.outputs): + output_row[slot] = _pack(output.terms, num_words) + real = np.finfo(dtype).dtype return StatevectorComponent( output_indices=layout.output_indices, width=plan.width, - num_words=num_words, - num_outputs=plan.num_outputs, dtype=dtype, f_selection=jnp.asarray(np.array(layout.f_selection, dtype=np.int32)), kind=jnp.asarray(kind), @@ -170,24 +180,26 @@ def lower(plan: Plan, layout: ComponentLayout, dtype: str) -> StatevectorCompone sin=jnp.asarray(np.sin(angle).astype(real)), pivot=jnp.asarray(pivot), branch=jnp.asarray(branch), - slot=jnp.asarray(slot), sign_const=jnp.asarray(sign_const), sign_row=jnp.asarray(sign_row), + output_const=jnp.asarray([output.constant for output in plan.outputs], bool), + output_row=jnp.asarray(output_row), ) def compile_statevector_component( - component: Graph, layout: ParamLayout, *, dtype: str | None = None + component: Graph, output_indices: Sequence[int], layout: ParamLayout ) -> StatevectorComponent: """Compile one non-direct component for ``strategy="statevector"``. + Amplitudes are complex128 when JAX has 64-bit types enabled, else complex64. + Args: component: A connected component of the prepared sampling graph. + output_indices: Global output index of each of the component's outputs. layout: The prepared graph's parameter layout. - dtype: Amplitude dtype; ``None`` picks complex128 when JAX has 64-bit - types enabled, else complex64. """ - resolved = resolve_dtype(dtype) - planned, component_layout = plan_component(component, layout) - return lower(planned, component_layout, resolved) + planned, component_layout = plan_component(component, output_indices, layout) + dtype = jax.dtypes.canonicalize_dtype(np.complex128).name + return lower(planned, component_layout, dtype) diff --git a/src/tsim/statevector/schedule.py b/src/tsim/statevector/schedule.py index a43ca4cd..39185b67 100644 --- a/src/tsim/statevector/schedule.py +++ b/src/tsim/statevector/schedule.py @@ -14,7 +14,7 @@ import numpy as np -from tsim.statevector.heisenberg import MEASUREMENT, HeisenbergOps +from tsim.statevector.heisenberg import MEASUREMENT, HeisenbergOps, anticommute class StabilizerGroup: @@ -32,28 +32,29 @@ def width(self) -> int: def anticommuting(self, x: np.ndarray, z: np.ndarray) -> np.ndarray: """Return which generators anticommute with the Pauli ``(x, z)``.""" - return np.logical_xor.reduce(self.x & z, axis=1) ^ np.logical_xor.reduce( - self.z & x, axis=1 - ) + return anticommute(self.x, self.z, x, z) - def rotate(self, x: np.ndarray, z: np.ndarray) -> None: - """Keep the generators that a rotation about ``(x, z)`` preserves.""" + def _eliminate(self, x: np.ndarray, z: np.ndarray) -> int | None: + """Leave one generator anticommuting with ``(x, z)``; return it, if any.""" anti = np.flatnonzero(self.anticommuting(x, z)) if len(anti) == 0: - return + return None first, rest = anti[0], anti[1:] self.x[rest] ^= self.x[first] self.z[rest] ^= self.z[first] - self.x = np.delete(self.x, first, axis=0) - self.z = np.delete(self.z, first, axis=0) + return int(first) + + def rotate(self, x: np.ndarray, z: np.ndarray) -> None: + """Keep the generators that a rotation about ``(x, z)`` preserves.""" + first = self._eliminate(x, z) + if first is not None: + self.x = np.delete(self.x, first, axis=0) + self.z = np.delete(self.z, first, axis=0) def measure(self, x: np.ndarray, z: np.ndarray) -> None: """Update the generators after measuring the Pauli ``(x, z)``.""" - anti = np.flatnonzero(self.anticommuting(x, z)) - if len(anti): - first, rest = anti[0], anti[1:] - self.x[rest] ^= self.x[first] - self.z[rest] ^= self.z[first] + first = self._eliminate(x, z) + if first is not None: self.x[first] = x self.z[first] = z elif not self.contains(x, z): @@ -61,10 +62,13 @@ def measure(self, x: np.ndarray, z: np.ndarray) -> None: self.z = np.vstack([self.z, z[None]]) def contains(self, x: np.ndarray, z: np.ndarray) -> bool: - """Return whether ``(x, z)`` is a product of the generators, up to sign.""" + """Return whether ``(x, z)`` is a product of the generators, up to sign. + + The generators commute and are independent, so their rank is their number. + """ matrix = np.hstack([self.x, self.z]) extended = np.vstack([matrix, np.concatenate([x, z])[None]]) - return _rank(extended) == _rank(matrix) + return _rank(extended) == len(matrix) def _rank(matrix: np.ndarray) -> int: @@ -88,9 +92,7 @@ def _rank(matrix: np.ndarray) -> int: def anticommutation(ops: HeisenbergOps) -> np.ndarray: """Return the symmetric matrix of which pairs of operations anticommute.""" - x = ops.x.astype(np.float32) - z = ops.z.astype(np.float32) - return (x @ z.T + z @ x.T) % 2 == 1 + return anticommute(ops.x, ops.z, ops.x, ops.z) def _ancestors(before: np.ndarray) -> np.ndarray: @@ -148,7 +150,7 @@ def run(j: int) -> None: while progressed: progressed = False for j in np.flatnonzero(~done): - if waiting[j] or done[j]: + if waiting[j]: continue keeps_width = not group.anticommuting(ops.x[j], ops.z[j]).any() if is_measurement[j] or keeps_width: diff --git a/test/helpers/statevector.py b/test/helpers/statevector.py index 602c4ba2..5e7283ea 100644 --- a/test/helpers/statevector.py +++ b/test/helpers/statevector.py @@ -6,8 +6,9 @@ from __future__ import annotations import itertools -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Mapping, Sequence from fractions import Fraction +from functools import cache from math import cos, pi, sin import jax.numpy as jnp @@ -27,7 +28,7 @@ HeisenbergOps, heisenberg_ops, ) -from tsim.statevector.plan import MeasureActive, Plan, Record, Rotate +from tsim.statevector.plan import MeasureActive, Plan, Rotate Distribution = dict[tuple[int, ...], float] @@ -167,7 +168,31 @@ def quantum_components( return components, prepared.layout -def hand_built_circuits() -> list[tuple[str, ExtractedCircuit, ComponentLayout]]: +def output_indices(component: Graph) -> list[int]: + """Return the global output index of each of ``component``'s outputs.""" + return [ + port for v in component.outputs() if (port := component.port(v)) is not None + ] + + +ONE_LINE_LAYOUT = ComponentLayout( + f_selection=(0,), f_position={7: 0}, output_indices=(0,) +) + + +def one_line(gates: list[Gate], *, line_start: int | None = None) -> ExtractedCircuit: + """Return a one-line circuit that ends at output 0; phase parameter 7 is f-bit 0.""" + return ExtractedCircuit( + gates=gates, + n_lines=1, + line_start={0: line_start}, + line_end={0: 0}, + scalar=Scalar.one(), + classical_ports=frozenset({0}), + ) + + +def hand_built_circuits() -> list[tuple[ExtractedCircuit, ComponentLayout]]: """Return small extracted circuits with a DISCARD and a reused MEASURE line.""" t = Phase(Fraction(1, 4)) noise = Phase(Fraction(0), frozenset({7})) @@ -207,23 +232,30 @@ def hand_built_circuits() -> list[tuple[str, ExtractedCircuit, ComponentLayout]] classical_ports=frozenset({0, 1}), ) return [ - ("discard", discard, ComponentLayout((0,), {7: 0}, (0,))), - ("measure", measure, ComponentLayout((0,), {7: 0}, (0, 1))), + (discard, ONE_LINE_LAYOUT), + (measure, ComponentLayout((0,), {7: 0}, (0, 1))), ] -def component_ops() -> Iterator[HeisenbergOps]: - """Yield the ops of every test component (both modes) and hand-built circuit.""" +@cache +def extracted_circuits() -> tuple[tuple[ExtractedCircuit, ComponentLayout], ...]: + """Return every test component's circuit (both modes), then the hand-built ones.""" + extracted = [] for name in sorted(CIRCUITS): for sample_detectors in (True, False): components, layout = quantum_components( CIRCUITS[name], sample_detectors=sample_detectors ) for component in components: - component_layout = ComponentLayout.of(component, layout) - yield heisenberg_ops(extract_circuit(component), component_layout) - for _, circuit, component_layout in hand_built_circuits(): - yield heisenberg_ops(circuit, component_layout) + indices = output_indices(component) + component_layout = ComponentLayout.of(component, indices, layout) + extracted.append((extract_circuit(component), component_layout)) + return (*extracted, *hand_built_circuits()) + + +def component_ops() -> list[HeisenbergOps]: + """Return the ops of every circuit in ``extracted_circuits``.""" + return [heisenberg_ops(circuit, layout) for circuit, layout in extracted_circuits()] def f_patterns(num_f: int, limit: int = 32) -> list[tuple[int, ...]]: @@ -408,32 +440,18 @@ def plan_distribution(plan: Plan, f: Sequence[int]) -> Distribution: """Return the exact output distribution of a plan by enumerating its branches.""" result: Distribution = {} - def written(record: list[int], slot: int, value: int) -> list[int]: - if slot < 0: - return record - updated = list(record) - updated[slot] = value - return updated - - def walk( - a: int, psi: np.ndarray, symbols: list[int], record: list[int], weight: float - ) -> None: + def walk(a: int, psi: np.ndarray, symbols: list[int], weight: float) -> None: if weight < 1e-14: return if a == len(plan.actions): - key = tuple(record) + key = tuple(output.evaluate(symbols) for output in plan.outputs) result[key] = result.get(key, 0.0) + weight return action = plan.actions[a] - sign = action.sign.evaluate(symbols) if isinstance(action, Rotate): - rotated = reference_rotate( - psi, action.x, action.z, action.angle, bool(sign) - ) - walk(a + 1, rotated, symbols, record, weight) - return - if isinstance(action, Record): - walk(a + 1, psi, symbols, written(record, action.slot, sign), weight) + negate = bool(action.sign.evaluate(symbols)) + rotated = reference_rotate(psi, action.x, action.z, action.angle, negate) + walk(a + 1, rotated, symbols, weight) return for outcome in (0, 1): if isinstance(action, MeasureActive): @@ -444,12 +462,11 @@ def walk( new, p = psi, 0.5 branched = list(symbols) branched[action.branch] = outcome - value = outcome ^ sign - walk(a + 1, new, branched, written(record, action.slot, value), weight * p) + walk(a + 1, new, branched, weight * p) psi = np.zeros(2**plan.width, dtype=complex) psi[0] = 1 - walk(0, psi, [*f, *[0] * plan.num_branches], [0] * plan.num_outputs, 1.0) + walk(0, psi, [*f, *[0] * plan.num_branches], 1.0) return result diff --git a/test/integration/test_statevector_sampler.py b/test/integration/test_statevector_sampler.py index b880ab88..567cea72 100644 --- a/test/integration/test_statevector_sampler.py +++ b/test/integration/test_statevector_sampler.py @@ -15,6 +15,7 @@ from tsim.core.graph import prepare_graph from tsim.core.types import SamplingStrategy from tsim.external.vec_sim.vec_sampler import VecSampler +from tsim.statevector.execute import lane_bytes from tsim.statevector.program import StatevectorComponent @@ -61,8 +62,6 @@ def test_statevector_samplers_hold_statevector_components(): assert components assert all(isinstance(c, StatevectorComponent) for c in components) assert "statevector components" in repr(sampler) - widest = max(c.width for c in components if isinstance(c, StatevectorComponent)) - assert sampler._peak_bytes_per_sample() >= 3 * 2**widest * 8 def test_postselection_and_partial_batches(): @@ -117,27 +116,14 @@ def test_complex128_sampling_matches_stabilizer_rank(): assert_counts_match(got, expected) -def test_a_lane_that_does_not_fit_in_memory_is_rejected(monkeypatch): - monkeypatch.setattr( - sampler_module._CompiledSamplerBase, "_available_memory", lambda self: 1 - ) - with pytest.raises(ValueError, match="cat5"): - Circuit(CIRCUITS["t_ghz"]).compile_detector_sampler(strategy="statevector") - - def _norm_warnings(monkeypatch, strategy, deviation, *, x64=False): - name = ( - "sample_statevector_component" - if strategy == "statevector" - else "sample_component" - ) - sample = getattr(sampler_module, name) + sample = sampler_module.sample_component def drifting(component, f_params, key): samples, key, _ = sample(component, f_params, key) return samples, key, jnp.asarray(deviation) - monkeypatch.setattr(sampler_module, name, drifting) + monkeypatch.setattr(sampler_module, "sample_component", drifting) with jax.enable_x64(x64): sampler = Circuit(CIRCUITS["t_ghz"]).compile_detector_sampler( strategy=strategy, seed=0 @@ -178,7 +164,7 @@ def test_a_vanishing_statevector_norm_raises(monkeypatch): def test_the_automatic_batch_size_of_a_statevector_sampler_is_capped(monkeypatch): monkeypatch.setattr( - sampler_module._CompiledSamplerBase, "_available_memory", lambda self: 2**50 + sampler_module._CompiledSamplerBase, "_memory_budget", lambda self: 2**49 ) circuit = Circuit(CIRCUITS["t_ghz"]) sampler = circuit.compile_detector_sampler(strategy="statevector", seed=0) @@ -205,15 +191,14 @@ def test_the_per_shot_estimate_counts_f_params_and_records(): sampler = Circuit(CIRCUITS["t_ghz"]).compile_detector_sampler( strategy="statevector" ) - widest = max( - c.width + lane = max( + lane_bytes(c) for c in sampler._program.components if isinstance(c, StatevectorComponent) ) num_f = sampler._channel_sampler.signature_matrix.shape[1] records = sampler._program.num_outputs - amplitudes = 3 * 2**widest * 8 - assert sampler._peak_bytes_per_sample() >= amplitudes + 2 * num_f + 2 * records + assert sampler._peak_bytes_per_sample() == lane + 2 * num_f + 2 * records def test_a_lane_must_fit_in_half_of_the_available_memory(monkeypatch): @@ -222,8 +207,8 @@ def test_a_lane_must_fit_in_half_of_the_available_memory(monkeypatch): strategy="statevector" )._peak_bytes_per_sample() base = sampler_module._CompiledSamplerBase - monkeypatch.setattr(base, "_available_memory", lambda self: 2 * peak - 1) + monkeypatch.setattr(base, "_memory_budget", lambda self: peak - 1) with pytest.raises(ValueError, match="cat5"): circuit.compile_detector_sampler(strategy="statevector") - monkeypatch.setattr(base, "_available_memory", lambda self: 2 * peak) + monkeypatch.setattr(base, "_memory_budget", lambda self: peak) circuit.compile_detector_sampler(strategy="statevector") diff --git a/test/unit/statevector/test_execute.py b/test/unit/statevector/test_execute.py index 3586de75..a7229501 100644 --- a/test/unit/statevector/test_execute.py +++ b/test/unit/statevector/test_execute.py @@ -3,6 +3,7 @@ apply_pauli, assert_matches_distribution, f_patterns, + output_indices, pauli_matrix, plan_distribution, quantum_components, @@ -17,40 +18,28 @@ import tsim.statevector.execute as execute from tsim.statevector.execute import ( initial_state, + outputs, pack_bits, run, sample_statevector_component, ) from tsim.statevector.heisenberg import ComponentLayout -from tsim.statevector.plan import ( - AffineRow, - MeasureActive, - MeasureDormant, - Plan, - Record, - Rotate, -) +from tsim.statevector.plan import AffineRow, MeasureActive, MeasureDormant, Plan, Rotate from tsim.statevector.program import ( MEASURE, MEASURE_DIAG, MEASURE_DORMANT, + compile_statevector_component, lower, plan_component, - resolve_dtype, ) -def _component( - actions, *, width, num_f=1, num_branches=0, num_outputs=0, dtype="complex64" -): +def _component(actions, *, width, num_f=1, outputs=(), dtype="complex64"): planned = Plan( - width=width, - num_f=num_f, - num_branches=num_branches, - num_outputs=num_outputs, - actions=tuple(actions), + width=width, num_f=num_f, actions=tuple(actions), outputs=tuple(outputs) ) - layout = ComponentLayout(tuple(range(num_f)), {}, tuple(range(num_outputs))) + layout = ComponentLayout(tuple(range(num_f)), {}, tuple(range(len(outputs)))) return lower(planned, layout, dtype) @@ -60,6 +49,27 @@ def _random_states(lanes, width, seed=0): return psi / np.linalg.norm(psi, axis=1, keepdims=True) +def _measuring(x, z, pivot, psi, f_bits=None): + """Return a component that measures ``x``, ``z``, and lanes in the states ``psi``. + + Its one output is the outcome XOR the f-bit. + """ + component = _component( + [MeasureActive(x, z, pivot, branch=1)], + width=3, + outputs=[AffineRow(False, frozenset({0, 1}))], + ) + if f_bits is None: + f_bits = jnp.zeros((len(psi), 1), dtype=bool) + amps = jnp.asarray(psi, jnp.complex64) + return component, initial_state(component, f_bits)._replace(amps=amps) + + +def _only_component(name): + (component,), layout = quantum_components(CIRCUITS[name], sample_detectors=True) + return component, output_indices(component), layout + + def test_pack_bits(): bits = np.zeros((2, 40), dtype=bool) bits[0, [0, 31, 32, 39]] = True @@ -90,45 +100,66 @@ def test_rotations_match_dense_matrices(x, z, dtype): ) -@pytest.mark.parametrize( - "x, z, pivot", [(0b011, 0b110, 0), (0b010, 0b011, 1), (0, 0b101, 2), (0, 0b011, 0)] -) +_MEASUREMENTS = [ + (0b011, 0b110, 0), + (0b010, 0b011, 1), + (0b110, 0b010, 2), + (0, 0b101, 2), + (0, 0b011, 0), +] + + +@pytest.mark.parametrize("x, z, pivot", _MEASUREMENTS) def test_measurements_match_the_reference_kernel(monkeypatch, x, z, pivot): - action = MeasureActive( - x, z, pivot, branch=1, sign=AffineRow(False, frozenset({0})), slot=0 - ) - component = _component([action], width=3, num_branches=1, num_outputs=1) - outcomes = jnp.array([False, True, True]) - monkeypatch.setattr(execute, "choose_outcome", lambda p1, count, key: outcomes) psi = _random_states(3, 3) f_bits = jnp.array([[False], [False], [True]]) - state = initial_state(component, f_bits)._replace( - amps=jnp.asarray(psi, jnp.complex64) - ) + component, state = _measuring(x, z, pivot, psi, f_bits) + outcomes = jnp.array([False, True, True]) + monkeypatch.setattr(execute, "choose_outcome", lambda p1, count, key: outcomes) out = run(component, state, jax.random.key(0), 0, 1) for lane, outcome in enumerate([0, 1, 1]): expected, _ = reference_measure(psi[lane], x, z, pivot, outcome) np.testing.assert_allclose(np.asarray(out.amps[lane]), expected, atol=1e-5) - assert np.asarray(out.records[:, 0]).tolist() == [False, True, False] + assert np.asarray(outputs(component, out)[:, 0]).tolist() == [False, True, False] assert ((np.asarray(out.symbols[:, 0]) >> 1) & 1).tolist() == [0, 1, 1] -def test_dormant_measurements_and_records_write_outcome_xor_sign(monkeypatch): +def test_outputs_are_parities_of_the_symbols(monkeypatch): component = _component( - [ - MeasureDormant(branch=1, sign=AffineRow(True, frozenset()), slot=0), - Record(sign=AffineRow(False, frozenset({0, 1})), slot=1), - ], + [MeasureDormant(branch=1)], width=0, - num_branches=1, - num_outputs=2, + outputs=[AffineRow(True, frozenset({1})), AffineRow(False, frozenset({0, 1}))], ) monkeypatch.setattr( execute, "choose_outcome", lambda p1, count, key: jnp.array([False, True]) ) state = initial_state(component, jnp.array([[True], [True]])) - out = run(component, state, jax.random.key(0), 0, 2) - assert np.asarray(out.records[:, :2]).tolist() == [[True, True], [False, False]] + out = run(component, state, jax.random.key(0), 0, 1) + assert np.asarray(outputs(component, out)).tolist() == [ + [True, True], + [False, False], + ] + + +def test_outputs_read_symbols_beyond_the_first_word(monkeypatch): + component = _component( + [MeasureDormant(branch=40)], + width=0, + num_f=40, + outputs=[ + AffineRow(False, frozenset({40})), + AffineRow(True, frozenset({3, 40})), + ], + ) + assert component.num_words == 2 + monkeypatch.setattr( + execute, "choose_outcome", lambda p1, count, key: jnp.array([False, True]) + ) + f_bits = np.zeros((2, 40), dtype=bool) + f_bits[1, 3] = True + state = initial_state(component, jnp.asarray(f_bits)) + out = run(component, state, jax.random.key(0), 0, 1) + assert np.asarray(outputs(component, out)).tolist() == [[False, True], [True, True]] def test_a_component_without_actions_samples_no_outputs(): @@ -159,13 +190,15 @@ def test_sampling_matches_the_plan_distribution(name): shots = 20_000 components, layout = quantum_components(CIRCUITS[name], sample_detectors=True) for component in components: - planned, component_layout = plan_component(component, layout) + indices = output_indices(component) + planned, component_layout = plan_component(component, indices, layout) + f_selection = component_layout.f_selection compiled = lower(planned, component_layout, "complex64") - num_columns = max(component_layout.f_selection, default=-1) + 1 + num_columns = max(f_selection, default=-1) + 1 for f in f_patterns(planned.num_f, limit=2): f_params = np.zeros((shots, num_columns), dtype=np.uint8) - if component_layout.f_selection: - f_params[:, list(component_layout.f_selection)] = f + if f_selection: + f_params[:, list(f_selection)] = f samples, _, deviation = sample_statevector_component( compiled, jnp.asarray(f_params), jax.random.key(7) ) @@ -175,18 +208,16 @@ def test_sampling_matches_the_plan_distribution(name): ) -def test_complex128_amplitudes(): - with jax.enable_x64(True): - (component,), layout = quantum_components( - CIRCUITS["t_ghz"], sample_detectors=True - ) - planned, component_layout = plan_component(component, layout) - compiled = lower(planned, component_layout, resolve_dtype(None)) - assert compiled.dtype == "complex128" - state = initial_state(compiled, jnp.zeros((4, planned.num_f), dtype=bool)) - assert state.amps.dtype == jnp.complex128 +@pytest.mark.parametrize("x64", [False, True]) +def test_amplitudes_are_complex128_with_64_bit_types(x64): + dtype = "complex128" if x64 else "complex64" + with jax.enable_x64(x64): + compiled = compile_statevector_component(*_only_component("t_ghz")) + assert compiled.dtype == dtype + num_f = compiled.f_selection.shape[0] + state = initial_state(compiled, jnp.zeros((4, num_f), dtype=bool)) out = run(compiled, state, jax.random.key(0), 0, compiled.num_actions) - assert out.amps.dtype == jnp.complex128 + assert out.amps.dtype == jnp.dtype(dtype) def _count_outcomes(p1, dtype, draws, chunk=2**22): @@ -242,29 +273,20 @@ def recording(p1, count, key): return p1, out -_MEASUREMENTS = [(0b011, 0b110, 0), (0b110, 0b010, 2), (0, 0b101, 2), (0, 0b011, 0)] - - @pytest.mark.parametrize("eigenvalue", [1, -1]) @pytest.mark.parametrize("x, z, pivot", _MEASUREMENTS) def test_measuring_an_eigenstate_never_samples_the_impossible_outcome( monkeypatch, x, z, pivot, eigenvalue ): lanes = 4096 - action = MeasureActive( - x, z, pivot, branch=1, sign=AffineRow(False, frozenset()), slot=0 - ) - component = _component([action], width=3, num_branches=1, num_outputs=1) # NOTE: the noise is of the size rotations leave in complex64 amplitudes. rounding = 1e-7 * _random_states(lanes, 3, seed=1) psi = _projected(_random_states(lanes, 3), x, z, eigenvalue) + rounding - state = initial_state(component, jnp.zeros((lanes, 1), dtype=bool))._replace( - amps=jnp.asarray(psi, jnp.complex64) - ) + component, state = _measuring(x, z, pivot, psi) p1, out = _run_recording_p1(monkeypatch, component, state) impossible = p1 if eigenvalue == 1 else 1 - p1 assert impossible.max() < 1e-12 - assert np.all(np.asarray(out.records[:, 0]) == (eigenvalue == -1)) + assert np.all(np.asarray(outputs(component, out)[:, 0]) == (eigenvalue == -1)) @pytest.mark.parametrize("unlikely", [0, 1]) @@ -273,37 +295,29 @@ def test_an_unlikely_outcome_leaves_a_normalized_state( monkeypatch, x, z, pivot, unlikely ): lanes = 64 - action = MeasureActive( - x, z, pivot, branch=1, sign=AffineRow(False, frozenset()), slot=0 - ) - component = _component([action], width=3, num_branches=1, num_outputs=1) vectors = _random_states(lanes, 3) likely = _projected(vectors, x, z, (-1) ** (1 - unlikely)) rare = _projected(vectors, x, z, (-1) ** unlikely) p = np.logspace(-5, -3, lanes)[:, None] - psi = np.sqrt(1 - p) * likely + np.sqrt(p) * rare + component, state = _measuring( + x, z, pivot, np.sqrt(1 - p) * likely + np.sqrt(p) * rare + ) outcomes = jnp.full(lanes, bool(unlikely)) monkeypatch.setattr(execute, "choose_outcome", lambda p1, count, key: outcomes) - state = initial_state(component, jnp.zeros((lanes, 1), dtype=bool))._replace( - amps=jnp.asarray(psi, jnp.complex64) - ) out = run(component, state, jax.random.key(0), 0, 1) norms = np.linalg.norm(np.asarray(out.amps, dtype=np.complex128), axis=1) ** 2 assert np.abs(norms - 1).max() < 1e-5 def test_running_a_tape_in_two_parts_equals_running_it_whole(): - (component,), layout = quantum_components( - CIRCUITS["midcircuit_feedforward"], sample_detectors=True - ) - planned, component_layout = plan_component(component, layout) - compiled = lower(planned, component_layout, "complex64") + compiled = compile_statevector_component(*_only_component("midcircuit_feedforward")) lanes, split, stop = 512, 4, compiled.num_actions kinds = np.asarray(compiled.kind) measurements = np.isin(kinds, [MEASURE_DIAG, MEASURE, MEASURE_DORMANT]) assert measurements[:split].any() and measurements[split:].any() rng = np.random.default_rng(0) - f_bits = jnp.asarray(rng.integers(0, 2, (lanes, planned.num_f)).astype(bool)) + num_f = compiled.f_selection.shape[0] + f_bits = jnp.asarray(rng.integers(0, 2, (lanes, num_f)).astype(bool)) state = initial_state(compiled, f_bits) key = jax.random.key(5) whole = run(compiled, state, key, 0, stop) diff --git a/test/unit/statevector/test_heisenberg.py b/test/unit/statevector/test_heisenberg.py index 7ac2009b..55989844 100644 --- a/test/unit/statevector/test_heisenberg.py +++ b/test/unit/statevector/test_heisenberg.py @@ -2,17 +2,20 @@ from math import pi from test.helpers.statevector import ( CIRCUITS, + ONE_LINE_LAYOUT, assert_distributions_close, + extracted_circuits, extracted_distribution, f_patterns, - hand_built_circuits, + one_line, ops_distribution, + output_indices, quantum_components, ) import pytest -from zxparam import Phase, Scalar -from zxparam.extract import ExtractedCircuit, Gate, extract_circuit +from zxparam import Phase +from zxparam.extract import Gate from tsim.statevector.heisenberg import ( MEASUREMENT, @@ -21,22 +24,11 @@ heisenberg_ops, ) -LAYOUT = ComponentLayout(f_selection=(0,), f_position={7: 0}, output_indices=(0,)) - - -def _one_line(gates, *, line_start=None): - return ExtractedCircuit( - gates=gates, - n_lines=1, - line_start={0: line_start}, - line_end={0: 0}, - scalar=Scalar.one(), - classical_ports=frozenset({0}), - ) - def test_t_on_plus_is_a_rotation_about_x(): - ops = heisenberg_ops(_one_line([Gate("Z", (0,), Phase(Fraction(1, 4)))]), LAYOUT) + ops = heisenberg_ops( + one_line([Gate("Z", (0,), Phase(Fraction(1, 4)))]), ONE_LINE_LAYOUT + ) assert ops.kind.tolist() == [ROTATION, MEASUREMENT] assert ops.x.tolist() == [[True], [True]] assert ops.z.tolist() == [[False], [False]] @@ -48,7 +40,9 @@ def test_t_on_plus_is_a_rotation_about_x(): def test_clifford_part_of_a_phase_goes_into_the_frame(): # 3π/4 = π (a Z, into the frame) − π/4 (a rotation by −π/8). - ops = heisenberg_ops(_one_line([Gate("Z", (0,), Phase(Fraction(3, 4)))]), LAYOUT) + ops = heisenberg_ops( + one_line([Gate("Z", (0,), Phase(Fraction(3, 4)))]), ONE_LINE_LAYOUT + ) assert ops.kind.tolist() == [ROTATION, MEASUREMENT] assert ops.angle[0] == pytest.approx(-pi / 8) assert ops.sign[:, 0].tolist() == [False, False] @@ -62,7 +56,7 @@ def test_noise_flips_the_ops_it_anticommutes_with(): Gate("H", (0,)), Gate("Z", (0,), Phase(Fraction(1, 4))), ] - ops = heisenberg_ops(_one_line(gates), LAYOUT) + ops = heisenberg_ops(one_line(gates), ONE_LINE_LAYOUT) assert ops.kind.tolist() == [ROTATION, MEASUREMENT] assert ops.x.tolist() == [[False], [False]] assert ops.z.tolist() == [[True], [True]] @@ -71,28 +65,20 @@ def test_noise_flips_the_ops_it_anticommutes_with(): def test_a_line_from_an_input_is_rejected(): with pytest.raises(ValueError, match="input port"): - heisenberg_ops(_one_line([], line_start=3), LAYOUT) + heisenberg_ops(one_line([], line_start=3), ONE_LINE_LAYOUT) def test_component_layout_reads_f_bits_and_ports(): (component,), layout = quantum_components(CIRCUITS["t_ghz"], sample_detectors=True) - component_layout = ComponentLayout.of(component, layout) + component_layout = ComponentLayout.of(component, output_indices(component), layout) assert list(component_layout.f_selection) == sorted(component_layout.f_selection) assert len(component_layout.f_selection) == 3 assert component_layout.output_indices == (0, 1, 2) assert sorted(component_layout.f_position.values()) == [0, 1, 2] -@pytest.mark.parametrize("name", sorted(CIRCUITS)) -@pytest.mark.parametrize("sample_detectors", [True, False]) -def test_ops_reproduce_the_extracted_circuit(name, sample_detectors): - components, layout = quantum_components( - CIRCUITS[name], sample_detectors=sample_detectors - ) - assert components - for component in components: - component_layout = ComponentLayout.of(component, layout) - circuit = extract_circuit(component) +def test_ops_reproduce_the_extracted_circuits(): + for circuit, component_layout in extracted_circuits(): ops = heisenberg_ops(circuit, component_layout) for f in f_patterns(ops.num_f): bindings = {p: f[i] for p, i in component_layout.f_position.items()} @@ -100,15 +86,3 @@ def test_ops_reproduce_the_extracted_circuit(name, sample_detectors): ops_distribution(ops, f), extracted_distribution(circuit, component_layout, bindings), ) - - -@pytest.mark.parametrize("index", [0, 1]) -def test_ops_reproduce_hand_built_circuits(index): - _, circuit, component_layout = hand_built_circuits()[index] - ops = heisenberg_ops(circuit, component_layout) - for f in f_patterns(ops.num_f): - bindings = {p: f[i] for p, i in component_layout.f_position.items()} - assert_distributions_close( - ops_distribution(ops, f), - extracted_distribution(circuit, component_layout, bindings), - ) diff --git a/test/unit/statevector/test_plan.py b/test/unit/statevector/test_plan.py index 667327c6..627b42bb 100644 --- a/test/unit/statevector/test_plan.py +++ b/test/unit/statevector/test_plan.py @@ -1,9 +1,11 @@ from fractions import Fraction from math import pi from test.helpers.statevector import ( + ONE_LINE_LAYOUT, assert_distributions_close, component_ops, f_patterns, + one_line, ops_distribution, plan_distribution, ) @@ -11,13 +13,12 @@ import numpy as np import pytest import stim -from zxparam import Phase, Scalar -from zxparam.extract import ExtractedCircuit, Gate +from zxparam import Phase +from zxparam.extract import Gate from tsim.statevector.heisenberg import ( MEASUREMENT, ROTATION, - ComponentLayout, HeisenbergOps, heisenberg_ops, ) @@ -35,15 +36,8 @@ def _t_then_h(): # |+⟩, T, H, measure Z: in the |0⟩ frame a rotation about X, then a # measurement of Z, which anticommute. - circuit = ExtractedCircuit( - gates=[Gate("Z", (0,), Phase(Fraction(1, 4))), Gate("H", (0,))], - n_lines=1, - line_start={0: None}, - line_end={0: 0}, - scalar=Scalar.one(), - classical_ports=frozenset({0}), - ) - return heisenberg_ops(circuit, ComponentLayout((0,), {7: 0}, (0,))) + gates = [Gate("Z", (0,), Phase(Fraction(1, 4))), Gate("H", (0,))] + return heisenberg_ops(one_line(gates), ONE_LINE_LAYOUT) def test_t_then_h_promotes_one_axis_and_measures_it(): @@ -56,9 +50,8 @@ def test_t_then_h_promotes_one_axis_and_measures_it(): assert (rotate.x, rotate.z) == (1, 0) assert rotate.angle == pytest.approx(pi / 8) assert rotate.sign == AffineRow(False, frozenset()) - assert measure == MeasureActive( - x=0, z=1, pivot=0, branch=1, sign=AffineRow(False, frozenset()), slot=0 - ) + assert measure == MeasureActive(x=0, z=1, pivot=0, branch=1) + assert planned.outputs == (AffineRow(False, frozenset({1})),) def test_affine_rows_xor_and_evaluate(): @@ -74,26 +67,31 @@ def test_a_schedule_that_is_too_narrow_is_rejected(): plan(ops, Schedule(order=np.arange(ops.num_ops, dtype=np.int32), width=0)) +def _checked_plan(ops): + """Plan ``ops`` and check that the plan reproduces the scheduled ops.""" + scheduled = schedule(ops) + planned = plan(ops, scheduled) + reordered = ops.take(scheduled.order) + for f in f_patterns(ops.num_f): + assert_distributions_close( + plan_distribution(planned, f), ops_distribution(reordered, f), atol=1e-9 + ) + return planned + + def test_plans_reproduce_the_scheduled_ops(): for ops in component_ops(): - scheduled = schedule(ops) - planned = plan(ops, scheduled) - reordered = ops.take(scheduled.order) - for f in f_patterns(ops.num_f): - assert_distributions_close( - plan_distribution(planned, f), ops_distribution(reordered, f), atol=1e-9 - ) + _checked_plan(ops) def test_plans_are_well_formed(): for ops in component_ops(): - scheduled = schedule(ops) - planned = plan(ops, scheduled) - assert planned.width == scheduled.width - assert planned.num_outputs == ops.num_outputs + planned = plan(ops, schedule(ops)) + assert len(planned.outputs) == ops.num_outputs defined = set(range(planned.num_f)) for action in planned.actions: - assert action.sign.terms <= defined + if isinstance(action, Rotate): + assert action.sign.terms <= defined if isinstance(action, (Rotate, MeasureActive)): assert 0 <= action.x < 2**planned.width assert 0 <= action.z < 2**planned.width @@ -101,14 +99,13 @@ def test_plans_are_well_formed(): assert action.branch not in defined defined.add(action.branch) assert defined == set(range(planned.num_symbols)) + assert all(output.terms <= defined for output in planned.outputs) def _ops(entries): """Return ops from (Pauli, angle or None for a measurement, f-sign, slot).""" paulis = [stim.PauliString(pauli).to_numpy() for pauli, _, _, _ in entries] return HeisenbergOps( - num_lines=len(entries[0][0]), - num_f=1, kind=np.array( [MEASUREMENT if angle is None else ROTATION for _, angle, _, _ in entries], dtype=np.uint8, @@ -131,14 +128,7 @@ def recording(self, gate, *targets): apply(self, gate, *targets) monkeypatch.setattr(_Planner, "_apply", recording) - scheduled = schedule(ops) - planned = plan(ops, scheduled) - reordered = ops.take(scheduled.order) - for f in f_patterns(ops.num_f): - assert_distributions_close( - plan_distribution(planned, f), ops_distribution(reordered, f), atol=1e-9 - ) - return planned, applied + return _checked_plan(ops), applied def test_a_dormant_measurement_of_y_turns_its_pivot_with_s_dag(monkeypatch): @@ -188,7 +178,7 @@ def test_a_diagonal_measurement_on_several_active_axes(monkeypatch): assert any( isinstance(action, MeasureActive) and action.x == 0 - and bin(action.z).count("1") == 2 + and action.z.bit_count() == 2 for action in planned.actions ) assert ("CX", (1, 0), frozenset({0, 1})) in applied diff --git a/test/unit/statevector/test_program.py b/test/unit/statevector/test_program.py index c8386378..50d8e04d 100644 --- a/test/unit/statevector/test_program.py +++ b/test/unit/statevector/test_program.py @@ -2,12 +2,12 @@ CIRCUITS, assert_distributions_close, f_patterns, + output_indices, plan_distribution, quantum_components, stabilizer_rank_distribution, ) -import jax import numpy as np import pytest @@ -17,26 +17,15 @@ from tsim.core.graph import prepare_graph from tsim.core.types import CompiledComponent from tsim.statevector.heisenberg import ComponentLayout -from tsim.statevector.plan import ( - AffineRow, - MeasureActive, - MeasureDormant, - Plan, - Record, - Rotate, -) +from tsim.statevector.plan import AffineRow, MeasureActive, MeasureDormant, Plan, Rotate from tsim.statevector.program import ( - MAX_WIDTH, MEASURE, MEASURE_DIAG, MEASURE_DORMANT, - RECORD, ROTATE, ROTATE_DIAG, - compile_statevector_component, lower, plan_component, - resolve_dtype, ) @@ -49,8 +38,9 @@ def test_plans_match_stabilizer_rank(name, sample_detectors): for component in prepared.graph.connected_components(): if classify_direct(component, prepared.layout) is not None: continue - planned, component_layout = plan_component(component, prepared.layout) - reference = by_outputs[component_layout.output_indices] + indices = output_indices(component) + planned, component_layout = plan_component(component, indices, prepared.layout) + reference = by_outputs[tuple(indices)] assert isinstance(reference, CompiledComponent) selection = tuple(np.asarray(reference.f_selection).tolist()) assert selection == component_layout.f_selection @@ -66,19 +56,16 @@ def test_lower_encodes_every_action(): planned = Plan( width=2, num_f=1, - num_branches=3, - num_outputs=2, actions=( Rotate(x=0, z=3, angle=0.25, sign=AffineRow(True, frozenset({0}))), Rotate(x=1, z=2, angle=-0.5, sign=AffineRow(False, frozenset())), - MeasureActive( - 0, 2, 1, branch=1, sign=AffineRow(False, frozenset({0})), slot=1 - ), - MeasureActive( - 1, 0, 0, branch=2, sign=AffineRow(True, frozenset({1})), slot=-1 - ), - MeasureDormant(branch=3, sign=AffineRow(False, frozenset({1, 2})), slot=0), - Record(sign=AffineRow(True, frozenset({3})), slot=1), + MeasureActive(0, 2, 1, branch=1), + MeasureActive(1, 0, 0, branch=2), + MeasureDormant(branch=3), + ), + outputs=( + AffineRow(False, frozenset({1, 3})), + AffineRow(True, frozenset({0, 2})), ), ) layout = ComponentLayout(f_selection=(4,), f_position={}, output_indices=(2, 5)) @@ -89,53 +76,41 @@ def test_lower_encodes_every_action(): MEASURE_DIAG, MEASURE, MEASURE_DORMANT, - RECORD, ] - assert compiled.x_mask.tolist() == [0, 1, 0, 1, 0, 0] - assert compiled.z_mask.tolist() == [3, 2, 2, 0, 0, 0] + assert compiled.x_mask.tolist() == [0, 1, 0, 1, 0] + assert compiled.z_mask.tolist() == [3, 2, 2, 0, 0] np.testing.assert_allclose(compiled.cos[:2], np.cos([0.25, -0.5]), rtol=1e-6) np.testing.assert_allclose(compiled.sin[:2], np.sin([0.25, -0.5]), rtol=1e-6) assert compiled.pivot.tolist()[2:4] == [1, 0] - assert compiled.branch.tolist() == [-1, -1, 1, 2, 3, -1] - assert compiled.slot.tolist() == [-1, -1, 1, -1, 0, 1] - assert compiled.sign_const.tolist() == [True, False, False, True, False, True] - assert compiled.sign_row[:, 0].tolist() == [1, 0, 1, 2, 6, 8] + assert compiled.branch.tolist() == [-1, -1, 1, 2, 3] + assert compiled.sign_const.tolist()[:2] == [True, False] + assert compiled.sign_row[:2, 0].tolist() == [1, 0] + assert compiled.output_const.tolist() == [False, True] + assert compiled.output_row[:, 0].tolist() == [0b1010, 0b0101] assert compiled.num_words == 1 assert compiled.f_selection.tolist() == [4] assert compiled.output_indices == (2, 5) assert (compiled.width, compiled.num_outputs, compiled.dtype) == (2, 2, "complex64") - assert compiled.num_actions == 6 + assert compiled.num_actions == 5 def test_lower_packs_symbols_beyond_32_into_more_words(): planned = Plan( width=0, num_f=40, - num_branches=0, - num_outputs=1, - actions=(Record(sign=AffineRow(False, frozenset({0, 33, 39})), slot=0),), + actions=(), + outputs=(AffineRow(False, frozenset({0, 33, 39})),), ) compiled = lower(planned, ComponentLayout(tuple(range(40)), {}, (0,)), "complex64") assert compiled.num_words == 2 - assert compiled.sign_row.tolist() == [[1, (1 << 1) | (1 << 7)]] - - -def test_compile_statevector_component(): - (component,), layout = quantum_components(CIRCUITS["t_ghz"], sample_detectors=True) - compiled = compile_statevector_component(component, layout) - planned, component_layout = plan_component(component, layout) - assert compiled.width == planned.width - assert compiled.num_actions == len(planned.actions) - assert compiled.output_indices == component_layout.output_indices - assert compiled.dtype == "complex64" + assert compiled.output_row.tolist() == [[1, (1 << 1) | (1 << 7)]] def test_too_wide_components_are_rejected(monkeypatch): - assert MAX_WIDTH == 30 monkeypatch.setattr(program, "MAX_WIDTH", 0) (component,), layout = quantum_components(CIRCUITS["t_ghz"], sample_detectors=True) with pytest.raises(ValueError, match="cat5"): - plan_component(component, layout) + plan_component(component, output_indices(component), layout) def test_extraction_failures_are_value_errors(monkeypatch): @@ -145,16 +120,4 @@ def fail(_): monkeypatch.setattr(program, "extract_circuit", fail) (component,), layout = quantum_components(CIRCUITS["t_ghz"], sample_detectors=True) with pytest.raises(ValueError, match="no gflow"): - plan_component(component, layout) - - -def test_resolve_dtype(): - assert resolve_dtype(None) == "complex64" - assert resolve_dtype("complex64") == "complex64" - with pytest.raises(ValueError, match="jax_enable_x64"): - resolve_dtype("complex128") - with pytest.raises(ValueError, match="complex32"): - resolve_dtype("complex32") - with jax.enable_x64(True): - assert resolve_dtype(None) == "complex128" - assert resolve_dtype("complex128") == "complex128" + plan_component(component, output_indices(component), layout) diff --git a/test/unit/statevector/test_schedule.py b/test/unit/statevector/test_schedule.py index 22f5044e..dc417e95 100644 --- a/test/unit/statevector/test_schedule.py +++ b/test/unit/statevector/test_schedule.py @@ -18,8 +18,6 @@ def _bits(*rows): def _two_expansions_then_two_measurements(): # Rotate about X0, rotate about X1, measure Z0, measure Z1. return HeisenbergOps( - num_lines=2, - num_f=0, kind=np.array([ROTATION, ROTATION, MEASUREMENT, MEASUREMENT], dtype=np.uint8), x=_bits([1, 0], [0, 1], [0, 0], [0, 0]), z=_bits([0, 0], [0, 0], [1, 0], [0, 1]), @@ -80,9 +78,7 @@ def test_width_is_the_largest_structural_width_of_the_order(): group = StabilizerGroup(ops.num_lines) widths = [0] for j in scheduled.order: - if ops.kind[j] == MEASUREMENT: - group.measure(ops.x[j], ops.z[j]) - else: - group.rotate(ops.x[j], ops.z[j]) + run = group.measure if ops.kind[j] == MEASUREMENT else group.rotate + run(ops.x[j], ops.z[j]) widths.append(group.width) assert scheduled.width == max(widths)