diff --git a/CHANGELOG.md b/CHANGELOG.md index b7bf371..cadca75 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/CLAUDE.md b/CLAUDE.md index ca769a9..e5c6d3f 100644 --- a/CLAUDE.md +++ b/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. diff --git a/benchmarks/bench_bloch.py b/benchmarks/bench_bloch.py new file mode 100644 index 0000000..a59f127 --- /dev/null +++ b/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() diff --git a/docs/guide_paper_figures.md b/docs/guide_paper_figures.md index d0446e6..8b3fc21 100644 --- a/docs/guide_paper_figures.md +++ b/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.`. +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: diff --git a/docs/user_manual.md b/docs/user_manual.md index ef543ec..2131fce 100644 --- a/docs/user_manual.md +++ b/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. diff --git a/examples/_figures.py b/examples/_figures.py index 42bb130..b060a20 100644 --- a/examples/_figures.py +++ b/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, diff --git a/examples/_methyl_water_common.py b/examples/_methyl_water_common.py index 2751e59..10cdf7c 100644 --- a/examples/_methyl_water_common.py +++ b/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: diff --git a/examples/compare_3919_reburp_myprofile.py b/examples/compare_3919_reburp_myprofile.py index 37d2a27..af83820 100644 --- a/examples/compare_3919_reburp_myprofile.py +++ b/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)") diff --git a/examples/grape_broadband_180.py b/examples/grape_broadband_180.py index 1f604b1..0ca7f60 100644 --- a/examples/grape_broadband_180.py +++ b/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, diff --git a/examples/methyl_pc9_sideband_180.py b/examples/methyl_pc9_sideband_180.py index 1904619..d80b727 100644 --- a/examples/methyl_pc9_sideband_180.py +++ b/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", diff --git a/examples/methyl_water_binary_symmetric_180.py b/examples/methyl_water_binary_symmetric_180.py index 348cef2..ea05a34 100644 --- a/examples/methyl_water_binary_symmetric_180.py +++ b/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: diff --git a/examples/methyl_water_reburp_180.py b/examples/methyl_water_reburp_180.py index 2fdd682..f41ca9d 100644 --- a/examples/methyl_water_reburp_180.py +++ b/examples/methyl_water_reburp_180.py @@ -59,7 +59,7 @@ from scipy.optimize import minimize from examples import _methyl_water_common as reburp_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 examples._methyl_water_common import _smoothness as _smoothness from examples._methyl_water_common import write_bruker_shape as write_bruker_shape @@ -349,22 +349,22 @@ def plot_diagnostics( phase_axis.set_yticks([0.0, 180.0]) axes[0].set_title(f"REBURP-style 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)") @@ -427,9 +427,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: diff --git a/examples/methyl_water_reburp_minlength_180.py b/examples/methyl_water_reburp_minlength_180.py index f45ddab..7112fa7 100644 --- a/examples/methyl_water_reburp_minlength_180.py +++ b/examples/methyl_water_reburp_minlength_180.py @@ -80,7 +80,7 @@ from scipy.optimize import minimize as minimize from examples import _methyl_water_common as minlength_helpers -from examples._figures import save_figure +from examples._figures import quantise, save_figure from examples._methyl_water_common import _smoothness as _smoothness from examples._methyl_water_common import write_bruker_shape as write_bruker_shape @@ -379,22 +379,22 @@ def plot_diagnostics( f"{DURATION_US:.0f} us, {RF_MAX_HZ / 1000.0:.1f} 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)") diff --git a/examples/methyl_water_reburp_minpower_180.py b/examples/methyl_water_reburp_minpower_180.py index 219abb0..ed645cf 100644 --- a/examples/methyl_water_reburp_minpower_180.py +++ b/examples/methyl_water_reburp_minpower_180.py @@ -74,7 +74,7 @@ from scipy.optimize import minimize as minimize from examples import _methyl_water_common as minpower_helpers -from examples._figures import save_figure +from examples._figures import quantise, save_figure from examples._methyl_water_common import _smoothness as _smoothness from examples._methyl_water_common import write_bruker_shape as write_bruker_shape @@ -372,22 +372,22 @@ def plot_diagnostics( f"Min-power REBURP-style methyl 180: {DURATION_US:.0f} us, {RF_MAX_HZ / 1000.0:.1f} 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)") diff --git a/examples/output/reburp_pulse_wide_components.png b/examples/output/reburp_pulse_wide_components.png index b33afe0..104a902 100644 Binary files a/examples/output/reburp_pulse_wide_components.png and b/examples/output/reburp_pulse_wide_components.png differ diff --git a/examples/phase_only_iy_inversion_10khz.py b/examples/phase_only_iy_inversion_10khz.py index 742b169..51457e8 100644 --- a/examples/phase_only_iy_inversion_10khz.py +++ b/examples/phase_only_iy_inversion_10khz.py @@ -28,7 +28,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 examples._phase_waveforms import phase_to_xy from optimalcontrol.bloch import propagate_bloch_ensemble from optimalcontrol.grape import ControlProblem, grape_gradient, grape_xy, phase_only_gradient @@ -367,7 +367,7 @@ def _plot_figure( ax_phase.set_ylim(0.0, 360.0) ax_amp.set_ylim(0.0, 100.0) - ax_profile.plot(offsets_khz, y_values, color="black", linewidth=1.6) + ax_profile.plot(offsets_khz, quantise(y_values), color="black", linewidth=1.6) ax_profile.axhline(PROFILE_THRESHOLD, color="gray", linestyle="--", linewidth=1.0) for edge in (-PROFILE_HALF_WIDTH_HZ, PROFILE_HALF_WIDTH_HZ): ax_profile.axvline(edge / 1.0e3, color="tab:blue", linestyle=":", linewidth=1.0) diff --git a/examples/reburp_pulse.py b/examples/reburp_pulse.py index 4bb8b0c..62d6f2e 100644 --- a/examples/reburp_pulse.py +++ b/examples/reburp_pulse.py @@ -30,7 +30,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 optimalcontrol.bloch import propagate_bloch_ensemble RealArray = npt.NDArray[np.float64] @@ -174,10 +174,18 @@ def time_evolution_profiles( for name, initial in initials.items(): trajectory = np.empty((n_steps + 1, offset_t.size, 3), dtype=np.float64) trajectory[0] = initial - for step in range(1, n_steps + 1): - trajectory[step] = propagate_bloch_ensemble( - initial, waveform_xy[:step], offset_t, scales, rf_time_product, dt - )[0] + for col in range(offset_t.size): + state = initial + for step in range(n_steps): + state = propagate_bloch_ensemble( + state, + waveform_xy[step : step + 1], + offset_t[col : col + 1], + scales, + rf_time_product, + dt, + )[0, 0] + trajectory[step + 1, col] = state trajectories[name] = trajectory return trajectories @@ -206,7 +214,7 @@ def plot_reburp( f"Generic REBURP 180 (B1_peak * T = {rf_time_product:.3f} for on-resonance flip)" ) - axes[1].plot(offset_t, profiles["iz_to_mz"], color="tab:green") + axes[1].plot(offset_t, quantise(profiles["iz_to_mz"]), color="tab:green") axes[1].axhline(-1.0, color="black", linestyle=":") axes[1].set_xlim(-PROFILE_HALF_WIDTH, PROFILE_HALF_WIDTH) axes[1].set_ylim(-1.05, 1.05) @@ -214,9 +222,9 @@ def plot_reburp( axes[1].set_xlabel("offset * T") axes[1].set_title("Band-selective inversion profile") - axes[2].plot(offset_t, profiles["x"], label="Ix -> Ix") - axes[2].plot(offset_t, profiles["y"], label="Iy -> -Iy") - axes[2].plot(offset_t, profiles["z"], label="Iz -> -Iz") + axes[2].plot(offset_t, quantise(profiles["x"]), label="Ix -> Ix") + axes[2].plot(offset_t, quantise(profiles["y"]), label="Iy -> -Iy") + axes[2].plot(offset_t, quantise(profiles["z"]), label="Iz -> -Iz") axes[2].axhline(1.0, color="black", linestyle=":") axes[2].set_xlim(-8.0, 8.0) axes[2].set_ylim(-1.05, 1.05) @@ -233,9 +241,9 @@ def plot_wide_profile(profiles: dict[str, RealArray], output_dir: Path) -> Path: """Write Ix, -Iy, Iz transfer over the wide ``offset * T`` window.""" offset_t = profiles["offset_t"] figure, axis = plt.subplots(1, 1, figsize=(9.0, 4.5), constrained_layout=True) - axis.plot(offset_t, profiles["x"], label="Ix -> Ix") - axis.plot(offset_t, profiles["y"], label="Iy -> -Iy") - axis.plot(offset_t, profiles["z"], label="Iz -> -Iz") + axis.plot(offset_t, quantise(profiles["x"]), label="Ix -> Ix") + axis.plot(offset_t, quantise(profiles["y"]), label="Iy -> -Iy") + axis.plot(offset_t, quantise(profiles["z"]), label="Iz -> -Iz") axis.axhline(1.0, color="black", linestyle=":") axis.set_xlim(-WIDE_HALF_WIDTH, WIDE_HALF_WIDTH) axis.set_ylim(-1.05, 1.05) @@ -255,9 +263,9 @@ def plot_wide_components(component_profile: dict[str, RealArray], output_dir: Pa 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(offset_t, final[:, 0], label="Mx", color="tab:blue") - axis.plot(offset_t, final[:, 1], label="My", color="tab:orange") - axis.plot(offset_t, final[:, 2], label="Mz", color="tab:green") + axis.plot(offset_t, quantise(final[:, 0]), label="Mx", color="tab:blue") + axis.plot(offset_t, quantise(final[:, 1]), label="My", color="tab:orange") + axis.plot(offset_t, quantise(final[:, 2]), label="Mz", color="tab:green") axis.axhline(0.0, color="black", linewidth=0.5) axis.set_xlim(-WIDE_HALF_WIDTH, WIDE_HALF_WIDTH) axis.set_ylim(-1.05, 1.05) @@ -291,9 +299,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(fraction, trajectory[:, col, 0], label="Mx", color="tab:blue") - axis.plot(fraction, trajectory[:, col, 1], label="My", color="tab:orange") - axis.plot(fraction, trajectory[:, col, 2], label="Mz", color="tab:green") + axis.plot(fraction, quantise(trajectory[:, col, 0]), label="Mx", color="tab:blue") + axis.plot(fraction, quantise(trajectory[:, col, 1]), label="My", color="tab:orange") + axis.plot(fraction, 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: diff --git a/examples/sciadv2023_fig1_ur180.py b/examples/sciadv2023_fig1_ur180.py index 9948609..eb3ea77 100644 --- a/examples/sciadv2023_fig1_ur180.py +++ b/examples/sciadv2023_fig1_ur180.py @@ -29,7 +29,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 examples._phase_waveforms import phase_to_xy from optimalcontrol.bloch import propagate_bloch_ensemble from optimalcontrol.grape import ControlProblem, grape_gradient, grape_xy, phase_only_gradient @@ -370,9 +370,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", @@ -382,12 +383,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, diff --git a/optimalcontrol/bloch.py b/optimalcontrol/bloch.py index bb2e9ab..d73c689 100644 --- a/optimalcontrol/bloch.py +++ b/optimalcontrol/bloch.py @@ -3,7 +3,6 @@ import math import numpy as np -from scipy.spatial.transform import Rotation from optimalcontrol._accelerator import _enabled, _rust from optimalcontrol._types import RealArray @@ -11,8 +10,22 @@ def _rotate(state: RealArray, field: RealArray, dt: float) -> RealArray: """Apply one Rodrigues rotation in the NumPy fallback path.""" - rotated = Rotation.from_rotvec(2.0 * math.pi * dt * field).apply(state) - return np.asarray(rotated, dtype=np.float64) + m = np.max(np.abs(field), axis=-1, keepdims=True) + safe_m = np.where(m == 0.0, 1.0, m) + q = field / safe_m + q0, q1, q2 = q[..., 0], q[..., 1], q[..., 2] + norm = m * np.sqrt((q0 * q0 + q1 * q1) + q2 * q2)[..., None] + safe_norm = np.where(norm == 0.0, 1.0, norm) + axis = field / safe_norm + angle = ((2.0 * math.pi) * norm) * dt + sn = np.sin(angle) + cs = np.cos(angle) + a0, a1, a2 = axis[..., 0], axis[..., 1], axis[..., 2] + s0, s1, s2 = state[..., 0], state[..., 1], state[..., 2] + dot = ((a0 * s0 + a1 * s1) + a2 * s2)[..., None] + cross = np.stack((a1 * s2 - a2 * s1, a2 * s0 - a0 * s2, a0 * s1 - a1 * s0), axis=-1) + rotated = (cs * state + sn * cross) + (((1.0 - cs) * axis) * dot) + return np.where(norm == 0.0, state, rotated) def propagate_bloch_ensemble( @@ -55,12 +68,12 @@ def propagate_bloch_ensemble( dtype=np.float64, ) - result = np.empty((scales.size, offsets.size, 3), dtype=np.float64) - for row, scale in enumerate(scales): - for col, offset in enumerate(offsets): - state = initial_array.copy() - for ux, uy in waveform: - field = np.array([scale * rf_hz * ux, scale * rf_hz * uy, offset], dtype=np.float64) - state = _rotate(state, field, dt) - result[row, col] = state - return result + state = np.broadcast_to(initial_array, (scales.size, offsets.size, 3)) + field = np.empty_like(state) + field[..., 2] = offsets + amp = scales[:, None] * rf_hz + for ux, uy in waveform: + field[..., 0] = amp * ux + field[..., 1] = amp * uy + state = _rotate(state, field, dt) + return np.ascontiguousarray(state) diff --git a/tests/test_bloch.py b/tests/test_bloch.py index 6c049c5..93fd478 100644 --- a/tests/test_bloch.py +++ b/tests/test_bloch.py @@ -4,7 +4,7 @@ import pytest from optimalcontrol._accelerator import RUST_ACCELERATOR_AVAILABLE -from optimalcontrol.bloch import propagate_bloch_ensemble +from optimalcontrol.bloch import _rotate, propagate_bloch_ensemble @pytest.mark.parametrize("disable_rust", ["0", "1"]) @@ -62,3 +62,77 @@ def test_bloch_ensemble_rejects_non_finite_arrays(monkeypatch) -> None: monkeypatch.setenv("OPTIMALCONTROL_DISABLE_RUST", "1") with pytest.raises(ValueError, match="finite"): propagate_bloch_ensemble(initial, waveform, offsets, scales, 100.0, 1e-5) + + +def test_fallback_matches_native_on_long_ensemble(monkeypatch: pytest.MonkeyPatch) -> None: + assert RUST_ACCELERATOR_AVAILABLE + initial = np.array([0.0, 0.0, 1.0]) + waveform = np.random.default_rng(0).uniform(-1.0, 1.0, (512, 2)) + offsets = np.linspace(-10.0, 10.0, 801) + scales = np.array([0.5, 1.0, 1.5]) + monkeypatch.setenv("OPTIMALCONTROL_DISABLE_RUST", "0") + native = propagate_bloch_ensemble(initial, waveform, offsets, scales, 6.0, 1.0 / 512) + + monkeypatch.setenv("OPTIMALCONTROL_DISABLE_RUST", "1") + fallback = propagate_bloch_ensemble(initial, waveform, offsets, scales, 6.0, 1.0 / 512) + + np.testing.assert_allclose(native, fallback, rtol=1e-13, atol=1e-13) + + +def test_fallback_single_step_matches_native_closely(monkeypatch: pytest.MonkeyPatch) -> None: + assert RUST_ACCELERATOR_AVAILABLE + rng = np.random.default_rng(1) + waveform = rng.uniform(-1.0, 1.0, (1, 2)) + offsets = rng.uniform(-5000.0, 5000.0, 2000) + initial = np.array([0.3, -0.4, np.sqrt(0.75)]) + scales = np.array([1.0]) + monkeypatch.setenv("OPTIMALCONTROL_DISABLE_RUST", "0") + native = propagate_bloch_ensemble(initial, waveform, offsets, scales, 5000.0, 7.5e-6) + + monkeypatch.setenv("OPTIMALCONTROL_DISABLE_RUST", "1") + fallback = propagate_bloch_ensemble(initial, waveform, offsets, scales, 5000.0, 7.5e-6) + + maximum = np.max(np.abs(native - fallback)) + assert maximum <= 2.0e-16, f"maximum absolute difference: {maximum:.17g}" + + +@pytest.mark.parametrize("disable_rust", ["0", "1"]) +def test_rotate_zero_field_leaves_state_unchanged( + monkeypatch: pytest.MonkeyPatch, disable_rust: str +) -> None: + initial = np.array([0.3, -0.4, np.sqrt(0.75)]) + monkeypatch.setenv("OPTIMALCONTROL_DISABLE_RUST", disable_rust) + + result = propagate_bloch_ensemble( + initial, np.ones((3, 2)), np.array([0.0]), np.array([1.0]), 0.0, 7.5e-6 + ) + + np.testing.assert_array_equal(result[0, 0], initial) + + +def test_rotate_broadcasts_over_leading_axes() -> None: + rng = np.random.default_rng(2) + state = rng.uniform(-1.0, 1.0, (4, 5, 3)) + field = rng.uniform(-5000.0, 5000.0, (4, 5, 3)) + + result = _rotate(state, field, 7.5e-6) + + expected = np.empty_like(state) + for row in range(4): + for col in range(5): + expected[row, col] = _rotate(state[row, col], field[row, col], 7.5e-6) + np.testing.assert_array_equal(result, expected) + np.testing.assert_allclose( + np.linalg.norm(result, axis=-1), np.linalg.norm(state, axis=-1), rtol=1e-13, atol=1e-13 + ) + + +def test_rotate_mixed_zero_fields_raise_nothing() -> None: + state = np.array([[0.3, -0.4, 0.5], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]]) + field = np.array([[0.0, 0.0, 0.0], [1200.0, -300.0, 700.0], [0.0, 0.0, 0.0]]) + + with np.errstate(all="raise"): + result = _rotate(state, field, 7.5e-6) + + np.testing.assert_array_equal(result[[0, 2]], state[[0, 2]]) + np.testing.assert_array_equal(result[1], _rotate(state[1], field[1], 7.5e-6)) diff --git a/tests/test_bloch_fallback_shape.py b/tests/test_bloch_fallback_shape.py new file mode 100644 index 0000000..6fde3e7 --- /dev/null +++ b/tests/test_bloch_fallback_shape.py @@ -0,0 +1,32 @@ +"""The NumPy Bloch fallback performs exactly one rotation call per waveform step.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from optimalcontrol import bloch +from optimalcontrol.bloch import propagate_bloch_ensemble + + +def test_fallback_calls_rotate_once_per_step(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("OPTIMALCONTROL_DISABLE_RUST", "1") + rng = np.random.default_rng(3) + waveform = rng.uniform(-1.0, 1.0, (17, 2)) + offsets = np.linspace(-3000.0, 3000.0, 23) + scales = np.array([0.8, 1.0, 1.2]) + initial = np.array([0.0, 0.0, 1.0]) + expected = propagate_bloch_ensemble(initial, waveform, offsets, scales, 5000.0, 1e-5) + + original = bloch._rotate + calls = 0 + + def counted(*args: object) -> object: + nonlocal calls + calls += 1 + return original(*args) + + monkeypatch.setattr(bloch, "_rotate", counted) + actual = propagate_bloch_ensemble(initial, waveform, offsets, scales, 5000.0, 1e-5) + assert calls == 17 + assert np.array_equal(actual, expected) diff --git a/tests/test_figures.py b/tests/test_figures.py new file mode 100644 index 0000000..4a6386e --- /dev/null +++ b/tests/test_figures.py @@ -0,0 +1,29 @@ +"""Tests for the example figure helpers.""" + +import matplotlib + +matplotlib.use("Agg") + +import matplotlib.pyplot as plt +import numpy as np +import numpy.testing as npt + +from examples._figures import quantise, save_figure + + +def test_quantise_erases_sub_12_decimal_noise_but_keeps_1e6_scale() -> None: + base = np.array([0.1234567890123456, -0.9999999999999]) + noisy = base + 1e-14 + npt.assert_array_equal(quantise(base), quantise(noisy)) + + scaled = np.array([1.0e-6]) + npt.assert_array_equal(quantise(scaled), scaled) + + +def test_save_figure_writes_png_without_matplotlib_version(tmp_path) -> None: + fig, ax = plt.subplots() + ax.plot([0.0, 1.0], [0.0, 1.0]) + out = save_figure(fig, tmp_path / "tiny.png") + data = out.read_bytes() + assert data.startswith(b"\x89PNG\r\n\x1a\n") + assert b"Matplotlib version" not in data diff --git a/tests/test_reburp_pulse.py b/tests/test_reburp_pulse.py new file mode 100644 index 0000000..ab69386 --- /dev/null +++ b/tests/test_reburp_pulse.py @@ -0,0 +1,37 @@ +import matplotlib + +matplotlib.use("Agg") + +import numpy as np +import numpy.testing as npt +import pytest + +from examples.reburp_pulse import TIME_EVOLUTION_OFFSETS, reburp_waveform, time_evolution_profiles +from optimalcontrol.bloch import propagate_bloch_ensemble + + +@pytest.mark.parametrize("disable_rust", ["0", "1"]) +def test_time_evolution_profiles_match_stepwise_propagation(monkeypatch, disable_rust): + monkeypatch.setenv("OPTIMALCONTROL_DISABLE_RUST", disable_rust) + n_steps = 16 + rf_time_product = 6.0 + signed = reburp_waveform(n_steps) + trajectories = time_evolution_profiles(signed, rf_time_product) + + waveform_xy = np.column_stack((signed, np.zeros_like(signed))) + offset_t = np.asarray(TIME_EVOLUTION_OFFSETS, dtype=np.float64) + scales = np.array([1.0]) + initials = { + "Ix": np.array([1.0, 0.0, 0.0]), + "Iy": np.array([0.0, 1.0, 0.0]), + "Iz": np.array([0.0, 0.0, 1.0]), + } + for name, initial in initials.items(): + trajectory = trajectories[name] + assert trajectory.shape == (n_steps + 1, offset_t.size, 3) + npt.assert_array_equal(trajectory[0], np.broadcast_to(initial, trajectory[0].shape)) + for k in range(1, n_steps + 1): + expected = propagate_bloch_ensemble( + initial, waveform_xy[:k], offset_t, scales, rf_time_product, 1.0 / n_steps + )[0] + assert np.array_equal(trajectory[k], expected), (name, k)