perf(bloch): vectorised NumPy fallback and backend-independent example figures

CHANGELOG.md

@@ -58,9 +58,17 @@
   (+ deriv) and `distortion_single_pole` are unchanged.
 - Internal cleanup with no behaviour change: shared scalar validation through
   `_validation`, one anti-Hermitian check shared by `grape` and `_accelerator`,
-  one adjoint step in `_seedless_kernel`, the Bloch fallback uses
-  `scipy.spatial.transform.Rotation`, and `ocseed.Band.state_pairs` derives from
-  `bloch_pairs`.
+  one adjoint step in `_seedless_kernel`, and `ocseed.Band.state_pairs` derives
+  from `bloch_pairs`.
+- Vectorise the NumPy Bloch fallback using the Rust kernel's Rodrigues formula.
+  A 512x801 ensemble drops from 3.0 s to 20 ms (more than 50x faster) and agrees
+  with Rust to ~1e-13.
+- Example figures quantise Bloch-derived curves to 12 decimals at the plot call,
+  so both backends write byte-identical PNGs. Committed figures were regenerated
+  with matplotlib 3.11.2 and Agg.
+- Add `benchmarks/bench_bloch.py`.
+- Add `tests/test_bloch_fallback_shape.py` and `tests/test_figures.py`, plus five
+  parity and edge-case tests in `tests/test_bloch.py`.
 
 ## v0.5.0 - 2026-08-15
 

CLAUDE.md

@@ -30,6 +30,7 @@ cargo clippy --locked --all-targets -- -D warnings  # Rust lint gate, warnings d
 .venv/bin/python -c "from optimalcontrol.ocseed import demo; demo()"               # quick smoke: "demo ok (...)"
 .venv/bin/python -m examples.methyl_water_binary_symmetric_180                     # run an example
 .venv/bin/python benchmarks/bench_grape_hotpath.py
+.venv/bin/python benchmarks/bench_bloch.py   # Bloch ensemble timing, both backends (OPTIMALCONTROL_DISABLE_RUST=1 for NumPy)
 ```
 
 - `tests/test_rust_accelerator.py` fails under `OPTIMALCONTROL_DISABLE_RUST=1` by design (it
@@ -85,6 +86,9 @@ cargo clippy --locked --all-targets -- -D warnings  # Rust lint gate, warnings d
    scaled-unitary kernel (Rust or NumPy); `fast=False` routes through the 4×4 Liouville
    GRAPE core as a cross-check. Ensemble members are offset-major (`offset[m // n_b1]`,
    `b1[m % n_b1]`). `bloch.py` is the independent forward model used for verification only.
+   The NumPy fallback of `propagate_bloch_ensemble` is a vectorised Rodrigues update that
+   mirrors `src/lib.rs` `rotate_bloch` (same expression order, agrees to ~1e-13), with one
+   `_rotate` call per step over the whole ensemble.
 7. `io.py` — `Waveform` container, CSV/JSON/Bruker JCAMP export and import. It's a facade
    over `_io_common.py` (the `Waveform` class and path/format helpers), `_io_bruker.py`
    (Bruker shape export) and `_io_jcamp.py` (JCAMP-DX and FAPT import).
@@ -94,8 +98,10 @@ cargo clippy --locked --all-targets -- -D warnings  # Rust lint gate, warnings d
    contract (units: ppm for bands/carrier, Hz for RF, seconds for duration).
 
 Example helpers shared across scripts live in `examples/_methyl_water_common.py` (the
-methyl/water search, refinement and plotting code), `examples/_figures.py` (deterministic
-PNG saving) and `examples/_phase_waveforms.py` (phase to x/y waveform conversion).
+methyl/water search, refinement and plotting code), `examples/_figures.py` (`save_figure`
+writes PNGs with version-free metadata; `quantise` rounds Bloch-derived curves to 12
+decimals at the plot call so both backends render identical PNGs, and library return values
+are never rounded) and `examples/_phase_waveforms.py` (phase to x/y waveform conversion).
 
 The public import paths are frozen by `tests/test_public_surface.py`: every pre-refactor
 import must still resolve, so move code into private modules and re-export it.

benchmarks/bench_bloch.py

