diff --git a/scripts/statevector/benchmark.py b/scripts/statevector/benchmark.py new file mode 100644 index 00000000..aa294bbd --- /dev/null +++ b/scripts/statevector/benchmark.py @@ -0,0 +1,86 @@ +"""Benchmark strategy="statevector" on the clifft-bench circuits. + +Run from the repository root: + + uv run python scripts/statevector/benchmark.py [--seconds S] [--seed N] [NAME ...] + +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 time + +import numpy as np +from workloads import CIRCUITS, WORKLOADS, acceptance_check + +import tsim + + +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) + parser.add_argument("--seed", type=int, default=0) + args = parser.parse_args() + print( + "| circuit | compile s | warmup s | shots/s | cores busy | 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=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 = time.process_time() + 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 + cores = (time.process_time() - cpu_start) / elapsed + acceptance = accepted / attempted + check = acceptance_check(acceptance, reference, attempted) + print( + f"| {name} | {compile_seconds:.2f} | {warmup_seconds:.2f} | " + f"{attempted / elapsed:,.0f} | {cores:.1f} | {acceptance:.4f} | " + f"{reference:.4f} | " + f"{check} | {errors} / {accepted} | `{sampler!r}` |", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/statevector/benchmark_clifft.py b/scripts/statevector/benchmark_clifft.py new file mode 100644 index 00000000..fd874f02 --- /dev/null +++ b/scripts/statevector/benchmark_clifft.py @@ -0,0 +1,95 @@ +"""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 time + +import clifft +from workloads import CIRCUITS, WORKLOADS, acceptance_check + +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 = time.process_time() + 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 = (time.process_time() - 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 + check = acceptance_check(acceptance, reference, attempted) + 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 new file mode 100644 index 00000000..b83b6916 --- /dev/null +++ b/scripts/statevector/results.md @@ -0,0 +1,200 @@ +# strategy="statevector" on the clifft-bench circuits + +- **Date:** 2026-10-10 +- **Machine:** Apple M4 Pro, 14 cores (10 performance, 4 efficiency) +- **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 + 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 + +### 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 + +| 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 | + +Clifft 0.10.0's figures are from clifft-bench release-v1-20260909-130000, on a +different machine. + +## Logical-error rate + +Tsim's counts pool the benchmark's run (seed 0) with independent runs: + + 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 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.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 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 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..ef470271 --- /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 +``workloads/circuits`` of a clifft-bench checkout next to this repository). +""" + +import math +import os +from pathlib import Path + +_SIBLING_CHECKOUT = Path(__file__).resolve().parents[3] / "clifft-bench" +CIRCUITS = Path( + os.environ.get( + "CLIFFT_BENCH_CIRCUITS", _SIBLING_CHECKOUT / "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 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/__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..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 @@ -23,7 +24,17 @@ 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 ( + STABILIZER_RANK_HINT, + StatevectorComponent, + compile_statevector_component, +) DecompositionMode = Literal["sequential", "joint"] @@ -32,30 +43,44 @@ 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": + 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 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 +91,7 @@ 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) - ) + 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 aa34bcb7..58ada1ae 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, ClassVar, 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) @@ -80,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: @@ -89,7 +97,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 +111,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..c1e29f58 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 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: @@ -24,6 +26,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, @@ -92,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. @@ -108,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) @@ -152,7 +161,7 @@ def sample_program( "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 > component.norm_warning_threshold: warnings.warn( "A marginal probability was not normalized correctly " f"(normalization deviated from 1 by {max_norm_deviation:.1e}). " @@ -167,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.""" @@ -176,7 +191,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 +200,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 +251,20 @@ def __init__( : self._num_detectors ].copy() + statevector = [ + c for c in self._program.components if isinstance(c, StatevectorComponent) + ] + 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. @@ -293,8 +323,14 @@ 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): + # 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: G = circuit.num_graphs max_a = circuit.node_phases.phases.shape[1] @@ -306,8 +342,8 @@ 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 _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() @@ -316,9 +352,18 @@ def _estimate_batch_size(self) -> int: ) 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. - half_of_available = int(available * 0.5) # conservative estimate - return max(1, half_of_available // self._peak_bytes_per_sample()) + 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( @@ -568,8 +613,16 @@ def __repr__(self) -> str: c_e_terms = [] total_memory_bytes = 0 num_outputs = [] + statevector_widths: list[int] = [] + statevector_actions = 0 for component in self._program.components: + if isinstance(component, StatevectorComponent): + statevector_widths.append(component.width) + statevector_actions += component.num_actions + num_outputs.append(len(component.output_indices)) + total_memory_bytes += _nbytes(component) + continue for circuit in component.compiled_scalar_graphs: num_outputs.append(len(component.output_indices)) c_graphs.append(circuit.num_graphs) @@ -588,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: @@ -606,8 +655,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"{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 +689,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 +744,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 +1003,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/src/tsim/statevector/__init__.py b/src/tsim/statevector/__init__.py new file mode 100644 index 00000000..5c66c918 --- /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 +(``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/execute.py b/src/tsim/statevector/execute.py new file mode 100644 index 00000000..490c7756 --- /dev/null +++ b/src/tsim/statevector/execute.py @@ -0,0 +1,294 @@ +"""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 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 + +from functools import partial +from typing import NamedTuple + +import jax +import jax.numpy as jnp +from jax import Array + +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 + + +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)``. + count: Shots each lane stands for, shape ``(L,)``. + max_norm_deviation: Largest ``|‖ψ‖² - 1|`` seen at a measurement. + + """ + + amps: Array + symbols: 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``. + + 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 + 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: + """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), + count=jnp.ones(lanes, dtype=jnp.int32), + max_norm_deviation=jnp.zeros((), dtype=jnp.finfo(dtype).dtype), + ) + + +def _parity(values: Array) -> Array: + return (jax.lax.population_count(values) & 1).astype(jnp.bool_) + + +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: + """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 _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, 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, 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 _squared(values: Array) -> Array: + return jnp.real(values) ** 2 + jnp.imag(values) ** 2 + + +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 _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))) + 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: + amps = state.amps + 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 + outcome, weight = _sample_outcome(state, a, key, weight_0, weight_1) + scale = jax.lax.rsqrt(jnp.maximum(2 * weight, _TINY))[:, None] + projected = jnp.where(outcome[:, None], minus, plus) * scale + return _collapse(component, state, a, outcome, projected, weight_0 + weight_1) + + +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 = _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) + 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] + projected = gathered * scale + return _collapse(component, state, a, outcome, projected, weight_0 + weight_1) + + +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 _set_branch(component, state, a, outcome) + + +_KERNELS = { + ROTATE_DIAG: _rotate_diag, + ROTATE: _rotate, + MEASURE_DIAG: _measure_diag, + MEASURE: _measure, + MEASURE_DORMANT: _measure_dormant, +} + + +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. + + A component without actions returns ``state`` unchanged. + """ + if component.num_actions == 0: + return state + 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) + + return jax.lax.fori_loop(start, stop, step, state) + + +@jax.jit +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) + 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 new file mode 100644 index 00000000..1ab0a834 --- /dev/null +++ b/src/tsim/statevector/heisenberg.py @@ -0,0 +1,252 @@ +"""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, Sequence +from dataclasses import dataclass, fields +from fractions import Fraction +from math import pi + +import numpy as np +import stim +from zxparam import Graph +from zxparam.extract import ExtractedCircuit, Gate + +from tsim.core.params import ParamLayout + +ROTATION = 0 +MEASUREMENT = 1 + +_CLIFFORD_GATES = {"H": "h", "CNOT": "cx", "CZ": "cz"} +_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. + + """ + + f_selection: tuple[int, ...] + f_position: Mapping[int, int] + output_indices: tuple[int, ...] + + @classmethod + 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(p) for p in component.parameters()) + ) + f_position = {layout.basis(i): pos for pos, i in enumerate(f_selection)} + 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)``. + + 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) +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: + 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,)``. + + """ + + 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_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.""" + 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(*(getattr(self, f.name)[order] for f in fields(self))) + + +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, layout) + for gate in circuit.gates: + 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] + 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, layout: ComponentLayout) -> None: + """Start every line in ``|+⟩`` (``H`` on ``|0⟩``).""" + self.num_lines = num_lines + 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.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.""" + 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) -> None: + """Absorb a Clifford gate, or record a rotation, noise or measurement.""" + if gate.kind in _CLIFFORD_GATES: + self._clifford(_CLIFFORD_GATES[gate.kind], *gate.lines) + elif gate.kind == "Z": + 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) -> 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.ops.append((ROTATION, x, z, negative, pi * float(remainder) / 2, -1)) + if phase.parameters: + x, z, _ = self._mapped(line) + 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.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. + + 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: + 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( + 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/src/tsim/statevector/plan.py b/src/tsim/statevector/plan.py new file mode 100644 index 00000000..46436558 --- /dev/null +++ b/src/tsim/statevector/plan.py @@ -0,0 +1,288 @@ +"""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. Each output is an +``AffineRow`` too: the measurement's sign, XOR its branch bit if it is random. +""" + +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, anticommute +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``. + """ + + x: int + z: int + pivot: int + branch: int + + +@dataclass(frozen=True) +class MeasureDormant: + """Sample a fair bit into symbol ``branch``.""" + + branch: int + + +Action = Rotate | MeasureActive | MeasureDormant + + +@dataclass(frozen=True) +class Plan: + """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``); 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 + 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: + """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.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, + actions=tuple(planner.actions), + outputs=tuple(planner.outputs[slot] for slot in range(ops.num_outputs)), + ) + + +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, width: int) -> None: + """Start with the identity frame and every coordinate dormant.""" + self.ops = ops + self.width = width + self.frame = stim.Tableau(ops.num_lines) + self.active: set[int] = set() + 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``.""" + 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 = 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 _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 _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(AffineRow(flip, frozenset({branch}))) + return branch + + def step(self, j: int) -> None: + """Plan op ``j``.""" + x, z, negative = self._virtual(j) + 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) + 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 + output = self._sign(j, negative) + if dormant_x: + branch = self._measure_dormant(j, dormant_x, negative) + elif on_active: + 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 activate it.""" + 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)) + self.active.add(axis) + + def _measure_dormant( + 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) + _, 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) + 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 + ) -> 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) + 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._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 new file mode 100644 index 00000000..2a48e4d7 --- /dev/null +++ b/src/tsim/statevector/program.py @@ -0,0 +1,205 @@ +"""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, 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 +import numpy as np +from jax import Array +from zxparam import Graph +from zxparam.extract import extract_circuit + +from tsim.core.params import ParamLayout +from tsim.statevector.heisenberg import ComponentLayout, heisenberg_ops +from tsim.statevector.plan import MeasureActive, Plan, Rotate, plan +from tsim.statevector.schedule import schedule + +ROTATE_DIAG = 0 +ROTATE = 1 +MEASURE_DIAG = 2 +MEASURE = 3 +MEASURE_DORMANT = 4 + +MAX_WIDTH = 30 + +# 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): + """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. + dtype: ``"complex64"`` or ``"complex128"``. + f_selection: Global f-bit index of each component f-bit. + 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 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) + 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 + 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, output_indices: Sequence[int], 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, 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}. {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}. " + f"{STABILIZER_RANK_HINT}" + ) + return plan(ops, scheduled), component_layout + + +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, 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) + 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] = action.pivot, action.branch + else: + 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, + 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), + 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, 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. + + """ + 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 new file mode 100644 index 00000000..39185b67 --- /dev/null +++ b/src/tsim/statevector/schedule.py @@ -0,0 +1,170 @@ +"""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, anticommute + + +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 anticommute(self.x, self.z, x, z) + + 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 None + first, rest = anti[0], anti[1:] + self.x[rest] ^= self.x[first] + self.z[rest] ^= self.z[first] + 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)``.""" + first = self._eliminate(x, z) + if first is not None: + 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. + + 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) == len(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.""" + return anticommute(ops.x, ops.z, ops.x, ops.z) + + +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]: + 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/helpers/statevector.py b/test/helpers/statevector.py new file mode 100644 index 00000000..5e7283ea --- /dev/null +++ b/test/helpers/statevector.py @@ -0,0 +1,535 @@ +"""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 Mapping, Sequence +from fractions import Fraction +from functools import cache +from math import cos, pi, sin + +import jax.numpy as jnp +import numpy as np +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 +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, + HeisenbergOps, + heisenberg_ops, +) +from tsim.statevector.plan import MeasureActive, Plan, Rotate + +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 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})) + 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, ONE_LINE_LAYOUT), + (measure, ComponentLayout((0,), {7: 0}, (0, 1))), + ] + + +@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: + 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, ...]]: + """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 + + +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 walk(a: int, psi: np.ndarray, symbols: list[int], weight: float) -> None: + if weight < 1e-14: + return + if a == len(plan.actions): + key = tuple(output.evaluate(symbols) for output in plan.outputs) + result[key] = result.get(key, 0.0) + weight + return + action = plan.actions[a] + if isinstance(action, Rotate): + 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): + 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 + 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], 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)} + + +_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, + ) + + +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..567cea72 --- /dev/null +++ b/test/integration/test_statevector_sampler.py @@ -0,0 +1,214 @@ +"""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 + +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.execute import lane_bytes +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) + + +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 _norm_warnings(monkeypatch, strategy, deviation, *, x64=False): + 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, "sample_component", 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) + + +def test_the_automatic_batch_size_of_a_statevector_sampler_is_capped(monkeypatch): + monkeypatch.setattr( + sampler_module._CompiledSamplerBase, "_memory_budget", lambda self: 2**49 + ) + 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" + ) + 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 + 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): + circuit = Circuit(CIRCUITS["t_ghz"]) + peak = circuit.compile_detector_sampler( + strategy="statevector" + )._peak_bytes_per_sample() + base = sampler_module._CompiledSamplerBase + monkeypatch.setattr(base, "_memory_budget", lambda self: peak - 1) + with pytest.raises(ValueError, match="cat5"): + circuit.compile_detector_sampler(strategy="statevector") + monkeypatch.setattr(base, "_memory_budget", lambda self: peak) + circuit.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/__init__.py b/test/unit/statevector/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/test/unit/statevector/test_execute.py b/test/unit/statevector/test_execute.py new file mode 100644 index 00000000..a7229501 --- /dev/null +++ b/test/unit/statevector/test_execute.py @@ -0,0 +1,326 @@ +from test.helpers.statevector import ( + CIRCUITS, + apply_pauli, + assert_matches_distribution, + f_patterns, + output_indices, + 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, + outputs, + pack_bits, + run, + sample_statevector_component, +) +from tsim.statevector.heisenberg import ComponentLayout +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, +) + + +def _component(actions, *, width, num_f=1, outputs=(), dtype="complex64"): + planned = Plan( + width=width, num_f=num_f, actions=tuple(actions), outputs=tuple(outputs) + ) + layout = ComponentLayout(tuple(range(num_f)), {}, tuple(range(len(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 _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 + 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 + ) + + +_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): + psi = _random_states(3, 3) + f_bits = jnp.array([[False], [False], [True]]) + 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(outputs(component, out)[:, 0]).tolist() == [False, True, False] + assert ((np.asarray(out.symbols[:, 0]) >> 1) & 1).tolist() == [0, 1, 1] + + +def test_outputs_are_parities_of_the_symbols(monkeypatch): + component = _component( + [MeasureDormant(branch=1)], + width=0, + 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, 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(): + 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", + [ + "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: + 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(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 f_selection: + f_params[:, list(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) + ) + + +@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.dtype(dtype) + + +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 + + +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 + + +@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 + # 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 + 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(outputs(component, out)[:, 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 + 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] + 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) + 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(): + 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) + 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) + 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_heisenberg.py b/test/unit/statevector/test_heisenberg.py new file mode 100644 index 00000000..55989844 --- /dev/null +++ b/test/unit/statevector/test_heisenberg.py @@ -0,0 +1,88 @@ +from fractions import Fraction +from math import pi +from test.helpers.statevector import ( + CIRCUITS, + ONE_LINE_LAYOUT, + assert_distributions_close, + extracted_circuits, + extracted_distribution, + f_patterns, + one_line, + ops_distribution, + output_indices, + quantum_components, +) + +import pytest +from zxparam import Phase +from zxparam.extract import Gate + +from tsim.statevector.heisenberg import ( + MEASUREMENT, + ROTATION, + ComponentLayout, + heisenberg_ops, +) + + +def test_t_on_plus_is_a_rotation_about_x(): + 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]] + 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)))]), ONE_LINE_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), ONE_LINE_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), 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, 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] + + +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()} + 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 new file mode 100644 index 00000000..627b42bb --- /dev/null +++ b/test/unit/statevector/test_plan.py @@ -0,0 +1,184 @@ +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, +) + +import numpy as np +import pytest +import stim +from zxparam import Phase +from zxparam.extract import Gate + +from tsim.statevector.heisenberg import ( + MEASUREMENT, + ROTATION, + HeisenbergOps, + heisenberg_ops, +) +from tsim.statevector.plan import ( + AffineRow, + MeasureActive, + MeasureDormant, + Rotate, + _Planner, + 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. + 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(): + 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) + assert planned.outputs == (AffineRow(False, frozenset({1})),) + + +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 _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(): + _checked_plan(ops) + + +def test_plans_are_well_formed(): + for ops in component_ops(): + planned = plan(ops, schedule(ops)) + assert len(planned.outputs) == ops.num_outputs + defined = set(range(planned.num_f)) + for action in planned.actions: + 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 + if isinstance(action, (MeasureActive, MeasureDormant)): + 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( + 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) + return _checked_plan(ops), 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 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 new file mode 100644 index 00000000..50d8e04d --- /dev/null +++ b/test/unit/statevector/test_program.py @@ -0,0 +1,123 @@ +from test.helpers.statevector import ( + CIRCUITS, + assert_distributions_close, + f_patterns, + output_indices, + plan_distribution, + quantum_components, + stabilizer_rank_distribution, +) + +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.core.types import CompiledComponent +from tsim.statevector.heisenberg import ComponentLayout +from tsim.statevector.plan import AffineRow, MeasureActive, MeasureDormant, Plan, Rotate +from tsim.statevector.program import ( + MEASURE, + MEASURE_DIAG, + MEASURE_DORMANT, + ROTATE, + ROTATE_DIAG, + lower, + plan_component, +) + + +@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 + 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 + 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, + 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), + 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)) + compiled = lower(planned, layout, "complex64") + assert compiled.kind.tolist() == [ + ROTATE_DIAG, + ROTATE, + MEASURE_DIAG, + MEASURE, + MEASURE_DORMANT, + ] + 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] + 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 == 5 + + +def test_lower_packs_symbols_beyond_32_into_more_words(): + planned = Plan( + width=0, + num_f=40, + actions=(), + outputs=(AffineRow(False, frozenset({0, 33, 39})),), + ) + compiled = lower(planned, ComponentLayout(tuple(range(40)), {}, (0,)), "complex64") + assert compiled.num_words == 2 + assert compiled.output_row.tolist() == [[1, (1 << 1) | (1 << 7)]] + + +def test_too_wide_components_are_rejected(monkeypatch): + 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, output_indices(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, output_indices(component), layout) diff --git a/test/unit/statevector/test_schedule.py b/test/unit/statevector/test_schedule.py new file mode 100644 index 00000000..dc417e95 --- /dev/null +++ b/test/unit/statevector/test_schedule.py @@ -0,0 +1,84 @@ +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( + 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: + 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)