feat(ebsd): Wave 3 — GPU dictionary indexing + pattern preprocessing - #83
Open
CSSFrancis wants to merge 3 commits into
Open
feat(ebsd): Wave 3 — GPU dictionary indexing + pattern preprocessing#83CSSFrancis wants to merge 3 commits into
CSSFrancis wants to merge 3 commits into
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.
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).