@@ -0,0 +1,36 @@
+"""Benchmark the Bloch ensemble forward model across ensemble sizes.
+
+Run from the repository root:
+    .venv/bin/python benchmarks/bench_bloch.py
+    OPTIMALCONTROL_DISABLE_RUST=1 .venv/bin/python benchmarks/bench_bloch.py
+"""
+
+from functools import partial
+from statistics import median
+from timeit import repeat
+
+import numpy as np
+
+from optimalcontrol._accelerator import _enabled
+from optimalcontrol.bloch import propagate_bloch_ensemble
+
+
+def main() -> None:
+    backend = "rust" if _enabled() else "numpy"
+    print("backend,steps,offsets,b1,seconds_per_call")
+    for steps, n_offsets, n_b1 in ((64, 101, 1), (512, 801, 1), (512, 801, 3), (5136, 2001, 1)):
+        waveform = np.random.default_rng(7).uniform(-1.0, 1.0, (steps, 2))
+        offsets = np.linspace(-10.0, 10.0, n_offsets) * 1000.0
+        scales = np.linspace(0.8, 1.2, n_b1)
+        initial = np.array([0.0, 0.0, 1.0])
+
+        call = partial(
+            propagate_bloch_ensemble, initial, waveform, offsets, scales, 6000.0, 1.0 / steps
+        )
+        call()
+        seconds = median(repeat(call, number=1, repeat=3))
+        print(f"{backend},{steps},{n_offsets},{n_b1},{seconds:.8e}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

docs/guide_paper_figures.md

@@ -4,6 +4,11 @@ Example scripts in `examples/` reproduce the numerical curves and pulse shapes
 used by the ROPE and CROP paper guides. Each script has a `run()` function for
 tests and saves a PNG under `examples/output/` when run from the repository root with
 `python -m examples.<name>`.
+The PNGs are byte-stable across the Rust and NumPy backends for the recorded
+rendering environment (matplotlib 3.11.2 with the Agg backend, used to regenerate
+the committed figures), because Bloch-derived curves are quantised to 12 decimals
+(`examples/_figures.quantise`) before plotting; a different matplotlib or FreeType
+version can still change the bytes.
 
 Run one figure script directly:
 

docs/user_manual.md

@@ -710,8 +710,13 @@ OPTIMALCONTROL_DISABLE_RUST=1 python3 my_first_pulse.py
 ```
 
 Both paths give the same answers to well within physical significance, but not
-bit-for-bit — expect agreement in the first several digits, not the last. Rust is
-roughly an order of magnitude faster on realistic problems (the [§4-3](#4-3-two-bands-invert-the-methyls-keep-the-water)
+bit-for-bit — expect agreement in the first several digits, not the last. The Bloch
+fallback uses the same Rodrigues formula as the Rust kernel and agrees to ~1e-13.
+Bloch-vector verification (`propagate_bloch_ensemble`, `SeedlessSpec.evaluate`, the
+MCP profile tool) is vectorised on the NumPy path, so a 512-step x 801-offset
+ensemble takes about 20 ms in fallback versus about 1 ms native. GRAPE ensemble
+optimisation is still roughly 15-20x slower in fallback, because members are
+propagated one by one with NumPy eigendecompositions (the [§4-3](#4-3-two-bands-invert-the-methyls-keep-the-water)
 optimisation takes ~0.8 s native versus ~10 s in fallback on one development
 machine). If a result *only* reproduces on one path, that is a bug worth
 reporting.

examples/_figures.py

@@ -3,10 +3,26 @@
 from __future__ import annotations
 
 from pathlib import Path
+from typing import Final
 
 import matplotlib.pyplot as plt
+import numpy as np
+import numpy.typing as npt
 from matplotlib.figure import Figure
 
+RealArray = npt.NDArray[np.float64]
+
+# 12 decimals sits 8 orders below the precision the figures convey and far above
+# the ~1e-14 Rust-vs-NumPy Bloch difference, which Agg's antialiaser can otherwise
+# turn into single-pixel PNG changes. Only plotted copies are rounded; library
+# return values are never rounded.
+QUANTISE_DECIMALS: Final = 12
+
+
+def quantise(values: RealArray) -> RealArray:
+    """Round Bloch-derived curves to ``QUANTISE_DECIMALS`` just before plotting."""
+    return np.round(np.asarray(values, dtype=np.float64), QUANTISE_DECIMALS)
+
 
 def save_figure(
     fig: Figure,

examples/_methyl_water_common.py

@@ -13,7 +13,7 @@
 import numpy.typing as npt
 from scipy.optimize import minimize
 
-from examples._figures import save_figure
+from examples._figures import quantise, save_figure
 from optimalcontrol.grape import ControlProblem, grape_xy_and_gradient
 from optimalcontrol.io import export_bruker_shape as _write_bruker_shape
 from optimalcontrol.operators import Ix, Iy, Iz, liouvillian_comm, vec
@@ -63,9 +63,9 @@ def plot_wide_profile(
     """Write Ix, -Iy, Iz transfer profiles over the full -6..6 ppm window."""
     ppm = wide_profile["ppm"]
     figure, axis = plt.subplots(1, 1, figsize=(9.0, 4.5), constrained_layout=True)
-    axis.plot(ppm, wide_profile["x"], label="Ix -> Ix")
-    axis.plot(ppm, wide_profile["y"], label="Iy -> -Iy")
-    axis.plot(ppm, wide_profile["z"], label="Iz -> -Iz")
+    axis.plot(ppm, quantise(wide_profile["x"]), label="Ix -> Ix")
+    axis.plot(ppm, quantise(wide_profile["y"]), label="Iy -> -Iy")
+    axis.plot(ppm, quantise(wide_profile["z"]), label="Iz -> -Iz")
     axis.axvline(METHYL_PPM_LO, color="black", linestyle=":", linewidth=0.8)
     axis.axvline(METHYL_PPM_HI, color="black", linestyle=":", linewidth=0.8)
     axis.axvline(WATER_PPM, color="tab:gray", linestyle="--", linewidth=0.8)
@@ -96,9 +96,9 @@ def plot_wide_components(
     figure, axes = plt.subplots(3, 1, figsize=(9.0, 9.0), sharex=True, constrained_layout=True)
     for axis, name in zip(axes, ("Ix", "Iy", "Iz"), strict=True):
         final = component_profile[name]
-        axis.plot(ppm, final[:, 0], label="Mx", color="tab:blue")
-        axis.plot(ppm, final[:, 1], label="My", color="tab:orange")
-        axis.plot(ppm, final[:, 2], label="Mz", color="tab:green")
+        axis.plot(ppm, quantise(final[:, 0]), label="Mx", color="tab:blue")
+        axis.plot(ppm, quantise(final[:, 1]), label="My", color="tab:orange")
+        axis.plot(ppm, quantise(final[:, 2]), label="Mz", color="tab:green")
         axis.axvline(METHYL_PPM_LO, color="black", linestyle=":", linewidth=0.8)
         axis.axvline(METHYL_PPM_HI, color="black", linestyle=":", linewidth=0.8)
         axis.axvline(WATER_PPM, color="tab:gray", linestyle="--", linewidth=0.8)
@@ -310,9 +310,9 @@ def plot_time_evolution(
         trajectory = time_evolution[name]
         for col in range(n_cols):
             axis = axes[row, col]
-            axis.plot(time_us, trajectory[:, col, 0], label="Mx", color="tab:blue")
-            axis.plot(time_us, trajectory[:, col, 1], label="My", color="tab:orange")
-            axis.plot(time_us, trajectory[:, col, 2], label="Mz", color="tab:green")
+            axis.plot(time_us, quantise(trajectory[:, col, 0]), label="Mx", color="tab:blue")
+            axis.plot(time_us, quantise(trajectory[:, col, 1]), label="My", color="tab:orange")
+            axis.plot(time_us, quantise(trajectory[:, col, 2]), label="Mz", color="tab:green")
             axis.axhline(0.0, color="black", linewidth=0.5)
             axis.set_ylim(-1.05, 1.05)
             if row == 0:

examples/compare_3919_reburp_myprofile.py

@@ -40,7 +40,7 @@
 import numpy as np
 import numpy.typing as npt
 
-from examples._figures import save_figure
+from examples._figures import quantise, save_figure
 from optimalcontrol.bloch import propagate_bloch_ensemble
 
 RealArray = npt.NDArray[np.float64]
@@ -145,8 +145,16 @@ def _draw_panel(
         WATER_PPM, color="tab:gray", linestyle="--", linewidth=0.9, label="water (4.7 ppm)"
     )
     axis.axvline(-4.0, color="tab:gray", linestyle=":", linewidth=0.9, label="3-9-19 null (-4 ppm)")
-    axis.plot(ppm, reburp, color="tab:blue", linewidth=1.8, label="min-length REBURP (10 kHz OC)")
-    axis.plot(ppm, w3919, color="tab:red", linewidth=1.4, label="3-9-19 WATERGATE (36 kHz hard)")
+    axis.plot(
+        ppm,
+        quantise(reburp),
+        color="tab:blue",
+        linewidth=1.8,
+        label="min-length REBURP (10 kHz OC)",
+    )
+    axis.plot(
+        ppm, quantise(w3919), color="tab:red", linewidth=1.4, label="3-9-19 WATERGATE (36 kHz hard)"
+    )
     axis.set_xlim(PPM_HI, PPM_LO)  # NMR convention: high ppm on the left
     axis.set_ylim(-1.15, 1.15)
     axis.set_xlabel("1H shift (ppm)")

examples/grape_broadband_180.py

@@ -29,7 +29,7 @@
 import numpy as np
 import numpy.typing as npt
 
-from examples._figures import save_figure
+from examples._figures import quantise, save_figure
 from optimalcontrol.bloch import propagate_bloch_ensemble
 from optimalcontrol.grape import ControlProblem
 from optimalcontrol.io import export_bruker_shape as write_bruker_shape
@@ -299,9 +299,10 @@ def _plot_figure(
         ),
     ]
     for ax, data, title in panels:
+        plotted = quantise(data)
         if data is phase_error:
             image = ax.imshow(
-                data,
+                plotted,
                 origin="lower",
                 extent=extent,
                 aspect="auto",
@@ -311,12 +312,12 @@ def _plot_figure(
             )
             levels = [-2.0, -1.0, 0.0, 1.0, 2.0]
         else:
-            image = ax.imshow(data, **image_kwargs)
+            image = ax.imshow(plotted, **image_kwargs)
             levels = contour_levels
         contours = ax.contour(
             offsets / 1000.0,
             deviations,
-            data,
+            plotted,
             levels=levels,
             colors="black",
             linewidths=0.55,

examples/methyl_pc9_sideband_180.py

@@ -63,7 +63,7 @@
 import numpy.typing as npt
 from scipy.optimize import minimize_scalar
 
-from examples._figures import save_figure
+from examples._figures import quantise, save_figure
 from examples.methyl_water_binary_symmetric_180 import (
     MAX_ARTIFACT_PERCENT,
     METHYL_PPM_HI,
@@ -249,16 +249,18 @@ def plot_diagnostics(
         f"180-degree PC9 (Pc9_4_90): {duration_s * 1e6:.0f} us, {rf_peak_hz / 1000.0:.2f} kHz peak"
     )
 
-    axes[1].plot(ppm, profiles["methyl_x"], label="Ix -> Ix")
-    axes[1].plot(ppm, profiles["methyl_y"], label="-Iy -> Iy")
-    axes[1].plot(ppm, profiles["methyl_z"], label="Iz -> -Iz")
+    axes[1].plot(ppm, quantise(profiles["methyl_x"]), label="Ix -> Ix")
+    axes[1].plot(ppm, quantise(profiles["methyl_y"]), label="-Iy -> Iy")
+    axes[1].plot(ppm, quantise(profiles["methyl_z"]), label="Iz -> -Iz")
     axes[1].axhline(MIN_METHYL_FIDELITY, color="black", linestyle=":")
     axes[1].set_xlim(METHYL_PPM_LO, METHYL_PPM_HI)
     axes[1].set_ylim(-1.05, 1.05)
     axes[1].set_ylabel("Transfer fidelity")
     axes[1].legend(loc="lower center", ncol=3, fontsize="small")
 
-    axes[2].semilogy(ppm, np.maximum(profiles["artifact_percent"], 1e-4), color="tab:green")
+    axes[2].semilogy(
+        ppm, quantise(np.maximum(profiles["artifact_percent"], 1e-4)), color="tab:green"
+    )
     axes[2].axhline(MAX_ARTIFACT_PERCENT, color="black", linestyle=":")
     axes[2].set_xlim(METHYL_PPM_LO, METHYL_PPM_HI)
     axes[2].set_ylabel("Kay sideband (%)")
@@ -267,7 +269,7 @@ def plot_diagnostics(
         f"Worst methyl sideband over -3..3 ppm = {sideband_max:.2f}% (target <= 0.1%)"
     )
 
-    axes[3].plot(profiles["water_offsets_hz"], profiles["water_z"], marker="o")
+    axes[3].plot(profiles["water_offsets_hz"], quantise(profiles["water_z"]), marker="o")
     axes[3].axhline(MIN_WATER_FIDELITY, color="black", linestyle=":")
     axes[3].set_ylim(-1.05, 1.05)
     axes[3].set_xlabel("Offset from carrier (Hz); water at 4.7 ppm")
@@ -284,7 +286,7 @@ def plot_regime(sweep: dict[str, RealArray], bwfac_us: float, output_dir: Path)
     figure, axis = plt.subplots(1, 1, figsize=(8.0, 4.8), constrained_layout=True)
     axis.semilogy(
         duration_us,
-        np.maximum(sweep["sideband_max_percent"], 1e-4),
+        quantise(np.maximum(sweep["sideband_max_percent"], 1e-4)),
         marker="o",
         color="tab:green",
         label="worst methyl sideband (%)",
@@ -303,7 +305,7 @@ def plot_regime(sweep: dict[str, RealArray], bwfac_us: float, output_dir: Path)
     water_axis = axis.twinx()
     water_axis.plot(
         duration_us,
-        sweep["water_z_min"],
+        quantise(sweep["water_z_min"]),
         marker="s",
         color="tab:blue",
         label="worst water Iz -> Iz",

examples/methyl_water_binary_symmetric_180.py

@@ -50,7 +50,7 @@
 from scipy.optimize import minimize
 
 from examples import _methyl_water_common as binary_helpers
-from examples._figures import save_figure
+from examples._figures import quantise, save_figure
 from examples._methyl_water_common import _constraint_score as _constraint_score
 from optimalcontrol.bloch import propagate_bloch_ensemble
 from optimalcontrol.grape import ControlProblem, grape_xy_and_gradient
@@ -611,22 +611,22 @@ def plot_diagnostics(
     phase_axis.set_yticks([0.0, 180.0])
     axes[0].set_title(f"Symmetric binary-phase methyl 180: {DURATION_US:.0f} us, 10 kHz max")
 
-    axes[1].plot(methyl_ppm, profiles["methyl_x"], label="Ix -> Ix")
-    axes[1].plot(methyl_ppm, profiles["methyl_y"], label="-Iy -> Iy")
-    axes[1].plot(methyl_ppm, profiles["methyl_z"], label="Iz -> -Iz")
+    axes[1].plot(methyl_ppm, quantise(profiles["methyl_x"]), label="Ix -> Ix")
+    axes[1].plot(methyl_ppm, quantise(profiles["methyl_y"]), label="-Iy -> Iy")
+    axes[1].plot(methyl_ppm, quantise(profiles["methyl_z"]), label="Iz -> -Iz")
     axes[1].axhline(MIN_METHYL_FIDELITY, color="black", linestyle=":")
     axes[1].set_xlim(METHYL_PPM_LO, METHYL_PPM_HI)
     axes[1].set_ylim(0.9985, 1.0001)
     axes[1].set_ylabel("Transfer fidelity")
     axes[1].legend()
 
-    axes[2].semilogy(methyl_ppm, profiles["artifact_percent"], color="tab:green")
+    axes[2].semilogy(methyl_ppm, quantise(profiles["artifact_percent"]), color="tab:green")
     axes[2].axhline(MAX_ARTIFACT_PERCENT, color="black", linestyle=":")
     axes[2].set_xlim(METHYL_PPM_LO, METHYL_PPM_HI)
     axes[2].set_ylabel("Kay artifact (%)")
     axes[2].set_xlabel("Methyl 1H shift (ppm)")
 
-    axes[3].plot(profiles["water_offsets_hz"], profiles["water_z"], marker="o")
+    axes[3].plot(profiles["water_offsets_hz"], quantise(profiles["water_z"]), marker="o")
     axes[3].axhline(MIN_WATER_FIDELITY, color="black", linestyle=":")
     axes[3].set_ylim(0.9985, 1.0001)
     axes[3].set_xlabel("Offset from water at 4.7 ppm (Hz)")
@@ -693,9 +693,9 @@ def plot_time_evolution(time_evolution: dict[str, RealArray], output_dir: Path)
         trajectory = time_evolution[name]
         for col in range(n_cols):
             axis = axes[row, col]
-            axis.plot(time_us, trajectory[:, col, 0], label="Mx", color="tab:blue")
-            axis.plot(time_us, trajectory[:, col, 1], label="My", color="tab:orange")
-            axis.plot(time_us, trajectory[:, col, 2], label="Mz", color="tab:green")
+            axis.plot(time_us, quantise(trajectory[:, col, 0]), label="Mx", color="tab:blue")
+            axis.plot(time_us, quantise(trajectory[:, col, 1]), label="My", color="tab:orange")
+            axis.plot(time_us, quantise(trajectory[:, col, 2]), label="Mz", color="tab:green")
             axis.axhline(0.0, color="black", linewidth=0.5)
             axis.set_ylim(-1.05, 1.05)
             if row == 0: