feat(atoms): Wave 4 — atom finding, GPU refinement, property maps - #85
Merged
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.
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).
Division of labour, and why this wave is small: atomap owns the STRUCTURE (peak finding, sublattices, neighbours, zone axes, dumbbell pairing) because reimplementing it would diverge from published atomap results. SpyDE owns the REFINEMENT, because refining atom positions IS a batched 2-D gaussian fit and spyde.fitting already does the whole field in one Levenberg-Marquardt where atomap fits one atom at a time with scipy. Only find_atoms needs the atoms extra; refinement and the property maps are plain numpy/torch, so they are testable and usable without it. Correctness comes from ground truth: the new atom_lattice() generator knows exactly where every atom is, so refinement is scored against those positions rather than against atomap's answer. The full pipeline (find -> COM -> gaussian) recovers them to a MEDIAN error of 0.2 px, and refinement alone to better than 0.1 px worst-case. Tolerances are tight deliberately — a test accepting +/-1 px would pass with the refinement removed entirely. Details that are easy to get wrong and are pinned by tests: - A pinned atom (the red/green toggle, #76) is EXCLUDED from the fit, not fitted and discarded — otherwise the work is wasted and a diverging neighbour can still perturb the shared solve. - A fit that leaves its own box keeps its INPUT position. A wild coordinate is worse than an unrefined one. - Gaussian A is a VOLUME, so the amplitude seed is scaled by 2*pi*sigma^2; seeding it with the peak height starts every fit an order of magnitude low. - Ellipticity is max/min sigma, never sigma_x/sigma_y: the latter drops below 1 for an atom elongated along y and puts a spurious boundary wherever the elongation direction changes. - Displacement uses a LOCAL reference (the centroid of k neighbours), so it measures real distortion and is immune to drift or a tilted scan, which a globally-fitted ideal lattice would report as displacement. atom_lattice() also grows dumbbell / displacement / ellipticity knobs so the downstream maps have something real to recover instead of a uniform zero.
Many structures image as PAIRS of closely-spaced atoms. Treating each pair as two independent atoms throws away what the pair actually measures — the separation and orientation that carry polarisation, tilt and local strain. Follows atomap's workflow: estimate the dumbbell vector, pair the atoms, then measure. The vector is estimated from the data rather than typed in, since it depends on the projected structure and the scan rotation. The part worth reading is refine_pairs, which exists because writing the tests exposed a systematic error rather than a bug. Fitting each atom of a dumbbell INDEPENDENTLY biases its centre TOWARDS its partner: the partner's tail is signal a single gaussian can only explain by moving. The separation therefore comes out too small — and uniformly so, which makes it look like a calibration rather than an error. Measured on a synthetic dumbbell of 6.0 px with sigma 2.0, independent fits return 4.21 px, a 30% underestimate. Fitting both atoms in ONE box as one two-gaussian model recovers 6.0, because each atom now explains the other's tail instead of absorbing it. Every pair is still one batched call — the engine does not care that the model gained a component. This is a good advert for the engine: the fix is a different MODEL, not different code. Three tests pin it: that independent fitting shows the bias (and says so if it ever stops), that joint fitting recovers the truth, and that joint beats independent by at least 5x — so a regression in either path shows up as the gap closing from the wrong side. Also pinned, because each is silently plausible if wrong: - The vector estimate must not CANCEL. Each pair contributes +v from one atom and -v from the other, so averaging without folding onto a half-plane gives exactly zero, which looks like a legitimate answer. - Pairing must not claim an atom twice, or there are more pairs than atoms. - The angle is folded to a half-turn, or two identical dumbbells differ by pi depending on which atom was listed first and an angle map shows a boundary that is not there.
Owner
Author
#78 dumbbell lattices — and a systematic error the tests exposedEstimate the dumbbell vector → pair the atoms → measure. The vector is estimated from the data rather than typed in, since it depends on the projected structure and the scan rotation. The part worth reading is Fitting each atom of a dumbbell independently biases its centre towards its partner — the partner's tail is signal a single gaussian can only explain by moving. The separation comes out too small, and uniformly so, which makes it look like a calibration rather than an error.
Fitting both atoms in one box as one two-gaussian model fixes it, because each atom now explains the other's tail instead of absorbing it. Every pair is still one batched call — the engine doesn't care that the model gained a component. This is a good advert for the engine, incidentally: the fix was a different model, not different code. Three tests pin it — that independent fitting shows the bias (and says so if that ever stops being true), that joint fitting recovers the truth, and that joint beats independent by ≥5×, so a regression in either path shows as the gap closing from the wrong side. Also pinned, because each is silently plausible if wrong
Wave 4's compute layer is now complete: #75, #77, #78, #79. 50 tests green. Still open here: #76 (the three atomap GUI functions as anyplotlib overlays) — UI, needs screenshots. |
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
1 participant
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.
Closes #75, #77, #79. Part of tracker #48.
Stacked on #82 (needs the 2-D fitting engine) — base retargets to
mainonce that lands.Division of labour
atomap owns the structure — peak finding, sublattices, neighbours, zone axes, dumbbell pairing — because reimplementing it would diverge from published atomap results.
SpyDE owns the refinement, because refining atom positions is a batched 2-D gaussian fit, and
spyde.fittingalready does the whole field in one Levenberg–Marquardt where atomap fits one atom at a time with scipy.Only
find_atomsneeds theatomsextra; refinement and the property maps are plain numpy/torch, so they're testable and usable without it.Accuracy, against ground truth
The new
atom_lattice()generator knows exactly where every atom is, so refinement is scored against those positions rather than against atomap's answer:Tolerances are tight deliberately — a test accepting ±1 px would pass with the refinement removed entirely.
Details pinned by tests because they're easy to get wrong
Ais a volume, so the amplitude seed is scaled by2πσ². Seeding with peak height starts every fit an order of magnitude low.Data
atom_lattice()grows dumbbell / displacement / ellipticity knobs so the downstream maps have something real to recover instead of a uniform zero. The grid is non-square and the ellipticity test requires the right-hand atoms to be measurably more elliptical than the left — a transposed or flipped result fails.25 tests, all green.
Still open in this wave
#76 (the three atomap GUI functions as anyplotlib overlays) and #78 (dumbbell lattices). #76 is UI — it gets verified by running the app and looking at pixels, which I've deliberately left for when you're around rather than doing unattended.