Skip to content

feat: 0.3.0 Wave 0 + Wave 1 — synthetic data, requires_package gate, batched GPU fitting - #82

Open
CSSFrancis wants to merge 6 commits into
mainfrom
feat/0.3.0-wave0-wave1
Open

feat: 0.3.0 Wave 0 + Wave 1 — synthetic data, requires_package gate, batched GPU fitting#82
CSSFrancis wants to merge 6 commits into
mainfrom
feat/0.3.0-wave0-wave1

Conversation

@CSSFrancis

Copy link
Copy Markdown
Owner

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.

Every spatial map differs from its transpose and both mirrors, so a flipped axis fails a test.

2. requires_package gate + 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 on spyde[all] so gated code actually runs.

Reversed on measurement: numexpr is NOT being added. HyperSpy warns it is missing and "slower to calculate model", but at spectrum sizes that is backwards:

numexpr spectra/s
absent 110
1 thread 77
4 threads 75
48 threads 70

Its per-call setup dominates ~1k-element arrays. pyproject.toml now 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/not spyde/models/, which is already the neural disk detector.

JᵀJ is (P, n, n) because n stays tiny; the only big object is the (P, C, n) Jacobian, so P is chunked. jacfwd not jacrev (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 with multifit, because "it converged" proves nothing when you are replacing a reference implementation.

Measured, 1024 spectra × 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.9×
batched cuda 3.97 s 257.8 spectra/s 54.1×

Two numerical traps found by measuring, both documented at the code:

  • Column scaling is load-bearing. A PowerLaw fits A ~ 1e6 beside r ~ 3, so cond(JᵀJ) ~ 1e12 and float64 Cholesky loses nearly every digit. 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. Step-size alone reads "stuck against a wall" as converged; requiring an accepted step never fires at all, since at the optimum nothing improves.

Known limitation

Convergence on hard many-parameter models is seed-limited — hyperspy hits maxfev on 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.

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.
@CSSFrancis

Copy link
Copy Markdown
Owner Author

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:

residual evaluation 3.4 ms
jacfwd Jacobian 51.6 ms — 94% of the solver

Autodiff 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). Components now carry analytic gradients:

before after
batched cpu 37.9 105.9 spectra/s (2.8×)
batched cuda 257.8 361.6 spectra/s (1.4×)

The GPU gains less because it is now dispatch-bound, not compute-bound — the remaining lever there is kernel fusion (torch.compile/CUDA graphs), not more math. Torch stays the right vehicle: already a dependency, and CPU + CUDA + MPS means it works on Mac too, which JAX-on-Windows and a hand-written CUDA kernel would not.

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:

  1. With exspy now installed, create_model() returns an EELSModel that already has a background — so the reference had two PowerLaws while the batched side had one. Both sides must fit the same model or the comparison means nothing.
  2. The EELS benchmark model cannot represent a hydrogenic edge, so neither engine converges and it was timing the iteration cap, not convergence. Added a well-posed EDS case (gaussian peaks, which the library expresses exactly) as the default; EELS stays as a deliberately-hard timing-only case, documented as such.

Well-posed case: 57.6% converged, 361 spectra/s, 11.6× multifit.

176 tests green.

Known, not yet addressed

The engine runs every pixel for max_iter even once most have converged, where hyperspy stops per-pixel — so the speedup ratio swings (11.6×–66×) with how hard the problem is. Compacting converged items out of the working set is the obvious next win.

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).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[1.1] spyde/fitting/ skeleton + serialisable ModelSpec

1 participant