feat(ebsd): Wave 3 — GPU dictionary indexing + pattern preprocessing - #83
Merged
Conversation
The reason Wave 3 exists. Normalised cross-correlation between an experimental and a simulated pattern is, once both are zero-mean and unit-norm, EXACTLY a dot product — so indexing P patterns against D dictionary entries is one matmul plus a top-k, which is the operation a GPU is built for. The (P, D) similarity must never exist (65k x 100k float32 = 26 GB), so both axes are tiled with a RUNNING top-k: each tile is reduced against the best-so-far and discarded. A test asserts tiled == untiled, because otherwise results would quietly depend on chunk size. preprocess.py adds static/dynamic background removal and the average dot-product map, all batched over the whole scan. Why background removal is not cosmetic, measured here: NCC is invariant to gain and offset but NOT to a spatial GRADIENT. The detector background is in every experimental pattern and in no simulated one, so it dilutes every score. An exact-orientation match scores ~0.92 with it left in and >0.99 once removed — and the correction must be applied to BOTH sides, or the dictionary keeps low-frequency content its counterpart no longer has. Correctness comes from ground truth, not from a fixture: the synthetic EBSD data stamps the exact Euler angles it was generated from, so the tests assert that indexing recovers them, that the two grains index to disjoint entries (a blank map would pass every score check), that a uniform grain indexes consistently, and that a finer dictionary matches monotonically better. One real bug caught by its own test: the shape guard compared pixel counts AFTER reshaping the experimental stack to the dictionary's pixel count, which silently succeeds whenever the sizes share a factor — a 40x40 scan against a 20x20 dictionary reshaped to 4x as many "patterns" and indexed garbage rather than raising. Nothing here needs the ebsd extra; indexing is plain torch + numpy. Only the IO/geometry/master-pattern layer (#69) will, and that is requires_package-gated.
Dictionary indexing can only ever return a dictionary ENTRY, so its accuracy is capped by the sampling step — a 5-degree dictionary gives 5-degree answers. Refinement lifts that cap by optimising each orientation continuously from its indexed match. Same shape of problem as actions/vector_orientation_gpu.py, solved the same way: the whole field at once. Every pattern's three Euler angles are one row of a (P, 3) tensor, the simulated patterns are one batched forward pass, and one Adam optimiser walks all P orientations together. No dask, no per-pattern loop. Two things carried over from the vector-orientation work rather than rediscovered (CLAUDE.md, GPU Computing): - backward() segfaults off the main thread under CUDA on Windows, so the loop runs INLINE with an on_yield callback instead of being pushed to a worker. - Yield INSIDE the step loop, not between stages, or a UI freezes for seconds while the progress bar looks stuck. A test asserts on_yield fires during the optimisation, not just around it. Refinement never returns a WORSE orientation than it was given: it is an improvement step on top of indexing, so if Adam wanders off (bad start, too large an lr) the indexed answer is still the better estimate and is what comes back. Tested by driving it to diverge on purpose with lr=5. The simulator is injectable and the built-in one renders the same band geometry as the synthetic data — asserted equal to the numpy generator, since refining against a DIFFERENT forward model than the one that made the data would be fitting the wrong thing. #69 swaps in kikuchipy master patterns without touching the optimiser. Orientations are compared by the PATTERNS they produce, never by Euler distance: Euler angles are not a metric space and cubic symmetry means different triples describe the same crystal.
The display half of Wave 3, and it is nearly free by design. SpyDE already depends on orix and already has IPF views, orientation maps and a 3D IPF toolbar from the 4D-STEM work, so the job is to hand the EBSD result over in the form that machinery already speaks — not to build a second orientation display beside the first. ipf_colors returns (H, W, 3) RGB, which commit_result_tree already accepts as an RGB primary. orientation_similarity_map is the one piece with non-obvious semantics, and it earns its place: it measures how much each position's RANKED match list agrees with its neighbours', which exposes indexing that is confidently WRONG. A raw score map cannot — a pattern that matched something strongly scores just as highly whether or not the answer is right. Tested on exactly that case: one position indexed to a completely different entry from its neighbourhood reads 0.0 while the rest of the map reads > 0.9. merge_phases combines per-phase runs by score, since multi-phase indexing runs one dictionary per phase. This commit is COMPUTE ONLY. Putting these on screen is UI work and gets verified by running the app and looking at pixels (CLAUDE.md), not by these tests. The closest headless proxy is asserted instead: index the synthetic scan, colour it, and require the two grains to come out distinguishable — identical colours would mean the map renders as a flat block and would still pass every range check. Two small fixes while writing the tests: - merge_phases validated coverage AFTER np.stack had already raised "all input arrays must have the same shape", which says nothing about which phase disagreed. Validate first, and name the sizes. - numpy's assert_allclose no longer broadcasts shapes, so comparing (N, 3) against row 0 fails on shape even when every row is identical.
…ad of spots (#74) Wave 3 had the compute and no way to drive it. This is the visualisation half: the 4D-STEM orientation wizard, stage for stage, on EBSD patterns. Orientation Mapping EBSD Indexing --------------------------- ------------------------------------- 1 Load .cif + voltage 1 Load .cif + voltage + detector PC 2 Library diffsims templates 2 Library simulated pattern dictionary 3 Refine matched SPOTS 3 Refine matched Kikuchi BANDS 4 Run whole-field match 4 Run whole-field index (+ refine) A band centre is a straight line on a flat detector, so `spyde/ebsd/bands.py` projects the matched orientation's lines (and the zone axes they cross at) onto the live pattern under the crosshair. Downstream nothing is new: the result is packed into a SpyDEOrientationMap, so the IPF colouring, the 3-D explorer, the point selector and the direction toggle work untouched, with NCC / orientation-similarity / ADP as chip views beside the IPF map. Two bugs found on the way, both of which every score-based test passes through: * the simulator rotated plane normals by the TRANSPOSE of the Bunge matrix. Indexing is self-consistent either way — the dictionary and the data shared the mistake and matched each other perfectly — but the Euler angles handed to orix to colour the IPF map were then the inverse orientation. The tell is symmetry, and it is now pinned: a crystal 4-fold must leave the pattern alone, a sample rotation must not. * the dictionary was not put through the same background filter as the data. Both sides of an NCC have to come through one filter or the correction makes the mismatch different rather than smaller. BandSimulator now high-passes as it renders, so the dictionary build and the refinement inherit it together. On the bundled synthetic scan: 1.8 deg misorientation from a 5 deg dictionary, 0.09 deg after refinement, 1.5 ms per live match. Also: orix fundamental-zone sampling (a quarter of the orientations for the same resolution), a resident SinglePatternIndexer the batch path reuses, and progressive IPF fill via dictionary_index's new pattern_chunk/on_chunk. Verified in the real app — ebsd_workflow.spec.ts drives the caret and asserts the band lines and zone-axis markers actually render on the pattern.
The run materialised the entire dataset to index it — the memory-safety rule, broken outright. It is a lazy chunked array and every stage is embarrassingly parallel over navigation, which is how kikuchipy does it too. So the run is now ONE dask graph over the chunked scan: scan.map_blocks(correct -> index -> refine) per nav chunk, no halo scan.map_overlap(ADP, depth=1 on nav) the one stage with neighbours Chunks go to a thread pool and land as they finish, which is both the parallelism and the progressive fill — no hand-rolled block loop, no manual ghost rows, no whole-scan read anywhere. The static background reference is a dask reduction for the same reason. An eager scan gets a byte-sized nav chunking so it parallelises too; a lazy one keeps what it was loaded with, and patterns split across chunks are merged (a chunk holding part of a pattern makes every read useless). Measured on a 64x64 scan of 60x60 patterns: 8 chunks of 7.4 MB against a 59 MB scan — 12.5% resident at peak — across 8 worker threads. Guarded, not just fixed: test_ebsd_wizard raises if anything computes a dask array the shape of the full dataset. It caught two real bugs while being written — chunk alignment that rounded UP could swallow a single-chunk scan whole, and slicing a map_overlap result per chunk re-derives the halo and mis-shapes the output (only map_blocks graphs may be sliced that way). Refinement takes a lock: torch CUDA backward from several threads at once is the Windows trap in CLAUDE.md, and the GPU serialises those kernels anyway.
_lazy_scan re-implemented the setup half of BaseSignal.map: eager -> lazy with a navigation chunking, plus a rechunk so the signal axes span whole patterns. hyperspy exposes exactly that as LazySignal.rechunk(nav_chunks=, sig_chunks=), which is what map() itself uses to prepare its input — and it takes the two off the axes manager, so it no longer assumes 2 nav + 2 signal dimensions the way indexing chunks[2]/chunks[3] did. nav_chunks=None for an already-lazy signal, so it keeps the chunking it was loaded with: storage alignment beats an after-the-fact reshuffle (CLAUDE.md, 419 s against 184 s). inplace=False throughout — the navigator reads through the same array, so preparing a compute must not repoint it. Both are now asserted rather than assumed. Not switching the compute itself to map(): it calls the function once per navigation POSITION (signal.py -> process_function_blockwise loops np.ndindex inside each chunk), and dictionary indexing is one matmul E @ D.T per chunk. Per pattern that is P mat-vecs and, on the GPU, P kernel launches — the failure CLAUDE.md already records from the vector-orientation work. kikuchipy splits it the same way: .map() for the per-pattern steps (static/dynamic background, downsampling, image quality), raw dask for the rest (its ADP is dask_array.map_overlap, same as ours), and dictionary indexing builds a dask similarity expression without map() at all.
…3-ebsd # Conflicts: # electron/src/renderer/src/components/FloatingToolbar.tsx # electron/src/renderer/src/kernel/SpyDEContext.tsx # electron/src/renderer/src/kernel/protocol.ts # spyde/data/synthetic.py
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 #70, #71. Part of tracker #48.
Stacked on #82 (needs the synthetic EBSD data) — review that one first, base retargets to
mainonce it lands.The idea
NCC between an experimental and a simulated pattern is, once both are zero-mean and unit-norm, exactly a dot product. So indexing P patterns against D dictionary entries is one matmul plus a top-k — the operation a GPU is built for.
The
(P, D)similarity must never exist (65k × 100k float32 = 26 GB), so both axes are tiled with a running top-k: each tile is reduced against the best-so-far and discarded. A test asserts tiled == untiled, or results would quietly depend on chunk size.Background removal is not cosmetic
Measured here: NCC is invariant to gain and offset but not to a spatial gradient. The detector background is in every experimental pattern and no simulated one, so it dilutes every score.
And it must be applied to both sides — correcting only the experimental patterns leaves the dictionary carrying low-frequency content its counterpart no longer has, which is a different mismatch rather than a fix. kikuchipy preprocesses its dictionary for the same reason.
Correctness comes from ground truth, not a fixture
The synthetic EBSD data stamps the exact Euler angles it was generated from, so the tests assert:
One real bug, caught by its own test
The shape guard compared pixel counts after reshaping the experimental stack to the dictionary's pixel count — which silently succeeds whenever the sizes share a factor. A 40×40 scan against a 20×20 dictionary reshaped to 4× as many "patterns" and indexed garbage instead of raising.
Scope note
Nothing here needs the
ebsdextra — indexing is plain torch + numpy, so it develops and tests without kikuchipy. Only the IO/geometry/master-pattern layer (#69) will, and that isrequires_package-gated.sample_orientationsis a deliberately simple Euler grid, not equal-area SO(3) — enough to develop and test indexing; #69 wires up orix/kikuchipy's correct sampling.17 tests, all green. Still open in this wave: #69 (IO + detector geometry), #72 (refinement), #73 (CrystalMap → existing IPF display).