feat: 0.3.0 Wave 0 + Wave 1 — synthetic data, requires_package gate, batched GPU fitting - #82
feat: 0.3.0 Wave 0 + Wave 1 — synthetic data, requires_package gate, batched GPU fitting#82CSSFrancis wants to merge 6 commits into
Conversation
Five waves, each independently shippable behind its own optional extra. Tracked in #48 with a sub-issue per task. Decisions recorded up front, because each one shapes the code: - spyde/models/ is already the neural disk detector, so the HyperSpy model port lives in spyde/fitting/. Unrelated meanings of "model". - exspy/kikuchipy/atomap ship as optional extras gated by a new requires_package key, mirroring requires_vectors. - The fitting engine is a batched GPU Levenberg-Marquardt over the whole nav grid (the vector_orientation_gpu.py playbook), with SAMFire's neighbour-seeding kept as the seed source rather than the scheduler. Measured baseline that motivates it: multifit on 1024 spectra x 1024 channels runs 10.83 s = 95 spectra/s single-threaded, which extrapolates to 11.6 min for a 256x256 SI. Parity with multifit is the acceptance test, not convergence. Also found while measuring: numexpr is not installed, so every hyperspy Expression component silently falls back to numpy.
Waves 1-4 need something loadable per modality from day one, before the real datasets land (#80). These are also the better thing to test against: the ground truth is known exactly, so a fit or an indexing run is scored against the numbers the data was built from rather than a golden file that quietly encodes yesterday's bug. eels_si power law + C/N/O K edges at real onsets, intensities following known concentration maps eds_si bremsstrahlung + Fe/Ni/Cu K families at real energies, with Fe-Kb/Ni-Ka overlapping ON PURPOSE so family-aware fitting is actually exercised rather than assumed ebsd_patterns real gnomonic Kikuchi-band geometry for a cubic crystal over a two-grain orientation field; the exact Euler angles are stamped on the signal, so dictionary indexing can be checked against truth instead of against a fixture of its own output Every spatial map is distinguishable from its transpose and both mirrors, so a flipped axis fails a test instead of looking fine — the si_grains/movie precedent. Ground truth is read with spyde.data.ground_truth(), not straight off the metadata: hyperspy turns nested metadata dicts into DictionaryTreeBrowser nodes, which have no .values(). None of this needs exspy or kikuchipy. set_signal_type is attempted and its "not understood" warning demoted to debug, so a user without the extras gets usable data and no scary warning. Reachable as load_test_data_eels / _eds / _ebsd.
Closes the framework half of #49. requires_package: hides a toolbar action until its optional extra is importable, mirroring requires_vectors — so a user without spyde[eels] never sees a button that would raise ImportError on click. Resolved with find_spec, so it never imports the package and never pays its import cost; cached, since an extra cannot appear mid-session. Applied in BOTH filter paths. get_toolbar_actions_for_plot and _action_matches_plot duplicate the gate list, and adding a key to one renders a button that never dispatches. A test asserts both contain it. Extras: eels (exspy), ebsd (kikuchipy), atoms (atomap), all. Deliberately not core deps — each pulls a large tree and the frozen app must stay installable for someone who only does 4D-STEM. A CI job installs spyde[all] so the gated code actually executes; the default job proves the app still works with none of it. NOT adding numexpr, and pyproject now says why. Hyperspy warns "Numexpr is not installed ... which is slower to calculate model" on every model build, but at spectrum sizes that advice is backwards. Measured on a 1024-channel multifit: absent 110 spectra/s numexpr, 1 thr 77 numexpr, 4 thr 75 numexpr, 48 thr 70 numexpr's per-call setup and thread dispatch dominate arrays of ~1k elements. The warning will keep inviting someone to add it, hence the comment.
Ports hyperspy's model concept and makes it fast. #51, #52, #53. Lives in spyde/fitting/, NOT spyde/models/ — that name is already the neural disk detector. Unrelated meanings of "model". spec.py ModelSpec: serialisable, round-trips against hyperspy's BaseModel.as_dictionary(), so a model built here opens in plain hyperspy. Packed parameter order IS the API — column j always means the same parameter. components.py batched autograd-differentiable ports of hyperspy's components; no Python loop over pixels anywhere engine.py batched Levenberg-Marquardt over the whole nav grid Why it works: n (free parameters) is tiny even when P and C are large, so the normal equations are (P, n, n) — a batch of small solves torch.linalg does natively. The one large object is the (P, C, n) Jacobian, so P is chunked. Jacobian by jacfwd not jacrev: forward-mode costs one pass per INPUT (n~10), reverse one per OUTPUT (C~2000). One implementation, not two. The "CPU fallback" is this code with device="cpu"; a separate scipy path would be a second thing to keep in agreement with hyperspy, and hyperspy already IS the reference — the acceptance test asserts parameter parity with multifit directly, because "it converged" proves nothing when you are replacing a reference implementation. Measured, 1024 spectra x 1024 channels, PowerLaw + 3 Erf (11 free params): hyperspy multifit 214.80 s 4.8 spectra/s batched cpu 27.02 s 37.9 spectra/s 7.9x batched cuda 3.97 s 257.8 spectra/s 54.1x Two numerical traps found by measurement, both documented at the code: - Column scaling is load-bearing, not a refinement. A PowerLaw fits A ~ 1e6 alongside r ~ 3, so cond(JtJ) ~ 1e12 and float64 Cholesky loses nearly every digit in the small-eigenvalue direction. Unscaled: 247 iterations for a TWO parameter fit, stopping 1.5-3.5% short of multifit. Scaled: 46 iterations to chisq ~1e-34. - Convergence must be a GRADIENT test. Testing step size alone reads "stuck against a wall with large lambda" as converged; requiring an ACCEPTED step instead never fires at all, because at the optimum nothing improves. Convergence on hard many-parameter models is still seed-limited (hyperspy hits maxfev on the same model) — that is what seeded propagation (#54) is for.
Answers "is PyTorch really the right way to fit?" — the honest answer was
"yes as a vehicle, no as I was using it". Profiling the solver:
residual evaluation 3.4 ms
jacfwd Jacobian 51.6 ms <- 94% of the solver
Forward-mode AD costs one forward pass per parameter, and these components have
closed-form derivatives that mostly reuse the value already computed
(df/dA IS the shape for every linear amplitude). So components now carry
analytic gradients and the engine builds the Jacobian directly; autodiff stays
as the fallback for any component without one.
batched cpu 37.9 -> 105.9 spectra/s (2.8x)
batched cuda 257.8 -> 361.6 spectra/s (1.4x)
The GPU gains less because it is now dispatch-bound, not compute-bound — the
remaining lever there is kernel fusion, not more math.
Also: the chunk heuristic was starving the GPU. Capping the Jacobian at 2**22
elements chopped a 1024-spectrum batch into pieces of 372 and cost a third of
the throughput (178 vs 267 spectra/s). Raised to 2**26 so typical scans run
unchunked.
Correctness is unchanged and still defined by hyperspy: every analytic gradient
is asserted equal to autodiff's, which is the ideal oracle since it is derived
mechanically from the value function and cannot share a mistake with a
hand-written formula.
seeding.py implements #54 — coarse strided fit, propagate to nearest neighbour,
one batched refine. SAMFire's insight (a neighbour's result is the best start)
without SAMFire's per-pixel scheduler. Only CONVERGED coarse fits propagate: a
failed one has wandered somewhere unphysical and seeding from it spreads the
failure rather than the answer.
Two benchmark corrections, both of which had been quietly measuring the wrong
thing:
- With exspy installed, create_model() returns an EELSModel that ALREADY has a
background, so the reference model had TWO PowerLaws while the batched side
had one. Both sides must fit the same model or the comparison means nothing.
- The EELS benchmark model cannot represent a hydrogenic edge, so neither
engine converges and it was timing the iteration cap. Added a well-posed EDS
case (gaussian peaks, which the library expresses exactly) as the default;
EELS stays as the deliberately-hard timing-only case, documented as such.
On the well-posed case: 57.6% converged, 361 spectra/s, 11.6x multifit.
Update: answering "is PyTorch really the right way to fit?"Profiled it rather than argued it. The answer was yes as a vehicle, no as I was using it:
Autodiff costs one forward pass per parameter, and these components have closed-form derivatives that mostly reuse the value already computed (
The GPU gains less because it is now dispatch-bound, not compute-bound — the remaining lever there is kernel fusion ( Also found: my chunk heuristic was starving the GPU. Capping the Jacobian at 2²² elements chopped a 1024-spectrum batch into pieces of 372 and cost a third of the throughput (178 vs 267 spectra/s). Raised to 2²⁶. Correctness is unchanged — every analytic gradient is asserted equal to autodiff's, which is the ideal oracle since it's derived mechanically from the value function and can't share a mistake with a hand-written formula. Also in this push#54 seeded propagation — coarse strided fit → propagate to nearest neighbour → one batched refine. Only converged coarse fits propagate; a failed one has wandered somewhere unphysical and seeding from it spreads the failure. Two benchmark corrections, both of which had been quietly measuring the wrong thing:
Well-posed case: 57.6% converged, 361 spectra/s, 11.6× multifit. 176 tests green. Known, not yet addressedThe engine runs every pixel for |
Gaussian2D through the SAME engine, not a second code path. An image is flattened to H*W sample points exactly as a spectrum is C channels; the only difference is that the "axis" carries (x, y) PAIRS instead of scalars, handled by one branch in TorchComponent._split. Packing, the Jacobian, bounds, fixed parameters and the LM solve are all the 1-D machinery unchanged — and the tests assert that rather than assume it (bounds and fixed-parameter tests are repeated in 2-D precisely because a different shaping path is where they would break). image_coordinates() builds the coordinate grid, row-major so it lines up with a .ravel()'d image. The test image is deliberately NON-SQUARE (24x28): a transposed coordinate grid would run happily on a square one and put every peak in the wrong place. As in 1-D, A is the VOLUME under the surface rather than the peak height — transcribed from hyperspy's expression, with parity, batching and analytic-vs-autodiff gradient tests to keep it honest. This also unblocks Wave 4: refining atom positions IS a batched 2-D gaussian fit (#77).
First code for the 0.3.0 tracker. Three reviewable commits.
Closes #51, #52, #53. Framework half of #49. Unblocks Waves 2-4 with data.
139 new tests, all green. Full suite: 1760 passed, 2 skipped, 1 xfailed — no regressions.
1. Synthetic EELS / EDS / EBSD data (
spyde/data/)Waves 1-4 need something loadable per modality before the real datasets land (#80) — and known ground truth beats a golden file, which quietly encodes yesterday's bug.
eels_si— power law + C/N/O K edges at real onsetseds_si— Fe/Ni/Cu K families at real energies, Fe-Kβ/Ni-Kα overlapping on purpose so family-aware fitting ([2.2] Composition → auto-populated components #62) is exercised, not assumedebsd_patterns— real gnomonic Kikuchi geometry, two-grain orientation field, exact Euler angles stamped on the signal so [3.3] GPU dictionary indexing — the reason this wave exists #71 can be checked against truthEvery spatial map differs from its transpose and both mirrors, so a flipped axis fails a test.
2.
requires_packagegate + domain extras (Wave 0)Hides an action until its extra is importable, mirroring
requires_vectors.find_spec, so it never imports the package. Applied in both filter paths — adding a gate to only one renders a button that never dispatches, and a test asserts both have it.Extras
eels/ebsd/atoms/all, plus a CI job onspyde[all]so gated code actually runs.Reversed on measurement:
numexpris NOT being added. HyperSpy warns it is missing and "slower to calculate model", but at spectrum sizes that is backwards:Its per-call setup dominates ~1k-element arrays.
pyproject.tomlnow carries a comment, because the warning will keep inviting someone to add it.3. Batched GPU fitting engine (Wave 1.1-1.3)
In
spyde/fitting/— notspyde/models/, which is already the neural disk detector.JᵀJis(P, n, n)because n stays tiny; the only big object is the(P, C, n)Jacobian, so P is chunked.jacfwdnotjacrev(one pass per input, n≈10, vs per output, C≈2000).One implementation, not two — the CPU path is the same code with
device="cpu". HyperSpy is the reference, and the acceptance test asserts parameter parity withmultifit, because "it converged" proves nothing when you are replacing a reference implementation.Measured, 1024 spectra × 1024 channels, PowerLaw + 3×Erf (11 free params):
Two numerical traps found by measuring, both documented at the code:
A ~ 1e6besider ~ 3, socond(JᵀJ) ~ 1e12and float64 Cholesky loses nearly every digit. Unscaled: 247 iterations for a two-parameter fit, stopping 1.5-3.5% short ofmultifit. Scaled: 46 iterations to chisq ~1e-34.Known limitation
Convergence on hard many-parameter models is seed-limited — hyperspy hits
maxfevon the same model. That is exactly what seeded propagation (#54) is for; the speedup above is measured at an equal iteration budget for both.Not yet done in Wave 1
#54 seeding, #55 wizard, #56 picker, #57 drag-to-fit, #58 component area maps, #59 2D. No UI in this PR, so nothing here needs screenshot verification yet.