Repository navigation
perf(bloch): vectorised NumPy fallback and backend-independent example figures - #21
Merged
Merged
Conversation
…mula The fallback looped scipy Rotation per (scale, offset, step) at ~7 us each. It now applies one Rodrigues update per step over the whole (n_b1, n_offsets) ensemble using the exact expression order of src/lib.rs rotate_bloch (scaled norm, guarded denominators, left-to-right sums), so a 512x801 ensemble drops from 3.0 s to 20 ms and agrees with the native kernel to 1e-13 after 512 steps and within 2e-16 on single steps. Five tests pin the long-ensemble parity, the single-step bound (captured failing at 3.33e-16 on the scipy path), zero fields, broadcasting over leading axes, and a mixed zero/non-zero batch under np.errstate(all='raise').
Bloch-derived curves differ between the Rust and NumPy backends at ~1e-14, which Agg's antialiaser can turn into single-pixel PNG changes. examples/_figures.quantise rounds a plotted copy to 12 decimals (8 orders below the precision the figures convey); library return values are never rounded.
A counting wrapper around bloch._rotate asserts the fallback issues exactly one batched rotation per waveform step over the whole (n_b1, n_offsets) ensemble and returns the same array as the unpatched run. Mutating the loop back to one call per offset fails the test with 391 calls instead of 17.
benchmarks/bench_bloch.py times propagate_bloch_ensemble on four ensemble sizes and prints CSV rows labelled by the active backend, so the Rust and NumPy paths can be compared with and without OPTIMALCONTROL_DISABLE_RUST=1. On this machine the vectorised fallback runs 512x801x1 in ~20 ms (Rust ~0.9 ms) and 5136x2001x1 in ~0.37 s (Rust ~13 ms).
…s render identical PNGs Thirteen Bloch-derived y-arrays in plot_reburp, plot_wide_profile, plot_wide_components and plot_time_evolution are wrapped with examples._figures.quantise at the matplotlib call. The profiles dict that run() uses for the passband decision and the returned summary is never modified.
Nine Bloch-derived y-arrays in plot_wide_profile, plot_wide_components and plot_time_evolution are wrapped with quantise at the matplotlib call; plot_pareto, signatures and the input dicts are untouched.
Seven Bloch-derived arrays in plot_diagnostics and plot_regime are wrapped with quantise at the matplotlib call (the semilogy arrays keep their np.maximum clamp inside the wrap). Metrics, labels and run() output are unchanged.
…e plotting The two map examples bind plotted = quantise(data) once per panel and hand it to both imshow calls and the contour while the phase-error identity test stays on the original array; the phase-only profile and the two WATERGATE comparison curves are wrapped at their plot calls. _profile_maps, text labels, summaries and cached OPTIMIZED_* arrays are untouched.
The four plot_diagnostics functions wrap the three methyl transfer curves, the artifact semilogy curve and the water curve with quantise at the matplotlib call; the two local plot_time_evolution implementations (binary, reburp) wrap the trajectory components. minlength/minpower keep delegating to the shared helpers quantised in the previous commit. Constants, evaluate_pulse, signed_amplitude and the profiles dicts are untouched.
…th backends reproduce them reburp_pulse_wide_components.png is the only committed figure whose bytes change once the plotted Bloch curves are quantised; it is regenerated natively (matplotlib 3.11.2, Agg, Python 3.12.14, NumPy 2.5.4) and the NumPy fallback now renders the identical file (sha256 add463c3...). All .shape files and the other 51 tracked outputs are unchanged.
…figures The user manual states the measured fallback cost model (Bloch verification 20 ms vs 1 ms native on a 512x801 ensemble; GRAPE ensembles still 15-20x slower), the paper-figure guide scopes PNG byte-stability to the recorded rendering environment, and CHANGELOG Unreleased lists the fallback vectorisation, figure quantisation, benchmark and tests.
…quantisation Architecture item 6 describes the Rust-matching Rodrigues fallback, the examples paragraph documents save_figure and quantise, and the Commands block lists benchmarks/bench_bloch.py. The local AGENTS.md drafts (git-excluded) are updated alongside with a reference check of REFS 20 STALE 0. Plan: .kbsi/plans/fallback-speed-png-noise.md
time_evolution_profiles re-propagated every waveform prefix (1536 propagate_bloch_ensemble calls, 393,984 rotations), which kept the fallback snapshot test at 7.8 s after the kernel vectorisation. It now carries the running state forward one step at a time per offset column, which is bitwise identical on both backends (tests/test_reburp_pulse.py pins equality with the prefix formulation; the committed time-evolution PNG is unchanged) and brings the fallback test to 1.3 s (34.5 s before this branch). Plan: .kbsi/plans/fallback-speed-png-noise.md
… bullet The Unreleased cleanup bullet still described the Bloch fallback as using scipy.spatial.transform.Rotation, which the vectorised Rodrigues fallback replaced (F2 review MINOR). Plan: .kbsi/plans/fallback-speed-png-noise.md
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fallback-path follow-ups tracked after #20: the NumPy Bloch fallback was ~3000x slower than the Rust kernel and a fallback run of one example left a committed PNG modified by a few pixels.
Wave 1 (this push):
perf(bloch):_rotateis now a vectorised Rodrigues update using the exact expression order ofsrc/lib.rs::rotate_bloch(scaled norm, guarded denominators, left-to-right sums);propagate_bloch_ensembleapplies one rotation per step over the whole (n_b1, n_offsets) ensemble. 512x801x1: 3.003 s -> 0.020 s on Apple M4 Pro (Rust: 0.9 ms). Rust == NumPy within 1e-13 after 512 steps and 1.11e-16 on single steps (scipy path was 3.33e-16).test(bloch): five parity/edge tests + a call-counting test pinning one rotation per step (mutation to per-offset calls fails with 391 != 17).bench:benchmarks/bench_bloch.pyprints CSV timings for both backends.feat(examples):examples/_figures.quantiserounds plotted Bloch curves to 12 decimals (library return values are never rounded).Waves 2-3 (to follow on this branch): wrap every Bloch-derived plotted curve with
quantise, regenerate the committed PNGs once natively so both backends leaveexamples/outputclean, docs/CHANGELOG/CLAUDE.md sync.Plan:
.kbsi/plans/fallback-speed-png-noise.md(local, not committed). Evidence under.kbsi/evidence/ulw/.../G001-.../a1/(local).Verification (wave 1, tree 6fce42b)