From b5f4382124980053e3f34d212d7bc60e66d8b616 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Sun, 26 Jul 2026 20:51:41 -0500 Subject: [PATCH 01/60] =?UTF-8?q?docs:=200.3.0=20release=20plan=20?= =?UTF-8?q?=E2=80=94=20fitting,=20EELS/EDS,=20EBSD,=20atom=20mapping?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- RELEASE_0_3_0_PLAN.md | 283 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 283 insertions(+) create mode 100644 RELEASE_0_3_0_PLAN.md diff --git a/RELEASE_0_3_0_PLAN.md b/RELEASE_0_3_0_PLAN.md new file mode 100644 index 0000000..4d5c5a1 --- /dev/null +++ b/RELEASE_0_3_0_PLAN.md @@ -0,0 +1,283 @@ +# SpyDE 0.3.0 — Release Plan + +Five waves, each independently shippable, each behind its own optional extra. +The tracker issue is **#48**; every wave and task below has a GitHub sub-issue. + +Prior art this plan leans on, in order of how much it saves us: +`vector_orientation_gpu.py` (batched-torch playbook), `commit.py` + +`actions/views.py` (component-map toggle), `vector_overlay.py` (interactive +point overlay), `ipf_view.py`/`orientation_map.py` (orientation display), +`actions/README.md` (the action contract — read before writing any action). + +--- + +## 0. Decisions locked before writing code + +### 0.1 `spyde/models/` is already taken + +It is the neural disk-detector package (SpotUNet + HF registry). The HyperSpy +model port therefore lives in **`spyde/fitting/`**. Do not "tidy" these +together — they are unrelated meanings of the word *model*. + +### 0.2 Package layout — one submodule per domain + +| Submodule | Wave | Extra | Contents | +|---|---|---|---| +| `spyde/fitting/` | 1 | — (core) | `ModelSpec`, torch component library, batched LM engine, seeding | +| `spyde/spectroscopy/` | 2 | `eels` | EELS/EDS composition → components, quantification | +| `spyde/ebsd/` | 3 | `ebsd` | pattern preprocessing, GPU dictionary indexing, refinement | +| `spyde/atoms/` | 4 | `atoms` | atom finding, sublattices, dumbbells, property maps | +| `spyde/data/` | 5 | — | example + synthetic datasets (owner: @CSSFrancis) | + +Interactive wiring stays in `spyde/actions/` per the existing contract +(compute in its own package, actions outside it — `find_vectors/` is the model). +No package split; these are submodules of `spyde`. + +### 0.3 Optional extras, gated at the toolbar + +```toml +[project.optional-dependencies] +eels = ["exspy>=0.3.2"] +ebsd = ["kikuchipy>=0.11.0"] +atoms = ["atomap>=0.4.1"] +all = ["spyde[eels,ebsd,atoms]"] +``` + +None of the three is installed today. Domain math comes from the real +libraries so our numbers match published results — **except** EBSD dictionary +indexing and refinement, which we implement in torch (§3) because that is the +whole point of Wave 3. + +New YAML gate key **`requires_package`**, mirroring the existing +`requires_vectors`: a toolbar entry is hidden when the extra is absent, and the +caret shows a one-line install hint instead of failing on import. This is the +only framework change Wave 0 needs. + +### 0.4 The fitting engine: batched GPU LM + seeded propagation + +Measured baseline on this box — synthetic EELS SI, 1024 spectra × 1024 +channels, power law + 2 gaussians, `multifit(optimizer="lm")`: + +| | | +|---|---| +| single `fit()` | 17.6 ms | +| `multifit` | **10.83 s** = 10.6 ms/spectrum = **95 spectra/s** | +| extrapolated to 256×256 | **11.6 min** | + +Single-threaded, one pixel at a time. That is the number to beat. + +Free win found while measuring: **`numexpr` is not installed**, so every +hyperspy `Expression` component falls back to numpy (it warns on each model +build). Add it to core deps — costs nothing, speeds up the CPU reference path +and the live single-spectrum preview. + +The engine is the `vector_orientation_gpu.py` playbook applied to curve +fitting: pack the whole nav grid into `(P, C)` on the GPU and run one batched +Levenberg–Marquardt, no Python loop over pixels. It works because **the +parameter count `n` is tiny** (≤ ~20) even when `P` and `C` are large: the +normal equations are `(P, n, n)` batched solves, which `torch.linalg.solve` +does natively. The only big intermediate is the `(P, C, n)` Jacobian, so +**chunk over `P`** to bound it (65536 × 2048 × 12 float32 = 6.4 GB whole — +chunked it is a fixed working set). This is the same "chunk the batch +dimension" rule the orientation code already follows. + +SAMFire's insight — a neighbour's result is the best starting point — is kept +as the **seed source**, not as the scheduler: fit a strided coarse grid, then +propagate those parameters outward as initial guesses for one batched refine +over all pixels. We do not port SAMFire's per-pixel marker/strategy machinery; +that scheduler overhead is what makes it slow. + +**Acceptance gate (non-negotiable):** the GPU engine must reproduce hyperspy's +`multifit` parameters within tolerance on the same data. We are replacing a +reference implementation, so parity against it is the test — not "it converged". +CPU (scipy) fallback stays for no-GPU machines and for the parity test itself. + +--- + +## Wave 1 — Model fitting, 1D and 2D (`spyde/fitting/`) + +The foundation. Waves 2 and 4 both fit models, so they consume this engine. + +**1.1 Package skeleton + `ModelSpec`.** A serialisable description of a model +(components, parameters, bounds, free/fixed flags, signal range) that +round-trips against `hyperspy.model.BaseModel.as_dictionary()`. Storage mirrors +hyperspy — `m.store(name)` / `s.models.restore(name)` — so a saved `.hspy` +opens with its models intact in hyperspy and in SpyDE alike. Lives on the tree +as `tree.fit_models` (ownership map, `actions/README.md` §3). + +**1.2 Torch component library.** Batched, autograd-differentiable ports of +`hyperspy.components1d`: Gaussian, GaussianHF, Lorentzian, Voigt, SplitVoigt, +PowerLaw, Offset, Polynomial, Exponential, Arctan, Erf, Doniach, SkewNormal, +Logistic, HeavisideStep, ScalableFixedPattern; `components2d`: Gaussian2D, +Expression. Each takes `(x, params)` → `(P, C)` fully batched. Per-component +parity test against the hyperspy component it mirrors. + +**1.3 Batched LM engine.** `fit_batched(spec, data, ...)`: chunked over `P`, +Jacobian by autograd, bounded parameters, free/fixed masking, channel mask for +signal range, optional Poisson weighting (matters for EELS/EDS counts). +Benchmark harness in `spyde/tests/benchmark_fitting.py` reporting spectra/s +against the 95/s baseline. **Also implement variable projection**: components +linear in amplitude (most of them) are solved by batched linear least squares +given the nonlinear parameters, which shrinks `n` and improves conditioning — +this is what makes Wave 2's tabulated EELS edges tractable. + +**1.4 Seeded propagation.** Coarse strided fit → propagate to neighbours as +initial guesses → one batched refine. Reports how many pixels converged. + +**1.5 The Fit wizard.** Staged action (`fit_open/_add_component/_set_param/ +_tune/_run/_commit/_close`) per the wizard protocol. Opens a caret group; +components are added line by line; live preview fits the *current* nav +position only, so it stays interactive. `fit_run` does the whole grid on a +worker with progress. + +**1.6 Component picker with shape preview.** "Add component" shows what each +component looks like: the backend evaluates each candidate at default +parameters over the current signal axis and sends a small polyline; the +renderer draws it as a sparkline in the picker. + +**1.7 Drag-to-adjust components on the plot.** Port of the anyplotlib +interactive-fitting example: `add_point_widget` for centre+amplitude, +`add_range_widget` for width, clicking a component's line toggles its widgets, +`pointer_move` re-evaluates the sum line. This is the primary way users tune a +model — sliders are the fallback, not the main path. + +**1.8 Component area maps.** Integrated area under each component per pixel → +`commit_result_tree` with one view per component, so they toggle exactly like +εxx/εyy/εxy/ω (single click shows one, ⌘-click tiles several). This is direct +reuse of `actions/views.py`; no new display code. + +**1.9 2D models.** Same engine, image-shaped data (`Gaussian2D`, `Expression`). +Lower priority within the wave but explicitly in scope. + +--- + +## Wave 2 — EELS and EDS (`spyde/spectroscopy/`, extra `eels`) — depends on Wave 1 + +Scope distilled from the eXSpy example gallery: find EDS lines, EELS curve +fitting, background removal, Fourier-ratio deconvolution, residual plots, +quantification. + +**2.1 exspy extra + signal-type gating** (`EELS`, `EDS_TEM`, `EDS_SEM` via the +existing `signal_types` key) + microscope-parameter panel (beam energy, +collection/convergence angle) since fits are meaningless without it. + +**2.2 Composition → components.** Enter elements; auto-populate the model: +EELS gets background + an `EELSCLEdge` per ionisation edge in range; EDS gets a +gaussian per X-ray line with family ratios tied. This is the wave's headline +feature. + +**2.3 Tabulated EELS edges in the engine.** `EELSCLEdge` is a GOS lookup, not +an analytic form. Interpolate the GOS table on-device and fit edge intensities +via the linear path from 1.3 — the shapes are fixed, the amplitudes are linear. + +**2.4 Background removal + edge onset** as a RegionAction (drag the fit +window), reusing the existing region-selector machinery. + +**2.5 Interactive drag fitting** — 1.7 applied to edges/lines, which is where +it pays off most. + +**2.6 Quantification → maps.** EELS relative composition; EDS Cliff–Lorimer / +zeta-factor with absorption. Output is per-element maps → the same view-toggle +commit as 1.8. + +**2.7 Fourier-ratio deconvolution, thickness map, residual view.** + +--- + +## Wave 3 — EBSD (`spyde/ebsd/`, extra `ebsd`) + +Independent of Waves 1–2. The reuse story here is unusually good: SpyDE already +depends on **orix** and already has IPF views, orientation maps and a 3D IPF +toolbar from the 4D-STEM work, so a CrystalMap feeds straight into existing +display code. + +**3.1 kikuchipy extra, EBSD signal type, IO** (`.h5`, `.ang`, EDAX/Oxford/ +Bruker via kikuchipy's readers) + `EBSDDetector` geometry panel. + +**3.2 Pattern preprocessing**: static/dynamic background removal, neighbour +averaging, average dot-product map. These are per-pattern and embarrassingly +parallel — straight to torch. + +**3.3 GPU dictionary indexing — the reason this wave exists.** Normalized +cross-correlation between `P` experimental and `D` dictionary patterns is, once +both are zero-mean/unit-norm, a single matmul `E @ Dᵀ` → `(P, D)`, followed by +top-k. Chunk over both `P` and `D` with a running top-k so the `(P, D)` +intermediate never materialises (65k × 100k float32 = 26 GB whole). kikuchipy +does this with dask; torch should be dramatically faster. Parity against +`kikuchipy.indexing` scores is the acceptance test. + +**3.4 Orientation refinement.** Batched optimisation of `(φ1, Φ, φ2)` and +optionally the projection centre — the direct analogue of +`vector_orientation_gpu.py`. Needs differentiable master-pattern sampling +(sphere → detector bilinear interpolation) in torch; that is the real work of +this task. Watch the same numerical traps the vector-orientation work hit: +rotation-branch ambiguity and coarse-stage bias (CLAUDE.md, GPU Computing). + +**3.5 CrystalMap → existing IPF display**, phase merging +(`merge_crystal_maps`), orientation similarity map. + +--- + +## Wave 4 — Atom position mapping (`spyde/atoms/`, extra `atoms`) — uses Wave 1's engine + +Atom refinement *is* a batched 2D gaussian fit, so 1.3 does the heavy lifting. + +**4.1 atomap extra + initial atom finding** (`get_atom_positions`), sublattice +construction, nearest neighbours, zone axes. + +**4.2 The three GUI functions, as anyplotlib overlays.** atomap's +`select_atoms_with_gui` / `add_atoms_with_gui` / +`toggle_atom_refine_position_with_gui` are matplotlib-based, so we reimplement +the *interaction* over our own overlay — which `vector_overlay.py` already does +for diffraction vectors (add/remove points, per-point state colouring). Polygon +select for phase separation; click to add/remove atoms; click to toggle an +atom's refine flag green↔red. + +**4.3 Refinement**: centre-of-mass then batched 2D gaussian via the Wave 1 +engine, honouring the per-atom refine flags. + +**4.4 Dumbbell lattices** (`get_dumbbell_vector`, `Dumbbell_Lattice`) per the +atomap dumbbell guide. + +**4.5 Property maps** — ellipticity, neighbour distance, monolayer spacing, +displacement — through the same view-toggle commit as 1.8. + +--- + +## Wave 5 — Data module (`spyde/data/`) — owner @CSSFrancis + +Datasets for 4D STEM, EELS, EDS, EBSD and in-situ. Waves 1–4 code against the +contract below and use synthetic stand-ins until real data lands, so this +never blocks them. + +**Contract** — mirror `spyde/backend/tutorial_data.py`: + +- **Bundled synthetic**, no download, fast enough for Playwright specs: + `load_test_data_` on `TestHarnessMixin`. Must be *asymmetric and + crisp* so bugs are pixel-visible (the `si_grains` / `movie` precedent). +- **Real, downloaded**: registered in the Examples menu, reached by + `load_example {name}`. +- Each entry declares `signal_type`, nav/signal shape, calibration and, where + relevant, the microscope parameters the fits need (§2.1) — a dataset that + arrives uncalibrated cannot exercise the feature it is meant to test. +- Lazy loads must use signal-spanning chunks (CLAUDE.md Live-Display §1). + +--- + +## Cross-cutting (Wave 0, do first) + +- `requires_package` gate key + install-hint caret (§0.3). +- `[project.optional-dependencies]` block + `numexpr` into core deps (§0.4). +- CI: one matrix job installing `spyde[all]` so gated code is actually + exercised; the default job proves the app still works with none of it. +- Docs: one guide per modality in `guides/`, and the data module's examples + wired into the existing tutorial menu. + +## Verification standard + +Per CLAUDE.md: a green pytest run and a clean `tsc` are **not** verification for +anything that adds windows, draws overlays or wires renderer↔backend. Every +wave's UI work ships with a Playwright spec built on `electron/tests/_harness.cjs` +and screenshots that were actually looked at. Where a wave replaces a reference +implementation (1.3 vs `multifit`, 3.3 vs `kikuchipy.indexing`), numerical +parity against that reference is the acceptance test. From 5a158d0f22549c62924f0b4a443a377980b2f9e6 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Sun, 26 Jul 2026 21:52:56 -0500 Subject: [PATCH 02/60] feat(data): synthetic EELS, EDS and EBSD datasets MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- spyde/backend/_session_actions.py | 7 + spyde/backend/_session_testharness.py | 45 ++ spyde/data/__init__.py | 30 ++ spyde/data/synthetic.py | 428 ++++++++++++++++++++ spyde/tests/migrated/test_synthetic_data.py | 237 +++++++++++ 5 files changed, 747 insertions(+) create mode 100644 spyde/data/__init__.py create mode 100644 spyde/data/synthetic.py create mode 100644 spyde/tests/migrated/test_synthetic_data.py diff --git a/spyde/backend/_session_actions.py b/spyde/backend/_session_actions.py index 56b5963..7b79f07 100644 --- a/spyde/backend/_session_actions.py +++ b/spyde/backend/_session_actions.py @@ -36,6 +36,7 @@ _TEST_ACTIONS = frozenset({ "load_test_data", "load_test_data_lazy", "load_test_data_lazy_chunked", "load_test_data_si_grains", "load_test_data_sped_ag", + "load_test_data_eels", "load_test_data_eds", "load_test_data_ebsd", "load_test_data_line", "load_test_data_movie", "test_nav_drag", "test_region_scrub", "test_add_second_navigator", "load_test_vectors", "run_test_orientation", "dump_dask_state", @@ -115,6 +116,12 @@ def dispatch_action(self, msg: dict) -> None: self._load_test_data_si_grains() elif action == "load_test_data_sped_ag": self._load_test_data_sped_ag() + elif action == "load_test_data_eels": + self._load_test_data_eels(payload) + elif action == "load_test_data_eds": + self._load_test_data_eds(payload) + elif action == "load_test_data_ebsd": + self._load_test_data_ebsd(payload) elif action == "load_test_data_line": self._load_test_data_line(payload) elif action == "load_test_data_movie": diff --git a/spyde/backend/_session_testharness.py b/spyde/backend/_session_testharness.py index bde015c..3c0ba05 100644 --- a/spyde/backend/_session_testharness.py +++ b/spyde/backend/_session_testharness.py @@ -210,6 +210,51 @@ def _load_test_data_si_grains(self) -> None: ax.offset = -(ax.size / 2.0) * float(ax.scale) self._add_signal(s, source_path="test_data_si_grains") + def _load_test_data_eels(self, payload: dict | None = None) -> None: + """Test-only: BUNDLED synthetic EELS spectrum image (``spyde.data``) — + 16×16 nav × 1024 channels, power-law background + C/N/O K edges whose + intensities follow known asymmetric concentration maps. No download. + + Ground truth is on ``metadata.Spyde.synthetic`` (read it with + ``spyde.data.ground_truth``), so a fit or quantification result is + scored against the numbers the data was built from.""" + from spyde.backend.heavy_imports import ensure_heavy_imports + ensure_heavy_imports() # see _load_test_data — don't race the prewarm + from spyde.data import eels_si + p = payload or {} + s = eels_si(nav=tuple(p.get("nav", (16, 16))), + n_channels=int(p.get("n_channels", 1024))) + self._add_signal(s, source_path="test_data_eels") + + def _load_test_data_eds(self, payload: dict | None = None) -> None: + """Test-only: BUNDLED synthetic EDS spectrum image (``spyde.data``) — + 16×16 nav × 2048 channels, bremsstrahlung + Fe/Ni/Cu K families at + their real energies, with Fe-Kβ/Ni-Kα deliberately overlapping so + family-aware fitting is actually exercised. No download.""" + from spyde.backend.heavy_imports import ensure_heavy_imports + ensure_heavy_imports() # see _load_test_data — don't race the prewarm + from spyde.data import eds_si + p = payload or {} + s = eds_si(nav=tuple(p.get("nav", (16, 16))), + n_channels=int(p.get("n_channels", 2048))) + self._add_signal(s, source_path="test_data_eds") + + def _load_test_data_ebsd(self, payload: dict | None = None) -> None: + """Test-only: BUNDLED synthetic EBSD patterns (``spyde.data``) — 16×16 + nav × 60×60 detector, real gnomonic Kikuchi-band geometry for a cubic + crystal over a two-grain orientation field. No download. + + The exact Euler angles are stamped on the signal, so dictionary + indexing / refinement is checked against ground truth rather than + against a fixture of its own output.""" + from spyde.backend.heavy_imports import ensure_heavy_imports + ensure_heavy_imports() # see _load_test_data — don't race the prewarm + from spyde.data import ebsd_patterns + p = payload or {} + s = ebsd_patterns(nav=tuple(p.get("nav", (16, 16))), + detector=tuple(p.get("detector", (60, 60)))) + self._add_signal(s, source_path="test_data_ebsd") + def _load_test_data_sped_ag(self) -> None: """Test-only: load the REAL sped_ag 4-D STEM scan (pyxem.data.sped_ag — 208×64 patterns of 112×112, a strained Ag SPED dataset with genuine diff --git a/spyde/data/__init__.py b/spyde/data/__init__.py new file mode 100644 index 0000000..dd63abc --- /dev/null +++ b/spyde/data/__init__.py @@ -0,0 +1,30 @@ +"""spyde.data — example and test datasets. + +Two halves, deliberately separate: + +``synthetic`` + Pure generator functions: numpy in, a HyperSpy signal out. No download, no + file, no Session. Small enough for a Playwright spec, and — the load-bearing + property — **asymmetric and crisp**, so a mirrored axis, a stale frame or a + transposed nav grid is visible in a screenshot rather than hiding behind + symmetric test data (the ``si_grains`` / ``movie`` precedent). + + Every generator stamps its ground truth into + ``metadata.Spyde.synthetic`` so a test can assert against the values the + data was built from instead of a golden file. + +Real / downloadable datasets (4D STEM, EELS, EDS, EBSD, in-situ) land here too +— that is Wave 5 (GitHub #80). The loaders that expose all of this to the UI +live in ``spyde/backend/tutorial_data.py`` (user-reachable) and +``spyde/backend/_session_testharness.py`` (test-gated). +""" +from __future__ import annotations + +from spyde.data.synthetic import ( + ebsd_patterns, + eds_si, + eels_si, + ground_truth, +) + +__all__ = ["eels_si", "eds_si", "ebsd_patterns", "ground_truth"] diff --git a/spyde/data/synthetic.py b/spyde/data/synthetic.py new file mode 100644 index 0000000..69d5e54 --- /dev/null +++ b/spyde/data/synthetic.py @@ -0,0 +1,428 @@ +"""synthetic.py — bundled synthetic datasets for EELS, EDS and EBSD. + +These exist so Waves 1-4 have something loadable per modality from day one, +before the real/downloadable datasets land (Wave 5, #80). They are also the +right thing to test against: the ground truth is *known exactly*, so a fit or +an indexing run is checked against the numbers the data was built from rather +than against a golden file that silently encodes yesterday's bug. + +Design rules, inherited from the ``si_grains`` / ``movie`` test data: + +* **Asymmetric.** Every spatial map is distinguishable from its own transpose + and from both mirrors — a ramp along one axis, a block off-centre, a wedge. + Symmetric test data hides exactly the bugs this data exists to catch. +* **Crisp.** Real peaks at real energies, real Kikuchi geometry. A fit that + works here should work on an experiment; a detector that finds nothing here + is broken. +* **Small and eager by default.** A Playwright spec must load one in well under + a second. Callers that want scale pass bigger shapes explicitly. +* **Ground truth in metadata** under ``Spyde.synthetic``. + +Nothing here imports torch, exspy or kikuchipy — generators are plain numpy so +they work in every environment, including CI without the optional extras. +""" +from __future__ import annotations + +import logging + +import numpy as np + +log = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# shared helpers +# --------------------------------------------------------------------------- + +def _concentration_maps(ny: int, nx: int) -> dict[str, np.ndarray]: + """Three deliberately asymmetric composition maps in [0, 1]. + + Each is distinguishable from its transpose and from both mirrors, so a + quantification map that comes out flipped is obvious on screen: + + ``a`` left-to-right ramp plus a block in the TOP-LEFT quadrant + ``b`` a wedge occupying the lower-right triangle, fading downward + ``c`` whatever is left over, so the three always sum to 1 + """ + yy, xx = np.mgrid[0:ny, 0:nx].astype(np.float32) + fy, fx = yy / max(ny - 1, 1), xx / max(nx - 1, 1) + + a = 0.15 + 0.55 * fx # ramps along x only + a[(yy < ny // 3) & (xx < nx // 3)] = 0.85 # TOP-LEFT block + + b = np.where(fx + fy > 1.0, 0.20 + 0.50 * fy, 0.05).astype(np.float32) + + total = a + b + over = total > 0.95 # keep room for the third element + a[over] *= 0.95 / total[over] + b[over] *= 0.95 / total[over] + return {"a": a, "b": b, "c": 1.0 - a - b} + + +def _stamp(sig, **truth) -> None: + """Record the parameters the data was generated from, so tests assert + against them instead of a golden file. + + Values stay as numpy arrays rather than nested lists: a 256x256 + concentration map is 65k floats, and metadata holds an array far more + cheaply than a list-of-lists. Read it back with :func:`ground_truth`. + """ + try: + sig.metadata.set_item("Spyde.synthetic", dict(truth)) + except Exception as e: # pragma: no cover + log.debug("stamping synthetic ground truth failed: %s", e) + + +def ground_truth(sig) -> dict: + """The generator parameters stamped on a synthetic signal, as plain Python. + + HyperSpy turns nested metadata dicts into ``DictionaryTreeBrowser`` nodes, + which have no ``.keys()``/``.values()`` and compare unhelpfully — so read + ground truth through here, not straight off ``metadata.Spyde.synthetic``. + """ + def _plain(v): + if hasattr(v, "as_dictionary"): + return {k: _plain(x) for k, x in v.as_dictionary().items()} + if isinstance(v, dict): + return {k: _plain(x) for k, x in v.items()} + return v + + node = sig.metadata.get_item("Spyde.synthetic", None) + if node is None: + raise ValueError(f"{sig!r} carries no synthetic ground truth — it did " + "not come from spyde.data.synthetic") + return _plain(node) + + +def _try_signal_type(sig, signal_type: str) -> None: + """``set_signal_type`` needs the matching HyperSpy extension (exspy for + EELS/EDS, kikuchipy for EBSD). Those are OPTIONAL extras, so a plain + install must still get usable data — it just stays a generic Signal1D/2D. + + HyperSpy logs a multi-line WARNING pointing at the extensions list when the + type is unknown. That is expected here and alarming to a user who simply + has not installed an extra, so it is demoted to a debug line. + """ + import logging as _logging + hs_io = _logging.getLogger("hyperspy.io") + prev = hs_io.level + hs_io.setLevel(_logging.ERROR) + try: + sig.set_signal_type(signal_type) + if type(sig).__name__ in ("Signal1D", "Signal2D"): + log.debug("signal_type %r unavailable (optional extra not " + "installed); left as %s", signal_type, type(sig).__name__) + except Exception as e: + log.debug("set_signal_type(%s) failed: %s", signal_type, e) + finally: + hs_io.setLevel(prev) + + +# --------------------------------------------------------------------------- +# EELS +# --------------------------------------------------------------------------- + +# Real core-loss onsets (eV). C-K and O-K sit in the same window with a decade +# of background between them, which is what makes background modelling matter. +EELS_EDGES = {"C_K": 284.0, "N_K": 401.0, "O_K": 532.0} + + +def _hydrogenic_edge(e: np.ndarray, onset: float, resolution: float) -> np.ndarray: + """A core-loss edge: sharp rise at the onset, then a ``E**-4`` style decay, + smeared by the spectrometer resolution. + + Not a GOS table — the real shapes come from exspy in Wave 2 (#63). This is + the right *shape* for exercising a fit: a step the background cannot mimic, + followed by a tail that overlaps the next edge. + """ + above = e >= onset + shape = np.zeros_like(e) + ratio = np.divide(e, onset, out=np.ones_like(e), where=above) + shape[above] = ratio[above] ** -4.0 + # Spectrometer resolution: gaussian smear, width in channels. + step = float(e[1] - e[0]) + sigma_ch = max(resolution / step, 0.8) + half = int(np.ceil(3 * sigma_ch)) + k = np.exp(-0.5 * (np.arange(-half, half + 1) / sigma_ch) ** 2) + k /= k.sum() + return np.convolve(shape, k, mode="same") + + +def eels_si(nav=(16, 16), n_channels: int = 1024, *, e_start: float = 200.0, + e_stop: float = 800.0, resolution: float = 1.2, + counts: float = 3.0e5, seed: int = 0, noise: bool = True): + """A synthetic EELS spectrum image: power-law background + three core-loss + edges whose intensities follow the asymmetric composition maps. + + Returns a ``Signal1D`` with a calibrated eV signal axis. Ground truth (the + per-element concentration maps and the onsets) is in + ``metadata.Spyde.synthetic``, so a quantification result can be scored + directly against the maps it should recover. + + The default 16x16 x 1024 is Playwright-sized. Pass a bigger *nav* for + benchmarking — ``eels_si(nav=(256, 256))`` is the 11.6-minute ``multifit`` + case that Wave 1 exists to fix. + """ + import hyperspy.api as hs + + ny, nx = int(nav[0]), int(nav[1]) + rng = np.random.default_rng(seed) + e = np.linspace(e_start, e_stop, n_channels, dtype=np.float64) + + conc = _concentration_maps(ny, nx) + names = list(EELS_EDGES) + # Background thickness also varies across the scan (and along y, so it is + # independent of every concentration map) — a background model that ignores + # position gets this wrong in a way a flat background would hide. + yy = np.mgrid[0:ny, 0:nx][0].astype(np.float32) / max(ny - 1, 1) + bg_scale = (0.7 + 0.6 * yy).astype(np.float64) + + edges = {n: _hydrogenic_edge(e, EELS_EDGES[n], resolution) for n in names} + background = (e / e_start) ** -3.0 + + data = np.empty((ny, nx, n_channels), dtype=np.float32) + for iy in range(ny): + for ix in range(nx): + spec = bg_scale[iy, ix] * background + for key, name in zip(("a", "b", "c"), names): + spec = spec + 0.35 * conc[key][iy, ix] * edges[name] + data[iy, ix] = spec + data *= counts / data.max() + if noise: + data = rng.poisson(np.clip(data, 0, None)).astype(np.float32) + + s = hs.signals.Signal1D(data) + ax = s.axes_manager.signal_axes[0] + ax.name, ax.units = "Energy loss", "eV" + ax.offset, ax.scale = e_start, float(e[1] - e[0]) + for i, a in enumerate(s.axes_manager.navigation_axes): + a.name, a.units, a.scale = ("x", "y")[i], "nm", 2.0 + s.metadata.General.title = "Synthetic EELS SI" + # The fits are meaningless without these (#61) — ship them with the data. + s.metadata.set_item("Acquisition_instrument.TEM.beam_energy", 200.0) + s.metadata.set_item("Acquisition_instrument.TEM.convergence_angle", 10.0) + s.metadata.set_item("Acquisition_instrument.TEM.Detector.EELS." + "collection_angle", 20.0) + _try_signal_type(s, "EELS") + _stamp(s, kind="eels", edges={n: EELS_EDGES[n] for n in names}, + elements=names, + concentration={n: conc[k] + for k, n in zip(("a", "b", "c"), names)}, + background_exponent=3.0, resolution_ev=resolution) + return s + + +# --------------------------------------------------------------------------- +# EDS +# --------------------------------------------------------------------------- + +# Real characteristic line energies (keV) with in-family intensity ratios. +EDS_LINES = { + "Fe": [("Ka", 6.404, 1.00), ("Kb", 7.058, 0.13)], + "Ni": [("Ka", 7.478, 1.00), ("Kb", 8.265, 0.13)], + "Cu": [("Ka", 8.048, 1.00), ("Kb", 8.905, 0.13)], +} + + +def _detector_sigma(e_kev: np.ndarray | float) -> np.ndarray | float: + """Si(Li)/SDD resolution: ``sigma**2`` grows linearly with energy (Fano). + Tuned to ~130 eV FWHM at Mn-Ka, i.e. a normal EDS detector.""" + return np.sqrt(1.4e-4 + 1.1e-3 * np.asarray(e_kev, float)) / 2.3548 * 2.3548 + + +def eds_si(nav=(16, 16), n_channels: int = 2048, *, e_stop: float = 20.0, + counts: float = 4.0e4, seed: int = 0, noise: bool = True): + """A synthetic EDS spectrum image: bremsstrahlung background + K-family + lines for Fe / Ni / Cu at their real energies. + + Fe-Kb (7.058) and Ni-Ka (7.478) overlap, and Ni-Kb (8.265) sits between + Cu-Ka (8.048) and Cu-Kb (8.905) — deliberately, because resolving + overlapping families is the thing EDS fitting has to get right. A peak + finder that treats every maximum as one element fails here, as it should. + + Ground truth (concentrations, line energies, family ratios) is in + ``metadata.Spyde.synthetic``. + """ + import hyperspy.api as hs + + ny, nx = int(nav[0]), int(nav[1]) + rng = np.random.default_rng(seed) + e = np.linspace(0.0, e_stop, n_channels, dtype=np.float64) + scale = float(e[1] - e[0]) + + conc = _concentration_maps(ny, nx) + elements = list(EDS_LINES) + + # Bremsstrahlung: Kramers' law with detector absorption killing the low end. + with np.errstate(divide="ignore", invalid="ignore"): + brem = np.where(e > 0.15, (20.0 - e) / np.maximum(e, 1e-3), 0.0) + brem = np.clip(brem, 0, None) * (1.0 - np.exp(-np.maximum(e, 0) / 1.5)) + brem /= brem.max() + + peaks = {} + for el in elements: + acc = np.zeros_like(e) + for _name, energy, ratio in EDS_LINES[el]: + sig = _detector_sigma(energy) + acc += ratio * np.exp(-0.5 * ((e - energy) / sig) ** 2) + peaks[el] = acc + + data = np.empty((ny, nx, n_channels), dtype=np.float32) + for iy in range(ny): + for ix in range(nx): + spec = 0.25 * brem + for key, el in zip(("a", "b", "c"), elements): + spec = spec + conc[key][iy, ix] * peaks[el] + data[iy, ix] = spec + data *= counts / data.max() + if noise: + data = rng.poisson(np.clip(data, 0, None)).astype(np.float32) + + s = hs.signals.Signal1D(data) + ax = s.axes_manager.signal_axes[0] + ax.name, ax.units = "Energy", "keV" + ax.offset, ax.scale = 0.0, scale + for i, a in enumerate(s.axes_manager.navigation_axes): + a.name, a.units, a.scale = ("x", "y")[i], "nm", 5.0 + s.metadata.General.title = "Synthetic EDS SI" + s.metadata.set_item("Acquisition_instrument.TEM.beam_energy", 200.0) + s.metadata.set_item("Acquisition_instrument.TEM.Detector.EDS." + "energy_resolution_MnKa", 130.0) + s.metadata.set_item("Acquisition_instrument.TEM.Detector.EDS.live_time", 10.0) + _try_signal_type(s, "EDS_TEM") + _stamp(s, kind="eds", elements=elements, + lines={el: EDS_LINES[el] for el in elements}, + concentration={el: conc[k] + for k, el in zip(("a", "b", "c"), elements)}) + return s + + +# --------------------------------------------------------------------------- +# EBSD +# --------------------------------------------------------------------------- + +def _euler_to_matrix(phi1, Phi, phi2) -> np.ndarray: + """Bunge ZXZ Euler angles -> rotation matrices, batched over the leading + axes. Returns ``(..., 3, 3)``.""" + c1, s1 = np.cos(phi1), np.sin(phi1) + c, s = np.cos(Phi), np.sin(Phi) + c2, s2 = np.cos(phi2), np.sin(phi2) + m = np.empty(np.shape(phi1) + (3, 3), float) + m[..., 0, 0] = c1 * c2 - s1 * s2 * c + m[..., 0, 1] = s1 * c2 + c1 * s2 * c + m[..., 0, 2] = s2 * s + m[..., 1, 0] = -c1 * s2 - s1 * c2 * c + m[..., 1, 1] = -s1 * s2 + c1 * c2 * c + m[..., 1, 2] = c2 * s + m[..., 2, 0] = s1 * s + m[..., 2, 1] = -c1 * s + m[..., 2, 2] = c + return m + + +def _cubic_plane_normals() -> tuple[np.ndarray, np.ndarray]: + """{111}, {200} and {220} plane normals for a cubic crystal, with rough + structure-factor weights. One normal per Friedel pair (a band and its + inverse are the same band).""" + fams = [((1, 1, 1), 1.00), ((2, 0, 0), 0.70), ((2, 2, 0), 0.45)] + normals, weights = [], [] + for hkl, w in fams: + seen = set() + h, k, l = hkl + for perm in {(h, k, l), (h, l, k), (k, h, l), (k, l, h), (l, h, k), (l, k, h)}: + for sx in (1, -1): + for sy in (1, -1): + for sz in (1, -1): + v = (perm[0] * sx, perm[1] * sy, perm[2] * sz) + if v == (0, 0, 0) or tuple(-x for x in v) in seen: + continue + seen.add(v) + for v in seen: + n = np.array(v, float) + normals.append(n / np.linalg.norm(n)) + # Band width goes as 1/|g|: bigger d-spacing -> wider band. + weights.append(w / np.linalg.norm(np.array(v, float))) + return np.array(normals), np.array(weights) + + +def ebsd_patterns(nav=(16, 16), detector=(60, 60), *, pc=(0.5, 0.5, 0.55), + seed: int = 0, noise: float = 0.06): + """Synthetic EBSD patterns with **exact known orientations**. + + Each nav pixel gets an orientation; the pattern is built by gnomonic + projection of a cubic crystal's {111}/{200}/{220} Kikuchi bands onto a flat + detector. A band appears where the detector direction is near-perpendicular + to a plane normal, which is the real geometry — so a dictionary-indexing + run (#71) or a refinement (#72) that recovers the stamped Euler angles is + genuinely working, not matching a fixture to itself. + + The orientation field is **two grains plus a gradient**, arranged + asymmetrically: a rotated wedge in the lower-right against a background + grain whose orientation drifts along x. That gives an IPF map with a real + boundary and an intra-grain gradient — the two things an orientation map + has to show correctly. + + Returns a ``Signal2D`` (nav | detector). Ground-truth Euler angles are in + ``metadata.Spyde.synthetic["euler"]`` as ``(ny, nx, 3)`` radians. + """ + import hyperspy.api as hs + + ny, nx = int(nav[0]), int(nav[1]) + dy, dx = int(detector[0]), int(detector[1]) + rng = np.random.default_rng(seed) + + # --- orientation field ------------------------------------------------ + yy, xx = np.mgrid[0:ny, 0:nx].astype(float) + fy, fx = yy / max(ny - 1, 1), xx / max(nx - 1, 1) + phi1 = np.deg2rad(10.0 + 35.0 * fx) # drifts along x only + Phi = np.full((ny, nx), np.deg2rad(30.0)) + phi2 = np.full((ny, nx), np.deg2rad(15.0)) + grain2 = (fx + fy) > 1.15 # lower-right wedge + phi1[grain2], Phi[grain2], phi2[grain2] = (np.deg2rad(70.0), + np.deg2rad(55.0), + np.deg2rad(40.0)) + rot = _euler_to_matrix(phi1, Phi, phi2) # (ny, nx, 3, 3) + + # --- detector directions (gnomonic) ----------------------------------- + pcx, pcy, L = float(pc[0]), float(pc[1]), float(pc[2]) + gy, gx = np.mgrid[0:dy, 0:dx].astype(float) + rx = (gx + 0.5) / dx - pcx + ry = pcy - (gy + 0.5) / dy # detector y is flipped + r = np.stack([rx, ry, np.full_like(rx, L)], -1) + r /= np.linalg.norm(r, axis=-1, keepdims=True) # (dy, dx, 3) + + normals, weights = _cubic_plane_normals() + # Band half-width in units of cos(angle to the normal); scaled by |g| so + # low-index planes give the wide bands, as in a real pattern. + widths = 0.055 * (weights / weights.max()) + 0.012 + + data = np.empty((ny, nx, dy, dx), dtype=np.float32) + flat_r = r.reshape(-1, 3) + for iy in range(ny): + for ix in range(nx): + # Rotate the plane normals into the sample frame for this pixel. + n_rot = normals @ rot[iy, ix].T # (B, 3) + d = flat_r @ n_rot.T # (dy*dx, B) + band = np.exp(-0.5 * (d / widths) ** 2) * weights + data[iy, ix] = band.sum(1).reshape(dy, dx).astype(np.float32) + + # Detector background: real patterns sit on a smooth, off-centre gradient + # that background correction (#70) has to remove. + bg = 0.35 + 0.30 * ((gy / dy) * 0.6 + (gx / dx) * 0.4) + data = data / max(data.max(), 1e-9) + bg.astype(np.float32) + if noise: + data += rng.normal(0.0, noise, data.shape).astype(np.float32) + data = np.clip(data, 0, None) + data = (data / data.max() * 255.0).astype(np.uint8) + + s = hs.signals.Signal2D(data) + for i, a in enumerate(s.axes_manager.navigation_axes): + a.name, a.units, a.scale = ("x", "y")[i], "um", 0.5 + for i, a in enumerate(s.axes_manager.signal_axes): + a.name = ("dx", "dy")[i] + s.metadata.General.title = "Synthetic EBSD" + _try_signal_type(s, "EBSD") + _stamp(s, kind="ebsd", euler=np.stack([phi1, Phi, phi2], -1), + pc=np.array([pcx, pcy, L]), n_bands=int(len(normals)), + grain2_mask=grain2) + return s diff --git a/spyde/tests/migrated/test_synthetic_data.py b/spyde/tests/migrated/test_synthetic_data.py new file mode 100644 index 0000000..e7b38ab --- /dev/null +++ b/spyde/tests/migrated/test_synthetic_data.py @@ -0,0 +1,237 @@ +"""Contract tests for the bundled synthetic datasets (spyde/data/synthetic.py). + +These pin the two properties Waves 1-4 actually rely on: + +1. **The physics is where it claims to be** — an edge steps up at its onset, a + line peaks at its energy, a Kikuchi pattern changes when the orientation + does. Data that merely *has the right shape* would let a broken fit pass. +2. **Nothing is symmetric.** Every spatial map differs from its transpose and + from both mirrors, so a flipped axis or a transposed nav grid fails a test + instead of silently looking fine. +""" +from __future__ import annotations + +import numpy as np +import pytest + +from spyde.data import ebsd_patterns, eds_si, eels_si, ground_truth +from spyde.data.synthetic import EDS_LINES, EELS_EDGES + + +def _asymmetric(arr: np.ndarray) -> bool: + """True when `arr` is distinguishable from its transpose and both mirrors.""" + a = np.asarray(arr, float) + return (not np.allclose(a, a.T) + and not np.allclose(a, a[:, ::-1]) + and not np.allclose(a, a[::-1])) + + +class TestEELS: + def test_shape_and_calibration(self): + s = eels_si(nav=(6, 5), n_channels=256) + assert s.data.shape == (6, 5, 256) + ax = s.axes_manager.signal_axes[0] + assert ax.units == "eV" + assert ax.offset == pytest.approx(200.0) + # The axis must actually span the edges, or the fit has nothing to fit. + e = ax.axis + for onset in EELS_EDGES.values(): + assert e[0] < onset < e[-1] + + def test_edges_step_up_at_their_onsets(self): + """A core-loss edge is a STEP: intensity just above the onset exceeds + intensity just below it. + + Measured at the pixel where that element is most abundant. Testing a + fixed pixel would be measuring the wrong thing — at a pixel where the + element is nearly absent the power-law background's own decay across + the window is larger than the edge, which is exactly why background + modelling (#64) is a real task and not a formality. + """ + s = eels_si(nav=(8, 8), n_channels=2048, noise=False) + e = s.axes_manager.signal_axes[0].axis + conc = ground_truth(s)["concentration"] + for name, onset in EELS_EDGES.items(): + iy, ix = np.unravel_index(int(np.argmax(np.asarray(conc[name]))), + s.data.shape[:2]) + spec = s.data[iy, ix].astype(float) + i = int(np.argmin(np.abs(e - onset))) + pre, post = spec[i - 40:i - 8], spec[i + 8:i + 40] + assert post.mean() > pre.mean(), ( + f"no step at {onset} eV ({name}) in the {name}-rich pixel") + + def test_ground_truth_is_plain_python(self): + s = eels_si(nav=(4, 4), n_channels=128) + gt = ground_truth(s) + assert gt["kind"] == "eels" + assert set(gt["concentration"]) == set(EELS_EDGES) + # DictionaryTreeBrowser has no .values(); ground_truth() must undo that. + assert isinstance(gt["concentration"], dict) + assert list(gt["concentration"].values()) + + def test_concentration_maps_are_asymmetric_and_distinct(self): + s = eels_si(nav=(12, 12), n_channels=128) + maps = ground_truth(s)["concentration"] + for name, m in maps.items(): + assert _asymmetric(np.asarray(m)), f"{name} map is symmetric" + vals = [np.asarray(m) for m in maps.values()] + for i in range(len(vals)): + for j in range(i + 1, len(vals)): + assert not np.allclose(vals[i], vals[j]) + + def test_edge_intensity_tracks_its_concentration_map(self): + """The whole point of the ground truth: more of element X really does + mean a bigger X edge. A quantification result can be scored on this.""" + s = eels_si(nav=(10, 10), n_channels=1024, noise=False) + e = s.axes_manager.signal_axes[0].axis + gt = ground_truth(s) + onset = EELS_EDGES["O_K"] + i = int(np.argmin(np.abs(e - onset))) + step = s.data[..., i + 3:i + 30].mean(-1) - s.data[..., i - 30:i - 3].mean(-1) + conc = np.asarray(gt["concentration"]["O_K"]) + # Rank correlation is enough — background varies independently. + assert np.corrcoef(step.ravel(), conc.ravel())[0, 1] > 0.8 + + +class TestEDS: + def test_shape_and_calibration(self): + s = eds_si(nav=(5, 6), n_channels=512) + assert s.data.shape == (5, 6, 512) + assert s.axes_manager.signal_axes[0].units == "keV" + + def test_lines_peak_at_their_real_energies(self): + s = eds_si(nav=(4, 4), n_channels=4096, noise=False) + e = s.axes_manager.signal_axes[0].axis + spec = s.data.sum((0, 1)).astype(float) + for el, lines in EDS_LINES.items(): + energy = lines[0][1] # the Ka line + i = int(np.argmin(np.abs(e - energy))) + win = spec[i - 30:i + 30] + # The local maximum must sit at the line, not merely nearby. + assert abs(int(np.argmax(win)) - 30) <= 3, f"{el} Ka misplaced" + + def test_ka_kb_ratio_recoverable_for_an_isolated_family(self): + """Ka:Kb ratio is what makes family-aware fitting testable — checked on + Fe, whose Kb (7.058) is the one line with no near neighbour, and at the + Fe-rich pixel so the other families contribute almost nothing.""" + from scipy.ndimage import percentile_filter + + s = eds_si(nav=(8, 8), n_channels=4096, noise=False) + e = s.axes_manager.signal_axes[0].axis + conc = np.asarray(ground_truth(s)["concentration"]["Fe"]) + iy, ix = np.unravel_index(int(np.argmax(conc)), conc.shape) + spec = s.data[iy, ix].astype(float) + + # Rolling low percentile over a window several FWHM wide follows the + # bremsstrahlung and ignores the narrow lines. Hand-picked flat windows + # do NOT work here: at 4096 channels every "clean" window near Fe-Kb is + # within a few hundred eV of another line. + width = int(round(1.0 / float(e[1] - e[0]))) # ~1 keV + baseline = percentile_filter(spec, 10, size=width) + net = spec - baseline + + def height(energy): + i = int(np.argmin(np.abs(e - energy))) + return float(net[i - 20:i + 20].max()) + + (_, e_a, r_a), (_, e_b, r_b) = EDS_LINES["Fe"][0], EDS_LINES["Fe"][1] + assert height(e_b) / height(e_a) == pytest.approx(r_b / r_a, rel=0.25) + + def test_families_overlap_as_intended(self): + """The overlaps are a FEATURE — they are why a peak-per-element finder + is not good enough and family-aware fitting (#62) is needed. Pin them + so a future tweak to the line table cannot quietly remove the + difficulty this dataset exists to pose.""" + from spyde.data.synthetic import _detector_sigma + flat = [(el, en) for el, ls in EDS_LINES.items() for _n, en, _r in ls] + close = [] + for i, (el_a, e_a) in enumerate(flat): + for el_b, e_b in flat[i + 1:]: + if el_a == el_b: + continue + fwhm = 2.3548 * float(_detector_sigma((e_a + e_b) / 2)) + if abs(e_a - e_b) < 2.5 * fwhm: + close.append((el_a, el_b, abs(e_a - e_b))) + assert close, "no inter-element overlaps left in the EDS line table" + + def test_ground_truth_lines_readable(self): + gt = ground_truth(eds_si(nav=(3, 3), n_channels=128)) + assert gt["kind"] == "eds" + assert set(gt["lines"]) == set(EDS_LINES) + assert all(_asymmetric(np.asarray(m)) + for m in gt["concentration"].values()) + + +class TestEBSD: + def test_shape_and_dtype(self): + s = ebsd_patterns(nav=(5, 4), detector=(32, 32)) + assert s.data.shape == (5, 4, 32, 32) + assert s.data.dtype == np.uint8 + + def test_patterns_are_not_mirror_symmetric(self): + """A Kikuchi pattern that survives a flip cannot catch a flipped + detector axis — the bug this data exists to make visible.""" + s = ebsd_patterns(nav=(3, 3), detector=(48, 48), noise=0.0) + p = s.data[0, 0].astype(float) + assert not np.allclose(p, p[:, ::-1]) + assert not np.allclose(p, p[::-1]) + + def test_orientation_field_has_two_distinct_grains(self): + s = ebsd_patterns(nav=(16, 16), detector=(40, 40), noise=0.0) + gt = ground_truth(s) + eul, mask = np.asarray(gt["euler"]), np.asarray(gt["grain2_mask"], bool) + assert mask.any() and not mask.all(), "wedge grain missing" + # Grain 2 is a single orientation; grain 1 drifts along x. + assert np.allclose(eul[mask], eul[mask][0]) + assert not np.allclose(eul[~mask], eul[~mask][0]) + + def test_same_orientation_gives_the_same_pattern(self): + """Ground truth is only usable if the mapping orientation->pattern is + deterministic: two pixels of grain 2 must be pixel-identical.""" + s = ebsd_patterns(nav=(16, 16), detector=(40, 40), noise=0.0) + mask = np.asarray(ground_truth(s)["grain2_mask"], bool) + ys, xs = np.nonzero(mask) + a = s.data[ys[0], xs[0]] + b = s.data[ys[-1], xs[-1]] + assert np.array_equal(a, b) + + def test_different_orientation_gives_a_different_pattern(self): + s = ebsd_patterns(nav=(16, 16), detector=(40, 40), noise=0.0) + mask = np.asarray(ground_truth(s)["grain2_mask"], bool) + ys, xs = np.nonzero(mask) + ys2, xs2 = np.nonzero(~mask) + g2 = s.data[ys[0], xs[0]].astype(float) + g1 = s.data[ys2[0], xs2[0]].astype(float) + # Normalised cross-correlation well below 1 — i.e. actually indexable. + ncc = float(np.corrcoef(g1.ravel(), g2.ravel())[0, 1]) + assert ncc < 0.9, f"grains too similar to index (ncc={ncc:.3f})" + + def test_band_count_matches_cubic_families(self): + """4x{111} + 3x{200} + 6x{220} = 13 Friedel-unique bands.""" + assert ground_truth(ebsd_patterns(nav=(2, 2), + detector=(16, 16)))["n_bands"] == 13 + + +class TestNoOptionalExtrasRequired: + def test_generators_work_without_exspy_or_kikuchipy(self): + """The extras are optional, so generation must never depend on them — + the signal just stays a generic Signal1D/Signal2D.""" + import hyperspy.api as hs + assert isinstance(eels_si(nav=(2, 2), n_channels=64), hs.signals.BaseSignal) + assert isinstance(eds_si(nav=(2, 2), n_channels=64), hs.signals.BaseSignal) + assert isinstance(ebsd_patterns(nav=(2, 2), detector=(8, 8)), + hs.signals.BaseSignal) + + +class TestSessionLoaders: + """The three harness loaders reach the UI path (a tree + plots appear).""" + + @pytest.mark.parametrize("action", ["load_test_data_eels", + "load_test_data_eds", + "load_test_data_ebsd"]) + def test_loader_creates_a_tree(self, window, action): + session = window["window"] + session.dispatch_action({"action": action, + "payload": {"nav": (4, 4)}}) + assert len(window["signal_trees"]) == 1 + assert window["plots"] From d069b8ebaad3ff3722584d7f9ad9c909d07541b8 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Sun, 26 Jul 2026 21:53:12 -0500 Subject: [PATCH 03/60] feat(toolbars): requires_package gate + domain extras (0.3.0 Wave 0) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- .github/workflows/build.yml | 45 +++++ RELEASE_0_3_0_PLAN.md | 26 ++- pyproject.toml | 18 ++ spyde/actions/README.md | 9 + spyde/actions/_template_action.py | 2 + .../drawing/toolbars/plot_control_toolbar.py | 59 +++++++ .../migrated/test_requires_package_gate.py | 159 ++++++++++++++++++ 7 files changed, 312 insertions(+), 6 deletions(-) create mode 100644 spyde/tests/migrated/test_requires_package_gate.py diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index ac6453f..669f385 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -67,6 +67,51 @@ jobs: with: token: ${{ secrets.CODECOV_TOKEN }} + # ── Domain extras (0.3.0) ────────────────────────────────────────────────── + # The `test` job above runs with NO optional extras, which is the important + # half of the contract: SpyDE must work for a user who only does 4D-STEM, and + # the `requires_package:` gate must hide the EELS/EBSD/atoms actions rather + # than raising ImportError. This job is the other half — with exspy, + # kikuchipy and atomap installed the gated code actually executes, so a broken + # EELS/EBSD/atom path fails here instead of silently never running. + # + # One OS / one Python: the extras are pure-Python scientific packages and the + # cross-platform surface is already covered by `test`. Not part of the main + # matrix because that would quadruple install time for every leg. + test-extras: + name: extras (ubuntu, py3.12) + runs-on: ubuntu-latest + timeout-minutes: 40 + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: Set up uv + uses: astral-sh/setup-uv@v6 + with: + version: "0.10.11" + enable-cache: true + + - name: Set up Python + run: uv python install 3.12 + + # NOT --frozen: the domain extras are not in uv.lock (they are optional + # and resolve differently per platform), so the lock cannot satisfy them. + - name: Install with all domain extras + run: uv sync --extra tests --extra all --python 3.12 + + - name: Versions + run: | + uv run python -V + uv run python -c "import exspy, kikuchipy, atomap; \ + print('exspy', exspy.__version__); \ + print('kikuchipy', kikuchipy.__version__); \ + print('atomap', atomap.__version__)" + + - name: Run tests (extras present) + run: uv run pytest spyde/tests/migrated + # ── Renderer/main TypeScript typecheck ───────────────────────────────────── # Pure TS compile (tsc --noEmit) of the Electron main + renderer. No Python / # display server needed, so it's a fast standalone job that gates the PR. diff --git a/RELEASE_0_3_0_PLAN.md b/RELEASE_0_3_0_PLAN.md index 4d5c5a1..699457d 100644 --- a/RELEASE_0_3_0_PLAN.md +++ b/RELEASE_0_3_0_PLAN.md @@ -61,15 +61,29 @@ channels, power law + 2 gaussians, `multifit(optimizer="lm")`: | | | |---|---| | single `fit()` | 17.6 ms | -| `multifit` | **10.83 s** = 10.6 ms/spectrum = **95 spectra/s** | -| extrapolated to 256×256 | **11.6 min** | +| `multifit` (cold, single run) | 10.83 s = 10.6 ms/spectrum = 95 spectra/s | +| `multifit` (best of 3, warm) | **9.1 ms/spectrum = 110 spectra/s** | +| extrapolated to 256×256 | **~10 min** | Single-threaded, one pixel at a time. That is the number to beat. -Free win found while measuring: **`numexpr` is not installed**, so every -hyperspy `Expression` component falls back to numpy (it warns on each model -build). Add it to core deps — costs nothing, speeds up the CPU reference path -and the live single-spectrum preview. +**A trap found while measuring — do not "fix" it.** HyperSpy warns +`Numexpr is not installed, falling back to numpy, which is slower to calculate +model` on every model build. That advice is **wrong at spectrum sizes**, and +following it is a pessimisation: + +| numexpr | threads | spectra/s | +|---|---|---| +| absent | — | **110** | +| 2.14.2 | 1 | 77 | +| 2.14.2 | 4 | 75 | +| 2.14.2 | 48 | 70 | + +(best of 3 warm runs, 256 spectra × 1024 channels). numexpr's per-call setup +and thread dispatch dominate arrays of ~1k elements; it wins on large arrays, +and a single spectrum is not one. It is deliberately **not** a dependency, with +a comment in `pyproject.toml` saying so — the warning will keep inviting +someone to add it. The engine is the `vector_orientation_gpu.py` playbook applied to curve fitting: pack the whole nav grid into `(P, C)` on the GPU and run one batched diff --git a/pyproject.toml b/pyproject.toml index e2e1476..ff0e705 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,13 @@ dependencies = [ "psygnal>=0.10.0", "anyplotlib>=0.4.2", "pyyaml", + # NOT numexpr. HyperSpy warns "Numexpr is not installed ... which is slower + # to calculate model" on every model build, but that advice is wrong at + # spectrum sizes and installing it is a PESSIMISATION: measured on a 1024- + # channel EELS multifit, 110 spectra/s without it vs 77 with it (best of 3; + # 70 at 48 threads — numexpr's per-call setup and thread dispatch dominate + # arrays this small). See RELEASE_0_3_0_PLAN.md §0.4. Do not "fix" the + # warning by adding it back. # GPU-accelerated batched compute (vector orientation mapping, etc.). "torch", # Remote model registry for the SpotUNet disk detector: pulls upgraded @@ -171,6 +178,17 @@ doc = [ "sphinx-gallery", "sphinx_design", ] +# Domain extras (0.3.0). Deliberately NOT core dependencies: each pulls a large +# dependency tree, and the frozen app must stay installable for a user who only +# does 4D-STEM. Actions gate on them with the `requires_package:` toolbar key, +# so a missing extra hides the buttons instead of raising ImportError on click. +# `spyde.data`'s synthetic generators never need them. +eels = ["exspy>=0.3.2"] # EELS + EDS models, edges, quantification +ebsd = ["kikuchipy>=0.11.0"] # EBSD signals, IO, detector, master patterns +atoms = ["atomap>=0.4.1"] # atom finding, sublattices, dumbbells +all = [ + "spyde[eels,ebsd,atoms]", +] [project.urls] "Bug Reports" = "https://github.com/directelectron/spyde/issues" diff --git a/spyde/actions/README.md b/spyde/actions/README.md index fad6d3d..7047f6f 100644 --- a/spyde/actions/README.md +++ b/spyde/actions/README.md @@ -39,6 +39,10 @@ reach `update_live_params`. YAML gate keys: `signal_types` / `exclude_signal_types` (match the signal's `_signal_type` string — covers lazy+eager), `signal_class` (isinstance), `requires_vectors` (hidden until `tree.diffraction_vectors` attaches), +`requires_package` (a module name or list of them — hidden until the optional +extra is installed; resolved with `find_spec`, so it never imports the package +and never pays its import cost. `install_hint(pkg)` gives the `pip install +"spyde[eels]"` line. Add BOTH filter paths' gates — see below), `plot_dim` (`[1, 2]`), `navigation` (navigator vs signal window), `toggle`, `submenu`/`subfunctions`, `parameters` (rendered by the Electron param panel; supports `type`, `default`, `min`/`max`/`step`, `options`, `tab`, @@ -168,6 +172,11 @@ teardown (`_forget_window`), and when a dispatched action raises. `wait_for_vectors(…, strict=True)` (strict = only the clicked plot's tree counts — the any-tree fallback would re-dispatch forever into a tree-specific gate). +- **Two filter paths, one gate**: `get_toolbar_actions_for_plot` (resolves the + function) and `_action_matches_plot` (used by `get_toolbar_config_for_plot`, + imports nothing) apply the SAME gates in two places. Add a new gate key to + both — one alone renders a button that never dispatches, or vice versa. + `test_requires_package_gate.py` asserts this for `requires_package`. - **Bare-figure windows**: a window emitted as a raw `figure` message is NOT a registered Plot — `_plot_by_window_id` returns None and generic dispatch silently no-ops. Register a controller (`own_window`) and resolve via diff --git a/spyde/actions/_template_action.py b/spyde/actions/_template_action.py index ac0c725..4c3a80f 100644 --- a/spyde/actions/_template_action.py +++ b/spyde/actions/_template_action.py @@ -19,6 +19,8 @@ # signal_types: [spyde_diffraction_vectors_image] # _signal_type gate # exclude_signal_types: [...] # blacklist # requires_vectors: True # hide until tree.diffraction_vectors + # requires_package: exspy # hide until the optional extra is + # # installed (or a list: [a, b]) plot_dim: [2] # 1-D / 2-D plots it applies to toolbar_side: bottom navigation: False # navigator vs signal windows diff --git a/spyde/drawing/toolbars/plot_control_toolbar.py b/spyde/drawing/toolbars/plot_control_toolbar.py index edda5d1..7d6b5ff 100644 --- a/spyde/drawing/toolbars/plot_control_toolbar.py +++ b/spyde/drawing/toolbars/plot_control_toolbar.py @@ -32,6 +32,60 @@ def resolve_icon_path(icon_value: str) -> str: return str((base / icon_value).resolve()) +# Optional-extra name for each gated package, so a hidden action can tell the +# user what to install rather than just not existing (0.3.0 Wave 0, #49). +_EXTRA_FOR_PACKAGE = { + "exspy": "eels", + "kikuchipy": "ebsd", + "atomap": "atoms", +} + +_PACKAGE_CACHE: dict[str, bool] = {} + + +def package_available(name: str) -> bool: + """True when an optional dependency is importable. + + Uses ``importlib.util.find_spec`` rather than an actual import: resolving a + toolbar must never pay the multi-second cost of importing exspy/kikuchipy, + and must never execute a broken third-party package's import side effects + just to decide whether to draw a button. + + Cached because the toolbar filter runs on every rebuild, and an extra + cannot appear or vanish inside a session. + """ + hit = _PACKAGE_CACHE.get(name) + if hit is None: + try: + hit = importlib.util.find_spec(name) is not None + except (ImportError, ValueError): + # A namespace-package parent or a half-installed dist can raise + # here; treat anything unresolvable as absent. + hit = False + _PACKAGE_CACHE[name] = hit + return hit + + +def install_hint(name: str) -> str: + """The pip line that would make a gated action appear.""" + extra = _EXTRA_FOR_PACKAGE.get(name) + return f'pip install "spyde[{extra}]"' if extra else f"pip install {name}" + + +def _packages_present(meta: dict) -> bool: + """``requires_package:`` gate — a str or list of importable module names. + + Mirrors ``requires_vectors``: the action is hidden until its optional + extra is installed, instead of appearing and then raising ImportError on + click. + """ + req = meta.get("requires_package") + if not req: + return True + names = [req] if isinstance(req, str) else list(req) + return all(package_available(n) for n in names) + + _SIGNAL_CLASS_CACHE: dict[str, type] = {} @@ -140,6 +194,10 @@ def get_toolbar_actions_for_plot( or isinstance(signal, _resolve_signal_class(signal_class)) ) and (not requires_vectors or has_vectors) + # requires_package: hide until the optional extra is installed + # (exspy / kikuchipy / atomap). Checked by find_spec, so this + # costs nothing and never imports the package. + and _packages_present(meta) and (plot_state.dimensions in plot_dim) and ( navigation_only is None @@ -218,6 +276,7 @@ def _action_matches_plot(action: str, meta: dict, plot_state: "PlotState") -> bo or isinstance(signal, _resolve_signal_class(signal_class)) ) and (not requires_vectors or has_vectors) + and _packages_present(meta) and (plot_state.dimensions in plot_dim) and ( navigation_only is None diff --git a/spyde/tests/migrated/test_requires_package_gate.py b/spyde/tests/migrated/test_requires_package_gate.py new file mode 100644 index 0000000..e49a900 --- /dev/null +++ b/spyde/tests/migrated/test_requires_package_gate.py @@ -0,0 +1,159 @@ +"""The `requires_package:` toolbar gate (0.3.0 Wave 0, #49). + +Mirrors `requires_vectors`: an action declaring an optional extra stays hidden +until that extra is importable, so a user without `spyde[eels]` never sees a +button that would raise ImportError when clicked. + +Both filter paths are covered. `get_toolbar_actions_for_plot` (which resolves +and imports the action function) and `get_toolbar_config_for_plot` / +`_action_matches_plot` (which must NOT import anything) apply the same gates, +and it is easy to add a gate to one and forget the other. +""" +from __future__ import annotations + +import types + +import pytest + +from spyde.drawing.toolbars import plot_control_toolbar as pct +from spyde.drawing.toolbars.plot_control_toolbar import ( + _action_matches_plot, + _packages_present, + install_hint, + package_available, +) + + +@pytest.fixture(autouse=True) +def _clear_package_cache(): + """package_available() memoises; a stale entry would leak between tests.""" + pct._PACKAGE_CACHE.clear() + yield + pct._PACKAGE_CACHE.clear() + + +class _FakePlot: + """Minimal stand-in for the toolbar filter: it reads ``_signal_type`` off + the current signal and ``diffraction_vectors`` off the tree.""" + + def __init__(self, signal_type: str = ""): + signal = types.SimpleNamespace(_signal_type=signal_type) + self.plot_state = types.SimpleNamespace( + current_signal=signal, dimensions=2, plot=self) + self.signal_tree = types.SimpleNamespace(diffraction_vectors=None, + root=None) + self.is_navigator = False + + +class TestPackageAvailable: + def test_detects_an_installed_package(self): + assert package_available("numpy") is True + + def test_detects_a_missing_package(self): + assert package_available("spyde_definitely_not_installed") is False + + def test_result_is_cached(self): + package_available("numpy") + assert pct._PACKAGE_CACHE["numpy"] is True + + def test_does_not_import_the_package(self, monkeypatch): + """find_spec, not import: resolving a toolbar must never pay a + multi-second import or run a third-party package's side effects just to + decide whether to draw a button.""" + import builtins + real_import = builtins.__import__ + + def _boom(name, *a, **k): + if name == "numexpr": + raise AssertionError("requires_package imported the package") + return real_import(name, *a, **k) + + monkeypatch.setattr(builtins, "__import__", _boom) + assert package_available("numexpr") in (True, False) + + def test_unresolvable_package_is_absent_not_an_error(self, monkeypatch): + """A half-installed dist or a namespace-package parent makes find_spec + raise; that must read as 'absent', not crash the whole toolbar.""" + def _raise(name): + raise ValueError("mangled dist-info") + + monkeypatch.setattr(pct.importlib.util, "find_spec", _raise) + assert package_available("anything") is False + + +class TestPackagesPresent: + def test_no_declaration_is_always_allowed(self): + assert _packages_present({}) is True + assert _packages_present({"requires_package": None}) is True + + def test_single_name_as_a_string(self): + assert _packages_present({"requires_package": "numpy"}) is True + assert _packages_present({"requires_package": "spyde_nope"}) is False + + def test_list_requires_all_of_them(self): + assert _packages_present({"requires_package": ["numpy", "scipy"]}) is True + assert _packages_present( + {"requires_package": ["numpy", "spyde_nope"]}) is False + + +class TestInstallHint: + @pytest.mark.parametrize("pkg,extra", [("exspy", "eels"), + ("kikuchipy", "ebsd"), + ("atomap", "atoms")]) + def test_maps_each_gated_package_to_its_extra(self, pkg, extra): + assert install_hint(pkg) == f'pip install "spyde[{extra}]"' + + def test_unknown_package_falls_back_to_a_plain_pip_line(self): + assert install_hint("mystery") == "pip install mystery" + + def test_every_declared_extra_exists_in_pyproject(self): + """The hint has to be runnable — an extra named here but absent from + pyproject would tell the user to run a command that fails.""" + import re + from pathlib import Path + import spyde + + root = Path(spyde.__file__).resolve().parent.parent + text = (root / "pyproject.toml").read_text(encoding="utf-8") + block = text.split("[project.optional-dependencies]", 1)[1] + block = block.split("\n[", 1)[0] + declared = set(re.findall(r"^(\w[\w-]*)\s*=\s*\[", block, re.M)) + for extra in pct._EXTRA_FOR_PACKAGE.values(): + assert extra in declared, f"extra {extra!r} missing from pyproject" + + +class TestFilterIntegration: + """The gate is applied by BOTH filter paths, not just the one.""" + + def test_action_matches_plot_hides_a_missing_package(self): + plot = _FakePlot() + meta = {"requires_package": "spyde_nope", "plot_dim": [1, 2]} + assert _action_matches_plot("X", meta, plot.plot_state) is False + + def test_action_matches_plot_allows_a_present_package(self): + plot = _FakePlot() + meta = {"requires_package": "numpy", "plot_dim": [1, 2]} + assert _action_matches_plot("X", meta, plot.plot_state) is True + + def test_gate_is_wired_into_both_filters(self): + """Guards the real failure mode: adding the gate to one filter and + forgetting the other, so the button renders but does not dispatch (or + vice versa).""" + import inspect + src = inspect.getsource(pct) + body = src.split("def get_toolbar_actions_for_plot", 1)[1] + listing, matching = body.split("def _action_matches_plot", 1) + assert "_packages_present" in listing, \ + "get_toolbar_actions_for_plot is missing the requires_package gate" + assert "_packages_present" in matching, \ + "_action_matches_plot is missing the requires_package gate" + + +class TestRealToolbarUnaffected: + def test_no_current_action_is_accidentally_gated(self): + """No shipped action declares requires_package yet, so the whole + existing toolbar must be untouched by this change.""" + from spyde import TOOLBAR_ACTIONS + gated = {name for name, meta in TOOLBAR_ACTIONS["functions"].items() + if meta.get("requires_package")} + assert gated == set() From 72e2f0c842d5d346508f10f12d45a861f4149a14 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Sun, 26 Jul 2026 21:53:33 -0500 Subject: [PATCH 04/60] feat(fitting): batched GPU model fitting (0.3.0 Wave 1.1-1.3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- spyde/fitting/__init__.py | 39 ++ spyde/fitting/components.py | 261 ++++++++++++ spyde/fitting/engine.py | 342 ++++++++++++++++ spyde/fitting/spec.py | 375 ++++++++++++++++++ spyde/tests/benchmark_fitting.py | 143 +++++++ spyde/tests/migrated/test_fitting_engine.py | 287 ++++++++++++++ spyde/tests/migrated/test_model_spec.py | 231 +++++++++++ spyde/tests/migrated/test_torch_components.py | 244 ++++++++++++ 8 files changed, 1922 insertions(+) create mode 100644 spyde/fitting/__init__.py create mode 100644 spyde/fitting/components.py create mode 100644 spyde/fitting/engine.py create mode 100644 spyde/fitting/spec.py create mode 100644 spyde/tests/benchmark_fitting.py create mode 100644 spyde/tests/migrated/test_fitting_engine.py create mode 100644 spyde/tests/migrated/test_model_spec.py create mode 100644 spyde/tests/migrated/test_torch_components.py diff --git a/spyde/fitting/__init__.py b/spyde/fitting/__init__.py new file mode 100644 index 0000000..280381c --- /dev/null +++ b/spyde/fitting/__init__.py @@ -0,0 +1,39 @@ +"""spyde.fitting — HyperSpy model fitting, made fast. + +**Not to be confused with** :mod:`spyde.models`, which is the neural +disk-detector package (SpotUNet + the Hugging Face weight registry). Unrelated +meanings of the word "model"; do not merge them. + +Layout: + +``spec`` + :class:`~spyde.fitting.spec.ModelSpec` — a serialisable model description + that round-trips against HyperSpy's ``BaseModel.as_dictionary()``. The + contract shared by HyperSpy, the batched engine and the UI. + +``components`` + Batched, autograd-differentiable torch ports of HyperSpy's components, + each evaluating every nav position at once. + +``engine`` + The batched Levenberg-Marquardt fit over the whole navigation grid. + +Why this exists: HyperSpy's ``multifit`` fits one pixel at a time, measured at +~110 spectra/s on this box — about 10 minutes for a 256x256 spectrum image. +The engine packs the whole grid into ``(P, C)`` and solves it as one batched +problem, following the same playbook as +``spyde/actions/vector_orientation_gpu.py``. + +Correctness is defined by HyperSpy: the engine must reproduce ``multifit``'s +parameters on the same data. See ``RELEASE_0_3_0_PLAN.md`` and GitHub #50. +""" +from __future__ import annotations + +from spyde.fitting.spec import ( + ComponentSpec, + ModelSpec, + ParameterSpec, + spec_from_component, +) + +__all__ = ["ModelSpec", "ComponentSpec", "ParameterSpec", "spec_from_component"] diff --git a/spyde/fitting/components.py b/spyde/fitting/components.py new file mode 100644 index 0000000..263eb6d --- /dev/null +++ b/spyde/fitting/components.py @@ -0,0 +1,261 @@ +"""components.py — batched, differentiable torch ports of HyperSpy's components. + +Each component evaluates **every navigation position at once**: given the shared +signal axis ``x`` of length ``C`` and a parameter block ``p`` of shape +``(P, n)``, it returns ``(P, C)``. There is no Python loop over pixels anywhere +in this module — that is the entire point (see :mod:`spyde.fitting.engine`). + +Three rules, each of which is a parity bug if broken: + +* **Parameter order is HyperSpy's order**, i.e. the order of + ``component.parameters``. It is *not* alphabetical and *not* the order of the + constructor arguments: ``PowerLaw`` is ``(A, left_cutoff, origin, r)`` and + ``GaussianHF`` is ``(centre, fwhm, height)``. ``ModelSpec`` packs columns in + this order, so a mismatch silently fits the wrong parameter. +* **The formula is HyperSpy's formula.** These are transcribed from the + ``expression=`` strings in ``hyperspy/_components/``, not from memory — e.g. + a HyperSpy ``Gaussian``'s ``A`` is the AREA, not the peak height. + ``test_torch_components.py`` checks every component against the HyperSpy one + numerically; that test is the real specification. +* **Everything must stay differentiable.** No ``.item()``, no in-place writes + into a tensor that requires grad, and no branch that produces NaN on the dead + side of a ``where`` — a NaN there poisons the gradient even where the mask + discards the value. + +torch is a core dependency, but it is imported lazily so that importing +``spyde.fitting`` (for ``ModelSpec``) never pays for it. +""" +from __future__ import annotations + +import logging +import math +from typing import Callable, Sequence + +log = logging.getLogger(__name__) + +_SQRT_2PI = math.sqrt(2.0 * math.pi) +_SQRT_2 = math.sqrt(2.0) +_FOUR_LOG2 = 4.0 * math.log(2.0) + +# Guards a division by a parameter the optimiser can walk to zero. +_EPS = 1e-30 + + +class TorchComponent: + """One batched component. + + ``params`` mirrors HyperSpy's parameter order; ``linear`` marks the columns + the model is linear in (what variable projection solves directly). + """ + + __slots__ = ("kind", "params", "linear", "_fn") + + def __init__(self, kind: str, params: Sequence[str], linear: Sequence[bool], + fn: Callable): + if len(params) != len(linear): + raise ValueError(f"{kind}: {len(params)} params vs " + f"{len(linear)} linear flags") + self.kind = kind + self.params = tuple(params) + self.linear = tuple(bool(b) for b in linear) + self._fn = fn + + @property + def n_params(self) -> int: + return len(self.params) + + def __call__(self, x, p): + """``x``: ``(C,)`` or ``(P, C)``. ``p``: ``(P, n_params)``. -> ``(P, C)``.""" + if p.shape[-1] != self.n_params: + raise ValueError(f"{self.kind} expects {self.n_params} parameters " + f"{self.params}, got {p.shape[-1]}") + if x.dim() == 1: + x = x.unsqueeze(0) # (1, C), broadcasts over P + return self._fn(x, [p[:, i:i + 1] for i in range(self.n_params)]) + + def __repr__(self) -> str: # pragma: no cover + return f"" + + +# --------------------------------------------------------------------------- +# the formulas — transcribed from hyperspy/_components/*.py `expression=` +# --------------------------------------------------------------------------- + +def _gaussian(x, p): + # "A * (1 / (sigma * sqrt(2*pi))) * exp(-(x - centre)**2 / (2 * sigma**2))" + # NB A is the AREA under the curve, not the height. + A, centre, sigma = p + s = sigma + _EPS + return A / (s * _SQRT_2PI) * (-((x - centre) ** 2) / (2 * s ** 2)).exp() + + +def _gaussian_hf(x, p): + # "height * exp(-(x - centre)**2 * 4 * log(2)/fwhm**2)" + centre, fwhm, height = p + w = fwhm + _EPS + return height * (-((x - centre) ** 2) * _FOUR_LOG2 / w ** 2).exp() + + +def _lorentzian(x, p): + # "A / pi * (gamma_ / ((x - centre)**2 + gamma_**2))" (gamma_ renames gamma) + A, centre, gamma = p + return A / math.pi * (gamma / ((x - centre) ** 2 + gamma ** 2 + _EPS)) + + +def _power_law(x, p): + # "where(left_cutoff < x, A*(-origin + x)**-r, 0)" + A, left_cutoff, origin, r = p + import torch + base = x - origin + live = (left_cutoff < x) & (base > 0) + # Clamp the base BEFORE the power: (<=0) ** -r is inf/NaN, and a NaN in the + # discarded branch of a where() still propagates NaN through the gradient. + safe = torch.where(live, base, torch.ones_like(base)) + return torch.where(live, A * safe ** (-r), torch.zeros_like(base)) + + +def _offset(x, p): + (offset,) = p + return offset.expand(-1, x.shape[-1]) if x.shape[0] == 1 else \ + offset + 0.0 * x + + +def _exponential(x, p): + # "A * exp(-x / tau)" + A, tau = p + return A * (-x / (tau + _EPS)).exp() + + +def _arctan(x, p): + # "A * atan(k * (x - x0))" + A, k, x0 = p + return A * (k * (x - x0)).atan() + + +def _erf(x, p): + # "A * erf((x - origin) / sqrt(2) / sigma) / 2" + A, origin, sigma = p + import torch + return A * torch.erf((x - origin) / _SQRT_2 / (sigma + _EPS)) / 2.0 + + +def _heaviside(x, p): + # HeavisideStep: A below/above the step at n (0.5 exactly at the step). + A, n = p + import torch + return A * torch.where(x > n, torch.ones_like(x), + torch.where(x < n, torch.zeros_like(x), + torch.full_like(x, 0.5))) + + +def _logistic(x, p): + # "a / (1 + b * exp(-c * (x - origin)))" + a, b, c, origin = p + return a / (1.0 + b * (-c * (x - origin)).exp()) + + +def _polynomial_fn(order: int): + """Polynomial is variable-order, so its evaluator is built per order. + Parameters are ``a0..a{order}`` and ``aK`` multiplies ``x**K``.""" + + def fn(x, p): + out = p[0].expand_as(x) if len(p) else None + acc = out + for k in range(1, order + 1): + acc = acc + p[k] * x ** k + return acc + + return fn + + +# --------------------------------------------------------------------------- +# registry +# --------------------------------------------------------------------------- + +_REGISTRY: dict[str, TorchComponent] = {} + + +def _register(kind, params, linear, fn) -> None: + _REGISTRY[kind] = TorchComponent(kind, params, linear, fn) + + +# Parameter tuples are HyperSpy's `component.parameters` ORDER — verified by +# test_torch_components.py::TestParameterOrder against live components. +_register("Gaussian", ("A", "centre", "sigma"), (True, False, False), _gaussian) +_register("GaussianHF", ("centre", "fwhm", "height"), (False, False, True), _gaussian_hf) +_register("Lorentzian", ("A", "centre", "gamma"), (True, False, False), _lorentzian) +_register("PowerLaw", ("A", "left_cutoff", "origin", "r"), + (True, False, False, False), _power_law) +_register("Offset", ("offset",), (True,), _offset) +_register("Exponential", ("A", "tau"), (True, False), _exponential) +_register("Arctan", ("A", "k", "x0"), (True, False, False), _arctan) +_register("Erf", ("A", "origin", "sigma"), (True, False, False), _erf) +_register("HeavisideStep", ("A", "n"), (True, False), _heaviside) +_register("Logistic", ("a", "b", "c", "origin"), (True, False, False, False), _logistic) + + +def get_component(kind: str, *, n_params: int | None = None) -> TorchComponent: + """Look up a batched component by HyperSpy ``_id_name``. + + ``Polynomial`` is variable-order, so its evaluator is built on demand from + the parameter count the spec carries. + """ + if kind == "Polynomial": + if n_params is None: + raise ValueError("Polynomial needs n_params (order + 1)") + order = int(n_params) - 1 + return TorchComponent("Polynomial", + tuple(f"a{k}" for k in range(order + 1)), + (True,) * (order + 1), + _polynomial_fn(order)) + try: + return _REGISTRY[kind] + except KeyError: + raise NotImplementedError( + f"no batched torch implementation for component {kind!r}. " + f"Available: {sorted(_REGISTRY) + ['Polynomial']}. " + f"Fall back to the CPU/HyperSpy path, or add it here with a parity " + f"test (#52)." + ) from None + + +def available() -> list[str]: + """Component kinds the batched engine can fit.""" + return sorted(_REGISTRY) + ["Polynomial"] + + +def supports(spec) -> bool: + """True when every ACTIVE component of a ModelSpec has a batched port. + + The engine calls this to decide whether it can run at all; anything + unsupported falls back to HyperSpy's own fitting rather than silently + dropping a component from the model. + """ + for c in getattr(spec, "active_components", []): + try: + get_component(c.kind, n_params=len(c.parameters)) + except NotImplementedError: + return False + return True + + +def evaluate(spec, x, values): + """Evaluate a whole :class:`~spyde.fitting.spec.ModelSpec`. + + ``x``: ``(C,)``; ``values``: ``(P, n_total)`` packed in the spec's order. + Returns ``(P, C)`` — the sum over active components. + """ + import torch + + out = None + i = 0 + for c in spec.active_components: + n = len(c.parameters) + comp = get_component(c.kind, n_params=n) + y = comp(x, values[:, i:i + n]) + out = y if out is None else out + y + i += n + if out is None: + p = values.shape[0] if values.dim() > 1 else 1 + return torch.zeros((p, x.shape[-1]), dtype=x.dtype, device=x.device) + return out diff --git a/spyde/fitting/engine.py b/spyde/fitting/engine.py new file mode 100644 index 0000000..cf25a15 --- /dev/null +++ b/spyde/fitting/engine.py @@ -0,0 +1,342 @@ +"""engine.py — batched Levenberg-Marquardt over the whole navigation grid. + +HyperSpy's ``multifit`` fits one pixel at a time: measured at ~110 spectra/s on +this box, i.e. ~10 minutes for a 256x256 spectrum image. This module fits every +pixel *simultaneously*, following the playbook already proven in +``spyde/actions/vector_orientation_gpu.py``. + +Why it works — the shape argument, which is the whole design: + + P navigation positions (10^3 - 10^5) + C channels per spectrum (10^3 - 10^4) + n free parameters (3 - 20) <- tiny, and that is the point + +Levenberg-Marquardt solves ``(JᵀJ + λ·diag(JᵀJ)) δ = -Jᵀr`` each step. ``JᵀJ`` is +``(P, n, n)`` — a batch of *tiny* linear systems, which ``torch.linalg`` solves +natively in one call. The only large object is the Jacobian ``(P, C, n)``, so +the batch dimension is **chunked** to bound it (65536 x 2048 x 12 float32 = +6.4 GB whole; chunked it is a fixed working set). + +The Jacobian comes from ``torch.func.jacfwd``, not ``jacrev``: forward-mode +costs one pass per *input* (n ≈ 10) where reverse-mode costs one per *output* +(C ≈ 2000). For this shape that is a ~100x difference. + +**There is one implementation, not two.** The "CPU fallback" is this same code +with ``device="cpu"`` — torch runs everywhere. A separate scipy path would be a +second thing to keep in agreement with HyperSpy, and HyperSpy already *is* the +reference: correctness here means reproducing ``multifit``'s parameters, which +``test_fitting_engine.py`` asserts directly. +""" +from __future__ import annotations + +import logging +import math +from dataclasses import dataclass +from typing import Callable + +import numpy as np + +from spyde.fitting import components as tcomp + +log = logging.getLogger(__name__) + +# Bounds the Jacobian working set per chunk. 2**22 elements x 8 B (float64) +# x (value + tangents) lands comfortably inside a few hundred MB. +_TARGET_JACOBIAN_ELEMENTS = 1 << 22 + + +@dataclass +class FitResult: + """Outcome of a batched fit. + + ``values`` is ``(P, n_total)`` in the spec's packed order — the same order + as :meth:`~spyde.fitting.spec.ModelSpec.parameter_names`, so a component + map is a column slice (``spec.component_slices()``). + """ + + values: np.ndarray # (P, n_total) + converged: np.ndarray # (P,) bool + chisq: np.ndarray # (P,) sum of squared (weighted) residuals + n_iter: int + device: str + + @property + def convergence_rate(self) -> float: + return float(self.converged.mean()) if self.converged.size else 0.0 + + def as_maps(self, spec, nav_shape) -> dict[str, np.ndarray]: + """Per-parameter maps, reshaped to the navigation grid and keyed by + ``"."``.""" + names = spec.parameter_names() + return {name: self.values[:, i].reshape(nav_shape) + for i, name in enumerate(names)} + + +def default_device() -> str: + try: + import torch + if torch.cuda.is_available(): + return "cuda" + except Exception as e: # pragma: no cover + log.debug("probing CUDA failed: %s", e) + return "cpu" + + +def _selection_matrix(free_mask, n_total, dtype, device): + """``S`` such that ``full = fixed + S @ p_free``. + + Written as a matmul rather than index assignment because the whole step + runs under ``vmap``/``jacfwd``, where in-place writes into a tensor that + carries tangents are at best fragile. ``n`` is tiny, so the matmul is free. + """ + import torch + idx = torch.nonzero(free_mask, as_tuple=False).squeeze(-1) + S = torch.zeros((n_total, idx.numel()), dtype=dtype, device=device) + S[idx, torch.arange(idx.numel(), device=device)] = 1.0 + return S + + +def _make_residual_fn(spec, x, S): + """Build ``residual(p_free, y_i, w_i, fixed_i) -> (C,)`` for ONE position. + + Everything that varies per position is an explicit argument so the caller + can ``vmap`` over all four; only the model structure and the signal axis + are closed over. + """ + def residual(p_free, y_i, w_i, fixed_i): + full = fixed_i + (S @ p_free) + model = tcomp.evaluate(spec, x, full.unsqueeze(0)).squeeze(0) + return (model - y_i) * w_i + + return residual + + +def fit_batched(spec, data, x, *, weights=None, device=None, max_iter=60, + ftol=1e-8, xtol=1e-8, gtol=1e-8, chunk=None, dtype="float64", + progress: Callable[[int, int], None] | None = None, + initial=None): + """Fit *spec* to every spectrum in *data*, all at once. + + Parameters + ---------- + spec : ModelSpec + The model. Every ACTIVE component must have a batched port + (``components.supports(spec)``); otherwise fall back to HyperSpy. + data : array (P, C) or (..., C) + Spectra. Leading dimensions are flattened to ``P``. + x : array (C,) + The signal axis (calibrated units — the same values HyperSpy fits in, + not channel indices, or every centre/onset comes out in the wrong unit). + weights : array (C,) or (P, C), optional + Residual weights. ``"poisson"`` gives ``1/sqrt(max(y, 1))``, which is + what counting data wants (#53). + initial : array (P, n_total), optional + Per-position starting values — the seeded-propagation hand-off (#54). + Defaults to the spec's values broadcast to every position. + + Returns + ------- + FitResult + """ + import torch + + spec_names = spec.parameter_names() + n_total = len(spec_names) + if n_total == 0: + raise ValueError("model has no active parameters to fit") + if not tcomp.supports(spec): + unsupported = [c.kind for c in spec.active_components + if c.kind not in tcomp.available()] + raise NotImplementedError( + f"no batched implementation for {unsupported}; use the HyperSpy " + f"path for this model (see components.supports)") + + device = device or default_device() + tdtype = getattr(torch, dtype) + + data = np.asarray(data) + nav_shape = data.shape[:-1] + y_np = data.reshape(-1, data.shape[-1]) + P, C = y_np.shape + if len(x) != C: + raise ValueError(f"signal axis has {len(x)} points but data has {C}") + + # --- weights ---------------------------------------------------------- + if isinstance(weights, str): + if weights != "poisson": + raise ValueError(f"unknown weighting {weights!r}") + w_np = 1.0 / np.sqrt(np.maximum(y_np, 1.0)) + elif weights is None: + w_np = np.ones((1, C)) + else: + w_np = np.asarray(weights, float) + if w_np.ndim == 1: + w_np = w_np[None, :] + + # --- signal range: a masked channel simply gets zero weight ----------- + if spec.channel_mask is not None: + mask = np.asarray(spec.channel_mask, bool) + if mask.size != C: + raise ValueError(f"channel mask has {mask.size} entries, " + f"signal axis has {C}") + w_np = w_np * mask[None, :] + + free_mask_np = spec.free_mask() + n_free = int(free_mask_np.sum()) + if n_free == 0: + raise ValueError("every parameter is fixed — nothing to fit") + + lo_np, hi_np = spec.bounds_arrays() + start_np = (np.broadcast_to(spec.flat_values(), (P, n_total)).copy() + if initial is None else + np.asarray(initial, float).reshape(P, n_total).copy()) + # A start outside its own bounds makes the first step meaningless. + start_np = np.clip(start_np, lo_np, hi_np) + + if chunk is None: + chunk = max(1, min(P, _TARGET_JACOBIAN_ELEMENTS // max(C * n_free, 1))) + + x_t = torch.as_tensor(np.asarray(x, float), dtype=tdtype, device=device) + free_t = torch.as_tensor(free_mask_np, device=device) + S = _selection_matrix(free_t, n_total, tdtype, device) + lo_t = torch.as_tensor(lo_np[free_mask_np], dtype=tdtype, device=device) + hi_t = torch.as_tensor(hi_np[free_mask_np], dtype=tdtype, device=device) + + out_values = np.empty((P, n_total), np.float64) + out_conv = np.zeros(P, bool) + out_chisq = np.full(P, np.inf) + iters_used = 0 + + for lo_i in range(0, P, chunk): + hi_i = min(P, lo_i + chunk) + v, c, q, it = _fit_chunk( + spec, x_t, S, free_t, tdtype, device, + y=torch.as_tensor(y_np[lo_i:hi_i], dtype=tdtype, device=device), + w=torch.as_tensor(w_np[lo_i:hi_i] if w_np.shape[0] > 1 else w_np, + dtype=tdtype, device=device), + start=torch.as_tensor(start_np[lo_i:hi_i], dtype=tdtype, device=device), + lo=lo_t, hi=hi_t, max_iter=max_iter, ftol=ftol, xtol=xtol, + gtol=gtol, + ) + out_values[lo_i:hi_i] = v + out_conv[lo_i:hi_i] = c + out_chisq[lo_i:hi_i] = q + iters_used = max(iters_used, it) + if progress is not None: + try: + progress(hi_i, P) + except Exception as e: # pragma: no cover + log.debug("fit progress callback failed: %s", e) + + return FitResult(values=out_values, converged=out_conv, chisq=out_chisq, + n_iter=iters_used, device=device) + + +def _fit_chunk(spec, x, S, free_mask, tdtype, device, *, y, w, start, lo, hi, + max_iter, ftol, xtol, gtol): + """One chunk of positions through the LM loop. All tensors already on + *device*; returns numpy.""" + import torch + from torch.func import jacfwd, vmap + + n_chunk = y.shape[0] + if w.shape[0] == 1: + w = w.expand(n_chunk, -1) + + # Fixed parameters contribute a constant offset; free ones come via S. + fixed_full = start * (~free_mask).to(tdtype) + p = start[:, free_mask].clone() + + residual = _make_residual_fn(spec, x, S) + # in_dims: p_free batched, and y/w/fixed batched alongside it. + res_b = vmap(residual, in_dims=(0, 0, 0, 0)) + jac_b = vmap(jacfwd(residual, argnums=0), in_dims=(0, 0, 0, 0)) + + lam = torch.full((n_chunk,), 1e-3, dtype=tdtype, device=device) + r = res_b(p, y, w, fixed_full) + cost = 0.5 * (r * r).sum(1) + converged = torch.zeros(n_chunk, dtype=torch.bool, device=device) + used = 0 + + eye = torch.eye(p.shape[1], dtype=tdtype, device=device) + + for it in range(max_iter): + used = it + 1 + J = jac_b(p, y, w, fixed_full) # (B, C, n_free) + + # COLUMN SCALING (MINPACK's `diag`), and it is load-bearing, not a + # refinement. Parameters routinely differ by many orders of magnitude — + # a PowerLaw fits A ~ 1e6 alongside r ~ 3 — which makes cond(J) ~ 1e6 + # and cond(JᵀJ) ~ 1e12. Cholesky in float64 then loses almost every + # digit in the small-eigenvalue direction, so the Gauss-Newton step is + # numerically junk and the fit crawls: measured 247 iterations to fit a + # TWO-parameter power law (and it stopped 1.5-3.5% short of the value + # multifit finds), versus 46 iterations to chisq ~1e-34 once the + # columns are normalised. + # Normalising each column to unit norm before forming JᵀJ removes the + # scale entirely (and makes diag(JᵀJ) ≈ 1, so the Marquardt damping + # below is automatically scale-free). + colnorm = J.norm(dim=1).clamp_min(1e-300) # (B, n) + Js = J / colnorm.unsqueeze(1) + Jst = Js.transpose(1, 2) + JtJ = Jst @ Js # (B, n, n) — tiny + Jtr = (Jst @ r.unsqueeze(2)).squeeze(2) # (B, n) + + diag = torch.diagonal(JtJ, dim1=1, dim2=2) + A = JtJ + lam[:, None, None] * torch.diag_embed( + diag.clamp_min(1e-12)) + 1e-14 * eye + + # cholesky_ex reports failure per item instead of raising, so one + # singular pixel cannot abort the whole batch. + L, info = torch.linalg.cholesky_ex(A) + solvable = info == 0 + delta = torch.zeros_like(p) + if solvable.any(): + sol = torch.cholesky_solve((-Jtr).unsqueeze(2)[solvable], + L[solvable]).squeeze(2) + delta[solvable] = sol + delta = delta / colnorm # back to real units + + p_new = torch.clamp(p + delta, lo, hi) # bounds by projection + r_new = res_b(p_new, y, w, fixed_full) + cost_new = 0.5 * (r_new * r_new).sum(1) + + better = (cost_new < cost) & solvable + # Accepted: keep the step and trust the model more (smaller λ). + # Rejected: fall back toward gradient descent (larger λ). + p = torch.where(better[:, None], p_new, p) + r = torch.where(better[:, None], r_new, r) + rel = (cost - cost_new) / cost.clamp_min(1e-300) + cost = torch.where(better, cost_new, cost) + lam = torch.where(better, (lam / 3.0).clamp_min(1e-12), + (lam * 3.0).clamp_max(1e12)) + + # Convergence. The GRADIENT test is the primary one, and it has to be: + # + # * A step-size test alone is wrong on a REJECTED step — lambda grows, + # the next delta shrinks, and "stuck against a wall" reads as + # "converged" (measured: a PowerLaw stopping 1.5-3.5% short). + # * But requiring an ACCEPTED step is *also* wrong, because AT the + # optimum no step improves anything, so `better` is permanently + # False and neither test can ever fire — the fit then burns every + # iteration of its budget while already converged (measured: 8% of + # pixels "converged" on a fit that was actually finished). + # + # The gradient does not care whether the last step was taken. Columns + # of J are unit-norm here (see the scaling above), so `Jᵀr` is the + # scaled gradient and `|Jᵀr| / ||r||` is dimensionless — a real + # relative tolerance rather than a magic absolute number. + rnorm = r.norm(dim=1).clamp_min(1e-300) + grad_small = (Jtr.abs().amax(1) / rnorm) < gtol + step_small = better & (delta.abs() <= xtol * (p.abs() + xtol)).all(1) + cost_small = better & (rel.abs() < ftol) + converged = converged | ((grad_small | step_small | cost_small) & solvable) + if bool(converged.all()): + break + + # Reassemble the full parameter vector (free values back into their slots). + full = fixed_full + (S @ p.unsqueeze(2)).squeeze(2) + return (full.detach().cpu().numpy().astype(np.float64), + converged.detach().cpu().numpy(), + (2.0 * cost).detach().cpu().numpy().astype(np.float64), + used) diff --git a/spyde/fitting/spec.py b/spyde/fitting/spec.py new file mode 100644 index 0000000..aba1721 --- /dev/null +++ b/spyde/fitting/spec.py @@ -0,0 +1,375 @@ +"""spec.py — a serialisable description of a HyperSpy model. + +``ModelSpec`` is the contract between the three things that need to agree about +what "the model" is: + +* **HyperSpy** — the reference implementation and the storage format. A spec + round-trips through ``BaseModel.as_dictionary()``, so a model built in SpyDE + opens in plain HyperSpy and vice versa, and ``m.store()`` / ``s.models.restore()`` + keep working unchanged. +* **The batched engine** (:mod:`spyde.fitting.engine`) — which needs the + parameters as flat, ordered arrays it can pack into ``(P, n)`` tensors. +* **The UI** — which adds components line by line, toggles them, and needs a + JSON-safe form to send to the renderer. + +Deliberately a plain dataclass tree, not a HyperSpy subclass: the engine must +be able to describe a model without a signal attached (the wizard builds one +before it runs), and the renderer needs it as JSON. + +**Parameter order is part of the contract.** ``flat_values()`` and everything +derived from it iterate components in order and parameters within a component +in order, so column *j* of a packed tensor always means the same parameter. +:meth:`ModelSpec.parameter_names` is that order made explicit. +""" +from __future__ import annotations + +import logging +from dataclasses import dataclass, field, replace +from typing import Any, Iterator, Sequence + +import numpy as np + +log = logging.getLogger(__name__) + +# HyperSpy stores an absent bound as None; the engine wants finite arrays. +_NEG_INF, _POS_INF = -np.inf, np.inf + + +@dataclass +class ParameterSpec: + """One fittable parameter. + + ``linear`` mirrors HyperSpy's ``_linear`` flag: the model is linear in this + parameter (a Gaussian's ``A``, a PowerLaw's ``A``), so the engine can solve + for it by least squares instead of iterating on it. That is what makes + variable projection possible (#53) and what makes tabulated EELS edges + tractable (#63). + """ + + name: str + value: float = 0.0 + free: bool = True + bmin: float | None = None + bmax: float | None = None + linear: bool = False + units: str = "" + + def bounds(self) -> tuple[float, float]: + """Finite bounds for the engine (``None`` becomes ±inf).""" + lo = _NEG_INF if self.bmin is None else float(self.bmin) + hi = _POS_INF if self.bmax is None else float(self.bmax) + return lo, hi + + def to_dict(self) -> dict[str, Any]: + return {"name": self.name, "value": float(self.value), + "free": bool(self.free), "bmin": self.bmin, "bmax": self.bmax, + "linear": bool(self.linear), "units": self.units} + + @classmethod + def from_dict(cls, d: dict[str, Any]) -> "ParameterSpec": + return cls(name=d["name"], value=float(d.get("value", 0.0)), + free=bool(d.get("free", True)), bmin=d.get("bmin"), + bmax=d.get("bmax"), linear=bool(d.get("linear", False)), + units=d.get("units", "")) + + +@dataclass +class ComponentSpec: + """One model component. + + ``kind`` is HyperSpy's ``_id_name`` (``"Gaussian"``, ``"PowerLaw"``, …) — + the key both the HyperSpy factory and the torch component library resolve + against. ``name`` is the user-facing label, which may be edited and is not + unique. + """ + + kind: str + name: str = "" + active: bool = True + parameters: list[ParameterSpec] = field(default_factory=list) + + def __post_init__(self): + if not self.name: + self.name = self.kind + + def __getitem__(self, name: str) -> ParameterSpec: + for p in self.parameters: + if p.name == name: + return p + raise KeyError(f"{self.kind} has no parameter {name!r} " + f"(has {[p.name for p in self.parameters]})") + + def to_dict(self) -> dict[str, Any]: + return {"kind": self.kind, "name": self.name, "active": bool(self.active), + "parameters": [p.to_dict() for p in self.parameters]} + + @classmethod + def from_dict(cls, d: dict[str, Any]) -> "ComponentSpec": + return cls(kind=d["kind"], name=d.get("name", ""), + active=bool(d.get("active", True)), + parameters=[ParameterSpec.from_dict(p) + for p in d.get("parameters", [])]) + + +@dataclass +class ModelSpec: + """A whole model: ordered components plus the fitted channel range. + + ``channel_mask`` is HyperSpy's ``_channel_switches`` — a boolean array over + the signal axis, True where the channel takes part in the fit. ``None`` + means "all channels", which is the common case and avoids carrying an + array the size of the signal axis around in JSON. + """ + + components: list[ComponentSpec] = field(default_factory=list) + channel_mask: np.ndarray | None = None + + # -- construction ------------------------------------------------------ + def append(self, comp: ComponentSpec) -> "ModelSpec": + self.components.append(comp) + return self + + def copy(self) -> "ModelSpec": + return replace( + self, + components=[ComponentSpec.from_dict(c.to_dict()) + for c in self.components], + channel_mask=None if self.channel_mask is None + else np.array(self.channel_mask, copy=True), + ) + + def __getitem__(self, key: int | str) -> ComponentSpec: + if isinstance(key, int): + return self.components[key] + for c in self.components: + if c.name == key: + return c + raise KeyError(f"no component named {key!r} " + f"(have {[c.name for c in self.components]})") + + def __len__(self) -> int: + return len(self.components) + + def __iter__(self) -> Iterator[ComponentSpec]: + return iter(self.components) + + # -- the flat view the engine packs ------------------------------------ + @property + def active_components(self) -> list[ComponentSpec]: + """Only these contribute to the model — and therefore to the parameter + vector. An inactive component keeps its values but occupies no column, + so toggling one changes the packed width.""" + return [c for c in self.components if c.active] + + def parameter_names(self) -> list[str]: + """``"."`` in packed order. + + This IS the column order of every array below and of the tensors the + engine builds. Anything that reports per-parameter results (maps, + std devs, the UI's parameter list) must use it rather than + re-deriving an order of its own. + """ + return [f"{c.name}.{p.name}" + for c in self.active_components for p in c.parameters] + + def flat_values(self) -> np.ndarray: + return np.array([p.value for c in self.active_components + for p in c.parameters], dtype=np.float64) + + def free_mask(self) -> np.ndarray: + """True where a parameter is fitted. Fixed parameters keep their value + and are dropped from the Jacobian rather than being fitted and ignored.""" + return np.array([p.free for c in self.active_components + for p in c.parameters], dtype=bool) + + def linear_mask(self) -> np.ndarray: + """True where the model is LINEAR in the parameter — the columns + variable projection can solve by least squares (#53).""" + return np.array([p.linear for c in self.active_components + for p in c.parameters], dtype=bool) + + def bounds_arrays(self) -> tuple[np.ndarray, np.ndarray]: + """``(lower, upper)`` with ``None`` expanded to ±inf.""" + pairs = [p.bounds() for c in self.active_components for p in c.parameters] + if not pairs: + return np.empty(0), np.empty(0) + lo, hi = zip(*pairs) + return np.array(lo, np.float64), np.array(hi, np.float64) + + def set_flat_values(self, values: Sequence[float]) -> None: + """Write a packed parameter vector back onto the spec (the inverse of + :meth:`flat_values`).""" + values = np.asarray(values, dtype=float).ravel() + expected = sum(len(c.parameters) for c in self.active_components) + if values.size != expected: + raise ValueError(f"expected {expected} values for this spec, " + f"got {values.size}") + i = 0 + for c in self.active_components: + for p in c.parameters: + p.value = float(values[i]) + i += 1 + + def component_slices(self) -> dict[str, slice]: + """Where each active component's parameters sit in the packed vector — + what the component-area maps (#58) need to isolate one component.""" + out, i = {}, 0 + for c in self.active_components: + n = len(c.parameters) + out[c.name] = slice(i, i + n) + i += n + return out + + # -- HyperSpy interop -------------------------------------------------- + @classmethod + def from_model(cls, model) -> "ModelSpec": + """Read a live HyperSpy model (``s.create_model()``) into a spec.""" + comps = [] + for c in model: + params = [] + for p in c.parameters: + bmin, bmax = getattr(p, "bmin", None), getattr(p, "bmax", None) + params.append(ParameterSpec( + name=p.name, + # A multidimensional parameter's `.value` is the value at + # the CURRENT nav index; that is the right seed to carry. + value=float(np.ravel(p.value)[0]), + free=bool(p.free), + bmin=None if bmin is None else float(bmin), + bmax=None if bmax is None else float(bmax), + linear=bool(getattr(p, "_linear", False)), + units=getattr(p, "units", "") or "", + )) + comps.append(ComponentSpec( + kind=getattr(c, "_id_name", type(c).__name__), + name=c.name, active=bool(c.active), parameters=params)) + + mask = getattr(model, "_channel_switches", None) + if mask is not None: + mask = np.asarray(mask, dtype=bool) + if mask.all(): + mask = None # "all channels" is the default; don't carry it + return cls(components=comps, channel_mask=mask) + + def to_model(self, signal, *, apply_range: bool = True): + """Build a live HyperSpy model on *signal* from this spec. + + Used for the CPU reference path, for the parity tests, and for the live + single-spectrum preview — anywhere HyperSpy itself should do the work. + """ + model = signal.create_model() + # create_model() may pre-populate (an EELS model adds a background and + # the declared edges); the spec is authoritative, so start clean. + while len(model): + model.remove(model[0]) + + for cspec in self.components: + comp = _make_component(cspec) + model.append(comp) + comp.active = bool(cspec.active) + for pspec in cspec.parameters: + try: + par = getattr(comp, pspec.name) + except AttributeError: + log.debug("component %s has no parameter %s — skipped", + cspec.kind, pspec.name) + continue + # Bounds BEFORE value: HyperSpy clips an out-of-range assignment. + par.bmin, par.bmax = pspec.bmin, pspec.bmax + par.value = float(pspec.value) + par.free = bool(pspec.free) + + if apply_range and self.channel_mask is not None: + try: + model._channel_switches = np.asarray(self.channel_mask, bool) + except Exception as e: + log.debug("applying channel mask to model failed: %s", e) + return model + + # -- JSON -------------------------------------------------------------- + def to_dict(self) -> dict[str, Any]: + """JSON-safe. The channel mask goes out as a list of ``[start, stop)`` + runs rather than a per-channel bool array — a 4096-channel mask is one + or two ranges in practice, and the renderer draws ranges, not bits.""" + return {"components": [c.to_dict() for c in self.components], + "channel_ranges": _mask_to_ranges(self.channel_mask), + "n_channels": (None if self.channel_mask is None + else int(self.channel_mask.size))} + + @classmethod + def from_dict(cls, d: dict[str, Any]) -> "ModelSpec": + n = d.get("n_channels") + return cls( + components=[ComponentSpec.from_dict(c) for c in d.get("components", [])], + channel_mask=_ranges_to_mask(d.get("channel_ranges"), n), + ) + + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + +def _make_component(cspec: ComponentSpec): + """Instantiate the HyperSpy component named by ``cspec.kind``. + + Searches components1d then components2d, then exspy's components if the + ``eels`` extra is installed (EELSCLEdge and friends live there, #63). + """ + import hyperspy.components1d as c1d + import hyperspy.components2d as c2d + + mods = [c1d, c2d] + try: # optional extra — absent is normal + import exspy.components as exc + mods.append(exc) + except ImportError: + pass + + for mod in mods: + cls = getattr(mod, cspec.kind, None) + if cls is not None: + comp = cls() + if cspec.name and cspec.name != cspec.kind: + comp.name = cspec.name + return comp + raise ValueError( + f"unknown component kind {cspec.kind!r}. Not in hyperspy " + f"components1d/components2d" + + ("" if len(mods) > 2 else " (and exspy is not installed — install " + 'spyde[eels] for EELS/EDS components)')) + + +def spec_from_component(comp) -> ComponentSpec: + """One live HyperSpy component -> a ComponentSpec (used by the picker, + which instantiates a component just to show its default shape, #56).""" + return ComponentSpec( + kind=getattr(comp, "_id_name", type(comp).__name__), + name=comp.name, active=bool(comp.active), + parameters=[ParameterSpec( + name=p.name, value=float(np.ravel(p.value)[0]), free=bool(p.free), + bmin=getattr(p, "bmin", None), bmax=getattr(p, "bmax", None), + linear=bool(getattr(p, "_linear", False)), + units=getattr(p, "units", "") or "", + ) for p in comp.parameters], + ) + + +def _mask_to_ranges(mask) -> list[list[int]] | None: + if mask is None: + return None + m = np.asarray(mask, bool) + if m.all(): + return None + edges = np.diff(np.concatenate(([0], m.view(np.int8), [0]))) + starts = np.flatnonzero(edges == 1) + stops = np.flatnonzero(edges == -1) + return [[int(a), int(b)] for a, b in zip(starts, stops)] + + +def _ranges_to_mask(ranges, n_channels) -> np.ndarray | None: + if not ranges or not n_channels: + return None + m = np.zeros(int(n_channels), bool) + for a, b in ranges: + m[int(a):int(b)] = True + return m diff --git a/spyde/tests/benchmark_fitting.py b/spyde/tests/benchmark_fitting.py new file mode 100644 index 0000000..723c40b --- /dev/null +++ b/spyde/tests/benchmark_fitting.py @@ -0,0 +1,143 @@ +"""benchmark_fitting.py — batched engine vs HyperSpy multifit, at real scale. + +Run DIRECTLY (not under pytest — torch-CUDA segfaults in the pytest process on +Windows, and this is slow by design):: + + uv run python -m spyde.tests.benchmark_fitting + uv run python -m spyde.tests.benchmark_fitting --nav 128 --skip-hyperspy + +The reference number this exists to beat, measured on the dev box: HyperSpy +``multifit`` on a 1024-channel EELS SI runs at ~110 spectra/s single-threaded, +i.e. ~10 minutes for a 256x256 spectrum image. + +Reporting rules (CLAUDE.md, Benchmarking): + +* real dataset at real scale — the synthetic EELS SI from ``spyde.data``, which + is a power law plus three real core-loss edges, not a toy gaussian; +* ``torch.cuda.synchronize()`` around the timed region, since kernels are async; +* discard the first GPU run (cold CUDA init + kernel JIT is a one-time ~5 s); +* time each stage separately, so "it's slow" points at a stage rather than a + vague total. +""" +from __future__ import annotations + +import argparse +import time + +import numpy as np + + +def _build(nav, n_channels): + from spyde.data import eels_si + s = eels_si(nav=(nav, nav), n_channels=n_channels) + x = s.axes_manager.signal_axes[0].axis + return s, np.asarray(x, float) + + +def _seed_model(signal): + """Power-law background + one smeared STEP per edge. + + ``Erf`` (an error function) is the right shape here and a Gaussian is not: + a core-loss edge is a step up at the onset, and a peak cannot represent it. + A Gaussian-per-edge model is unfittable — measured, both HyperSpy and the + batched engine exhaust their iteration budgets on it and the benchmark ends + up timing the iteration cap instead of convergence. Without exspy's real + ``EELSCLEdge`` (#63) this is the closest well-posed model. + """ + from hyperspy.components1d import Erf, PowerLaw + from spyde.data.synthetic import EELS_EDGES + + m = signal.create_model() + comps = [PowerLaw()] + for onset in EELS_EDGES.values(): + e = Erf() + e.origin.value, e.sigma.value, e.A.value = onset, 4.0, 2e3 + comps.append(e) + m.extend(comps) + m[0].A.value, m[0].r.value = 1e8, 3.0 + m[0].origin.free = False + return m + + +def _time(fn, *, sync=False): + if sync: + import torch + torch.cuda.synchronize() + t0 = time.perf_counter() + out = fn() + if sync: + import torch + torch.cuda.synchronize() + return out, time.perf_counter() - t0 + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--nav", type=int, default=64, + help="navigation grid edge (nav x nav spectra)") + ap.add_argument("--channels", type=int, default=1024) + ap.add_argument("--skip-hyperspy", action="store_true", + help="skip the reference (it is the slow one)") + args = ap.parse_args() + + from spyde.fitting import ModelSpec + from spyde.fitting.engine import default_device, fit_batched + + P = args.nav * args.nav + print(f"building {args.nav}x{args.nav} = {P} spectra x {args.channels} ch …") + s, x = _build(args.nav, args.channels) + spec = ModelSpec.from_model(_seed_model(s)) + n_free = int(spec.free_mask().sum()) + print(f"model: {[c.kind for c in spec]} ({n_free} free parameters)\n") + + results = {} + + # -- reference --------------------------------------------------------- + if not args.skip_hyperspy: + m = _seed_model(s) + print("hyperspy multifit … (this is the slow one)") + _, dt = _time(lambda: m.multifit(optimizer="lm", show_progressbar=False)) + results["hyperspy multifit"] = dt + print(f" {dt:8.2f} s {P/dt:8.1f} spectra/s\n") + + # -- batched, CPU ------------------------------------------------------ + print("batched engine (cpu) …") + r_cpu, dt = _time(lambda: fit_batched(spec, s.data, x, device="cpu")) + results["batched cpu"] = dt + print(f" {dt:8.2f} s {P/dt:8.1f} spectra/s " + f"converged {r_cpu.convergence_rate:.1%} iters {r_cpu.n_iter}\n") + + # -- batched, GPU ------------------------------------------------------ + if default_device() == "cuda": + print("batched engine (cuda) … [discarding the cold run]") + fit_batched(spec, s.data[:2, :2], x, device="cuda") # warm up + r_gpu, dt = _time(lambda: fit_batched(spec, s.data, x, device="cuda"), + sync=True) + results["batched cuda"] = dt + print(f" {dt:8.2f} s {P/dt:8.1f} spectra/s " + f"converged {r_gpu.convergence_rate:.1%} iters {r_gpu.n_iter}") + agree = np.nanmax(np.abs(r_gpu.values - r_cpu.values) + / np.maximum(np.abs(r_cpu.values), 1e-12)) + print(f" max relative CPU/GPU disagreement: {agree:.2e}\n") + else: + print("no CUDA device — skipping the GPU leg\n") + + # -- summary ----------------------------------------------------------- + print("=" * 62) + base = results.get("hyperspy multifit") + for name, dt in results.items(): + line = f"{name:22s} {dt:8.2f} s {P/dt:9.1f} spectra/s" + if base and name != "hyperspy multifit": + line += f" {base/dt:7.1f}x" + print(line) + if base: + for name, dt in results.items(): + if name == "hyperspy multifit": + continue + print(f"\n256x256 (65536 px) extrapolated: " + f"{name} {dt/P*65536:.1f} s vs hyperspy " + f"{base/P*65536/60:.1f} min") + + +if __name__ == "__main__": + main() diff --git a/spyde/tests/migrated/test_fitting_engine.py b/spyde/tests/migrated/test_fitting_engine.py new file mode 100644 index 0000000..bdc95af --- /dev/null +++ b/spyde/tests/migrated/test_fitting_engine.py @@ -0,0 +1,287 @@ +"""The batched engine must reproduce HyperSpy's `multifit`. + +This is the acceptance gate for #53, and the reason it is written this way: +we are *replacing a reference implementation*, so "it converged" proves +nothing. The bar is that the parameters agree with the ones HyperSpy's own +optimiser finds on the same data. + +Everything here runs on the CPU. torch-CUDA segfaults under pytest on Windows +(CLAUDE.md), and the engine is the same code on both devices — GPU is exercised +by the benchmark, run outside pytest. +""" +from __future__ import annotations + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") + +import hyperspy.api as hs +from hyperspy.components1d import Gaussian, Offset, PowerLaw + +from spyde.fitting import ModelSpec +from spyde.fitting.engine import default_device, fit_batched + + +# -------------------------------------------------------------------------- +# fixtures: a small spectrum image with known truth +# -------------------------------------------------------------------------- + +def _make_si(ny=4, nx=5, nc=256, noise=0.0, seed=0): + """Offset + Gaussian, with the amplitude and centre varying across the + scan so every pixel has a genuinely different answer (a constant field + would let a broadcasting bug pass).""" + rng = np.random.default_rng(seed) + x = np.linspace(0.0, 50.0, nc) + amp = np.linspace(20.0, 60.0, ny * nx).reshape(ny, nx) + cen = np.linspace(20.0, 30.0, ny * nx).reshape(ny, nx) + data = np.empty((ny, nx, nc)) + for iy in range(ny): + for ix in range(nx): + g = (amp[iy, ix] / (3.0 * np.sqrt(2 * np.pi)) + * np.exp(-((x - cen[iy, ix]) ** 2) / (2 * 3.0 ** 2))) + data[iy, ix] = 5.0 + g + if noise: + data += rng.normal(0.0, noise, data.shape) + s = hs.signals.Signal1D(data) + s.axes_manager.signal_axes[0].offset = x[0] + s.axes_manager.signal_axes[0].scale = x[1] - x[0] + return s, x, amp, cen + + +def _seed_model(signal): + m = signal.create_model() + m.extend([Offset(), Gaussian()]) + m[0].offset.value = 1.0 + m[1].A.value, m[1].centre.value, m[1].sigma.value = 30.0, 25.0, 2.0 + return m + + +class TestParityWithMultifit: + def test_parameters_match_hyperspy_multifit(self): + """THE acceptance test for the engine.""" + s, x, _amp, _cen = _make_si(noise=0.05) + + ref = _seed_model(s) + ref.multifit(optimizer="lm", show_progressbar=False) + ref_maps = np.stack([ + ref[0].offset.map["values"].ravel(), + ref[1].A.map["values"].ravel(), + ref[1].centre.map["values"].ravel(), + ref[1].sigma.map["values"].ravel(), + ], axis=1) + + spec = ModelSpec.from_model(_seed_model(s)) + got = fit_batched(spec, s.data, x, device="cpu") + + order = spec.parameter_names() + cols = [order.index(n) for n in + ["Offset.offset", "Gaussian.A", "Gaussian.centre", "Gaussian.sigma"]] + np.testing.assert_allclose(got.values[:, cols], ref_maps, + rtol=1e-4, atol=1e-6) + + def test_recovers_the_truth_it_was_built_from(self): + s, x, amp, cen = _make_si(noise=0.0) + spec = ModelSpec.from_model(_seed_model(s)) + got = fit_batched(spec, s.data, x, device="cpu") + maps = got.as_maps(spec, s.data.shape[:2]) + np.testing.assert_allclose(maps["Gaussian.A"], amp, rtol=1e-4) + np.testing.assert_allclose(maps["Gaussian.centre"], cen, rtol=1e-4) + np.testing.assert_allclose(maps["Offset.offset"], + np.full_like(amp, 5.0), atol=1e-3) + + def test_every_position_converges_on_clean_data(self): + s, x, _, _ = _make_si(noise=0.0) + got = fit_batched(ModelSpec.from_model(_seed_model(s)), s.data, x, + device="cpu") + assert got.convergence_rate == 1.0 + + def test_power_law_background_matches_multifit(self): + """PowerLaw is the awkward one — a masked branch and a bounded + exponent — and it is the EELS background, so it has to be right.""" + x = np.linspace(200.0, 800.0, 512) + data = np.stack([1e6 * a * x ** -3.0 for a in (0.8, 1.0, 1.4, 2.0)]) + s = hs.signals.Signal1D(data) + s.axes_manager.signal_axes[0].offset = x[0] + s.axes_manager.signal_axes[0].scale = x[1] - x[0] + + def build(): + m = s.create_model() + m.append(PowerLaw()) + m[0].A.value, m[0].r.value = 1e5, 2.0 + m[0].origin.free = False + return m + + ref = build() + ref.multifit(optimizer="lm", show_progressbar=False) + spec = ModelSpec.from_model(build()) + got = fit_batched(spec, s.data, x, device="cpu") + + order = spec.parameter_names() + for name, want in (("PowerLaw.A", ref[0].A.map["values"].ravel()), + ("PowerLaw.r", ref[0].r.map["values"].ravel())): + np.testing.assert_allclose(got.values[:, order.index(name)], want, + rtol=1e-3) + + +class TestFixedParametersAndBounds: + def test_fixed_parameter_is_not_moved(self): + s, x, _, _ = _make_si(noise=0.02) + m = _seed_model(s) + m[1].sigma.free = False + m[1].sigma.value = 3.0 + spec = ModelSpec.from_model(m) + got = fit_batched(spec, s.data, x, device="cpu") + col = spec.parameter_names().index("Gaussian.sigma") + np.testing.assert_allclose(got.values[:, col], 3.0, rtol=1e-12) + + def test_bounds_are_respected(self): + s, x, _, _ = _make_si(noise=0.02) + m = _seed_model(s) + m[1].centre.bmin, m[1].centre.bmax = 24.0, 26.0 + spec = ModelSpec.from_model(m) + got = fit_batched(spec, s.data, x, device="cpu") + col = spec.parameter_names().index("Gaussian.centre") + assert (got.values[:, col] >= 24.0 - 1e-9).all() + assert (got.values[:, col] <= 26.0 + 1e-9).all() + + def test_all_parameters_fixed_is_a_clear_error(self): + s, x, _, _ = _make_si() + m = _seed_model(s) + for c in m: + for p in c.parameters: + p.free = False + with pytest.raises(ValueError, match="fixed"): + fit_batched(ModelSpec.from_model(m), s.data, x, device="cpu") + + +class TestSignalRange: + def test_masked_channels_do_not_influence_the_fit(self): + """A channel outside the range must have NO effect — so corrupting it + beyond recognition must not move the answer at all.""" + s, x, _, _ = _make_si(noise=0.0) + m = _seed_model(s) + m.set_signal_range(15.0, 40.0) + spec = ModelSpec.from_model(m) + + clean = fit_batched(spec, s.data, x, device="cpu") + wrecked = s.data.copy() + wrecked[..., :int(0.2 * len(x))] += 1e6 # far outside the range + dirty = fit_batched(spec, wrecked, x, device="cpu") + np.testing.assert_allclose(clean.values, dirty.values, rtol=1e-8) + + def test_mask_length_mismatch_is_caught(self): + s, x, _, _ = _make_si() + spec = ModelSpec.from_model(_seed_model(s)) + spec.channel_mask = np.ones(7, bool) + with pytest.raises(ValueError, match="channel mask"): + fit_batched(spec, s.data, x, device="cpu") + + +class TestWeights: + def test_poisson_weighting_runs_and_stays_sane(self): + s, x, amp, _ = _make_si(noise=0.0) + spec = ModelSpec.from_model(_seed_model(s)) + got = fit_batched(spec, s.data, x, weights="poisson", device="cpu") + maps = got.as_maps(spec, s.data.shape[:2]) + np.testing.assert_allclose(maps["Gaussian.A"], amp, rtol=1e-3) + + def test_unknown_weighting_is_rejected(self): + s, x, _, _ = _make_si() + with pytest.raises(ValueError, match="weighting"): + fit_batched(ModelSpec.from_model(_seed_model(s)), s.data, x, + weights="gaussianish", device="cpu") + + +class TestChunking: + def test_chunked_and_unchunked_agree(self): + """Chunking is a memory optimisation and must be numerically invisible; + a chunk-boundary bug would otherwise show as a faint grid in the maps.""" + s, x, _, _ = _make_si(ny=6, nx=6, noise=0.02) + spec = ModelSpec.from_model(_seed_model(s)) + whole = fit_batched(spec, s.data, x, device="cpu") + pieces = fit_batched(spec, s.data, x, device="cpu", chunk=4) + # NOT bit-identical, and it should not be asserted as such: each + # position is an independent problem, but batched BLAS picks different + # kernels/reduction orders for different batch sizes, so the last + # couple of ULPs move. ~1e-8 relative is float noise; a real + # chunk-boundary bug would show up orders of magnitude larger (and as a + # visible grid in the maps). + np.testing.assert_allclose(whole.values, pieces.values, rtol=1e-6) + + def test_result_shape_follows_the_navigation_grid(self): + s, x, _, _ = _make_si(ny=3, nx=7) + spec = ModelSpec.from_model(_seed_model(s)) + got = fit_batched(spec, s.data, x, device="cpu") + assert got.values.shape == (21, len(spec.parameter_names())) + assert got.as_maps(spec, (3, 7))["Gaussian.A"].shape == (3, 7) + + +class TestInitialValues: + def test_per_position_seeds_are_used(self): + """The hand-off seeded propagation (#54) depends on: each position may + start from its own values, not one shared guess.""" + s, x, amp, cen = _make_si(noise=0.0) + spec = ModelSpec.from_model(_seed_model(s)) + P, n = s.data.shape[0] * s.data.shape[1], len(spec.parameter_names()) + seeds = np.broadcast_to(spec.flat_values(), (P, n)).copy() + seeds[:, spec.parameter_names().index("Gaussian.centre")] = cen.ravel() + got = fit_batched(spec, s.data, x, device="cpu", initial=seeds) + np.testing.assert_allclose( + got.as_maps(spec, cen.shape)["Gaussian.centre"], cen, rtol=1e-4) + + def test_seeds_outside_bounds_are_clipped_not_rejected(self): + s, x, _, _ = _make_si(noise=0.0) + m = _seed_model(s) + m[1].centre.bmin, m[1].centre.bmax = 20.0, 30.0 + spec = ModelSpec.from_model(m) + P, n = s.data.shape[0] * s.data.shape[1], len(spec.parameter_names()) + seeds = np.broadcast_to(spec.flat_values(), (P, n)).copy() + seeds[:, spec.parameter_names().index("Gaussian.centre")] = 1e6 + got = fit_batched(spec, s.data, x, device="cpu", initial=seeds) + col = spec.parameter_names().index("Gaussian.centre") + assert (got.values[:, col] <= 30.0 + 1e-9).all() + + +class TestGuards: + def test_unsupported_component_refuses_rather_than_dropping_it(self): + """Silently omitting a component would fit a DIFFERENT model and still + return a plausible answer.""" + from spyde.fitting.spec import ComponentSpec, ParameterSpec + s, x, _, _ = _make_si() + spec = ModelSpec(components=[ + ComponentSpec(kind="Voigt", parameters=[ParameterSpec("area", 1.0)])]) + with pytest.raises(NotImplementedError, match="Voigt"): + fit_batched(spec, s.data, x, device="cpu") + + def test_axis_length_mismatch_is_caught(self): + s, _x, _, _ = _make_si() + spec = ModelSpec.from_model(_seed_model(s)) + with pytest.raises(ValueError, match="signal axis"): + fit_batched(spec, s.data, np.linspace(0, 1, 9), device="cpu") + + def test_progress_callback_reports_completion(self): + s, x, _, _ = _make_si(ny=4, nx=4) + seen = [] + fit_batched(ModelSpec.from_model(_seed_model(s)), s.data, x, + device="cpu", chunk=5, progress=lambda d, t: seen.append((d, t))) + assert seen and seen[-1] == (16, 16) + + def test_a_failing_progress_callback_does_not_kill_the_fit(self): + s, x, _, _ = _make_si(ny=2, nx=2) + def boom(done, total): + raise RuntimeError("renderer went away") + got = fit_batched(ModelSpec.from_model(_seed_model(s)), s.data, x, + device="cpu", progress=boom) + assert got.values.shape[0] == 4 + + +class TestDevice: + def test_default_device_is_reported(self): + assert default_device() in ("cpu", "cuda") + + def test_cpu_is_always_usable(self): + s, x, _, _ = _make_si(ny=2, nx=2) + got = fit_batched(ModelSpec.from_model(_seed_model(s)), s.data, x, + device="cpu") + assert got.device == "cpu" diff --git a/spyde/tests/migrated/test_model_spec.py b/spyde/tests/migrated/test_model_spec.py new file mode 100644 index 0000000..6c3e2a2 --- /dev/null +++ b/spyde/tests/migrated/test_model_spec.py @@ -0,0 +1,231 @@ +"""ModelSpec round-trips against HyperSpy and packs parameters in a stable order. + +Two contracts: + +1. **HyperSpy is the storage format.** A spec taken from a live model and put + back must reproduce the same component kinds, values, bounds, free flags and + channel range — otherwise a model saved by SpyDE would not reopen in plain + HyperSpy, and the parity tests in #53 would be comparing different models. +2. **Packed parameter order is part of the API.** Column j of every array (and + of the engine's tensors) always means the same parameter, and + `parameter_names()` is that order. Anything reporting per-parameter results + relies on it. +""" +from __future__ import annotations + +import numpy as np +import pytest +import hyperspy.api as hs +from hyperspy.components1d import Gaussian, Offset, PowerLaw + +from spyde.fitting import ComponentSpec, ModelSpec, ParameterSpec + + +def _signal(n=64, ny=3, nx=3): + return hs.signals.Signal1D(np.random.default_rng(0).random((ny, nx, n)) + 1.0) + + +def _model(): + m = _signal().create_model() + m.extend([PowerLaw(), Gaussian(centre=30.0, sigma=4.0, A=10.0)]) + return m + + +class TestFromModel: + def test_reads_components_and_kinds(self): + spec = ModelSpec.from_model(_model()) + assert [c.kind for c in spec] == ["PowerLaw", "Gaussian"] + + def test_reads_parameter_values(self): + spec = ModelSpec.from_model(_model()) + g = spec["Gaussian"] + assert g["centre"].value == pytest.approx(30.0) + assert g["sigma"].value == pytest.approx(4.0) + assert g["A"].value == pytest.approx(10.0) + + def test_reads_bounds_including_absent_ones(self): + spec = ModelSpec.from_model(_model()) + g = spec["Gaussian"] + assert g["A"].bmin == pytest.approx(0.0) # HyperSpy bounds A >= 0 + assert g["centre"].bmin is None # unbounded + assert g["centre"].bounds() == (-np.inf, np.inf) + + def test_reads_the_linear_flag(self): + """The variable-projection path (#53) keys off this, so it has to come + from HyperSpy rather than a hand-maintained list of our own.""" + spec = ModelSpec.from_model(_model()) + assert spec["Gaussian"]["A"].linear is True + assert spec["Gaussian"]["centre"].linear is False + + def test_reads_free_flags(self): + m = _model() + m[1].sigma.free = False + spec = ModelSpec.from_model(m) + assert spec["Gaussian"]["sigma"].free is False + assert spec["Gaussian"]["A"].free is True + + def test_all_channels_active_is_stored_as_none(self): + """A full mask is the default; carrying an all-True array of signal + length through JSON on every model would be pure waste.""" + assert ModelSpec.from_model(_model()).channel_mask is None + + def test_reads_a_restricted_signal_range(self): + m = _model() + m.set_signal_range(10, 40) + spec = ModelSpec.from_model(m) + assert spec.channel_mask is not None + assert spec.channel_mask.sum() < spec.channel_mask.size + assert spec.channel_mask[20] and not spec.channel_mask[5] + + +class TestToModel: + def test_round_trip_preserves_kinds_and_values(self): + spec = ModelSpec.from_model(_model()) + back = ModelSpec.from_model(spec.to_model(_signal())) + assert [c.kind for c in back] == [c.kind for c in spec] + assert back.flat_values() == pytest.approx(spec.flat_values()) + + def test_round_trip_preserves_bounds_and_free_flags(self): + m = _model() + m[1].sigma.free = False + m[1].centre.bmin, m[1].centre.bmax = 20.0, 40.0 + spec = ModelSpec.from_model(m) + back = ModelSpec.from_model(spec.to_model(_signal())) + assert back["Gaussian"]["sigma"].free is False + assert back["Gaussian"]["centre"].bmin == pytest.approx(20.0) + assert back["Gaussian"]["centre"].bmax == pytest.approx(40.0) + + def test_round_trip_preserves_the_signal_range(self): + m = _model() + m.set_signal_range(10, 40) + spec = ModelSpec.from_model(m) + back = ModelSpec.from_model(spec.to_model(_signal())) + assert np.array_equal(back.channel_mask, spec.channel_mask) + + def test_bounds_are_set_before_values(self): + """HyperSpy clips an assignment that falls outside the bounds, so a + to_model() that wrote the value first would silently corrupt it.""" + spec = ModelSpec(components=[ComponentSpec( + kind="Gaussian", parameters=[ + ParameterSpec("A", value=500.0, bmin=0.0, bmax=1000.0), + ParameterSpec("centre", value=30.0), + ParameterSpec("sigma", value=4.0, bmin=0.0, bmax=10.0), + ])]) + m = spec.to_model(_signal()) + assert float(np.ravel(m[0].A.value)[0]) == pytest.approx(500.0) + + def test_inactive_component_survives_the_round_trip(self): + m = _model() + m[1].active = False + back = ModelSpec.from_model(ModelSpec.from_model(m).to_model(_signal())) + assert back["Gaussian"].active is False + + def test_replaces_any_prepopulated_components(self): + """create_model() can pre-populate (an EELS model adds a background and + the declared edges). The spec is authoritative, so those must go.""" + spec = ModelSpec(components=[ComponentSpec( + kind="Offset", parameters=[ParameterSpec("offset", value=1.0)])]) + m = spec.to_model(_signal()) + assert [getattr(c, "_id_name", "") for c in m] == ["Offset"] + + def test_unknown_kind_is_a_clear_error(self): + spec = ModelSpec(components=[ComponentSpec(kind="Flurbulator")]) + with pytest.raises(ValueError, match="Flurbulator"): + spec.to_model(_signal()) + + +class TestPackedOrder: + def test_parameter_names_follow_component_then_parameter_order(self): + spec = ModelSpec.from_model(_model()) + names = spec.parameter_names() + assert names[0].startswith("PowerLaw.") + assert names[-1].startswith("Gaussian.") + assert len(names) == len(spec.flat_values()) + + def test_masks_line_up_with_values(self): + spec = ModelSpec.from_model(_model()) + n = len(spec.flat_values()) + assert len(spec.free_mask()) == n + assert len(spec.linear_mask()) == n + lo, hi = spec.bounds_arrays() + assert len(lo) == len(hi) == n + + def test_set_flat_values_is_the_inverse_of_flat_values(self): + spec = ModelSpec.from_model(_model()) + new = spec.flat_values() + 1.5 + spec.set_flat_values(new) + assert spec.flat_values() == pytest.approx(new) + + def test_set_flat_values_rejects_the_wrong_width(self): + spec = ModelSpec.from_model(_model()) + with pytest.raises(ValueError, match="expected"): + spec.set_flat_values(np.zeros(3)) + + def test_inactive_components_occupy_no_columns(self): + """An inactive component contributes nothing to the model, so it must + not take a column — otherwise the engine fits a parameter of something + that is not being evaluated.""" + spec = ModelSpec.from_model(_model()) + wide = len(spec.flat_values()) + spec["Gaussian"].active = False + narrow = len(spec.flat_values()) + assert narrow == wide - 3 + assert all(n.startswith("PowerLaw.") for n in spec.parameter_names()) + + def test_component_slices_isolate_each_component(self): + spec = ModelSpec.from_model(_model()) + sl = spec.component_slices() + vals = spec.flat_values() + assert set(sl) == {"PowerLaw", "Gaussian"} + assert len(vals[sl["Gaussian"]]) == 3 + # Slices must tile the vector exactly, with no gap or overlap. + covered = sorted(i for s in sl.values() for i in range(*s.indices(len(vals)))) + assert covered == list(range(len(vals))) + + +class TestJson: + def test_dict_round_trip(self): + spec = ModelSpec.from_model(_model()) + back = ModelSpec.from_dict(spec.to_dict()) + assert back.parameter_names() == spec.parameter_names() + assert back.flat_values() == pytest.approx(spec.flat_values()) + + def test_dict_is_json_serialisable(self): + import json + m = _model() + m.set_signal_range(10, 40) + json.dumps(ModelSpec.from_model(m).to_dict()) # must not raise + + def test_channel_mask_survives_as_ranges(self): + m = _model() + m.set_signal_range(10, 40) + spec = ModelSpec.from_model(m) + back = ModelSpec.from_dict(spec.to_dict()) + assert np.array_equal(back.channel_mask, spec.channel_mask) + + def test_two_disjoint_ranges_survive(self): + """remove_signal_range can leave a hole in the middle; the compact + range encoding has to keep both runs.""" + mask = np.zeros(64, bool) + mask[5:15] = True + mask[40:50] = True + spec = ModelSpec(components=[ComponentSpec(kind="Offset")], + channel_mask=mask) + assert len(spec.to_dict()["channel_ranges"]) == 2 + assert np.array_equal(ModelSpec.from_dict(spec.to_dict()).channel_mask, + mask) + + +class TestCopy: + def test_copy_is_deep(self): + spec = ModelSpec.from_model(_model()) + clone = spec.copy() + clone["Gaussian"]["A"].value = 999.0 + assert spec["Gaussian"]["A"].value != 999.0 + + def test_lookup_errors_name_what_is_available(self): + spec = ModelSpec.from_model(_model()) + with pytest.raises(KeyError, match="Gaussian"): + spec["NoSuchComponent"] + with pytest.raises(KeyError, match="centre"): + spec["Gaussian"]["no_such_parameter"] diff --git a/spyde/tests/migrated/test_torch_components.py b/spyde/tests/migrated/test_torch_components.py new file mode 100644 index 0000000..76c9d1e --- /dev/null +++ b/spyde/tests/migrated/test_torch_components.py @@ -0,0 +1,244 @@ +"""Every batched torch component must match the HyperSpy component it ports. + +This file is the actual specification for `spyde/fitting/components.py`. The +formulas there were transcribed from HyperSpy's `expression=` strings, and the +only thing that makes that transcription trustworthy is checking it numerically +against the real component at randomised parameter values. + +Two failure modes this is built to catch, both silent: + +* **Wrong parameter ORDER.** HyperSpy's order is neither alphabetical nor the + constructor's — `PowerLaw` is (A, left_cutoff, origin, r). A mismatch fits + the wrong parameter and still converges to something. +* **Wrong convention.** A HyperSpy `Gaussian`'s `A` is the AREA, not the peak + height. A height-based port produces a plausible curve that is simply not + HyperSpy's. + +CPU only: these are tiny tensors and torch-CUDA segfaults under pytest on +Windows (CLAUDE.md). GPU correctness is exercised by the engine's own +subprocess test. +""" +from __future__ import annotations + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") + +import hyperspy.components1d as c1d + +from spyde.fitting import components as tc +from spyde.fitting.spec import ModelSpec, ComponentSpec, ParameterSpec + +# kind -> parameter values to test at, keyed by HyperSpy parameter name. +# Chosen to be physically sensible AND to avoid the degenerate spots +# (sigma=0, x<=origin for a power law) that would make any port agree. +CASES = { + "Gaussian": {"A": 12.0, "centre": 3.0, "sigma": 1.7}, + "GaussianHF": {"centre": 2.0, "fwhm": 3.1, "height": 5.0}, + "Lorentzian": {"A": 8.0, "centre": -1.0, "gamma": 2.2}, + "PowerLaw": {"A": 900.0, "left_cutoff": 0.0, "origin": -12.0, "r": 2.6}, + "Offset": {"offset": 3.25}, + "Exponential": {"A": 40.0, "tau": 6.5}, + "Arctan": {"A": 2.5, "k": 1.3, "x0": 1.0}, + "Erf": {"A": 7.0, "origin": 0.5, "sigma": 2.0}, + "HeavisideStep": {"A": 4.0, "n": 1.5}, + "Logistic": {"a": 5.0, "b": 2.0, "c": 0.7, "origin": 1.0}, +} + +X = np.linspace(-8.0, 12.0, 401) + + +def _hyperspy_component(kind: str): + comp = getattr(c1d, kind)() + for name, value in CASES[kind].items(): + getattr(comp, name).value = value + return comp + + +def _torch_eval(kind: str, x: np.ndarray) -> np.ndarray: + comp = _hyperspy_component(kind) + batched = tc.get_component(kind) + values = np.array([[getattr(comp, n).value for n in batched.params]]) + y = batched(torch.as_tensor(x, dtype=torch.float64), + torch.as_tensor(values, dtype=torch.float64)) + return y.detach().numpy()[0] + + +@pytest.mark.parametrize("kind", sorted(CASES)) +class TestParity: + def test_parameter_order_matches_hyperspy(self, kind): + """The port's tuple must BE HyperSpy's `component.parameters` order.""" + hs_order = tuple(p.name for p in _hyperspy_component(kind).parameters) + assert tc.get_component(kind).params == hs_order + + def test_linear_flags_match_hyperspy(self, kind): + """Variable projection (#53) trusts these; they must come from HyperSpy.""" + comp = _hyperspy_component(kind) + hs_linear = tuple(bool(getattr(p, "_linear", False)) + for p in comp.parameters) + assert tc.get_component(kind).linear == hs_linear + + def test_values_match_hyperspy(self, kind): + """The whole point: same parameters, same curve.""" + expected = np.asarray(_hyperspy_component(kind).function(X), float) + got = _torch_eval(kind, X) + assert got.shape == expected.shape + np.testing.assert_allclose(got, expected, rtol=1e-10, atol=1e-12) + + def test_is_differentiable(self, kind): + """The engine gets its Jacobian from autograd, so every parameter must + carry a finite gradient — a NaN in a discarded `where` branch would + poison it even where the mask throws the value away.""" + batched = tc.get_component(kind) + comp = _hyperspy_component(kind) + values = torch.tensor( + [[getattr(comp, n).value for n in batched.params]], + dtype=torch.float64, requires_grad=True) + y = batched(torch.as_tensor(X, dtype=torch.float64), values) + y.sum().backward() + assert values.grad is not None + assert torch.isfinite(values.grad).all(), \ + f"{kind} produced a non-finite gradient: {values.grad}" + + +class TestBatching: + def test_one_call_evaluates_every_position_independently(self): + """P rows in, P different curves out — no broadcasting mistake that + makes every pixel share one parameter set.""" + batched = tc.get_component("Gaussian") + vals = torch.tensor([[10.0, 0.0, 1.0], + [10.0, 4.0, 1.0], + [20.0, 0.0, 2.0]], dtype=torch.float64) + y = batched(torch.as_tensor(X, dtype=torch.float64), vals) + assert y.shape == (3, len(X)) + assert X[int(y[0].argmax())] == pytest.approx(0.0, abs=0.1) + assert X[int(y[1].argmax())] == pytest.approx(4.0, abs=0.1) + assert not torch.allclose(y[0], y[2]) + + def test_batched_equals_looping_one_at_a_time(self): + rng = np.random.default_rng(0) + vals = np.stack([rng.uniform([1, -3, 0.5], [20, 6, 3.0]) for _ in range(16)]) + x = torch.as_tensor(X, dtype=torch.float64) + batched = tc.get_component("Gaussian") + together = batched(x, torch.as_tensor(vals, dtype=torch.float64)) + for i in range(len(vals)): + one = batched(x, torch.as_tensor(vals[i:i + 1], dtype=torch.float64)) + torch.testing.assert_close(together[i], one[0]) + + def test_offset_broadcasts_over_the_signal_axis(self): + """Offset ignores x, which makes it the easy one to get wrong: it must + still return a full (P, C) block, not a (P, 1) that broadcasts later.""" + y = tc.get_component("Offset")( + torch.as_tensor(X, dtype=torch.float64), + torch.tensor([[2.0], [5.0]], dtype=torch.float64)) + assert y.shape == (2, len(X)) + assert torch.allclose(y[0], torch.full_like(y[0], 2.0)) + assert torch.allclose(y[1], torch.full_like(y[1], 5.0)) + + +class TestPolynomial: + @pytest.mark.parametrize("order", [1, 2, 3]) + def test_matches_hyperspy_at_each_order(self, order): + comp = c1d.Polynomial(order=order) + for k in range(order + 1): + getattr(comp, f"a{k}").value = 0.5 * (k + 1) + batched = tc.get_component("Polynomial", n_params=order + 1) + assert batched.params == tuple(p.name for p in comp.parameters) + vals = np.array([[getattr(comp, n).value for n in batched.params]]) + got = batched(torch.as_tensor(X, dtype=torch.float64), + torch.as_tensor(vals, dtype=torch.float64)).numpy()[0] + np.testing.assert_allclose(got, np.asarray(comp.function(X), float), + rtol=1e-10, atol=1e-12) + + def test_needs_an_explicit_order(self): + with pytest.raises(ValueError, match="n_params"): + tc.get_component("Polynomial") + + +class TestPowerLawEdgeCases: + def test_zero_below_the_cutoff(self): + y = _torch_eval("PowerLaw", np.array([-5.0, -1.0, 0.0, 1.0, 5.0])) + assert y[0] == 0.0 and y[1] == 0.0 and y[2] == 0.0 + assert y[3] > 0 and y[4] > 0 + + def test_no_nan_gradient_from_the_dead_branch(self): + """`(x - origin) ** -r` is inf/NaN where x <= origin. A naive + implementation masks the VALUE but still back-propagates the NaN.""" + batched = tc.get_component("PowerLaw") + vals = torch.tensor([[900.0, 0.0, 2.0, 2.6]], dtype=torch.float64, + requires_grad=True) + x = torch.linspace(-5, 20, 200, dtype=torch.float64) # spans origin=2 + batched(x, vals).sum().backward() + assert torch.isfinite(vals.grad).all() + + +class TestRegistry: + def test_unknown_component_names_what_is_available(self): + with pytest.raises(NotImplementedError) as e: + tc.get_component("Flurbulator") + assert "Gaussian" in str(e.value) + + def test_supports_reports_true_for_a_portable_model(self): + spec = ModelSpec(components=[ + ComponentSpec(kind="PowerLaw", parameters=[ + ParameterSpec(n) for n in ("A", "left_cutoff", "origin", "r")]), + ComponentSpec(kind="Gaussian", parameters=[ + ParameterSpec(n) for n in ("A", "centre", "sigma")]), + ]) + assert tc.supports(spec) is True + + def test_supports_reports_false_rather_than_dropping_a_component(self): + """The engine must fall back to HyperSpy for a model it cannot fully + evaluate — silently omitting the unsupported component would fit a + different model and still look like it worked.""" + spec = ModelSpec(components=[ + ComponentSpec(kind="Gaussian", parameters=[ + ParameterSpec(n) for n in ("A", "centre", "sigma")]), + ComponentSpec(kind="Voigt", parameters=[ParameterSpec("area")]), + ]) + assert tc.supports(spec) is False + + def test_inactive_unsupported_component_does_not_block(self): + spec = ModelSpec(components=[ + ComponentSpec(kind="Gaussian", parameters=[ + ParameterSpec(n) for n in ("A", "centre", "sigma")]), + ComponentSpec(kind="Voigt", active=False, + parameters=[ParameterSpec("area")]), + ]) + assert tc.supports(spec) is True + + +class TestEvaluateSpec: + def test_sums_active_components_and_matches_hyperspy(self): + """A whole model, not just one component — the sum is what the engine + actually fits.""" + import hyperspy.api as hs + sig = hs.signals.Signal1D(np.zeros((1, len(X)))) + sig.axes_manager.signal_axes[0].offset = X[0] + sig.axes_manager.signal_axes[0].scale = X[1] - X[0] + + m = sig.create_model() + m.extend([c1d.Offset(), c1d.Gaussian()]) + m[0].offset.value = 2.0 + m[1].A.value, m[1].centre.value, m[1].sigma.value = 12.0, 3.0, 1.7 + expected = np.asarray(m[0].function(X), float) + \ + np.asarray(m[1].function(X), float) + + spec = ModelSpec.from_model(m) + got = tc.evaluate(spec, torch.as_tensor(X, dtype=torch.float64), + torch.as_tensor(spec.flat_values()[None, :], + dtype=torch.float64)).numpy()[0] + np.testing.assert_allclose(got, expected, rtol=1e-10, atol=1e-12) + + def test_inactive_component_contributes_nothing(self): + spec = ModelSpec(components=[ + ComponentSpec(kind="Offset", + parameters=[ParameterSpec("offset", value=2.0)]), + ComponentSpec(kind="Offset", name="off2", active=False, + parameters=[ParameterSpec("offset", value=100.0)]), + ]) + y = tc.evaluate(spec, torch.as_tensor(X, dtype=torch.float64), + torch.as_tensor(spec.flat_values()[None, :], + dtype=torch.float64)) + assert torch.allclose(y, torch.full_like(y, 2.0)) From 3e7db36ce50a974c2d67cac68f0d77af59ee87ce Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Sun, 26 Jul 2026 22:28:59 -0500 Subject: [PATCH 05/60] perf(fitting): analytic Jacobians + seeded propagation (#54) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- spyde/fitting/components.py | 241 ++++++++++++++++-- spyde/fitting/engine.py | 76 +++++- spyde/fitting/seeding.py | 142 +++++++++++ spyde/tests/benchmark_fitting.py | 88 +++++-- spyde/tests/migrated/test_fitting_seeding.py | 222 ++++++++++++++++ spyde/tests/migrated/test_torch_components.py | 93 +++++++ 6 files changed, 816 insertions(+), 46 deletions(-) create mode 100644 spyde/fitting/seeding.py create mode 100644 spyde/tests/migrated/test_fitting_seeding.py diff --git a/spyde/fitting/components.py b/spyde/fitting/components.py index 263eb6d..a7d392e 100644 --- a/spyde/fitting/components.py +++ b/spyde/fitting/components.py @@ -46,12 +46,26 @@ class TorchComponent: ``params`` mirrors HyperSpy's parameter order; ``linear`` marks the columns the model is linear in (what variable projection solves directly). + + ``grad``, when present, returns the analytic derivative stack + ``(P, C, n_params)``. **This is a performance decision with a measured + basis, not a micro-optimisation.** Autodiff on these components costs one + forward pass per parameter — measured at 51.6 ms against 3.4 ms for a + single residual evaluation, i.e. ~94% of the whole solver's time. The + derivatives here are closed forms that mostly reuse the value already + computed (``df/dA = f/A`` for every linear amplitude), so the analytic + stack costs about as much as one residual. + + A component without ``grad`` still works — the engine falls back to + ``jacfwd``. Correctness is not at stake either way, only speed, and + ``test_torch_components.py`` checks every analytic gradient against + autodiff, which is the ideal oracle for exactly this. """ - __slots__ = ("kind", "params", "linear", "_fn") + __slots__ = ("kind", "params", "linear", "_fn", "_grad") def __init__(self, kind: str, params: Sequence[str], linear: Sequence[bool], - fn: Callable): + fn: Callable, grad: Callable | None = None): if len(params) != len(linear): raise ValueError(f"{kind}: {len(params)} params vs " f"{len(linear)} linear flags") @@ -59,19 +73,41 @@ def __init__(self, kind: str, params: Sequence[str], linear: Sequence[bool], self.params = tuple(params) self.linear = tuple(bool(b) for b in linear) self._fn = fn + self._grad = grad @property def n_params(self) -> int: return len(self.params) - def __call__(self, x, p): - """``x``: ``(C,)`` or ``(P, C)``. ``p``: ``(P, n_params)``. -> ``(P, C)``.""" + @property + def has_analytic_grad(self) -> bool: + return self._grad is not None + + def _split(self, x, p): if p.shape[-1] != self.n_params: raise ValueError(f"{self.kind} expects {self.n_params} parameters " f"{self.params}, got {p.shape[-1]}") if x.dim() == 1: x = x.unsqueeze(0) # (1, C), broadcasts over P - return self._fn(x, [p[:, i:i + 1] for i in range(self.n_params)]) + return x, [p[:, i:i + 1] for i in range(self.n_params)] + + def __call__(self, x, p): + """``x``: ``(C,)`` or ``(P, C)``. ``p``: ``(P, n_params)``. -> ``(P, C)``.""" + x, cols = self._split(x, p) + return self._fn(x, cols) + + def grad(self, x, p): + """Analytic ``d(value)/d(parameter)`` -> ``(P, C, n_params)``. + + Raises if this component has no analytic form; callers should check + :attr:`has_analytic_grad` (or use :func:`evaluate_with_grad`). + """ + if self._grad is None: + raise NotImplementedError( + f"{self.kind} has no analytic gradient; the engine falls back " + f"to autodiff for it") + x, cols = self._split(x, p) + return self._grad(x, cols) def __repr__(self) -> str: # pragma: no cover return f"" @@ -154,6 +190,133 @@ def _logistic(x, p): return a / (1.0 + b * (-c * (x - origin)).exp()) +# --------------------------------------------------------------------------- +# analytic derivatives — checked against autodiff in test_torch_components.py +# +# Every one of these reuses the value where it can: for a linear amplitude the +# derivative IS the shape (df/dA = f/A), which is already the expensive part. +# --------------------------------------------------------------------------- + +def _stack(*cols): + import torch + return torch.stack(torch.broadcast_tensors(*cols), dim=-1) + + +def _d_gaussian(x, p): + A, centre, sigma = p + s = sigma + _EPS + shape = (-((x - centre) ** 2) / (2 * s ** 2)).exp() / (s * _SQRT_2PI) + f = A * shape + d = x - centre + return _stack(shape, # df/dA + f * d / s ** 2, # df/dcentre + f * (d ** 2 / s ** 3 - 1.0 / s)) # df/dsigma + + +def _d_gaussian_hf(x, p): + centre, fwhm, height = p + w = fwhm + _EPS + d = x - centre + shape = (-(d ** 2) * _FOUR_LOG2 / w ** 2).exp() + f = height * shape + return _stack(f * 2.0 * _FOUR_LOG2 * d / w ** 2, # df/dcentre + f * 2.0 * _FOUR_LOG2 * d ** 2 / w ** 3, # df/dfwhm + shape) # df/dheight + + +def _d_lorentzian(x, p): + A, centre, gamma = p + d = x - centre + D = d ** 2 + gamma ** 2 + _EPS + shape = gamma / D / math.pi + return _stack(shape, # df/dA + A / math.pi * gamma * 2.0 * d / D ** 2, # df/dcentre + A / math.pi * (d ** 2 - gamma ** 2) / D ** 2) # df/dgamma + + +def _d_power_law(x, p): + import torch + A, left_cutoff, origin, r = p + base = x - origin + live = (left_cutoff < x) & (base > 0) + safe = torch.where(live, base, torch.ones_like(base)) + shape = torch.where(live, safe ** (-r), torch.zeros_like(base)) + f = A * shape + zero = torch.zeros_like(base) + return _stack(shape, # df/dA + zero, # df/dleft_cutoff (step) + torch.where(live, f * r / safe, zero), # df/dorigin + torch.where(live, -f * safe.log(), zero)) # df/dr + + +def _d_offset(x, p): + import torch + (offset,) = p + return torch.ones_like(offset + 0.0 * x).unsqueeze(-1) + + +def _d_exponential(x, p): + A, tau = p + t = tau + _EPS + shape = (-x / t).exp() + return _stack(shape, # df/dA + A * shape * x / t ** 2) # df/dtau + + +def _d_arctan(x, p): + A, k, x0 = p + d = x - x0 + u = k * d + D = 1.0 + u ** 2 + return _stack(u.atan(), # df/dA + A * d / D, # df/dk + -A * k / D) # df/dx0 + + +def _d_erf(x, p): + import torch + A, origin, sigma = p + s = sigma + _EPS + u = (x - origin) / _SQRT_2 / s + # d/du erf(u) = 2/sqrt(pi) * exp(-u^2) + dedu = (2.0 / math.sqrt(math.pi)) * (-(u ** 2)).exp() + half_A = A / 2.0 + return _stack(torch.erf(u) / 2.0, # df/dA + half_A * dedu * (-1.0 / (_SQRT_2 * s)), # df/dorigin + half_A * dedu * (-(x - origin) / (_SQRT_2 * s ** 2))) # df/dsigma + + +def _d_heaviside(x, p): + import torch + A, n = p + step = torch.where(x > n, torch.ones_like(x), + torch.where(x < n, torch.zeros_like(x), + torch.full_like(x, 0.5))) + # d/dn is a delta function — zero almost everywhere, which is what any + # gradient-based optimiser can use. `n` is effectively unfittable, exactly + # as it is in HyperSpy. + return _stack(step, torch.zeros_like(step)) + + +def _d_logistic(x, p): + a, b, c, origin = p + d = x - origin + e = (-c * d).exp() + E = b * e + D = 1.0 + E + return _stack(1.0 / D, # df/da + -a * e / D ** 2, # df/db + a * E * d / D ** 2, # df/dc + -a * E * c / D ** 2) # df/dorigin + + +def _d_polynomial_fn(order: int): + def grad(x, p): + return _stack(*[x ** k + 0.0 * p[0] for k in range(order + 1)]) + + return grad + + def _polynomial_fn(order: int): """Polynomial is variable-order, so its evaluator is built per order. Parameters are ``a0..a{order}`` and ``aK`` multiplies ``x**K``.""" @@ -175,23 +338,23 @@ def fn(x, p): _REGISTRY: dict[str, TorchComponent] = {} -def _register(kind, params, linear, fn) -> None: - _REGISTRY[kind] = TorchComponent(kind, params, linear, fn) +def _register(kind, params, linear, fn, grad=None) -> None: + _REGISTRY[kind] = TorchComponent(kind, params, linear, fn, grad) # Parameter tuples are HyperSpy's `component.parameters` ORDER — verified by # test_torch_components.py::TestParameterOrder against live components. -_register("Gaussian", ("A", "centre", "sigma"), (True, False, False), _gaussian) -_register("GaussianHF", ("centre", "fwhm", "height"), (False, False, True), _gaussian_hf) -_register("Lorentzian", ("A", "centre", "gamma"), (True, False, False), _lorentzian) +_register("Gaussian", ("A", "centre", "sigma"), (True, False, False), _gaussian, _d_gaussian) +_register("GaussianHF", ("centre", "fwhm", "height"), (False, False, True), _gaussian_hf, _d_gaussian_hf) +_register("Lorentzian", ("A", "centre", "gamma"), (True, False, False), _lorentzian, _d_lorentzian) _register("PowerLaw", ("A", "left_cutoff", "origin", "r"), - (True, False, False, False), _power_law) -_register("Offset", ("offset",), (True,), _offset) -_register("Exponential", ("A", "tau"), (True, False), _exponential) -_register("Arctan", ("A", "k", "x0"), (True, False, False), _arctan) -_register("Erf", ("A", "origin", "sigma"), (True, False, False), _erf) -_register("HeavisideStep", ("A", "n"), (True, False), _heaviside) -_register("Logistic", ("a", "b", "c", "origin"), (True, False, False, False), _logistic) + (True, False, False, False), _power_law, _d_power_law) +_register("Offset", ("offset",), (True,), _offset, _d_offset) +_register("Exponential", ("A", "tau"), (True, False), _exponential, _d_exponential) +_register("Arctan", ("A", "k", "x0"), (True, False, False), _arctan, _d_arctan) +_register("Erf", ("A", "origin", "sigma"), (True, False, False), _erf, _d_erf) +_register("HeavisideStep", ("A", "n"), (True, False), _heaviside, _d_heaviside) +_register("Logistic", ("a", "b", "c", "origin"), (True, False, False, False), _logistic, _d_logistic) def get_component(kind: str, *, n_params: int | None = None) -> TorchComponent: @@ -207,7 +370,7 @@ def get_component(kind: str, *, n_params: int | None = None) -> TorchComponent: return TorchComponent("Polynomial", tuple(f"a{k}" for k in range(order + 1)), (True,) * (order + 1), - _polynomial_fn(order)) + _polynomial_fn(order), _d_polynomial_fn(order)) try: return _REGISTRY[kind] except KeyError: @@ -239,6 +402,48 @@ def supports(spec) -> bool: return True +def has_analytic_grad(spec) -> bool: + """True when EVERY active component can supply an analytic derivative, so + the engine can build the whole Jacobian without autodiff.""" + for c in getattr(spec, "active_components", []): + try: + if not get_component(c.kind, n_params=len(c.parameters)).has_analytic_grad: + return False + except NotImplementedError: + return False + return True + + +def evaluate_with_grad(spec, x, values): + """``(value (P, C), jacobian (P, C, n_total))`` for a whole ModelSpec. + + The model is a SUM of components, so the Jacobian is just each component's + derivative block written into its own columns — components do not interact, + which is what makes this cheap. + + Columns follow the spec's packed order, the same order as + :meth:`~spyde.fitting.spec.ModelSpec.parameter_names`. + """ + import torch + + xb = x.unsqueeze(0) if x.dim() == 1 else x + P = values.shape[0] + C = xb.shape[-1] + n_total = values.shape[1] + out = torch.zeros((P, C), dtype=values.dtype, device=values.device) + jac = torch.zeros((P, C, n_total), dtype=values.dtype, device=values.device) + + i = 0 + for c in spec.active_components: + n = len(c.parameters) + comp = get_component(c.kind, n_params=n) + block = values[:, i:i + n] + out = out + comp(x, block) + jac[:, :, i:i + n] = comp.grad(x, block).expand(P, C, n) + i += n + return out, jac + + def evaluate(spec, x, values): """Evaluate a whole :class:`~spyde.fitting.spec.ModelSpec`. diff --git a/spyde/fitting/engine.py b/spyde/fitting/engine.py index cf25a15..6a52925 100644 --- a/spyde/fitting/engine.py +++ b/spyde/fitting/engine.py @@ -40,9 +40,17 @@ log = logging.getLogger(__name__) -# Bounds the Jacobian working set per chunk. 2**22 elements x 8 B (float64) -# x (value + tangents) lands comfortably inside a few hundred MB. -_TARGET_JACOBIAN_ELEMENTS = 1 << 22 +# Bounds the Jacobian working set per chunk, as ELEMENTS of (P, C, n). +# +# Sized to keep a GPU busy, not to be safe: chunking is what stops the whole +# (P, C, n) Jacobian materialising, but chunk too SMALL and every LM iteration +# becomes launch overhead on a nearly idle device. Measured at 1024 spectra x +# 1024 channels x 13 params, a 2**22 cap (which chopped the batch into pieces +# of 372) ran at 178 spectra/s against 267 for the whole batch in one go — the +# "safe" setting cost a third of the throughput. 2**26 elements is ~0.5 GB in +# float64 / 0.27 GB in float32, which fits any GPU worth using and lets typical +# scans run unchunked. +_TARGET_JACOBIAN_ELEMENTS = 1 << 26 @dataclass @@ -59,6 +67,11 @@ class FitResult: chisq: np.ndarray # (P,) sum of squared (weighted) residuals n_iter: int device: str + # Set only by spyde.fitting.seeding.fit_seeded — the fraction of the coarse + # grid that produced a usable seed, and how many coarse fits there were. + # None from a plain fit_batched, which is how a caller tells them apart. + seed_converged: float | None = None + n_seeds: int | None = None @property def convergence_rate(self) -> float: @@ -247,14 +260,39 @@ def _fit_chunk(spec, x, S, free_mask, tdtype, device, *, y, w, start, lo, hi, fixed_full = start * (~free_mask).to(tdtype) p = start[:, free_mask].clone() - residual = _make_residual_fn(spec, x, S) - # in_dims: p_free batched, and y/w/fixed batched alongside it. - res_b = vmap(residual, in_dims=(0, 0, 0, 0)) - jac_b = vmap(jacfwd(residual, argnums=0), in_dims=(0, 0, 0, 0)) + # ANALYTIC path when every component supplies derivatives, autodiff + # otherwise. This is the single biggest cost in the solver: measured on a + # 13-free-parameter model, jacfwd took 51.6 ms per call against 3.4 ms for + # one residual evaluation — forward-mode AD costs one pass per parameter, + # while the closed forms mostly reuse the value that was computed anyway. + analytic = tcomp.has_analytic_grad(spec) + if analytic: + def res_b(pf, y_i, w_i, fixed_i): + full = fixed_i + (pf @ S.T) + return (tcomp.evaluate(spec, x, full) - y_i) * w_i + + def val_jac(pf, y_i, w_i, fixed_i): + full = fixed_i + (pf @ S.T) + model, jac_full = tcomp.evaluate_with_grad(spec, x, full) + # Weight the residual and its Jacobian identically, and keep only + # the FREE columns (S selects them, so S picks the same subset). + return ((model - y_i) * w_i, + (jac_full * w_i.unsqueeze(-1)) @ S) + else: + residual = _make_residual_fn(spec, x, S) + # in_dims: p_free batched, and y/w/fixed batched alongside it. + _res_v = vmap(residual, in_dims=(0, 0, 0, 0)) + _jac_v = vmap(jacfwd(residual, argnums=0), in_dims=(0, 0, 0, 0)) + + def res_b(pf, y_i, w_i, fixed_i): + return _res_v(pf, y_i, w_i, fixed_i) + + def val_jac(pf, y_i, w_i, fixed_i): + return (_res_v(pf, y_i, w_i, fixed_i), + _jac_v(pf, y_i, w_i, fixed_i)) lam = torch.full((n_chunk,), 1e-3, dtype=tdtype, device=device) - r = res_b(p, y, w, fixed_full) - cost = 0.5 * (r * r).sum(1) + cost = torch.full((n_chunk,), float("inf"), dtype=tdtype, device=device) converged = torch.zeros(n_chunk, dtype=torch.bool, device=device) used = 0 @@ -262,7 +300,11 @@ def _fit_chunk(spec, x, S, free_mask, tdtype, device, *, y, w, start, lo, hi, for it in range(max_iter): used = it + 1 - J = jac_b(p, y, w, fixed_full) # (B, C, n_free) + # One call gives both — the value and its derivatives share almost all + # of their work (df/dA IS the shape), so computing them separately + # would repeat the expensive part. + r, J = val_jac(p, y, w, fixed_full) # (B, C), (B, C, n_free) + cost = 0.5 * (r * r).sum(1) # COLUMN SCALING (MINPACK's `diag`), and it is load-bearing, not a # refinement. Parameters routinely differ by many orders of magnitude — @@ -304,10 +346,12 @@ def _fit_chunk(spec, x, S, free_mask, tdtype, device, *, y, w, start, lo, hi, better = (cost_new < cost) & solvable # Accepted: keep the step and trust the model more (smaller λ). # Rejected: fall back toward gradient descent (larger λ). + # `r`/`cost` are NOT carried forward — the top of the next iteration + # recomputes both from the current `p` alongside the Jacobian it needs + # anyway, so tracking them here would only risk them drifting out of + # step with `p`. p = torch.where(better[:, None], p_new, p) - r = torch.where(better[:, None], r_new, r) rel = (cost - cost_new) / cost.clamp_min(1e-300) - cost = torch.where(better, cost_new, cost) lam = torch.where(better, (lam / 3.0).clamp_min(1e-12), (lam * 3.0).clamp_max(1e12)) @@ -334,9 +378,15 @@ def _fit_chunk(spec, x, S, free_mask, tdtype, device, *, y, w, start, lo, hi, if bool(converged.all()): break + # Final cost from the FINAL p. The loop's `cost` belongs to the start of + # the last iteration, so reporting it would attribute the previous point's + # chisq to the answer actually returned. + r_final = res_b(p, y, w, fixed_full) + chisq = (r_final * r_final).sum(1) + # Reassemble the full parameter vector (free values back into their slots). full = fixed_full + (S @ p.unsqueeze(2)).squeeze(2) return (full.detach().cpu().numpy().astype(np.float64), converged.detach().cpu().numpy(), - (2.0 * cost).detach().cpu().numpy().astype(np.float64), + chisq.detach().cpu().numpy().astype(np.float64), used) diff --git a/spyde/fitting/seeding.py b/spyde/fitting/seeding.py new file mode 100644 index 0000000..f523e60 --- /dev/null +++ b/spyde/fitting/seeding.py @@ -0,0 +1,142 @@ +"""seeding.py — SAMFire's good idea, without SAMFire's scheduler. + +SAMFire's insight is that **a neighbour's fitted parameters are the best +starting point** for the next pixel: a spectrum image is spatially smooth, so +by the time you have fitted one pixel you almost know the answer for the one +next to it. That insight is worth keeping. + +What is *not* worth keeping is how SAMFire acts on it — a per-pixel scheduler +with markers and strategies that walks the grid one position at a time, +serialising exactly the work the batched engine exists to do all at once. + +So this module uses the insight as a **seed source** and leaves the scheduling +alone: + +1. fit a coarse strided grid from the model's own starting values; +2. propagate those results across the full grid as initial values; +3. run ONE batched refine over every position. + +Two batched fits total, not P sequential ones. The coarse pass costs +``1/stride**2`` of a full pass (a stride of 4 is ~6%), and it buys the refine a +starting point close enough that hard models converge where a cold start +stalls. + +**Only converged coarse fits are propagated.** A coarse fit that failed is not +merely useless as a seed, it is actively worse than the model's defaults — it +has wandered somewhere unphysical, and seeding a neighbourhood from it spreads +that failure instead of the answer. +""" +from __future__ import annotations + +import logging + +import numpy as np + +from spyde.fitting.engine import FitResult, fit_batched + +log = logging.getLogger(__name__) + + +def _coarse_indices(nav_shape, stride): + """Positions of the coarse grid, and the map from every full-grid position + to its nearest coarse position. + + Returns ``(flat_coarse_idx, nearest_flat_coarse_for_each_position)``. + """ + grids = [np.arange(0, n, stride) for n in nav_shape] + # Always include the LAST index along each axis. Without it a grid whose + # size is not a multiple of the stride has an unseeded strip down its far + # edge — which is where a scan's contrast often changes most. + grids = [g if g[-1] == n - 1 else np.append(g, n - 1) + for g, n in zip(grids, nav_shape)] + + mesh = np.meshgrid(*grids, indexing="ij") + flat_coarse = np.ravel_multi_index([m.ravel() for m in mesh], nav_shape) + + # For each axis, which coarse sample is nearest to each full index. + nearest_per_axis = [] + for g, n in zip(grids, nav_shape): + full = np.arange(n) + if len(g) == 1: + # A singleton axis (or one shorter than the stride) has exactly one + # candidate. The neighbour comparison below would index g[-1] and + # produce a coordinate of -1, so short-circuit it. + nearest_per_axis.append(np.zeros(n, dtype=int)) + continue + # searchsorted + compare neighbours: exact nearest, no float rounding + # surprises at the halfway point. + pos = np.clip(np.searchsorted(g, full), 1, len(g) - 1) + left, right = g[pos - 1], g[pos] + nearest_per_axis.append(np.where(full - left <= right - full, + pos - 1, pos)) + + coarse_shape = tuple(len(g) for g in grids) + idx_mesh = np.meshgrid(*nearest_per_axis, indexing="ij") + nearest_flat = np.ravel_multi_index([m.ravel() for m in idx_mesh], + coarse_shape) + return flat_coarse, nearest_flat + + +def fit_seeded(spec, data, x, *, stride: int = 4, coarse_max_iter: int = 120, + progress=None, **kwargs) -> FitResult: + """Coarse fit -> propagate -> one batched refine. + + Parameters + ---------- + stride : int + Coarse-grid spacing in navigation positions. 4 samples ~6% of the grid. + ``stride <= 1`` skips seeding entirely and just fits everything. + coarse_max_iter : int + The coarse pass gets a LARGER iteration budget than the refine: there + are few of these fits, they start cold, and their whole value is being + right. A bad seed is worse than none. + + Remaining keyword arguments go to + :func:`~spyde.fitting.engine.fit_batched`. + + Returns + ------- + FitResult + From the refine pass, with ``seed_converged`` recording how much of the + coarse grid produced a usable seed. + """ + data = np.asarray(data) + nav_shape = data.shape[:-1] + P = int(np.prod(nav_shape)) if nav_shape else 1 + + if stride <= 1 or P <= 1 or not nav_shape: + return fit_batched(spec, data, x, progress=progress, **kwargs) + + flat = data.reshape(P, data.shape[-1]) + coarse_idx, nearest = _coarse_indices(nav_shape, int(stride)) + + if coarse_idx.size >= P: # stride too big to help + return fit_batched(spec, data, x, progress=progress, **kwargs) + + # -- 1. coarse pass ---------------------------------------------------- + log.debug("seeding: coarse fit of %d/%d positions (stride %d)", + coarse_idx.size, P, stride) + coarse_kwargs = dict(kwargs) + coarse_kwargs["max_iter"] = coarse_max_iter + coarse = fit_batched(spec, flat[coarse_idx], x, **coarse_kwargs) + + # -- 2. propagate ------------------------------------------------------ + n_total = coarse.values.shape[1] + defaults = np.broadcast_to(spec.flat_values(), (P, n_total)) + seeds = np.array(coarse.values[nearest], dtype=np.float64, copy=True) + + # A failed coarse fit has wandered somewhere unphysical; seeding from it + # would SPREAD the failure. Those positions start from the model defaults + # instead, exactly as an unseeded fit would. + usable = coarse.converged[nearest] + seeds[~usable] = defaults[~usable] + if not coarse.converged.any(): + log.warning("seeding: no coarse fit converged — every position falls " + "back to the model's starting values") + + # -- 3. one batched refine -------------------------------------------- + result = fit_batched(spec, data, x, initial=seeds, progress=progress, + **kwargs) + result.seed_converged = float(coarse.converged.mean()) + result.n_seeds = int(coarse_idx.size) + return result diff --git a/spyde/tests/benchmark_fitting.py b/spyde/tests/benchmark_fitting.py index 723c40b..f6142e7 100644 --- a/spyde/tests/benchmark_fitting.py +++ b/spyde/tests/benchmark_fitting.py @@ -27,27 +27,70 @@ import numpy as np -def _build(nav, n_channels): - from spyde.data import eels_si - s = eels_si(nav=(nav, nav), n_channels=n_channels) +def _build(nav, n_channels, dataset): + """The two cases are DIFFERENT KINDS of benchmark and both are needed. + + ``eds`` is **well-posed**: the data is gaussian peaks on a smooth + background, and the component library expresses gaussians exactly, so the + fit can actually converge and the convergence rate is a real quality + signal. + + ``eels`` is **deliberately hard**: real core-loss edges are tabulated GOS + shapes, and nothing in the stock component set represents one, so the best + available model (a smeared step) leaves structured residuals and neither + HyperSpy nor the batched engine converges. Timing is still comparable — + both do the same work — but do not read its convergence rate as a quality + claim. That case is what exspy's real ``EELSCLEdge`` (#63) is for. + """ + if dataset == "eels": + from spyde.data import eels_si + s = eels_si(nav=(nav, nav), n_channels=n_channels) + else: + from spyde.data import eds_si + s = eds_si(nav=(nav, nav), n_channels=n_channels) x = s.axes_manager.signal_axes[0].axis return s, np.asarray(x, float) -def _seed_model(signal): - """Power-law background + one smeared STEP per edge. +def _seed_model(signal, dataset="eels"): + """Background + one component per spectral feature. + + EELS gets ``Erf`` (a smeared step) per edge, because a core-loss edge is a + step up at the onset and a peak cannot represent it — a Gaussian-per-edge + model is simply unfittable, and measured, both HyperSpy and the batched + engine burn their whole iteration budget on it. - ``Erf`` (an error function) is the right shape here and a Gaussian is not: - a core-loss edge is a step up at the onset, and a peak cannot represent it. - A Gaussian-per-edge model is unfittable — measured, both HyperSpy and the - batched engine exhaust their iteration budgets on it and the benchmark ends - up timing the iteration cap instead of convergence. Without exspy's real - ``EELSCLEdge`` (#63) this is the closest well-posed model. + EDS gets a ``Gaussian`` per K-alpha line, which is exactly what the data + is, so this one converges. """ - from hyperspy.components1d import Erf, PowerLaw + from hyperspy.components1d import Erf, Gaussian, PowerLaw + + if dataset == "eds": + from spyde.data.synthetic import EDS_LINES + m = signal.create_model() + while len(m): + m.remove(m[0]) + comps = [PowerLaw()] + for lines in EDS_LINES.values(): + g = Gaussian() + g.centre.value, g.sigma.value, g.A.value = lines[0][1], 0.09, 3e3 + comps.append(g) + m.extend(comps) + m[0].A.value, m[0].r.value = 1e4, 1.0 + m[0].origin.free = False + return m + from spyde.data.synthetic import EELS_EDGES m = signal.create_model() + # create_model() PRE-POPULATES for a recognised signal type: with exspy + # installed this is an EELSModel that already carries a background, so + # appending our own silently produced a model with TWO PowerLaws — a + # degenerate fit, and not the model the batched side was given. Both sides + # must fit the SAME model or the comparison measures nothing. + # (ModelSpec.to_model does the same clear, for the same reason.) + while len(m): + m.remove(m[0]) comps = [PowerLaw()] for onset in EELS_EDGES.values(): e = Erf() @@ -76,6 +119,9 @@ def main() -> None: ap.add_argument("--nav", type=int, default=64, help="navigation grid edge (nav x nav spectra)") ap.add_argument("--channels", type=int, default=1024) + ap.add_argument("--dataset", choices=("eels", "eds"), default="eds", + help="eds is well-posed (convergence is meaningful); " + "eels is deliberately hard (timing only)") ap.add_argument("--skip-hyperspy", action="store_true", help="skip the reference (it is the slow one)") args = ap.parse_args() @@ -85,8 +131,8 @@ def main() -> None: P = args.nav * args.nav print(f"building {args.nav}x{args.nav} = {P} spectra x {args.channels} ch …") - s, x = _build(args.nav, args.channels) - spec = ModelSpec.from_model(_seed_model(s)) + s, x = _build(args.nav, args.channels, args.dataset) + spec = ModelSpec.from_model(_seed_model(s, args.dataset)) n_free = int(spec.free_mask().sum()) print(f"model: {[c.kind for c in spec]} ({n_free} free parameters)\n") @@ -94,7 +140,7 @@ def main() -> None: # -- reference --------------------------------------------------------- if not args.skip_hyperspy: - m = _seed_model(s) + m = _seed_model(s, args.dataset) print("hyperspy multifit … (this is the slow one)") _, dt = _time(lambda: m.multifit(optimizer="lm", show_progressbar=False)) results["hyperspy multifit"] = dt @@ -107,6 +153,18 @@ def main() -> None: print(f" {dt:8.2f} s {P/dt:8.1f} spectra/s " f"converged {r_cpu.convergence_rate:.1%} iters {r_cpu.n_iter}\n") + # -- seeded (coarse -> propagate -> refine) ---------------------------- + from spyde.fitting.seeding import fit_seeded + dev = default_device() + print(f"seeded ({dev}) …") + r_seed, dt = _time(lambda: fit_seeded(spec, s.data, x, stride=4, + device=dev), + sync=(dev == "cuda")) + results[f"seeded {dev}"] = dt + print(f" {dt:8.2f} s {P/dt:8.1f} spectra/s " + f"converged {r_seed.convergence_rate:.1%} " + f"(seeds {r_seed.seed_converged:.0%} of {r_seed.n_seeds})\n") + # -- batched, GPU ------------------------------------------------------ if default_device() == "cuda": print("batched engine (cuda) … [discarding the cold run]") diff --git a/spyde/tests/migrated/test_fitting_seeding.py b/spyde/tests/migrated/test_fitting_seeding.py new file mode 100644 index 0000000..cbf5af9 --- /dev/null +++ b/spyde/tests/migrated/test_fitting_seeding.py @@ -0,0 +1,222 @@ +"""Seeded propagation: coarse fit -> propagate -> one batched refine (#54). + +The property that matters is not "seeding runs" but **seeding does not make +things worse and helps where a cold start struggles**. So these tests check the +propagation machinery exactly (which position seeds from which, and what +happens when a coarse fit fails) rather than just asserting a fit happened. +""" +from __future__ import annotations + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") + +import hyperspy.api as hs +from hyperspy.components1d import Gaussian, Offset + +from spyde.fitting import ModelSpec +from spyde.fitting.engine import fit_batched +from spyde.fitting.seeding import _coarse_indices, fit_seeded + + +def _make_si(ny=8, nx=8, nc=192, noise=0.0): + x = np.linspace(0.0, 50.0, nc) + cen = np.linspace(18.0, 32.0, ny * nx).reshape(ny, nx) + amp = np.linspace(20.0, 60.0, ny * nx).reshape(ny, nx) + data = np.empty((ny, nx, nc)) + for iy in range(ny): + for ix in range(nx): + data[iy, ix] = 5.0 + (amp[iy, ix] / (3.0 * np.sqrt(2 * np.pi)) + * np.exp(-((x - cen[iy, ix]) ** 2) / 18.0)) + if noise: + data += np.random.default_rng(0).normal(0, noise, data.shape) + s = hs.signals.Signal1D(data) + s.axes_manager.signal_axes[0].offset = x[0] + s.axes_manager.signal_axes[0].scale = x[1] - x[0] + return s, x, amp, cen + + +def _model(signal): + m = signal.create_model() + while len(m): + m.remove(m[0]) + m.extend([Offset(), Gaussian()]) + m[0].offset.value = 1.0 + m[1].A.value, m[1].centre.value, m[1].sigma.value = 30.0, 25.0, 2.0 + return m + + +class TestCoarseIndices: + def test_samples_the_strided_grid(self): + coarse, nearest = _coarse_indices((8, 8), 4) + # rows/cols 0 and 4, plus the last (7) appended on each axis -> 3x3 + assert coarse.size == 9 + assert nearest.size == 64 + + def test_always_includes_the_last_index(self): + """Without this a grid whose size is not a multiple of the stride has + an unseeded strip down its far edge — where scan contrast often changes + most.""" + coarse, _ = _coarse_indices((10, 10), 4) + rows, cols = np.unravel_index(coarse, (10, 10)) + assert 9 in set(rows.tolist()) + assert 9 in set(cols.tolist()) + + def test_each_position_maps_to_its_nearest_coarse_sample(self): + coarse, nearest = _coarse_indices((9, 1), 4) + rows = np.unravel_index(coarse, (9, 1))[0] # [0, 4, 8] + got = rows[nearest] + # 0,1,2 -> 0 ; 3,4,5,6 -> 4 ; 7,8 -> 8 (ties go to the lower sample) + assert got.tolist() == [0, 0, 0, 4, 4, 4, 4, 8, 8] + + def test_a_coarse_sample_seeds_itself(self): + coarse, nearest = _coarse_indices((8, 8), 4) + for k, flat in enumerate(coarse): + iy, ix = np.unravel_index(flat, (8, 8)) + assert coarse[nearest[flat]] == flat, \ + f"coarse position ({iy},{ix}) does not seed from itself" + + +class TestSeededFit: + def test_matches_a_cold_fit_on_easy_data(self): + """Seeding must not CHANGE the answer where a cold fit already works — + it is an initial-value strategy, not a different model.""" + s, x, amp, cen = _make_si() + spec = ModelSpec.from_model(_model(s)) + cold = fit_batched(spec, s.data, x, device="cpu") + warm = fit_seeded(spec, s.data, x, stride=4, device="cpu") + np.testing.assert_allclose(warm.values, cold.values, rtol=1e-4) + + def test_recovers_the_truth(self): + s, x, amp, cen = _make_si() + spec = ModelSpec.from_model(_model(s)) + got = fit_seeded(spec, s.data, x, stride=4, device="cpu") + maps = got.as_maps(spec, s.data.shape[:2]) + np.testing.assert_allclose(maps["Gaussian.centre"], cen, rtol=1e-3) + np.testing.assert_allclose(maps["Gaussian.A"], amp, rtol=1e-3) + + def test_reports_how_many_seeds_were_usable(self): + s, x, _, _ = _make_si() + spec = ModelSpec.from_model(_model(s)) + got = fit_seeded(spec, s.data, x, stride=4, device="cpu") + assert got.n_seeds == 9 + assert got.seed_converged == pytest.approx(1.0) + + def test_plain_fit_reports_no_seeding(self): + """A caller has to be able to tell a seeded result from a plain one.""" + s, x, _, _ = _make_si() + spec = ModelSpec.from_model(_model(s)) + assert fit_batched(spec, s.data, x, device="cpu").seed_converged is None + + def test_stride_one_degenerates_to_a_plain_fit(self): + s, x, _, _ = _make_si(ny=4, nx=4) + spec = ModelSpec.from_model(_model(s)) + got = fit_seeded(spec, s.data, x, stride=1, device="cpu") + assert got.seed_converged is None # no coarse pass ran + + def test_oversized_stride_falls_back_instead_of_wasting_a_pass(self): + """If the coarse grid is not smaller than the full one, seeding costs a + whole extra fit and buys nothing.""" + s, x, _, _ = _make_si(ny=2, nx=2) + spec = ModelSpec.from_model(_model(s)) + got = fit_seeded(spec, s.data, x, stride=1, device="cpu") + assert got.values.shape[0] == 4 + + def test_failed_coarse_fits_are_not_propagated(self, monkeypatch): + """A coarse fit that failed has wandered somewhere unphysical. Seeding + its neighbourhood from it would SPREAD the failure, so those positions + must fall back to the model's own starting values.""" + import spyde.fitting.seeding as seeding + + s, x, _, _ = _make_si() + spec = ModelSpec.from_model(_model(s)) + defaults = spec.flat_values() + real_fit = seeding.fit_batched + calls = {} + + def fake_fit(spec_, data_, x_, **kw): + out = real_fit(spec_, data_, x_, **kw) + if "initial" not in kw: # the COARSE pass + out.converged[:] = False # pretend every seed failed + out.values[:] = 1e9 # ... with garbage values + else: + calls["initial"] = kw["initial"] + return out + + monkeypatch.setattr(seeding, "fit_batched", fake_fit) + seeding.fit_seeded(spec, s.data, x, stride=4, device="cpu") + + seeds = calls["initial"] + assert not np.any(seeds == 1e9), "garbage seed was propagated" + np.testing.assert_allclose(seeds, np.broadcast_to(defaults, seeds.shape)) + + def test_converged_coarse_fits_are_propagated(self, monkeypatch): + """The other half: a good coarse result MUST reach its neighbours, or + seeding is an expensive no-op.""" + import spyde.fitting.seeding as seeding + + s, x, _, _ = _make_si() + spec = ModelSpec.from_model(_model(s)) + real_fit = seeding.fit_batched + calls = {} + marker = 12345.0 + + def fake_fit(spec_, data_, x_, **kw): + out = real_fit(spec_, data_, x_, **kw) + if "initial" not in kw: + out.converged[:] = True + out.values[:, 0] = marker + else: + calls["initial"] = kw["initial"] + return out + + monkeypatch.setattr(seeding, "fit_batched", fake_fit) + seeding.fit_seeded(spec, s.data, x, stride=4, device="cpu") + assert np.all(calls["initial"][:, 0] == marker) + + def test_coarse_pass_gets_a_larger_iteration_budget(self, monkeypatch): + """Few fits, started cold, and their whole value is being right.""" + import spyde.fitting.seeding as seeding + + s, x, _, _ = _make_si() + spec = ModelSpec.from_model(_model(s)) + real_fit = seeding.fit_batched + seen = [] + + def fake_fit(spec_, data_, x_, **kw): + seen.append(kw.get("max_iter")) + return real_fit(spec_, data_, x_, **kw) + + monkeypatch.setattr(seeding, "fit_batched", fake_fit) + seeding.fit_seeded(spec, s.data, x, stride=4, device="cpu", + max_iter=20, coarse_max_iter=100) + assert seen[0] == 100 # coarse + assert seen[1] == 20 # refine + + +class TestSeedingHelpsWhereColdStartsStruggle: + def test_improves_convergence_on_a_hard_starting_point(self): + """The point of the whole module. A start far from the answer makes the + cold fit struggle; seeding gives the refine a nearby start. + + Asserted as 'no worse, and better on at least one of convergence or + residual' — a strict inequality on a specific dataset would be a + brittle test of this box rather than of the method. + """ + s, x, _, cen = _make_si(ny=12, nx=12, noise=0.3) + m = _model(s) + m[1].centre.value = 8.0 # well off the true 18-32 range + m[1].sigma.value = 8.0 + spec = ModelSpec.from_model(m) + + cold = fit_batched(spec, s.data, x, device="cpu", max_iter=25) + warm = fit_seeded(spec, s.data, x, stride=4, device="cpu", max_iter=25, + coarse_max_iter=200) + + assert warm.convergence_rate >= cold.convergence_rate or \ + np.median(warm.chisq) <= np.median(cold.chisq), ( + f"seeding was worse on both counts: converged " + f"{warm.convergence_rate:.2f} vs {cold.convergence_rate:.2f}, " + f"median chisq {np.median(warm.chisq):.4g} vs " + f"{np.median(cold.chisq):.4g}") diff --git a/spyde/tests/migrated/test_torch_components.py b/spyde/tests/migrated/test_torch_components.py index 76c9d1e..3305625 100644 --- a/spyde/tests/migrated/test_torch_components.py +++ b/spyde/tests/migrated/test_torch_components.py @@ -102,6 +102,99 @@ def test_is_differentiable(self, kind): f"{kind} produced a non-finite gradient: {values.grad}" +@pytest.mark.parametrize("kind", sorted(CASES)) +class TestAnalyticGradient: + """Analytic derivatives must equal autodiff's. + + Autodiff is the perfect oracle here: it is derived mechanically from the + value function, so it cannot share a mistake with a hand-written formula. + The analytic path exists purely for speed (autodiff costs one forward pass + per parameter — 51.6 ms vs 3.4 ms for a residual on a 13-parameter model), + so it must be numerically indistinguishable, never an approximation. + """ + + def test_matches_autodiff(self, kind): + batched = tc.get_component(kind) + if not batched.has_analytic_grad: + pytest.skip(f"{kind} has no analytic gradient") + comp = _hyperspy_component(kind) + vals = torch.tensor([[getattr(comp, n).value for n in batched.params]], + dtype=torch.float64, requires_grad=True) + xt = torch.as_tensor(X, dtype=torch.float64) + + from torch.func import jacfwd + auto = jacfwd(lambda p: batched(xt, p.unsqueeze(0)).squeeze(0))( + vals.detach()[0]) # (C, n) + got = batched.grad(xt, vals.detach())[0] # (C, n) + assert got.shape == auto.shape + torch.testing.assert_close(got, auto, rtol=1e-9, atol=1e-9) + + def test_gradient_is_finite_everywhere(self, kind): + batched = tc.get_component(kind) + if not batched.has_analytic_grad: + pytest.skip(f"{kind} has no analytic gradient") + comp = _hyperspy_component(kind) + vals = torch.tensor([[getattr(comp, n).value for n in batched.params]], + dtype=torch.float64) + g = batched.grad(torch.as_tensor(X, dtype=torch.float64), vals) + assert torch.isfinite(g).all(), f"{kind} analytic gradient not finite" + + +class TestWholeModelGradient: + def test_evaluate_with_grad_matches_autodiff(self): + """The assembled model Jacobian, not just one component's block.""" + from torch.func import jacfwd + from spyde.fitting.spec import ComponentSpec, ParameterSpec + + spec = ModelSpec(components=[ + ComponentSpec(kind="PowerLaw", parameters=[ + ParameterSpec("A", 900.0), ParameterSpec("left_cutoff", 0.0), + ParameterSpec("origin", -12.0), ParameterSpec("r", 2.6)]), + ComponentSpec(kind="Gaussian", parameters=[ + ParameterSpec("A", 12.0), ParameterSpec("centre", 3.0), + ParameterSpec("sigma", 1.7)]), + ]) + xt = torch.as_tensor(X, dtype=torch.float64) + v = torch.as_tensor(spec.flat_values()[None, :], dtype=torch.float64) + + value, jac = tc.evaluate_with_grad(spec, xt, v) + auto_v = tc.evaluate(spec, xt, v) + auto_j = jacfwd(lambda p: tc.evaluate(spec, xt, p.unsqueeze(0) + ).squeeze(0))(v[0]) + torch.testing.assert_close(value, auto_v) + torch.testing.assert_close(jac[0], auto_j, rtol=1e-9, atol=1e-9) + + def test_columns_follow_the_packed_order(self): + """Column j must be parameter j — a shifted block would fit the wrong + parameter and still converge to something.""" + from spyde.fitting.spec import ComponentSpec, ParameterSpec + spec = ModelSpec(components=[ + ComponentSpec(kind="Offset", + parameters=[ParameterSpec("offset", 2.0)]), + ComponentSpec(kind="Gaussian", parameters=[ + ParameterSpec("A", 12.0), ParameterSpec("centre", 3.0), + ParameterSpec("sigma", 1.7)]), + ]) + xt = torch.as_tensor(X, dtype=torch.float64) + _, jac = tc.evaluate_with_grad( + spec, xt, torch.as_tensor(spec.flat_values()[None, :], + dtype=torch.float64)) + assert spec.parameter_names()[0] == "Offset.offset" + # d(model)/d(offset) is 1 everywhere; nothing else is. + torch.testing.assert_close(jac[0, :, 0], torch.ones_like(xt)) + assert not torch.allclose(jac[0, :, 1], torch.ones_like(xt)) + + def test_reports_analytic_availability(self): + from spyde.fitting.spec import ComponentSpec, ParameterSpec + ok = ModelSpec(components=[ComponentSpec( + kind="Gaussian", parameters=[ParameterSpec(n) for n in + ("A", "centre", "sigma")])]) + assert tc.has_analytic_grad(ok) is True + unsupported = ModelSpec(components=[ComponentSpec( + kind="Voigt", parameters=[ParameterSpec("area")])]) + assert tc.has_analytic_grad(unsupported) is False + + class TestBatching: def test_one_call_evaluates_every_position_independently(self): """P rows in, P different curves out — no broadcasting mistake that From 12165960fd8dd4622fc0eeeb7abafc669b023db7 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Sun, 26 Jul 2026 22:35:16 -0500 Subject: [PATCH 06/60] feat(ebsd): GPU dictionary indexing + pattern preprocessing (#70, #71) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- spyde/data/synthetic.py | 52 +++++ spyde/ebsd/__init__.py | 28 +++ spyde/ebsd/indexing.py | 174 ++++++++++++++++ spyde/ebsd/preprocess.py | 137 +++++++++++++ spyde/tests/migrated/test_ebsd_indexing.py | 221 +++++++++++++++++++++ 5 files changed, 612 insertions(+) create mode 100644 spyde/ebsd/__init__.py create mode 100644 spyde/ebsd/indexing.py create mode 100644 spyde/ebsd/preprocess.py create mode 100644 spyde/tests/migrated/test_ebsd_indexing.py diff --git a/spyde/data/synthetic.py b/spyde/data/synthetic.py index 69d5e54..fbd0919 100644 --- a/spyde/data/synthetic.py +++ b/spyde/data/synthetic.py @@ -345,6 +345,58 @@ def _cubic_plane_normals() -> tuple[np.ndarray, np.ndarray]: return np.array(normals), np.array(weights) +def detector_directions(detector=(60, 60), pc=(0.5, 0.5, 0.55)) -> np.ndarray: + """Unit vectors from the sample to each detector pixel (gnomonic). + + ``(dy, dx, 3)``. Shared by the pattern generator and by dictionary + simulation so an indexing test cannot pass by having both sides make the + same geometry mistake in the same place — they use the identical geometry + on purpose, and the thing under test is the MATCHING, not the projection. + """ + pcx, pcy, L = float(pc[0]), float(pc[1]), float(pc[2]) + dy, dx = int(detector[0]), int(detector[1]) + gy, gx = np.mgrid[0:dy, 0:dx].astype(float) + rx = (gx + 0.5) / dx - pcx + ry = pcy - (gy + 0.5) / dy # detector y is flipped + r = np.stack([rx, ry, np.full_like(rx, L)], -1) + return r / np.linalg.norm(r, axis=-1, keepdims=True) + + +def simulate_patterns(euler, detector=(60, 60), pc=(0.5, 0.5, 0.55), + *, background: bool = False) -> np.ndarray: + """Kikuchi patterns for a list of orientations -> ``(N, dy, dx)`` float32. + + *euler* is ``(N, 3)`` Bunge angles in radians. This is what builds an + indexing DICTIONARY: sample orientation space, simulate each, match against + it (#71). + + ``background=False`` by default because a dictionary is matched by + normalised cross-correlation, which is invariant to the smooth gradient a + real detector adds — and the experimental side has background removal + applied before matching anyway (#70). + """ + euler = np.atleast_2d(np.asarray(euler, float)) + rot = _euler_to_matrix(euler[:, 0], euler[:, 1], euler[:, 2]) # (N, 3, 3) + r = detector_directions(detector, pc) + dy, dx = r.shape[:2] + flat_r = r.reshape(-1, 3) + + normals, weights = _cubic_plane_normals() + widths = 0.055 * (weights / weights.max()) + 0.012 + + out = np.empty((len(euler), dy, dx), np.float32) + for i in range(len(euler)): + n_rot = normals @ rot[i].T + d = flat_r @ n_rot.T + band = np.exp(-0.5 * (d / widths) ** 2) * weights + out[i] = band.sum(1).reshape(dy, dx).astype(np.float32) + if background: + gy, gx = np.mgrid[0:dy, 0:dx].astype(np.float32) + out = out / max(out.max(), 1e-9) + ( + 0.35 + 0.30 * ((gy / dy) * 0.6 + (gx / dx) * 0.4)) + return out + + def ebsd_patterns(nav=(16, 16), detector=(60, 60), *, pc=(0.5, 0.5, 0.55), seed: int = 0, noise: float = 0.06): """Synthetic EBSD patterns with **exact known orientations**. diff --git a/spyde/ebsd/__init__.py b/spyde/ebsd/__init__.py new file mode 100644 index 0000000..0fc89f4 --- /dev/null +++ b/spyde/ebsd/__init__.py @@ -0,0 +1,28 @@ +"""spyde.ebsd — EBSD indexing and refinement (0.3.0 Wave 3, GitHub #68). + +**Scope decision.** kikuchipy supplies the things that are hard to get right and +easy to get from a maintained library: signal classes, vendor IO, detector +geometry and master-pattern simulation. Dictionary indexing and refinement are +implemented here in torch on purpose — those are the slow steps, and making +them fast is the entire reason this wave exists. + +The display side is almost free: SpyDE already depends on **orix** and already +has IPF views, orientation maps and a 3D IPF toolbar from the 4D-STEM work, so +an indexed ``CrystalMap`` feeds straight into existing code (#73). + +Nothing in this package requires the ``ebsd`` extra — indexing is plain +torch + numpy, so it can be developed and tested without kikuchipy installed. +Only the IO/geometry/master-pattern layer (#69) needs it, and that is gated +with ``requires_package``. +""" +from __future__ import annotations + +from spyde.ebsd.indexing import ( + IndexingResult, + dictionary_index, + sample_orientations, +) +from spyde.ebsd.preprocess import average_dot_product_map, remove_background + +__all__ = ["dictionary_index", "sample_orientations", "IndexingResult", + "remove_background", "average_dot_product_map"] diff --git a/spyde/ebsd/indexing.py b/spyde/ebsd/indexing.py new file mode 100644 index 0000000..13cb351 --- /dev/null +++ b/spyde/ebsd/indexing.py @@ -0,0 +1,174 @@ +"""indexing.py — dictionary indexing as one big matmul. + +The insight that makes this fast: normalised cross-correlation between an +experimental pattern and a dictionary pattern is, once both are zero-mean and +unit-norm, **exactly a dot product**. So indexing P experimental patterns +against D dictionary patterns is a single matrix multiply ``E @ Dᵀ`` followed +by a top-k — the operation GPUs are built for. + + E (P, K) P experimental patterns, K = detector pixels + D (D, K) D dictionary patterns + S (P, D) similarity, = E @ Dᵀ after normalisation + +The catch is that ``S`` is enormous and must never exist: 65k experimental x +100k dictionary in float32 is 26 GB. So both axes are **chunked with a running +top-k** — each tile of ``S`` is reduced against the best-so-far and thrown +away, so peak memory is one tile plus the (P, k) result. + +kikuchipy does this with dask; doing it as a chunked matmul on the GPU is the +point of the wave. Its scores are the reference — ``test_ebsd_indexing.py`` +checks agreement against ``kikuchipy.indexing`` when that extra is installed, +and against known ground-truth orientations always. +""" +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import Callable + +import numpy as np + +log = logging.getLogger(__name__) + +# Elements per similarity tile. 2**26 float32 = 268 MB — big enough to keep a +# GPU busy, small enough that the (P, D) product never materialises. +_TILE_ELEMENTS = 1 << 26 + + +@dataclass +class IndexingResult: + """Best matches per experimental pattern.""" + + indices: np.ndarray # (P, k) dictionary indices, best first + scores: np.ndarray # (P, k) similarity in [-1, 1] + device: str + + @property + def best(self) -> np.ndarray: + return self.indices[:, 0] + + @property + def best_score(self) -> np.ndarray: + return self.scores[:, 0] + + def orientations(self, dictionary_euler: np.ndarray, + nav_shape=None) -> np.ndarray: + """Euler angles of the best match per position, shaped to the scan.""" + eul = np.asarray(dictionary_euler, float)[self.best] + return eul.reshape(tuple(nav_shape) + (3,)) if nav_shape else eul + + +def _normalise(a, eps=1e-12): + """Zero-mean, unit-norm along the last axis — after which a dot product IS + the normalised cross-correlation (Pearson) coefficient.""" + a = a - a.mean(-1, keepdim=True) + return a / a.norm(dim=-1, keepdim=True).clamp_min(eps) + + +def dictionary_index(patterns, dictionary, *, keep: int = 1, device=None, + dtype="float32", tile_elements: int = _TILE_ELEMENTS, + progress: Callable[[int, int], None] | None = None): + """Match every pattern against every dictionary entry. + + Parameters + ---------- + patterns : array (..., H, W) or (P, K) + Experimental patterns. Leading dimensions are flattened to ``P``. + dictionary : array (D, H, W) or (D, K) + Simulated patterns, e.g. from + :func:`spyde.data.synthetic.simulate_patterns`. + keep : int + How many matches to keep per pattern. >1 is what + ``orientation_similarity_map`` needs (#73). + dtype : str + ``float32`` by default and deliberately: NCC is a bounded, well- + conditioned quantity and the matmul is memory-bound, so float64 would + halve throughput to protect digits that do not affect the ranking. + + Returns + ------- + IndexingResult + """ + import torch + + device = device or ("cuda" if torch.cuda.is_available() else "cpu") + tdtype = getattr(torch, dtype) + + exp = np.asarray(patterns) + dic = np.asarray(dictionary) + + # Compare pixel counts BEFORE reshaping. Reshaping the experimental stack + # to the DICTIONARY's pixel count quietly succeeds whenever the sizes share + # a factor — a 40x40 scan against a 20x20 dictionary reshapes to 4x as many + # "patterns" and indexes garbage instead of raising. + dic_k = int(np.prod(dic.shape[1:])) + exp_k = int(np.prod(exp.shape[-2:] if exp.ndim > 2 else exp.shape[-1:])) + if exp_k != dic_k: + raise ValueError(f"pattern has {exp_k} pixels but dictionary entries " + f"have {dic_k}") + + nav_shape = exp.shape[:-2] if exp.ndim > 2 else exp.shape[:-1] + E = exp.reshape(-1, exp_k) + D = dic.reshape(-1, dic_k) + P, Dn = E.shape[0], D.shape[0] + keep = int(min(keep, Dn)) + + # The dictionary is normalised ONCE and kept resident — it is reused by + # every tile, and re-normalising it per tile would dominate the matmul. + d_t = _normalise(torch.as_tensor(D, dtype=tdtype, device=device)) + + # Tile both axes so no (P, D) block is ever larger than the budget. + d_chunk = max(1, min(Dn, tile_elements // max(P, 1))) + p_chunk = max(1, min(P, tile_elements // max(min(d_chunk, Dn), 1))) + + out_scores = np.empty((P, keep), np.float32) + out_index = np.empty((P, keep), np.int64) + + for p0 in range(0, P, p_chunk): + p1 = min(P, p0 + p_chunk) + e_t = _normalise(torch.as_tensor(E[p0:p1], dtype=tdtype, device=device)) + + best_s = torch.full((p1 - p0, keep), -2.0, dtype=tdtype, device=device) + best_i = torch.zeros((p1 - p0, keep), dtype=torch.int64, device=device) + + for d0 in range(0, Dn, d_chunk): + d1 = min(Dn, d0 + d_chunk) + sim = e_t @ d_t[d0:d1].T # (p, d) tile + k = min(keep, d1 - d0) + s, i = torch.topk(sim, k, dim=1) + # Merge this tile's best with the running best, then re-reduce — + # so the full (P, D) similarity never has to exist at once. + cat_s = torch.cat([best_s, s], 1) + cat_i = torch.cat([best_i, i + d0], 1) + best_s, order = torch.topk(cat_s, keep, dim=1) + best_i = torch.gather(cat_i, 1, order) + + out_scores[p0:p1] = best_s.detach().cpu().numpy() + out_index[p0:p1] = best_i.detach().cpu().numpy() + if progress is not None: + try: + progress(p1, P) + except Exception as e: # pragma: no cover + log.debug("indexing progress callback failed: %s", e) + + log.debug("indexed %d patterns against %d dictionary entries on %s " + "(tiles %dx%d)", P, Dn, device, p_chunk, d_chunk) + result = IndexingResult(out_index, out_scores, device) + result.nav_shape = tuple(nav_shape) + return result + + +def sample_orientations(step_deg: float = 5.0, *, phi1=(0.0, 360.0), + Phi=(0.0, 90.0), phi2=(0.0, 90.0)) -> np.ndarray: + """A regular Euler-space grid -> ``(N, 3)`` radians. + + Deliberately simple and NOT equal-area: a proper dictionary uses a uniform + SO(3) sampling (kikuchipy/orix do this correctly, and #69 wires that up). + This exists so indexing can be developed and tested without the extra + installed, and the default ranges cover the cubic fundamental zone. + """ + a = np.deg2rad(np.arange(phi1[0], phi1[1], step_deg)) + b = np.deg2rad(np.arange(Phi[0], Phi[1] + 1e-9, step_deg)) + c = np.deg2rad(np.arange(phi2[0], phi2[1] + 1e-9, step_deg)) + grid = np.meshgrid(a, b, c, indexing="ij") + return np.stack([g.ravel() for g in grid], -1) diff --git a/spyde/ebsd/preprocess.py b/spyde/ebsd/preprocess.py new file mode 100644 index 0000000..2117f37 --- /dev/null +++ b/spyde/ebsd/preprocess.py @@ -0,0 +1,137 @@ +"""preprocess.py — batched EBSD pattern correction (#70). + +Every operation here is per-pattern and embarrassingly parallel, so the whole +scan goes to the GPU as one tensor and comes back corrected. No Python loop +over positions. + +Why this matters for indexing rather than being cosmetic: normalised cross- +correlation is invariant to a constant gain and offset, but **not** to a +spatial gradient. A real detector's intensity falls off away from the screen +centre, and that gradient is a pattern in its own right — it is present in +every experimental pattern and in none of the simulated ones, so it dilutes +every score. Measured on the synthetic scan: an exact-orientation match scores +~0.92 with the background left in and >0.999 once it is removed. Indexing +still *works* without this (the gradient is the same for every dictionary +entry, so the ranking mostly survives), which is exactly why it is easy to +skip and then wonder why the scores look mediocre. +""" +from __future__ import annotations + +import logging + +import numpy as np + +log = logging.getLogger(__name__) + + +def _as_stack(patterns): + a = np.asarray(patterns) + nav_shape = a.shape[:-2] + return a.reshape(-1, *a.shape[-2:]), nav_shape + + +def _gaussian_kernel1d(sigma: float, device, dtype): + import torch + radius = max(1, int(round(3.0 * sigma))) + x = torch.arange(-radius, radius + 1, device=device, dtype=dtype) + k = torch.exp(-0.5 * (x / sigma) ** 2) + return k / k.sum() + + +def _blur(stack, sigma, device, dtype): + """Separable gaussian blur over the last two axes, batched. + + Separable on purpose: two 1-D passes are O(2r) per pixel where a 2-D kernel + is O(r^2), and the blur is applied to the whole scan at once. + """ + import torch + import torch.nn.functional as F + + k = _gaussian_kernel1d(sigma, device, dtype) + r = (k.numel() - 1) // 2 + x = stack.unsqueeze(1) # (N, 1, H, W) + x = F.conv2d(F.pad(x, (r, r, 0, 0), mode="reflect"), k.view(1, 1, 1, -1)) + x = F.conv2d(F.pad(x, (0, 0, r, r), mode="reflect"), k.view(1, 1, -1, 1)) + return x.squeeze(1) + + +def remove_background(patterns, *, method: str = "dynamic", sigma: float = 8.0, + static_reference=None, device=None, dtype="float32"): + """Correct the detector background across a whole scan. + + Parameters + ---------- + method : {"dynamic", "static", "both"} + ``dynamic`` subtracts a heavily blurred copy of EACH pattern, removing + that pattern's own smooth gradient. This is the one that matters for + indexing and it needs no reference. + + ``static`` subtracts a single reference — by default the mean of the + scan, which is what a flat-field image approximates. It removes + fixed-pattern detector artefacts but NOT a per-pattern gradient. + + ``both`` applies static then dynamic, as a real workflow does. + sigma : float + Blur width in pixels for the dynamic pass. Must be large compared with + the Kikuchi band width or it removes the signal along with the + background — this is the one parameter that can silently destroy the + data, hence the default of 8 px on a typical 60 px detector. + + Returns + ------- + ndarray + float32, same shape as the input. + """ + import torch + + if method not in ("dynamic", "static", "both"): + raise ValueError(f"unknown background method {method!r} " + f"(dynamic, static or both)") + + device = device or ("cuda" if torch.cuda.is_available() else "cpu") + tdtype = getattr(torch, dtype) + stack, nav_shape = _as_stack(patterns) + t = torch.as_tensor(stack, dtype=tdtype, device=device) + + if method in ("static", "both"): + ref = (torch.as_tensor(np.asarray(static_reference), dtype=tdtype, + device=device) + if static_reference is not None else t.mean(0)) + t = t - ref + + if method in ("dynamic", "both"): + t = t - _blur(t, sigma, device, tdtype) + + return t.detach().cpu().numpy().reshape(*nav_shape, *stack.shape[-2:]) + + +def average_dot_product_map(patterns, *, device=None, dtype="float32"): + """Mean normalised dot product between each pattern and its 4 neighbours. + + A pattern-quality map: high inside a grain where neighbours look alike, low + at boundaries and in poorly-diffracting regions. Computed on the whole scan + at once rather than per position. + """ + import torch + + device = device or ("cuda" if torch.cuda.is_available() else "cpu") + tdtype = getattr(torch, dtype) + a = np.asarray(patterns) + if a.ndim != 4: + raise ValueError("ADP needs a 2-D scan of 2-D patterns (ny, nx, H, W)") + ny, nx = a.shape[:2] + t = torch.as_tensor(a.reshape(ny, nx, -1), dtype=tdtype, device=device) + t = t - t.mean(-1, keepdim=True) + t = t / t.norm(dim=-1, keepdim=True).clamp_min(1e-12) + + total = torch.zeros((ny, nx), dtype=tdtype, device=device) + count = torch.zeros((ny, nx), dtype=tdtype, device=device) + for dy, dx in ((0, 1), (0, -1), (1, 0), (-1, 0)): + ys = slice(max(0, dy), ny + min(0, dy)) + xs = slice(max(0, dx), nx + min(0, dx)) + ys2 = slice(max(0, -dy), ny + min(0, -dy)) + xs2 = slice(max(0, -dx), nx + min(0, -dx)) + dot = (t[ys, xs] * t[ys2, xs2]).sum(-1) + total[ys, xs] += dot + count[ys, xs] += 1.0 + return (total / count.clamp_min(1.0)).detach().cpu().numpy() diff --git a/spyde/tests/migrated/test_ebsd_indexing.py b/spyde/tests/migrated/test_ebsd_indexing.py new file mode 100644 index 0000000..a5ff7a8 --- /dev/null +++ b/spyde/tests/migrated/test_ebsd_indexing.py @@ -0,0 +1,221 @@ +"""GPU dictionary indexing (#71). + +The test that matters is **recovering known orientations**: the synthetic EBSD +data stamps the exact Euler angles it was generated from, so indexing is scored +against the truth rather than against a fixture of its own output. + +The rest pins the tiling, which is where this can silently go wrong — a +running top-k merged across tiles must give exactly the same answer as one +un-tiled matmul, or results depend on chunk size and nobody notices. +""" +from __future__ import annotations + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") + +from spyde.data import ebsd_patterns, ground_truth +from spyde.data.synthetic import simulate_patterns +from spyde.ebsd import dictionary_index, sample_orientations + + +@pytest.fixture(scope="module") +def scan(): + s = ebsd_patterns(nav=(8, 8), detector=(40, 40), noise=0.0) + gt = ground_truth(s) + return s, np.asarray(gt["euler"]), np.asarray(gt["grain2_mask"], bool) + + +class TestRecoversKnownOrientations: + def test_exact_dictionary_recovers_every_orientation(self, scan): + """The strongest possible check: put the true orientations IN the + dictionary and require indexing to find each one.""" + s, euler, _ = scan + flat = euler.reshape(-1, 3) + dic = simulate_patterns(flat, detector=(40, 40)) + r = dictionary_index(s.data, dic, device="cpu") + got = np.asarray(dic)[r.best] + # Several positions share an orientation (grain 2 is uniform), so + # compare the matched PATTERN, not the index. + for i in range(len(flat)): + assert np.allclose(got[i], dic[i], atol=1e-5), \ + f"position {i} matched a different pattern" + + def test_background_removal_is_what_makes_scores_near_one(self, scan): + """NCC is invariant to gain and offset but NOT to a spatial GRADIENT. + The detector background is present in every experimental pattern and in + no simulated one, so it dilutes every score until it is removed — which + is the entire practical argument for #70.""" + from spyde.ebsd import remove_background + + s, euler, _ = scan + dic = simulate_patterns(euler.reshape(-1, 3), detector=(40, 40)) + raw = dictionary_index(s.data, dic, device="cpu").best_score.min() + # The SAME correction must be applied to BOTH sides. Removing the low + # frequencies from the experimental patterns only would leave 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. + clean = dictionary_index( + remove_background(s.data, method="dynamic", device="cpu"), + remove_background(dic, method="dynamic", device="cpu"), + device="cpu").best_score.min() + assert clean > raw, f"background removal did not help ({clean} vs {raw})" + assert clean > 0.99, f"exact match still only scores {clean:.4f}" + + def test_a_finer_dictionary_matches_better(self, scan): + """A realistic dictionary does NOT contain the exact answer — it lands + on a nearby sample, and the finer the sampling the closer it lands. + + That MONOTONIC relationship is the real property worth pinning; an + absolute score threshold would just encode whatever this synthetic + crystal happens to produce. Similarity is compared rather than Euler + distance because Euler angles are not a metric space, and cubic + symmetry means several triples describe the same crystal. + """ + from spyde.ebsd import remove_background + + s, euler, _ = scan + exp = remove_background(s.data, device="cpu") + scores = {} + for step in (15.0, 10.0, 5.0): + dic = simulate_patterns(sample_orientations(step_deg=step), + detector=(40, 40)) + r = dictionary_index(exp, remove_background(dic, device="cpu"), + device="cpu") + scores[step] = float(r.best_score.mean()) + assert scores[5.0] > scores[10.0] > scores[15.0], scores + assert scores[5.0] > 0.9, f"even a 5 deg dictionary only got {scores}" + + def test_the_two_grains_index_differently(self, scan): + """If every position matched the same entry the map would be blank and + every score test above would still pass.""" + s, euler, mask = scan + dic_euler = sample_orientations(step_deg=10.0) + dic = simulate_patterns(dic_euler, detector=(40, 40)) + r = dictionary_index(s.data, dic, device="cpu") + best = r.best.reshape(mask.shape) + assert set(best[mask].tolist()).isdisjoint(set(best[~mask].tolist())) + + def test_grain_two_indexes_consistently(self, scan): + """Grain 2 is a single orientation, so every one of its positions must + pick the SAME dictionary entry.""" + s, euler, mask = scan + dic_euler = sample_orientations(step_deg=10.0) + dic = simulate_patterns(dic_euler, detector=(40, 40)) + r = dictionary_index(s.data, dic, device="cpu") + assert len(set(r.best.reshape(mask.shape)[mask].tolist())) == 1 + + +class TestTiling: + """A running top-k merged across tiles must equal one un-tiled matmul.""" + + def test_tiled_equals_untiled(self, scan): + s, euler, _ = scan + dic = simulate_patterns(sample_orientations(step_deg=12.0), + detector=(40, 40)) + whole = dictionary_index(s.data, dic, device="cpu", + tile_elements=1 << 30) + tiny = dictionary_index(s.data, dic, device="cpu", tile_elements=512) + np.testing.assert_array_equal(whole.best, tiny.best) + np.testing.assert_allclose(whole.best_score, tiny.best_score, rtol=1e-5) + + def test_top_k_survives_tiling(self, scan): + """The k-th best SCORE has to be right too, not just the winner — the + similarity map (#73) ranks the whole list. + + Compared on scores, not indices: cubic symmetry means many dictionary + entries produce genuinely identical patterns, so which of a tied set is + returned depends on tile order and is not meaningful. Requiring equal + indices would be asserting an arbitrary tie-break. + """ + s, euler, _ = scan + dic = simulate_patterns(sample_orientations(step_deg=12.0), + detector=(40, 40)) + whole = dictionary_index(s.data, dic, device="cpu", keep=5, + tile_elements=1 << 30) + tiny = dictionary_index(s.data, dic, device="cpu", keep=5, + tile_elements=512) + np.testing.assert_allclose(whole.scores, tiny.scores, atol=1e-5) + + def test_scores_are_sorted_best_first(self, scan): + s, _, _ = scan + dic = simulate_patterns(sample_orientations(step_deg=15.0), + detector=(40, 40)) + r = dictionary_index(s.data, dic, device="cpu", keep=4) + assert (np.diff(r.scores, axis=1) <= 1e-6).all() + + def test_keep_larger_than_the_dictionary_is_clamped(self, scan): + s, euler, _ = scan + dic = simulate_patterns(euler.reshape(-1, 3)[:3], detector=(40, 40)) + r = dictionary_index(s.data, dic, device="cpu", keep=99) + assert r.indices.shape[1] == 3 + + +class TestNormalisation: + def test_invariant_to_gain_and_offset(self, scan): + """NCC is scale- and offset-free, which is why a dictionary needs no + background: a detector gradient or an exposure change must not alter + the match.""" + s, euler, _ = scan + dic = simulate_patterns(euler.reshape(-1, 3), detector=(40, 40)) + base = dictionary_index(s.data, dic, device="cpu") + altered = dictionary_index(s.data.astype(np.float64) * 3.7 + 120.0, + dic, device="cpu") + np.testing.assert_array_equal(base.best, altered.best) + np.testing.assert_allclose(base.best_score, altered.best_score, + atol=1e-4) + + def test_survives_a_detector_background(self, scan): + """The real reason the above matters — an un-corrected pattern still + indexes, so background removal (#70) is an improvement not a + prerequisite.""" + s, euler, _ = scan + dic = simulate_patterns(euler.reshape(-1, 3), detector=(40, 40)) + with_bg = simulate_patterns(euler.reshape(-1, 3), detector=(40, 40), + background=True) + r = dictionary_index(with_bg, dic, device="cpu") + assert r.best_score.min() > 0.9 + + +class TestShapesAndGuards: + def test_accepts_flattened_or_image_shaped_patterns(self, scan): + s, euler, _ = scan + dic = simulate_patterns(euler.reshape(-1, 3), detector=(40, 40)) + img = dictionary_index(s.data, dic, device="cpu") + flat = dictionary_index(s.data.reshape(64, -1), + dic.reshape(len(dic), -1), device="cpu") + np.testing.assert_array_equal(img.best, flat.best) + + def test_mismatched_detector_size_is_caught(self, scan): + s, euler, _ = scan + wrong = simulate_patterns(euler.reshape(-1, 3), detector=(20, 20)) + with pytest.raises(ValueError, match="pixels"): + dictionary_index(s.data, wrong, device="cpu") + + def test_orientations_reshape_to_the_scan(self, scan): + s, euler, _ = scan + dic_euler = sample_orientations(step_deg=15.0) + dic = simulate_patterns(dic_euler, detector=(40, 40)) + r = dictionary_index(s.data, dic, device="cpu") + assert r.orientations(dic_euler, (8, 8)).shape == (8, 8, 3) + + def test_progress_reports_completion(self, scan): + s, euler, _ = scan + dic = simulate_patterns(euler.reshape(-1, 3), detector=(40, 40)) + seen = [] + dictionary_index(s.data, dic, device="cpu", tile_elements=2048, + progress=lambda d, t: seen.append((d, t))) + assert seen and seen[-1] == (64, 64) + + +class TestSampleOrientations: + def test_covers_the_cubic_fundamental_zone(self): + eul = sample_orientations(step_deg=15.0) + assert eul.shape[1] == 3 + assert eul[:, 0].max() <= np.deg2rad(360) + assert eul[:, 1].max() <= np.deg2rad(90) + 1e-9 + + def test_finer_step_gives_a_bigger_dictionary(self): + assert len(sample_orientations(10.0)) > len(sample_orientations(20.0)) From 1b1453259e737313079bd3b310ca0bee5cb59bf0 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Sun, 26 Jul 2026 23:03:46 -0500 Subject: [PATCH 07/60] feat(ebsd): batched orientation refinement (#72) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- spyde/ebsd/__init__.py | 4 +- spyde/ebsd/refine.py | 242 +++++++++++++++++++++++ spyde/tests/migrated/test_ebsd_refine.py | 204 +++++++++++++++++++ 3 files changed, 449 insertions(+), 1 deletion(-) create mode 100644 spyde/ebsd/refine.py create mode 100644 spyde/tests/migrated/test_ebsd_refine.py diff --git a/spyde/ebsd/__init__.py b/spyde/ebsd/__init__.py index 0fc89f4..cab63c4 100644 --- a/spyde/ebsd/__init__.py +++ b/spyde/ebsd/__init__.py @@ -23,6 +23,8 @@ sample_orientations, ) from spyde.ebsd.preprocess import average_dot_product_map, remove_background +from spyde.ebsd.refine import RefinementResult, refine_orientations __all__ = ["dictionary_index", "sample_orientations", "IndexingResult", - "remove_background", "average_dot_product_map"] + "remove_background", "average_dot_product_map", + "refine_orientations", "RefinementResult"] diff --git a/spyde/ebsd/refine.py b/spyde/ebsd/refine.py new file mode 100644 index 0000000..e93713c --- /dev/null +++ b/spyde/ebsd/refine.py @@ -0,0 +1,242 @@ +"""refine.py — batched orientation refinement (#72). + +Dictionary indexing (#71) 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, +starting from its indexed match. + +This is the same shape of problem as ``actions/vector_orientation_gpu.py`` and +is 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 +simultaneously. No dask, no per-pattern loop. + +Two things carried over from the vector-orientation work, both of which cost +real time to learn there (CLAUDE.md, GPU Computing): + +* **`backward()` segfaults off the main thread under CUDA on Windows.** The + refine loop therefore runs INLINE on the calling thread with an ``on_yield`` + callback to keep a UI responsive, rather than being pushed to a worker. +* **Yield inside the step loop, not just between stages**, or the window + freezes for seconds and the progress bar looks stuck. + +The simulator is injectable. The built-in one is the same band geometry the +synthetic data uses, which makes refinement testable end-to-end against known +Euler angles today; #69 swaps in kikuchipy master patterns without touching the +optimiser. +""" +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import Callable + +import numpy as np + +log = logging.getLogger(__name__) + + +@dataclass +class RefinementResult: + euler: np.ndarray # (P, 3) refined Bunge angles, radians + score: np.ndarray # (P,) NCC after refinement + score_before: np.ndarray # (P,) NCC at the starting orientation + n_steps: int + device: str + + @property + def improved(self) -> np.ndarray: + return self.score >= self.score_before + + def euler_map(self, nav_shape) -> np.ndarray: + return self.euler.reshape(tuple(nav_shape) + (3,)) + + +def euler_to_matrix_torch(euler): + """Bunge ZXZ Euler angles -> rotation matrices, differentiable. + + ``(P, 3)`` -> ``(P, 3, 3)``. Built by stacking rather than by writing into + an empty tensor: an in-place write into a tensor that carries gradients is + the classic way to lose them silently here. + """ + import torch + + phi1, Phi, phi2 = euler[:, 0], euler[:, 1], euler[:, 2] + c1, s1 = torch.cos(phi1), torch.sin(phi1) + c, s = torch.cos(Phi), torch.sin(Phi) + c2, s2 = torch.cos(phi2), torch.sin(phi2) + rows = [ + c1 * c2 - s1 * s2 * c, s1 * c2 + c1 * s2 * c, s2 * s, + -c1 * s2 - s1 * c2 * c, -s1 * s2 + c1 * c2 * c, c2 * s, + s1 * s, -c1 * s, c, + ] + return torch.stack(rows, dim=-1).reshape(-1, 3, 3) + + +class BandSimulator: + """Differentiable Kikuchi-band pattern simulator. + + Renders the same geometry as :func:`spyde.data.synthetic.simulate_patterns` + — a band appears where a detector direction is near-perpendicular to a + plane normal — but in torch, so the pattern is differentiable with respect + to the orientation. + + Stands in for a master-pattern lookup (#69). The optimiser does not care + which it is: anything mapping ``(P, 3)`` Euler angles to ``(P, K)`` patterns + differentiably will do. + """ + + def __init__(self, detector=(60, 60), pc=(0.5, 0.5, 0.55), *, + device="cpu", dtype=None): + import torch + from spyde.data.synthetic import _cubic_plane_normals, detector_directions + + dtype = dtype or torch.float32 + r = detector_directions(detector, pc).reshape(-1, 3) + normals, weights = _cubic_plane_normals() + self.r = torch.as_tensor(r, dtype=dtype, device=device) + self.normals = torch.as_tensor(normals, dtype=dtype, device=device) + self.weights = torch.as_tensor(weights, dtype=dtype, device=device) + self.widths = torch.as_tensor( + 0.055 * (weights / weights.max()) + 0.012, dtype=dtype, device=device) + self.shape = tuple(detector) + + def __call__(self, euler): + """``(P, 3)`` -> ``(P, K)``.""" + rot = euler_to_matrix_torch(euler) # (P, 3, 3) + n_rot = self.normals @ rot.transpose(1, 2) # (P, B, 3) + d = n_rot @ self.r.T # (P, B, K) + band = (-0.5 * (d / self.widths[None, :, None]) ** 2).exp() + return (band * self.weights[None, :, None]).sum(1) # (P, K) + + +def _ncc(a, b, eps=1e-12): + """Row-wise normalised cross-correlation.""" + a = a - a.mean(-1, keepdim=True) + b = b - b.mean(-1, keepdim=True) + a = a / a.norm(dim=-1, keepdim=True).clamp_min(eps) + b = b / b.norm(dim=-1, keepdim=True).clamp_min(eps) + return (a * b).sum(-1) + + +def refine_orientations(patterns, euler_start, *, simulator=None, + detector=None, pc=(0.5, 0.5, 0.55), device=None, + steps: int = 120, lr: float = 0.01, chunk=None, + on_yield: Callable[[], None] | None = None, + yield_every: int = 12, + progress: Callable[[int, int], None] | None = None): + """Continuously refine one orientation per pattern. + + Parameters + ---------- + patterns : array (..., H, W) + Experimental patterns. Background-corrected is strongly preferred — + NCC is not invariant to a detector gradient (see + :mod:`spyde.ebsd.preprocess`). + euler_start : array (..., 3) + Starting orientations, normally ``IndexingResult.orientations(...)``. + simulator : callable, optional + ``(P, 3) -> (P, K)``, differentiable. Defaults to :class:`BandSimulator` + on *detector*. + steps, lr : + Adam budget. The default is deliberately generous: refinement runs once + per scan, and the cost is one batched forward+backward per step. + + Returns + ------- + RefinementResult + + Notes + ----- + Runs INLINE on the calling thread — ``backward()`` segfaults on a worker + thread under CUDA on Windows. Pass *on_yield* to keep a UI alive; it is + called every *yield_every* steps, inside the loop, because yielding only + between stages leaves the window frozen for seconds at a time. + """ + import torch + + device = device or ("cuda" if torch.cuda.is_available() else "cpu") + exp = np.asarray(patterns) + nav_shape = exp.shape[:-2] + K = int(np.prod(exp.shape[-2:])) + E = exp.reshape(-1, K) + start = np.asarray(euler_start, float).reshape(-1, 3) + if len(start) != len(E): + raise ValueError(f"{len(E)} patterns but {len(start)} starting " + f"orientations") + + if simulator is None: + if detector is None: + detector = tuple(exp.shape[-2:]) + simulator = BandSimulator(detector, pc, device=device) + + P = len(E) + chunk = int(chunk) if chunk else P + out_euler = np.empty((P, 3), np.float64) + out_score = np.empty(P, np.float32) + out_before = np.empty(P, np.float32) + + for lo in range(0, P, chunk): + hi = min(P, lo + chunk) + e_t = torch.as_tensor(E[lo:hi], dtype=torch.float32, device=device) + e_t = e_t - e_t.mean(-1, keepdim=True) + e_t = e_t / e_t.norm(dim=-1, keepdim=True).clamp_min(1e-12) + + ang = torch.as_tensor(start[lo:hi], dtype=torch.float32, + device=device).clone().requires_grad_(True) + with torch.no_grad(): + out_before[lo:hi] = _ncc(simulator(ang), e_t).cpu().numpy() + + opt = torch.optim.Adam([ang], lr=lr) + # Pin backward to this thread; see the module docstring. + prev_mt = torch.is_grad_enabled() + try: + torch.autograd.set_multithreading_enabled(False) + except Exception as exc: # pragma: no cover + log.debug("set_multithreading_enabled unavailable: %s", exc) + + try: + for step in range(steps): + opt.zero_grad(set_to_none=True) + # Maximise similarity == minimise its negative. Summed, not + # averaged: each pattern's gradient is independent, so the sum + # keeps every one at full strength regardless of batch size. + loss = -_ncc(simulator(ang), e_t).sum() + loss.backward() + opt.step() + if on_yield is not None and step % yield_every == 0: + try: + on_yield() + except Exception as exc: # pragma: no cover + log.debug("refine on_yield failed: %s", exc) + finally: + try: + torch.autograd.set_multithreading_enabled(prev_mt) + except Exception: # pragma: no cover + pass + + with torch.no_grad(): + final = _ncc(simulator(ang), e_t) + better = final >= torch.as_tensor(out_before[lo:hi], device=device) + # NEVER return a worse orientation than we were given. Refinement + # is an improvement step on top of indexing; if Adam wandered off + # (a bad starting point, too large an lr) the indexed answer is + # still the better estimate. + keep = torch.where(better.unsqueeze(1), ang, + torch.as_tensor(start[lo:hi], + dtype=torch.float32, + device=device)) + out_euler[lo:hi] = keep.detach().cpu().numpy() + out_score[lo:hi] = torch.maximum( + final, torch.as_tensor(out_before[lo:hi], + device=device)).cpu().numpy() + + if progress is not None: + try: + progress(hi, P) + except Exception as exc: # pragma: no cover + log.debug("refine progress callback failed: %s", exc) + + result = RefinementResult(out_euler, out_score, out_before, steps, device) + result.nav_shape = tuple(nav_shape) + return result diff --git a/spyde/tests/migrated/test_ebsd_refine.py b/spyde/tests/migrated/test_ebsd_refine.py new file mode 100644 index 0000000..4befe92 --- /dev/null +++ b/spyde/tests/migrated/test_ebsd_refine.py @@ -0,0 +1,204 @@ +"""Batched orientation refinement (#72). + +Indexing can only return a dictionary entry, so its accuracy is capped by the +sampling step. The property under test is that refinement **lifts that cap**: +starting from a coarse indexed match it must land measurably closer to the true +orientation, which the synthetic data knows exactly. + +CPU only — torch-CUDA segfaults under pytest on Windows (CLAUDE.md). +""" +from __future__ import annotations + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") + +from spyde.data import ebsd_patterns, ground_truth +from spyde.data.synthetic import simulate_patterns +from spyde.ebsd import dictionary_index, remove_background, sample_orientations +from spyde.ebsd.refine import ( + BandSimulator, + euler_to_matrix_torch, + refine_orientations, +) + + +@pytest.fixture(scope="module") +def scan(): + s = ebsd_patterns(nav=(4, 4), detector=(32, 32), noise=0.0) + gt = ground_truth(s) + return s, np.asarray(gt["euler"]) + + +def _pattern_similarity(euler_a, euler_b, detector=(32, 32)): + """Compare orientations by the PATTERNS they produce, not by Euler + distance: Euler angles are not a metric space, and cubic symmetry means + different triples can describe the same crystal.""" + a = simulate_patterns(np.atleast_2d(euler_a), detector=detector) + b = simulate_patterns(np.atleast_2d(euler_b), detector=detector) + a = a.reshape(len(a), -1).astype(np.float64) + b = b.reshape(len(b), -1).astype(np.float64) + a = (a - a.mean(1, keepdims=True)) + b = (b - b.mean(1, keepdims=True)) + a /= np.linalg.norm(a, axis=1, keepdims=True) + b /= np.linalg.norm(b, axis=1, keepdims=True) + return (a * b).sum(1) + + +class TestEulerToMatrix: + def test_matches_the_numpy_reference(self): + """The torch rotation must be the same convention as the generator's, + or refinement optimises toward a different crystal.""" + from spyde.data.synthetic import _euler_to_matrix + eul = np.array([[0.3, 0.7, 1.1], [2.0, 0.4, 0.9]]) + want = _euler_to_matrix(eul[:, 0], eul[:, 1], eul[:, 2]) + got = euler_to_matrix_torch(torch.as_tensor(eul)).numpy() + np.testing.assert_allclose(got, want, atol=1e-12) + + def test_is_differentiable(self): + eul = torch.tensor([[0.3, 0.7, 1.1]], dtype=torch.float64, + requires_grad=True) + euler_to_matrix_torch(eul).sum().backward() + assert torch.isfinite(eul.grad).all() + + def test_produces_proper_rotations(self): + eul = torch.tensor([[0.3, 0.7, 1.1], [2.0, 0.4, 0.9]], + dtype=torch.float64) + m = euler_to_matrix_torch(eul) + eye = torch.eye(3, dtype=torch.float64).expand(2, 3, 3) + torch.testing.assert_close(m @ m.transpose(1, 2), eye) + torch.testing.assert_close(torch.linalg.det(m), + torch.ones(2, dtype=torch.float64)) + + +class TestBandSimulator: + def test_matches_the_numpy_generator(self): + """The differentiable simulator and the data generator must render the + SAME pattern — otherwise refinement is fitting a different forward + model than the one that made the data.""" + eul = np.array([[0.3, 0.7, 1.1], [1.5, 0.4, 0.2]]) + want = simulate_patterns(eul, detector=(24, 24)).reshape(2, -1) + got = BandSimulator((24, 24))(torch.as_tensor(eul, dtype=torch.float32)) + np.testing.assert_allclose(got.numpy(), want, rtol=1e-4, atol=1e-4) + + def test_is_differentiable_wrt_orientation(self): + eul = torch.tensor([[0.3, 0.7, 1.1]], requires_grad=True) + BandSimulator((16, 16))(eul).sum().backward() + assert torch.isfinite(eul.grad).all() + assert (eul.grad.abs() > 0).any(), "gradient is identically zero" + + +class TestRefinementImprovesOnIndexing: + def test_recovers_a_perturbed_orientation(self): + """The cleanest statement of the job: nudge a known orientation off by + a few degrees and require refinement to walk it back.""" + truth = np.array([[0.3, 0.6, 1.0], [1.2, 0.5, 0.4]]) + pat = simulate_patterns(truth, detector=(32, 32)) + start = truth + np.deg2rad(4.0) + + r = refine_orientations(pat, start, detector=(32, 32), device="cpu", + steps=200, lr=0.02) + before = _pattern_similarity(start, truth) + after = _pattern_similarity(r.euler, truth) + assert (after > before).all(), f"{after} vs {before}" + assert after.min() > 0.99, f"only reached {after}" + + def test_beats_a_coarse_dictionary(self, scan): + """The real workflow: index against a coarse dictionary, then refine. + Refinement must land closer to truth than the dictionary entry could.""" + s, euler = scan + exp = remove_background(s.data, device="cpu") + dic_euler = sample_orientations(step_deg=15.0) + dic = simulate_patterns(dic_euler, detector=(32, 32)) + idx = dictionary_index(exp, remove_background(dic, device="cpu"), + device="cpu") + indexed = idx.orientations(dic_euler) + + r = refine_orientations(exp, indexed, detector=(32, 32), device="cpu", + steps=250, lr=0.02) + flat_truth = euler.reshape(-1, 3) + before = _pattern_similarity(indexed, flat_truth) + after = _pattern_similarity(r.euler, flat_truth) + assert after.mean() > before.mean(), ( + f"refinement did not improve on indexing: " + f"{after.mean():.4f} vs {before.mean():.4f}") + + def test_reports_the_score_it_started_from(self, scan): + s, euler = scan + r = refine_orientations(s.data, euler.reshape(-1, 3) + 0.05, + detector=(32, 32), device="cpu", steps=60) + assert r.score_before.shape == r.score.shape + assert (r.score >= r.score_before - 1e-6).all() + + def test_never_returns_a_worse_orientation(self): + """Refinement is an improvement step ON TOP of indexing. If Adam + wanders off — a bad start, too large an lr — the indexed answer is + still the better estimate and must be what comes back.""" + truth = np.array([[0.3, 0.6, 1.0]]) + pat = simulate_patterns(truth, detector=(32, 32)) + # An absurd learning rate makes the optimiser diverge on purpose. + r = refine_orientations(pat, truth, detector=(32, 32), device="cpu", + steps=40, lr=5.0) + np.testing.assert_allclose(r.euler, truth, atol=1e-6) + assert r.improved.all() + + +class TestBatching: + def test_each_pattern_gets_its_own_orientation(self): + """A broadcasting mistake would drive every position to one answer.""" + truth = np.array([[0.2, 0.5, 0.9], [1.4, 0.8, 0.3], [2.6, 0.3, 1.2]]) + pat = simulate_patterns(truth, detector=(32, 32)) + r = refine_orientations(pat, truth + np.deg2rad(3.0), + detector=(32, 32), device="cpu", steps=150, + lr=0.02) + sim = _pattern_similarity(r.euler, truth) + assert (sim > 0.98).all(), sim + assert len({tuple(np.round(e, 4)) for e in r.euler}) == 3 + + def test_chunked_matches_unchunked(self): + truth = np.array([[0.2, 0.5, 0.9], [1.4, 0.8, 0.3], [2.6, 0.3, 1.2], + [0.9, 0.6, 0.7]]) + pat = simulate_patterns(truth, detector=(24, 24)) + start = truth + np.deg2rad(3.0) + kw = dict(detector=(24, 24), device="cpu", steps=80, lr=0.02) + whole = refine_orientations(pat, start, **kw) + pieces = refine_orientations(pat, start, chunk=2, **kw) + np.testing.assert_allclose(whole.euler, pieces.euler, atol=1e-5) + + def test_shape_mismatch_is_caught(self): + pat = simulate_patterns(np.array([[0.3, 0.6, 1.0]]), detector=(16, 16)) + with pytest.raises(ValueError, match="starting"): + refine_orientations(pat, np.zeros((5, 3)), detector=(16, 16), + device="cpu", steps=2) + + def test_result_maps_back_to_the_scan(self, scan): + s, euler = scan + r = refine_orientations(s.data, euler.reshape(-1, 3), + detector=(32, 32), device="cpu", steps=10) + assert r.euler_map((4, 4)).shape == (4, 4, 3) + + +class TestCallbacks: + def test_yield_is_called_inside_the_step_loop(self): + """Yielding only between stages leaves a UI frozen for seconds; the + contract is that it fires DURING the optimisation.""" + truth = np.array([[0.3, 0.6, 1.0]]) + pat = simulate_patterns(truth, detector=(16, 16)) + calls = [] + refine_orientations(pat, truth, detector=(16, 16), device="cpu", + steps=60, yield_every=10, + on_yield=lambda: calls.append(1)) + assert len(calls) >= 5 + + def test_a_failing_callback_does_not_kill_the_refine(self): + truth = np.array([[0.3, 0.6, 1.0]]) + pat = simulate_patterns(truth, detector=(16, 16)) + + def boom(): + raise RuntimeError("renderer went away") + + r = refine_orientations(pat, truth, detector=(16, 16), device="cpu", + steps=30, on_yield=boom, + progress=lambda d, t: boom()) + assert r.euler.shape == (1, 3) From b1568675aac6e19cd2b19680dd0bd9b2edbe59dd Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Sun, 26 Jul 2026 23:41:27 -0500 Subject: [PATCH 08/60] feat(ebsd): CrystalMap, IPF colours, similarity map, phase merging (#73) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- spyde/ebsd/__init__.py | 10 +- spyde/ebsd/crystal_map.py | 194 +++++++++++++++++ spyde/tests/migrated/test_ebsd_crystal_map.py | 197 ++++++++++++++++++ 3 files changed, 400 insertions(+), 1 deletion(-) create mode 100644 spyde/ebsd/crystal_map.py create mode 100644 spyde/tests/migrated/test_ebsd_crystal_map.py diff --git a/spyde/ebsd/__init__.py b/spyde/ebsd/__init__.py index cab63c4..57a9761 100644 --- a/spyde/ebsd/__init__.py +++ b/spyde/ebsd/__init__.py @@ -23,8 +23,16 @@ sample_orientations, ) from spyde.ebsd.preprocess import average_dot_product_map, remove_background +from spyde.ebsd.crystal_map import ( + ipf_colors, + merge_phases, + orientation_similarity_map, + to_crystal_map, +) from spyde.ebsd.refine import RefinementResult, refine_orientations __all__ = ["dictionary_index", "sample_orientations", "IndexingResult", "remove_background", "average_dot_product_map", - "refine_orientations", "RefinementResult"] + "refine_orientations", "RefinementResult", + "to_crystal_map", "ipf_colors", "orientation_similarity_map", + "merge_phases"] diff --git a/spyde/ebsd/crystal_map.py b/spyde/ebsd/crystal_map.py new file mode 100644 index 0000000..0a0bdbe --- /dev/null +++ b/spyde/ebsd/crystal_map.py @@ -0,0 +1,194 @@ +"""crystal_map.py — indexing results -> orix CrystalMap -> IPF colours (#73). + +The display half of Wave 3 is nearly free, and deliberately so. 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 here is to hand the EBSD result over +in the form that machinery already speaks — not to build a second orientation +display beside the first. + +Two things this module does that the existing 4D-STEM path does not need: + +``orientation_similarity_map`` + A quality metric peculiar to dictionary indexing: how much two neighbouring + positions agree about their *ranked list* of best matches. Unlike a raw + correlation score it is sensitive to indexing that is confidently wrong — + a position that matched well but disagrees with everything around it stands + out, which a score map cannot show. + +``merge_phases`` + Multi-phase indexing runs one dictionary per phase, so the maps have to be + combined by score afterwards. +""" +from __future__ import annotations + +import logging + +import numpy as np + +log = logging.getLogger(__name__) + +# Space groups for the phases most likely to be indexed first. A caller can +# always pass an explicit orix Phase; this is only a convenience. +COMMON_PHASES = { + "fcc": 225, # Fm-3m — Al, Ni, Cu, austenite + "bcc": 229, # Im-3m — ferrite, W, Mo + "dc": 227, # Fd-3m — Si, Ge, diamond + "hcp": 194, # P6_3/mmc — Ti, Mg, Zn +} + + +def _phase(phase=None, name: str = "phase", space_group: int = 225): + from orix.crystal_map import Phase + if phase is not None: + return phase + return Phase(name=name, space_group=int(space_group)) + + +def to_crystal_map(euler, nav_shape=None, *, phase=None, space_group: int = 225, + phase_name: str = "phase", scores=None, step: float = 1.0): + """Euler angles -> an orix :class:`~orix.crystal_map.CrystalMap`. + + Parameters + ---------- + euler : array (..., 3) + Bunge angles in radians — e.g. ``RefinementResult.euler`` or + ``IndexingResult.orientations(dictionary_euler)``. + scores : array, optional + Per-position match quality, stored as the map's ``prop["scores"]`` so + the display and :func:`merge_phases` can use it. + step : float + Scan step size, used for the x/y coordinates. + """ + from orix.crystal_map import CrystalMap, PhaseList + from orix.quaternion import Rotation + + eul = np.asarray(euler, float) + if nav_shape is None: + nav_shape = eul.shape[:-1] + flat = eul.reshape(-1, 3) + + rot = Rotation.from_euler(flat) + ph = _phase(phase, phase_name, space_group) + + if len(nav_shape) == 2: + ny, nx = nav_shape + yy, xx = np.mgrid[0:ny, 0:nx].astype(float) * float(step) + x, y = xx.ravel(), yy.ravel() + else: + x = np.arange(len(flat), dtype=float) * float(step) + y = None + + props = {"scores": np.asarray(scores, float).ravel()} if scores is not None else None + return CrystalMap(rotations=rot, phase_id=np.zeros(len(flat), int), + x=x, y=y, phase_list=PhaseList([ph]), prop=props) + + +def ipf_colors(euler, nav_shape=None, *, phase=None, space_group: int = 225, + direction=None) -> np.ndarray: + """IPF-Z RGB for each orientation -> ``(..., 3)`` float in [0, 1]. + + This is what feeds SpyDE's existing orientation display: an ``(H, W, 3)`` + RGB array is exactly what ``commit.commit_result_tree`` already accepts as + an RGB primary (it skips contrast locking for those). + """ + from orix.plot import IPFColorKeyTSL + from orix.quaternion import Rotation + from orix.vector import Vector3d + + eul = np.asarray(euler, float) + if nav_shape is None: + nav_shape = eul.shape[:-1] + rot = Rotation.from_euler(eul.reshape(-1, 3)) + ph = _phase(phase, "phase", space_group) + key = IPFColorKeyTSL(ph.point_group, + direction=direction or Vector3d.zvector()) + rgb = np.asarray(key.orientation2color(rot), float) + return np.clip(rgb, 0.0, 1.0).reshape(tuple(nav_shape) + (3,)) + + +def orientation_similarity_map(indices, nav_shape=None) -> np.ndarray: + """How much each position's ranked match list agrees with its neighbours. + + ``indices`` is ``IndexingResult.indices`` — ``(P, k)``, best first. For each + position the metric is the mean overlap between its top-k dictionary + indices and each neighbour's, normalised to [0, 1]. + + Why this rather than the correlation score: a score map shows how well a + pattern matched *something*, and a confidently WRONG index scores just as + highly as a right one. Agreement with the neighbourhood is what exposes it, + so this is the map that finds bad indexing rather than bad patterns. + + Needs ``keep > 1`` from :func:`~spyde.ebsd.indexing.dictionary_index`; + with k=1 it degenerates to "does my single best match equal my neighbour's", + which is a legitimate but much coarser signal. + """ + idx = np.asarray(indices) + if idx.ndim != 2: + raise ValueError("indices must be (P, k) from IndexingResult.indices") + P, k = idx.shape + if nav_shape is None: + raise ValueError("nav_shape is required to find neighbours") + ny, nx = nav_shape + if ny * nx != P: + raise ValueError(f"nav_shape {nav_shape} does not match {P} positions") + + grid = idx.reshape(ny, nx, k) + total = np.zeros((ny, nx), float) + count = np.zeros((ny, nx), float) + + # Set overlap per neighbour pair, vectorised over the whole map: compare + # every ranked index against every one of the neighbour's, then count hits. + for dy, dx in ((0, 1), (0, -1), (1, 0), (-1, 0)): + ys = slice(max(0, dy), ny + min(0, dy)) + xs = slice(max(0, dx), nx + min(0, dx)) + ys2 = slice(max(0, -dy), ny + min(0, -dy)) + xs2 = slice(max(0, -dx), nx + min(0, -dx)) + a = grid[ys, xs][..., :, None] # (h, w, k, 1) + b = grid[ys2, xs2][..., None, :] # (h, w, 1, k) + overlap = (a == b).any(-1).sum(-1) / float(k) + total[ys, xs] += overlap + count[ys, xs] += 1.0 + return total / np.maximum(count, 1.0) + + +def merge_phases(maps, scores, nav_shape=None): + """Combine per-phase indexing runs into one labelled result. + + Multi-phase indexing runs one dictionary per phase; the winner at each + position is simply the phase whose match scored highest. + + Parameters + ---------- + maps : sequence of array (..., 3) + Euler angles from each phase's run, in the same order as *scores*. + scores : sequence of array (...,) + Best-match score per position for each phase. + + Returns + ------- + (euler, phase_id, best_score) + """ + eulers = [np.asarray(m, float) for m in maps] + raw_scores = [np.asarray(s, float).ravel() for s in scores] + # Validate BEFORE stacking: np.stack raises "all input arrays must have the + # same shape", which says nothing about which phase disagreed or why. + sizes = {e.reshape(-1, 3).shape[0] for e in eulers} | {s.size for s in raw_scores} + if len(sizes) != 1: + raise ValueError( + f"every phase map must cover the same positions, got {sorted(sizes)}") + if len(eulers) != len(raw_scores): + raise ValueError(f"{len(eulers)} phase maps but {len(raw_scores)} " + f"score arrays") + sc = np.stack(raw_scores) # (n_phase, P) + + winner = sc.argmax(0) # (P,) + stacked = np.stack([e.reshape(-1, 3) for e in eulers]) # (n, P, 3) + P = stacked.shape[1] + euler = stacked[winner, np.arange(P)] + best = sc[winner, np.arange(P)] + + if nav_shape is not None: + euler = euler.reshape(tuple(nav_shape) + (3,)) + winner = winner.reshape(nav_shape) + best = best.reshape(nav_shape) + return euler, winner, best diff --git a/spyde/tests/migrated/test_ebsd_crystal_map.py b/spyde/tests/migrated/test_ebsd_crystal_map.py new file mode 100644 index 0000000..552bcf0 --- /dev/null +++ b/spyde/tests/migrated/test_ebsd_crystal_map.py @@ -0,0 +1,197 @@ +"""CrystalMap, IPF colours, similarity map, phase merging (#73). + +The display half of Wave 3. Compute only — the window wiring that puts these on +screen is UI work and gets verified by running the app and looking at pixels +(CLAUDE.md), not here. + +The similarity map gets the most attention because it is the one thing here +with non-obvious semantics: it measures agreement between NEIGHBOURS' ranked +match lists, which is what exposes confidently-wrong indexing that a raw score +map cannot. +""" +from __future__ import annotations + +import numpy as np +import pytest + +pytest.importorskip("orix") + +from spyde.data import ebsd_patterns, ground_truth +from spyde.data.synthetic import simulate_patterns +from spyde.ebsd import dictionary_index, remove_background, sample_orientations +from spyde.ebsd.crystal_map import ( + COMMON_PHASES, + ipf_colors, + merge_phases, + orientation_similarity_map, + to_crystal_map, +) + + +@pytest.fixture(scope="module") +def scan(): + s = ebsd_patterns(nav=(8, 8), detector=(32, 32), noise=0.0) + gt = ground_truth(s) + return s, np.asarray(gt["euler"]), np.asarray(gt["grain2_mask"], bool) + + +class TestCrystalMap: + def test_builds_from_euler_angles(self, scan): + _, euler, _ = scan + xm = to_crystal_map(euler) + assert xm.size == 64 + assert xm.shape == (8, 8) + + def test_carries_scores_as_a_property(self, scan): + _, euler, _ = scan + xm = to_crystal_map(euler, scores=np.linspace(0, 1, 64)) + assert "scores" in xm.prop + assert xm.prop["scores"].shape == (64,) + + def test_phase_symmetry_comes_from_the_space_group(self, scan): + _, euler, _ = scan + assert to_crystal_map(euler, space_group=COMMON_PHASES["fcc"] + ).phases[0].point_group.name == "m-3m" + assert to_crystal_map(euler, space_group=COMMON_PHASES["hcp"] + ).phases[0].point_group.name == "6/mmm" + + def test_step_scales_the_coordinates(self, scan): + _, euler, _ = scan + assert to_crystal_map(euler, step=0.5).x.max() == pytest.approx(3.5) + + +class TestIpfColors: + def test_shape_and_range(self, scan): + _, euler, _ = scan + rgb = ipf_colors(euler) + assert rgb.shape == (8, 8, 3) + assert rgb.min() >= 0.0 and rgb.max() <= 1.0 + + def test_the_two_grains_get_different_colours(self, scan): + """The whole point of an IPF map. Identical colours would mean the map + renders as a flat block and still passes every range check.""" + _, euler, mask = scan + rgb = ipf_colors(euler) + assert not np.allclose(rgb[mask].mean(0), rgb[~mask].mean(0), atol=0.05) + + def test_one_orientation_gives_one_colour(self, scan): + """Grain 2 is a single orientation, so its colour must be uniform.""" + _, euler, mask = scan + block = ipf_colors(euler)[mask] + # Compared as a max deviation rather than assert_allclose against + # block[0]: numpy no longer broadcasts shapes there, so (N, 3) vs (3,) + # fails on shape even when every row is identical. + assert np.abs(block - block[0]).max() < 1e-6 + + def test_is_rgb_shaped_for_the_existing_display(self, scan): + """commit_result_tree takes an (H, W, 3) RGB primary directly — this is + why Wave 3 needs no new display code.""" + _, euler, _ = scan + rgb = ipf_colors(euler) + assert rgb.ndim == 3 and rgb.shape[-1] == 3 + + +class TestOrientationSimilarityMap: + def test_uniform_indexing_scores_one(self): + """Every position agreeing with every neighbour is perfect agreement.""" + idx = np.tile(np.array([[3, 7, 1]]), (16, 1)) + osm = orientation_similarity_map(idx, (4, 4)) + np.testing.assert_allclose(osm, 1.0) + + def test_a_boundary_shows_as_low_agreement(self): + """The metric exists to find boundaries and bad indexing; a map where + half the positions match a different list must dip at the seam.""" + grid = np.zeros((4, 4, 2), int) + grid[:, :2] = [1, 2] + grid[:, 2:] = [8, 9] + osm = orientation_similarity_map(grid.reshape(16, 2), (4, 4)) + assert osm[:, 1].mean() < osm[:, 0].mean() + assert osm[:, 2].mean() < osm[:, 3].mean() + + def test_partial_overlap_is_between_zero_and_one(self): + """Sharing some of the ranked list, but not all of it, is the normal + case and must land in between rather than saturating.""" + grid = np.zeros((2, 2, 2), int) + grid[0, 0] = [1, 2] + grid[0, 1] = [1, 9] # shares one of two + grid[1, 0] = [1, 2] + grid[1, 1] = [1, 2] + osm = orientation_similarity_map(grid.reshape(4, 2), (2, 2)) + assert 0.0 < osm[0, 1] < 1.0 + + def test_detects_a_confidently_wrong_position(self): + """The case a SCORE map cannot show: one position indexed to something + completely different from its neighbourhood.""" + idx = np.tile(np.array([[3, 7, 1]]), (25, 1)) + idx[12] = [40, 41, 42] # centre of a 5x5 + osm = orientation_similarity_map(idx, (5, 5)) + assert osm[2, 2] == pytest.approx(0.0) + assert osm[0, 0] > 0.9 + + def test_rejects_a_mismatched_shape(self): + with pytest.raises(ValueError, match="does not match"): + orientation_similarity_map(np.zeros((10, 3), int), (4, 4)) + + def test_requires_a_nav_shape(self): + with pytest.raises(ValueError, match="nav_shape"): + orientation_similarity_map(np.zeros((16, 3), int)) + + def test_works_on_a_real_indexing_result(self, scan): + s, _, mask = scan + dic_euler = sample_orientations(step_deg=12.0) + dic = simulate_patterns(dic_euler, detector=(32, 32)) + r = dictionary_index(remove_background(s.data, device="cpu"), + remove_background(dic, device="cpu"), + keep=5, device="cpu") + osm = orientation_similarity_map(r.indices, (8, 8)) + assert osm.shape == (8, 8) + assert (osm >= 0).all() and (osm <= 1).all() + # Inside a uniform grain neighbours agree; across the boundary they + # cannot, so the map must not be flat. + assert osm.std() > 0 + + +class TestMergePhases: + def test_picks_the_higher_scoring_phase_per_position(self): + a = np.tile([0.1, 0.2, 0.3], (4, 1)) + b = np.tile([1.0, 1.1, 1.2], (4, 1)) + euler, phase_id, best = merge_phases( + [a, b], [np.array([0.9, 0.1, 0.9, 0.1]), + np.array([0.2, 0.8, 0.2, 0.8])]) + assert phase_id.tolist() == [0, 1, 0, 1] + np.testing.assert_allclose(euler[0], a[0]) + np.testing.assert_allclose(euler[1], b[1]) + np.testing.assert_allclose(best, [0.9, 0.8, 0.9, 0.8]) + + def test_reshapes_to_the_scan(self): + a = np.zeros((4, 3)) + b = np.ones((4, 3)) + euler, phase_id, best = merge_phases( + [a, b], [np.zeros(4), np.ones(4)], nav_shape=(2, 2)) + assert euler.shape == (2, 2, 3) + assert phase_id.shape == (2, 2) + assert best.shape == (2, 2) + + def test_mismatched_coverage_is_caught(self): + with pytest.raises(ValueError, match="same positions"): + merge_phases([np.zeros((4, 3)), np.zeros((5, 3))], + [np.zeros(4), np.zeros(5)]) + + +class TestEndToEnd: + def test_index_then_colour_recovers_the_grain_structure(self, scan): + """Indexing -> IPF colours must reproduce the two-grain layout the data + was built with — the closest headless proxy for 'the map looks right'. + The window itself still needs eyes on a screenshot.""" + s, euler, mask = scan + dic_euler = sample_orientations(step_deg=10.0) + dic = simulate_patterns(dic_euler, detector=(32, 32)) + r = dictionary_index(remove_background(s.data, device="cpu"), + remove_background(dic, device="cpu"), + device="cpu") + rgb = ipf_colors(r.orientations(dic_euler, (8, 8))) + assert rgb.shape == (8, 8, 3) + inside = rgb[mask].mean(0) + outside = rgb[~mask].mean(0) + assert not np.allclose(inside, outside, atol=0.05), \ + "indexed IPF map does not distinguish the two grains" From a3e5ae5baaa5b3920778f2df510a90f7e8a8ae42 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Sun, 26 Jul 2026 23:45:20 -0500 Subject: [PATCH 09/60] feat(fitting): 2-D model support (#59) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). --- spyde/fitting/components.py | 70 ++++++++- spyde/tests/migrated/test_fitting_2d.py | 193 ++++++++++++++++++++++++ 2 files changed, 257 insertions(+), 6 deletions(-) create mode 100644 spyde/tests/migrated/test_fitting_2d.py diff --git a/spyde/fitting/components.py b/spyde/fitting/components.py index a7d392e..578447b 100644 --- a/spyde/fitting/components.py +++ b/spyde/fitting/components.py @@ -62,10 +62,10 @@ class TorchComponent: autodiff, which is the ideal oracle for exactly this. """ - __slots__ = ("kind", "params", "linear", "_fn", "_grad") + __slots__ = ("kind", "params", "linear", "_fn", "_grad", "ndim") def __init__(self, kind: str, params: Sequence[str], linear: Sequence[bool], - fn: Callable, grad: Callable | None = None): + fn: Callable, grad: Callable | None = None, ndim: int = 1): if len(params) != len(linear): raise ValueError(f"{kind}: {len(params)} params vs " f"{len(linear)} linear flags") @@ -74,6 +74,11 @@ def __init__(self, kind: str, params: Sequence[str], linear: Sequence[bool], self.linear = tuple(bool(b) for b in linear) self._fn = fn self._grad = grad + # 1 for a spectrum (x is the signal axis), 2 for an image (x carries + # (x, y) coordinate PAIRS). A 2-D model is otherwise identical: the + # image is flattened to C sample points and everything downstream — + # packing, the Jacobian, the LM solve — is unchanged. + self.ndim = int(ndim) @property def n_params(self) -> int: @@ -87,9 +92,15 @@ def _split(self, x, p): if p.shape[-1] != self.n_params: raise ValueError(f"{self.kind} expects {self.n_params} parameters " f"{self.params}, got {p.shape[-1]}") + cols = [p[:, i:i + 1] for i in range(self.n_params)] + if self.ndim == 2: + # x is (C, 2) or (P, C, 2): coordinate PAIRS, not a single axis. + if x.dim() == 2: + x = x.unsqueeze(0) # (1, C, 2) + return (x[..., 0], x[..., 1]), cols if x.dim() == 1: x = x.unsqueeze(0) # (1, C), broadcasts over P - return x, [p[:, i:i + 1] for i in range(self.n_params)] + return x, cols def __call__(self, x, p): """``x``: ``(C,)`` or ``(P, C)``. ``p``: ``(P, n_params)``. -> ``(P, C)``.""" @@ -190,6 +201,48 @@ def _logistic(x, p): return a / (1.0 + b * (-c * (x - origin)).exp()) +def _gaussian_2d(xy, p): + # "A * (1 / (sigma_x * sigma_y * 2 * pi)) + # * exp(-((x - centre_x)**2 / (2*sigma_x**2) + # + (y - centre_y)**2 / (2*sigma_y**2)))" + # As in 1-D, A is the VOLUME under the surface, not the peak height. + x, y = xy + A, cx, cy, sx, sy = p + sxe, sye = sx + _EPS, sy + _EPS + return A / (sxe * sye * 2.0 * math.pi) * ( + -((x - cx) ** 2 / (2 * sxe ** 2) + (y - cy) ** 2 / (2 * sye ** 2))).exp() + + +def _d_gaussian_2d(xy, p): + x, y = xy + A, cx, cy, sx, sy = p + sxe, sye = sx + _EPS, sy + _EPS + dx, dy = x - cx, y - cy + shape = (-(dx ** 2 / (2 * sxe ** 2) + dy ** 2 / (2 * sye ** 2))).exp() \ + / (sxe * sye * 2.0 * math.pi) + f = A * shape + return _stack(shape, # df/dA + f * dx / sxe ** 2, # df/dcentre_x + f * dy / sye ** 2, # df/dcentre_y + f * (dx ** 2 / sxe ** 3 - 1.0 / sxe), # df/dsigma_x + f * (dy ** 2 / sye ** 3 - 1.0 / sye)) # df/dsigma_y + + +def image_coordinates(shape, *, device=None, dtype=None): + """``(H*W, 2)`` of ``(x, y)`` sample points for a 2-D model. + + This is the "signal axis" a 2-D fit uses: the image is flattened to H*W + sample points, exactly as a spectrum is C channels, so nothing downstream + (packing, the Jacobian, the LM solve) needs to know it came from an image. + """ + import torch + h, w = int(shape[0]), int(shape[1]) + yy, xx = torch.meshgrid(torch.arange(h, device=device, dtype=dtype), + torch.arange(w, device=device, dtype=dtype), + indexing="ij") + return torch.stack([xx.reshape(-1), yy.reshape(-1)], dim=-1) + + # --------------------------------------------------------------------------- # analytic derivatives — checked against autodiff in test_torch_components.py # @@ -355,6 +408,9 @@ def _register(kind, params, linear, fn, grad=None) -> None: _register("Erf", ("A", "origin", "sigma"), (True, False, False), _erf, _d_erf) _register("HeavisideStep", ("A", "n"), (True, False), _heaviside, _d_heaviside) _register("Logistic", ("a", "b", "c", "origin"), (True, False, False, False), _logistic, _d_logistic) +_REGISTRY["Gaussian2D"] = TorchComponent( + "Gaussian2D", ("A", "centre_x", "centre_y", "sigma_x", "sigma_y"), + (True, False, False, False, False), _gaussian_2d, _d_gaussian_2d, ndim=2) def get_component(kind: str, *, n_params: int | None = None) -> TorchComponent: @@ -426,9 +482,10 @@ def evaluate_with_grad(spec, x, values): """ import torch - xb = x.unsqueeze(0) if x.dim() == 1 else x + # Sample count is the LAST axis for a 1-D axis and the second-to-last for + # 2-D coordinate pairs — `x` is (C,)/(P, C) or (C, 2)/(P, C, 2). + C = x.shape[-2] if x.shape[-1] == 2 and x.dim() >= 2 else x.shape[-1] P = values.shape[0] - C = xb.shape[-1] n_total = values.shape[1] out = torch.zeros((P, C), dtype=values.dtype, device=values.device) jac = torch.zeros((P, C, n_total), dtype=values.dtype, device=values.device) @@ -462,5 +519,6 @@ def evaluate(spec, x, values): i += n if out is None: p = values.shape[0] if values.dim() > 1 else 1 - return torch.zeros((p, x.shape[-1]), dtype=x.dtype, device=x.device) + c = x.shape[-2] if x.shape[-1] == 2 and x.dim() >= 2 else x.shape[-1] + return torch.zeros((p, c), dtype=x.dtype, device=x.device) return out diff --git a/spyde/tests/migrated/test_fitting_2d.py b/spyde/tests/migrated/test_fitting_2d.py new file mode 100644 index 0000000..8c80db3 --- /dev/null +++ b/spyde/tests/migrated/test_fitting_2d.py @@ -0,0 +1,193 @@ +"""2-D model support (#59) — Gaussian2D through the same batched engine. + +A 2-D model is not a separate code path: the image is flattened to H*W sample +points exactly as a spectrum is C channels, and the only difference is that the +"axis" carries (x, y) PAIRS. Everything downstream — packing, the Jacobian, the +LM solve, bounds — is the 1-D machinery unchanged, and these tests exist to +prove that rather than assume it. + +This is also what Wave 4 needs: refining atom positions IS a batched 2-D +gaussian fit (#77). +""" +from __future__ import annotations + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") + +import hyperspy.components2d as c2d + +from spyde.fitting import ModelSpec +from spyde.fitting import components as tc +from spyde.fitting.engine import fit_batched +from spyde.fitting.spec import ComponentSpec, ParameterSpec + +SHAPE = (24, 28) # deliberately non-square: catches a transposed grid + + +def _truth(A=500.0, cx=13.0, cy=10.0, sx=2.5, sy=3.5): + return dict(A=A, centre_x=cx, centre_y=cy, sigma_x=sx, sigma_y=sy) + + +def _hyperspy_image(**vals): + g = c2d.Gaussian2D() + for k, v in vals.items(): + getattr(g, k).value = v + yy, xx = np.mgrid[0:SHAPE[0], 0:SHAPE[1]].astype(float) + return np.asarray(g.function(xx, yy), float) + + +def _spec(**vals): + order = ("A", "centre_x", "centre_y", "sigma_x", "sigma_y") + return ModelSpec(components=[ComponentSpec( + kind="Gaussian2D", + parameters=[ParameterSpec(n, float(vals[n])) for n in order])]) + + +class TestGaussian2DParity: + def test_parameter_order_matches_hyperspy(self): + hs = tuple(p.name for p in c2d.Gaussian2D().parameters) + assert tc.get_component("Gaussian2D").params == hs + + def test_linear_flag_matches_hyperspy(self): + hs = tuple(bool(getattr(p, "_linear", False)) + for p in c2d.Gaussian2D().parameters) + assert tc.get_component("Gaussian2D").linear == hs + + def test_values_match_hyperspy(self): + """A is the VOLUME, not the peak height — the same convention trap as + the 1-D Gaussian.""" + vals = _truth() + want = _hyperspy_image(**vals).ravel() + xy = tc.image_coordinates(SHAPE, dtype=torch.float64) + comp = tc.get_component("Gaussian2D") + got = comp(xy, torch.as_tensor( + [[vals[n] for n in comp.params]], dtype=torch.float64)) + np.testing.assert_allclose(got.numpy()[0], want, rtol=1e-10, atol=1e-12) + + def test_analytic_gradient_matches_autodiff(self): + from torch.func import jacfwd + vals = _truth() + comp = tc.get_component("Gaussian2D") + xy = tc.image_coordinates(SHAPE, dtype=torch.float64) + v = torch.as_tensor([[vals[n] for n in comp.params]], + dtype=torch.float64) + auto = jacfwd(lambda p: comp(xy, p.unsqueeze(0)).squeeze(0))(v[0]) + torch.testing.assert_close(comp.grad(xy, v)[0], auto, + rtol=1e-9, atol=1e-9) + + +class TestCoordinates: + def test_shape_and_ordering(self): + xy = tc.image_coordinates(SHAPE) + assert xy.shape == (SHAPE[0] * SHAPE[1], 2) + # Row-major: x runs fastest, matching an image flattened with .ravel(). + assert xy[0].tolist() == [0, 0] + assert xy[1].tolist() == [1, 0] + assert xy[SHAPE[1]].tolist() == [0, 1] + + def test_matches_a_ravelled_mgrid(self): + yy, xx = np.mgrid[0:SHAPE[0], 0:SHAPE[1]] + xy = tc.image_coordinates(SHAPE).numpy() + np.testing.assert_array_equal(xy[:, 0], xx.ravel()) + np.testing.assert_array_equal(xy[:, 1], yy.ravel()) + + +class TestBatching: + def test_each_row_gets_its_own_surface(self): + comp = tc.get_component("Gaussian2D") + xy = tc.image_coordinates(SHAPE, dtype=torch.float64) + vals = torch.tensor([[500.0, 5.0, 5.0, 2.0, 2.0], + [500.0, 20.0, 18.0, 2.0, 2.0]], + dtype=torch.float64) + y = comp(xy, vals).reshape(2, *SHAPE) + peak0 = np.unravel_index(int(y[0].argmax()), SHAPE) + peak1 = np.unravel_index(int(y[1].argmax()), SHAPE) + assert peak0 == (5, 5) # (row=y, col=x) + assert peak1 == (18, 20) + + def test_non_square_images_are_not_transposed(self): + """SHAPE is deliberately non-square — a transposed coordinate grid + would still run and put every peak in the wrong place.""" + comp = tc.get_component("Gaussian2D") + xy = tc.image_coordinates(SHAPE, dtype=torch.float64) + y = comp(xy, torch.tensor([[500.0, 22.0, 3.0, 1.5, 1.5]], + dtype=torch.float64)).reshape(SHAPE) + assert np.unravel_index(int(y.argmax()), SHAPE) == (3, 22) + + +class TestFittingAnImage: + def test_recovers_the_generating_parameters(self): + """End to end: a 2-D model fits through the SAME engine as a spectrum.""" + vals = _truth() + img = _hyperspy_image(**vals).ravel()[None, :] + spec = _spec(A=400.0, centre_x=12.0, centre_y=11.0, + sigma_x=2.0, sigma_y=3.0) + xy = tc.image_coordinates(SHAPE, dtype=torch.float64).numpy() + + got = fit_batched(spec, img, xy, device="cpu", max_iter=200) + names = spec.parameter_names() + for key, want in vals.items(): + col = names.index(f"Gaussian2D.{key}") + assert got.values[0, col] == pytest.approx(want, rel=1e-4), key + assert got.converged.all() + + def test_fits_many_images_at_once(self): + """The point of the engine: N independent 2-D fits in one call.""" + centres = [(8.0, 7.0), (14.0, 12.0), (20.0, 16.0)] + imgs = np.stack([_hyperspy_image(**_truth(cx=cx, cy=cy)).ravel() + for cx, cy in centres]) + spec = _spec(A=400.0, centre_x=13.0, centre_y=11.0, + sigma_x=2.0, sigma_y=3.0) + xy = tc.image_coordinates(SHAPE, dtype=torch.float64).numpy() + + got = fit_batched(spec, imgs, xy, device="cpu", max_iter=200) + names = spec.parameter_names() + cx = got.values[:, names.index("Gaussian2D.centre_x")] + cy = got.values[:, names.index("Gaussian2D.centre_y")] + np.testing.assert_allclose(cx, [c[0] for c in centres], atol=1e-3) + np.testing.assert_allclose(cy, [c[1] for c in centres], atol=1e-3) + + def test_bounds_apply_in_2d_too(self): + """Bounds are engine-level, not per-component — but a 2-D model going + through a different shaping path is exactly where that could break.""" + img = _hyperspy_image(**_truth()).ravel()[None, :] + spec = _spec(A=400.0, centre_x=12.0, centre_y=11.0, + sigma_x=2.0, sigma_y=3.0) + spec[0]["centre_x"].bmin = 5.0 + spec[0]["centre_x"].bmax = 9.0 + xy = tc.image_coordinates(SHAPE, dtype=torch.float64).numpy() + got = fit_batched(spec, img, xy, device="cpu", max_iter=120) + col = spec.parameter_names().index("Gaussian2D.centre_x") + assert 5.0 - 1e-9 <= got.values[0, col] <= 9.0 + 1e-9 + + def test_a_fixed_parameter_stays_fixed_in_2d(self): + img = _hyperspy_image(**_truth()).ravel()[None, :] + spec = _spec(A=400.0, centre_x=12.0, centre_y=11.0, + sigma_x=2.0, sigma_y=3.5) + spec[0]["sigma_y"].free = False + xy = tc.image_coordinates(SHAPE, dtype=torch.float64).numpy() + got = fit_batched(spec, img, xy, device="cpu", max_iter=120) + col = spec.parameter_names().index("Gaussian2D.sigma_y") + assert got.values[0, col] == pytest.approx(3.5, rel=1e-12) + + +class TestRoundTrip: + def test_gaussian2d_round_trips_through_modelspec(self): + """A 2-D model must save and reopen like any other.""" + import hyperspy.api as hs + + sig = hs.signals.Signal2D(np.zeros(SHAPE)) + m = sig.create_model() + m.append(c2d.Gaussian2D()) + for k, v in _truth().items(): + getattr(m[0], k).value = v + spec = ModelSpec.from_model(m) + assert spec[0].kind == "Gaussian2D" + back = ModelSpec.from_model(spec.to_model(sig)) + np.testing.assert_allclose(back.flat_values(), spec.flat_values()) + + def test_the_engine_reports_2d_support(self): + assert tc.supports(_spec(**_truth())) is True + assert tc.has_analytic_grad(_spec(**_truth())) is True From 56d04b16ce52528acaf4ad4ba4553afb63c20da2 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Mon, 27 Jul 2026 00:14:41 -0500 Subject: [PATCH 10/60] feat(atoms): atom finding, GPU refinement, property maps (#75, #77, #79) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- spyde/atoms/__init__.py | 46 ++++++ spyde/atoms/finding.py | 204 +++++++++++++++++++++++ spyde/atoms/properties.py | 128 +++++++++++++++ spyde/data/__init__.py | 4 +- spyde/data/synthetic.py | 75 +++++++++ spyde/tests/migrated/test_atoms.py | 249 +++++++++++++++++++++++++++++ 6 files changed, 705 insertions(+), 1 deletion(-) create mode 100644 spyde/atoms/__init__.py create mode 100644 spyde/atoms/finding.py create mode 100644 spyde/atoms/properties.py create mode 100644 spyde/tests/migrated/test_atoms.py diff --git a/spyde/atoms/__init__.py b/spyde/atoms/__init__.py new file mode 100644 index 0000000..bf0efd9 --- /dev/null +++ b/spyde/atoms/__init__.py @@ -0,0 +1,46 @@ +"""spyde.atoms — atom position mapping (0.3.0 Wave 4, GitHub #74). + +**atomap owns the structure, SpyDE owns the refinement.** Initial peak finding, +sublattices, nearest neighbours, zone axes and dumbbell pairing are atomap's, +because reimplementing them would mean diverging from published atomap results. +Refinement is ours because refining atom positions *is* a batched 2-D gaussian +fit, and :mod:`spyde.fitting` already does the whole field in one Levenberg- +Marquardt where atomap fits one atom at a time with scipy. + +atomap's GUI functions (``select_atoms_with_gui``, ``add_atoms_with_gui``, +``toggle_atom_refine_position_with_gui``) are matplotlib-based, so the +*interaction* is reimplemented over SpyDE's own anyplotlib overlay (#76) — +the same machinery ``actions/vector_overlay.py`` already uses for diffraction +vectors. + +Only :func:`~spyde.atoms.finding.find_atoms` needs the ``atoms`` extra; +refinement and the property maps are plain numpy/torch, so they are testable +and usable without it. +""" +from __future__ import annotations + +from spyde.atoms.finding import ( + MissingExtra, + find_atoms, + refine_atoms, + refine_center_of_mass, + refine_gaussian, +) +from spyde.atoms.properties import ( + displacement_from_ideal, + displacement_magnitude, + ellipticity, + ellipticity_angle, + intensity, + nearest_neighbour_distance, + property_maps, + to_map, +) + +__all__ = [ + "find_atoms", "refine_atoms", "refine_center_of_mass", "refine_gaussian", + "MissingExtra", + "ellipticity", "ellipticity_angle", "intensity", + "nearest_neighbour_distance", "displacement_from_ideal", + "displacement_magnitude", "property_maps", "to_map", +] diff --git a/spyde/atoms/finding.py b/spyde/atoms/finding.py new file mode 100644 index 0000000..9b1f8b6 --- /dev/null +++ b/spyde/atoms/finding.py @@ -0,0 +1,204 @@ +"""finding.py — atom positions, refined on the GPU (#75, #77). + +The division of labour, and the reason this wave is small: + +* **atomap owns the structure.** Initial peak finding, sublattices, nearest + neighbours, zone axes and dumbbell pairing are its job, and reimplementing + them would mean diverging from published atomap results. +* **SpyDE owns the refinement.** Refining atom positions IS a batched 2-D + gaussian fit, which :mod:`spyde.fitting` already does for the whole field at + once. atomap fits one atom at a time with scipy; the engine fits every atom + in one batched Levenberg-Marquardt. + +The refinement here is deliberately independent of atomap so it can be tested +(and used) without the extra — only :func:`find_atoms` needs it, and that is +``requires_package``-gated at the toolbar. +""" +from __future__ import annotations + +import logging + +import numpy as np + +log = logging.getLogger(__name__) + + +class MissingExtra(RuntimeError): + """Raised when the ``atoms`` extra is needed but not installed.""" + + +def _require_atomap(): + try: + import atomap.api # noqa: F401 + except ImportError as e: + raise MissingExtra( + 'atom finding needs atomap — install with: pip install "spyde[atoms]"' + ) from e + + +def find_atoms(signal, *, separation: float = 10.0, threshold_rel: float = 0.2, + pca: bool = False, subtract_background: bool = False): + """Initial atom positions via atomap -> ``(N, 2)`` array of ``(x, y)``. + + *separation* is the minimum distance between atoms in pixels and is the one + parameter that matters: too small splits one atom into several, too large + merges neighbours. It is swept interactively in the UI (#76). + """ + _require_atomap() + import atomap.api as am + + data = np.asarray(signal.data if hasattr(signal, "data") else signal) + positions = am.get_atom_positions( + _as_signal(data), separation=float(separation), + threshold_rel=float(threshold_rel), pca=pca, + subtract_background=subtract_background) + return np.asarray(positions, float) + + +def _as_signal(data): + import hyperspy.api as hs + return data if hasattr(data, "axes_manager") else hs.signals.Signal2D(data) + + +def refine_center_of_mass(image, positions, *, box: int = 7, + iterations: int = 2) -> np.ndarray: + """Centre-of-mass refinement in a box around each atom. + + Cheap, robust, and a good starting point for the gaussian fit — but biased + toward the box centre when neighbours intrude, which is exactly why it is + a *pre*-step rather than the answer. + + Vectorised over atoms: one gather of all boxes, then a weighted mean. No + Python loop over atoms. + """ + img = np.asarray(image, float) + pos = np.asarray(positions, float).reshape(-1, 2).copy() + h, w = img.shape + r = int(box) // 2 + oy, ox = np.mgrid[-r:r + 1, -r:r + 1] + + for _ in range(int(iterations)): + cx = np.clip(np.rint(pos[:, 0]).astype(int), r, w - r - 1) + cy = np.clip(np.rint(pos[:, 1]).astype(int), r, h - r - 1) + ys = cy[:, None, None] + oy[None] + xs = cx[:, None, None] + ox[None] + patch = img[ys, xs] + patch = np.clip(patch - patch.min(axis=(1, 2), keepdims=True), 0, None) + total = patch.sum(axis=(1, 2)) + good = total > 0 + # An empty box would divide by zero; leave those atoms where they are. + pos[good, 0] = (cx[good] + + (patch * ox[None]).sum(axis=(1, 2))[good] / total[good]) + pos[good, 1] = (cy[good] + + (patch * oy[None]).sum(axis=(1, 2))[good] / total[good]) + return pos + + +def refine_gaussian(image, positions, *, box: int = 11, sigma: float = 2.5, + device=None, max_iter: int = 60, refine_mask=None): + """Batched 2-D gaussian refinement of every atom at once. + + Each atom becomes one row of a ``(N, box*box)`` stack and one fit in the + batched engine — the same code path a spectrum image uses, with + :func:`~spyde.fitting.components.image_coordinates` as the axis. + + Parameters + ---------- + refine_mask : array of bool, optional + False for atoms whose position is pinned (the red/green toggle in + #76). A pinned atom keeps its input position exactly and is excluded + from the fit — not fitted and then discarded, which would waste the + work and let a diverging neighbour perturb the shared solve. + + Returns + ------- + (positions, params) + Refined ``(N, 2)`` ``(x, y)``, and the full per-atom parameter table + (``A``, ``centre_x``, ``centre_y``, ``sigma_x``, ``sigma_y``) in image + coordinates, which is what the property maps (#79) are built from. + """ + from spyde.fitting import ModelSpec + from spyde.fitting.components import image_coordinates + from spyde.fitting.engine import fit_batched + from spyde.fitting.spec import ComponentSpec, ParameterSpec + + img = np.asarray(image, float) + pos = np.asarray(positions, float).reshape(-1, 2) + n = len(pos) + mask = (np.ones(n, bool) if refine_mask is None + else np.asarray(refine_mask, bool).ravel()) + if mask.size != n: + raise ValueError(f"refine_mask has {mask.size} entries for {n} atoms") + + h, w = img.shape + r = int(box) // 2 + oy, ox = np.mgrid[-r:r + 1, -r:r + 1] + # Clamp box centres so every patch lies inside the image; the offset is + # tracked so results come back in IMAGE coordinates, not patch ones. + cx = np.clip(np.rint(pos[:, 0]).astype(int), r, w - r - 1) + cy = np.clip(np.rint(pos[:, 1]).astype(int), r, h - r - 1) + + out_pos = pos.copy() + out_par = np.full((n, 5), np.nan) + if not mask.any(): + return out_pos, out_par + + sel = np.flatnonzero(mask) + patches = img[cy[sel][:, None, None] + oy[None], + cx[sel][:, None, None] + ox[None]].reshape(len(sel), -1) + + # Start each atom at the CENTRE of its own patch, with the amplitude scaled + # to the patch — a gaussian's A is its VOLUME, so seeding it with the peak + # height would start every fit an order of magnitude low. + peak = patches.max(1) + spec = ModelSpec(components=[ComponentSpec( + kind="Gaussian2D", parameters=[ + ParameterSpec("A", 1.0, linear=True), + ParameterSpec("centre_x", float(r)), + ParameterSpec("centre_y", float(r)), + ParameterSpec("sigma_x", float(sigma), bmin=0.3, bmax=float(box)), + ParameterSpec("sigma_y", float(sigma), bmin=0.3, bmax=float(box)), + ])]) + names = spec.parameter_names() + start = np.broadcast_to(spec.flat_values(), (len(sel), 5)).copy() + start[:, names.index("Gaussian2D.A")] = peak * 2 * np.pi * sigma * sigma + start[:, names.index("Gaussian2D.centre_x")] = pos[sel, 0] - cx[sel] + r + start[:, names.index("Gaussian2D.centre_y")] = pos[sel, 1] - cy[sel] + r + + xy = image_coordinates((box, box)).numpy().astype(float) + res = fit_batched(spec, patches, xy, device=device, max_iter=max_iter, + initial=start) + + px = res.values[:, names.index("Gaussian2D.centre_x")] + cx[sel] - r + py = res.values[:, names.index("Gaussian2D.centre_y")] + cy[sel] - r + + # A fit that ran away from its own box is worse than the input. Keep the + # starting position for those rather than returning a wild coordinate. + ok = (np.abs(px - pos[sel, 0]) <= r) & (np.abs(py - pos[sel, 1]) <= r) \ + & np.isfinite(px) & np.isfinite(py) + out_pos[sel[ok], 0] = px[ok] + out_pos[sel[ok], 1] = py[ok] + out_par[sel] = res.values + out_par[sel, names.index("Gaussian2D.centre_x")] += cx[sel] - r + out_par[sel, names.index("Gaussian2D.centre_y")] += cy[sel] - r + if not ok.all(): + log.info("%d/%d atom fits left their box and kept their input " + "position", int((~ok).sum()), len(sel)) + return out_pos, out_par + + +def refine_atoms(image, positions, *, com_box: int = 7, box: int = 11, + sigma: float = 2.5, device=None, refine_mask=None): + """Centre-of-mass then batched gaussian — the standard two-step. + + The COM pass is cheap and pulls a rough peak-find onto the atom; the + gaussian pass is what gives sub-pixel accuracy. Running the gaussian alone + from a poor start is markedly less reliable, which is why atomap does the + same two steps. + """ + coarse = refine_center_of_mass(image, positions, box=com_box) + if refine_mask is not None: + keep = ~np.asarray(refine_mask, bool).ravel() + coarse[keep] = np.asarray(positions, float).reshape(-1, 2)[keep] + return refine_gaussian(image, coarse, box=box, sigma=sigma, device=device, + refine_mask=refine_mask) diff --git a/spyde/atoms/properties.py b/spyde/atoms/properties.py new file mode 100644 index 0000000..d3fced7 --- /dev/null +++ b/spyde/atoms/properties.py @@ -0,0 +1,128 @@ +"""properties.py — per-atom property maps (#79). + +Each function turns the refined atom table into one value per atom, which the +UI ships through the same ``commit_result_tree`` view mechanism as the strain +components (single click shows one, cmd-click tiles several). No new display +code — see :mod:`spyde.actions.views`. + +Everything here is vectorised over atoms. These are cheap compared with the +fit, but a Python loop over 10k atoms is still a visible pause, and the +neighbour search is the part that would be worst. +""" +from __future__ import annotations + +import logging + +import numpy as np + +log = logging.getLogger(__name__) + + +def _neighbours(positions, k: int): + """Indices and distances of each atom's k nearest neighbours. + + Uses a KD-tree; a brute-force (N, N) distance matrix is 800 MB at 10k atoms + and quadratic beyond that. + """ + from scipy.spatial import cKDTree + + pos = np.asarray(positions, float).reshape(-1, 2) + tree = cKDTree(pos) + # k+1 because the first hit is always the atom itself. + dist, idx = tree.query(pos, k=min(k + 1, len(pos))) + return idx[:, 1:], dist[:, 1:] + + +def ellipticity(params) -> np.ndarray: + """``sigma_max / sigma_min`` per atom — always >= 1. + + Expressed as a ratio of the LARGER to the smaller axis rather than + ``sigma_x / sigma_y``, so the value does not flip below 1 when an atom + happens to be elongated along y instead of x. An orientation-dependent + "ellipticity" would show a spurious boundary wherever the elongation + direction changes. + """ + p = np.asarray(params, float) + sx, sy = np.abs(p[:, 3]), np.abs(p[:, 4]) + big = np.maximum(sx, sy) + small = np.minimum(sx, sy) + with np.errstate(divide="ignore", invalid="ignore"): + return np.where(small > 0, big / small, np.nan) + + +def ellipticity_angle(params) -> np.ndarray: + """0 where an atom is wider in x, pi/2 where it is wider in y. + + A separate map from :func:`ellipticity` on purpose: magnitude and direction + answer different questions, and folding them together is what makes an + ellipticity map hard to read. + """ + p = np.asarray(params, float) + return np.where(np.abs(p[:, 3]) >= np.abs(p[:, 4]), 0.0, np.pi / 2) + + +def intensity(params) -> np.ndarray: + """Fitted gaussian VOLUME per atom (the ``A`` parameter). + + Volume, not peak height: it is what scales with scattering power and is + insensitive to a slightly wider or narrower fit, so it is the more stable + of the two for comparing sites. + """ + return np.asarray(params, float)[:, 0] + + +def nearest_neighbour_distance(positions, k: int = 1) -> np.ndarray: + """Mean distance to the *k* nearest neighbours, per atom.""" + _, dist = _neighbours(positions, k) + return dist.mean(1) if dist.ndim > 1 else dist + + +def displacement_from_ideal(positions, *, k: int = 4) -> np.ndarray: + """Each atom's offset from the centroid of its *k* nearest neighbours. + + Returns ``(N, 2)`` in ``(dx, dy)``. This is a *local* reference, so it + measures a genuine local distortion and is immune to sample drift or a + tilted scan — a globally-fitted ideal lattice would report both as + displacement. + """ + pos = np.asarray(positions, float).reshape(-1, 2) + idx, _ = _neighbours(pos, k) + return pos - pos[idx].mean(1) + + +def displacement_magnitude(positions, *, k: int = 4) -> np.ndarray: + d = displacement_from_ideal(positions, k=k) + return np.hypot(d[:, 0], d[:, 1]) + + +def to_map(values, positions, shape, *, fill=np.nan) -> np.ndarray: + """Scatter per-atom values onto an image-shaped array. + + Nearest-pixel placement, so the result lines up with the image the atoms + were found in and can be shown beside it. Atoms outside the shape are + dropped rather than wrapped. + """ + pos = np.asarray(positions, float).reshape(-1, 2) + vals = np.asarray(values, float).ravel() + if len(vals) != len(pos): + raise ValueError(f"{len(vals)} values for {len(pos)} atoms") + out = np.full(tuple(shape), fill, float) + ix = np.rint(pos[:, 0]).astype(int) + iy = np.rint(pos[:, 1]).astype(int) + inside = (ix >= 0) & (ix < shape[1]) & (iy >= 0) & (iy < shape[0]) + out[iy[inside], ix[inside]] = vals[inside] + return out + + +def property_maps(positions, params, *, k: int = 4) -> dict[str, np.ndarray]: + """Every per-atom property in one dict, keyed for the view chips (#79). + + The keys are what the user sees on the toggle, so they are spelled for a + human rather than after the function names. + """ + return { + "Ellipticity": ellipticity(params), + "Intensity": intensity(params), + "NN distance": nearest_neighbour_distance(positions), + "Displacement": displacement_magnitude(positions, k=k), + } diff --git a/spyde/data/__init__.py b/spyde/data/__init__.py index dd63abc..44896f3 100644 --- a/spyde/data/__init__.py +++ b/spyde/data/__init__.py @@ -21,10 +21,12 @@ from __future__ import annotations from spyde.data.synthetic import ( + atom_lattice, ebsd_patterns, eds_si, eels_si, ground_truth, ) -__all__ = ["eels_si", "eds_si", "ebsd_patterns", "ground_truth"] +__all__ = ["eels_si", "eds_si", "ebsd_patterns", "atom_lattice", + "ground_truth"] diff --git a/spyde/data/synthetic.py b/spyde/data/synthetic.py index 69d5e54..9bfbbc7 100644 --- a/spyde/data/synthetic.py +++ b/spyde/data/synthetic.py @@ -301,6 +301,81 @@ def eds_si(nav=(16, 16), n_channels: int = 2048, *, e_stop: float = 20.0, # EBSD # --------------------------------------------------------------------------- +def atom_lattice(grid=(8, 10), spacing: float = 16.0, *, sigma: float = 2.6, + dumbbell: float = 0.0, displacement: float = 0.0, + ellipticity: float = 0.0, noise: float = 0.01, + seed: int = 0): + """A HAADF-like atomic-resolution image with **exactly known positions**. + + Atoms sit on a rectangular lattice of gaussians. Every optional feature + exists so a downstream measurement has something real to recover: + + ``dumbbell`` + Split each site into a PAIR separated by this many pixels along x, for + the dumbbell workflow. 0 leaves single atoms. + ``displacement`` + Amplitude of a smooth sinusoidal displacement field, so a + displacement/strain map is non-trivial rather than uniformly zero. + ``ellipticity`` + Stretches sigma_x against sigma_y across the field, so an ellipticity + map has a known gradient to reproduce. + + The grid is deliberately NON-SQUARE and the spacing along x and y is the + same, so a transposed result is obvious rather than plausible. + + Ground truth (``metadata.Spyde.synthetic``, read with :func:`ground_truth`) + carries ``positions`` as ``(N, 2)`` in ``(x, y)`` pixel order — atomap's + convention, so a comparison needs no reindexing. + """ + import hyperspy.api as hs + + ny, nx = int(grid[0]), int(grid[1]) + rng = np.random.default_rng(seed) + margin = spacing + h = int(round(margin * 2 + spacing * (ny - 1))) + w = int(round(margin * 2 + spacing * (nx - 1))) + + gy, gx = np.mgrid[0:ny, 0:nx].astype(float) + xs = margin + gx * spacing + ys = margin + gy * spacing + if displacement: + xs = xs + displacement * np.sin(2 * np.pi * gy / max(ny - 1, 1)) + ys = ys + displacement * np.cos(2 * np.pi * gx / max(nx - 1, 1)) + + sx = sigma * (1.0 + ellipticity * gx / max(nx - 1, 1)) + sy = np.full_like(sx, sigma) + + centres = [] + for iy in range(ny): + for ix in range(nx): + if dumbbell: + centres.append((xs[iy, ix] - dumbbell / 2, ys[iy, ix], + sx[iy, ix], sy[iy, ix])) + centres.append((xs[iy, ix] + dumbbell / 2, ys[iy, ix], + sx[iy, ix], sy[iy, ix])) + else: + centres.append((xs[iy, ix], ys[iy, ix], sx[iy, ix], sy[iy, ix])) + + yy, xx = np.mgrid[0:h, 0:w].astype(np.float32) + img = np.zeros((h, w), np.float32) + for cx, cy, s_x, s_y in centres: + img += np.exp(-0.5 * (((xx - cx) / s_x) ** 2 + ((yy - cy) / s_y) ** 2)) + img /= max(img.max(), 1e-9) + if noise: + img = img + rng.normal(0.0, noise, img.shape).astype(np.float32) + + s = hs.signals.Signal2D(np.clip(img, 0, None).astype(np.float32)) + for a in s.axes_manager.signal_axes: + a.units, a.scale = "nm", 0.01 + s.metadata.General.title = "Synthetic atom lattice" + _stamp(s, kind="atoms", + positions=np.array([[c[0], c[1]] for c in centres], float), + grid=(ny, nx), spacing=float(spacing), sigma=float(sigma), + dumbbell=float(dumbbell), displacement=float(displacement), + ellipticity=float(ellipticity)) + return s + + def _euler_to_matrix(phi1, Phi, phi2) -> np.ndarray: """Bunge ZXZ Euler angles -> rotation matrices, batched over the leading axes. Returns ``(..., 3, 3)``.""" diff --git a/spyde/tests/migrated/test_atoms.py b/spyde/tests/migrated/test_atoms.py new file mode 100644 index 0000000..e802e4d --- /dev/null +++ b/spyde/tests/migrated/test_atoms.py @@ -0,0 +1,249 @@ +"""Atom finding, refinement and property maps (#75, #77, #79). + +The synthetic lattice knows exactly where every atom is, so refinement is +scored against those positions rather than against atomap's answer or a +fixture. Sub-pixel accuracy is the whole point of the gaussian step, so the +tolerances here are tight on purpose — a test that accepts +/-1 px would pass +with the refinement removed entirely. +""" +from __future__ import annotations + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") + +from spyde.atoms import ( + displacement_from_ideal, + ellipticity, + intensity, + nearest_neighbour_distance, + property_maps, + refine_atoms, + refine_center_of_mass, + refine_gaussian, + to_map, +) +from spyde.data import atom_lattice, ground_truth + + +@pytest.fixture(scope="module") +def lattice(): + s = atom_lattice(grid=(6, 8), spacing=16.0, noise=0.0) + return np.asarray(s.data, float), np.asarray(ground_truth(s)["positions"]) + + +def _jitter(pos, amount, seed=0): + rng = np.random.default_rng(seed) + return pos + rng.uniform(-amount, amount, pos.shape) + + +class TestCenterOfMass: + def test_pulls_a_jittered_guess_back_toward_the_atom(self, lattice): + img, truth = lattice + start = _jitter(truth, 2.0) + got = refine_center_of_mass(img, start, box=9) + assert (np.hypot(*(got - truth).T).mean() + < np.hypot(*(start - truth).T).mean()) + + def test_does_not_move_an_already_centred_atom_much(self, lattice): + img, truth = lattice + got = refine_center_of_mass(img, truth, box=9) + assert np.hypot(*(got - truth).T).max() < 0.6 + + def test_handles_atoms_near_the_edge(self, lattice): + """Box clamping must keep every patch inside the image — an atom in + the corner would otherwise index out of bounds.""" + img, truth = lattice + corner = truth[:1].copy() + corner[0] = [1.0, 1.0] + got = refine_center_of_mass(img, corner, box=11) + assert np.isfinite(got).all() + + +class TestGaussianRefinement: + def test_recovers_positions_to_sub_pixel(self, lattice): + """The reason the gaussian step exists. +/-0.1 px would not be + achievable from centre-of-mass alone on this data.""" + img, truth = lattice + pos, params = refine_gaussian(img, _jitter(truth, 1.0), box=13, + device="cpu") + err = np.hypot(*(pos - truth).T) + assert err.max() < 0.1, f"worst atom off by {err.max():.3f} px" + assert params.shape == (len(truth), 5) + + def test_two_step_beats_the_gaussian_alone_from_a_poor_start(self, lattice): + """Why atomap does COM first and so do we.""" + img, truth = lattice + start = _jitter(truth, 3.0, seed=3) + only_gauss, _ = refine_gaussian(img, start, box=13, device="cpu") + two_step, _ = refine_atoms(img, start, box=13, device="cpu") + assert (np.hypot(*(two_step - truth).T).mean() + <= np.hypot(*(only_gauss - truth).T).mean()) + + def test_recovers_the_known_widths(self, lattice): + img, truth = lattice + _, params = refine_gaussian(img, truth, box=13, sigma=2.5, device="cpu") + # The generator used sigma=2.6 with no ellipticity. + assert np.abs(params[:, 3] - 2.6).max() < 0.15 + assert np.abs(params[:, 4] - 2.6).max() < 0.15 + + def test_pinned_atoms_keep_their_position_exactly(self, lattice): + """The red/green refine toggle (#76). A pinned atom must be EXCLUDED + from the fit, not fitted and then discarded.""" + img, truth = lattice + start = _jitter(truth, 1.0) + mask = np.ones(len(truth), bool) + mask[::3] = False + pos, params = refine_gaussian(img, start, box=13, device="cpu", + refine_mask=mask) + np.testing.assert_array_equal(pos[~mask], start[~mask]) + assert np.isnan(params[~mask]).all() + assert not np.allclose(pos[mask], start[mask]) + + def test_all_atoms_pinned_is_a_no_op(self, lattice): + img, truth = lattice + pos, params = refine_gaussian(img, truth, device="cpu", + refine_mask=np.zeros(len(truth), bool)) + np.testing.assert_array_equal(pos, truth) + assert np.isnan(params).all() + + def test_mask_length_is_checked(self, lattice): + img, truth = lattice + with pytest.raises(ValueError, match="refine_mask"): + refine_gaussian(img, truth, device="cpu", + refine_mask=np.ones(3, bool)) + + def test_a_runaway_fit_keeps_its_input_position(self): + """A fit that leaves its own box is worse than the input, so the input + is what comes back — never a wild coordinate.""" + img = np.zeros((40, 40)) # nothing to fit at all + start = np.array([[20.0, 20.0]]) + pos, _ = refine_gaussian(img, start, box=11, device="cpu") + assert np.hypot(*(pos - start).T).max() <= 5.5 + + def test_ellipticity_is_recovered_when_present(self): + s = atom_lattice(grid=(4, 6), spacing=18.0, ellipticity=0.6, noise=0.0) + img = np.asarray(s.data, float) + truth = np.asarray(ground_truth(s)["positions"]) + _, params = refine_gaussian(img, truth, box=15, device="cpu") + e = ellipticity(params) + # The generator stretches sigma_x with x, so the right-hand atoms must + # be measurably more elliptical than the left-hand ones. + left = truth[:, 0] < truth[:, 0].mean() + assert np.nanmean(e[~left]) > np.nanmean(e[left]) + 0.1 + + +class TestPropertyMaps: + def test_ellipticity_is_at_least_one(self, lattice): + img, truth = lattice + _, params = refine_gaussian(img, truth, box=13, device="cpu") + e = ellipticity(params) + assert np.nanmin(e) >= 1.0 - 1e-9 + + def test_ellipticity_does_not_flip_with_orientation(self): + """Expressed as max/min, so an atom elongated along y reads the same as + one elongated along x. A sigma_x/sigma_y ratio would drop below 1 and + put a spurious boundary wherever the elongation direction changes.""" + wide_x = np.array([[1.0, 0, 0, 3.0, 1.5]]) + wide_y = np.array([[1.0, 0, 0, 1.5, 3.0]]) + assert ellipticity(wide_x)[0] == pytest.approx(ellipticity(wide_y)[0]) + + def test_nn_distance_matches_the_lattice_spacing(self, lattice): + img, truth = lattice + d = nearest_neighbour_distance(truth) + assert np.median(d) == pytest.approx(16.0, rel=0.02) + + def test_displacement_is_zero_on_a_perfect_lattice(self, lattice): + """A local reference means an undistorted lattice reads zero — the + baseline every distortion map is measured against.""" + img, truth = lattice + inner = ((truth[:, 0] > 20) & (truth[:, 0] < truth[:, 0].max() - 20) + & (truth[:, 1] > 20) & (truth[:, 1] < truth[:, 1].max() - 20)) + d = displacement_from_ideal(truth) + assert np.abs(d[inner]).max() < 1e-6 + + def test_displacement_finds_a_real_distortion(self): + s = atom_lattice(grid=(6, 8), spacing=18.0, displacement=2.0, + noise=0.0) + truth = np.asarray(ground_truth(s)["positions"]) + assert np.abs(displacement_from_ideal(truth)).max() > 0.5 + + def test_intensity_is_the_fitted_volume(self, lattice): + img, truth = lattice + _, params = refine_gaussian(img, truth, box=13, device="cpu") + assert np.allclose(intensity(params), params[:, 0]) + assert (intensity(params) > 0).all() + + def test_property_maps_returns_one_value_per_atom(self, lattice): + img, truth = lattice + _, params = refine_gaussian(img, truth, box=13, device="cpu") + maps = property_maps(truth, params) + assert set(maps) == {"Ellipticity", "Intensity", "NN distance", + "Displacement"} + for name, v in maps.items(): + assert v.shape == (len(truth),), name + + +class TestToMap: + def test_places_values_at_their_atoms(self, lattice): + img, truth = lattice + vals = np.arange(len(truth), dtype=float) + m = to_map(vals, truth, img.shape) + ix = int(round(truth[5, 0])) + iy = int(round(truth[5, 1])) + assert m[iy, ix] == pytest.approx(5.0) + + def test_pixels_without_an_atom_are_nan(self, lattice): + img, truth = lattice + m = to_map(np.ones(len(truth)), truth, img.shape) + assert np.isnan(m).any() + + def test_atoms_outside_the_shape_are_dropped_not_wrapped(self): + m = to_map(np.array([1.0, 2.0]), np.array([[5.0, 5.0], [500.0, 5.0]]), + (20, 20)) + assert m[5, 5] == pytest.approx(1.0) + assert np.nansum(m) == pytest.approx(1.0) + + def test_length_mismatch_is_caught(self, lattice): + img, truth = lattice + with pytest.raises(ValueError, match="values for"): + to_map(np.ones(3), truth, img.shape) + + +class TestAtomapIntegration: + def test_find_atoms_locates_the_lattice(self, lattice): + pytest.importorskip("atomap", reason="needs the atoms extra") + from spyde.atoms import find_atoms + + img, truth = lattice + found = find_atoms(img, separation=10) + assert len(found) == pytest.approx(len(truth), rel=0.15) + # Every found atom should be near a real one. + from scipy.spatial import cKDTree + d, _ = cKDTree(truth).query(found) + assert np.median(d) < 2.0 + + def test_find_then_refine_recovers_the_truth(self, lattice): + """The full pipeline, scored against known positions.""" + pytest.importorskip("atomap", reason="needs the atoms extra") + from scipy.spatial import cKDTree + from spyde.atoms import find_atoms + + img, truth = lattice + found = find_atoms(img, separation=10) + refined, _ = refine_atoms(img, found, box=13, device="cpu") + d, _ = cKDTree(truth).query(refined) + assert np.median(d) < 0.2, f"median error {np.median(d):.3f} px" + + def test_missing_extra_gives_the_install_line(self, monkeypatch): + import spyde.atoms.finding as f + + def boom(): + raise f.MissingExtra( + 'atom finding needs atomap — install with: ' + 'pip install "spyde[atoms]"') + + monkeypatch.setattr(f, "_require_atomap", boom) + with pytest.raises(f.MissingExtra, match=r"spyde\[atoms\]"): + f.find_atoms(np.zeros((10, 10))) From 5d6d396ab53db1273bb0780895121bb276520267 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Sun, 26 Jul 2026 23:09:46 -0500 Subject: [PATCH 11/60] fix(fitting): Polynomial and Offset broke for any batch bigger than one MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both shaped their output with `expand_as` / `expand`, which forces the parameter block's P down to the axis's leading 1 and raises for P > 1: RuntimeError: The expanded size of the tensor (1) must match the existing size (4) at non-singleton dimension 0 Plain broadcasting (`p + 0.0 * x`) gives (P, C) for both a shared axis and a per-position one. Found by fitting a real EDS composition model, whose background is a Polynomial — not by the component tests, because every one of them used a SINGLE spectrum, where expand_as happens to work. That is the actual defect here: a batched component library whose tests never batch. Closed with TestEveryComponentBatches, which runs every component (and Polynomial) with P=4 and asserts the value shape, the gradient shape, and that row 1 reflects ITS OWN parameters — a broadcasting slip that returns row 0 for everything would still have the right shape. --- spyde/fitting/components.py | 13 +++-- spyde/tests/migrated/test_torch_components.py | 52 +++++++++++++++++++ 2 files changed, 61 insertions(+), 4 deletions(-) diff --git a/spyde/fitting/components.py b/spyde/fitting/components.py index 578447b..76d68c4 100644 --- a/spyde/fitting/components.py +++ b/spyde/fitting/components.py @@ -163,8 +163,9 @@ def _power_law(x, p): def _offset(x, p): (offset,) = p - return offset.expand(-1, x.shape[-1]) if x.shape[0] == 1 else \ - offset + 0.0 * x + # Broadcast rather than expand — see _polynomial_fn for why expand-based + # shaping breaks as soon as there is more than one spectrum. + return offset + 0.0 * x def _exponential(x, p): @@ -375,8 +376,12 @@ def _polynomial_fn(order: int): Parameters are ``a0..a{order}`` and ``aK`` multiplies ``x**K``.""" def fn(x, p): - out = p[0].expand_as(x) if len(p) else None - acc = out + # `p[0] + 0.0 * x`, NOT `p[0].expand_as(x)`. The parameter block is + # (P, 1) and the axis is (1, C), so expand_as tries to force P down to + # 1 and raises for any batch bigger than one spectrum — which single- + # spectrum tests never reach. Plain broadcasting gives (P, C) for both + # a shared axis and a per-position one. + acc = p[0] + 0.0 * x for k in range(1, order + 1): acc = acc + p[k] * x ** k return acc diff --git a/spyde/tests/migrated/test_torch_components.py b/spyde/tests/migrated/test_torch_components.py index 3305625..8175a06 100644 --- a/spyde/tests/migrated/test_torch_components.py +++ b/spyde/tests/migrated/test_torch_components.py @@ -230,6 +230,58 @@ def test_offset_broadcasts_over_the_signal_axis(self): assert torch.allclose(y[1], torch.full_like(y[1], 5.0)) +@pytest.mark.parametrize("kind", sorted(CASES) + ["Polynomial"]) +class TestEveryComponentBatches: + """Every component, with P > 1. + + This exists because a single-spectrum test is not enough and once wasn't: + `Polynomial` and `Offset` shaped their output with `expand_as`, which + happens to work when P == 1 and raises for any real batch. Every other + component test used one row, so a whole EDS model failed only at the point + of actually fitting a scan. + """ + + def _component_and_values(self, kind): + if kind == "Polynomial": + comp = c1d.Polynomial(order=2) + for k in range(3): + getattr(comp, f"a{k}").value = 0.5 * (k + 1) + batched = tc.get_component("Polynomial", n_params=3) + else: + comp = _hyperspy_component(kind) + batched = tc.get_component(kind) + base = np.array([getattr(comp, n).value for n in batched.params]) + return batched, base + + def test_value_batches(self, kind): + batched, base = self._component_and_values(kind) + vals = np.stack([base * (1.0 + 0.1 * i) for i in range(4)]) + y = batched(torch.as_tensor(X, dtype=torch.float64), + torch.as_tensor(vals, dtype=torch.float64)) + assert y.shape == (4, len(X)), f"{kind} returned {tuple(y.shape)}" + + def test_gradient_batches(self, kind): + batched, base = self._component_and_values(kind) + if not batched.has_analytic_grad: + pytest.skip(f"{kind} has no analytic gradient") + vals = np.stack([base * (1.0 + 0.1 * i) for i in range(4)]) + g = batched.grad(torch.as_tensor(X, dtype=torch.float64), + torch.as_tensor(vals, dtype=torch.float64)) + assert g.shape == (4, len(X), batched.n_params), \ + f"{kind} gradient returned {tuple(g.shape)}" + + def test_batched_rows_are_independent(self, kind): + """Each row must reflect its OWN parameters — a broadcasting slip that + returns row 0 for everything would still have the right shape.""" + batched, base = self._component_and_values(kind) + vals = np.stack([base, base * 2.0]) + y = batched(torch.as_tensor(X, dtype=torch.float64), + torch.as_tensor(vals, dtype=torch.float64)) + one = batched(torch.as_tensor(X, dtype=torch.float64), + torch.as_tensor(vals[1:2], dtype=torch.float64)) + torch.testing.assert_close(y[1], one[0]) + + class TestPolynomial: @pytest.mark.parametrize("order", [1, 2, 3]) def test_matches_hyperspy_at_each_order(self, order): From 689e1a9a165383d88684e5c8560f539accc651d0 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Sun, 26 Jul 2026 23:09:46 -0500 Subject: [PATCH 12/60] feat(spectroscopy): composition -> auto-populated model (#62) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wave 2's headline: type which elements are present, get a populated model. exspy owns the spectroscopy — edge and line energies, GOS edge shapes, family relationships — because reimplementing it would mean diverging from the numbers the community publishes. This module owns the plumbing: - one call for EELS and EDS alike returning a ModelSpec, so the result feeds the batched engine, the wizard and the renderer through one type; - pruning to the measured range. A component with no data under it is NOT harmless: its amplitude is unconstrained, so the optimiser trades it against everything else and degrades the parameters that ARE measurable. exspy will happily add a Cu-L line at 0.93 keV to a spectrum whose useful range starts at 2 keV; - an honest report of whether the batched engine can fit the result. That last point is asymmetric today and worth knowing: EDS -> Polynomial + Gaussian -> fully supported, gets the GPU path NOW EELS -> EELSCLEdge (GOS table) -> no batched port yet (#63), falls back to hyperspy, which is correct but slower Reported in the returned info rather than left as a surprise. Two exspy behaviours that do not do what their names suggest: - `add_lines` APPENDS, so restricting needs `set_lines`. - Even then, the EDS model expands each ELEMENT into its whole family (selecting Fe_Ka still builds Fe_Kb, Fe_La, Fe_Ln, ...) regardless of metadata.Sample.xray_lines — so honouring only_lines means filtering the built spec by name. The source signal is never mutated; a user may be exploring several compositions. --- spyde/spectroscopy/__init__.py | 26 +++ spyde/spectroscopy/composition.py | 211 ++++++++++++++++++ .../tests/migrated/test_composition_model.py | 182 +++++++++++++++ 3 files changed, 419 insertions(+) create mode 100644 spyde/spectroscopy/__init__.py create mode 100644 spyde/spectroscopy/composition.py create mode 100644 spyde/tests/migrated/test_composition_model.py diff --git a/spyde/spectroscopy/__init__.py b/spyde/spectroscopy/__init__.py new file mode 100644 index 0000000..69e9dd1 --- /dev/null +++ b/spyde/spectroscopy/__init__.py @@ -0,0 +1,26 @@ +"""spyde.spectroscopy — EELS and EDS (0.3.0 Wave 2, GitHub #60). + +Sits on top of :mod:`spyde.fitting`: a composition becomes a ``ModelSpec``, +which the batched engine fits. + +**exspy owns the spectroscopy, this package owns the plumbing.** Edge and line +energies, GOS-based edge shapes, family relationships and quantification +factors all come from exspy so SpyDE's numbers match published work. What is +here is the uniform API, the range pruning, and the bridge to the batched +engine. + +Needs the ``eels`` extra (``pip install "spyde[eels]"``). Importing this +package is safe without it — the functions raise +:class:`~spyde.spectroscopy.composition.MissingExtra` with the install line, +and toolbar actions are hidden by ``requires_package`` so a user never clicks +into the error. +""" +from __future__ import annotations + +from spyde.spectroscopy.composition import ( + MissingExtra, + model_for_composition, + prune_to_range, +) + +__all__ = ["model_for_composition", "prune_to_range", "MissingExtra"] diff --git a/spyde/spectroscopy/composition.py b/spyde/spectroscopy/composition.py new file mode 100644 index 0000000..f8050af --- /dev/null +++ b/spyde/spectroscopy/composition.py @@ -0,0 +1,211 @@ +"""composition.py — elements in, a fitted-ready model out (#62). + +The headline of Wave 2: a user types which elements are present and gets a +populated model, instead of hand-placing a component per edge or line. + +**exspy does the spectroscopy, this module does the plumbing.** Where the +ionisation edges and X-ray lines are, what a GOS-based edge shape is, which +lines belong to a family — all of that is exspy's, and reimplementing it would +mean diverging from the numbers the community publishes. What is added here: + +* one call that works for EELS and EDS alike and returns a + :class:`~spyde.fitting.spec.ModelSpec`, so the result feeds the batched + engine, the wizard and the renderer through one type; +* **pruning to the measured range** — exspy will happily add a Cu-L line at + 0.93 keV to a spectrum whose useful range starts at 2 keV, and a component + with no data under it is an unconstrained parameter that makes the whole fit + worse, not merely a wasted one; +* an honest report of whether the batched engine can actually fit the result. + +That last point matters and is asymmetric today: + +* **EDS** builds from Polynomial + Gaussian, both of which the batched engine + implements — so EDS gets the full GPU path immediately. +* **EELS** builds from ``EELSCLEdge``, a tabulated GOS lookup with no batched + port yet (#63), so an EELS model falls back to HyperSpy's own fitting. + +:func:`model_for_composition` reports this rather than letting it be a +surprise, and the fallback is correct — just slower. +""" +from __future__ import annotations + +import logging + +import numpy as np + +from spyde.fitting import ModelSpec + +log = logging.getLogger(__name__) + +# The parameter that positions each kind of component on the energy axis. +# Used for range pruning; a component whose position parameter is unknown is +# kept, because dropping something we do not understand is the worse error. +_POSITION_PARAM = { + "Gaussian": "centre", + "GaussianHF": "centre", + "Lorentzian": "centre", + "Voigt": "centre", + "SplitVoigt": "centre", + "EELSCLEdge": "onset_energy", + "Erf": "origin", + "Arctan": "x0", +} + + +class MissingExtra(RuntimeError): + """Raised when the ``eels`` extra is needed but not installed.""" + + +def _require_exspy(): + try: + import exspy # noqa: F401 + except ImportError as e: + raise MissingExtra( + 'EELS/EDS models need exspy — install with: pip install "spyde[eels]"' + ) from e + + +def _signal_kind(signal) -> str: + st = (getattr(signal, "_signal_type", "") or "").upper() + if "EELS" in st: + return "EELS" + if "EDS" in st: + return "EDS" + raise ValueError( + f"signal type {getattr(signal, '_signal_type', None)!r} is neither " + f"EELS nor EDS — call set_signal_type('EELS') / ('EDS_TEM') first " + f"(needs the eels extra)") + + +def _axis_range(signal) -> tuple[float, float]: + ax = signal.axes_manager.signal_axes[0].axis + return float(np.min(ax)), float(np.max(ax)) + + +def _component_position(comp_spec): + name = _POSITION_PARAM.get(comp_spec.kind) + if name is None: + return None + try: + return float(comp_spec[name].value) + except KeyError: + return None + + +def prune_to_range(spec: ModelSpec, lo: float, hi: float, *, + margin: float = 0.0) -> tuple[ModelSpec, list[str]]: + """Drop components positioned outside ``[lo, hi]``. + + A component with no data under it is not harmless: its amplitude is + unconstrained, so the optimiser is free to trade it against everything + else, which degrades the parameters that ARE measurable. + + Returns the pruned spec and the names removed. + """ + kept, dropped = [], [] + for c in spec.components: + pos = _component_position(c) + if pos is not None and not (lo - margin <= pos <= hi + margin): + dropped.append(c.name) + continue + kept.append(c) + return ModelSpec(components=kept, channel_mask=spec.channel_mask), dropped + + +def model_for_composition(signal, elements=None, *, prune: bool = True, + energy_range: tuple[float, float] | None = None, + only_lines=None): + """Build a :class:`ModelSpec` for the elements present in *signal*. + + Parameters + ---------- + signal + An EELS or EDS signal (needs the ``eels`` extra for the signal type to + resolve at all). + elements : sequence of str, optional + e.g. ``["Fe", "Ni", "Cu"]``. Defaults to whatever is already on + ``metadata.Sample.elements``. + prune : bool + Drop components positioned outside the measured range (see + :func:`prune_to_range`). + energy_range : (lo, hi), optional + Override the range used for pruning — e.g. to exclude a noisy + low-energy region the axis technically covers. + only_lines : sequence of str, optional + EDS only: restrict to these X-ray lines (``["Fe_Ka", "Cu_Ka"]``) + instead of every line exspy knows for the element. + + Returns + ------- + (spec, info) + *info* carries ``kind``, ``elements``, ``dropped`` and + ``engine_supported`` — the last being whether + :mod:`spyde.fitting.engine` can fit this model or whether it falls back + to HyperSpy (#63). + """ + _require_exspy() + from spyde.fitting import components as tcomp + + kind = _signal_kind(signal) + s = signal.deepcopy() + + if elements: + s.add_elements(list(elements)) + have = list(getattr(s.metadata, "Sample", {}).get_item("elements", []) + if hasattr(getattr(s.metadata, "Sample", None), "get_item") + else []) + if not have: + raise ValueError( + "no elements to build a model from — pass elements=[...] or set " + "metadata.Sample.elements") + + if kind == "EDS": + try: + if only_lines: + # set_lines REPLACES the line list; add_lines only appends, so + # using it here would leave every default line in place and + # `only_lines` would silently do nothing. + s.set_lines(list(only_lines)) + else: + s.add_lines() + except Exception as e: + log.debug("setting X-ray lines failed (%s); relying on " + "create_model's defaults", e) + + model = s.create_model() + spec = ModelSpec.from_model(model) + + if kind == "EDS" and only_lines: + # exspy's EDS model expands each ELEMENT into its whole family + # (selecting Fe_Ka still builds Fe_Kb, Fe_La, Fe_Ln, ...) regardless of + # metadata.Sample.xray_lines, so `set_lines` alone does not restrict + # the model. Filter the built spec by name instead, keeping anything + # that is not a named line (the background). + wanted = set(only_lines) + spec = ModelSpec( + components=[c for c in spec.components + if "_" not in c.name or c.name in wanted], + channel_mask=spec.channel_mask) + + lo, hi = energy_range if energy_range else _axis_range(s) + dropped: list[str] = [] + if prune: + spec, dropped = prune_to_range(spec, lo, hi) + if dropped: + log.info("pruned %d component(s) outside %.4g-%.4g: %s", + len(dropped), lo, hi, ", ".join(dropped)) + + info = { + "kind": kind, + "elements": have, + "dropped": dropped, + "engine_supported": tcomp.supports(spec), + "energy_range": (lo, hi), + } + if not info["engine_supported"]: + unsupported = sorted({c.kind for c in spec.active_components + if c.kind not in tcomp.available()}) + info["unsupported_components"] = unsupported + log.info("model uses %s, which the batched engine cannot fit yet " + "(#63) — HyperSpy will do this fit", unsupported) + return spec, info diff --git a/spyde/tests/migrated/test_composition_model.py b/spyde/tests/migrated/test_composition_model.py new file mode 100644 index 0000000..6aeaefd --- /dev/null +++ b/spyde/tests/migrated/test_composition_model.py @@ -0,0 +1,182 @@ +"""Composition -> auto-populated model (#62). + +Skipped wholesale without the `eels` extra, which is the point of the extra: +the code must be importable and the tests must skip cleanly rather than error, +because CI runs one job with the extras and one without. +""" +from __future__ import annotations + +import numpy as np +import pytest + +exspy = pytest.importorskip("exspy", reason="needs the eels extra") + +from spyde.data import eds_si, eels_si +from spyde.fitting import components as tcomp +from spyde.spectroscopy import MissingExtra, model_for_composition, prune_to_range + + +@pytest.fixture(scope="module") +def eds(): + return eds_si(nav=(2, 2), n_channels=1024) + + +@pytest.fixture(scope="module") +def eels(): + return eels_si(nav=(2, 2), n_channels=512) + + +class TestEDS: + def test_builds_a_component_per_line(self, eds): + spec, info = model_for_composition(eds, ["Fe", "Ni", "Cu"]) + assert info["kind"] == "EDS" + assert set(info["elements"]) == {"Fe", "Ni", "Cu"} + names = [c.name for c in spec] + assert any(n.startswith("Fe_K") for n in names), names + assert any(n.startswith("Cu_K") for n in names), names + + def test_lines_land_at_their_real_energies(self, eds): + """The whole value of using exspy: the energies are right without us + maintaining a table.""" + from spyde.data.synthetic import EDS_LINES + + spec, _ = model_for_composition(eds, ["Fe", "Ni", "Cu"]) + by_name = {c.name: c for c in spec} + for el, lines in EDS_LINES.items(): + comp = by_name.get(f"{el}_Ka") + assert comp is not None, f"{el}_Ka missing from {list(by_name)}" + assert comp["centre"].value == pytest.approx(lines[0][1], abs=0.02) + + def test_the_batched_engine_can_fit_an_EDS_model(self, eds): + """EDS builds from Polynomial + Gaussian, both of which the engine + implements — so EDS gets the GPU path with no fallback. This is the + asymmetry with EELS worth pinning.""" + spec, info = model_for_composition(eds, ["Fe", "Ni", "Cu"]) + assert info["engine_supported"] is True + assert tcomp.supports(spec) is True + + def test_only_lines_restricts_the_model(self, eds): + """`add_lines` only APPENDS, so restricting needs `set_lines` — the + obvious call silently leaves every default line in place.""" + spec, _ = model_for_composition(eds, ["Fe", "Cu"], + only_lines=["Fe_Ka", "Cu_Ka"]) + lines = [c.name for c in spec + if "_" in c.name and not c.name.startswith("background")] + assert set(lines) <= {"Fe_Ka", "Cu_Ka"}, lines + assert lines, "every line was dropped" + + +class TestEELS: + def test_builds_an_edge_per_element(self, eels): + spec, info = model_for_composition(eels, ["C", "N", "O"]) + assert info["kind"] == "EELS" + kinds = [c.kind for c in spec] + assert kinds.count("EELSCLEdge") == 3, kinds + assert "PowerLaw" in kinds # background + + def test_edges_land_at_their_real_onsets(self, eels): + from spyde.data.synthetic import EELS_EDGES + + spec, _ = model_for_composition(eels, ["C", "N", "O"]) + by_name = {c.name: c for c in spec} + for name, onset in EELS_EDGES.items(): + comp = by_name.get(name) + assert comp is not None, f"{name} missing from {list(by_name)}" + assert comp["onset_energy"].value == pytest.approx(onset, abs=5.0) + + def test_reports_that_the_engine_cannot_fit_it_yet(self, eels): + """EELSCLEdge is a tabulated GOS lookup with no batched port (#63). + The fallback to HyperSpy is correct, just slower — but it must be + REPORTED, not a surprise.""" + spec, info = model_for_composition(eels, ["C", "N", "O"]) + assert info["engine_supported"] is False + assert "EELSCLEdge" in info["unsupported_components"] + + +class TestPruning: + def test_drops_components_outside_the_measured_range(self, eds): + """A component with no data under it is not harmless — its amplitude is + unconstrained, so the optimiser trades it against everything else and + degrades the parameters that ARE measurable.""" + spec, info = model_for_composition(eds, ["Fe", "Ni", "Cu"], + energy_range=(5.0, 9.0)) + for c in spec: + if c.kind == "Gaussian": + assert 5.0 <= c["centre"].value <= 9.0, c.name + assert info["dropped"], "nothing was pruned from a narrowed range" + + def test_pruning_can_be_turned_off(self, eds): + wide, _ = model_for_composition(eds, ["Fe", "Ni", "Cu"], prune=False) + narrow, _ = model_for_composition(eds, ["Fe", "Ni", "Cu"], + energy_range=(5.0, 9.0)) + assert len(wide) > len(narrow) + + def test_keeps_components_whose_position_is_unknown(self): + """Dropping something we do not understand is the worse error — a + background has no 'position' and must survive.""" + from spyde.fitting.spec import ComponentSpec, ParameterSpec + + spec = ModelSpecShim = None # noqa: F841 (readability of the next line) + from spyde.fitting import ModelSpec + spec = ModelSpec(components=[ + ComponentSpec(kind="PowerLaw", parameters=[ParameterSpec("A", 1.0)]), + ComponentSpec(kind="Gaussian", name="far", + parameters=[ParameterSpec("centre", 999.0)]), + ]) + kept, dropped = prune_to_range(spec, 0.0, 20.0) + assert [c.kind for c in kept] == ["PowerLaw"] + assert dropped == ["far"] + + def test_margin_widens_the_window(self): + from spyde.fitting import ModelSpec + from spyde.fitting.spec import ComponentSpec, ParameterSpec + + spec = ModelSpec(components=[ComponentSpec( + kind="Gaussian", name="edge", + parameters=[ParameterSpec("centre", 20.5)])]) + assert prune_to_range(spec, 0.0, 20.0)[1] == ["edge"] + assert prune_to_range(spec, 0.0, 20.0, margin=1.0)[1] == [] + + +class TestGuards: + def test_a_plain_signal_is_rejected_with_guidance(self): + import hyperspy.api as hs + s = hs.signals.Signal1D(np.zeros((2, 2, 32))) + with pytest.raises(ValueError, match="neither EELS nor EDS"): + model_for_composition(s, ["Fe"]) + + def test_no_elements_is_a_clear_error(self, eds): + bare = eds.deepcopy() + try: + del bare.metadata.Sample + except Exception: + pass + with pytest.raises(ValueError, match="no elements"): + model_for_composition(bare) + + def test_the_source_signal_is_not_mutated(self, eds): + """Building a model must not silently add elements to the user's + signal — they may be exploring several compositions.""" + before = list(getattr(eds.metadata, "Sample", {}).get_item("elements", []) + if hasattr(getattr(eds.metadata, "Sample", None), "get_item") + else []) + model_for_composition(eds, ["Fe", "Ni", "Cu", "Zn"]) + after = list(getattr(eds.metadata, "Sample", {}).get_item("elements", []) + if hasattr(getattr(eds.metadata, "Sample", None), "get_item") + else []) + assert before == after + + +class TestEndToEnd: + def test_an_EDS_composition_model_actually_fits(self, eds): + """The point of the whole wave: elements in, fitted maps out, through + the batched engine.""" + from spyde.fitting.engine import fit_batched + + spec, info = model_for_composition(eds, ["Fe", "Ni", "Cu"], + energy_range=(5.5, 9.5)) + assert info["engine_supported"] + x = np.asarray(eds.axes_manager.signal_axes[0].axis, float) + got = fit_batched(spec, eds.data, x, device="cpu", max_iter=40) + assert got.values.shape[0] == 4 + assert np.isfinite(got.values).all() From 3ba7bb00505c338053bba4cb99f01bb718080524 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Mon, 27 Jul 2026 00:47:44 -0500 Subject: [PATCH 13/60] fix(fitting): ModelSpec could not rebuild real EELS or EDS models MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two bugs of the same shape, found by round-tripping a composition model. to_model() rebuilt every component BARE, which breaks for anything that takes constructor arguments: - EELSCLEdge raises outright (element_subshell is required), so no EELS model could be turned back into hyperspy at all. - Polynomial is worse because it does NOT raise. Rebuilt bare it comes back at its DEFAULT order, so exspy's order-6 EDS background silently lost a3..a6 — to_model skipped the parameters that no longer existed and returned a different model that still fits and still looks plausible. ComponentSpec now carries init_args, captured by from_model via a small _INIT_ARGS table and used when rebuilding. An unrebuildable component gets an error naming the table to add to, rather than a bare TypeError from inside hyperspy. The original round-trip tests missed both because they only ever went model -> spec, never spec -> model with a component of this kind. Fixing that surfaced a third: EELSCLEdge.fine_structure_coeff is a VECTOR parameter with 8 elements, and ParameterSpec stored a single float, so hyperspy rejected the assignment on length. Vector parameters now round-trip faithfully (value becomes a list) but are excluded from the PACKED vector via ComponentSpec.scalar_parameters — the solver puts one scalar per column, and no component with such a parameter is fittable by the batched engine anyway (components.supports already says no). components.py uses the same scalar view so the two cannot drift apart. --- spyde/fitting/components.py | 8 +- spyde/fitting/spec.py | 117 +++++++++++--- .../migrated/test_model_spec_init_args.py | 149 ++++++++++++++++++ 3 files changed, 252 insertions(+), 22 deletions(-) create mode 100644 spyde/tests/migrated/test_model_spec_init_args.py diff --git a/spyde/fitting/components.py b/spyde/fitting/components.py index 76d68c4..1c5a53f 100644 --- a/spyde/fitting/components.py +++ b/spyde/fitting/components.py @@ -457,7 +457,7 @@ def supports(spec) -> bool: """ for c in getattr(spec, "active_components", []): try: - get_component(c.kind, n_params=len(c.parameters)) + get_component(c.kind, n_params=len(c.scalar_parameters)) except NotImplementedError: return False return True @@ -468,7 +468,7 @@ def has_analytic_grad(spec) -> bool: the engine can build the whole Jacobian without autodiff.""" for c in getattr(spec, "active_components", []): try: - if not get_component(c.kind, n_params=len(c.parameters)).has_analytic_grad: + if not get_component(c.kind, n_params=len(c.scalar_parameters)).has_analytic_grad: return False except NotImplementedError: return False @@ -497,7 +497,7 @@ def evaluate_with_grad(spec, x, values): i = 0 for c in spec.active_components: - n = len(c.parameters) + n = len(c.scalar_parameters) comp = get_component(c.kind, n_params=n) block = values[:, i:i + n] out = out + comp(x, block) @@ -517,7 +517,7 @@ def evaluate(spec, x, values): out = None i = 0 for c in spec.active_components: - n = len(c.parameters) + n = len(c.scalar_parameters) comp = get_component(c.kind, n_params=n) y = comp(x, values[:, i:i + n]) out = y if out is None else out + y diff --git a/spyde/fitting/spec.py b/spyde/fitting/spec.py index aba1721..7456384 100644 --- a/spyde/fitting/spec.py +++ b/spyde/fitting/spec.py @@ -53,6 +53,17 @@ class ParameterSpec: bmax: float | None = None linear: bool = False units: str = "" + # Some HyperSpy parameters are VECTORS — EELSCLEdge's fine_structure_coeff + # holds 8 values. Those round-trip faithfully (``value`` becomes a list) but + # are excluded from the packed parameter vector, because the batched engine + # solves one scalar per column. Nothing that uses the flat view can fit such + # a component anyway (``components.supports`` already says no), so this + # keeps storage honest without complicating the solver. + n_elements: int = 1 + + @property + def is_scalar(self) -> bool: + return self.n_elements == 1 def bounds(self) -> tuple[float, float]: """Finite bounds for the engine (``None`` becomes ±inf).""" @@ -61,16 +72,22 @@ def bounds(self) -> tuple[float, float]: return lo, hi def to_dict(self) -> dict[str, Any]: - return {"name": self.name, "value": float(self.value), + value = ([float(v) for v in np.ravel(self.value)] if self.n_elements > 1 + else float(self.value)) + return {"name": self.name, "value": value, "free": bool(self.free), "bmin": self.bmin, "bmax": self.bmax, - "linear": bool(self.linear), "units": self.units} + "linear": bool(self.linear), "units": self.units, + "n_elements": int(self.n_elements)} @classmethod def from_dict(cls, d: dict[str, Any]) -> "ParameterSpec": - return cls(name=d["name"], value=float(d.get("value", 0.0)), + raw = d.get("value", 0.0) + n = int(d.get("n_elements", 1)) + value = [float(v) for v in raw] if n > 1 else float(raw) + return cls(name=d["name"], value=value, free=bool(d.get("free", True)), bmin=d.get("bmin"), bmax=d.get("bmax"), linear=bool(d.get("linear", False)), - units=d.get("units", "")) + units=d.get("units", ""), n_elements=n) @dataclass @@ -87,11 +104,28 @@ class ComponentSpec: name: str = "" active: bool = True parameters: list[ParameterSpec] = field(default_factory=list) + # CONSTRUCTOR arguments, for components that cannot be built bare. + # Without these a spec cannot be turned back into a HyperSpy model: + # ``EELSCLEdge`` raises (element_subshell is required) and ``Polynomial`` + # silently rebuilds at its DEFAULT order, quietly dropping every + # coefficient beyond it. See ``_INIT_ARGS``. + init_args: dict[str, Any] = field(default_factory=dict) def __post_init__(self): if not self.name: self.name = self.kind + @property + def scalar_parameters(self) -> list[ParameterSpec]: + """The parameters that occupy a column in the packed vector. + + Vector-valued parameters (EELSCLEdge.fine_structure_coeff) are stored + and round-tripped but never packed — the solver puts one scalar per + column, and no component with such a parameter is fittable by the + batched engine anyway. + """ + return [p for p in self.parameters if p.is_scalar] + def __getitem__(self, name: str) -> ParameterSpec: for p in self.parameters: if p.name == name: @@ -101,14 +135,16 @@ def __getitem__(self, name: str) -> ParameterSpec: def to_dict(self) -> dict[str, Any]: return {"kind": self.kind, "name": self.name, "active": bool(self.active), - "parameters": [p.to_dict() for p in self.parameters]} + "parameters": [p.to_dict() for p in self.parameters], + "init_args": dict(self.init_args)} @classmethod def from_dict(cls, d: dict[str, Any]) -> "ComponentSpec": return cls(kind=d["kind"], name=d.get("name", ""), active=bool(d.get("active", True)), parameters=[ParameterSpec.from_dict(p) - for p in d.get("parameters", [])]) + for p in d.get("parameters", [])], + init_args=dict(d.get("init_args") or {})) @dataclass @@ -170,27 +206,27 @@ def parameter_names(self) -> list[str]: re-deriving an order of its own. """ return [f"{c.name}.{p.name}" - for c in self.active_components for p in c.parameters] + for c in self.active_components for p in c.scalar_parameters] def flat_values(self) -> np.ndarray: return np.array([p.value for c in self.active_components - for p in c.parameters], dtype=np.float64) + for p in c.scalar_parameters], dtype=np.float64) def free_mask(self) -> np.ndarray: """True where a parameter is fitted. Fixed parameters keep their value and are dropped from the Jacobian rather than being fitted and ignored.""" return np.array([p.free for c in self.active_components - for p in c.parameters], dtype=bool) + for p in c.scalar_parameters], dtype=bool) def linear_mask(self) -> np.ndarray: """True where the model is LINEAR in the parameter — the columns variable projection can solve by least squares (#53).""" return np.array([p.linear for c in self.active_components - for p in c.parameters], dtype=bool) + for p in c.scalar_parameters], dtype=bool) def bounds_arrays(self) -> tuple[np.ndarray, np.ndarray]: """``(lower, upper)`` with ``None`` expanded to ±inf.""" - pairs = [p.bounds() for c in self.active_components for p in c.parameters] + pairs = [p.bounds() for c in self.active_components for p in c.scalar_parameters] if not pairs: return np.empty(0), np.empty(0) lo, hi = zip(*pairs) @@ -200,13 +236,13 @@ def set_flat_values(self, values: Sequence[float]) -> None: """Write a packed parameter vector back onto the spec (the inverse of :meth:`flat_values`).""" values = np.asarray(values, dtype=float).ravel() - expected = sum(len(c.parameters) for c in self.active_components) + expected = sum(len(c.scalar_parameters) for c in self.active_components) if values.size != expected: raise ValueError(f"expected {expected} values for this spec, " f"got {values.size}") i = 0 for c in self.active_components: - for p in c.parameters: + for p in c.scalar_parameters: p.value = float(values[i]) i += 1 @@ -215,7 +251,7 @@ def component_slices(self) -> dict[str, slice]: what the component-area maps (#58) need to isolate one component.""" out, i = {}, 0 for c in self.active_components: - n = len(c.parameters) + n = len(c.scalar_parameters) out[c.name] = slice(i, i + n) i += n return out @@ -229,11 +265,17 @@ def from_model(cls, model) -> "ModelSpec": params = [] for p in c.parameters: bmin, bmax = getattr(p, "bmin", None), getattr(p, "bmax", None) + n_el = int(getattr(p, "_number_of_elements", 1) or 1) params.append(ParameterSpec( name=p.name, # A multidimensional parameter's `.value` is the value at # the CURRENT nav index; that is the right seed to carry. - value=float(np.ravel(p.value)[0]), + # A VECTOR parameter (EELSCLEdge.fine_structure_coeff has 8 + # elements) keeps all of them, or to_model would assign a + # scalar and hyperspy would reject the length. + value=([float(v) for v in np.ravel(p.value)] if n_el > 1 + else float(np.ravel(p.value)[0])), + n_elements=n_el, free=bool(p.free), bmin=None if bmin is None else float(bmin), bmax=None if bmax is None else float(bmax), @@ -242,7 +284,8 @@ def from_model(cls, model) -> "ModelSpec": )) comps.append(ComponentSpec( kind=getattr(c, "_id_name", type(c).__name__), - name=c.name, active=bool(c.active), parameters=params)) + name=c.name, active=bool(c.active), parameters=params, + init_args=_init_args_for(c))) mask = getattr(model, "_channel_switches", None) if mask is not None: @@ -276,7 +319,8 @@ def to_model(self, signal, *, apply_range: bool = True): continue # Bounds BEFORE value: HyperSpy clips an out-of-range assignment. par.bmin, par.bmax = pspec.bmin, pspec.bmax - par.value = float(pspec.value) + par.value = (list(np.ravel(pspec.value)) if pspec.n_elements > 1 + else float(pspec.value)) par.free = bool(pspec.free) if apply_range and self.channel_mask is not None: @@ -309,6 +353,35 @@ def from_dict(cls, d: dict[str, Any]) -> "ModelSpec": # helpers # --------------------------------------------------------------------------- +# How to recover the CONSTRUCTOR arguments of components that cannot be built +# bare. Each entry maps a component kind to a function of the live HyperSpy +# component returning the kwargs needed to rebuild it. +# +# This is not a nicety: without it `to_model` either raises (EELSCLEdge needs +# element_subshell) or — far worse — silently succeeds with the WRONG model. +# A `Polynomial` rebuilt at its default order simply has no a3..a6, so +# `to_model` skips those parameters and an order-6 EDS background becomes an +# order-2 one that still fits and still looks plausible. +_INIT_ARGS = { + "EELSCLEdge": lambda c: {"element_subshell": getattr( + c, "element_subshell", None) or c.name}, + "Polynomial": lambda c: {"order": max(0, len(c.parameters) - 1)}, + "Expression": lambda c: {"expression": getattr(c, "_expression", ""), + "name": c.name}, +} + + +def _init_args_for(comp) -> dict[str, Any]: + fn = _INIT_ARGS.get(getattr(comp, "_id_name", type(comp).__name__)) + if fn is None: + return {} + try: + return {k: v for k, v in fn(comp).items() if v is not None} + except Exception as e: # pragma: no cover + log.debug("could not capture init args for %r: %s", comp, e) + return {} + + def _make_component(cspec: ComponentSpec): """Instantiate the HyperSpy component named by ``cspec.kind``. @@ -328,7 +401,14 @@ def _make_component(cspec: ComponentSpec): for mod in mods: cls = getattr(mod, cspec.kind, None) if cls is not None: - comp = cls() + try: + comp = cls(**cspec.init_args) if cspec.init_args else cls() + except TypeError as e: + raise ValueError( + f"cannot rebuild {cspec.kind!r} from the spec: {e}. It " + f"needs constructor arguments that were not captured — add " + f"an entry to spec._INIT_ARGS so from_model() records them." + ) from e if cspec.name and cspec.name != cspec.kind: comp.name = cspec.name return comp @@ -345,6 +425,7 @@ def spec_from_component(comp) -> ComponentSpec: return ComponentSpec( kind=getattr(comp, "_id_name", type(comp).__name__), name=comp.name, active=bool(comp.active), + init_args=_init_args_for(comp), parameters=[ParameterSpec( name=p.name, value=float(np.ravel(p.value)[0]), free=bool(p.free), bmin=getattr(p, "bmin", None), bmax=getattr(p, "bmax", None), diff --git a/spyde/tests/migrated/test_model_spec_init_args.py b/spyde/tests/migrated/test_model_spec_init_args.py new file mode 100644 index 0000000..8ff8bfa --- /dev/null +++ b/spyde/tests/migrated/test_model_spec_init_args.py @@ -0,0 +1,149 @@ +"""ModelSpec must rebuild components that take CONSTRUCTOR arguments. + +Two real bugs this closes, found by trying to round-trip a real EELS model: + +* ``EELSCLEdge`` raises outright — ``element_subshell`` is required, so + ``to_model`` could not rebuild any EELS model at all. +* ``Polynomial`` is worse because it does NOT raise. Rebuilt bare it comes back + at its default order, so an order-6 EDS background silently loses a3..a6 — + ``to_model`` skips the parameters that no longer exist and returns a + different model that still fits and still looks plausible. + +The original round-trip tests missed both because they only ever went +model -> spec, never spec -> model with a component of this kind. +""" +from __future__ import annotations + +import numpy as np +import pytest + +import hyperspy.api as hs +from hyperspy.components1d import Polynomial + +from spyde.fitting import ModelSpec +from spyde.fitting.spec import ComponentSpec, ParameterSpec + + +def _signal(n=512, lo=200.0, hi=800.0): + s = hs.signals.Signal1D(np.random.default_rng(0).random((2, 2, n)) + 1.0) + s.axes_manager.signal_axes[0].offset = lo + s.axes_manager.signal_axes[0].scale = (hi - lo) / n + return s + + +class TestPolynomialOrder: + @pytest.mark.parametrize("order", [1, 2, 6]) + def test_order_survives_the_round_trip(self, order): + sig = _signal() + m = sig.create_model() + m.append(Polynomial(order=order)) + for k in range(order + 1): + getattr(m[0], f"a{k}").value = 0.5 * (k + 1) + + spec = ModelSpec.from_model(m) + assert spec[0].init_args == {"order": order} + + back = ModelSpec.from_model(spec.to_model(sig)) + assert len(back[0].parameters) == order + 1, \ + "rebuilt Polynomial lost coefficients" + np.testing.assert_allclose(back.flat_values(), spec.flat_values()) + + def test_high_order_coefficients_are_not_silently_dropped(self): + """The specific failure: a bare rebuild keeps only the default order's + coefficients, and every later one is quietly skipped.""" + sig = _signal() + m = sig.create_model() + m.append(Polynomial(order=6)) + for k in range(7): + getattr(m[0], f"a{k}").value = float(k + 1) + spec = ModelSpec.from_model(m) + rebuilt = spec.to_model(sig) + names = [p.name for p in rebuilt[0].parameters] + assert "a6" in names, names + assert float(np.ravel(rebuilt[0].a6.value)[0]) == pytest.approx(7.0) + + def test_order_survives_json(self): + sig = _signal() + m = sig.create_model() + m.append(Polynomial(order=4)) + spec = ModelSpec.from_model(m) + back = ModelSpec.from_dict(spec.to_dict()) + assert back[0].init_args == {"order": 4} + assert len(back.to_model(sig)[0].parameters) == 5 + + +class TestEelsEdge: + def test_edge_round_trips(self): + pytest.importorskip("exspy", reason="needs the eels extra") + from spyde.data import eels_si + from spyde.spectroscopy import model_for_composition + + s = eels_si(nav=(2, 2), n_channels=256) + spec, _ = model_for_composition(s, ["C", "N", "O"]) + edges = [c for c in spec if c.kind == "EELSCLEdge"] + assert edges, "no edges in the composition model" + for e in edges: + assert e.init_args.get("element_subshell"), e.name + + # The whole point: this used to raise. + rebuilt = spec.to_model(s) + kinds = [getattr(c, "_id_name", "") for c in rebuilt] + assert kinds.count("EELSCLEdge") == len(edges) + + def test_edge_keeps_its_subshell_identity(self): + pytest.importorskip("exspy", reason="needs the eels extra") + from spyde.data import eels_si + from spyde.spectroscopy import model_for_composition + + s = eels_si(nav=(2, 2), n_channels=256) + spec, _ = model_for_composition(s, ["C", "N", "O"]) + rebuilt = ModelSpec.from_model(spec.to_model(s)) + assert {c.name for c in rebuilt if c.kind == "EELSCLEdge"} == \ + {c.name for c in spec if c.kind == "EELSCLEdge"} + + def test_edge_onsets_survive(self): + pytest.importorskip("exspy", reason="needs the eels extra") + from spyde.data import eels_si + from spyde.spectroscopy import model_for_composition + + s = eels_si(nav=(2, 2), n_channels=256) + spec, _ = model_for_composition(s, ["C", "N", "O"]) + before = {c.name: c["onset_energy"].value + for c in spec if c.kind == "EELSCLEdge"} + back = ModelSpec.from_model(spec.to_model(s)) + after = {c.name: c["onset_energy"].value + for c in back if c.kind == "EELSCLEdge"} + assert before == pytest.approx(after) + + def test_edge_round_trips_through_json(self): + pytest.importorskip("exspy", reason="needs the eels extra") + import json + from spyde.data import eels_si + from spyde.spectroscopy import model_for_composition + + s = eels_si(nav=(2, 2), n_channels=256) + spec, _ = model_for_composition(s, ["C", "O"]) + revived = ModelSpec.from_dict(json.loads(json.dumps(spec.to_dict()))) + assert revived.to_model(s) # must not raise + + +class TestUnrebuildableComponent: + def test_missing_init_args_give_an_actionable_error(self): + """If a component needs constructor arguments nobody captured, say so + and say where to fix it — not a bare TypeError from deep inside + hyperspy.""" + pytest.importorskip("exspy", reason="needs the eels extra") + spec = ModelSpec(components=[ComponentSpec( + kind="EELSCLEdge", name="O_K", + parameters=[ParameterSpec("intensity", 1.0)])]) + with pytest.raises(ValueError, match="_INIT_ARGS"): + spec.to_model(_signal()) + + def test_components_without_init_args_are_unaffected(self): + sig = _signal() + m = sig.create_model() + from hyperspy.components1d import Gaussian, Offset + m.extend([Offset(), Gaussian()]) + spec = ModelSpec.from_model(m) + assert all(c.init_args == {} for c in spec) + assert len(ModelSpec.from_model(spec.to_model(sig))) == 2 From 3105585f12cdefef8d4980a24c643f3b502a2e04 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Mon, 27 Jul 2026 01:22:46 -0500 Subject: [PATCH 14/60] feat(spectroscopy): tabulated EELS edges, so EELS gets the GPU path too (#63) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An EELS core-loss edge is not a formula — EELSCLEdge looks its shape up in a GOS table, which is why it had no batched port and why an EELS model fell back to hyperspy's one-pixel-at-a-time fitting while an EDS model (gaussians) got the GPU. The way out is to notice what is actually being fitted. Across a spectrum image the edge SHAPE is the same everywhere — atomic physics, fixed by the element and the microscope. What varies pixel to pixel is HOW MUCH of the element there is and exactly WHERE its edge starts. So the edge is sampled once on the signal axis and replaced by a tabulated component fitting intensity (linear) and onset_shift, batched like everything else. The approximation, stated rather than buried: fine structure and effective angle are FROZEN at the values the table was sampled with. They are what make an edge a lookup in the first place, and refitting them would mean re-sampling the table every iteration — the per-item Python work the batched engine exists to avoid. tabulate_model() reports what it froze. Right trade for quantification; wrong one for fine-structure analysis, which is why it is an explicit call and not automatic. Two deliberate choices in the component: - Outside the table the value HOLDS at the end sample rather than extrapolating. Extrapolating a GOS tail would invent signal where the measurement has none. - The onset-shift derivative is a central difference over one channel, not the exact piecewise-linear slope. The exact one is a step function that jumps at every segment boundary, so LM chatters as the shift crosses a channel. It differs from autodiff by ~1.5% of the gradient's own scale, only within a channel of a kink, and the test states that as the property rather than asserting an equality that is not true. Tested end to end: the table reproduces hyperspy's edge to 1e-9, a fit recovers both a known intensity and a known 12 eV onset shift, and on the synthetic SI the fitted O_K intensity tracks the concentration map the data was built from. --- spyde/fitting/components.py | 120 ++++++++- spyde/fitting/spec.py | 16 +- spyde/spectroscopy/__init__.py | 5 +- spyde/spectroscopy/tabulate.py | 168 +++++++++++++ spyde/tests/migrated/test_eels_tabulate.py | 271 +++++++++++++++++++++ 5 files changed, 570 insertions(+), 10 deletions(-) create mode 100644 spyde/spectroscopy/tabulate.py create mode 100644 spyde/tests/migrated/test_eels_tabulate.py diff --git a/spyde/fitting/components.py b/spyde/fitting/components.py index 1c5a53f..777aabc 100644 --- a/spyde/fitting/components.py +++ b/spyde/fitting/components.py @@ -31,6 +31,8 @@ import math from typing import Callable, Sequence +import numpy as np + log = logging.getLogger(__name__) _SQRT_2PI = math.sqrt(2.0 * math.pi) @@ -364,6 +366,61 @@ def _d_logistic(x, p): -a * E * c / D ** 2) # df/dorigin +def _interp_uniform(xq, x0, dx, table): + """Linear interpolation of ``table`` (sampled uniformly from ``x0``) at + ``xq``, differentiable with respect to ``xq``. + + ``torch.searchsorted`` is not needed because the table is on the signal + axis, which is uniform — index arithmetic is exact and far cheaper. The + index itself is a floor and carries no gradient, which is correct: the + function is piecewise linear, so the gradient comes entirely from the + interpolation weight. + + Outside the table the value is held at the end sample rather than + extrapolated. Extrapolating a GOS tail would invent signal where the + measurement has none. + """ + import torch + + n = table.shape[-1] + pos = (xq - x0) / dx + idx = torch.clamp(pos.floor(), 0, n - 2) + t = torch.clamp(pos - idx, 0.0, 1.0) + i0 = idx.long() + y0 = table[i0] + y1 = table[torch.clamp(i0 + 1, max=n - 1)] + return y0 + t * (y1 - y0) + + +def _tabulated_fn(table, x0, dx): + def fn(x, p): + intensity, onset_shift = p + # Shifting the SAMPLE point left is the same as moving the edge right, + # so a positive onset_shift moves the edge up in energy as a user + # expects. + return intensity * _interp_uniform(x - onset_shift, x0, dx, table) + + return fn + + +def _tabulated_grad(table, x0, dx): + def grad(x, p): + intensity, onset_shift = p + shape = _interp_uniform(x - onset_shift, x0, dx, table) + # d/d(shift) is minus the local slope. Taken as a CENTRAL difference + # over one sample rather than the exact piecewise-linear derivative, + # deliberately: the exact one is a step function that jumps at every + # segment boundary, so LM chatters as the shift crosses a channel. The + # smoothed version differs from autodiff only within one channel of a + # kink and converges to the same place. + half = dx / 2 + slope = (_interp_uniform(x - onset_shift + half, x0, dx, table) + - _interp_uniform(x - onset_shift - half, x0, dx, table)) / dx + return _stack(shape, -intensity * slope) + + return grad + + def _d_polynomial_fn(order: int): def grad(x, p): return _stack(*[x ** k + 0.0 * p[0] for k in range(order + 1)]) @@ -443,9 +500,58 @@ def get_component(kind: str, *, n_params: int | None = None) -> TorchComponent: ) from None +TABULATED_KIND = "TabulatedShape" + + +def tabulated_component(table, x0: float, dx: float, *, device=None, + dtype=None) -> TorchComponent: + """A component whose SHAPE comes from a lookup table (#63). + + Fits two parameters against it: ``intensity`` (linear — the amplitude that + scales the whole shape) and ``onset_shift`` (how far the shape slides along + the energy axis). That is the right pair for an EELS core-loss edge, whose + form is a GOS table rather than a formula: the shape is fixed physics and + the fit is asking "how much of this element, and exactly where does its + edge start". + + What this deliberately does NOT fit is the fine structure and effective + angle, which are FROZEN at the values the table was sampled with. Those are + what make an edge a tabulated lookup in the first place, and refitting them + would mean re-sampling the table every iteration — which is exactly the + per-item Python work the batched engine exists to avoid. + """ + import torch + + t = torch.as_tensor(np.asarray(table, float), device=device, + dtype=dtype or torch.float64) + return TorchComponent(TABULATED_KIND, ("intensity", "onset_shift"), + (True, False), + _tabulated_fn(t, float(x0), float(dx)), + _tabulated_grad(t, float(x0), float(dx))) + + +def component_for(cspec, *, device=None, dtype=None) -> TorchComponent: + """Resolve the batched component for a :class:`ComponentSpec`. + + Everything analytic resolves by ``kind`` alone; a tabulated component also + needs its own table, which lives on the spec. Callers that have a spec + should use this rather than :func:`get_component`, so a data-bound + component is never silently looked up as if it were stateless. + """ + if cspec.kind == TABULATED_KIND: + if cspec.data is None: + raise ValueError(f"{cspec.name}: {TABULATED_KIND} has no table — " + f"build it with spyde.spectroscopy.tabulate") + x0 = float(cspec.init_args.get("x0", 0.0)) + dx = float(cspec.init_args.get("dx", 1.0)) + return tabulated_component(cspec.data, x0, dx, device=device, + dtype=dtype) + return get_component(cspec.kind, n_params=len(cspec.scalar_parameters)) + + def available() -> list[str]: """Component kinds the batched engine can fit.""" - return sorted(_REGISTRY) + ["Polynomial"] + return sorted(_REGISTRY) + ["Polynomial", TABULATED_KIND] def supports(spec) -> bool: @@ -457,8 +563,8 @@ def supports(spec) -> bool: """ for c in getattr(spec, "active_components", []): try: - get_component(c.kind, n_params=len(c.scalar_parameters)) - except NotImplementedError: + component_for(c) + except (NotImplementedError, ValueError): return False return True @@ -468,9 +574,9 @@ def has_analytic_grad(spec) -> bool: the engine can build the whole Jacobian without autodiff.""" for c in getattr(spec, "active_components", []): try: - if not get_component(c.kind, n_params=len(c.scalar_parameters)).has_analytic_grad: + if not component_for(c).has_analytic_grad: return False - except NotImplementedError: + except (NotImplementedError, ValueError): return False return True @@ -498,7 +604,7 @@ def evaluate_with_grad(spec, x, values): i = 0 for c in spec.active_components: n = len(c.scalar_parameters) - comp = get_component(c.kind, n_params=n) + comp = component_for(c, device=values.device, dtype=values.dtype) block = values[:, i:i + n] out = out + comp(x, block) jac[:, :, i:i + n] = comp.grad(x, block).expand(P, C, n) @@ -518,7 +624,7 @@ def evaluate(spec, x, values): i = 0 for c in spec.active_components: n = len(c.scalar_parameters) - comp = get_component(c.kind, n_params=n) + comp = component_for(c, device=values.device, dtype=values.dtype) y = comp(x, values[:, i:i + n]) out = y if out is None else out + y i += n diff --git a/spyde/fitting/spec.py b/spyde/fitting/spec.py index 7456384..316c5ab 100644 --- a/spyde/fitting/spec.py +++ b/spyde/fitting/spec.py @@ -110,6 +110,12 @@ class ComponentSpec: # silently rebuilds at its DEFAULT order, quietly dropping every # coefficient beyond it. See ``_INIT_ARGS``. init_args: dict[str, Any] = field(default_factory=dict) + # Tabulated SHAPE for components whose form is a lookup rather than a + # formula — an EELS core-loss edge is a GOS table, not an expression + # (#63). Sampled on the signal axis; the component then fits an intensity + # (linear) and an onset shift against it. None for every analytic + # component, which is nearly all of them. + data: np.ndarray | None = None def __post_init__(self): if not self.name: @@ -136,15 +142,21 @@ def __getitem__(self, name: str) -> ParameterSpec: def to_dict(self) -> dict[str, Any]: return {"kind": self.kind, "name": self.name, "active": bool(self.active), "parameters": [p.to_dict() for p in self.parameters], - "init_args": dict(self.init_args)} + "init_args": dict(self.init_args), + # A tabulated shape is one float per channel, so it does go + # through JSON — but only for the components that have one. + "data": (None if self.data is None + else [float(v) for v in np.ravel(self.data)])} @classmethod def from_dict(cls, d: dict[str, Any]) -> "ComponentSpec": + raw = d.get("data") return cls(kind=d["kind"], name=d.get("name", ""), active=bool(d.get("active", True)), parameters=[ParameterSpec.from_dict(p) for p in d.get("parameters", [])], - init_args=dict(d.get("init_args") or {})) + init_args=dict(d.get("init_args") or {}), + data=None if raw is None else np.asarray(raw, float)) @dataclass diff --git a/spyde/spectroscopy/__init__.py b/spyde/spectroscopy/__init__.py index 69e9dd1..7502b35 100644 --- a/spyde/spectroscopy/__init__.py +++ b/spyde/spectroscopy/__init__.py @@ -23,4 +23,7 @@ prune_to_range, ) -__all__ = ["model_for_composition", "prune_to_range", "MissingExtra"] +from spyde.spectroscopy.tabulate import onset_energies, tabulate_model + +__all__ = ["model_for_composition", "prune_to_range", "MissingExtra", + "tabulate_model", "onset_energies"] diff --git a/spyde/spectroscopy/tabulate.py b/spyde/spectroscopy/tabulate.py new file mode 100644 index 0000000..e7a1799 --- /dev/null +++ b/spyde/spectroscopy/tabulate.py @@ -0,0 +1,168 @@ +"""tabulate.py — make GOS-based EELS edges fittable by the batched engine (#63). + +An EELS core-loss edge is not a formula. ``EELSCLEdge`` looks its shape up in a +generalised-oscillator-strength table, which is why it has no batched torch +port and why an EELS model falls back to HyperSpy's one-pixel-at-a-time fitting +while an EDS model (gaussians) gets the GPU. + +The way out is to notice what is actually being fitted. Across a spectrum image +the edge SHAPE is the same everywhere — it is atomic physics, fixed by the +element and the microscope. What varies pixel to pixel is **how much** of the +element there is and **exactly where** its edge starts. So: + +1. sample each edge once, on the signal axis, at its current parameters; +2. replace it with a tabulated component carrying that shape; +3. fit ``intensity`` (linear) and ``onset_shift`` against it, batched. + +**The approximation, stated plainly.** Fine structure and effective angle are +frozen at the values the table was sampled with. They are what make an edge a +lookup in the first place, and refitting them would mean re-sampling the table +every iteration — the per-item Python work the batched engine exists to avoid. +For quantification, which asks "how much of each element", that is the right +trade. For fine-structure analysis it is not, and +:func:`tabulate_model` says so by leaving the original model alone unless asked. +""" +from __future__ import annotations + +import logging + +import numpy as np + +from spyde.fitting import ModelSpec +from spyde.fitting.components import TABULATED_KIND +from spyde.fitting.spec import ComponentSpec, ParameterSpec + +log = logging.getLogger(__name__) + +# Components worth tabulating: form is a lookup, amplitude is linear. +TABULATABLE = ("EELSCLEdge",) + + +def _axis(signal): + ax = signal.axes_manager.signal_axes[0] + x = np.asarray(ax.axis, float) + if len(x) < 2: + raise ValueError("signal axis needs at least two channels to tabulate") + dx = float(x[1] - x[0]) + if not np.allclose(np.diff(x), dx, rtol=1e-6): + raise ValueError("tabulation needs a UNIFORM signal axis; this one is " + "not evenly spaced") + return x, float(x[0]), dx + + +def tabulate_model(spec: ModelSpec, signal, *, kinds=TABULATABLE, + intensity_name: str = "intensity"): + """Replace lookup-shaped components with tabulated ones. + + Returns ``(new_spec, info)``. *info* lists what was tabulated and what was + left alone, so a caller can tell the user which parameters are no longer + being fitted. + + The returned spec is fittable by :mod:`spyde.fitting.engine` whenever every + remaining component has a batched port — check with + ``components.supports``. + """ + x, x0, dx = _axis(signal) + model = spec.to_model(signal) + + by_name = {c.name: c for c in model} + out, tabulated, skipped = [], [], [] + + for cspec in spec.components: + if cspec.kind not in kinds: + out.append(cspec) + continue + + comp = by_name.get(cspec.name) + if comp is None: + skipped.append(cspec.name) + out.append(cspec) + continue + + try: + table, onset = _sample(comp, x, intensity_name) + except Exception as e: + log.info("could not tabulate %s (%s); left as-is", cspec.name, e) + skipped.append(cspec.name) + out.append(cspec) + continue + + intensity = _value(cspec, intensity_name, default=1.0) + out.append(ComponentSpec( + kind=TABULATED_KIND, name=cspec.name, active=cspec.active, + # The table is sampled ON the signal axis, so the component's own + # x0/dx are the axis's — recorded here because the component needs + # them to interpolate and they are not derivable from the table. + init_args={"x0": x0, "dx": dx, "source_kind": cspec.kind, + "onset_energy": onset}, + data=table, + parameters=[ + ParameterSpec(intensity_name, float(intensity), linear=True, + bmin=0.0), + # Bounded to a few channels: the onset is known from the + # element, and a shift larger than this is not a refinement but + # the fit sliding onto a neighbouring edge. + ParameterSpec("onset_shift", 0.0, + bmin=-20.0 * dx, bmax=20.0 * dx), + ])) + tabulated.append(cspec.name) + + info = {"tabulated": tabulated, "skipped": skipped, + "frozen": ["fine_structure_coeff", "effective_angle"], + "x0": x0, "dx": dx} + if tabulated: + log.info("tabulated %d edge(s): %s — fine structure and effective " + "angle are now FROZEN", len(tabulated), ", ".join(tabulated)) + return ModelSpec(components=out, channel_mask=spec.channel_mask), info + + +def _sample(comp, x, intensity_name): + """Sample one component's shape at UNIT intensity. + + Unit intensity is what makes the table a pure shape, so the fitted + ``intensity`` means the same thing it did in HyperSpy rather than being + scaled by whatever the seed happened to be. + """ + par = getattr(comp, intensity_name, None) + original = None if par is None else float(np.ravel(par.value)[0]) + try: + if par is not None: + par.value = 1.0 + y = np.nan_to_num(np.asarray(comp.function(x), float), + nan=0.0, posinf=0.0, neginf=0.0) + finally: + if par is not None and original is not None: + par.value = original + + if not np.isfinite(y).all() or not np.any(y): + raise ValueError("component sampled to an empty or non-finite shape") + onset = getattr(getattr(comp, "onset_energy", None), "value", None) + return y, (None if onset is None else float(np.ravel(onset)[0])) + + +def _value(cspec, name, default=1.0): + try: + return float(cspec[name].value) + except (KeyError, TypeError): + return default + + +def onset_energies(spec: ModelSpec) -> dict[str, float]: + """Absolute fitted onset per tabulated component. + + ``onset_shift`` is relative to where the table was sampled, which is not + what a user wants to read — this adds the original onset back so the number + is an energy in eV. + """ + out = {} + for c in spec.components: + if c.kind != TABULATED_KIND: + continue + base = c.init_args.get("onset_energy") + if base is None: + continue + try: + out[c.name] = float(base) + float(c["onset_shift"].value) + except KeyError: + continue + return out diff --git a/spyde/tests/migrated/test_eels_tabulate.py b/spyde/tests/migrated/test_eels_tabulate.py new file mode 100644 index 0000000..d6ce924 --- /dev/null +++ b/spyde/tests/migrated/test_eels_tabulate.py @@ -0,0 +1,271 @@ +"""Tabulated EELS edges through the batched engine (#63). + +The claim under test: an EELS model that could only be fitted by HyperSpy, one +pixel at a time, becomes fittable by the batched engine — and the answer still +means what it meant before. + +So the tests check both halves. That the tabulated shape REPRODUCES the +hyperspy edge (otherwise the speed is worthless), and that fitting it recovers +the intensities the data was built from. +""" +from __future__ import annotations + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("exspy", reason="needs the eels extra") + +from spyde.data import eels_si +from spyde.fitting import ModelSpec +from spyde.fitting import components as tc +from spyde.fitting.components import TABULATED_KIND, tabulated_component +from spyde.fitting.engine import fit_batched +from spyde.fitting.spec import ComponentSpec, ParameterSpec +from spyde.spectroscopy import model_for_composition, onset_energies, tabulate_model + + +@pytest.fixture(scope="module") +def eels(): + return eels_si(nav=(3, 3), n_channels=512) + + +class TestTabulatedComponent: + def test_reproduces_the_table_at_unit_intensity(self): + x = np.linspace(0.0, 10.0, 101) + table = np.sin(x) ** 2 + comp = tabulated_component(table, x[0], x[1] - x[0]) + got = comp(torch.as_tensor(x), torch.tensor([[1.0, 0.0]], + dtype=torch.float64)) + np.testing.assert_allclose(got.numpy()[0], table, atol=1e-12) + + def test_intensity_scales_linearly(self): + x = np.linspace(0.0, 10.0, 101) + comp = tabulated_component(np.sin(x) ** 2, x[0], x[1] - x[0]) + xt = torch.as_tensor(x) + one = comp(xt, torch.tensor([[1.0, 0.0]], dtype=torch.float64)) + three = comp(xt, torch.tensor([[3.0, 0.0]], dtype=torch.float64)) + torch.testing.assert_close(three, 3 * one) + + def test_onset_shift_moves_the_shape_up_in_energy(self): + """A POSITIVE shift must move the edge to higher energy — the sign a + user expects. Getting this backwards still fits, just with the shift + inverted, and nothing else would catch it.""" + x = np.linspace(0.0, 100.0, 501) + table = (x > 50.0).astype(float) # a step at 50 + comp = tabulated_component(table, x[0], x[1] - x[0]) + y = comp(torch.as_tensor(x), + torch.tensor([[1.0, 10.0]], dtype=torch.float64)).numpy()[0] + assert x[int(np.argmax(y > 0.5))] == pytest.approx(60.0, abs=0.5) + + def test_holds_the_end_value_outside_the_table(self): + """Extrapolating a GOS tail would invent signal where the measurement + has none.""" + x = np.linspace(0.0, 10.0, 51) + comp = tabulated_component(np.ones_like(x) * 2.0, x[0], x[1] - x[0]) + far = comp(torch.as_tensor(np.array([-50.0, 60.0])), + torch.tensor([[1.0, 0.0]], dtype=torch.float64)) + np.testing.assert_allclose(far.numpy()[0], [2.0, 2.0]) + + def test_analytic_gradient_matches_autodiff_on_a_smooth_table(self): + """Checked on a SMOOTH table, where the two definitions coincide. + + The shift derivative here is a deliberate central difference, not the + exact piecewise-linear slope: the exact one is a step function that + jumps at every segment boundary, so LM chatters as the shift crosses a + channel. On a table with a KINK the two therefore differ within one + channel of it — by design — so a kinked table would be testing the + artefact rather than the machinery. + """ + from torch.func import jacfwd + x = np.linspace(200.0, 800.0, 301) + table = np.exp(-0.5 * ((x - 500.0) / 40.0) ** 2) + comp = tabulated_component(table, x[0], x[1] - x[0]) + xt = torch.as_tensor(x) + v = torch.tensor([[2.0, 1.5]], dtype=torch.float64) + auto = jacfwd(lambda p: comp(xt, p.unsqueeze(0)).squeeze(0))(v[0]) + got = comp.grad(xt, v)[0] + + # The INTENSITY column is exact — it is just the table. + torch.testing.assert_close(got[:, 0], auto[:, 0], rtol=1e-9, atol=1e-9) + + # The SHIFT column is smoothed over one channel, so it is compared as a + # small perturbation of the autodiff gradient rather than as an + # equality. The largest disagreement sits at the peak, where the left + # and right segment slopes cancel in the central difference but + # autodiff returns one of them — a real property of the smoothing, not + # an error. + # Measured at ~1.5% of the gradient's own scale on this table; the + # bound is the O(dx * f'') error a one-channel central difference must + # have, not a number tuned until it passed. + scale = auto[:, 1].abs().max() + assert (got[:, 1] - auto[:, 1]).abs().max() < 0.05 * scale + # And it must still point the same way wherever the slope is real. + steep = auto[:, 1].abs() > 0.1 * scale + assert (torch.sign(got[steep, 1]) == torch.sign(auto[steep, 1])).all() + + def test_shift_gradient_has_the_right_sign_on_a_kinked_table(self): + """A kinked (edge-like) table is the real case, so the gradient still + has to point the right way there even though it is smoothed.""" + x = np.linspace(200.0, 800.0, 301) + table = np.clip(x - 400.0, 0, None) ** 0.5 + comp = tabulated_component(table, x[0], x[1] - x[0]) + g = comp.grad(torch.as_tensor(x), + torch.tensor([[2.0, 0.0]], dtype=torch.float64))[0] + above = x > 420.0 + # Shifting the edge UP in energy reduces intensity above the onset, + # where the table is rising. + assert (g[above, 1] < 0).all() + + def test_missing_table_is_an_actionable_error(self): + spec = ComponentSpec(kind=TABULATED_KIND, name="O_K", + parameters=[ParameterSpec("intensity", 1.0)]) + with pytest.raises(ValueError, match="tabulate"): + tc.component_for(spec) + + +class TestTabulateModel: + def test_edges_become_tabulated_components(self, eels): + spec, info = model_for_composition(eels, ["C", "N", "O"]) + assert info["engine_supported"] is False # before + + tab, tinfo = tabulate_model(spec, eels) + assert set(tinfo["tabulated"]) == {"C_K", "N_K", "O_K"} + assert not tinfo["skipped"] + assert [c.kind for c in tab].count(TABULATED_KIND) == 3 + + def test_the_batched_engine_can_now_fit_it(self, eels): + """The headline: an EELS model that HyperSpy alone could fit now goes + through the batched engine.""" + spec, _ = model_for_composition(eels, ["C", "N", "O"]) + tab, _ = tabulate_model(spec, eels) + assert tc.supports(tab) is True + assert tc.has_analytic_grad(tab) is True + + def test_the_table_reproduces_the_hyperspy_edge(self, eels): + """If the tabulated shape is not the edge, the speed is worthless.""" + spec, _ = model_for_composition(eels, ["O"]) + tab, tinfo = tabulate_model(spec, eels) + x = np.asarray(eels.axes_manager.signal_axes[0].axis, float) + + model = spec.to_model(eels) + edge = [c for c in model if getattr(c, "_id_name", "") == "EELSCLEdge"][0] + edge.intensity.value = 1.0 + want = np.nan_to_num(np.asarray(edge.function(x), float)) + + comp = [c for c in tab if c.kind == TABULATED_KIND][0] + got = tc.component_for(comp)( + torch.as_tensor(x), torch.tensor([[1.0, 0.0]], dtype=torch.float64)) + np.testing.assert_allclose(got.numpy()[0], want, rtol=1e-9, atol=1e-12) + + def test_intensity_is_marked_linear(self, eels): + """Variable projection and the seeding path both key off this.""" + spec, _ = model_for_composition(eels, ["O"]) + tab, _ = tabulate_model(spec, eels) + comp = [c for c in tab if c.kind == TABULATED_KIND][0] + assert comp["intensity"].linear is True + assert comp["onset_shift"].linear is False + + def test_non_edge_components_are_left_alone(self, eels): + spec, _ = model_for_composition(eels, ["C", "N", "O"]) + tab, _ = tabulate_model(spec, eels) + assert any(c.kind == "PowerLaw" for c in tab), \ + "the background was tabulated or lost" + + def test_reports_what_it_froze(self, eels): + """The approximation must be stated, not buried — fine structure and + effective angle stop being fitted.""" + spec, _ = model_for_composition(eels, ["O"]) + _, info = tabulate_model(spec, eels) + assert "fine_structure_coeff" in info["frozen"] + assert "effective_angle" in info["frozen"] + + def test_a_non_uniform_axis_is_rejected(self, eels): + """Tabulation interpolates by index arithmetic, which silently gives + wrong energies on a non-uniform axis.""" + s = eels.deepcopy() + spec, _ = model_for_composition(s, ["O"]) + ax = s.axes_manager.signal_axes[0] + import hyperspy.axes as hax + s.axes_manager.remove(ax) + s.axes_manager._axes.append(hax.DataAxis(axis=np.geomspace(200, 800, 512))) + with pytest.raises(ValueError, match="UNIFORM"): + tabulate_model(spec, s) + + +class TestFittingATabulatedModel: + def test_recovers_the_intensity_it_was_given(self): + """End to end on data built from the tabulated shape itself, so the + only thing under test is the FIT.""" + x = np.linspace(200.0, 800.0, 512) + table = np.clip(x - 400.0, 0, None) ** 0.5 + truth = np.array([2.0, 5.0, 9.0]) + data = np.stack([a * table for a in truth]) + + spec = ModelSpec(components=[ComponentSpec( + kind=TABULATED_KIND, name="edge", + init_args={"x0": float(x[0]), "dx": float(x[1] - x[0])}, + data=table, + parameters=[ParameterSpec("intensity", 1.0, linear=True, bmin=0.0), + ParameterSpec("onset_shift", 0.0, bmin=-20.0, bmax=20.0)])]) + got = fit_batched(spec, data, x, device="cpu", max_iter=80) + col = spec.parameter_names().index("edge.intensity") + np.testing.assert_allclose(got.values[:, col], truth, rtol=1e-4) + + def test_recovers_an_onset_shift(self): + x = np.linspace(200.0, 800.0, 512) + table = np.clip(x - 400.0, 0, None) ** 0.5 + shifted = np.clip(x - 412.0, 0, None) ** 0.5 # +12 eV + + spec = ModelSpec(components=[ComponentSpec( + kind=TABULATED_KIND, name="edge", + init_args={"x0": float(x[0]), "dx": float(x[1] - x[0])}, + data=table, + parameters=[ParameterSpec("intensity", 1.0, linear=True, bmin=0.0), + ParameterSpec("onset_shift", 0.0, bmin=-30.0, bmax=30.0)])]) + got = fit_batched(spec, shifted[None, :], x, device="cpu", max_iter=150) + col = spec.parameter_names().index("edge.onset_shift") + assert got.values[0, col] == pytest.approx(12.0, abs=1.5) + + def test_fits_a_real_composition_model_on_real_synthetic_data(self, eels): + """The whole path: composition -> model -> tabulate -> batched fit.""" + spec, _ = model_for_composition(eels, ["C", "N", "O"]) + tab, _ = tabulate_model(spec, eels) + x = np.asarray(eels.axes_manager.signal_axes[0].axis, float) + got = fit_batched(tab, eels.data, x, device="cpu", max_iter=60) + assert got.values.shape[0] == 9 + assert np.isfinite(got.values).all() + + def test_edge_intensity_tracks_the_known_concentration(self, eels): + """The measurement that matters: more of an element must fit a bigger + edge intensity, scored against the map the data was built from.""" + from spyde.data import ground_truth + + spec, _ = model_for_composition(eels, ["C", "N", "O"]) + tab, _ = tabulate_model(spec, eels) + x = np.asarray(eels.axes_manager.signal_axes[0].axis, float) + got = fit_batched(tab, eels.data, x, device="cpu", max_iter=80) + + names = tab.parameter_names() + conc = np.asarray(ground_truth(eels)["concentration"]["O_K"]).ravel() + fitted = got.values[:, names.index("O_K.intensity")] + assert np.corrcoef(fitted, conc)[0, 1] > 0.7, \ + f"fitted O_K intensity does not track its concentration map" + + +class TestOnsetEnergies: + def test_reports_an_absolute_energy(self, eels): + """onset_shift alone is relative to where the table was sampled, which + is not a number a user can read.""" + spec, _ = model_for_composition(eels, ["O"]) + tab, _ = tabulate_model(spec, eels) + comp = [c for c in tab if c.kind == TABULATED_KIND][0] + comp["onset_shift"].value = 5.0 + got = onset_energies(tab) + assert got["O_K"] == pytest.approx( + comp.init_args["onset_energy"] + 5.0) + + def test_ignores_untabulated_components(self, eels): + spec, _ = model_for_composition(eels, ["O"]) + tab, _ = tabulate_model(spec, eels) + assert set(onset_energies(tab)) == {"O_K"} From 7202f3aeaa229df68993bffa15ab3b4691bb010e Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Mon, 27 Jul 2026 01:52:44 -0500 Subject: [PATCH 15/60] feat(spectroscopy): quantification -> per-element composition maps (#66) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two steps, kept separate because they fail differently. element_intensity_maps() is the BRIDGE: fitted parameters -> one intensity map per element, reading the component names the model was built with (Fe_Ka, O_K) and summing a family's lines. Pure array work, no physics, no optional extra. One bridge serves EDS and EELS because the naming shape is the same and the intensity parameter is looked up per component kind (Gaussian's A, TabulatedShape's intensity) — otherwise the EELS path silently produces no maps. quantify() is the PHYSICS: relative, Cliff-Lorimer, or EELS cross-sections. The split matters because the bridge is where a naming or ordering mistake sends iron's intensity into copper's map, and that is SILENT — every downstream number stays plausible. A wrong k-factor is at least a number a microscopist can sanity-check. They deserve separate tests, and the bridge gets more of them. Three things that would otherwise be quiet: - Negative fitted intensities are clamped at SOURCE. A fit can drive a line slightly negative on noise, and a negative "amount of an element" is not just unphysical for that element — it corrupts the normalisation for every other one. - Missing k-factors default to 1.0 and are REPORTED, because defaulting silently makes an uncalibrated number look calibrated. - The result records that composition is RELATIVE to the fitted elements. Cliff-Lorimer cannot know about an element you did not fit, so leaving one out does not cause a small error — it redistributes that element's share across everything else. End to end on the synthetic EDS SI: composition -> model -> batched fit -> quantify recovers the generator's concentration maps at r > 0.9 per element. "Quantification works" means it recovers what we put in, not that the numbers sum to 1. --- spyde/spectroscopy/__init__.py | 8 +- spyde/spectroscopy/quantify.py | 185 ++++++++++++++++++++++ spyde/tests/migrated/test_quantify.py | 217 ++++++++++++++++++++++++++ 3 files changed, 409 insertions(+), 1 deletion(-) create mode 100644 spyde/spectroscopy/quantify.py create mode 100644 spyde/tests/migrated/test_quantify.py diff --git a/spyde/spectroscopy/__init__.py b/spyde/spectroscopy/__init__.py index 7502b35..9cab235 100644 --- a/spyde/spectroscopy/__init__.py +++ b/spyde/spectroscopy/__init__.py @@ -23,7 +23,13 @@ prune_to_range, ) +from spyde.spectroscopy.quantify import ( + element_intensity_maps, + quantify, + quantify_result, +) from spyde.spectroscopy.tabulate import onset_energies, tabulate_model __all__ = ["model_for_composition", "prune_to_range", "MissingExtra", - "tabulate_model", "onset_energies"] + "tabulate_model", "onset_energies", + "element_intensity_maps", "quantify", "quantify_result"] diff --git a/spyde/spectroscopy/quantify.py b/spyde/spectroscopy/quantify.py new file mode 100644 index 0000000..31a6012 --- /dev/null +++ b/spyde/spectroscopy/quantify.py @@ -0,0 +1,185 @@ +"""quantify.py — fitted intensities -> per-element composition maps (#66). + +Two steps, kept separate because they fail differently: + +:func:`element_intensity_maps` + The BRIDGE. Turns a :class:`~spyde.fitting.engine.FitResult` into one + intensity map per element, by reading the component names the model was + built with (``Fe_Ka``, ``O_K``) and summing a family's lines. Pure array + work, no physics — and no optional extra. + +:func:`quantify` + The PHYSICS. Cliff-Lorimer for EDS, relative cross-sections for EELS. + +Why the split matters: the bridge is where a naming or ordering mistake sends +iron's intensity into copper's map, and that is silent — every downstream number +stays plausible. The physics is where a wrong k-factor lives, and that is at +least a number a microscopist can sanity-check. They deserve separate tests. + +**Composition is normalised and therefore relative.** Cliff-Lorimer gives +ratios; it cannot know about an element you did not fit, so the fractions sum +to 1 over the elements PRESENT IN THE MODEL. Leaving an element out does not +produce a small error, it redistributes its share across everything else. +:func:`quantify` says so in its result rather than letting the number look +absolute. +""" +from __future__ import annotations + +import logging +import re + +import numpy as np + +log = logging.getLogger(__name__) + +# "Fe_Ka" / "Cu_Kb1" / "O_K" -> element, line. The EELS edge naming ("O_K") +# and the EDS line naming ("Fe_Ka") share this shape, which is why one bridge +# serves both. +_LINE_RE = re.compile(r"^([A-Z][a-z]?)_([A-Za-z0-9]+)$") + +# The parameter carrying "how much" for each component kind. +_INTENSITY_PARAM = { + "Gaussian": "A", + "GaussianHF": "height", + "Lorentzian": "A", + "TabulatedShape": "intensity", + "EELSCLEdge": "intensity", +} + + +def parse_line(name: str) -> tuple[str, str] | None: + """``"Fe_Ka"`` -> ``("Fe", "Ka")``; ``None`` for anything else.""" + m = _LINE_RE.match(name or "") + return (m.group(1), m.group(2)) if m else None + + +def element_intensity_maps(spec, result, nav_shape=None, *, + lines: dict[str, str] | None = None): + """One intensity map per element, summed over that element's lines. + + Parameters + ---------- + spec : ModelSpec + The fitted model — its component names carry the element identity. + result : FitResult + From :func:`~spyde.fitting.engine.fit_batched`. + lines : dict, optional + Restrict to one line per element, e.g. ``{"Fe": "Ka"}``. Summing a + whole K family is usually right, but if a family member overlaps + another element badly, using the clean line alone is more accurate than + a contaminated sum. + + Returns + ------- + dict of element -> map + Shaped to *nav_shape* when given, otherwise flat. + """ + names = spec.parameter_names() + out: dict[str, np.ndarray] = {} + + for comp in spec.active_components: + parsed = parse_line(comp.name) + if parsed is None: + continue # background, etc. + element, line = parsed + if lines and lines.get(element) not in (None, line): + continue + + pname = _INTENSITY_PARAM.get(comp.kind) + if pname is None: + log.debug("no intensity parameter known for %s (%s) — skipped", + comp.name, comp.kind) + continue + key = f"{comp.name}.{pname}" + if key not in names: + log.debug("%s not in the fitted parameters — skipped", key) + continue + + col = result.values[:, names.index(key)] + # A fit can drive a line slightly negative on noise; a negative + # "amount of an element" is not physical and would corrupt the + # normalisation for every OTHER element, so clamp at the source. + out[element] = out.get(element, 0.0) + np.clip(col, 0.0, None) + + if nav_shape is not None: + out = {k: v.reshape(nav_shape) for k, v in out.items()} + return out + + +def quantify(intensity_maps, *, method: str = "relative", + kfactors: dict[str, float] | None = None, + cross_sections: dict[str, float] | None = None): + """Intensity maps -> atomic-fraction maps. + + Parameters + ---------- + method : {"relative", "cliff_lorimer", "eels"} + ``relative`` + Normalise raw intensities. Honest only when the elements have + similar sensitivity — offered because it needs no factors and is + the right first look. + ``cliff_lorimer`` + EDS: ``C_A/C_B = k_AB * I_A/I_B``. *kfactors* is per element, + relative to a common reference; missing ones default to 1.0 and + are reported. + ``eels`` + Divide by partial cross-sections, then normalise. + + Returns + ------- + (fractions, info) + *fractions* maps element -> atomic fraction in [0, 1] summing to 1 + across the elements present. *info* records the method, which factors + were defaulted, and that the result is RELATIVE to the fitted elements. + """ + if not intensity_maps: + raise ValueError("no intensity maps to quantify") + + elements = sorted(intensity_maps) + stack = np.stack([np.asarray(intensity_maps[e], float) for e in elements]) + defaulted: list[str] = [] + + def _factors(table) -> np.ndarray: + """Per-element factor, broadcast over whatever map shape we have.""" + vals = [] + for e in elements: + v = (table or {}).get(e) + if v is None: + defaulted.append(e) + v = 1.0 + vals.append(float(v)) + return np.asarray(vals).reshape((-1,) + (1,) * (stack.ndim - 1)) + + if method == "relative": + weighted = stack + elif method == "cliff_lorimer": + weighted = stack * _factors(kfactors) + elif method == "eels": + weighted = stack / _factors(cross_sections) + else: + raise ValueError(f"unknown quantification method {method!r} " + f"(relative, cliff_lorimer or eels)") + + total = weighted.sum(0) + with np.errstate(invalid="ignore", divide="ignore"): + frac = np.where(total > 0, weighted / total, np.nan) + + if defaulted: + log.info("quantification: no factor for %s — defaulted to 1.0, so " + "those fractions are uncalibrated", ", ".join(defaulted)) + + return ({e: frac[i] for i, e in enumerate(elements)}, + {"method": method, "elements": elements, + "defaulted_factors": defaulted, + "relative_to": elements, + "note": "fractions are normalised over the FITTED elements only; " + "an element left out of the model has its share " + "redistributed across the rest"}) + + +def quantify_result(spec, result, nav_shape=None, **kwargs): + """:func:`element_intensity_maps` then :func:`quantify`, in one call.""" + maps = element_intensity_maps(spec, result, nav_shape) + fractions, info = quantify(maps, **kwargs) + info["intensity_maps"] = maps + return fractions, info diff --git a/spyde/tests/migrated/test_quantify.py b/spyde/tests/migrated/test_quantify.py new file mode 100644 index 0000000..29eae59 --- /dev/null +++ b/spyde/tests/migrated/test_quantify.py @@ -0,0 +1,217 @@ +"""Fitted intensities -> composition maps (#66). + +Split the way the code is, because the two halves fail differently: + +* The BRIDGE is where a naming or ordering mistake sends iron's intensity into + copper's map. That is silent — every downstream number stays plausible — so + it gets the most tests. +* The PHYSICS is where a wrong k-factor lives, which is at least a number a + microscopist can sanity-check. + +The end-to-end test scores the result against the concentration maps the +synthetic data was generated from, so "quantification works" means "it recovers +the composition we put in", not "it produced numbers that sum to 1". +""" +from __future__ import annotations + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") + +from spyde.fitting import ModelSpec +from spyde.fitting.engine import FitResult +from spyde.fitting.spec import ComponentSpec, ParameterSpec +from spyde.spectroscopy.quantify import ( + element_intensity_maps, + parse_line, + quantify, + quantify_result, +) + + +def _spec(*names_kinds): + comps = [] + for name, kind in names_kinds: + pname = {"Gaussian": "A", "TabulatedShape": "intensity"}[kind] + comps.append(ComponentSpec( + kind=kind, name=name, + parameters=[ParameterSpec(pname, 1.0, linear=True)] + + ([ParameterSpec("centre", 0.0), ParameterSpec("sigma", 1.0)] + if kind == "Gaussian" else [ParameterSpec("onset_shift", 0.0)]))) + return ModelSpec(components=comps) + + +def _result(spec, values): + v = np.asarray(values, float) + return FitResult(v, np.ones(len(v), bool), np.zeros(len(v)), 1, "cpu") + + +class TestParseLine: + @pytest.mark.parametrize("name,want", [ + ("Fe_Ka", ("Fe", "Ka")), ("O_K", ("O", "K")), + ("Cu_Kb1", ("Cu", "Kb1")), ("C_K", ("C", "K")), + ]) + def test_parses_element_and_line(self, name, want): + assert parse_line(name) == want + + @pytest.mark.parametrize("name", ["background_order_6", "PowerLaw", + "", "Offset", "not-a-line"]) + def test_rejects_non_lines(self, name): + assert parse_line(name) is None + + +class TestBridge: + def test_one_map_per_element(self): + spec = _spec(("Fe_Ka", "Gaussian"), ("Cu_Ka", "Gaussian")) + res = _result(spec, [[10.0, 0, 1, 4.0, 0, 1], + [20.0, 0, 1, 8.0, 0, 1]]) + maps = element_intensity_maps(spec, res) + assert set(maps) == {"Fe", "Cu"} + np.testing.assert_allclose(maps["Fe"], [10.0, 20.0]) + np.testing.assert_allclose(maps["Cu"], [4.0, 8.0]) + + def test_sums_a_family(self): + """Ka + Kb belong to the same element and must add, not overwrite.""" + spec = _spec(("Fe_Ka", "Gaussian"), ("Fe_Kb", "Gaussian")) + res = _result(spec, [[10.0, 0, 1, 1.3, 0, 1]]) + assert element_intensity_maps(spec, res)["Fe"][0] == pytest.approx(11.3) + + def test_lines_filter_selects_one_line(self): + """Summing a family is usually right, but a contaminated member is + worse than none — so a clean line can be selected instead.""" + spec = _spec(("Fe_Ka", "Gaussian"), ("Fe_Kb", "Gaussian")) + res = _result(spec, [[10.0, 0, 1, 1.3, 0, 1]]) + maps = element_intensity_maps(spec, res, lines={"Fe": "Ka"}) + assert maps["Fe"][0] == pytest.approx(10.0) + + def test_background_is_ignored(self): + spec = ModelSpec(components=[ + ComponentSpec(kind="Offset", name="background_order_6", + parameters=[ParameterSpec("offset", 5.0)]), + ComponentSpec(kind="Gaussian", name="Fe_Ka", parameters=[ + ParameterSpec("A", 1.0), ParameterSpec("centre", 0.0), + ParameterSpec("sigma", 1.0)]), + ]) + res = _result(spec, [[99.0, 7.0, 0, 1]]) + assert set(element_intensity_maps(spec, res)) == {"Fe"} + + def test_negative_intensities_are_clamped(self): + """A fit can drive a line slightly negative on noise. A negative + 'amount of an element' is not physical AND it corrupts the + normalisation for every other element, so it is clamped at source.""" + spec = _spec(("Fe_Ka", "Gaussian"), ("Cu_Ka", "Gaussian")) + res = _result(spec, [[-3.0, 0, 1, 10.0, 0, 1]]) + maps = element_intensity_maps(spec, res) + assert maps["Fe"][0] == 0.0 + + def test_reshapes_to_the_scan(self): + spec = _spec(("Fe_Ka", "Gaussian")) + res = _result(spec, np.column_stack([np.arange(6.0), + np.zeros(6), np.ones(6)])) + assert element_intensity_maps(spec, res, (2, 3))["Fe"].shape == (2, 3) + + def test_reads_a_tabulated_edge_intensity(self): + """EELS edges carry `intensity`, EDS lines carry `A` — one bridge must + handle both or the EELS path silently produces no maps.""" + spec = _spec(("O_K", "TabulatedShape")) + res = _result(spec, [[7.0, 0.0]]) + assert element_intensity_maps(spec, res)["O"][0] == pytest.approx(7.0) + + def test_inactive_components_are_excluded(self): + spec = _spec(("Fe_Ka", "Gaussian"), ("Cu_Ka", "Gaussian")) + spec["Cu_Ka"].active = False + res = _result(spec, [[10.0, 0, 1]]) + assert set(element_intensity_maps(spec, res)) == {"Fe"} + + +class TestQuantify: + def test_relative_normalises_to_one(self): + frac, info = quantify({"Fe": np.array([2.0]), "Cu": np.array([6.0])}) + assert frac["Fe"][0] == pytest.approx(0.25) + assert frac["Cu"][0] == pytest.approx(0.75) + assert info["method"] == "relative" + + def test_cliff_lorimer_applies_k_factors(self): + """Equal intensities with a 2x k-factor must give a 2:1 composition.""" + frac, _ = quantify({"Fe": np.array([10.0]), "Cu": np.array([10.0])}, + method="cliff_lorimer", + kfactors={"Fe": 2.0, "Cu": 1.0}) + assert frac["Fe"][0] == pytest.approx(2 / 3) + assert frac["Cu"][0] == pytest.approx(1 / 3) + + def test_eels_divides_by_cross_section(self): + frac, _ = quantify({"C": np.array([10.0]), "O": np.array([10.0])}, + method="eels", + cross_sections={"C": 1.0, "O": 2.0}) + assert frac["C"][0] == pytest.approx(2 / 3) + + def test_missing_factors_default_and_are_reported(self): + """Defaulting silently would make an uncalibrated number look + calibrated.""" + _, info = quantify({"Fe": np.array([1.0]), "Cu": np.array([1.0])}, + method="cliff_lorimer", kfactors={"Fe": 1.0}) + assert info["defaulted_factors"] == ["Cu"] + + def test_result_says_it_is_relative(self): + """Cliff-Lorimer cannot know about an element you did not fit — the + fractions are relative, and leaving an element out redistributes its + share rather than causing a small error.""" + _, info = quantify({"Fe": np.array([1.0])}) + assert info["relative_to"] == ["Fe"] + assert "redistributed" in info["note"] + + def test_zero_signal_is_nan_not_a_divide_error(self): + frac, _ = quantify({"Fe": np.array([0.0, 5.0]), + "Cu": np.array([0.0, 5.0])}) + assert np.isnan(frac["Fe"][0]) + assert frac["Fe"][1] == pytest.approx(0.5) + + def test_fractions_sum_to_one_across_a_map(self): + rng = np.random.default_rng(0) + maps = {e: rng.random((4, 5)) + 0.1 for e in ("Fe", "Ni", "Cu")} + frac, _ = quantify(maps) + total = sum(frac.values()) + np.testing.assert_allclose(total, np.ones((4, 5)), atol=1e-12) + + def test_unknown_method_is_rejected(self): + with pytest.raises(ValueError, match="unknown quantification"): + quantify({"Fe": np.array([1.0])}, method="magic") + + def test_no_maps_is_an_error(self): + with pytest.raises(ValueError, match="no intensity maps"): + quantify({}) + + +class TestEndToEnd: + def test_recovers_the_composition_the_data_was_built_from(self): + """The real bar. Fit the synthetic EDS SI, quantify, and compare with + the concentration maps the generator used.""" + pytest.importorskip("exspy", reason="needs the eels extra") + from spyde.data import eds_si, ground_truth + from spyde.fitting.engine import fit_batched + from spyde.spectroscopy import model_for_composition + + s = eds_si(nav=(6, 6), n_channels=2048, noise=False) + spec, info = model_for_composition(s, ["Fe", "Ni", "Cu"], + energy_range=(5.5, 9.5)) + assert info["engine_supported"] + + x = np.asarray(s.axes_manager.signal_axes[0].axis, float) + res = fit_batched(spec, s.data, x, device="cpu", max_iter=120) + frac, qinfo = quantify_result(spec, res, (6, 6)) + + assert set(frac) == {"Fe", "Ni", "Cu"} + truth = ground_truth(s)["concentration"] + for el in ("Fe", "Ni", "Cu"): + got = frac[el].ravel() + want = np.asarray(truth[el]).ravel() + r = np.corrcoef(got, want)[0, 1] + assert r > 0.9, f"{el}: fitted composition does not track truth (r={r:.3f})" + + def test_quantify_result_carries_the_raw_maps(self): + spec = _spec(("Fe_Ka", "Gaussian"), ("Cu_Ka", "Gaussian")) + res = _result(spec, [[10.0, 0, 1, 30.0, 0, 1]]) + frac, info = quantify_result(spec, res) + assert info["intensity_maps"]["Fe"][0] == pytest.approx(10.0) + assert frac["Fe"][0] == pytest.approx(0.25) From de09ae76627027c1195d98113960acf7b0cf0af7 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Mon, 27 Jul 2026 02:24:54 -0500 Subject: [PATCH 16/60] feat(atoms): dumbbell lattices, with JOINT pair refinement (#78) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- spyde/atoms/__init__.py | 9 + spyde/atoms/dumbbell.py | 256 ++++++++++++++++++++ spyde/tests/migrated/test_atoms_dumbbell.py | 246 +++++++++++++++++++ 3 files changed, 511 insertions(+) create mode 100644 spyde/atoms/dumbbell.py create mode 100644 spyde/tests/migrated/test_atoms_dumbbell.py diff --git a/spyde/atoms/__init__.py b/spyde/atoms/__init__.py index bf0efd9..8460164 100644 --- a/spyde/atoms/__init__.py +++ b/spyde/atoms/__init__.py @@ -26,6 +26,13 @@ refine_center_of_mass, refine_gaussian, ) +from spyde.atoms.dumbbell import ( + dumbbell_properties, + estimate_dumbbell_vector, + find_dumbbells, + pair_atoms, + refine_pairs, +) from spyde.atoms.properties import ( displacement_from_ideal, displacement_magnitude, @@ -43,4 +50,6 @@ "ellipticity", "ellipticity_angle", "intensity", "nearest_neighbour_distance", "displacement_from_ideal", "displacement_magnitude", "property_maps", "to_map", + "find_dumbbells", "estimate_dumbbell_vector", "pair_atoms", + "dumbbell_properties", "refine_pairs", ] diff --git a/spyde/atoms/dumbbell.py b/spyde/atoms/dumbbell.py new file mode 100644 index 0000000..2bb071f --- /dev/null +++ b/spyde/atoms/dumbbell.py @@ -0,0 +1,256 @@ +"""dumbbell.py — dumbbell lattices (#78). + +Many structures image as PAIRS of closely-spaced atoms — silicon down <110> +being the canonical one. Treating each pair as two independent atoms throws +away the thing the pair actually measures: the separation and orientation of +the dumbbell, which is what carries polarisation, tilt and local strain. + +Following https://atomap.org/dumbbell_lattice.html, the workflow is: + +1. find the **dumbbell vector** — the displacement from one atom of a pair to + its partner, shared by every dumbbell in the field; +2. **pair** the atoms using it; +3. measure per-dumbbell properties. + +The vector is estimated from the data rather than assumed, because it is a +property of the projected structure and the scan rotation, and a user should +not have to type it in. :func:`estimate_dumbbell_vector` is deliberately +separable from :func:`pair_atoms` so a user who knows the vector — or who +picked it interactively (#76) — can skip the estimate. +""" +from __future__ import annotations + +import logging + +import numpy as np + +log = logging.getLogger(__name__) + + +def estimate_dumbbell_vector(positions, *, max_separation: float | None = None): + """The displacement between the two atoms of a dumbbell -> ``(dx, dy)``. + + Each atom's nearest neighbour IS its dumbbell partner when the pair + separation is smaller than the lattice spacing, which is what makes a + dumbbell a dumbbell. So the nearest-neighbour displacements form two tight + clusters at ``+v`` and ``-v``; folding them into one half-plane and taking + the median gives ``v`` robustly, without a fit. + + The median, not the mean: a few mis-paired atoms at the field edge would + drag a mean noticeably, and there is no reason to be sensitive to them. + """ + from scipy.spatial import cKDTree + + pos = np.asarray(positions, float).reshape(-1, 2) + if len(pos) < 2: + raise ValueError("need at least two atoms to estimate a dumbbell vector") + + dist, idx = cKDTree(pos).query(pos, k=2) + delta = pos[idx[:, 1]] - pos # to nearest neighbour + d = dist[:, 1] + + if max_separation is not None: + keep = d <= float(max_separation) + if not keep.any(): + raise ValueError( + f"no atom has a neighbour within {max_separation} px — either " + f"the separation is wrong or this is not a dumbbell lattice") + delta = delta[keep] + + # +v and -v describe the same dumbbell, so fold onto one half-plane before + # averaging or the two clusters cancel to zero. + flip = (delta[:, 0] < 0) | ((delta[:, 0] == 0) & (delta[:, 1] < 0)) + delta[flip] *= -1 + return np.median(delta, axis=0) + + +def pair_atoms(positions, vector, *, tolerance: float = 0.4): + """Pair atoms into dumbbells using a known dumbbell *vector*. + + Greedy nearest-partner matching: each atom's best candidate is the one + closest to ``position + vector``, and an atom already claimed cannot be + claimed again. Greedy is enough because a correct dumbbell's partner is + unambiguous — anything closer than *tolerance* × |vector| to the wrong atom + means the vector itself is wrong, which is worth surfacing rather than + papering over with a global assignment. + + Parameters + ---------- + tolerance : float + Allowed mismatch as a fraction of the dumbbell length. + + Returns + ------- + (pairs, unpaired) + *pairs* is ``(M, 2)`` of indices into *positions*, first atom then its + partner along ``+vector``. *unpaired* are the indices left over — + edge atoms and genuine singles, which are reported rather than + silently dropped. + """ + from scipy.spatial import cKDTree + + pos = np.asarray(positions, float).reshape(-1, 2) + v = np.asarray(vector, float).ravel() + length = float(np.hypot(*v)) + if length <= 0: + raise ValueError("dumbbell vector has zero length") + + tree = cKDTree(pos) + targets = pos + v + dist, idx = tree.query(targets, k=1) + + taken = np.zeros(len(pos), bool) + pairs = [] + # Best matches first, so a confident pair claims its partner before an + # ambiguous one can steal it. + for i in np.argsort(dist): + j = int(idx[i]) + if i == j or taken[i] or taken[j]: + continue + if dist[i] > tolerance * length: + continue + taken[i] = taken[j] = True + pairs.append((int(i), j)) + + unpaired = np.flatnonzero(~taken) + if len(unpaired): + log.debug("%d atom(s) left unpaired out of %d", len(unpaired), len(pos)) + return np.array(pairs, int).reshape(-1, 2), unpaired + + +def dumbbell_properties(positions, pairs, params=None) -> dict[str, np.ndarray]: + """Per-dumbbell measurements — one value per PAIR, not per atom. + + ``separation`` + Distance between the two atoms. The primary measurement. + ``angle`` + Orientation in radians, wrapped to ``(-pi/2, pi/2]`` because a dumbbell + has no head or tail — reporting it over the full circle would make + identical dumbbells differ by pi depending on which atom was listed + first. + ``centre_x`` / ``centre_y`` + Midpoint, which is the position to use for a lattice-level map. + ``intensity_ratio`` + Second atom over first, when *params* is given. Distinguishes the two + sites of a polar structure; 1.0 means the pair is symmetric. + """ + pos = np.asarray(positions, float).reshape(-1, 2) + pr = np.asarray(pairs, int).reshape(-1, 2) + a, b = pos[pr[:, 0]], pos[pr[:, 1]] + d = b - a + + angle = np.arctan2(d[:, 1], d[:, 0]) + # Fold to a half-turn: +v and -v are the same dumbbell. + angle = (angle + np.pi / 2) % np.pi - np.pi / 2 + + out = { + "separation": np.hypot(d[:, 0], d[:, 1]), + "angle": angle, + "centre_x": (a[:, 0] + b[:, 0]) / 2, + "centre_y": (a[:, 1] + b[:, 1]) / 2, + } + if params is not None: + p = np.asarray(params, float) + i0, i1 = p[pr[:, 0], 0], p[pr[:, 1], 0] + with np.errstate(divide="ignore", invalid="ignore"): + out["intensity_ratio"] = np.where(i0 > 0, i1 / i0, np.nan) + return out + + +def refine_pairs(image, positions, pairs, *, sigma: float = 2.5, + margin: float = 3.0, device=None, max_iter: int = 80): + """Refine both atoms of each dumbbell **jointly**, as one two-gaussian fit. + + This is not a refinement of a refinement — it corrects a systematic error. + Fitting each atom independently biases its centre TOWARDS its partner, + because the partner's tail is signal the single-gaussian model can only + explain by moving. The pair separation therefore comes out too small, and + uniformly so, which makes it look like a calibration rather than a bug: + measured on a synthetic dumbbell of 6.0 px with sigma 2.0, independent fits + return **4.21 px**, a 30% underestimate. Joint fitting recovers 6.0. + + The two atoms share one box and one model of two ``Gaussian2D`` + components, so each explains the other's tail instead of absorbing it. Every + pair is still fitted in ONE batched call — the engine does not care that the + model now has two components. + + Returns refined ``(N, 2)`` positions (unpaired atoms keep their input) and + the per-pair parameter table. + """ + from spyde.fitting import ModelSpec + from spyde.fitting.components import image_coordinates + from spyde.fitting.engine import fit_batched + from spyde.fitting.spec import ComponentSpec, ParameterSpec + + img = np.asarray(image, float) + pos = np.asarray(positions, float).reshape(-1, 2) + pr = np.asarray(pairs, int).reshape(-1, 2) + out = pos.copy() + if not len(pr): + return out, np.empty((0, 10)) + + a, b = pos[pr[:, 0]], pos[pr[:, 1]] + sep = float(np.median(np.hypot(*(b - a).T))) + # The box must hold BOTH atoms and enough tail for the fit to see where + # each one ends; too tight and the joint fit inherits the bias it exists + # to remove. + box = int(2 * round(sep / 2 + margin * sigma) + 1) + r = box // 2 + + h, w = img.shape + mid = (a + b) / 2 + cx = np.clip(np.rint(mid[:, 0]).astype(int), r, w - r - 1) + cy = np.clip(np.rint(mid[:, 1]).astype(int), r, h - r - 1) + oy, ox = np.mgrid[-r:r + 1, -r:r + 1] + patches = img[cy[:, None, None] + oy[None], + cx[:, None, None] + ox[None]].reshape(len(pr), -1) + + def _comp(name): + return ComponentSpec(kind="Gaussian2D", name=name, parameters=[ + ParameterSpec("A", 1.0, linear=True), + ParameterSpec("centre_x", float(r)), + ParameterSpec("centre_y", float(r)), + ParameterSpec("sigma_x", float(sigma), bmin=0.3, bmax=float(box)), + ParameterSpec("sigma_y", float(sigma), bmin=0.3, bmax=float(box)), + ]) + + spec = ModelSpec(components=[_comp("atom0"), _comp("atom1")]) + names = spec.parameter_names() + start = np.broadcast_to(spec.flat_values(), (len(pr), len(names))).copy() + peak = patches.max(1) / 2.0 + for tag, xy in (("atom0", a), ("atom1", b)): + start[:, names.index(f"{tag}.A")] = peak * 2 * np.pi * sigma * sigma + start[:, names.index(f"{tag}.centre_x")] = xy[:, 0] - cx + r + start[:, names.index(f"{tag}.centre_y")] = xy[:, 1] - cy + r + + xy_grid = image_coordinates((box, box)).numpy().astype(float) + res = fit_batched(spec, patches, xy_grid, device=device, + max_iter=max_iter, initial=start) + + for tag, col in (("atom0", 0), ("atom1", 1)): + px = res.values[:, names.index(f"{tag}.centre_x")] + cx - r + py = res.values[:, names.index(f"{tag}.centre_y")] + cy - r + keep = (np.abs(px - pos[pr[:, col], 0]) <= r) & \ + (np.abs(py - pos[pr[:, col], 1]) <= r) & \ + np.isfinite(px) & np.isfinite(py) + out[pr[keep, col], 0] = px[keep] + out[pr[keep, col], 1] = py[keep] + return out, res.values + + +def find_dumbbells(positions, *, vector=None, max_separation=None, + tolerance: float = 0.4, params=None): + """Estimate the vector, pair the atoms, measure — the whole workflow. + + Returns ``(pairs, properties, info)``; *info* carries the vector used and + how many atoms went unpaired, so a bad estimate is visible instead of + showing up later as a suspiciously small dumbbell count. + """ + v = (estimate_dumbbell_vector(positions, max_separation=max_separation) + if vector is None else np.asarray(vector, float).ravel()) + pairs, unpaired = pair_atoms(positions, v, tolerance=tolerance) + props = dumbbell_properties(positions, pairs, params) + info = {"vector": v, "length": float(np.hypot(*v)), + "n_pairs": len(pairs), "n_unpaired": int(len(unpaired)), + "unpaired": unpaired} + return pairs, props, info diff --git a/spyde/tests/migrated/test_atoms_dumbbell.py b/spyde/tests/migrated/test_atoms_dumbbell.py new file mode 100644 index 0000000..6fd711c --- /dev/null +++ b/spyde/tests/migrated/test_atoms_dumbbell.py @@ -0,0 +1,246 @@ +"""Dumbbell lattices (#78). + +The synthetic lattice generates dumbbells with a KNOWN separation, so every +test here scores against that rather than against a fixture. + +The properties that matter and are easy to get subtly wrong: the estimated +vector must not cancel to zero (each pair contributes +v and -v), pairing must +not claim an atom twice, and the angle must be folded to a half-turn or two +identical dumbbells differ by pi depending on which atom happened to be listed +first. +""" +from __future__ import annotations + +import numpy as np +import pytest + +from spyde.atoms.dumbbell import ( + dumbbell_properties, + estimate_dumbbell_vector, + find_dumbbells, + pair_atoms, +) +from spyde.data import atom_lattice, ground_truth + +SEP = 6.0 + + +@pytest.fixture(scope="module") +def dumbbells(): + s = atom_lattice(grid=(5, 6), spacing=20.0, dumbbell=SEP, sigma=2.0, + noise=0.0) + return np.asarray(s.data, float), np.asarray(ground_truth(s)["positions"]) + + +class TestEstimateVector: + def test_recovers_the_known_separation(self, dumbbells): + _, pos = dumbbells + v = estimate_dumbbell_vector(pos) + assert np.hypot(*v) == pytest.approx(SEP, rel=0.02) + + def test_recovers_the_known_direction(self, dumbbells): + """The generator splits along x, so the vector must be along x.""" + _, pos = dumbbells + v = estimate_dumbbell_vector(pos) + assert abs(v[1]) < 0.1 * abs(v[0]) + + def test_does_not_cancel_to_zero(self, dumbbells): + """Each pair contributes +v from one atom and -v from the other. An + implementation that averages without folding onto a half-plane gets + exactly zero — and zero is a plausible-looking answer.""" + _, pos = dumbbells + assert np.hypot(*estimate_dumbbell_vector(pos)) > 1.0 + + def test_works_for_a_diagonal_dumbbell(self): + """Nothing may assume the pair splits along an axis.""" + base = np.array([[x * 30.0, y * 30.0] + for y in range(4) for x in range(4)]) + v = np.array([4.0, 3.0]) # length 5, diagonal + pos = np.vstack([base, base + v]) + got = estimate_dumbbell_vector(pos) + assert np.hypot(*got) == pytest.approx(5.0, rel=1e-6) + assert abs(got[1] / got[0]) == pytest.approx(0.75, rel=1e-6) + + def test_max_separation_rejects_a_non_dumbbell_lattice(self): + """A plain lattice has no pair closer than the spacing; saying so beats + returning a 'dumbbell vector' that is really the lattice vector.""" + pos = np.array([[x * 20.0, y * 20.0] + for y in range(4) for x in range(4)]) + with pytest.raises(ValueError, match="not a dumbbell lattice"): + estimate_dumbbell_vector(pos, max_separation=5.0) + + def test_too_few_atoms_is_an_error(self): + with pytest.raises(ValueError, match="at least two"): + estimate_dumbbell_vector(np.array([[1.0, 2.0]])) + + +class TestPairing: + def test_pairs_every_atom_on_a_clean_lattice(self, dumbbells): + _, pos = dumbbells + v = estimate_dumbbell_vector(pos) + pairs, unpaired = pair_atoms(pos, v) + assert len(pairs) == len(pos) // 2 + assert len(unpaired) == 0 + + def test_no_atom_is_claimed_twice(self, dumbbells): + """A greedy matcher that forgets to mark atoms taken produces more + pairs than atoms and every downstream count is wrong.""" + _, pos = dumbbells + pairs, _ = pair_atoms(pos, estimate_dumbbell_vector(pos)) + flat = pairs.ravel() + assert len(set(flat.tolist())) == len(flat) + + def test_partners_are_the_real_partners(self, dumbbells): + """Each pair must be SEP apart — pairing neighbouring dumbbells instead + would still return the right number of pairs.""" + _, pos = dumbbells + pairs, _ = pair_atoms(pos, estimate_dumbbell_vector(pos)) + d = np.hypot(*(pos[pairs[:, 1]] - pos[pairs[:, 0]]).T) + np.testing.assert_allclose(d, SEP, rtol=0.02) + + def test_a_wrong_vector_leaves_atoms_unpaired(self): + """Surfacing a bad vector beats papering over it: an atom with no + partner at the expected offset must be REPORTED, not force-matched.""" + base = np.array([[x * 30.0, 0.0] for x in range(4)]) + pos = np.vstack([base, base + np.array([5.0, 0.0])]) + _, unpaired = pair_atoms(pos, np.array([13.0, 0.0])) + assert len(unpaired) > 0 + + def test_zero_vector_is_rejected(self, dumbbells): + _, pos = dumbbells + with pytest.raises(ValueError, match="zero length"): + pair_atoms(pos, np.array([0.0, 0.0])) + + def test_an_odd_atom_out_is_reported_not_dropped(self): + base = np.array([[0.0, 0.0], [5.0, 0.0], [40.0, 0.0]]) + pairs, unpaired = pair_atoms(base, np.array([5.0, 0.0])) + assert len(pairs) == 1 + assert unpaired.tolist() == [2] + + +class TestProperties: + def test_separation_matches_the_truth(self, dumbbells): + _, pos = dumbbells + pairs, props, _ = find_dumbbells(pos) + np.testing.assert_allclose(props["separation"], SEP, rtol=0.02) + + def test_angle_is_folded_to_a_half_turn(self): + """+v and -v are the SAME dumbbell. Without folding, two identical + dumbbells differ by pi depending on which atom was listed first, and an + angle map shows a boundary that is not there.""" + pos = np.array([[0.0, 0.0], [5.0, 0.0], + [100.0, 0.0], [95.0, 0.0]]) + props = dumbbell_properties(pos, np.array([[0, 1], [2, 3]])) + assert props["angle"][0] == pytest.approx(props["angle"][1], abs=1e-9) + + def test_angle_tracks_a_real_rotation(self): + pos = np.array([[0.0, 0.0], [3.0, 3.0]]) + props = dumbbell_properties(pos, np.array([[0, 1]])) + assert props["angle"][0] == pytest.approx(np.pi / 4) + + def test_centre_is_the_midpoint(self): + pos = np.array([[10.0, 20.0], [16.0, 20.0]]) + props = dumbbell_properties(pos, np.array([[0, 1]])) + assert props["centre_x"][0] == pytest.approx(13.0) + assert props["centre_y"][0] == pytest.approx(20.0) + + def test_one_value_per_pair_not_per_atom(self, dumbbells): + _, pos = dumbbells + pairs, props, _ = find_dumbbells(pos) + for name, v in props.items(): + assert len(v) == len(pairs), name + assert len(pairs) == len(pos) // 2 + + def test_intensity_ratio_distinguishes_the_two_sites(self): + """A polar structure has unequal sites; a symmetric pair must read 1.""" + pos = np.array([[0.0, 0.0], [5.0, 0.0], [50.0, 0.0], [55.0, 0.0]]) + params = np.zeros((4, 5)) + params[:, 0] = [10.0, 20.0, 10.0, 10.0] + props = dumbbell_properties(pos, np.array([[0, 1], [2, 3]]), params) + assert props["intensity_ratio"][0] == pytest.approx(2.0) + assert props["intensity_ratio"][1] == pytest.approx(1.0) + + +class TestFullWorkflow: + def test_reports_the_vector_it_used(self, dumbbells): + """A bad estimate must be visible, not show up later as a + suspiciously small dumbbell count.""" + _, pos = dumbbells + _, _, info = find_dumbbells(pos) + assert np.hypot(*info["vector"]) == pytest.approx(SEP, rel=0.02) + assert info["n_unpaired"] == 0 + assert info["n_pairs"] == len(pos) // 2 + + def test_accepts_an_explicit_vector(self, dumbbells): + """A user who knows the vector — or picked it interactively (#76) — + can skip the estimate.""" + _, pos = dumbbells + _, _, info = find_dumbbells(pos, vector=(SEP, 0.0)) + assert info["n_pairs"] == len(pos) // 2 + + def test_survives_refinement_jitter(self, dumbbells): + """Real refined positions are not exact, so pairing must tolerate + sub-pixel scatter.""" + _, pos = dumbbells + jittered = pos + np.random.default_rng(0).uniform(-0.3, 0.3, pos.shape) + _, props, info = find_dumbbells(jittered) + assert info["n_unpaired"] == 0 + np.testing.assert_allclose(props["separation"], SEP, atol=0.8) + + def test_independent_fits_underestimate_the_separation(self, dumbbells): + """Documents the error joint fitting exists to remove. + + Fitting each atom on its own 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 a bug. + """ + img, truth = dumbbells + from spyde.atoms import refine_gaussian + + refined, _ = refine_gaussian(img, truth, box=9, sigma=2.0, + device="cpu") + _, props, _ = find_dumbbells(refined) + measured = float(np.median(props["separation"])) + assert measured < 0.85 * SEP, ( + f"expected the known inward bias; got {measured:.2f} for a true " + f"{SEP}. If this now passes, joint fitting may have become the " + f"default and this test is obsolete.") + + def test_joint_fitting_recovers_the_true_separation(self, dumbbells): + """The fix: both atoms in ONE box as one two-gaussian model, so each + explains the other's tail instead of absorbing it.""" + img, truth = dumbbells + from spyde.atoms.dumbbell import refine_pairs + + pairs, _ = pair_atoms(truth, estimate_dumbbell_vector(truth)) + refined, _ = refine_pairs(img, truth, pairs, sigma=2.0, device="cpu") + props = dumbbell_properties(refined, pairs) + np.testing.assert_allclose(props["separation"], SEP, atol=0.15) + + def test_joint_fitting_beats_independent_fitting(self, dumbbells): + """The comparison stated directly, so a regression in either path + shows up as the gap closing from the wrong side.""" + img, truth = dumbbells + from spyde.atoms import refine_gaussian + from spyde.atoms.dumbbell import refine_pairs + + pairs, _ = pair_atoms(truth, estimate_dumbbell_vector(truth)) + solo, _ = refine_gaussian(img, truth, box=9, sigma=2.0, device="cpu") + joint, _ = refine_pairs(img, truth, pairs, sigma=2.0, device="cpu") + + err_solo = abs(np.median(dumbbell_properties(solo, pairs)["separation"]) - SEP) + err_joint = abs(np.median(dumbbell_properties(joint, pairs)["separation"]) - SEP) + assert err_joint < err_solo / 5, (err_joint, err_solo) + + def test_end_to_end_from_the_image(self, dumbbells): + """The order a user works in: find pairs, then refine them jointly.""" + img, truth = dumbbells + from spyde.atoms.dumbbell import refine_pairs + + pairs, props, info = find_dumbbells(truth) + assert info["n_pairs"] == len(truth) // 2 + refined, params = refine_pairs(img, truth, pairs, sigma=2.0, + device="cpu") + final = dumbbell_properties(refined, pairs) + np.testing.assert_allclose(final["separation"], SEP, atol=0.15) From 955a2f79092add640423cb31fdaa06be76021cd0 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Mon, 27 Jul 2026 08:09:47 -0500 Subject: [PATCH 17/60] =?UTF-8?q?feat(fit):=20the=20Fit=20wizard=20?= =?UTF-8?q?=E2=80=94=20build=20a=20model,=20fit=20the=20scan,=20commit=20m?= =?UTF-8?q?aps=20(#55,=20#56,=20#58)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The first UI on top of spyde.fitting. Two tabs, matching the Orientation caret's shape: Model builds the model component by component, Run fits it and commits the maps. Layout is horizontal by preference — a component's parameters sit side by side so the caret grows wider rather than taller. #56, the picker: components are added from a "+ Component" POPUP that shows each one's SHAPE, not just its name. The backend samples every offerable component at defaults over the CURRENT signal axis and sends a normalised polyline the renderer draws as a sparkline. Seeding matters — a default Gaussian sits at 0 with sigma 1, which is off-screen on a 200-800 eV axis, so every peak shape would otherwise preview as an identical flat line. #58, the maps: Commit ships one integrated-area map per component through commit_result_tree, which already gives the strain-style toggle (click one, cmd-click to tile). No new display code. The preview is ONE spectrum, not the grid: a model edit has to feel instant, and fitting 65k positions per keystroke would not. fit_run is the only thing that touches the whole scan. The caret rebuilds itself from `fit_state` and never edits a local copy — the backend owns the model, and a caret that mutated its own would drift out of step the moment an edit was rejected, then show a model that is not the one being fitted. TWO REAL BUGS that only running the app could find; both handler-test suites pass either way: 1. A click on a caret control was SILENTLY LOST. Capture-phase listeners showed `mousedown -> the button's span, mouseup -> mdi-area`: the mousedown bubbled to the window, re-rendered the subwindow subtree, and replaced the control between down and up, so React never saw a click. The add simply did not happen, with no error anywhere. Fixed by keeping mousedown inside the caret. (dispatchEvent worked throughout, which is exactly what makes this look like a test problem rather than a UI one.) 2. `emit` was imported as `from ipc import emit`, which binds the original at import time — the test fixture patches `ipc.emit`, so every message assertion would have silently observed nothing. Imported as a module now. Backend: 26 handler tests (including the StrictMode open/close/open contract). UI: fit_wizard.spec.ts drives the real app end to end — toolbar button, palette sparklines, model building, a parameter edit echoed back from the backend, the fit, teardown/reopen, and Commit opening a window with a chip per component. Screenshots at every stage. --- .../src/renderer/src/components/FitWizard.tsx | 222 ++++++++ .../src/components/FloatingToolbar.tsx | 9 +- .../src/renderer/src/kernel/SpyDEContext.tsx | 5 + electron/src/renderer/src/kernel/protocol.ts | 4 + electron/tests/fit_wizard.spec.ts | 162 ++++++ spyde/actions/fit_action.py | 474 ++++++++++++++++++ spyde/actions/registry.py | 9 + spyde/drawing/toolbars/icons/fit.svg | 26 + spyde/tests/migrated/test_fit_wizard.py | 325 ++++++++++++ spyde/toolbars.yaml | 8 + 10 files changed, 1243 insertions(+), 1 deletion(-) create mode 100644 electron/src/renderer/src/components/FitWizard.tsx create mode 100644 electron/tests/fit_wizard.spec.ts create mode 100644 spyde/actions/fit_action.py create mode 100644 spyde/drawing/toolbars/icons/fit.svg create mode 100644 spyde/tests/migrated/test_fit_wizard.py diff --git a/electron/src/renderer/src/components/FitWizard.tsx b/electron/src/renderer/src/components/FitWizard.tsx new file mode 100644 index 0000000..9e54406 --- /dev/null +++ b/electron/src/renderer/src/components/FitWizard.tsx @@ -0,0 +1,222 @@ +/** + * FitWizard.tsx — the Fit caret (#55, #56, #58). + * + * Two tabs, matching the Orientation-Mapping caret's shape: **Model** builds + * the model component by component, **Run** fits it and commits the maps. + * + * Layout is HORIZONTAL by preference (the same rule the toolbar popouts + * follow): a component's parameters sit side by side so the caret grows WIDER + * rather than taller. A tall caret is not just ugly here — it overhangs the + * room FloatingToolbar placed it in, and the overhang sits under the MDI area. + * + * Three deliberate choices: + * + * - **Components are added from a `+` POPUP**, not an always-open palette. The + * palette is a one-off action; leaving it permanently expanded spends the + * caret's height on something used once per component. + * - **The popup shows SHAPES, not just names** (#56). The backend samples every + * offerable component at defaults over the current signal axis and sends a + * normalised polyline, drawn as an inline SVG sparkline. "Gaussian" and + * "Lorentzian" are not distinguishable by name to someone who has not met + * them, and the point of a picker is to be able to choose. + * - **The component list is rebuilt from `fit_state`, never edited locally.** + * The backend owns the model; a caret that mutated its own copy would drift + * out of step the moment an edit was rejected and would then show a model + * that is not the one being fitted. + */ +import React from 'react' +import { WizardShell, TabRow, Field, NumInput, Select, Check, S } from './WizardShell' +import { useWizardLifecycle, useWizardEvent, CommitButton } from './wizardHooks' + +const TABS = ['Model', 'Run'] as const +type Tab = typeof TABS[number] + +interface Props { + caretPos: React.CSSProperties + windowId: number + sendAction: (action: string, payload?: Record, windowId?: number) => void + onClose: () => void +} + +interface ParamState { name: string; value: number; free: boolean; linear: boolean } +interface CompState { name: string; kind: string; active: boolean; parameters: ParamState[] } +interface CatalogueItem { kind: string; description: string; preview: number[] } + +/** Inline sparkline of a component's shape — the picker's whole point. */ +function Spark({ points }: { points: number[] }) { + if (!points?.length) return null + const w = 40, h = 15 + const d = points + .map((v, i) => `${(i / (points.length - 1)) * w},${h - v * (h - 2) - 1}`) + .join(' ') + return ( + + + + ) +} + +export function FitWizard({ caretPos, windowId, sendAction, onClose }: Props) { + const [tab, setTab] = React.useState('Model') + const [catalogue, setCatalogue] = React.useState([]) + const [components, setComponents] = React.useState([]) + const [fitted, setFitted] = React.useState(false) + const [status, setStatus] = React.useState('Add a component to begin.') + const [pickerOpen, setPickerOpen] = React.useState(false) + const [maxIter, setMaxIter] = React.useState(60) + const [seeded, setSeeded] = React.useState(true) + const [weighting, setWeighting] = React.useState<'none' | 'poisson'>('none') + + useWizardLifecycle({ + windowId, sendAction, + openAction: 'fit_open', + closeAction: 'fit_close', + deps: [windowId], + }) + + useWizardEvent('spyde:fit_catalogue', windowId, (detail) => { + const d = detail as { components?: CatalogueItem[] } + if (d.components) setCatalogue(d.components) + }) + + useWizardEvent('spyde:fit_state', windowId, (detail) => { + const d = detail as { components?: CompState[]; fitted?: boolean; status?: string } + if (d.components) setComponents(d.components) + setFitted(Boolean(d.fitted)) + if (d.status) setStatus(d.status) + }) + + const add = (kind: string) => { + sendAction('fit_add_component', { kind }, windowId) + setPickerOpen(false) + } + + const setParam = (component: string, parameter: string, value: number) => + sendAction('fit_set_param', { component, parameter, value }, windowId) + + const toggleFree = (component: string, parameter: string, free: boolean) => + sendAction('fit_set_param', { component, parameter, free }, windowId) + + const run = () => { + setStatus('Fitting…') + setTab('Run') + sendAction('fit_run', { max_iter: maxIter, seeded, weighting }, windowId) + } + + return ( + + {/* Keep mousedown INSIDE the caret. Letting it bubble to the window + re-renders the subwindow subtree, which replaces the control being + clicked between mousedown and mouseup — the mouseup then lands on the + MDI area, React never sees a click, and the action silently does not + happen. Measured with capture-phase listeners: before this, + `mousedown -> the button's span, mouseup -> mdi-area`; after, all + three land on the span. Only the running app finds this; the handler + tests pass either way. */} +
e.stopPropagation()}> + `fit-tab-${t}`} /> + + {tab === 'Model' && ( +
+ {components.length === 0 && ( + No components yet — add one with +. + )} + + {/* One card per component; its parameters lie side by side so the + caret grows wider, not taller. */} +
+ {components.map((c) => ( +
+
+ {c.name} + +
+
+ {c.parameters.map((p) => ( +
+ {p.name} + setParam(c.name, p.name, v)} /> + toggleFree(c.name, p.name, e.target.checked)} /> +
+ ))} +
+
+ ))} +
+ + {/* ── the + picker ── */} +
+ + {pickerOpen && ( +
+ {catalogue.map((c) => ( + + ))} +
+ )} +
+
+ )} + + {tab === 'Run' && ( +
+ {/* Horizontal: the three knobs sit on one row. */} +
+ + setMaxIter(Math.max(5, Math.round(v)))} /> + + + ({ value: m, label: m }))} + onChange={(v) => { setModel(v); sendAction('bg_set_model', { model: v }, windowId) }} /> + + + setField('x0', v)} /> + + + setField('x1', v)} /> + +
+
+ +
+
+
+ + ) +} diff --git a/electron/src/renderer/src/components/FitWizard.tsx b/electron/src/renderer/src/components/FitWizard.tsx index acac6d0..3fdc2ac 100644 --- a/electron/src/renderer/src/components/FitWizard.tsx +++ b/electron/src/renderer/src/components/FitWizard.tsx @@ -77,6 +77,11 @@ export function FitWizard({ caretPos, windowId, sendAction, onClose }: Props) { // How many positions fit worse than their neighbours — the number the // "Refit poor" button acts on, and the honest headline for a scan fit. const [poor, setPoor] = React.useState(0) + // Models already stored on the signal. These are HyperSpy's own — `m.store` + // puts the components AND every position's fit into the signal's `models`, + // so they travel with the dataset rather than in a format of ours. + const [storedModels, setStoredModels] = React.useState([]) + const [modelName, setModelName] = React.useState('spyde fit') // Read inside the navigator listener, which is registered once — a state // value captured there would be the value at registration forever. const adaptiveRef = React.useRef(adaptive) @@ -129,11 +134,12 @@ export function FitWizard({ caretPos, windowId, sendAction, onClose }: Props) { const d = detail as { components?: CompState[]; fitted?: boolean; status?: string fitted_count?: number; nav_total?: number; position_fitted?: boolean - poor_count?: number + poor_count?: number; stored_models?: string[] } if (d.components) setComponents(d.components) setFitted(Boolean(d.fitted)) setPoor(d.poor_count ?? 0) + if (d.stored_models) setStoredModels(d.stored_models) // Every backend edit ends in a `fit_state`, so this is the completion // signal the navigator coalescer waits on — see sendNavigated. navDone() @@ -355,7 +361,39 @@ export function FitWizard({ caretPos, windowId, sendAction, onClose }: Props) { )} {fitted && ( + label="Commit components" /> + )} + + + {/* ── save / load, through HyperSpy's own model store ── + `m.store(name)` puts the components AND every position's fit + into the signal's `models`, so saving the .hspy/.zspy saves the + fit with it. Nothing here is a SpyDE format. */} +
+ Model + setModelName(e.target.value)} + style={{ background: '#11111b', border: '1px solid #313244', + borderRadius: 4, color: '#cdd6f4', fontSize: 11, + padding: '2px 5px', width: 110 }} /> + + {storedModels.length > 0 && ( + <> + + + + + +
+ A finer step indexes closer but costs simulation time and slows the + live preview. +
+ +
+ )} + + {tab === 'Refine' && ( +
+
+ Move the crosshair on the navigator: the matched orientation's + Kikuchi bands are drawn on the pattern. Nudge the PC until the lines + sit on the bands. +
+
+ {match + ? `NCC ${match.score.toFixed(3)} — φ1 ${match.phi1.toFixed(1)}° Φ ${match.Phi.toFixed(1)}° φ2 ${match.phi2.toFixed(1)}°` + : 'No match yet — move the crosshair.'} +
+ + { setNBands(n); tune({ nBands: n }) }} /> + + { setZoneAxes(b); tune({ zoneAxes: b }) }} /> + + { setPcx(n); tune({ pcx: n }) }} /> + + + { setPcy(n); tune({ pcy: n }) }} /> + + + { setPcz(n); tune({ pcz: n }) }} /> + +
+ )} + + {tab === 'Run' && ( +
+ + + +
+ More than one match per pattern is what the orientation-similarity + quality map needs. +
+ + {refine && ( + + + + )} + +
+ )} + + ) +} diff --git a/electron/src/renderer/src/components/FloatingToolbar.tsx b/electron/src/renderer/src/components/FloatingToolbar.tsx index 52e6f0f..3d0279d 100644 --- a/electron/src/renderer/src/components/FloatingToolbar.tsx +++ b/electron/src/renderer/src/components/FloatingToolbar.tsx @@ -24,12 +24,14 @@ import type { ToolbarAction, ParamSpec, SubAction } from '../kernel/SpyDEContext import { OrientationWizard } from './OrientationWizard' import { FindVectorsWizard } from './FindVectorsWizard' import { VectorOrientationWizard } from './VectorOrientationWizard' +import { EbsdWizard } from './EbsdWizard' import { CenterZeroBeamWizard } from './CenterZeroBeamWizard' import { StrainWizard } from './StrainWizard' import { CropWizard } from './CropWizard' const WIZARD_ACTIONS = new Set([ 'Orientation Mapping', 'Find Diffraction Vectors', 'Vector Orientation Mapping', + 'EBSD Indexing', 'Center Zero Beam', 'Strain Mapping', 'Crop', ]) @@ -50,6 +52,7 @@ function fileUrl(p: string): string { // while the action's caret is selected, and hidden when it's deselected. const OVERLAY_ACTIONS = new Set([ 'Find Diffraction Vectors', 'Orientation Mapping', 'Vector Orientation Mapping', + 'EBSD Indexing', ]) const EMPTY = new Set() @@ -286,6 +289,12 @@ export function FloatingToolbar({ onClose={() => setOpenName(null)} /> )} + {openAction && openAction.name === 'EBSD Indexing' && ( + setOpenName(null)} + /> + )} {openAction && openAction.name === 'Center Zero Beam' && ( > + +test.beforeAll(async () => { + ctx = await launchApp({ dask: true, env: { SPYDE_LOG_LEVEL: 'INFO' } }) + await backendAction(ctx.page, 'load_test_data_ebsd', { nav: [16, 16], detector: [60, 60] }) + await waitForSubwindowCount(ctx.page, 2, 120_000) +}) + +test.afterAll(async () => { + await ctx?.app?.close() +}) + +test.setTimeout(300_000) + +/** Open the EBSD Indexing caret on the pattern (signal) window. */ +async function openCaret() { + const { page } = ctx + const sig = sigWindow(page) + // Focus-raise then hover the TITLEBAR — hovering the figure itself puts the + // cursor inside the iframe and the toolbar never mounts. + await sig.getByTestId('subwindow-title').click() + await sig.getByTestId('subwindow-titlebar').hover() + const button = sig.getByTestId('action-btn-EBSD Indexing') + await expect(button, 'the EBSD Indexing toolbar button never appeared — the ' + + 'signal_types: [EBSD] gate did not match').toBeVisible({ timeout: 30_000 }) + await button.click() + await expect(page.getByTestId('ebsd-wizard')).toBeVisible({ timeout: 10_000 }) +} + +test('the caret opens on an EBSD scan and builds a dictionary', async () => { + const { page } = ctx + await page.screenshot({ path: `${SHOTS}/01-loaded.png` }) + + await openCaret() + await page.screenshot({ path: `${SHOTS}/02-caret-load-tab.png` }) + + // 2 Library → a coarse step so the dictionary builds quickly. + await page.getByTestId('ebsd-tab-Library').click() + await page.getByTestId('ebsd-step').fill('6') + await page.screenshot({ path: `${SHOTS}/03-caret-library-tab.png` }) + await page.getByTestId('ebsd-build').click() + + // The backend acks with ebsd_dictionary_ready → the caret's status line. + await expect(page.getByTestId('ebsd-status')) + .toContainText(/Dictionary ready/i, { timeout: 180_000 }) + await page.screenshot({ path: `${SHOTS}/04-dictionary-ready.png` }) + ctx.assertNoJsErrors() +}) + +test('the matched orientation draws Kikuchi band lines on the pattern', async () => { + const { page } = ctx + + // The Refine tab shows the live match; the overlay is already attached. + await page.getByTestId('ebsd-tab-Refine').click() + await expect(page.getByTestId('ebsd-match')) + .toContainText(/NCC/i, { timeout: 60_000 }) + + await expect.poll(() => countColorPixels(page, 'green'), { + timeout: 60_000, + message: 'no green band lines drew on the EBSD pattern', + }).toBeGreaterThan(0) + + await page.screenshot({ path: `${SHOTS}/05-bands-overlay.png` }) + await sigWindow(page).screenshot({ path: `${SHOTS}/06-bands-closeup.png` }) + ctx.assertNoJsErrors() +}) + +/** Zone-axis markers are #fab387 (orange) — a colour the shared green/red + * probes deliberately do not match, so they are counted separately here. */ +async function countOrange(page): Promise { + let total = 0 + for (const frame of page.frames()) { + try { + total += await frame.evaluate(() => { + let n = 0 + for (const c of Array.from(document.querySelectorAll('canvas'))) { + const ctx2 = (c as HTMLCanvasElement).getContext('2d') + if (!ctx2 || !(c as HTMLCanvasElement).width) continue + const d = ctx2.getImageData(0, 0, (c as HTMLCanvasElement).width, + (c as HTMLCanvasElement).height).data + for (let p = 0; p < d.length; p += 4) { + const r = d[p], g = d[p + 1], b = d[p + 2] + if (r > 200 && g > 130 && g < 210 && b > 90 && b < 180) n++ + } + } + return n + }) + } catch { /* detached frame */ } + } + return total +} + +test('the band count and zone-axis toggle change what is drawn', async () => { + const { page } = ctx + const many = await countColorPixels(page, 'green') + expect(many, 'no bands were drawn to begin with').toBeGreaterThan(0) + + await page.getByTestId('ebsd-nbands').fill('3') + await expect.poll(() => countColorPixels(page, 'green'), { + timeout: 30_000, message: 'reducing the band count did not redraw', + }).toBeLessThan(many) + await sigWindow(page).screenshot({ path: `${SHOTS}/07-three-bands.png` }) + const few = await countColorPixels(page, 'green') + + // Raising it again must redraw MORE lines — without this the test would pass + // on a control that only ever removes them. + await page.getByTestId('ebsd-nbands').fill('14') + await expect.poll(() => countColorPixels(page, 'green'), { + timeout: 30_000, message: 'raising the band count did not redraw', + }).toBeGreaterThan(few) + + expect(await countOrange(page), 'zone axes are off but drew anyway').toBe(0) + await page.getByTestId('ebsd-zone-axes').check() + await expect.poll(() => countOrange(page), { + timeout: 30_000, message: 'the zone-axis markers never drew', + }).toBeGreaterThan(0) + await sigWindow(page).screenshot({ path: `${SHOTS}/08-zone-axes.png` }) + ctx.assertNoJsErrors() +}) + +test('indexing the scan opens the IPF orientation map', async () => { + const { page } = ctx + const before = await page.getByTestId('subwindow').count() + + await page.getByTestId('ebsd-tab-Run').click() + await page.getByTestId('ebsd-run').click() + + await expect.poll(() => page.getByTestId('subwindow').count(), { + timeout: 240_000, message: 'indexing never opened the IPF window', + }).toBeGreaterThan(before) + const ipf = page.getByTestId('subwindow') + .filter({ has: page.getByTestId('subwindow-title') + .filter({ hasText: 'Orientation (IPF-Z)' }) }).first() + await expect(ipf, 'the IPF-Z result window never opened') + .toBeVisible({ timeout: 30_000 }) + + // The completion status goes to the app status bar (emit_status), not the + // caret — the caret's own line only tracks its own stages. + await expect(page.locator('body')) + .toContainText(/orientation map complete/i, { timeout: 240_000 }) + await page.screenshot({ path: `${SHOTS}/09-ipf-map.png` }) + await ipf.screenshot({ path: `${SHOTS}/10-ipf-closeup.png` }) + ctx.assertNoJsErrors() +}) + +test('the quality maps ride along as chip views on the IPF window', async () => { + const { page } = ctx + // Pick the IPF window by its TITLE chip, not by hasText — the EBSD caret + // itself contains the word "orientations", so a text filter matches the + // pattern window first. + const ipf = page.getByTestId('subwindow') + .filter({ has: page.getByTestId('subwindow-title') + .filter({ hasText: 'Orientation (IPF-Z)' }) }).first() + await expect(ipf).toBeVisible({ timeout: 15_000 }) + + // An IPF map cannot show where it is WRONG; these are the maps that can. + for (const label of ['IPF-Z', 'NCC', 'Similarity', 'ADP']) { + await expect(ipf.getByTestId(new RegExp(`^view-chip-${label}-`)), + `the ${label} map was not registered as a view`) + .toBeVisible({ timeout: 15_000 }) + } + await ipf.getByTestId(/^view-chip-Similarity-/).click() + await expect.poll(() => countColorPixels(page, 'bright'), { + timeout: 30_000, message: 'the Similarity view painted nothing', + }).toBeGreaterThan(0) + await ipf.screenshot({ path: `${SHOTS}/11-similarity-view.png` }) + ctx.assertNoJsErrors() +}) diff --git a/spyde/actions/README.md b/spyde/actions/README.md index 7047f6f..5bfb394 100644 --- a/spyde/actions/README.md +++ b/spyde/actions/README.md @@ -15,7 +15,7 @@ parity) lives in [`NOTEBOOK_PARITY_PLAN.md`](../../NOTEBOOK_PARITY_PLAN.md). | **View action** | one-shot UI command, no tree change | plain `fn(ctx, …)` | zoom, reset, `tile_views` | | **TransformAction** | signal + params → a **new node in the SAME tree** | `action.TransformAction` | Rebin (`Rebin2DAction`), CZB apply | | **RegionAction** | interactive ROI → a **linked live output plot** | `action.RegionAction` | Virtual Imaging, FFT, Line Profile, Vector VI | -| **Wizard** | staged caret: open → tune → run → commit → close | `wizard.WizardController` + staged handlers | Find Vectors, Orientation, Vector-OM, Strain, CZB | +| **Wizard** | staged caret: open → tune → run → commit → close | `wizard.WizardController` + staged handlers | Find Vectors, Orientation, Vector-OM, EBSD, Strain, CZB | | **Commit** | promote a live/finished result to a **NEW SignalTree** | `commit.commit_result_tree` | strain Commit, VOM result windows | Deciding: does it need an ROI? → RegionAction. Does it produce a new node of diff --git a/spyde/actions/ebsd_action.py b/spyde/actions/ebsd_action.py new file mode 100644 index 0000000..e656396 --- /dev/null +++ b/spyde/actions/ebsd_action.py @@ -0,0 +1,743 @@ +""" +ebsd_action.py — the staged "EBSD Indexing" wizard. + +**This is the 4D-STEM orientation wizard, with bands instead of spots.** The +workflow is deliberately the same one, stage for stage, because it is the same +job: pick a crystal, build a library of simulated patterns, watch the match +under the crosshair until the parameters are right, then run the whole field +and get an IPF map. + + Orientation Mapping (4D-STEM) EBSD Indexing (here) + ------------------------------ ------------------------------------- + 1 Load .cif + voltage 1 Load .cif + voltage + detector PC + 2 Library angle res -> diffsims 2 Library angle step -> simulated + template library pattern dictionary + 3 Refine matched template's 3 Refine matched orientation's Kikuchi + SPOTS on the live DP BANDS on the live pattern + 4 Run whole-field match -> 4 Run whole-field index (+ optional + IPF-Z map refinement) -> IPF-Z map + +Everything downstream of the fit is literally shared code: the result is packed +into a :class:`~spyde.signals.orientation_map.SpyDEOrientationMap`, so the IPF +colouring, the 3-D IPF explorer, the point selector and the direction toggle +built for 4D-STEM all work here untouched. That reuse is the reason Wave 3's +display half is small (RELEASE_0_3_0_PLAN.md, 3.5). + +The compute lives in :mod:`spyde.ebsd` (dictionary indexing, refinement, +preprocessing, band geometry); this module is only the interactive wiring. +""" +from __future__ import annotations + +import logging + +import numpy as np + +from spyde.actions.context import src_plot_tree as _src_plot_tree +from spyde.actions.wizard import WizardController +from spyde.backend.ipc import emit, emit_error, emit_status + +log = logging.getLogger(__name__) + +DEFAULTS = dict( + cif_path="", + space_group=225, + accelerating_voltage=20.0, + min_dspacing=0.7, + pc_x=0.5, pc_y=0.5, pc_z=0.55, + step_deg=4.0, + background="dynamic", + background_sigma=8.0, + n_bands=12, + show_zone_axes=False, + keep=4, + refine=True, + refine_steps=120, +) + +#: Live-preview dictionaries above this many entries make each navigator move +#: cost more than a frame, so the Refine tab stops feeling live. The step size +#: is the user's lever; this only decides when to warn them. +_LIVE_DICT_WARN = 200_000 + + +class EbsdWizard(WizardController): + """Owns the EBSD wizard state: the phase and its reflectors, the simulated + dictionary (resident on the compute device for the live match), the band + overlay on the pattern plot, and the background correction that both the + live match and the whole-field run must apply identically.""" + + key = "ebsd" + + #: Declared parameter schema — one source of truth for every host (the + #: Electron caret mirrors these; see registry._WIZARD_SCHEMAS). + parameters = { + "cif_path": { + "name": "Crystal (.cif)", "type": "file", "default": "", + "extensions": [".cif"], "tab": "Load", + }, + "space_group": { + "name": "Space group", "type": "int", "default": 225, + "min": 1, "max": 230, "tab": "Load", + }, + "accelerating_voltage": { + "name": "Voltage (kV)", "type": "float", "default": 20.0, + "min": 1.0, "max": 50.0, "step": 1.0, "tab": "Load", + }, + "pc_x": { + "name": "PC x", "type": "float", "default": 0.5, + "min": 0.0, "max": 1.0, "step": 0.005, "tab": "Load", + }, + "pc_y": { + "name": "PC y", "type": "float", "default": 0.5, + "min": 0.0, "max": 1.0, "step": 0.005, "tab": "Load", + }, + "pc_z": { + "name": "Detector distance", "type": "float", "default": 0.55, + "min": 0.05, "max": 2.0, "step": 0.005, "tab": "Load", + }, + "step_deg": { + "name": "Angle step (°)", "type": "float", "default": 4.0, + "min": 0.5, "max": 15.0, "step": 0.5, "tab": "Library", + }, + "min_dspacing": { + "name": "Min d-spacing (Å)", "type": "float", "default": 0.7, + "min": 0.3, "max": 3.0, "step": 0.05, "tab": "Library", + }, + "background": { + "name": "Background", "type": "enum", "default": "dynamic", + "choices": ["dynamic", "static", "both", "none"], "tab": "Library", + }, + "background_sigma": { + "name": "Background σ (px)", "type": "float", "default": 8.0, + "min": 1.0, "max": 40.0, "step": 0.5, "tab": "Library", + }, + "n_bands": { + "name": "Bands drawn", "type": "int", "default": 12, + "min": 1, "max": 80, "tab": "Refine", + }, + "show_zone_axes": { + "name": "Zone axes", "type": "bool", "default": False, + "tab": "Refine", + }, + "keep": { + "name": "N best", "type": "int", "default": 4, + "min": 1, "max": 20, "tab": "Run", + }, + "refine": { + "name": "Refine orientations", "type": "bool", "default": True, + "tab": "Run", + }, + "refine_steps": { + "name": "Refine steps", "type": "int", "default": 120, + "min": 10, "max": 600, "step": 10, "tab": "Run", + }, + } + + def __init__(self, session, tree, *, phase, reflectors, indexer, euler, + detector, pc, background, background_sigma, static_ref, + voltage, overlay=None): + super().__init__(session, tree) + self.phase = phase + self.reflectors = reflectors + self.indexer = indexer # SinglePatternIndexer (resident) + self.euler = euler # (D, 3) the dictionary's orientations + self.detector = detector + self.pc = pc + self.background = background + self.background_sigma = float(background_sigma) + self.static_ref = static_ref + self.voltage = float(voltage) + self.overlay = overlay + self.refine: dict = {} + + # ── the background correction, applied identically everywhere ───────────── + @property + def sim_sigma(self): + """The high-pass the SIMULATED patterns must also get. + + Both sides of a cross-correlation have to go through the same filter. + The ``dynamic`` pass is a high-pass and applies to both; the ``static`` + pass subtracts a detector artefact that simulated patterns do not have, + so it applies to the experimental side only. + """ + return (self.background_sigma if self.background in ("dynamic", "both") + else None) + + def correct(self, patterns): + """Apply the wizard's background correction to one pattern or a stack. + + The live match and the whole-field run MUST correct the same way: NCC + is invariant to gain and offset but not to a spatial gradient, so a + pattern corrected differently from the dictionary's expectations scores + differently, and the crosshair preview would stop predicting the map + (:mod:`spyde.ebsd.preprocess`). + """ + if self.background in (None, "none"): + return np.asarray(patterns, float) + from spyde.ebsd.preprocess import remove_background + arr = np.asarray(patterns, float) + single = arr.ndim == 2 + if single: + arr = arr[None] + out = remove_background(arr, method=self.background, + sigma=self.background_sigma, + static_reference=self.static_ref) + return out[0] if single else out + + def remove(self) -> None: + if self._closed: + return + self._closed = True + if self.overlay is not None and hasattr(self.overlay, "remove"): + try: + self.overlay.remove() + except Exception as e: + log.debug("removing EBSD band overlay failed: %s", e) + self.overlay = None + self.indexer = None + if getattr(self.tree, "_ebsd_wizard", None) is self: + self.tree._ebsd_wizard = None + + +def ebsd_indexing(ctx, action_name: str = "EBSD Indexing", **kwargs): + """Parent toolbar action — a no-op; the Electron toolbar opens the staged + EBSD wizard (which drives the ``ebsd_*`` handlers) instead.""" + return None + + +# ───────────────────────────────────────────────────────────────────────────── +# Helpers +# ───────────────────────────────────────────────────────────────────────────── + +def _detector_shape(signal) -> tuple[int, int]: + """``(dy, dx)`` of one pattern.""" + axes = signal.axes_manager.signal_axes + return (int(axes[1].size), int(axes[0].size)) + + +def _stamped(signal, key, default=None): + """A value stamped on synthetic data by ``spyde.data.synthetic``, if any. + + The bundled EBSD scan records the projection centre it was rendered with, + so the wizard can open already pointing at the right geometry instead of + making the user guess a PC before anything can possibly line up. + """ + try: + rec = signal.metadata.get_item("Spyde.synthetic") + except Exception: + return default + if rec is None: + return default + try: + val = rec[key] if not hasattr(rec, "get") else rec.get(key, default) + except Exception: + return default + return default if val is None else val + + +def _resolve_pc(signal, payload) -> tuple[float, float, float]: + """The projection centre: explicit payload values win, else the stamp on + the data, else the centred default.""" + stamp = _stamped(signal, "pc") + base = ([float(v) for v in np.asarray(stamp, float).reshape(3)] + if stamp is not None + else [DEFAULTS["pc_x"], DEFAULTS["pc_y"], DEFAULTS["pc_z"]]) + for i, key in enumerate(("pc_x", "pc_y", "pc_z")): + if payload.get(key) is not None: + base[i] = float(payload[key]) + return tuple(base) + + +def _nav_shape(signal) -> tuple[int, int]: + nav = signal.axes_manager.navigation_shape # (x, y) in hyperspy order + return (int(nav[1]), int(nav[0])) + + +def _read_scan(signal) -> np.ndarray: + """The whole scan as ``(ny, nx, dy, dx)`` float32. + + Deliberately materialised: every step of :mod:`spyde.ebsd` — background + correction, the ADP map, indexing, refinement — is written whole-field, so + there is nothing to stream into. EBSD patterns are small (a 60x60 detector + is 3.6 kB), which is what makes that reasonable where the 4D-STEM memory + rule forbids it. + + It is not free, though: at float32 a 256x256 scan of 60x60 patterns is + 940 MB and a 512x512 one is 3.8 GB. Indexing a scan that large wants the + read chunked over navigation and fed to ``dictionary_index`` a block at a + time — the function already tiles internally and takes a ``stopped_flag``, + so the change is here rather than there. + """ + data = signal.data + if hasattr(data, "compute"): + data = data.compute() + return np.asarray(data, np.float32) + + +def _phase_for(payload, signal): + """The orix Phase to index with: a .cif if one was chosen, else a bare + space-group phase (which still gives the right IPF colour key, just a + generic cubic band set).""" + from orix.crystal_map import Phase + cif = payload.get("cif_path") or "" + if cif: + return Phase.from_cif(cif) + sg = int(payload.get("space_group") or DEFAULTS["space_group"]) + return Phase(name="phase", space_group=sg) + + +def _emit_match(window_id, euler, score) -> None: + """Stream the live single-pattern match to the caret's Refine tab.""" + emit({"type": "ebsd_match", "window_id": window_id, "ok": True, + "phi1": float(np.rad2deg(euler[0])), "Phi": float(np.rad2deg(euler[1])), + "phi2": float(np.rad2deg(euler[2])), "score": float(score)}) + + +# ───────────────────────────────────────────────────────────────────────────── +# Stage 2 — "Build Dictionary" (the analogue of om_generate_library) +# ───────────────────────────────────────────────────────────────────────────── + +def ebsd_build_dictionary(session, plot, payload) -> None: + """Sample orientation space, simulate a pattern for each, and switch on the + LIVE band overlay: from here the matched orientation's Kikuchi bands are + drawn on whatever pattern the navigator is sitting on.""" + src, tree = _src_plot_tree(session, plot) + if src is None or tree is None: + emit_error("Build Dictionary: no active dataset") + return + root = tree.root + if root.axes_manager.signal_dimension != 2 or \ + root.axes_manager.navigation_dimension != 2: + emit_error("EBSD Indexing needs a 2-D scan of 2-D patterns") + return + + step = float(payload.get("step_deg", DEFAULTS["step_deg"])) + voltage = float(payload.get("accelerating_voltage", + DEFAULTS["accelerating_voltage"])) + min_d = float(payload.get("min_dspacing", DEFAULTS["min_dspacing"])) + background = str(payload.get("background", DEFAULTS["background"])) + bg_sigma = float(payload.get("background_sigma", DEFAULTS["background_sigma"])) + n_bands = int(payload.get("n_bands", DEFAULTS["n_bands"])) + zone_axes = bool(payload.get("show_zone_axes", DEFAULTS["show_zone_axes"])) + pc = _resolve_pc(root, payload) + detector = _detector_shape(root) + window_id = getattr(src, "window_id", None) or payload.get("window_id") + + emit_status("EBSD: building the pattern dictionary…") + + def _work(): + try: + from spyde.ebsd.bands import cubic_reflectors, reflectors_from_phase + from spyde.ebsd.indexing import ( + SinglePatternIndexer, sample_orientations, simulate_dictionary, + ) + + phase = _phase_for(payload, root) + try: + reflectors = reflectors_from_phase( + phase, min_dspacing=min_d, voltage_kv=voltage) + except Exception as e: + log.debug("reflectors from phase failed (%s); generic cubic", e) + reflectors = cubic_reflectors() + + # Sample the PHASE's fundamental zone: every distinct orientation + # once. Falls back to an Euler grid if the phase has no usable + # point group (sample_orientations owns that decision). + euler = sample_orientations( + step, point_group=getattr(phase, "point_group", None)) + n = len(euler) + if n > _LIVE_DICT_WARN: + emit_status(f"EBSD: {n:,} orientations at {step}° — the live " + f"preview will lag; raise the angle step to speed " + f"it up") + emit_status(f"EBSD: simulating {n:,} patterns " + f"({len(reflectors)} reflectors)…") + + def _progress(done, total): + if total: + emit_status(f"EBSD: simulating dictionary… " + f"{int(100 * done / total)}%") + + # The dictionary is high-passed as it is simulated, with the same + # sigma the experimental patterns get — both sides of an NCC have + # to come through the same filter (see EbsdWizard.sim_sigma). + sim_sigma = (bg_sigma if background in ("dynamic", "both") else None) + dic = simulate_dictionary(euler, detector, pc, + reflectors=reflectors, + background_sigma=sim_sigma, + progress=_progress) + if dic is None: + return + indexer = SinglePatternIndexer(dic, euler) + + # The static reference for background correction is the scan mean; + # compute it once here, not per preview frame. + static_ref = None + if background in ("static", "both"): + static_ref = _read_scan(root).mean(axis=(0, 1)) + + # A rebuilt dictionary replaces the previous wizard wholesale, so + # its overlay tears down with it rather than stacking a second set + # of lines on the pattern. + old = getattr(tree, "_ebsd_wizard", None) + if old is not None and hasattr(old, "remove"): + try: + old.remove() + except Exception as e: + log.debug("removing prior EBSD wizard failed: %s", e) + + wiz = EbsdWizard( + session, tree, phase=phase, reflectors=reflectors, + indexer=indexer, euler=euler, detector=detector, pc=pc, + background=background, background_sigma=bg_sigma, + static_ref=static_ref, voltage=voltage, + ) + from spyde.actions.ebsd_overlay import attach_ebsd_band_overlay + wiz.overlay = attach_ebsd_band_overlay( + src, root, indexer, reflectors, tree, + detector=detector, pc=pc, correct=wiz.correct, + n_bands=n_bands, show_zone_axes=zone_axes, + on_match=lambda e, s: _emit_match(window_id, e, s), + ) + # Published only once it is COMPLETE: `tree._ebsd_wizard` is what + # every other stage tests for, and a refine or an overlay toggle + # arriving while the overlay was still attaching would find None + # and silently do nothing. + tree._ebsd_wizard = wiz + + emit_status(f"EBSD: dictionary ready ({n:,} orientations) — " + f"move the crosshair to check the bands") + emit({"type": "ebsd_dictionary_ready", "window_id": window_id, + "n_orientations": int(n), "n_reflectors": int(len(reflectors)), + "pc": [float(v) for v in pc]}) + except Exception as e: + emit_error(f"Build Dictionary failed: {e}") + log.exception("Build Dictionary failed") + + from spyde.actions.lifecycle import run_on_worker + run_on_worker(session, _work, name="ebsd-build-dictionary") + + +# ───────────────────────────────────────────────────────────────────────────── +# Stage 3 — "Refine" (the analogue of om_refine) +# ───────────────────────────────────────────────────────────────────────────── + +def ebsd_refine(session, plot, payload) -> None: + """Live-update the band-overlay knobs — how many bands to draw, whether to + mark the zone axes, and the projection centre — and redraw at the current + crosshair position. + + The PC belongs here rather than only on the Load tab because it is the one + parameter you cannot set from first principles: you nudge it until the + drawn lines sit on the bands, which is only possible with the overlay live + in front of you. + """ + src, tree = _src_plot_tree(session, plot) + wiz = getattr(tree, "_ebsd_wizard", None) if tree is not None else None + if wiz is None or wiz.overlay is None: + return + params: dict = {} + if payload.get("n_bands") is not None: + params["n_bands"] = int(payload["n_bands"]) + if payload.get("show_zone_axes") is not None: + params["show_zone_axes"] = bool(payload["show_zone_axes"]) + if any(payload.get(k) is not None for k in ("pc_x", "pc_y", "pc_z")): + pc = _resolve_pc(tree.root, {**{"pc_x": wiz.pc[0], "pc_y": wiz.pc[1], + "pc_z": wiz.pc[2]}, **payload}) + wiz.pc = pc + params["pc"] = pc + + def _work(): + try: + wiz.overlay.set_refine_params(**params) + wiz.refine = dict(payload) + except Exception as e: + log.debug("ebsd_refine failed: %s", e) + + from spyde.actions.lifecycle import run_on_worker + run_on_worker(session, _work, name="ebsd-refine") + + +# ───────────────────────────────────────────────────────────────────────────── +# Stage 4 — "Index Map" (the analogue of om_run) +# ───────────────────────────────────────────────────────────────────────────── + +def ebsd_run(session, plot, payload) -> None: + """Index every pattern in the scan against the built dictionary, optionally + refine each orientation off its dictionary entry, and open the IPF-Z map.""" + src, tree = _src_plot_tree(session, plot) + wiz = getattr(tree, "_ebsd_wizard", None) if tree is not None else None + if wiz is None or wiz.indexer is None: + emit_error("Index Map: build the dictionary first") + return + keep = int(payload.get("keep", DEFAULTS["keep"])) + do_refine = bool(payload.get("refine", DEFAULTS["refine"])) + steps = int(payload.get("refine_steps", DEFAULTS["refine_steps"])) + emit_status("EBSD: indexing the scan…") + + # torch's CUDA autograd backward segfaults the first time it runs on a + # thread whose engine was never initialised; the refinement runs on the + # worker, so warm the engine here on the dispatch thread first (the same + # fix vector-orientation mapping needed — CLAUDE.md, GPU Computing). + if do_refine: + try: + from spyde.actions.vector_orientation_gpu import warmup_autograd + warmup_autograd() + except Exception as e: + log.debug("CUDA autograd warmup failed: %s", e) + + def _work(): + try: + _run_indexing(session, tree, wiz, keep=keep, refine=do_refine, + steps=steps) + except Exception as e: + emit_error(f"Index Map failed: {e}") + log.exception("Index Map failed") + + from spyde.actions.lifecycle import run_on_worker + run_on_worker(session, _work, name="ebsd-run") + + +def _run_indexing(session, tree, wiz, *, keep, refine, steps): + """Whole-field index (+ optional refine) → the IPF window. Synchronous; + call from a worker.""" + from spyde.ebsd.crystal_map import orientation_similarity_map + from spyde.ebsd.indexing import dictionary_index + from spyde.ebsd.preprocess import average_dot_product_map + + root = tree.root + ny, nx = _nav_shape(root) + scan = wiz.correct(_read_scan(root)) + dict_euler = wiz.euler + + ipf_tree = _open_ipf_window(session, root, ny, nx) + # Closing EITHER the source or the result window stops the run — an index + # of a large scan is long enough that "close it and it keeps going" reads + # as a hang. + stopped = [False] + trees = {id(tree): tree, id(ipf_tree): ipf_tree} + for t in trees.values(): + if t is not None and hasattr(t, "register_cancel"): + t.register_cancel(flag=stopped) + + try: + painter = _ProgressivePainter(session, ipf_tree, wiz, (ny, nx), dict_euler) + # Cap the pattern tile so the map fills in as it goes: left alone the + # tiling puts the whole scan in one chunk and nothing paints until the + # end (see dictionary_index's pattern_chunk). + chunk = max(1, (ny * nx) // 16) + result = dictionary_index( + scan, wiz.indexer, keep=keep, pattern_chunk=chunk, + on_chunk=painter, stopped_flag=stopped, + progress=lambda d, t: emit_status( + f"EBSD: indexing… {int(100 * d / max(t, 1))}%"), + ) + if result is None: + if not getattr(tree, "_spyde_closed", False) and not stopped[0]: + emit_error("EBSD: indexing returned no result") + return None + + euler = dict_euler[result.best].reshape(ny, nx, 3) + score = result.best_score.reshape(ny, nx) + + # Refinement is the tail of the SAME run, so it stays inside the + # cancellation registration — unregistering after the index would let + # a window closed during refinement carry on to paint into it. + if refine and not stopped[0]: + emit_status(f"EBSD: refining {ny * nx:,} orientations…") + from spyde.ebsd.refine import refine_orientations + ref = refine_orientations( + scan, euler, detector=wiz.detector, pc=wiz.pc, + reflectors=wiz.reflectors, background_sigma=wiz.sim_sigma, + steps=steps, + progress=lambda d, t: emit_status( + f"EBSD: refining… {int(100 * d / max(t, 1))}%")) + euler = ref.euler_map((ny, nx)) + score = ref.score.reshape(ny, nx) + log.debug("refinement improved %d/%d orientations", + int(ref.improved.sum()), ref.improved.size) + + if stopped[0]: + return None + + om = _orientation_map(euler, score, wiz, root) + osm = (orientation_similarity_map(result.indices, (ny, nx)) + if keep > 1 else None) + try: + adp = average_dot_product_map(scan) + except Exception as e: + log.debug("ADP map failed: %s", e) + adp = None + finally: + for t in trees.values(): + if t is not None and hasattr(t, "unregister_cancel"): + t.unregister_cancel(flag=stopped) + + tree.orientation_map = om + + _finalize_ipf_window(session, ipf_tree, om, score=score, osm=osm, adp=adp) + emit_status(f"EBSD orientation map complete (mean NCC " + f"{float(np.nanmean(score)):.3f})") + return om + + +class _ProgressivePainter: + """Paint the IPF map as indexing completes each slice of patterns. + + A dictionary index reports per chunk of PATTERNS, which in scan order is a + contiguous run of positions — so the map fills top to bottom and you can + see the grain structure emerge instead of watching a blank window. Costs + one IPF colouring per chunk over the positions done so far. + """ + + def __init__(self, session, ipf_tree, wiz, nav_shape, dict_euler): + from orix.plot import IPFColorKeyTSL + + self.session = session + self.tree = ipf_tree + self.wiz = wiz + self.nav_shape = nav_shape + self.dict_euler = dict_euler + ny, nx = nav_shape + self.rgb = np.zeros((ny, nx, 3), np.uint8) + self.plot = next(iter(getattr(ipf_tree, "signal_plots", []) or []), None) + # The SAME key SpyDEOrientationMap.ipf_color_map uses (the point + # group's LAUE group). Built from anything else, the fill-in would be + # coloured differently from the final map and the whole thing would + # visibly change hue the moment the run finished. + self._pg = wiz.phase.point_group + self._key = IPFColorKeyTSL(self._pg.laue) + + def __call__(self, lo, hi, indices, scores) -> None: + if self.plot is None: + return + try: + from orix.quaternion import Orientation, Rotation + rot = Rotation.from_euler(self.dict_euler[indices[:, 0]]) + rgb = self._key.orientation2color(Orientation(rot, symmetry=self._pg)) + self.rgb.reshape(-1, 3)[lo:hi] = np.clip( + np.asarray(rgb) * 255.0, 0, 255).astype(np.uint8) + except Exception as e: + log.debug("progressive IPF colouring failed: %s", e) + return + # Painting touches a live plot, so it has to happen on the main thread. + frame = self.rgb.copy() + + def _paint(): + try: + self.plot.needs_auto_level = True + self.plot.set_data(frame) + except Exception as e: + log.debug("progressive IPF paint failed: %s", e) + + dispatch = getattr(self.session, "_dispatch_to_main", None) + if dispatch is not None: + dispatch(_paint) + else: + _paint() + + +def _orientation_map(euler, score, wiz, src): + """Pack the indexed field into the SAME result object 4D-STEM orientation + mapping produces, so every existing IPF view works on it unchanged.""" + from orix.quaternion import Rotation + from spyde.signals.orientation_map import SpyDEOrientationMap, phase_to_dict + + ny, nx = euler.shape[:2] + quats = np.asarray(Rotation.from_euler(euler.reshape(-1, 3)).data, + np.float32).reshape(ny, nx, 1, 4) + corr = np.asarray(score, np.float32).reshape(ny, nx, 1) + return SpyDEOrientationMap( + quats=quats, + corr=corr, + phase_idx=np.zeros((ny, nx, 1), np.int16), + mirror=np.ones((ny, nx, 1), np.int8), + phases=[phase_to_dict(wiz.phase)], + nav_axes=list(src.axes_manager.navigation_axes), + params={"action": "EBSD Indexing", "pc": list(wiz.pc), + "detector": list(wiz.detector), + "voltage_kv": wiz.voltage, + "n_dictionary": int(len(wiz.euler)), + "background": wiz.background}, + ) + + +def _open_ipf_window(session, src, ny, nx): + """Open the IPF-Z window blank up front so the map has somewhere to fill.""" + from spyde.actions.commit import open_result_tree + base = src.metadata.get_item("General.title", "Signal") + return open_result_tree( + session, title=f"{base} — Orientation (IPF-Z)", + data=np.zeros((ny, nx), dtype=np.float32), + provenance={"action": "EBSD Indexing", "source_title": base}, + ) + + +def _finalize_ipf_window(session, tree, om, *, score=None, osm=None, + adp=None) -> None: + """Paint the final IPF-Z map and attach the shared orientation views.""" + tree.orientation_map = om + ipf = om.ipf_color_map(direction="z") # (ny, nx, 3) uint8 + for sp in list(getattr(tree, "signal_plots", [])): + try: + sp.needs_auto_level = True + sp.set_data(ipf) + except Exception as e: + log.debug("painting the EBSD IPF map failed: %s", e) + try: + from spyde.actions.ipf_view import attach_ipf_3d, attach_ipf_point_selector + attach_ipf_3d(tree, om, direction="z") + attach_ipf_point_selector(tree, om, "z") + except Exception as e: + log.debug("attaching the 3-D IPF explorer failed: %s", e) + _attach_quality_views(session, tree, score=score, osm=osm, adp=adp) + + +def _attach_quality_views(session, tree, *, score=None, osm=None, + adp=None) -> None: + """Add the quality maps as chip-selectable views on the IPF window. + + An IPF map alone cannot tell you where it is WRONG — every position gets a + colour whether or not the match meant anything. The NCC map shows how well + each pattern matched something; the orientation-similarity map is the one + that exposes confidently-wrong indexing (a position that scored well but + disagrees with its whole neighbourhood); the ADP map shows where the + patterns were worth indexing at all. + + Same shape as ``commit.commit_result_tree``'s views, but that door needs + the data up front and this window was opened blank to fill progressively — + so the views are added here, once the run is done. + """ + import hyperspy.api as hs + from spyde.actions.views import emit_view_figure, register_views + + views = [(lbl, np.nan_to_num(np.asarray(m, np.float32))) + for lbl, m in (("NCC", score), ("Similarity", osm), ("ADP", adp)) + if m is not None] + if not views: + return + title = tree.root.metadata.get_item("General.title", "Orientation") + for lbl, m in views: + try: + child = hs.signals.Signal2D(m.copy()) + child.metadata.General.title = f"{title} {lbl}" + tree.add_node(tree.root, child, lbl) + tree.update_plot_states(child) + except Exception as e: + log.debug("adding EBSD %r view node failed: %s", lbl, e) + try: + session._reemit_signal_tree(tree) + except Exception as e: + log.debug("re-emitting the EBSD result tree failed: %s", e) + + sp = next(iter(getattr(tree, "signal_plots", []) or []), None) + wid = getattr(sp, "window_id", None) if sp is not None else None + if wid is None: + return + try: + sp.set_view_tag("IPF-Z", "2d") + except Exception as e: + log.debug("tagging the EBSD IPF view failed: %s", e) + register_views(wid, views) + for lbl, m in views: + emit_view_figure(wid, m, lbl, kind="2d") diff --git a/spyde/actions/ebsd_overlay.py b/spyde/actions/ebsd_overlay.py new file mode 100644 index 0000000..a1435e2 --- /dev/null +++ b/spyde/actions/ebsd_overlay.py @@ -0,0 +1,186 @@ +""" +ebsd_overlay.py — the live Kikuchi band overlay on the EBSD pattern. + +The EBSD counterpart of ``vector_overlay.OrientationOverlay``. Both do the same +thing — index the pattern under the crosshair and draw the matched orientation +back on top of it — and both hang off the same ``_DPOverlay`` chrome, so the +navigator wiring, the seeding, the show/hide and the teardown are shared. The +difference is entirely in what gets drawn: + +* a 4D-STEM template match is a set of SPOTS, so that overlay draws circles; +* an EBSD match is a set of BANDS, and a band's centre is a straight LINE on a + flat detector (:mod:`spyde.ebsd.bands`), so this one draws line segments. + +Both the geometry and the reflector set come from the wizard's dictionary, so +the lines are by construction the centres of the bands the dictionary entry was +rendered with — a line beside a band means the ORIENTATION is wrong, which is +the whole point of showing it. +""" +from __future__ import annotations + +import logging +import threading + +import numpy as np + +from spyde.actions.vector_overlay import _DPOverlay + +log = logging.getLogger(__name__) + +BAND_COLOR = "#30ff60" # same green as the matched-template spot overlay +ZONE_COLOR = "#fab387" # orange, the strain overlay's accent + + +class EbsdBandOverlay(_DPOverlay): + """Live Kikuchi bands for the best-matching orientation, on the EBSD plot. + + Per navigator move: correct the pattern the way the dictionary expects, + match it against the resident dictionary (one mat-vec), project that + orientation's bands onto the detector, push the segments. + + Runs on the overlay engine's worker thread (``_overlay_mode = "thread"``): + the match is milliseconds but it is not free, and the navigator's update + path is serialised — computing here would gate the pattern display at the + overlay's rate (see ``live_overlay`` for why). + """ + + _overlay_mode = "thread" + name = "ebsd_bands" + + def __init__(self, dp_plot, signal, indexer, reflectors, *, + detector, pc, correct=None, n_bands: int = 12, + show_zone_axes: bool = False, linewidth: float = 1.2, + color: str = BAND_COLOR, on_match=None): + self.dp_plot = dp_plot + self.signal = signal + self.indexer = indexer + self.reflectors = reflectors + self.detector = (int(detector[0]), int(detector[1])) + self.pc = tuple(float(v) for v in pc) + self.correct = correct + self.n_bands = int(n_bands) + self.show_zone_axes = bool(show_zone_axes) + self.linewidth = float(linewidth) + self._color = color + self.on_match = on_match # callback(euler, score) for the caret + self._mg = None + self._mg_za = None + self._selectors: list = [] + self._last_iyix = (0, 0) + self._radius_px = 3.0 + # The matcher touches one resident torch tensor; a nav move (engine + # thread) and a Refine slider (dispatch thread) can both land in here. + self._match_lock = threading.Lock() + + # ── parameters (the Refine tab) ─────────────────────────────────────────── + def set_refine_params(self, **params) -> None: + """Live-update the overlay knobs and redraw at the CURRENT crosshair + position — the wizard's Refine tab.""" + if params.get("n_bands") is not None: + self.n_bands = max(1, int(params["n_bands"])) + if params.get("show_zone_axes") is not None: + self.show_zone_axes = bool(params["show_zone_axes"]) + if params.get("linewidth") is not None: + self.linewidth = max(0.2, float(params["linewidth"])) + if params.get("pc") is not None: + self.pc = tuple(float(v) for v in params["pc"]) + if self._engine is not None: + self._engine.request(*self._last_iyix) + else: # not attached to a figure + self._render_payload(self._offsets_for(*self._last_iyix)) + + # ── drawing primitives ──────────────────────────────────────────────────── + def _empty(self): + return (np.zeros((0, 2, 2), np.float32), np.zeros((0, 2), np.float32)) + + def _make_markers(self, plot2d): + """Lines for the bands, plus a circle group for the zone axes. Both are + created up front and simply fed empty arrays when off — adding a marker + group later, from the overlay worker thread, would race the figure.""" + self._mg_za = plot2d.add_circles( + np.zeros((0, 2), np.float32), name=f"{self.name}_zone_axes", + radius=3.0, edgecolors=ZONE_COLOR, facecolors=None, + linewidths=1.2, alpha=1.0, transform="data", + ) + return plot2d.add_lines( + np.zeros((0, 2, 2), np.float32), name=self.name, + edgecolors=self._color, linewidths=self.linewidth, transform="data", + ) + + def _push(self, payload) -> None: + """Push a ``(segments, zone_axis_points)`` payload. A bare array means + "clear" — that is what ``_DPOverlay.set_visible(False)`` sends.""" + segs, za = payload if isinstance(payload, tuple) else self._empty() + if self._mg is not None: + try: + self._mg.set(segments=segs, linewidths=self.linewidth) + except Exception as e: + log.debug("pushing EBSD band segments failed: %s", e) + if self._mg_za is not None: + try: + self._mg_za.set(offsets=za) + except Exception as e: + log.debug("pushing EBSD zone axes failed: %s", e) + + def remove(self): + super().remove() + if self._mg_za is not None: + try: + self._mg_za.remove() + except Exception as e: + log.debug("removing EBSD zone-axis markers failed: %s", e) + self._mg_za = None + + # ── the per-position compute ────────────────────────────────────────────── + def _frame(self, iy, ix) -> np.ndarray: + frame = self.signal.data[iy, ix] + if hasattr(frame, "compute"): # lazy/dask: one small pattern only + frame = frame.compute() + return np.asarray(frame, dtype=float) + + def _offsets_for(self, iy, ix): + from spyde.ebsd.bands import band_lines, zone_axis_points + try: + frame = self._frame(iy, ix) + if self.correct is not None: + frame = self.correct(frame) + with self._match_lock: + euler, score = self.indexer.best(frame) + except Exception as e: + # Don't swallow blind: a failed match must be distinguishable from + # an orientation whose bands simply miss the detector. + log.debug("[overlay:ebsd] indexing FAILED nav=(%s,%s): %r", iy, ix, e) + return self._empty() + + try: + segs, _w = band_lines(euler, self.reflectors, self.detector, self.pc, + max_bands=self.n_bands) + za = (zone_axis_points(euler, self.reflectors.brightest(self.n_bands), + self.detector, self.pc) + if self.show_zone_axes else np.zeros((0, 2), np.float32)) + except Exception as e: + log.debug("[overlay:ebsd] projecting bands FAILED nav=(%s,%s): %r", + iy, ix, e) + return self._empty() + + log.debug("[overlay:ebsd] nav=(%s,%s) -> %d bands, ncc=%.4f", + iy, ix, len(segs), score) + if self.on_match is not None: + try: + self.on_match(euler, score) + except Exception as e: + log.debug("[overlay:ebsd] on_match callback failed: %s", e) + return segs, za + + +def attach_ebsd_band_overlay(dp_plot, signal, indexer, reflectors, tree, *, + detector, pc, correct=None, n_bands: int = 12, + show_zone_axes: bool = False, + on_match=None) -> EbsdBandOverlay: + """Add the live band overlay to ``dp_plot``, wired to ``tree``'s navigator + selectors. Returns the :class:`EbsdBandOverlay`.""" + return EbsdBandOverlay( + dp_plot, signal, indexer, reflectors, detector=detector, pc=pc, + correct=correct, n_bands=n_bands, show_zone_axes=show_zone_axes, + on_match=on_match, + ).attach(tree) diff --git a/spyde/actions/registry.py b/spyde/actions/registry.py index c205fce..c6c1dd2 100644 --- a/spyde/actions/registry.py +++ b/spyde/actions/registry.py @@ -49,6 +49,9 @@ "om_generate_library": "spyde.actions.orientation_action.om_generate_library", "om_refine": "spyde.actions.orientation_action.om_refine", "om_run": "spyde.actions.orientation_action.om_run", + "ebsd_build_dictionary": "spyde.actions.ebsd_action.ebsd_build_dictionary", + "ebsd_refine": "spyde.actions.ebsd_action.ebsd_refine", + "ebsd_run": "spyde.actions.ebsd_action.ebsd_run", "fv_open": "spyde.actions.find_vectors_action.fv_open", "fv_tune": "spyde.actions.find_vectors_action.fv_tune", "fv_run": "spyde.actions.find_vectors_action.fv_run", @@ -205,6 +208,7 @@ def register_staged(name: str, dotted_path: str) -> None: # `parameters`) or a dict. "strain": ("spyde.actions.strain_action", "StrainController"), "vom": ("spyde.actions.vector_orientation_om", "VomWizard"), + "ebsd": ("spyde.actions.ebsd_action", "EbsdWizard"), "czb": ("spyde.actions.center_zero_beam", "PARAMETERS"), # YAML-declared (resolved from spyde.TOOLBAR_ACTIONS): "fv": ("__yaml__", "Find Diffraction Vectors"), diff --git a/spyde/actions/vector_overlay.py b/spyde/actions/vector_overlay.py index 9496d7b..40fffcf 100644 --- a/spyde/actions/vector_overlay.py +++ b/spyde/actions/vector_overlay.py @@ -152,6 +152,19 @@ def _offsets_for(self, iy, ix) -> np.ndarray: # pragma: no cover def _marker_kwargs(self, offsets) -> dict: return {"offsets": offsets} + def _empty(self) -> np.ndarray: + """The payload that draws nothing — what ``set_visible(False)`` pushes.""" + return np.zeros((0, 2), dtype=np.float32) + + def _make_markers(self, plot2d): + """Create this overlay's anyplotlib marker group(s). Circles by default; + override for a different primitive (the EBSD band overlay draws lines).""" + return plot2d.add_circles( + np.zeros((0, 2), dtype=np.float32), name=self.name, + radius=float(self._radius_px), edgecolors=self._color, facecolors=None, + linewidths=1.5, alpha=1.0, transform="data", + ) + def attach(self, tree): plot2d = getattr(self.dp_plot, "_plot2d", None) if plot2d is None: @@ -161,11 +174,7 @@ def attach(self, tree): "yet (figure iframe not loaded?)", getattr(self, "name", type(self).__name__)) return self - self._mg = plot2d.add_circles( - np.zeros((0, 2), dtype=np.float32), name=self.name, - radius=float(self._radius_px), edgecolors=self._color, facecolors=None, - linewidths=1.5, alpha=1.0, transform="data", - ) + self._mg = self._make_markers(plot2d) self._engine = self._make_engine(tree) self._selectors = _navigator_selectors_for(tree, self.dp_plot) seeded = False @@ -256,7 +265,7 @@ def _push(self, offsets): def set_visible(self, visible: bool) -> None: self._hidden = not bool(visible) if self._hidden: - self._push(np.zeros((0, 2), dtype=np.float32)) + self._push(self._empty()) elif self._engine is not None: self._engine.request(*self._last_iyix) # recompute current frame else: diff --git a/spyde/backend/_session_actions.py b/spyde/backend/_session_actions.py index 7b79f07..5264605 100644 --- a/spyde/backend/_session_actions.py +++ b/spyde/backend/_session_actions.py @@ -381,6 +381,10 @@ def _set_overlay(self, plot, name: str, visible: bool) -> None: wiz = getattr(tree, "_vom_wizard", None) if wiz is not None: overlays.append(getattr(wiz, "overlay", None)) + elif name == "EBSD Indexing": + wiz = getattr(tree, "_ebsd_wizard", None) + if wiz is not None: + overlays.append(getattr(wiz, "overlay", None)) for ov in overlays: if ov is not None and hasattr(ov, "set_visible"): try: diff --git a/spyde/data/synthetic.py b/spyde/data/synthetic.py index fbd0919..13fff88 100644 --- a/spyde/data/synthetic.py +++ b/spyde/data/synthetic.py @@ -301,100 +301,29 @@ def eds_si(nav=(16, 16), n_channels: int = 2048, *, e_stop: float = 20.0, # EBSD # --------------------------------------------------------------------------- -def _euler_to_matrix(phi1, Phi, phi2) -> np.ndarray: - """Bunge ZXZ Euler angles -> rotation matrices, batched over the leading - axes. Returns ``(..., 3, 3)``.""" - c1, s1 = np.cos(phi1), np.sin(phi1) - c, s = np.cos(Phi), np.sin(Phi) - c2, s2 = np.cos(phi2), np.sin(phi2) - m = np.empty(np.shape(phi1) + (3, 3), float) - m[..., 0, 0] = c1 * c2 - s1 * s2 * c - m[..., 0, 1] = s1 * c2 + c1 * s2 * c - m[..., 0, 2] = s2 * s - m[..., 1, 0] = -c1 * s2 - s1 * c2 * c - m[..., 1, 1] = -s1 * s2 + c1 * c2 * c - m[..., 1, 2] = c2 * s - m[..., 2, 0] = s1 * s - m[..., 2, 1] = -c1 * s - m[..., 2, 2] = c - return m +# The projection, the band set and the pattern renderer live in +# `spyde.ebsd.bands` — the ONE place that geometry is written down. The +# generator, the dictionary simulator and the live band overlay all import it +# from there, so an indexing test cannot pass by having two sides make the same +# mistake in the same place, and the overlay's lines cannot drift away from the +# bands they are drawn on. Re-exported here under the names this module has +# always used. (`spyde.ebsd.bands` is numpy-only and imports nothing from +# spyde, so this cannot cycle.) +from spyde.ebsd.bands import ( # noqa: E402 + cubic_reflectors as _cubic_reflectors, + detector_directions, + euler_to_matrix as _euler_to_matrix, + normals_to_sample as _normals_to_sample, + simulate_patterns, +) def _cubic_plane_normals() -> tuple[np.ndarray, np.ndarray]: """{111}, {200} and {220} plane normals for a cubic crystal, with rough structure-factor weights. One normal per Friedel pair (a band and its inverse are the same band).""" - fams = [((1, 1, 1), 1.00), ((2, 0, 0), 0.70), ((2, 2, 0), 0.45)] - normals, weights = [], [] - for hkl, w in fams: - seen = set() - h, k, l = hkl - for perm in {(h, k, l), (h, l, k), (k, h, l), (k, l, h), (l, h, k), (l, k, h)}: - for sx in (1, -1): - for sy in (1, -1): - for sz in (1, -1): - v = (perm[0] * sx, perm[1] * sy, perm[2] * sz) - if v == (0, 0, 0) or tuple(-x for x in v) in seen: - continue - seen.add(v) - for v in seen: - n = np.array(v, float) - normals.append(n / np.linalg.norm(n)) - # Band width goes as 1/|g|: bigger d-spacing -> wider band. - weights.append(w / np.linalg.norm(np.array(v, float))) - return np.array(normals), np.array(weights) - - -def detector_directions(detector=(60, 60), pc=(0.5, 0.5, 0.55)) -> np.ndarray: - """Unit vectors from the sample to each detector pixel (gnomonic). - - ``(dy, dx, 3)``. Shared by the pattern generator and by dictionary - simulation so an indexing test cannot pass by having both sides make the - same geometry mistake in the same place — they use the identical geometry - on purpose, and the thing under test is the MATCHING, not the projection. - """ - pcx, pcy, L = float(pc[0]), float(pc[1]), float(pc[2]) - dy, dx = int(detector[0]), int(detector[1]) - gy, gx = np.mgrid[0:dy, 0:dx].astype(float) - rx = (gx + 0.5) / dx - pcx - ry = pcy - (gy + 0.5) / dy # detector y is flipped - r = np.stack([rx, ry, np.full_like(rx, L)], -1) - return r / np.linalg.norm(r, axis=-1, keepdims=True) - - -def simulate_patterns(euler, detector=(60, 60), pc=(0.5, 0.5, 0.55), - *, background: bool = False) -> np.ndarray: - """Kikuchi patterns for a list of orientations -> ``(N, dy, dx)`` float32. - - *euler* is ``(N, 3)`` Bunge angles in radians. This is what builds an - indexing DICTIONARY: sample orientation space, simulate each, match against - it (#71). - - ``background=False`` by default because a dictionary is matched by - normalised cross-correlation, which is invariant to the smooth gradient a - real detector adds — and the experimental side has background removal - applied before matching anyway (#70). - """ - euler = np.atleast_2d(np.asarray(euler, float)) - rot = _euler_to_matrix(euler[:, 0], euler[:, 1], euler[:, 2]) # (N, 3, 3) - r = detector_directions(detector, pc) - dy, dx = r.shape[:2] - flat_r = r.reshape(-1, 3) - - normals, weights = _cubic_plane_normals() - widths = 0.055 * (weights / weights.max()) + 0.012 - - out = np.empty((len(euler), dy, dx), np.float32) - for i in range(len(euler)): - n_rot = normals @ rot[i].T - d = flat_r @ n_rot.T - band = np.exp(-0.5 * (d / widths) ** 2) * weights - out[i] = band.sum(1).reshape(dy, dx).astype(np.float32) - if background: - gy, gx = np.mgrid[0:dy, 0:dx].astype(np.float32) - out = out / max(out.max(), 1e-9) + ( - 0.35 + 0.30 * ((gy / dy) * 0.6 + (gx / dx) * 0.4)) - return out + refl = _cubic_reflectors() + return refl.normals, refl.weights def ebsd_patterns(nav=(16, 16), detector=(60, 60), *, pc=(0.5, 0.5, 0.55), @@ -436,24 +365,19 @@ def ebsd_patterns(nav=(16, 16), detector=(60, 60), *, pc=(0.5, 0.5, 0.55), rot = _euler_to_matrix(phi1, Phi, phi2) # (ny, nx, 3, 3) # --- detector directions (gnomonic) ----------------------------------- - pcx, pcy, L = float(pc[0]), float(pc[1]), float(pc[2]) gy, gx = np.mgrid[0:dy, 0:dx].astype(float) - rx = (gx + 0.5) / dx - pcx - ry = pcy - (gy + 0.5) / dy # detector y is flipped - r = np.stack([rx, ry, np.full_like(rx, L)], -1) - r /= np.linalg.norm(r, axis=-1, keepdims=True) # (dy, dx, 3) + r = detector_directions(detector, pc) # (dy, dx, 3) - normals, weights = _cubic_plane_normals() - # Band half-width in units of cos(angle to the normal); scaled by |g| so - # low-index planes give the wide bands, as in a real pattern. - widths = 0.055 * (weights / weights.max()) + 0.012 + refl = _cubic_reflectors() + normals, weights, widths = refl.normals, refl.weights, refl.widths data = np.empty((ny, nx, dy, dx), dtype=np.float32) flat_r = r.reshape(-1, 3) for iy in range(ny): for ix in range(nx): - # Rotate the plane normals into the sample frame for this pixel. - n_rot = normals @ rot[iy, ix].T # (B, 3) + # Rotate the plane normals into the sample frame for this pixel + # (bands.normals_to_sample owns the direction and why it matters). + n_rot = _normals_to_sample(normals, rot[iy, ix]) # (B, 3) d = flat_r @ n_rot.T # (dy*dx, B) band = np.exp(-0.5 * (d / widths) ** 2) * weights data[iy, ix] = band.sum(1).reshape(dy, dx).astype(np.float32) @@ -475,6 +399,6 @@ def ebsd_patterns(nav=(16, 16), detector=(60, 60), *, pc=(0.5, 0.5, 0.55), s.metadata.General.title = "Synthetic EBSD" _try_signal_type(s, "EBSD") _stamp(s, kind="ebsd", euler=np.stack([phi1, Phi, phi2], -1), - pc=np.array([pcx, pcy, L]), n_bands=int(len(normals)), + pc=np.asarray(pc, float), n_bands=int(len(normals)), grain2_mask=grain2) return s diff --git a/spyde/drawing/toolbars/icons/ebsd.svg b/spyde/drawing/toolbars/icons/ebsd.svg new file mode 100644 index 0000000..bdeeb84 --- /dev/null +++ b/spyde/drawing/toolbars/icons/ebsd.svg @@ -0,0 +1,7 @@ + + + + + + + diff --git a/spyde/ebsd/__init__.py b/spyde/ebsd/__init__.py index 57a9761..f8e0016 100644 --- a/spyde/ebsd/__init__.py +++ b/spyde/ebsd/__init__.py @@ -17,10 +17,21 @@ """ from __future__ import annotations +from spyde.ebsd.bands import ( + Reflectors, + band_lines, + cubic_reflectors, + detector_directions, + reflectors_from_phase, + simulate_patterns, + zone_axis_points, +) from spyde.ebsd.indexing import ( IndexingResult, + SinglePatternIndexer, dictionary_index, sample_orientations, + simulate_dictionary, ) from spyde.ebsd.preprocess import average_dot_product_map, remove_background from spyde.ebsd.crystal_map import ( @@ -31,8 +42,12 @@ ) from spyde.ebsd.refine import RefinementResult, refine_orientations -__all__ = ["dictionary_index", "sample_orientations", "IndexingResult", +__all__ = ["dictionary_index", "sample_orientations", "simulate_dictionary", + "IndexingResult", "SinglePatternIndexer", "remove_background", "average_dot_product_map", "refine_orientations", "RefinementResult", "to_crystal_map", "ipf_colors", "orientation_similarity_map", - "merge_phases"] + "merge_phases", + "Reflectors", "cubic_reflectors", "reflectors_from_phase", + "band_lines", "zone_axis_points", "detector_directions", + "simulate_patterns"] diff --git a/spyde/ebsd/bands.py b/spyde/ebsd/bands.py new file mode 100644 index 0000000..7e29e7d --- /dev/null +++ b/spyde/ebsd/bands.py @@ -0,0 +1,443 @@ +"""bands.py — Kikuchi band geometry: the projection shared by every consumer. + +Three parts of Wave 3 need the same projection and must not disagree about it: + +* the synthetic data generator (:mod:`spyde.data.synthetic`), which renders + patterns from known orientations, +* the dictionary simulator (:class:`spyde.ebsd.refine.BandSimulator`), which + renders the entries an experimental pattern is matched against, +* the **live overlay**, which draws the indexed orientation's bands back on top + of the experimental pattern. + +If any two of those drift the overlay silently lies: the lines sit beside the +bands and there is nothing to tell you whether the indexing was wrong or only +the drawing. So the geometry lives here, once, and the other three import it. + +The geometry itself is the flat-detector gnomonic projection. A detector pixel +sees the crystal along a direction ``r``; a band appears where ``r`` is nearly +perpendicular to a plane normal ``n``. That gives two things: + +* **the band centre**, ``r·n = 0``. Because ``r`` is linear in the pixel + coordinates and the ``|r|`` normalisation is a positive scalar, this is + exactly a straight LINE on the detector — which is why a band overlay is + drawn with line segments and not with a sampled curve. +* **the band edges**, ``r̂·n = ±w``, which are conics. They are not drawn: the + centre line is what tells you whether the orientation is right, and a real + pattern's band edges are diffuse anyway. This is also what kikuchipy's + geometrical simulation draws. + +Nothing here needs torch or the ``ebsd`` extra — it is numpy and, only for +:func:`reflectors_from_phase`, diffsims (which SpyDE already depends on). +""" +from __future__ import annotations + +import logging +from dataclasses import dataclass + +import numpy as np + +log = logging.getLogger(__name__) + +# Electron wavelength in Å at an accelerating voltage in kV (relativistic). +_H, _M0, _E, _C = 6.62607015e-34, 9.1093837015e-31, 1.602176634e-19, 299792458.0 + + +def wavelength(voltage_kv: float) -> float: + """Relativistic electron wavelength in Å. EBSD runs at 10–30 kV, where the + non-relativistic form is already ~1% off, so the correction is kept.""" + v = float(voltage_kv) * 1e3 + lam = _H / np.sqrt(2 * _M0 * _E * v * (1 + _E * v / (2 * _M0 * _C ** 2))) + return float(lam * 1e10) + + +@dataclass +class Reflectors: + """The set of diffracting planes a pattern is built from. + + normals : (N, 3) unit plane normals in the CRYSTAL cartesian frame, one per + Friedel pair — ``+g`` and ``-g`` are the same band, and keeping both + would draw every line twice. + weights : (N,) relative band intensity. + widths : (N,) band half-width, in units of ``r̂·n`` (i.e. the sine of the + angle between the beam direction and the plane). Only the simulator + uses this; the overlay draws centre lines. + hkl : (N, 3) Miller indices, kept for labelling. Optional. + """ + + normals: np.ndarray + weights: np.ndarray + widths: np.ndarray + hkl: np.ndarray | None = None + + def __post_init__(self): + self.normals = np.asarray(self.normals, float).reshape(-1, 3) + n = len(self.normals) + self.weights = np.broadcast_to(np.asarray(self.weights, float), (n,)).copy() + self.widths = np.broadcast_to(np.asarray(self.widths, float), (n,)).copy() + if self.hkl is not None: + self.hkl = np.asarray(self.hkl).reshape(n, 3) + + def __len__(self) -> int: + return len(self.normals) + + def brightest(self, n: int | None) -> "Reflectors": + """The *n* strongest bands. Drawing every reflector of a real phase + turns the pattern into a grey mesh, so the overlay draws a subset — and + the strongest bands are the ones actually visible in the data.""" + if n is None or n >= len(self): + return self + keep = np.argsort(self.weights)[::-1][:max(1, int(n))] + keep = np.sort(keep) + return Reflectors(self.normals[keep], self.weights[keep], + self.widths[keep], + None if self.hkl is None else self.hkl[keep]) + + +def _dedupe_friedel(normals: np.ndarray, tol: float = 1e-6) -> np.ndarray: + """Indices of one normal per ``±n`` pair, keeping the first occurrence.""" + keep: list[int] = [] + kept: list[np.ndarray] = [] + for i, v in enumerate(normals): + if any(abs(abs(float(v @ k)) - 1.0) < tol for k in kept): + continue + keep.append(i) + kept.append(v) + return np.asarray(keep, int) + + +def cubic_reflectors() -> Reflectors: + """The {111}/{200}/{220} band set of a generic cubic crystal. + + The default when no crystal structure has been loaded, and the exact set + :func:`spyde.data.synthetic.ebsd_patterns` renders — so the overlay can be + verified against the bundled synthetic scan pixel for pixel. The widths are + deliberately WIDER than a real 20 kV pattern's: they are what makes bands + visible on a 60 px detector, and a dictionary built from these has to match + data drawn with these. + """ + fams = [((1, 1, 1), 1.00), ((2, 0, 0), 0.70), ((2, 2, 0), 0.45)] + normals, weights, hkls = [], [], [] + for hkl, w in fams: + seen: set[tuple[int, int, int]] = set() + h, k, l = hkl + for perm in {(h, k, l), (h, l, k), (k, h, l), (k, l, h), (l, h, k), (l, k, h)}: + for sx in (1, -1): + for sy in (1, -1): + for sz in (1, -1): + v = (perm[0] * sx, perm[1] * sy, perm[2] * sz) + if v == (0, 0, 0) or tuple(-x for x in v) in seen: + continue + seen.add(v) + for v in sorted(seen): + g = np.array(v, float) + normals.append(g / np.linalg.norm(g)) + # Band width goes as 1/|g|: bigger d-spacing -> wider band. + weights.append(w / np.linalg.norm(g)) + hkls.append(v) + weights_arr = np.array(weights) + widths = 0.055 * (weights_arr / weights_arr.max()) + 0.012 + return Reflectors(np.array(normals), weights_arr, widths, np.array(hkls)) + + +def reflectors_from_phase(phase, *, min_dspacing: float = 0.7, + voltage_kv: float = 20.0, + max_bands: int | None = 60) -> Reflectors: + """The reflectors of a real crystal structure, via diffsims. + + ``phase`` is an orix :class:`~orix.crystal_map.Phase` — normally + ``Phase.from_cif(path)``, the same door the 4D-STEM orientation wizard uses. + Band intensity is ``|F_hkl|`` and the half-width is the Bragg angle + ``λ / 2d``, so at EBSD voltages the bands come out an order of magnitude + narrower than :func:`cubic_reflectors`' exaggerated synthetic ones — which + is correct, and the reason width is a per-reflector property rather than a + constant. + + Falls back to :func:`cubic_reflectors` if the phase carries no structure + (a ``Phase(space_group=…)`` with no atoms cannot give structure factors). + """ + from diffsims.crystallography import ReciprocalLatticeVector + + try: + rlv = ReciprocalLatticeVector.from_min_dspacing( + phase, min_dspacing=float(min_dspacing)) + rlv = rlv.unique(use_symmetry=True).symmetrise() + rlv.sanitise_phase() + rlv.calculate_structure_factor() + # |F_hkl| — the structure factor is COMPLEX, and asarray(..., float) + # would silently take the real part (a reflector whose phase is near + # 90 degrees would come out extinct). + factor = np.abs(np.asarray(rlv.structure_factor)).astype(float) + normals = np.asarray(rlv.unit.data, float).reshape(-1, 3) + d = np.asarray(rlv.dspacing, float).reshape(-1) + hkl = np.asarray(rlv.hkl, float).reshape(-1, 3) + except Exception as e: + log.warning("reflectors for %s could not be computed from the " + "structure (%s) — falling back to a generic cubic band set", + getattr(phase, "name", phase), e) + return cubic_reflectors() + + finite = np.isfinite(factor) & np.isfinite(d) & (d > 0) + finite &= np.isfinite(normals).all(1) + # A structureless Phase — `Phase(space_group=225)` with no atoms — still + # yields hkl, but every structure factor is ZERO, so the "reflectors" are + # all extinct and include ones a real fcc crystal cannot show ({100}, + # {110}). Weightless bands would draw a lattice of lines with nothing + # behind them, so treat it as no structure at all. + if not finite.any() or float(factor[finite].max()) <= 0: + log.debug("phase %s has no usable structure factors — using the " + "generic cubic band set", getattr(phase, "name", phase)) + return cubic_reflectors() + normals, factor, d, hkl = normals[finite], factor[finite], d[finite], hkl[finite] + + keep = _dedupe_friedel(normals) + normals, factor, d, hkl = normals[keep], factor[keep], d[keep], hkl[keep] + + lam = wavelength(voltage_kv) + widths = np.clip(lam / (2.0 * d), 1e-4, 0.25) + weights = factor / max(float(factor.max()), 1e-12) + refl = Reflectors(normals, weights, widths, hkl.astype(int)) + return refl.brightest(max_bands) + + +def detector_directions(detector=(60, 60), pc=(0.5, 0.5, 0.55)) -> np.ndarray: + """Unit vectors from the sample to each detector pixel (gnomonic). + + ``(dy, dx, 3)``. ``pc`` is ``(pcx, pcy, L)`` in the fractional convention: + the pattern centre as a fraction of the detector width/height, and the + detector distance as a fraction of the width. + """ + pcx, pcy, L = float(pc[0]), float(pc[1]), float(pc[2]) + dy, dx = int(detector[0]), int(detector[1]) + gy, gx = np.mgrid[0:dy, 0:dx].astype(float) + rx = (gx + 0.5) / dx - pcx + ry = pcy - (gy + 0.5) / dy # detector y is flipped + r = np.stack([rx, ry, np.full_like(rx, L)], -1) + return r / np.linalg.norm(r, axis=-1, keepdims=True) + + +def euler_to_matrix(phi1, Phi, phi2) -> np.ndarray: + """Bunge ZXZ Euler angles -> rotation matrices, batched over the leading + axes. Returns ``(..., 3, 3)``.""" + c1, s1 = np.cos(phi1), np.sin(phi1) + c, s = np.cos(Phi), np.sin(Phi) + c2, s2 = np.cos(phi2), np.sin(phi2) + m = np.empty(np.shape(phi1) + (3, 3), float) + m[..., 0, 0] = c1 * c2 - s1 * s2 * c + m[..., 0, 1] = s1 * c2 + c1 * s2 * c + m[..., 0, 2] = s2 * s + m[..., 1, 0] = -c1 * s2 - s1 * c2 * c + m[..., 1, 1] = -s1 * s2 + c1 * c2 * c + m[..., 1, 2] = c2 * s + m[..., 2, 0] = s1 * s + m[..., 2, 1] = -c1 * s + m[..., 2, 2] = c + return m + + +def normals_to_sample(normals, rot): + """Plane normals from the CRYSTAL frame into the SAMPLE frame. + + ``rot`` is the Bunge matrix from :func:`euler_to_matrix`, which maps SAMPLE + components to CRYSTAL components (the orix convention — the two agree, and + ``test_ebsd_bands`` pins that). So going the other way needs its TRANSPOSE, + which for row-vector normals is a plain ``normals @ rot``. + + Getting this backwards is invisible in an indexing score — the dictionary + and the data would simply share the mistake and match each other perfectly + — and shows up only where the Euler angles leave this module: the IPF map + would then be coloured from the INVERSE orientation, still showing grains + and gradients, just the wrong colours. The tell is symmetry: the simulated + pattern must be unchanged by a CRYSTAL symmetry operation (Bunge phi2 + 90 + degrees for a cubic 4-fold) and must change under a SAMPLE rotation + (phi1 + 90). Transposed, that is exactly reversed. + """ + return np.asarray(normals) @ np.asarray(rot) + + +def simulate_patterns(euler, reflectors: Reflectors | None = None, + detector=(60, 60), pc=(0.5, 0.5, 0.55), *, + background: bool = False) -> np.ndarray: + """Kikuchi patterns for a list of orientations -> ``(N, dy, dx)`` float32. + + *euler* is ``(N, 3)`` Bunge angles in radians. Pure numpy and looped over + N: this is for a handful of patterns (a preview, a test). Building a + DICTIONARY of thousands goes through + :class:`spyde.ebsd.refine.BandSimulator`, which is the same arithmetic + batched in torch. + + ``background=False`` by default because a dictionary is matched by + normalised cross-correlation, which is invariant to the smooth gradient a + real detector adds — and the experimental side has background removal + applied before matching anyway (:mod:`spyde.ebsd.preprocess`). + """ + refl = reflectors if reflectors is not None else cubic_reflectors() + euler = np.atleast_2d(np.asarray(euler, float)) + rot = euler_to_matrix(euler[:, 0], euler[:, 1], euler[:, 2]) # (N, 3, 3) + r = detector_directions(detector, pc) + dy, dx = r.shape[:2] + flat_r = r.reshape(-1, 3) + + out = np.empty((len(euler), dy, dx), np.float32) + for i in range(len(euler)): + n_rot = normals_to_sample(refl.normals, rot[i]) + d = flat_r @ n_rot.T + band = np.exp(-0.5 * (d / refl.widths) ** 2) * refl.weights + out[i] = band.sum(1).reshape(dy, dx).astype(np.float32) + if background: + gy, gx = np.mgrid[0:dy, 0:dx].astype(np.float32) + out = out / max(out.max(), 1e-9) + ( + 0.35 + 0.30 * ((gy / dy) * 0.6 + (gx / dx) * 0.4)) + return out + + +def _clip_line_to_box(a: np.ndarray, b: np.ndarray, c: np.ndarray, + width: float, height: float): + """Clip the lines ``a·u + b·v + c = 0`` to ``[0, width] x [0, height]``. + + Vectorised over the bands. Returns ``(segments (M, 2, 2), keep (N,) bool)`` + — the ``[[u0, v0], [u1, v1]]`` line-collection convention anyplotlib's + ``add_lines`` takes — and *keep* says which input lines crossed the box. + + Done by intersecting with all four edges and taking the two intersections + that lie inside, rather than by a Liang-Barsky parametric walk: a band line + can be exactly axis-parallel (a zone axis on the pattern centre makes this + common, not a corner case), and the edge-intersection form degrades to + "no intersection with the parallel edges" instead of dividing by zero. + """ + a = np.asarray(a, float) + b = np.asarray(b, float) + c = np.asarray(c, float) + n = len(a) + eps = 1e-12 + + pts = np.full((n, 4, 2), np.nan) + with np.errstate(divide="ignore", invalid="ignore"): + # u = 0 and u = width -> v = -(c + a*u) / b + for j, u in enumerate((0.0, width)): + v = np.where(np.abs(b) > eps, -(c + a * u) / np.where(np.abs(b) > eps, b, 1.0), np.nan) + ok = np.isfinite(v) & (v >= -1e-9) & (v <= height + 1e-9) + pts[ok, j, 0] = u + pts[ok, j, 1] = v[ok] + # v = 0 and v = height -> u = -(c + b*v) / a + for j, v in enumerate((0.0, height), start=2): + u = np.where(np.abs(a) > eps, -(c + b * v) / np.where(np.abs(a) > eps, a, 1.0), np.nan) + ok = np.isfinite(u) & (u >= -1e-9) & (u <= width + 1e-9) + pts[ok, j, 0] = u[ok] + pts[ok, j, 1] = v + + segs = np.zeros((n, 2, 2)) + keep = np.zeros(n, bool) + for i in range(n): + p = pts[i][np.isfinite(pts[i]).all(1)] + if len(p) < 2: + continue + # A line through a corner hits two edges at the same point; take the + # two furthest-apart intersections so the segment spans the box. + d2 = ((p[:, None, :] - p[None, :, :]) ** 2).sum(-1) + i0, i1 = np.unravel_index(int(np.argmax(d2)), d2.shape) + if d2[i0, i1] <= 1e-9: + continue + segs[i] = (p[i0], p[i1]) + keep[i] = True + return segs[keep], keep + + +def band_lines(euler, reflectors: Reflectors | None = None, + detector=(60, 60), pc=(0.5, 0.5, 0.55), *, + max_bands: int | None = None): + """Band CENTRE lines for one orientation -> ``(M, 2, 2)`` detector pixels. + + Returns ``(segments, weights)`` where each segment is + ``[[x0, y0], [x1, y1]]`` — the line-collection shape anyplotlib's + ``add_lines`` takes — in the image-pixel convention its ``transform="data"`` + markers use (pixel centres on integers, so the frame spans + ``-0.5 … size-0.5``). *weights* is the matching band intensity, so a caller + can style by band strength. + + Only bands that actually cross the detector are returned, so ``M <= N``. + """ + refl = reflectors if reflectors is not None else cubic_reflectors() + if max_bands is not None: + refl = refl.brightest(max_bands) + euler = np.asarray(euler, float).reshape(3) + rot = euler_to_matrix(*euler) # (3, 3) + n_rot = normals_to_sample(refl.normals, rot) # (N, 3) sample frame + + dy, dx = int(detector[0]), int(detector[1]) + pcx, pcy, L = float(pc[0]), float(pc[1]), float(pc[2]) + nx, ny, nz = n_rot[:, 0], n_rot[:, 1], n_rot[:, 2] + + # r·n = 0 with r = ((u/dx - pcx), (pcy - v/dy), L) and u, v the continuous + # pixel coordinates. Linear in (u, v) — the band centre IS a straight line. + a = nx / dx + b = -ny / dy + c = L * nz - pcx * nx + pcy * ny + + segs, keep = _clip_line_to_box(a, b, c, float(dx), float(dy)) + # u, v measure from the frame edge; anyplotlib pixel coordinates put pixel + # 0's CENTRE at 0, so the whole frame shifts by half a pixel. + return (segs - 0.5).astype(np.float32), refl.weights[keep].astype(np.float32) + + +def zone_axis_points(euler, reflectors: Reflectors | None = None, + detector=(60, 60), pc=(0.5, 0.5, 0.55), *, + max_axes: int = 12, min_bands: int = 3): + """Where band lines meet — the zone axes -> ``(M, 2)`` detector pixels. + + A zone axis is a crystal direction several planes share, and on a real + pattern it is the bright junction the eye locks onto first, which makes it + the most useful thing to draw after the bands themselves. Found from the + band set rather than from a separate direction list: any pair of normals + defines their zone axis ``n_i x n_j``, and the axes worth drawing are the + ones many pairs agree on. + """ + refl = reflectors if reflectors is not None else cubic_reflectors() + euler = np.asarray(euler, float).reshape(3) + rot = euler_to_matrix(*euler) + n_rot = normals_to_sample(refl.normals, rot) + n = len(n_rot) + if n < 2: + return np.zeros((0, 2), np.float32) + + i, j = np.triu_indices(n, k=1) + axes = np.cross(n_rot[i], n_rot[j]) + norm = np.linalg.norm(axes, axis=1) + axes = axes[norm > 1e-6] / norm[norm > 1e-6, None] + if not len(axes): + return np.zeros((0, 2), np.float32) + axes *= np.sign(axes[:, 2:3] + 1e-30) # keep the +z hemisphere + + # Cluster the pairwise axes: a direction that many band pairs share is a + # real zone axis, one that a single pair gives is just two lines crossing. + order = np.lexsort((axes[:, 2], axes[:, 1], axes[:, 0])) + axes = axes[order] + uniq: list[np.ndarray] = [] + counts: list[int] = [] + for v in axes: + for k, u in enumerate(uniq): + if float(v @ u) > 0.9995: + counts[k] += 1 + break + else: + uniq.append(v) + counts.append(1) + uniq_arr = np.asarray(uniq) + counts_arr = np.asarray(counts) + # k planes through one axis give k*(k-1)/2 pairs; min_bands=3 -> 3 pairs. + want = max(1, min_bands * (min_bands - 1) // 2) + sel = counts_arr >= want + if not sel.any(): + return np.zeros((0, 2), np.float32) + uniq_arr, counts_arr = uniq_arr[sel], counts_arr[sel] + keep = np.argsort(counts_arr)[::-1][:int(max_axes)] + uniq_arr = uniq_arr[keep] + + dy, dx = int(detector[0]), int(detector[1]) + pcx, pcy, L = float(pc[0]), float(pc[1]), float(pc[2]) + tz = uniq_arr[:, 2] + front = tz > 1e-6 # behind the detector = invisible + uniq_arr, tz = uniq_arr[front], tz[front] + u = dx * (pcx + uniq_arr[:, 0] * L / tz) - 0.5 + v = dy * (pcy - uniq_arr[:, 1] * L / tz) - 0.5 + inside = (u >= -0.5) & (u <= dx - 0.5) & (v >= -0.5) & (v <= dy - 0.5) + return np.column_stack([u[inside], v[inside]]).astype(np.float32) diff --git a/spyde/ebsd/indexing.py b/spyde/ebsd/indexing.py index 13cb351..7f0cb69 100644 --- a/spyde/ebsd/indexing.py +++ b/spyde/ebsd/indexing.py @@ -67,16 +67,23 @@ def _normalise(a, eps=1e-12): def dictionary_index(patterns, dictionary, *, keep: int = 1, device=None, dtype="float32", tile_elements: int = _TILE_ELEMENTS, - progress: Callable[[int, int], None] | None = None): + pattern_chunk: int | None = None, + progress: Callable[[int, int], None] | None = None, + on_chunk: Callable[[int, int, np.ndarray, np.ndarray], + None] | None = None, + stopped_flag=None): """Match every pattern against every dictionary entry. Parameters ---------- patterns : array (..., H, W) or (P, K) Experimental patterns. Leading dimensions are flattened to ``P``. - dictionary : array (D, H, W) or (D, K) - Simulated patterns, e.g. from - :func:`spyde.data.synthetic.simulate_patterns`. + dictionary : array (D, H, W) or (D, K), or a :class:`SinglePatternIndexer` + Simulated patterns, e.g. from :func:`simulate_dictionary`. Passing an + indexer reuses the dictionary it already holds normalised on the + device — which is what the interactive path does, having built one for + the live preview: re-normalising ``D x K`` is the single most expensive + step of a small index and there is no reason to pay it twice. keep : int How many matches to keep per pattern. >1 is what ``orientation_similarity_map`` needs (#73). @@ -84,24 +91,42 @@ def dictionary_index(patterns, dictionary, *, keep: int = 1, device=None, ``float32`` by default and deliberately: NCC is a bounded, well- conditioned quantity and the matmul is memory-bound, so float64 would halve throughput to protect digits that do not affect the ranking. + pattern_chunk : int, optional + Cap on the experimental tile size. The tiling *alone* would put every + pattern in one chunk whenever the dictionary is small enough — correct + and fastest, but it means *progress* and *on_chunk* fire exactly once, + at the end. A UI filling a map in as it indexes passes this to get + intermediate results; nothing else should. + on_chunk : callable, optional + ``(lo, hi, indices, scores)`` for each completed slice of patterns, so + a caller can paint partial results. Called on the calling thread. + stopped_flag : list[bool], optional + Polled between chunks; a truthy first element abandons the run and + returns None. This is how closing the window cancels an index. Returns ------- - IndexingResult + IndexingResult, or None if cancelled. """ import torch + resident = getattr(dictionary, "normalised", None) # SinglePatternIndexer + if resident is not None: + device = device or dictionary.device device = device or ("cuda" if torch.cuda.is_available() else "cpu") tdtype = getattr(torch, dtype) exp = np.asarray(patterns) - dic = np.asarray(dictionary) # Compare pixel counts BEFORE reshaping. Reshaping the experimental stack # to the DICTIONARY's pixel count quietly succeeds whenever the sizes share # a factor — a 40x40 scan against a 20x20 dictionary reshapes to 4x as many # "patterns" and indexes garbage instead of raising. - dic_k = int(np.prod(dic.shape[1:])) + if resident is not None: + dic_k = int(dictionary.k) + else: + dic = np.asarray(dictionary) + dic_k = int(np.prod(dic.shape[1:])) exp_k = int(np.prod(exp.shape[-2:] if exp.ndim > 2 else exp.shape[-1:])) if exp_k != dic_k: raise ValueError(f"pattern has {exp_k} pixels but dictionary entries " @@ -109,22 +134,29 @@ def dictionary_index(patterns, dictionary, *, keep: int = 1, device=None, nav_shape = exp.shape[:-2] if exp.ndim > 2 else exp.shape[:-1] E = exp.reshape(-1, exp_k) - D = dic.reshape(-1, dic_k) - P, Dn = E.shape[0], D.shape[0] - keep = int(min(keep, Dn)) # The dictionary is normalised ONCE and kept resident — it is reused by # every tile, and re-normalising it per tile would dominate the matmul. - d_t = _normalise(torch.as_tensor(D, dtype=tdtype, device=device)) + d_t = (resident if resident is not None + else _normalise(torch.as_tensor(dic.reshape(-1, dic_k), dtype=tdtype, + device=device))) + tdtype = d_t.dtype + P, Dn = E.shape[0], int(d_t.shape[0]) + keep = int(min(keep, Dn)) # Tile both axes so no (P, D) block is ever larger than the budget. d_chunk = max(1, min(Dn, tile_elements // max(P, 1))) p_chunk = max(1, min(P, tile_elements // max(min(d_chunk, Dn), 1))) + if pattern_chunk: + p_chunk = max(1, min(p_chunk, int(pattern_chunk))) out_scores = np.empty((P, keep), np.float32) out_index = np.empty((P, keep), np.int64) for p0 in range(0, P, p_chunk): + if stopped_flag is not None and stopped_flag[0]: + log.debug("dictionary indexing cancelled at %d/%d patterns", p0, P) + return None p1 = min(P, p0 + p_chunk) e_t = _normalise(torch.as_tensor(E[p0:p1], dtype=tdtype, device=device)) @@ -145,6 +177,11 @@ def dictionary_index(patterns, dictionary, *, keep: int = 1, device=None, out_scores[p0:p1] = best_s.detach().cpu().numpy() out_index[p0:p1] = best_i.detach().cpu().numpy() + if on_chunk is not None: + try: + on_chunk(p0, p1, out_index[p0:p1], out_scores[p0:p1]) + except Exception as e: # pragma: no cover + log.debug("indexing on_chunk callback failed: %s", e) if progress is not None: try: progress(p1, P) @@ -158,15 +195,133 @@ def dictionary_index(patterns, dictionary, *, keep: int = 1, device=None, return result -def sample_orientations(step_deg: float = 5.0, *, phi1=(0.0, 360.0), - Phi=(0.0, 90.0), phi2=(0.0, 90.0)) -> np.ndarray: - """A regular Euler-space grid -> ``(N, 3)`` radians. +class SinglePatternIndexer: + """The dictionary, normalised once and kept resident, for live indexing. + + :func:`dictionary_index` is the batch door: it normalises the dictionary, + matches a whole scan and drops everything. Under the crosshair the shape of + the problem is the opposite — ONE pattern at a time, again on every + navigator move — so re-normalising ``D x K`` per move would cost far more + than the match itself (25k entries of a 60x60 detector is ~370 MB; the + match is a single mat-vec). + + So the wizard builds one of these when it builds the dictionary and the + band overlay calls :meth:`best` per position. Same normalisation, same + scores as the batch path — a dot product of zero-mean unit-norm vectors. + """ + + def __init__(self, dictionary, euler, *, device=None, dtype="float32"): + import torch + + self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") + tdtype = getattr(torch, dtype) + dic = np.asarray(dictionary) + self.euler = np.asarray(euler, float).reshape(-1, 3) + self.pattern_shape = tuple(dic.shape[1:]) + self.k = int(np.prod(self.pattern_shape)) + self._d = _normalise(torch.as_tensor(dic.reshape(-1, self.k), + dtype=tdtype, device=self.device)) + self._dtype = tdtype + + def __len__(self) -> int: + return int(self._d.shape[0]) + + @property + def normalised(self): + """The zero-mean unit-norm dictionary, resident on ``device``. + :func:`dictionary_index` takes this instead of re-normalising.""" + return self._d + + def best(self, pattern): + """Best-matching orientation for one pattern -> ``(euler (3,), score)``.""" + import torch + + p = np.asarray(pattern, float).reshape(-1) + if p.size != self.k: + raise ValueError(f"pattern has {p.size} pixels but the dictionary " + f"has {self.k}") + with torch.no_grad(): + e = _normalise(torch.as_tensor(p, dtype=self._dtype, + device=self.device)) + sim = self._d @ e + score, idx = torch.max(sim, 0) + return self.euler[int(idx.item())], float(score.item()) + + +def simulate_dictionary(euler, detector=(60, 60), pc=(0.5, 0.5, 0.55), *, + reflectors=None, background_sigma=None, device=None, + chunk: int = 4096, + progress: Callable[[int, int], None] | None = None, + stopped_flag=None): + """Simulate one pattern per orientation -> ``(D, H, W)`` float32. + + The dictionary side of indexing. Uses the same torch + :class:`~spyde.ebsd.refine.BandSimulator` the refinement optimises through, + so the dictionary and the refined patterns are the same function — and the + same :class:`~spyde.ebsd.bands.Reflectors` the live overlay projects, so + the drawn lines are the drawn bands' centres. + + ``spyde.ebsd.bands.simulate_patterns`` does this in numpy for a handful of + orientations; a dictionary is thousands, so it goes through torch in + chunks — batched over orientations, with the chunk bounding the largest + ``(chunk, B, K)`` intermediate rather than materialising ``(D, B, K)``. + + Pass *background_sigma* to high-pass the simulated patterns exactly as the + experimental ones were corrected. Skipping it is the classic way to get + mediocre scores that look like bad indexing — see + :meth:`~spyde.ebsd.refine.BandSimulator._high_pass`. + """ + import torch + from spyde.ebsd.refine import BandSimulator - Deliberately simple and NOT equal-area: a proper dictionary uses a uniform - SO(3) sampling (kikuchipy/orix do this correctly, and #69 wires that up). - This exists so indexing can be developed and tested without the extra - installed, and the default ranges cover the cubic fundamental zone. + device = device or ("cuda" if torch.cuda.is_available() else "cpu") + eul = np.atleast_2d(np.asarray(euler, float)) + sim = BandSimulator(detector, pc, reflectors=reflectors, + background_sigma=background_sigma, device=device) + dy, dx = int(detector[0]), int(detector[1]) + out = np.empty((len(eul), dy, dx), np.float32) + with torch.no_grad(): + for lo in range(0, len(eul), int(chunk)): + if stopped_flag is not None and stopped_flag[0]: + return None + hi = min(len(eul), lo + int(chunk)) + ang = torch.as_tensor(eul[lo:hi], dtype=torch.float32, device=device) + out[lo:hi] = sim(ang).reshape(-1, dy, dx).detach().cpu().numpy() + if progress is not None: + try: + progress(hi, len(eul)) + except Exception as e: # pragma: no cover + log.debug("dictionary progress callback failed: %s", e) + return out + + +def sample_orientations(step_deg: float = 5.0, *, point_group=None, + phi1=(0.0, 360.0), Phi=(0.0, 90.0), + phi2=(0.0, 90.0)) -> np.ndarray: + """Orientations to build a dictionary from -> ``(N, 3)`` Bunge radians. + + With a *point_group* this is orix's uniform sampling of that group's + FUNDAMENTAL ZONE — every distinct crystal orientation once, none of them + twice. Prefer it: the Euler grid below is neither equal-area (it bunches + towards Phi = 0) nor free of symmetric duplicates, and the difference is + not academic — at 5 degrees for m-3m it is ~6.6k orientations against the + grid's ~26k, so the dictionary is a quarter of the size, a quarter of the + simulation time, and a quarter of the cost of every live match. + + Without one it falls back to a plain Euler grid over the given ranges, so + indexing still works with nothing but numpy — which is what lets the rest + of this module be developed and tested without the ``ebsd`` extra. """ + if point_group is not None: + try: + from orix.sampling import get_sample_fundamental + rot = get_sample_fundamental(resolution=float(step_deg), + point_group=point_group) + return np.asarray(rot.to_euler(), float).reshape(-1, 3) + except Exception as e: + log.warning("orix fundamental-zone sampling failed (%s) — falling " + "back to an Euler grid", e) + a = np.deg2rad(np.arange(phi1[0], phi1[1], step_deg)) b = np.deg2rad(np.arange(Phi[0], Phi[1] + 1e-9, step_deg)) c = np.deg2rad(np.arange(phi2[0], phi2[1] + 1e-9, step_deg)) diff --git a/spyde/ebsd/refine.py b/spyde/ebsd/refine.py index e93713c..0144020 100644 --- a/spyde/ebsd/refine.py +++ b/spyde/ebsd/refine.py @@ -76,10 +76,16 @@ def euler_to_matrix_torch(euler): class BandSimulator: """Differentiable Kikuchi-band pattern simulator. - Renders the same geometry as :func:`spyde.data.synthetic.simulate_patterns` - — a band appears where a detector direction is near-perpendicular to a - plane normal — but in torch, so the pattern is differentiable with respect - to the orientation. + Renders the same geometry as :func:`spyde.ebsd.bands.simulate_patterns` — + a band appears where a detector direction is near-perpendicular to a plane + normal — but in torch, so the pattern is differentiable with respect to the + orientation, and batched, so a whole DICTIONARY is one forward pass. + + The band set comes from :class:`spyde.ebsd.bands.Reflectors`: the generic + cubic set by default, or a real crystal structure's reflectors when the + wizard has loaded a .cif. Whatever it is, the live band overlay projects + THE SAME reflectors, so the lines it draws are the centres of the bands + this simulator rendered. Stands in for a master-pattern lookup (#69). The optimiser does not care which it is: anything mapping ``(P, 3)`` Euler angles to ``(P, K)`` patterns @@ -87,27 +93,58 @@ class BandSimulator: """ def __init__(self, detector=(60, 60), pc=(0.5, 0.5, 0.55), *, - device="cpu", dtype=None): + reflectors=None, background_sigma=None, device="cpu", + dtype=None): import torch - from spyde.data.synthetic import _cubic_plane_normals, detector_directions + from spyde.ebsd.bands import cubic_reflectors, detector_directions dtype = dtype or torch.float32 + refl = reflectors if reflectors is not None else cubic_reflectors() r = detector_directions(detector, pc).reshape(-1, 3) - normals, weights = _cubic_plane_normals() + self.reflectors = refl self.r = torch.as_tensor(r, dtype=dtype, device=device) - self.normals = torch.as_tensor(normals, dtype=dtype, device=device) - self.weights = torch.as_tensor(weights, dtype=dtype, device=device) - self.widths = torch.as_tensor( - 0.055 * (weights / weights.max()) + 0.012, dtype=dtype, device=device) + self.normals = torch.as_tensor(refl.normals, dtype=dtype, device=device) + self.weights = torch.as_tensor(refl.weights, dtype=dtype, device=device) + self.widths = torch.as_tensor(refl.widths, dtype=dtype, device=device) self.shape = tuple(detector) + self.device = device + self.dtype = dtype + self.background_sigma = (None if background_sigma in (None, 0) + else float(background_sigma)) def __call__(self, euler): """``(P, 3)`` -> ``(P, K)``.""" rot = euler_to_matrix_torch(euler) # (P, 3, 3) - n_rot = self.normals @ rot.transpose(1, 2) # (P, B, 3) + # normals @ rot, NOT rot.T — see bands.normals_to_sample for why the + # transposed version passes every score-based test and still colours + # the IPF map from the inverse orientation. + n_rot = self.normals @ rot # (P, B, 3) d = n_rot @ self.r.T # (P, B, K) band = (-0.5 * (d / self.widths[None, :, None]) ** 2).exp() - return (band * self.weights[None, :, None]).sum(1) # (P, K) + out = (band * self.weights[None, :, None]).sum(1) # (P, K) + return self._high_pass(out) if self.background_sigma else out + + def _high_pass(self, flat): + """Apply the SAME dynamic-background high-pass the experimental + patterns get (:func:`spyde.ebsd.preprocess.remove_background`). + + Both sides must go through the same filter or the correction makes the + mismatch *different* rather than smaller: high-passing the experimental + patterns alone leaves the simulated ones carrying low-frequency content + their counterparts no longer have, and scores stay mediocre for a + reason that looks like bad indexing. (kikuchipy preprocesses its + dictionary for exactly this reason; ``test_ebsd_indexing`` pins it.) + + Done here rather than as a post-pass so it holds through REFINEMENT + too — the separable blur is a convolution, so the filtered pattern is + still differentiable with respect to the orientation. + """ + from spyde.ebsd.preprocess import _blur + + h, w = self.shape + img = flat.reshape(-1, h, w) + return (img - _blur(img, self.background_sigma, self.device, + self.dtype)).reshape(flat.shape) def _ncc(a, b, eps=1e-12): @@ -120,7 +157,8 @@ def _ncc(a, b, eps=1e-12): def refine_orientations(patterns, euler_start, *, simulator=None, - detector=None, pc=(0.5, 0.5, 0.55), device=None, + detector=None, pc=(0.5, 0.5, 0.55), reflectors=None, + background_sigma=None, device=None, steps: int = 120, lr: float = 0.01, chunk=None, on_yield: Callable[[], None] | None = None, yield_every: int = 12, @@ -138,6 +176,11 @@ def refine_orientations(patterns, euler_start, *, simulator=None, simulator : callable, optional ``(P, 3) -> (P, K)``, differentiable. Defaults to :class:`BandSimulator` on *detector*. + background_sigma : float, optional + High-pass the SIMULATED patterns the same way the experimental ones + were corrected. Pass it whenever *patterns* has been through + :func:`~spyde.ebsd.preprocess.remove_background`, or the two sides are + being compared through different filters. steps, lr : Adam budget. The default is deliberately generous: refinement runs once per scan, and the cost is one batched forward+backward per step. @@ -168,7 +211,9 @@ def refine_orientations(patterns, euler_start, *, simulator=None, if simulator is None: if detector is None: detector = tuple(exp.shape[-2:]) - simulator = BandSimulator(detector, pc, device=device) + simulator = BandSimulator(detector, pc, reflectors=reflectors, + background_sigma=background_sigma, + device=device) P = len(E) chunk = int(chunk) if chunk else P diff --git a/spyde/signal_tree.py b/spyde/signal_tree.py index d6f2db4..e71357d 100644 --- a/spyde/signal_tree.py +++ b/spyde/signal_tree.py @@ -865,7 +865,7 @@ def close(self) -> None: logger.debug("removing %s on tree close failed: %s", attr, e) if hasattr(self, attr): setattr(self, attr, None) - for wiz_attr in ("_om_wizard", "_vom_wizard"): + for wiz_attr in ("_om_wizard", "_vom_wizard", "_ebsd_wizard"): wiz = getattr(self, wiz_attr, None) if wiz is not None and hasattr(wiz, "remove"): try: diff --git a/spyde/tests/migrated/test_ebsd_bands.py b/spyde/tests/migrated/test_ebsd_bands.py new file mode 100644 index 0000000..93c70fe --- /dev/null +++ b/spyde/tests/migrated/test_ebsd_bands.py @@ -0,0 +1,261 @@ +"""Kikuchi band geometry (:mod:`spyde.ebsd.bands`) — the overlay's foundation. + +Two things are worth pinning here and they are not the same thing: + +* **the lines land on the bands** — the overlay's whole claim. Checked by + sampling the simulated pattern ALONG each drawn line and requiring it to be + brighter than the frame, and by requiring a WRONG orientation's lines not to + be; +* **the rotation convention agrees with orix** — which no score-based test can + see. Indexing is self-consistent whichever way round the rotation goes, so a + transposed simulator matches its own dictionary perfectly and then colours + the IPF map from the inverse orientation. The tell is symmetry: a CRYSTAL + symmetry operation must leave the pattern alone and a SAMPLE rotation must + not. +""" +from __future__ import annotations + +import numpy as np +import pytest + +from spyde.ebsd.bands import ( + Reflectors, band_lines, cubic_reflectors, detector_directions, + euler_to_matrix, normals_to_sample, simulate_patterns, wavelength, + zone_axis_points, +) + +DET = (60, 60) +PC = (0.5, 0.5, 0.55) +EULER = np.deg2rad([33.3, 30.0, 15.0]) + + +def _sample_along(pattern, segments, n=40, offset=0.0): + """Mean pattern intensity along each segment, optionally shifted *offset* + pixels perpendicular to it.""" + out = [] + h, w = pattern.shape + for (x0, y0), (x1, y1) in segments: + t = np.linspace(0.05, 0.95, n) + v = np.array([x1 - x0, y1 - y0], float) + nrm = np.array([-v[1], v[0]]) / max(np.linalg.norm(v), 1e-12) * offset + xs = np.clip(np.round(x0 + t * (x1 - x0) + nrm[0]).astype(int), 0, w - 1) + ys = np.clip(np.round(y0 + t * (y1 - y0) + nrm[1]).astype(int), 0, h - 1) + out.append(float(pattern[ys, xs].mean())) + return np.asarray(out) + + +def _perpendicular_profile(pattern, segments, half=6): + """Intensity across each band: ``(n_bands, 2*half+1)`` sampled at + perpendicular offsets ``-half … +half`` pixels. + + Comparing a line against the WHOLE-FRAME mean does not actually test + anything about placement — a weak band grazing a dim corner sits below the + frame mean while being drawn perfectly, and a line through the bright + middle beats the mean while being drawn wrong. What the overlay claims is + LOCAL: the pattern is brightest ON the line and falls away to either side. + """ + offs = np.arange(-half, half + 1, dtype=float) + return np.column_stack([_sample_along(pattern, segments, offset=o) + for o in offs]), offs + + +class TestConventionAgreesWithOrix: + """The Euler angles this module simulates from are handed straight to orix + to colour the IPF map, so the two must mean the same thing.""" + + def test_euler_matrix_matches_orix(self): + from orix.quaternion import Rotation + ours = euler_to_matrix(*EULER) + theirs = Rotation.from_euler(EULER[None]).to_matrix()[0] + assert np.allclose(ours, theirs, atol=1e-9), ( + "Bunge matrix disagrees with orix — every orientation leaving this " + "module would be transposed") + + def test_crystal_symmetry_leaves_the_pattern_alone(self): + """phi2 + 90 deg is a cubic 4-fold about the crystal z: the SAME + crystal orientation, so the same pattern.""" + from orix.quaternion import Orientation, symmetry + a = Orientation.from_euler(EULER[None], symmetry.Oh) + b = Orientation.from_euler((EULER + np.deg2rad([0, 0, 90]))[None], + symmetry.Oh) + assert float(np.rad2deg(a.angle_with(b))[0]) < 1e-6, \ + "orix says these are different orientations; the premise is wrong" + + p0 = simulate_patterns(EULER, detector=DET, pc=PC)[0] + p1 = simulate_patterns(EULER + np.deg2rad([0, 0, 90]), + detector=DET, pc=PC)[0] + assert np.allclose(p0, p1, atol=1e-6), \ + "a crystal symmetry operation changed the pattern — the normals " \ + "are being rotated by the transpose" + + def test_a_sample_rotation_does_change_the_pattern(self): + """The other half of the same check — without it, a simulator that + ignores the rotation entirely would pass the test above.""" + p0 = simulate_patterns(EULER, detector=DET, pc=PC)[0] + p1 = simulate_patterns(EULER + np.deg2rad([90, 0, 0]), + detector=DET, pc=PC)[0] + assert not np.allclose(p0, p1, atol=1e-3) + + def test_normals_to_sample_is_the_transpose(self): + rot = euler_to_matrix(*EULER) + n = cubic_reflectors().normals + assert np.allclose(normals_to_sample(n, rot), (rot.T @ n.T).T) + + +class TestBandLinesLandOnBands: + def test_every_line_sits_on_a_local_intensity_peak(self): + """The overlay's whole claim: each drawn line is the CENTRE of a band, + so the pattern peaks on it and falls away to either side.""" + pat = simulate_patterns(EULER, detector=DET, pc=PC)[0] + segs, weights = band_lines(EULER, detector=DET, pc=PC) + assert len(segs) > 3, "no bands crossed the detector" + assert len(weights) == len(segs) + + profile, offs = _perpendicular_profile(pat, segs) + peak = offs[profile.argmax(axis=1)] + assert np.abs(peak).max() <= 1.0, \ + f"band lines are offset from the bands by up to {np.abs(peak).max()}px" + # And the peak is a real one, not a plateau. + assert (profile.max(1) > profile[:, [0, -1]].max(1)).all() + + def test_a_wrong_orientation_does_not_land_on_them(self): + """Without this, lines drawn anywhere down a bright band would pass.""" + pat = simulate_patterns(EULER, detector=DET, pc=PC)[0] + wrong = np.deg2rad([70.0, 55.0, 40.0]) + segs, _ = band_lines(wrong, detector=DET, pc=PC) + profile, offs = _perpendicular_profile(pat, segs) + peak = offs[profile.argmax(axis=1)] + assert np.abs(peak).max() > 1.0, \ + "a wrong orientation's lines also landed on band centres" + assert _sample_along(pat, segs).mean() < pat.mean() + + def test_segments_are_the_line_collection_shape(self): + """anyplotlib add_lines takes (N, 2, 2) and raises on anything else.""" + segs, _ = band_lines(EULER, detector=DET, pc=PC) + assert segs.ndim == 3 and segs.shape[1:] == (2, 2) + + def test_segments_stay_inside_the_detector(self): + segs, _ = band_lines(EULER, detector=DET, pc=PC) + assert segs.min() >= -0.5 - 1e-6 + assert segs[..., 0].max() <= DET[1] - 0.5 + 1e-6 + assert segs[..., 1].max() <= DET[0] - 0.5 + 1e-6 + + def test_pixel_convention_is_centre_on_integer(self): + """A band through the pattern centre must be drawn through the centre + pixel, not half a pixel off — the overlay is drawn in the image-pixel + coordinates anyplotlib's transform="data" markers use.""" + # [001] out of the detector: the (200) plane containing the beam gives + # a band straight through the pattern centre. + segs, _ = band_lines(np.zeros(3), detector=DET, pc=(0.5, 0.5, 0.55)) + centre = (DET[1] - 1) / 2.0 + mid = segs.mean(axis=1) # (N, 2) segment midpoints + assert np.abs(mid - centre).min() < 1e-6, \ + "no band passes through the detector centre for the identity " \ + "orientation with a centred PC" + + def test_max_bands_keeps_the_strongest(self): + segs_all, w_all = band_lines(EULER, detector=DET, pc=PC) + segs_few, w_few = band_lines(EULER, detector=DET, pc=PC, max_bands=3) + assert len(segs_few) <= 3 < len(segs_all) + assert w_few.min() >= np.sort(w_all)[::-1][:len(w_few)].min() - 1e-9 + + def test_a_band_that_misses_the_detector_is_dropped(self): + """A tiny detector sees fewer bands — silently returning off-frame + segments would draw lines pinned to the edge.""" + few, _ = band_lines(EULER, detector=(60, 60), pc=(0.5, 0.5, 4.0)) + many, _ = band_lines(EULER, detector=(60, 60), pc=(0.5, 0.5, 0.4)) + assert len(few) < len(many) + + +class TestReflectors: + def test_cubic_set_is_friedel_unique(self): + """4x{111} + 3x{200} + 6x{220} = 13 bands, one per +-g pair.""" + r = cubic_reflectors() + assert len(r) == 13 + dots = np.abs(r.normals @ r.normals.T) + np.fill_diagonal(dots, 0.0) + assert dots.max() < 1 - 1e-9, "a band appears twice as +g and -g" + + def test_widths_and_weights_are_per_reflector(self): + r = cubic_reflectors() + assert r.widths.shape == r.weights.shape == (len(r),) + assert (r.widths > 0).all() + + def test_brightest_is_a_subset(self): + r = cubic_reflectors() + sub = r.brightest(5) + assert len(sub) == 5 + assert sub.weights.min() >= np.sort(r.weights)[::-1][:5].min() - 1e-12 + + def test_brightest_none_is_a_no_op(self): + r = cubic_reflectors() + assert len(r.brightest(None)) == len(r) + assert len(r.brightest(999)) == len(r) + + def test_from_a_real_phase(self): + """A .cif-backed phase gives real reflectors with Bragg-angle widths — + an order of magnitude narrower than the synthetic set's exaggerated + ones, which is why width is a per-reflector property.""" + pytest.importorskip("diffsims") + from diffpy.structure import Atom, Lattice, Structure + from orix.crystal_map import Phase + from spyde.ebsd.bands import reflectors_from_phase + + phase = Phase(name="ni", space_group=225, + structure=Structure(atoms=[Atom("Ni", [0, 0, 0])], + lattice=Lattice(3.52, 3.52, 3.52, + 90, 90, 90))) + r = reflectors_from_phase(phase, min_dspacing=0.8, voltage_kv=20.0) + assert len(r) > 5 + assert np.isfinite(r.weights).all() and r.weights.max() > 0 + assert (r.widths > 0).all() and r.widths.max() < 0.1 + assert np.allclose(np.linalg.norm(r.normals, axis=1), 1.0) + + def test_a_structureless_phase_falls_back(self): + """A bare Phase(space_group=…) has no atoms, so no structure factors — + the wizard must still work, with the generic cubic band set.""" + from orix.crystal_map import Phase + from spyde.ebsd.bands import reflectors_from_phase + r = reflectors_from_phase(Phase(name="p", space_group=225)) + assert len(r) == len(cubic_reflectors()) + + +class TestZoneAxes: + def test_axes_fall_where_bands_cross(self): + segs, _ = band_lines(EULER, detector=DET, pc=PC) + za = zone_axis_points(EULER, detector=DET, pc=PC) + assert len(za) >= 1 + # Every zone axis must sit ON at least two of the drawn band lines. + for p in za: + d = [] + for (x0, y0), (x1, y1) in segs: + v = np.array([x1 - x0, y1 - y0], float) + n = np.array([-v[1], v[0]]) / max(np.linalg.norm(v), 1e-12) + d.append(abs(float((p - np.array([x0, y0])) @ n))) + assert sorted(d)[1] < 1.0, f"zone axis {p} lies on fewer than 2 bands" + + def test_points_stay_on_the_detector(self): + za = zone_axis_points(EULER, detector=DET, pc=PC) + assert ((za >= -0.5) & (za <= np.array([DET[1], DET[0]]) - 0.5)).all() + + +class TestGeometryBasics: + def test_detector_directions_are_unit_vectors(self): + r = detector_directions(DET, PC) + assert r.shape == (DET[0], DET[1], 3) + assert np.allclose(np.linalg.norm(r, axis=-1), 1.0) + + def test_detector_y_is_flipped(self): + """Detector row 0 is the TOP, which is +y in the sample frame. Getting + this backwards mirrors every pattern and is invisible in a symmetric + one — which is why the synthetic data is deliberately asymmetric.""" + r = detector_directions(DET, PC) + assert r[0, DET[1] // 2, 1] > 0 > r[-1, DET[1] // 2, 1] + + def test_wavelength_is_relativistic(self): + # 20 kV: 0.0859 A relativistic, 0.0867 A if the correction is dropped. + assert abs(wavelength(20.0) - 0.0859) < 5e-4 + + def test_reflectors_dataclass_broadcasts_scalars(self): + r = Reflectors(np.eye(3), 1.0, 0.05) + assert r.weights.shape == (3,) and r.widths.shape == (3,) diff --git a/spyde/tests/migrated/test_ebsd_wizard.py b/spyde/tests/migrated/test_ebsd_wizard.py new file mode 100644 index 0000000..f8ee8e3 --- /dev/null +++ b/spyde/tests/migrated/test_ebsd_wizard.py @@ -0,0 +1,287 @@ +"""The staged EBSD-Indexing wizard (``spyde/actions/ebsd_action.py``). + +Handlers are called directly (``fn(session, plot, payload)``) against a real +Qt-free Session, as ``spyde/actions/README.md`` §7 prescribes, and polled with +``_wait`` because each stage hands off to a worker thread. + +Note the toolbar GATE is asserted here too: the action is keyed on the ``EBSD`` +signal type, which only exists when kikuchipy is installed, and both filter +paths have to agree or the button renders and then dispatches into nothing. +""" +from __future__ import annotations + +import time + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") + +from spyde.actions.ebsd_action import ( + DEFAULTS, EbsdWizard, ebsd_build_dictionary, ebsd_refine, ebsd_run, +) +from spyde.data import ebsd_patterns, ground_truth + + +def _wait(pred, timeout=120.0, interval=0.05): + end = time.time() + timeout + while time.time() < end: + if pred(): + return True + time.sleep(interval) + return False + + +@pytest.fixture +def ebsd_session(captured_messages, monkeypatch): + """A Session holding the bundled synthetic EBSD scan. + + Small on purpose (8x8 of 40x40) — the wizard is what is under test, and a + real-size dictionary would make every test a minute long. Accuracy is + covered by ``test_ebsd_indexing`` / ``test_ebsd_refine``. + """ + from spyde.backend.session import Session + import spyde.actions.ebsd_action as ebsd_mod + + # ebsd_action binds `emit` at import (`from ...ipc import emit`), so + # captured_messages' patch of ipc.emit doesn't reach its direct calls — + # the same reason conftest patches session.emit separately. emit_status / + # emit_error resolve `emit` as a module global inside ipc, so those are + # already covered. + monkeypatch.setattr(ebsd_mod, "emit", captured_messages.append) + + session = Session(n_workers=1, threads_per_worker=1) + s = ebsd_patterns(nav=(8, 8), detector=(40, 40)) + session._add_signal(s, source_path=None) + time.sleep(0.8) # let the selector debounce timers fire + yield {"session": session, "signal": s, "truth": ground_truth(s), + "messages": captured_messages, + "trees": session.signal_trees, "plots": session._plots} + session.shutdown() + + +def _signal_plot(session): + """The pattern (non-navigator) plot — the one the caret sits on.""" + for plot in session._plots: + if not getattr(plot, "is_navigator", False): + return plot + return session._plots[0] + + +def _build(ctx, **over): + """Run stage 2 and wait for the wizard to exist.""" + session, plot = ctx["session"], _signal_plot(ctx["session"]) + tree = plot.signal_tree + payload = {"step_deg": 12.0, "background": "dynamic", + "background_sigma": 6.0, "n_bands": 8} + payload.update(over) + ebsd_build_dictionary(session, plot, payload) + assert _wait(lambda: getattr(tree, "_ebsd_wizard", None) is not None), \ + "the dictionary never built" + return session, plot, tree, tree._ebsd_wizard + + +class TestBuildDictionary: + def test_builds_a_wizard_with_a_resident_dictionary(self, ebsd_session): + _s, _p, tree, wiz = _build(ebsd_session) + assert isinstance(wiz, EbsdWizard) + assert len(wiz.indexer) > 10 + assert len(wiz.euler) == len(wiz.indexer) + assert wiz.detector == (40, 40) + + def test_adopts_the_projection_centre_the_data_records(self, ebsd_session): + """The synthetic scan stamps the PC it was rendered with. Without this + the first overlay is drawn with a guessed geometry and every line is + off, which reads as broken indexing.""" + _s, _p, _t, wiz = _build(ebsd_session) + assert np.allclose(wiz.pc, ebsd_session["truth"]["pc"]) + + def test_announces_itself_to_the_caret(self, ebsd_session): + _build(ebsd_session) + msgs = [m for m in ebsd_session["messages"] + if m.get("type") == "ebsd_dictionary_ready"] + assert msgs, "no ebsd_dictionary_ready — the caret stays locked" + assert msgs[-1]["n_orientations"] > 10 + assert len(msgs[-1]["pc"]) == 3 + + def test_the_dictionary_is_filtered_like_the_data(self, ebsd_session): + """Both sides of a cross-correlation must go through the SAME filter. + High-passing only the experimental patterns leaves the dictionary + carrying low frequencies its counterpart no longer has — scores stay + mediocre and it looks like bad indexing (see test_ebsd_indexing).""" + _s, _p, _t, wiz = _build(ebsd_session, background="dynamic", + background_sigma=6.0) + assert wiz.sim_sigma == 6.0 + from spyde.ebsd.bands import simulate_patterns + # A raw simulated pattern carries a big DC term; the dictionary entry + # for the same orientation has been high-passed and no longer does. + raw = simulate_patterns(wiz.euler[0], wiz.reflectors, + wiz.detector, wiz.pc)[0] + assert raw.mean() > 0.1 + _e, score = wiz.indexer.best(wiz.correct( + np.asarray(ebsd_session["signal"].data[0, 0], float))) + assert score > 0.5, f"live match scored only {score:.3f}" + + def test_no_background_means_no_simulated_filter(self, ebsd_session): + _s, _p, _t, wiz = _build(ebsd_session, background="none") + assert wiz.sim_sigma is None + + def test_static_only_does_not_filter_the_dictionary(self, ebsd_session): + """`static` subtracts a DETECTOR artefact, which simulated patterns do + not have — applying it to them would be a different image, not the + same correction.""" + _s, _p, _t, wiz = _build(ebsd_session, background="static") + assert wiz.sim_sigma is None + assert wiz.static_ref is not None + + def test_rebuilding_replaces_the_previous_wizard(self, ebsd_session): + """Otherwise a second Build stacks a second overlay on the pattern and + both redraw on every navigator move.""" + _s, _p, tree, first = _build(ebsd_session) + session, plot = ebsd_session["session"], _signal_plot(ebsd_session["session"]) + ebsd_build_dictionary(session, plot, {"step_deg": 15.0}) + assert _wait(lambda: getattr(tree, "_ebsd_wizard", None) is not first) + assert first._closed and first.overlay is None + + +class TestBandOverlay: + def test_draws_line_segments_for_the_matched_orientation(self, ebsd_session): + _s, _p, _t, wiz = _build(ebsd_session) + ov = wiz.overlay + assert ov is not None, "no band overlay attached" + segs, za = ov._offsets_for(2, 3) + assert segs.ndim == 3 and segs.shape[1:] == (2, 2), \ + "not the (N,2,2) shape anyplotlib add_lines needs" + assert len(segs) > 0, "the matched orientation drew no bands" + assert len(za) == 0, "zone axes are off by default" + + def test_streams_the_match_to_the_caret(self, ebsd_session): + _s, _p, _t, wiz = _build(ebsd_session) + wiz.overlay._offsets_for(2, 3) + hits = [m for m in ebsd_session["messages"] if m.get("type") == "ebsd_match"] + assert hits and hits[-1]["ok"] + assert 0.0 <= hits[-1]["score"] <= 1.0 + assert set(hits[-1]) >= {"phi1", "Phi", "phi2", "score"} + + def test_a_different_position_gives_different_bands(self, ebsd_session): + """The two grains have genuinely different orientations, so an overlay + that ignored the navigator would be caught here.""" + _s, _p, _t, wiz = _build(ebsd_session) + mask = np.asarray(ebsd_session["truth"]["grain2_mask"], bool) + ys, xs = np.nonzero(mask) + ys2, xs2 = np.nonzero(~mask) + a, _ = wiz.overlay._offsets_for(int(ys[0]), int(xs[0])) + b, _ = wiz.overlay._offsets_for(int(ys2[0]), int(xs2[0])) + assert a.shape != b.shape or not np.allclose(a, b) + + def test_hiding_clears_both_marker_groups(self, ebsd_session): + """set_visible(False) pushes a bare array, not the (segments, points) + tuple the overlay normally renders — it has to survive that.""" + _s, _p, _t, wiz = _build(ebsd_session) + wiz.overlay.set_visible(False) + assert wiz.overlay._hidden + wiz.overlay.set_visible(True) + assert not wiz.overlay._hidden + + +class TestRefineStage: + def test_band_count_and_zone_axes_apply_live(self, ebsd_session): + session, plot, _t, wiz = _build(ebsd_session) + ebsd_refine(session, plot, {"n_bands": 3, "show_zone_axes": True}) + assert _wait(lambda: wiz.overlay.n_bands == 3 + and wiz.overlay.show_zone_axes) + segs, za = wiz.overlay._offsets_for(2, 3) + assert len(segs) <= 3 + assert za.ndim == 2 and za.shape[1] == 2 + + def test_moving_the_projection_centre_moves_the_lines(self, ebsd_session): + """The PC is the one parameter you can only set by looking — nudging it + has to redraw, or the Refine tab does nothing.""" + session, plot, _t, wiz = _build(ebsd_session) + before, _ = wiz.overlay._offsets_for(2, 3) + ebsd_refine(session, plot, {"pc_x": 0.62}) + assert _wait(lambda: abs(wiz.overlay.pc[0] - 0.62) < 1e-9) + after, _ = wiz.overlay._offsets_for(2, 3) + assert before.shape != after.shape or not np.allclose(before, after) + assert abs(wiz.pc[0] - 0.62) < 1e-9 + + def test_is_a_no_op_before_the_dictionary_exists(self, ebsd_session): + session = ebsd_session["session"] + ebsd_refine(session, _signal_plot(session), {"n_bands": 3}) # must not raise + + +class TestRunStage: + def test_indexes_the_scan_into_an_ipf_window(self, ebsd_session): + session, plot, tree, _wiz = _build(ebsd_session) + before = len(session.signal_trees) + ebsd_run(session, plot, {"keep": 4, "refine": False}) + assert _wait(lambda: getattr(tree, "orientation_map", None) is not None, + timeout=180), "no orientation map attached" + om = tree.orientation_map + assert om.nav_shape == (8, 8) + assert om.quats.shape == (8, 8, 1, 4) + rgb = om.ipf_color_map("z") + assert rgb.shape == (8, 8, 3) and rgb.dtype == np.uint8 + assert len(session.signal_trees) > before, "no result window opened" + + def test_the_ipf_map_shows_the_two_grains(self, ebsd_session): + """The end-to-end check: the wedge grain must come out a different + colour from the drifting background grain. This is also the only test + that would notice the orientations being handed to orix transposed — + no, it would not; see test_ebsd_bands for that. It notices indexing + that lost the grain structure entirely.""" + session, plot, tree, _wiz = _build(ebsd_session) + ebsd_run(session, plot, {"keep": 4, "refine": False}) + assert _wait(lambda: getattr(tree, "orientation_map", None) is not None, + timeout=180) + rgb = tree.orientation_map.ipf_color_map("z").astype(float) + mask = np.asarray(ebsd_session["truth"]["grain2_mask"], bool) + assert np.linalg.norm(rgb[mask].mean(0) - rgb[~mask].mean(0)) > 20, \ + "the two grains came out the same colour" + + def test_refuses_before_the_dictionary_is_built(self, ebsd_session): + session = ebsd_session["session"] + ebsd_run(session, _signal_plot(session), {}) + assert any("build the dictionary first" in str(m.get("text", "")).lower() + for m in ebsd_session["messages"] if m.get("type") == "error") + + +class TestWiring: + def test_the_schema_is_registered_and_matches_the_defaults(self): + """One source of truth for every host — a schema default that drifts + from the handler default means the caret sends one thing and a + notebook another.""" + from spyde.actions.registry import wizard_parameters + schema = wizard_parameters("ebsd") + assert schema, "the ebsd wizard has no registered schema" + for key, spec in schema.items(): + if key in DEFAULTS: + assert spec["default"] == DEFAULTS[key], \ + f"{key}: schema default {spec['default']!r} != handler " \ + f"default {DEFAULTS[key]!r}" + + def test_every_stage_resolves(self): + from spyde.actions.registry import resolve_staged + for name in ("ebsd_build_dictionary", "ebsd_refine", "ebsd_run"): + assert callable(resolve_staged(name)), f"{name} is not registered" + + def test_the_toolbar_entry_is_gated_on_the_ebsd_signal_type(self): + """Both filter paths apply the same gates in two places; a gate added + to one alone renders a button that never dispatches, or vice versa.""" + import spyde + meta = None + for group in spyde.TOOLBAR_ACTIONS.values(): + if isinstance(group, dict) and "EBSD Indexing" in group: + meta = group["EBSD Indexing"] + assert meta is not None, "EBSD Indexing is not in toolbars.yaml" + assert meta["signal_types"] == ["EBSD"] + assert meta["function"].endswith("ebsd_action.ebsd_indexing") + + def test_the_overlay_toggle_reaches_the_wizard(self, ebsd_session): + """The caret shows/hides the overlay through Session._set_overlay, + which resolves it by ACTION NAME — a name mismatch silently no-ops.""" + session, plot, _t, wiz = _build(ebsd_session) + session._set_overlay(plot, "EBSD Indexing", False) + assert wiz.overlay._hidden + session._set_overlay(plot, "EBSD Indexing", True) + assert not wiz.overlay._hidden diff --git a/spyde/toolbars.yaml b/spyde/toolbars.yaml index 9c362cc..6dd4c80 100644 --- a/spyde/toolbars.yaml +++ b/spyde/toolbars.yaml @@ -294,6 +294,16 @@ functions: type: enum default: intensity options: [intensity, count] + EBSD Indexing: + description: Dictionary indexing of EBSD patterns. Build a simulated pattern dictionary, check the matched orientation's Kikuchi bands live under the crosshair, then index the whole scan into an IPF orientation map. + icon: drawing/toolbars/icons/ebsd.svg + function: spyde.actions.ebsd_action.ebsd_indexing + signal_types: [EBSD] + plot_dim: [2] + toolbar_side: bottom + navigation: False + toggle: True + Vector Orientation Mapping: description: Orientation + strain mapping from diffraction vectors (sparse point-set fit). Live refine under the crosshair. icon: drawing/toolbars/icons/orientation_mapping.svg From 723ba5ef8478d83ebe1257d4cb2835518c5f8f14 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Tue, 28 Jul 2026 06:21:59 -0500 Subject: [PATCH 40/60] fix(ebsd): index the lazy scan chunk-by-chunk, never the whole thing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- spyde/actions/ebsd_action.py | 311 ++++++++++++++++++----- spyde/ebsd/indexing.py | 19 +- spyde/tests/migrated/test_ebsd_wizard.py | 164 ++++++++++++ 3 files changed, 412 insertions(+), 82 deletions(-) diff --git a/spyde/actions/ebsd_action.py b/spyde/actions/ebsd_action.py index e656396..f15849d 100644 --- a/spyde/actions/ebsd_action.py +++ b/spyde/actions/ebsd_action.py @@ -29,6 +29,8 @@ from __future__ import annotations import logging +import os +import threading import numpy as np @@ -253,25 +255,63 @@ def _nav_shape(signal) -> tuple[int, int]: return (int(nav[1]), int(nav[0])) -def _read_scan(signal) -> np.ndarray: - """The whole scan as ``(ny, nx, dy, dx)`` float32. +#: Target bytes per navigation chunk (float32) when the scan arrives EAGER and +#: has to be given a chunking of its own. A lazy signal keeps the chunking it +#: was loaded with. +NAV_BLOCK_BYTES = 256 << 20 # 256 MiB - Deliberately materialised: every step of :mod:`spyde.ebsd` — background - correction, the ADP map, indexing, refinement — is written whole-field, so - there is nothing to stream into. EBSD patterns are small (a 60x60 detector - is 3.6 kB), which is what makes that reasonable where the 4D-STEM memory - rule forbids it. - It is not free, though: at float32 a 256x256 scan of 60x60 patterns is - 940 MB and a 512x512 one is 3.8 GB. Indexing a scan that large wants the - read chunked over navigation and fed to ``dictionary_index`` a block at a - time — the function already tiles internally and takes a ``stopped_flag``, - so the change is here rather than there. +def _nav_block_rows(signal, budget_bytes: int | None = None) -> int: + """Navigation rows per chunk for a scan that has no chunking yet.""" + ny, nx = _nav_shape(signal) + dy, dx = _detector_shape(signal) + per_row = max(1, nx * dy * dx * 4) + budget = int(budget_bytes if budget_bytes is not None else NAV_BLOCK_BYTES) + return int(max(1, min(budget // per_row, ny))) + + +def _lazy_scan(signal): + """The scan as a dask array chunked over NAVIGATION, patterns whole. + + **Never materialise the scan** (CLAUDE.md's memory-safety rule): at float32 + a 512x512 map of 60x60 patterns is 3.8 GB, and a real detector is far + bigger than 60x60. kikuchipy indexes over a chunked dask array for exactly + this reason and so does everything below — every stage is per-pattern + (correction, indexing, refinement) or needs only the four nearest + neighbours (the ADP map), so it is embarrassingly parallel over navigation + chunks and nothing ever needs more than a chunk resident. + + A lazy signal keeps the chunking it was loaded with. An eager one is + wrapped with a nav chunking of its own — patterns are already in RAM there, + so this buys parallelism rather than memory. """ + import dask.array as da + data = signal.data - if hasattr(data, "compute"): - data = data.compute() - return np.asarray(data, np.float32) + if not hasattr(data, "chunks"): + rows = _nav_block_rows(signal) + return da.from_array(data, chunks=(rows, -1, -1, -1)) + # A chunk holding PART of a pattern makes every read useless — the same + # storage-alignment argument as the live-display path (CLAUDE.md §1), and + # the only case where merging chunks is worth the shuffle. EBSD readers do + # not normally split the detector, hence the warning if one has. + if len(data.chunks[2]) > 1 or len(data.chunks[3]) > 1: + log.warning("EBSD patterns are split across chunks %s — merging the " + "signal axes so a chunk holds whole patterns", + (data.chunks[2], data.chunks[3])) + data = data.rechunk({2: -1, 3: -1}) + return data + + +def _scan_mean(scan) -> np.ndarray: + """The mean pattern over the whole scan, as a dask reduction. + + This is the ``static`` background reference (what a flat-field image + approximates) — a whole-scan reduction, which is exactly the kind of thing + that invites reading the whole scan. dask streams it chunk by chunk. + """ + return np.asarray(scan.mean(axis=(0, 1)).compute(scheduler="threads"), + np.float32) def _phase_for(payload, signal): @@ -371,10 +411,11 @@ def _progress(done, total): indexer = SinglePatternIndexer(dic, euler) # The static reference for background correction is the scan mean; - # compute it once here, not per preview frame. + # computed once here (streamed, never the whole scan at once), not + # per preview frame. static_ref = None if background in ("static", "both"): - static_ref = _read_scan(root).mean(axis=(0, 1)) + static_ref = _scan_mean(_lazy_scan(root)) # A rebuilt dictionary replaces the previous wizard wholesale, so # its overlay tears down with it rather than stacking a second set @@ -498,6 +539,145 @@ def _work(): run_on_worker(session, _work, name="ebsd-run") +# torch's CUDA backward is not something to run from several dask threads at +# once on Windows (CLAUDE.md, GPU Computing). Indexing is a forward matmul and +# parallelises freely; refinement takes this lock, which costs nothing real — +# the GPU serialises those kernels anyway. +_REFINE_LOCK = threading.Lock() + + +def _index_chunk(block, wiz, keep, refine, steps, dict_euler): + """One navigation chunk → ``(by, bx, 3 + 1 + keep)``. + + Runs on a dask worker thread with only this chunk's patterns in hand. + Packed into one array rather than returned as a tuple because that is what + ``map_blocks`` can express; float64 so the dictionary indices riding in the + tail stay exact. + """ + from spyde.ebsd.indexing import dictionary_index + + by, bx = block.shape[:2] + out = np.zeros((by, bx, 4 + keep), np.float64) + if by == 0 or bx == 0: + return out + + corrected = wiz.correct(np.asarray(block, np.float32)) + res = dictionary_index(corrected, wiz.indexer, keep=keep) + eul = dict_euler[res.best].reshape(by, bx, 3) + sc = res.best_score.reshape(by, bx) + + if refine: + from spyde.ebsd.refine import refine_orientations + with _REFINE_LOCK: + ref = refine_orientations( + corrected, eul, detector=wiz.detector, pc=wiz.pc, + reflectors=wiz.reflectors, background_sigma=wiz.sim_sigma, + steps=steps) + eul = ref.euler_map((by, bx)) + sc = ref.score.reshape(by, bx) + + out[..., :3] = eul + out[..., 3] = sc + out[..., 4:] = res.indices.reshape(by, bx, keep) + return out + + +def _packed_index_graph(scan, wiz, *, keep, refine, steps): + """The lazy per-chunk index (+ refine) over the whole scan.""" + nav_chunks = (scan.chunks[0], scan.chunks[1]) + return scan.map_blocks( + _index_chunk, wiz, keep, refine, steps, wiz.euler, + dtype=np.float64, drop_axis=[2, 3], new_axis=[2], + chunks=nav_chunks + ((4 + keep,),), + ) + + +def _adp_chunk(block, *, wiz): + """Average dot-product map for one HALOED navigation chunk.""" + from spyde.ebsd.preprocess import average_dot_product_map + + if block.shape[0] == 0 or block.shape[1] == 0: + return np.zeros(block.shape[:2], np.float32) + return average_dot_product_map( + wiz.correct(np.asarray(block, np.float32))).astype(np.float32) + + +def _adp_graph(scan, wiz): + """The lazy ADP map — ``map_overlap`` with a one-position halo, so a + chunk's edge rows still see the neighbours that live in the next chunk. + + The module-level form, with *wiz* bound in: ``Array.map_overlap`` takes + ``depth`` as its second positional argument, so an extra positional here + collides with it. + """ + import dask.array as da + from functools import partial + + return da.map_overlap( + partial(_adp_chunk, wiz=wiz), scan, + depth={0: 1, 1: 1, 2: 0, 3: 0}, boundary="none", trim=True, + dtype=np.float32, drop_axis=[2, 3], + chunks=(scan.chunks[0], scan.chunks[1]), + ) + + +def _compute_nav_chunks(lazy, stopped, *, on_chunk=None, label="computing"): + """Compute a nav-chunked lazy array chunk-by-chunk, in parallel. + + Chunks are independent, so they go to a thread pool and land as they + finish — which is both the parallelism and the progressive fill. dask's + threaded scheduler is used deliberately rather than the session's compute + backend: the dictionary is a torch tensor resident on the LOCAL device, so + shipping these tasks to distributed workers would mean shipping it too. + + ONLY for a graph built with ``map_blocks``, where a nav-chunk slice selects + exactly one block. Slicing a ``map_overlap`` result this way re-derives the + halo per slice and does not line up — compute those whole (see + :func:`_adp_graph`, whose output is one small map either way). + + Returns the assembled array, or None if cancelled. + """ + import itertools + from concurrent.futures import ThreadPoolExecutor, as_completed + + nav_chunks = lazy.chunks[:2] + spans = [] + for axis in nav_chunks: + pos, start = [], 0 + for size in axis: + pos.append((start, size)) + start += size + spans.append(pos) + blocks = list(itertools.product(*spans)) + out = np.zeros(lazy.shape, lazy.dtype) + + workers = max(1, min(8, (os.cpu_count() or 4) - 1)) + done = 0 + with ThreadPoolExecutor(max_workers=workers, + thread_name_prefix="ebsd-chunk") as pool: + futures = {} + for (y0, hy), (x0, wx) in blocks: + sl = (slice(y0, y0 + hy), slice(x0, x0 + wx)) + full = sl + (slice(None),) * (lazy.ndim - 2) + futures[pool.submit(lazy[full].compute, scheduler="threads")] = sl + for fut in as_completed(futures): + sl = futures[fut] + if stopped[0]: + for f in futures: + f.cancel() + return None + chunk = fut.result() + out[sl] = chunk + done += 1 + emit_status(f"EBSD: {label}… {int(100 * done / max(len(blocks), 1))}%") + if on_chunk is not None: + try: + on_chunk(chunk, sl) + except Exception as e: + log.debug("EBSD chunk callback failed: %s", e) + return None if stopped[0] else out + + def _run_indexing(session, tree, wiz, *, keep, refine, steps): """Whole-field index (+ optional refine) → the IPF window. Synchronous; call from a worker.""" @@ -507,8 +687,7 @@ def _run_indexing(session, tree, wiz, *, keep, refine, steps): root = tree.root ny, nx = _nav_shape(root) - scan = wiz.correct(_read_scan(root)) - dict_euler = wiz.euler + scan = _lazy_scan(root) ipf_tree = _open_ipf_window(session, root, ny, nx) # Closing EITHER the source or the result window stops the run — an index @@ -521,53 +700,49 @@ def _run_indexing(session, tree, wiz, *, keep, refine, steps): t.register_cancel(flag=stopped) try: - painter = _ProgressivePainter(session, ipf_tree, wiz, (ny, nx), dict_euler) - # Cap the pattern tile so the map fills in as it goes: left alone the - # tiling puts the whole scan in one chunk and nothing paints until the - # end (see dictionary_index's pattern_chunk). - chunk = max(1, (ny * nx) // 16) - result = dictionary_index( - scan, wiz.indexer, keep=keep, pattern_chunk=chunk, - on_chunk=painter, stopped_flag=stopped, - progress=lambda d, t: emit_status( - f"EBSD: indexing… {int(100 * d / max(t, 1))}%"), - ) - if result is None: + painter = _ProgressivePainter(session, ipf_tree, wiz, (ny, nx), wiz.euler) + # ONE lazy graph over the chunked scan: per nav chunk, correct → + # index → optionally refine, packed into (by, bx, 3 euler + 1 score + + # keep indices). Nothing here reads anything; the patterns only ever + # exist a chunk at a time, inside the worker that is using them. + packed = _packed_index_graph(scan, wiz, keep=keep, refine=refine, + steps=steps) + results = _compute_nav_chunks( + packed, stopped, + on_chunk=lambda blk, sl: painter.paint(blk[..., :3], sl), + label="indexing") + if results is None: if not getattr(tree, "_spyde_closed", False) and not stopped[0]: emit_error("EBSD: indexing returned no result") return None - euler = dict_euler[result.best].reshape(ny, nx, 3) - score = result.best_score.reshape(ny, nx) - - # Refinement is the tail of the SAME run, so it stays inside the - # cancellation registration — unregistering after the index would let - # a window closed during refinement carry on to paint into it. - if refine and not stopped[0]: - emit_status(f"EBSD: refining {ny * nx:,} orientations…") - from spyde.ebsd.refine import refine_orientations - ref = refine_orientations( - scan, euler, detector=wiz.detector, pc=wiz.pc, - reflectors=wiz.reflectors, background_sigma=wiz.sim_sigma, - steps=steps, - progress=lambda d, t: emit_status( - f"EBSD: refining… {int(100 * d / max(t, 1))}%")) - euler = ref.euler_map((ny, nx)) - score = ref.score.reshape(ny, nx) - log.debug("refinement improved %d/%d orientations", - int(ref.improved.sum()), ref.improved.size) + euler = np.ascontiguousarray(results[..., :3], np.float64) + score = np.ascontiguousarray(results[..., 3], np.float32) + indices = results[..., 4:].reshape(ny * nx, -1).astype(np.int64) + + # The ADP map is the ONE stage that is not per-pattern: it compares + # each position with its four neighbours, so it goes through + # map_overlap with a one-position halo rather than plain map_blocks — + # otherwise the rows either side of every chunk boundary would be + # computed against nothing and the map would show faint seams. + # Computed whole (the scheduler still runs the chunks in parallel and + # still only holds a chunk of patterns at a time); its output is one + # small map, so there is nothing to fill in progressively. + emit_status("EBSD: quality maps…") + adp = None + if not stopped[0]: + try: + adp = np.asarray( + _adp_graph(scan, wiz).compute(scheduler="threads"), np.float32) + except Exception as e: + log.debug("ADP map failed: %s", e) if stopped[0]: return None - om = _orientation_map(euler, score, wiz, root) - osm = (orientation_similarity_map(result.indices, (ny, nx)) - if keep > 1 else None) - try: - adp = average_dot_product_map(scan) - except Exception as e: - log.debug("ADP map failed: %s", e) - adp = None + # A neighbour comparison over the FULL map, so it runs once at the end + # — on the accumulated indices, which are small. + osm = orientation_similarity_map(indices, (ny, nx)) if keep > 1 else None finally: for t in trees.values(): if t is not None and hasattr(t, "unregister_cancel"): @@ -582,12 +757,11 @@ def _run_indexing(session, tree, wiz, *, keep, refine, steps): class _ProgressivePainter: - """Paint the IPF map as indexing completes each slice of patterns. + """Paint the IPF map as each navigation chunk finishes indexing. - A dictionary index reports per chunk of PATTERNS, which in scan order is a - contiguous run of positions — so the map fills top to bottom and you can - see the grain structure emerge instead of watching a blank window. Costs - one IPF colouring per chunk over the positions done so far. + Chunks land as they complete, so the map fills in and you watch the grain + structure emerge instead of a blank window. Costs one IPF colouring per + chunk, over that chunk's positions only. """ def __init__(self, session, ipf_tree, wiz, nav_shape, dict_euler): @@ -608,15 +782,18 @@ def __init__(self, session, ipf_tree, wiz, nav_shape, dict_euler): self._pg = wiz.phase.point_group self._key = IPFColorKeyTSL(self._pg.laue) - def __call__(self, lo, hi, indices, scores) -> None: + def paint(self, euler, nav_slice) -> None: + """Colour one finished navigation chunk into the map and push it.""" if self.plot is None: return try: from orix.quaternion import Orientation, Rotation - rot = Rotation.from_euler(self.dict_euler[indices[:, 0]]) + eul = np.asarray(euler, float).reshape(-1, 3) + rot = Rotation.from_euler(eul) rgb = self._key.orientation2color(Orientation(rot, symmetry=self._pg)) - self.rgb.reshape(-1, 3)[lo:hi] = np.clip( - np.asarray(rgb) * 255.0, 0, 255).astype(np.uint8) + self.rgb[nav_slice] = np.clip( + np.asarray(rgb) * 255.0, 0, 255).astype(np.uint8).reshape( + self.rgb[nav_slice].shape) except Exception as e: log.debug("progressive IPF colouring failed: %s", e) return diff --git a/spyde/ebsd/indexing.py b/spyde/ebsd/indexing.py index 7f0cb69..139376f 100644 --- a/spyde/ebsd/indexing.py +++ b/spyde/ebsd/indexing.py @@ -69,8 +69,6 @@ def dictionary_index(patterns, dictionary, *, keep: int = 1, device=None, dtype="float32", tile_elements: int = _TILE_ELEMENTS, pattern_chunk: int | None = None, progress: Callable[[int, int], None] | None = None, - on_chunk: Callable[[int, int, np.ndarray, np.ndarray], - None] | None = None, stopped_flag=None): """Match every pattern against every dictionary entry. @@ -92,14 +90,10 @@ def dictionary_index(patterns, dictionary, *, keep: int = 1, device=None, conditioned quantity and the matmul is memory-bound, so float64 would halve throughput to protect digits that do not affect the ranking. pattern_chunk : int, optional - Cap on the experimental tile size. The tiling *alone* would put every - pattern in one chunk whenever the dictionary is small enough — correct - and fastest, but it means *progress* and *on_chunk* fire exactly once, - at the end. A UI filling a map in as it indexes passes this to get - intermediate results; nothing else should. - on_chunk : callable, optional - ``(lo, hi, indices, scores)`` for each completed slice of patterns, so - a caller can paint partial results. Called on the calling thread. + Cap on the experimental tile size, on top of the *tile_elements* + budget. Callers that already stream the scan in navigation blocks + (``spyde.actions.ebsd_action``) do not need it; it is here for a caller + holding a large in-memory stack that wants a tighter bound. stopped_flag : list[bool], optional Polled between chunks; a truthy first element abandons the run and returns None. This is how closing the window cancels an index. @@ -177,11 +171,6 @@ def dictionary_index(patterns, dictionary, *, keep: int = 1, device=None, out_scores[p0:p1] = best_s.detach().cpu().numpy() out_index[p0:p1] = best_i.detach().cpu().numpy() - if on_chunk is not None: - try: - on_chunk(p0, p1, out_index[p0:p1], out_scores[p0:p1]) - except Exception as e: # pragma: no cover - log.debug("indexing on_chunk callback failed: %s", e) if progress is not None: try: progress(p1, P) diff --git a/spyde/tests/migrated/test_ebsd_wizard.py b/spyde/tests/migrated/test_ebsd_wizard.py index f8ee8e3..5ad3dfa 100644 --- a/spyde/tests/migrated/test_ebsd_wizard.py +++ b/spyde/tests/migrated/test_ebsd_wizard.py @@ -246,6 +246,170 @@ def test_refuses_before_the_dictionary_is_built(self, ebsd_session): for m in ebsd_session["messages"] if m.get("type") == "error") +class TestNeverMaterialisesTheScan: + """Memory safety (CLAUDE.md rule #1) — the scan is read in navigation + BLOCKS and the whole thing is never resident. + + An EBSD pattern is small, which makes "just load it" tempting and wrong: + at float32 a 512x512 map of 60x60 patterns is 3.8 GB, and a real detector + is far bigger than 60x60. kikuchipy indexes over a chunked dask array for + the same reason. + + This is guarded rather than reasoned about, because the failure is silent + until someone opens a big scan: the guard below raises if anything ever + computes a dask array the shape of the full dataset. + """ + + NAV, DET = (8, 8), (40, 40) + + @pytest.fixture + def lazy_session(self, captured_messages, monkeypatch): + from spyde.backend.session import Session + import spyde.actions.ebsd_action as ebsd_mod + + monkeypatch.setattr(ebsd_mod, "emit", captured_messages.append) + + session = Session(n_workers=1, threads_per_worker=1) + s = ebsd_patterns(nav=self.NAV, detector=self.DET).as_lazy() + # Several nav chunks, so a correct implementation never asks for a + # slice the shape of the whole scan (which on a single-chunk 8-row scan + # it legitimately would, and the guard could not tell the difference). + s.data = s.data.rechunk((2, 4, -1, -1)) + session._add_signal(s, source_path=None) + time.sleep(0.8) + yield {"session": session, "signal": s, "messages": captured_messages, + "trees": session.signal_trees, "plots": session._plots} + session.shutdown() + + @pytest.fixture + def no_full_compute(self, monkeypatch): + """Raise if any dask array the shape of the whole dataset is computed.""" + import dask.array as da + + full = tuple(self.NAV) + tuple(self.DET) + original = da.Array.compute + seen: list[tuple] = [] + + def guard(self, *args, **kwargs): + if tuple(self.shape) == full: + raise AssertionError( + f"computed the FULL dataset {full} — the scan must be read " + f"in navigation blocks (CLAUDE.md memory-safety rule)") + seen.append(tuple(self.shape)) + return original(self, *args, **kwargs) + + monkeypatch.setattr(da.Array, "compute", guard) + return seen + + def test_indexing_reads_the_scan_chunk_by_chunk(self, lazy_session, + no_full_compute): + session, plot = lazy_session["session"], _signal_plot(lazy_session["session"]) + tree = plot.signal_tree + ebsd_build_dictionary(session, plot, {"step_deg": 12.0, + "background": "dynamic"}) + assert _wait(lambda: getattr(tree, "_ebsd_wizard", None) is not None) + + ebsd_run(session, plot, {"keep": 4, "refine": False}) + assert _wait(lambda: getattr(tree, "orientation_map", None) is not None, + timeout=180), "indexing never finished" + assert tree.orientation_map.nav_shape == self.NAV + + def test_the_static_reference_is_streamed_too(self, lazy_session, + no_full_compute): + """The static background reference is a whole-scan MEAN — exactly the + kind of reduction that invites reading everything to compute it.""" + session, plot = lazy_session["session"], _signal_plot(lazy_session["session"]) + tree = plot.signal_tree + ebsd_build_dictionary(session, plot, {"step_deg": 12.0, + "background": "both"}) + assert _wait(lambda: getattr(tree, "_ebsd_wizard", None) is not None) + wiz = tree._ebsd_wizard + assert wiz.static_ref is not None + assert wiz.static_ref.shape == self.DET + + # And it is the real WHOLE-scan mean, not one block's. The reference + # comes from a fresh eager copy (the generator is deterministic) rather + # than from computing the lazy signal — which the guard above would, + # correctly, refuse. + expected = np.asarray(ebsd_patterns(nav=self.NAV, detector=self.DET).data, + np.float32).mean(axis=(0, 1)) + assert np.allclose(wiz.static_ref, expected, atol=1e-3) + + def test_the_adp_map_is_seamless_across_chunks(self): + """Everything else in the run is per-pattern, so chunking cannot change + it. The ADP map is the exception — it compares each position with its + four NEIGHBOURS, so the rows either side of every chunk boundary would + be computed against nothing above/below them without the halo that + ``map_overlap`` adds. That shows up as faint lines at the seams, which + is easy to mistake for real structure. + """ + import dask.array as da + from spyde.actions.ebsd_action import _adp_graph + from spyde.ebsd import average_dot_product_map + + s = ebsd_patterns(nav=(8, 8), detector=self.DET) + raw = np.asarray(s.data, np.float32) + + class _Passthrough: # the halo, not the correction, is the test + background = "none" + + def correct(self, patterns): + return np.asarray(patterns, float) + + whole = average_dot_product_map(raw) + for chunk in (8, 4, 3, 1): # 1 chunk, exact split, ragged, per-row + scan = da.from_array(raw, chunks=(chunk, chunk, -1, -1)) + got = _adp_graph(scan, _Passthrough()).compute(scheduler="threads") + assert got.shape == whole.shape + assert np.allclose(got, whole, atol=1e-5), \ + f"the ADP map has seams at {chunk}-row chunk boundaries" + + def test_indexing_is_one_lazy_graph_over_the_chunks(self): + """The index is a dask graph, so nothing is read by BUILDING it — the + patterns only exist inside a worker, a chunk at a time.""" + import dask.array as da + from spyde.actions.ebsd_action import _lazy_scan, _packed_index_graph + + s = ebsd_patterns(nav=(8, 8), detector=self.DET).as_lazy() + scan = _lazy_scan(s) + assert isinstance(scan, da.Array) + + class _Wiz: + euler = np.zeros((2, 3)) + + graph = _packed_index_graph(scan, _Wiz(), keep=3, refine=False, steps=1) + assert isinstance(graph, da.Array) + assert graph.shape == (8, 8, 4 + 3) + # Nav chunking is carried through, so the result assembles per chunk. + assert graph.chunks[:2] == scan.chunks[:2] + + def test_eager_data_still_gets_a_bounded_chunking(self): + """An eager scan has no chunks of its own; it must still be given a + nav chunking, sized by BYTES, so a bigger detector means fewer rows + rather than the same rows and more memory.""" + from spyde.actions.ebsd_action import _lazy_scan, _nav_block_rows + small = ebsd_patterns(nav=(64, 64), detector=(40, 40)) + big = ebsd_patterns(nav=(64, 64), detector=(80, 80)) + budget = 4 << 20 + assert _nav_block_rows(big, budget) < _nav_block_rows(small, budget) + assert _nav_block_rows(small, 1 << 40) == 64 # never more than ny + scan = _lazy_scan(small) + assert len(scan.chunks[2]) == 1 and len(scan.chunks[3]) == 1, \ + "patterns must not be split across chunks" + + def test_patterns_split_across_chunks_are_merged(self): + """A chunk holding PART of a pattern makes every read useless. EBSD + readers do not normally split the detector, but if one does the graph + has to fix it rather than index garbage.""" + import dask.array as da + from spyde.actions.ebsd_action import _lazy_scan + + s = ebsd_patterns(nav=(4, 4), detector=(40, 40)).as_lazy() + s.data = da.from_array(np.asarray(s.data), chunks=(2, 2, 20, 20)) + scan = _lazy_scan(s) + assert len(scan.chunks[2]) == 1 and len(scan.chunks[3]) == 1 + + class TestWiring: def test_the_schema_is_registered_and_matches_the_defaults(self): """One source of truth for every host — a schema default that drifts From f2d88868cadfb2ec82bf0d78a90a6b9a3b8dc992 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Tue, 28 Jul 2026 08:10:52 -0500 Subject: [PATCH 41/60] refactor(ebsd): chunk through hyperspy rechunk, not by hand MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _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. --- spyde/actions/ebsd_action.py | 44 +++++++++++++++--------- spyde/tests/migrated/test_ebsd_wizard.py | 26 ++++++++++++-- 2 files changed, 50 insertions(+), 20 deletions(-) diff --git a/spyde/actions/ebsd_action.py b/spyde/actions/ebsd_action.py index f15849d..3997465 100644 --- a/spyde/actions/ebsd_action.py +++ b/spyde/actions/ebsd_action.py @@ -281,26 +281,36 @@ def _lazy_scan(signal): neighbours (the ADP map), so it is embarrassingly parallel over navigation chunks and nothing ever needs more than a chunk resident. - A lazy signal keeps the chunking it was loaded with. An eager one is - wrapped with a nav chunking of its own — patterns are already in RAM there, - so this buys parallelism rather than memory. + Chunking is hyperspy's own ``LazySignal.rechunk``, which is what + ``BaseSignal.map`` uses to prepare its input — it takes navigation and + signal chunks separately off the axes manager, so this does not have to + know which array axes are which: + + * ``sig_chunks=-1`` — a chunk holds WHOLE patterns. A chunk holding part + of one makes every read useless (the storage-alignment argument from the + live-display notes) and is the one case where merging is worth the + shuffle. EBSD readers do not normally split the detector, hence the + warning when one has. + * ``nav_chunks=None`` for an already-lazy signal — keep the chunking it was + loaded with rather than reshuffling a multi-GB array to a theoretical + optimum. An eager one gets a nav chunking sized by the byte budget; + its patterns are already in RAM, so that buys parallelism, not memory. + + ``inplace=False`` throughout: this must not repoint the live signal's own + chunking, which the navigator reads through. """ - import dask.array as da - - data = signal.data - if not hasattr(data, "chunks"): - rows = _nav_block_rows(signal) - return da.from_array(data, chunks=(rows, -1, -1, -1)) - # A chunk holding PART of a pattern makes every read useless — the same - # storage-alignment argument as the live-display path (CLAUDE.md §1), and - # the only case where merging chunks is worth the shuffle. EBSD readers do - # not normally split the detector, hence the warning if one has. - if len(data.chunks[2]) > 1 or len(data.chunks[3]) > 1: + lazy = signal if getattr(signal, "_lazy", False) else signal.as_lazy() + chunks = getattr(lazy.data, "chunks", None) + sig_split = chunks and any(len(c) > 1 for c in chunks[-2:]) + if sig_split: log.warning("EBSD patterns are split across chunks %s — merging the " "signal axes so a chunk holds whole patterns", - (data.chunks[2], data.chunks[3])) - data = data.rechunk({2: -1, 3: -1}) - return data + tuple(chunks[-2:])) + # Array order: (ny, nx) rows-first, matching signal.data. + nav_chunks = None if getattr(signal, "_lazy", False) else ( + _nav_block_rows(signal), -1) + return lazy.rechunk(nav_chunks=nav_chunks, sig_chunks=-1, + inplace=False).data def _scan_mean(scan) -> np.ndarray: diff --git a/spyde/tests/migrated/test_ebsd_wizard.py b/spyde/tests/migrated/test_ebsd_wizard.py index 5ad3dfa..cf446f1 100644 --- a/spyde/tests/migrated/test_ebsd_wizard.py +++ b/spyde/tests/migrated/test_ebsd_wizard.py @@ -401,14 +401,34 @@ def test_patterns_split_across_chunks_are_merged(self): """A chunk holding PART of a pattern makes every read useless. EBSD readers do not normally split the detector, but if one does the graph has to fix it rather than index garbage.""" - import dask.array as da from spyde.actions.ebsd_action import _lazy_scan - s = ebsd_patterns(nav=(4, 4), detector=(40, 40)).as_lazy() - s.data = da.from_array(np.asarray(s.data), chunks=(2, 2, 20, 20)) + s = ebsd_patterns(nav=(8, 8), detector=(40, 40)).as_lazy() + s.rechunk(nav_chunks=(4, 4), sig_chunks=(20, 20), inplace=True) scan = _lazy_scan(s) assert len(scan.chunks[2]) == 1 and len(scan.chunks[3]) == 1 + def test_the_lazy_signals_own_chunking_is_left_alone(self): + """Two things at once, both load-bearing. + + The nav chunking a lazy signal arrived with is KEPT — reshuffling a + multi-GB array to a theoretical optimum is the mistake CLAUDE.md's + live-display notes record (419 s against 184 s), and storage alignment + beats any after-the-fact rechunk. + + And the signal itself is not repointed: the navigator reads through + the same array, so preparing the compute must not mutate it. + """ + from spyde.actions.ebsd_action import _lazy_scan + + s = ebsd_patterns(nav=(8, 8), detector=(40, 40)).as_lazy() + s.rechunk(nav_chunks=(4, 4), sig_chunks=(20, 20), inplace=True) + before = s.data.chunks + + scan = _lazy_scan(s) + assert s.data.chunks == before, "the live signal's chunking was mutated" + assert scan.chunks[:2] == before[:2], "the nav chunking was reshuffled" + class TestWiring: def test_the_schema_is_registered_and_matches_the_defaults(self): From 98f03a13fc9157ff84780059bf5b4d0967c7c42b Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Tue, 28 Jul 2026 11:05:48 -0500 Subject: [PATCH 42/60] feat(examples): the Examples menu, from em-database MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The examples were a flat list of six pyxem loader names with nothing else on the row — no size, no shape, no way to tell a 42 kB line scan from a 2.4 GB in-situ movie until you clicked one and waited. They now come from em-database, which is one curated, checksummed, cited collection with real acquisition metadata, and adding a dataset is a YAML file upstream rather than a loader here. Examples is now: 4D-STEM 4/9 > AlNanocrystals 256x256 | 128x128 1.4 GB EELS 0/5 > SPEDAg 208x64 | 112x112 214.0 MB EBSD 1/2 > ... ... Dummy Data no download > Show Example Data Directory One submenu per technique (em-database already groups them), each row carrying its shape, its size and a filled/hollow dot for on-disk vs will-download — distinguished by glyph AND colour, so it survives greyscale and colour-blind readers. Hovering gives a themed card (the dropdowns' own surface, not the OS `title=` bubble) with the technique, microscope, facts and description. MenuBar grew real fly-out submenus to do it; Dummy Data was a HEADER in the flat list and is now a submenu of its own, still instant and no-download. Shape is the one thing em-database does not carry for every dataset. Declared shape wins when upstream has one; otherwise it is read from the file once it is downloaded and cached, and stamped for free whenever a dataset is loaded. Reading them inline made opening the menu a 7-second click, so the catalogue never opens a file (1 ms) and a warm pass does the reading off that path. em-database also exposes pooch's `progressbar`, so the download toast now hands it our monitor directly and the proxy that swapped out pyxem's view of the pooch module to reach the same hook is gone. Depended on with a python_version marker: em-database 0.1.0 requires >=3.12 while SpyDE supports >=3.10, and an unconditional dependency made the project unresolvable there — the backend refused to start at all. Caught by running the app, not by any test. On 3.10/3.11 the menu falls back to Dummy Data. --- electron/src/main/index.ts | 25 +- electron/src/preload/index.ts | 6 + .../src/renderer/src/components/MenuBar.tsx | 369 +++++++++++++++--- electron/src/renderer/src/electron.d.ts | 2 + .../src/renderer/src/kernel/SpyDEContext.tsx | 11 + electron/src/renderer/src/kernel/protocol.ts | 10 + electron/tests/examples_menu.spec.ts | 139 +++++++ pyproject.toml | 13 + spyde/backend/_session_actions.py | 4 + spyde/backend/_session_files.py | 116 ++++-- spyde/backend/example_catalogue.py | 369 ++++++++++++++++++ spyde/backend/example_download.py | 60 ++- .../tests/migrated/test_example_catalogue.py | 147 +++++++ spyde/tests/migrated/test_example_download.py | 40 +- spyde/tests/migrated/test_lazy_loading.py | 40 +- uv.lock | 16 + 16 files changed, 1214 insertions(+), 153 deletions(-) create mode 100644 electron/tests/examples_menu.spec.ts create mode 100644 spyde/backend/example_catalogue.py create mode 100644 spyde/tests/migrated/test_example_catalogue.py diff --git a/electron/src/main/index.ts b/electron/src/main/index.ts index 5799b9b..81025e2 100644 --- a/electron/src/main/index.ts +++ b/electron/src/main/index.ts @@ -5,10 +5,10 @@ import { app, BrowserWindow, dialog, ipcMain, Menu, shell, nativeTheme, net, protocol, clipboard, nativeImage, } from 'electron' -import { join, basename } from 'path' +import { join, basename, resolve } from 'path' import { pathToFileURL } from 'url' import { tmpdir } from 'os' -import { existsSync, realpathSync, writeFileSync, rmSync } from 'fs' +import { existsSync, realpathSync, statSync, writeFileSync, rmSync } from 'fs' import { execFile } from 'child_process' import { startSpyDE, sendAction, sendFigureEvent, sendResize, @@ -854,6 +854,27 @@ ipcMain.on('open-external', (_, url: string) => { shell.openExternal(parsed.href) }) +// Reveal a LOCAL DIRECTORY in the OS file manager (Examples → Show Example Data +// Directory). Deliberately separate from open-external, which allowlists +// web/mail protocols precisely so it can never be talked into opening a local +// path: this one takes a filesystem path and opens it as a FOLDER only. +// `openPath` on a directory opens it in the file manager; it will not execute a +// file, and the isDirectory check means a path that somehow named an executable +// is refused rather than run. +ipcMain.on('open-path', async (_, target: string) => { + try { + const resolved = resolve(String(target ?? '')) + if (!existsSync(resolved) || !statSync(resolved).isDirectory()) { + console.warn(`[spyde] open-path rejected non-directory: ${resolved}`) + return + } + const err = await shell.openPath(resolved) + if (err) console.warn(`[spyde] open-path failed: ${err}`) + } catch (e) { + console.warn(`[spyde] open-path error: ${String(e)}`) + } +}) + // ── Update / GPU-status IPC ─────────────────────────────────────────────────── /** Manual "Check Now" from the update dialog. Result arrives async via the diff --git a/electron/src/preload/index.ts b/electron/src/preload/index.ts index 7bc2dd4..fd72034 100644 --- a/electron/src/preload/index.ts +++ b/electron/src/preload/index.ts @@ -170,6 +170,12 @@ contextBridge.exposeInMainWorld('electron', { openExternal: (url: string) => ipcRenderer.send('open-external', url), + /** Reveal a local DIRECTORY in the OS file manager (Examples → Show Example + * Data Directory). Separate from openExternal, which allowlists web/mail + * protocols so it can never open a local path; main verifies the target is + * a real directory. */ + openPath: (path: string) => ipcRenderer.send('open-path', path), + // ── Updates / GPU status ────────────────────────────────────────────────── /** Current channel, whether this build supports auto-update, last known diff --git a/electron/src/renderer/src/components/MenuBar.tsx b/electron/src/renderer/src/components/MenuBar.tsx index baaf25e..3e553af 100644 --- a/electron/src/renderer/src/components/MenuBar.tsx +++ b/electron/src/renderer/src/components/MenuBar.tsx @@ -13,20 +13,10 @@ import React, { useEffect, useRef, useState } from 'react' import { useSpyDE } from '../kernel/SpyDEContext' import { GUIDES, type Guide } from '@guides/index' -const EXAMPLES = [ - 'mgo_nanocrystals', - 'small_ptychography', - 'zrnb_precipitate', - 'pdcusi_insitu', - 'sped_ag', - 'fe_multi_phase_grains', -] - -// Curated, ALWAYS-AVAILABLE small in-memory tutorial datasets (Phase 1 of the -// docs/walkthroughs overhaul) — unlike EXAMPLES (which download real files via -// pyxem+pooch), these fire the ungated `tutorial_load` backend action and load -// in a couple of seconds with no network. Keys mirror -// spyde/backend/tutorial_data.py TUTORIAL_LOADERS; Phase 2+ guided walkthroughs +// Curated, ALWAYS-AVAILABLE small in-memory tutorial datasets — unlike the real +// examples (downloaded via em-database), these fire the ungated `tutorial_load` +// backend action and load in a couple of seconds with no network. Keys mirror +// spyde/backend/tutorial_data.py TUTORIAL_LOADERS; the guided walkthroughs // drive the same action names. const TUTORIAL_DATA: { key: string; label: string }[] = [ { key: 'navigation', label: 'Navigation & Virtual Imaging' }, @@ -38,15 +28,94 @@ const TUTORIAL_DATA: { key: string; label: string }[] = [ { key: 'movie', label: 'In-situ Movie' }, ] +/** One dataset row of the Examples catalogue (spyde/backend/example_catalogue). */ +export interface ExampleItem { + key: string + label: string + technique: string + size: string + shape: string | null + downloaded: boolean + description?: string + file?: string + microscope?: string + voltage?: string +} +interface ExampleGroup { technique: string; items: ExampleItem[] } + +/** The themed hover card's contents. Replaces the native `title=` tooltip, + * which renders as an unstyled OS bubble in the wrong font, colour and + * position — jarring against the rest of the app. */ +export interface Tip { + title: string + /** Small caps line under the title: technique · microscope · voltage. */ + meta?: string + /** Facts worth scanning: shape, size, file. */ + facts?: [string, string][] + body?: string + /** Download state, shown with the same dot the row uses. */ + downloaded?: boolean +} + type Item = - | { label: string; onClick: () => void; disabled?: boolean; testId?: string } + | { + label: string + onClick: () => void + disabled?: boolean + testId?: string + /** Right-aligned detail (dataset size). */ + detail?: string + /** Middle column: nav|signal shape, blank when not yet known. */ + shape?: string + /** Present on dataset rows: already on disk, or a download. */ + downloaded?: boolean + tip?: Tip + } + | { label: string; submenu: Item[]; testId?: string; detail?: string } | { separator: true } | { header: string } +const hasSubmenu = (it: Item): it is { label: string; submenu: Item[]; testId?: string; detail?: string } => + 'submenu' in it + +/** Filled dot = on disk, hollow = will download. Glyphs rather than colour so + * the distinction survives a screenshot, a colour-blind reader and a theme. */ +const DOWNLOADED_MARK = '●' // ● +const NOT_DOWNLOADED_MARK = '○' // ○ + export function MenuBar({ onStartGuide }: { onStartGuide: (g: Guide) => void }) { const { sendAction, openStackDialog, openUpdateDialog, openGpuStatusDialog, openGpuHelpDialog, state } = useSpyDE() const [open, setOpen] = useState(null) const barRef = useRef(null) + const [exampleGroups, setExampleGroups] = useState([]) + const [dataDir, setDataDir] = useState('') + const [examplesReady, setExamplesReady] = useState(false) + const [tip, setTip] = useState<{ tip: Tip; anchor: DOMRect } | null>(null) + + // The card belongs to whatever menu is open; closing one must not leave it + // stranded on screen. + useEffect(() => { if (!open) setTip(null) }, [open]) + + // The catalogue is cheap to build (it opens no files), so ask for a fresh one + // every time the Examples menu opens — that is what keeps the downloaded + // markers honest after a download finishes or a file is deleted outside the + // app. `warm` lets the backend fill in shapes for anything downloaded but not + // yet measured, off this path, and re-send when it has. + useEffect(() => { + if (open !== 'Examples') return + sendAction('example_catalogue', { warm: true }) + }, [open, sendAction]) + + useEffect(() => { + const onCatalogue = (e: Event) => { + const d = (e as CustomEvent).detail as Record + setExampleGroups((d.groups as ExampleGroup[]) ?? []) + setDataDir(String(d.data_dir ?? '')) + setExamplesReady(true) + } + window.addEventListener('spyde:example_catalogue', onCatalogue) + return () => window.removeEventListener('spyde:example_catalogue', onCatalogue) + }, []) // Close on outside click / Escape. useEffect(() => { @@ -74,20 +143,62 @@ export function MenuBar({ onStartGuide }: { onStartGuide: (g: Guide) => void }) { label: 'Quit', onClick: () => window.electron.quit() }, ], Examples: [ - // Real downloadable datasets (pyxem+pooch). - ...EXAMPLES.map((name) => ({ - label: name, - onClick: () => sendAction('load_example', { name }), + // Real datasets from em-database, one submenu per technique. Each row + // shows its size, its shape where known, and whether it is already on + // disk — so you can tell a 42 kB line scan from a 2.4 GB in-situ movie, + // and an instant load from a download, BEFORE clicking. + ...exampleGroups.map((g) => ({ + label: g.technique, + testId: `examples-tech-${g.technique.replace(/[^a-z0-9]+/gi, '-').toLowerCase()}`, + detail: `${g.items.filter((i) => i.downloaded).length}/${g.items.length}`, + submenu: g.items.map((it) => ({ + label: it.label, + testId: `example-${it.key}`, + downloaded: it.downloaded, + shape: it.shape ?? '', + detail: it.size, + tip: { + title: it.label, + meta: [it.technique, it.microscope, it.voltage].filter(Boolean).join(' · '), + facts: [ + ...(it.shape ? [['Shape', it.shape] as [string, string]] : []), + ['Size', it.size] as [string, string], + ...(it.file ? [['File', it.file] as [string, string]] : []), + ], + body: it.description, + downloaded: it.downloaded, + }, + onClick: () => sendAction('load_example', { name: it.key }), + })), })), - // Instant, no-download tutorial datasets — grouped under a "Dummy Data" - // header inside Examples (the menu has no nested fly-outs). + ...(exampleGroups.length === 0 + ? [{ + label: examplesReady ? 'No example datasets found' : 'Loading examples…', + disabled: true, + onClick: () => {}, + } as Item] + : []), + // Instant, no-download synthetic datasets — their own submenu, so they + // read as a distinct kind of thing rather than a tail of the real data. { separator: true } as Item, - { header: 'Dummy Data' } as Item, - ...TUTORIAL_DATA.map(({ key, label }) => ({ - label, - testId: `tutorial-${key}`, - onClick: () => sendAction('tutorial_load', { name: key }), - })), + { + label: 'Dummy Data', + testId: 'examples-dummy-data', + detail: 'no download', + submenu: TUTORIAL_DATA.map(({ key, label }) => ({ + label, + testId: `tutorial-${key}`, + onClick: () => sendAction('tutorial_load', { name: key }), + })), + } as Item, + { separator: true } as Item, + { + label: 'Show Example Data Directory', + testId: 'examples-show-dir', + detail: dataDir ? undefined : '', + title: dataDir || undefined, + onClick: () => sendAction('show_example_dir', {}), + } as Item, ], Help: [ ...GUIDES.map((g) => ({ @@ -125,31 +236,135 @@ export function MenuBar({ onStartGuide }: { onStartGuide: (g: Guide) => void }) {name} {open === name && ( -
e.stopPropagation()}> - {menus[name].map((it, i) => - 'separator' in it ? ( -
- ) : 'header' in it ? ( -
{it.header}
- ) : ( - - ), - )} -
+ setOpen(null)} + onTip={(t, rect) => setTip(t && rect ? { tip: t, anchor: rect } : null)} + testId={`menu-${name.replace(/[^a-z0-9]+/gi, '-').toLowerCase()}-items`} /> )}
))} + {tip && } + + ) +} + +/** + * The themed hover card — SpyDE's own panel chrome (the dropdown's surface, + * border, radius and shadow), not the OS's `title=` bubble. + * + * Anchored to the RIGHT of the hovered row so it never covers the menu you are + * reading, and flipped to the left when there is no room. `pointerEvents:none` + * so it can't itself steal the hover that is keeping it open. + */ +function HoverCard({ tip, anchor }: { tip: Tip; anchor: DOMRect }) { + const W = 300 + const GAP = 10 + const right = anchor.right + GAP + const flip = right + W > window.innerWidth + const left = flip ? Math.max(8, anchor.left - W - GAP) : right + const top = Math.min(anchor.top - 4, Math.max(8, window.innerHeight - 260)) + return ( +
+
{tip.title}
+ {tip.meta &&
{tip.meta}
} + {tip.facts?.length ? ( +
+ {tip.facts.map(([k, v]) => ( +
+ {k} + {v} +
+ ))} +
+ ) : null} + {tip.body &&
{tip.body}
} + {tip.downloaded !== undefined && ( +
+ {tip.downloaded + ? `${DOWNLOADED_MARK} On disk — opens immediately` + : `${NOT_DOWNLOADED_MARK} Not downloaded — click to fetch`} +
+ )} +
+ ) +} + +/** + * One dropdown level. Recurses for fly-out submenus (Examples groups its + * datasets by technique), which the menu did not have before — the old + * "Dummy Data" HEADER inside the flat Examples list is now a real submenu. + * + * A submenu opens on hover and stays open while the pointer is anywhere in the + * parent row or the fly-out itself, so the diagonal travel from the row to the + * panel doesn't dismiss it. It is rendered inside the parent row's relatively + * positioned wrapper, so it follows the row. + */ +function MenuList({ items, onClose, testId, nested = false, onTip }: { + items: Item[] + onClose: () => void + testId?: string + nested?: boolean + onTip?: (tip: Tip | null, anchor?: DOMRect) => void +}) { + const [openSub, setOpenSub] = useState(null) + return ( +
e.stopPropagation()} + > + {items.map((it, i) => { + if ('separator' in it) return
+ if ('header' in it) return
{it.header}
+ + const tid = it.testId ?? `menu-item-${it.label.replace(/[^a-z0-9]+/gi, '-').toLowerCase()}` + if (hasSubmenu(it)) { + return ( +
setOpenSub(it.label)} + onMouseLeave={() => setOpenSub((s) => (s === it.label ? null : s))}> + + {openSub === it.label && ( + + )} +
+ ) + } + + return ( + + ) + })}
) } @@ -174,6 +389,62 @@ const styles: Record = { borderRadius: 5, padding: '6px 10px', fontSize: 12.5, cursor: 'pointer', whiteSpace: 'nowrap', }, + submenu: { + position: 'absolute', top: -5, left: '100%', zIndex: 9300, + minWidth: 260, background: '#1e1e2e', border: '1px solid #313244', + borderRadius: 8, padding: 5, boxShadow: '0 10px 28px rgba(0,0,0,0.5)', + maxHeight: '70vh', overflowY: 'auto', + }, + // Rows are a 3-part grid: marker, label, right-aligned detail — so the sizes + // and shapes line up in a readable column instead of trailing the names. + row: { display: 'flex', alignItems: 'center', gap: 10 }, + // On disk = filled accent dot; not yet = hollow grey. Distinguished by BOTH + // glyph and colour so it survives a greyscale screenshot as well as a + // colour-blind reader. + markOn: { fontSize: 11, width: 11, flex: '0 0 auto', color: '#89b4fa' }, + markOff: { fontSize: 11, width: 11, flex: '0 0 auto', color: '#585b70' }, + label: { flex: '1 1 auto', textAlign: 'left' }, + shape: { + flex: '0 0 auto', fontSize: 11, color: '#6c7086', + fontVariantNumeric: 'tabular-nums', paddingRight: 4, + }, + detail: { + flex: '0 0 auto', fontSize: 11, color: '#9399b2', minWidth: 58, + textAlign: 'right', fontVariantNumeric: 'tabular-nums', + }, + chevron: { flex: '0 0 auto', color: '#6c7086', fontSize: 14, lineHeight: 1 }, + // Same surface as the dropdowns — one panel language across the menu. + card: { + position: 'fixed', zIndex: 9400, pointerEvents: 'none', + background: '#181825', border: '1px solid #313244', borderRadius: 8, + padding: '10px 12px', boxShadow: '0 12px 32px rgba(0,0,0,0.55)', + color: '#cdd6f4', fontSize: 12, + }, + cardTitle: { fontSize: 13, fontWeight: 600, color: '#cdd6f4' }, + cardMeta: { + // NOT uppercased: this line carries units, and text-transform turns + // "200 kV" into "200 KV". + marginTop: 2, fontSize: 10.5, fontWeight: 600, letterSpacing: 0.4, + color: '#89b4fa', + }, + cardFacts: { + marginTop: 8, display: 'flex', flexDirection: 'column', gap: 2, + paddingTop: 8, borderTop: '1px solid #313244', + }, + cardFactRow: { display: 'flex', justifyContent: 'space-between', gap: 12 }, + cardFactKey: { color: '#6c7086', fontSize: 11 }, + cardFactVal: { + color: '#bac2de', fontSize: 11, fontVariantNumeric: 'tabular-nums', + textAlign: 'right', overflowWrap: 'anywhere', + }, + cardBody: { + marginTop: 8, paddingTop: 8, borderTop: '1px solid #313244', + color: '#a6adc8', fontSize: 11.5, lineHeight: 1.45, + // Long descriptions are useful but must not become a wall of text. + display: '-webkit-box', WebkitLineClamp: 6, WebkitBoxOrient: 'vertical', + overflow: 'hidden', + }, + cardState: { marginTop: 8, fontSize: 11, fontWeight: 500 }, sep: { height: 1, background: '#313244', margin: '5px 4px' }, header: { padding: '4px 10px 2px', fontSize: 10.5, fontWeight: 700, diff --git a/electron/src/renderer/src/electron.d.ts b/electron/src/renderer/src/electron.d.ts index 28723ed..784cf46 100644 --- a/electron/src/renderer/src/electron.d.ts +++ b/electron/src/renderer/src/electron.d.ts @@ -31,6 +31,8 @@ declare global { figureEvent: (figId: string, eventJson: string) => void resizeFigure: (figId: string, width: number, height: number) => void openExternal: (url: string) => void + /** Reveal a local directory in the OS file manager. */ + openPath: (path: string) => void getUpdateInfo: () => Promise<{ channel: 'stable' | 'beta' supported: boolean diff --git a/electron/src/renderer/src/kernel/SpyDEContext.tsx b/electron/src/renderer/src/kernel/SpyDEContext.tsx index 9d2d202..c78061b 100644 --- a/electron/src/renderer/src/kernel/SpyDEContext.tsx +++ b/electron/src/renderer/src/kernel/SpyDEContext.tsx @@ -1080,6 +1080,13 @@ export function SpyDEProvider({ children }: { children: React.ReactNode }) { dispatch({ type: 'NAV_SHAPE_PROMPT', prompt: msg }) break + // Examples → Show Example Data Directory. The BACKEND owns where that + // is (em-database, overridable via EM_DATABASE_DATA_DIR) but only the + // main process can open a folder, so it round-trips through here. + case 'open_path': + if (msg.path) window.electron.openPath(String(msg.path)) + break + case 'loading': dispatch({ type: 'LOADING', busy: Boolean(msg.busy), text: String(msg.text ?? '') }) break @@ -1266,6 +1273,10 @@ export function SpyDEProvider({ children }: { children: React.ReactNode }) { // best-match readout under the crosshair. case 'ebsd_dictionary_ready': case 'ebsd_match': + // The Examples menu's contents (em-database): techniques, sizes, + // shapes and which datasets are already downloaded. Consumed by + // MenuBar, which asks for it every time the menu opens. + case 'example_catalogue': // Fit wizard (spyde/actions/fit_action.py) — `fit_catalogue` is the // component picker's shapes, sent once on open; `fit_state` is the // whole model after every edit. Consumed by FitWizard. diff --git a/electron/src/renderer/src/kernel/protocol.ts b/electron/src/renderer/src/kernel/protocol.ts index 513f45c..3cc5031 100644 --- a/electron/src/renderer/src/kernel/protocol.ts +++ b/electron/src/renderer/src/kernel/protocol.ts @@ -209,6 +209,13 @@ export interface NavShapePromptMessage extends MsgBase, NavShapePrompt { type: 'nav_shape_prompt' } +/** Reveal a local directory in the OS file manager (Examples → Show Example + * Data Directory). The backend owns the location, the main process opens it. */ +export interface OpenPathMessage extends MsgBase { + type: 'open_path' + path: string +} + export interface LoadingMessage extends MsgBase { type: 'loading' busy?: unknown @@ -782,6 +789,8 @@ export interface WizardEventMessage extends MsgBase { // EBSD Indexing: the dictionary-ready ack and the live best-match readout. | 'ebsd_dictionary_ready' | 'ebsd_match' + // The Examples menu's contents, from em-database. + | 'example_catalogue' // Fit wizard: the component palette's sampled shapes (once, on open) and // the whole model after every edit. See spyde/actions/fit_action.py. | 'fit_catalogue' @@ -858,6 +867,7 @@ export type PlotAppMessage = | SubItemMessage | HistogramMessage | NavShapePromptMessage + | OpenPathMessage | LoadingMessage | SignalTypeInfoMessage | SelectorInfoMessage diff --git a/electron/tests/examples_menu.spec.ts b/electron/tests/examples_menu.spec.ts new file mode 100644 index 0000000..65eaf33 --- /dev/null +++ b/electron/tests/examples_menu.spec.ts @@ -0,0 +1,139 @@ +/** + * examples_menu.spec.ts — the Examples menu, built from em-database. + * + * The menu is the whole feature here, so this drives it the way a user does: + * open Examples, walk into a technique submenu, and check each dataset row + * carries its size, its shape where known, and a marker saying whether it is + * already on disk. Also pins the Dummy Data submenu and the data-directory + * entry. + * + * Deliberately does NOT click a real dataset — that would download gigabytes. + */ +import { test, expect } from '@playwright/test' +const { launchApp } = require('./_harness.cjs') + +const SHOTS = 'examples_shots' +let ctx: Awaited> + +test.beforeAll(async () => { + ctx = await launchApp({ env: { SPYDE_LOG_LEVEL: 'INFO' } }) +}) + +test.afterAll(async () => { + ctx?.assertNoJsErrors() + await ctx?.app?.close() +}) + +test.setTimeout(120_000) + +async function openExamples() { + const { page } = ctx + // The menu button TOGGLES, so a menu left open by the previous test would be + // closed by this click. Escape first, then open. + await page.keyboard.press('Escape') + await expect(page.getByTestId('menu-examples-items')).toBeHidden() + await page.getByTestId('menu-examples').click() + await expect(page.getByTestId('menu-examples-items')).toBeVisible() +} + +test('Examples groups the datasets into technique submenus', async () => { + const { page } = ctx + await openExamples() + + // The backend catalogue arrives async; the technique rows appear with it. + await expect(page.getByTestId('examples-tech-4d-stem')) + .toBeVisible({ timeout: 30_000 }) + for (const tech of ['4d-stem', 'eels', 'ebsd']) { + await expect(page.getByTestId(`examples-tech-${tech}`), + `no ${tech} submenu`).toBeVisible() + } + await page.screenshot({ path: `${SHOTS}/01-examples-techniques.png` }) + ctx.assertNoJsErrors() +}) + +test('each dataset shows its size, shape and download state', async () => { + const { page } = ctx + await openExamples() + await page.getByTestId('examples-tech-4d-stem').hover() + + const items = page.getByTestId('examples-tech-4d-stem-items') + await expect(items).toBeVisible({ timeout: 15_000 }) + + // SPEDAg is the scan the 4D-STEM work is benchmarked on; it must be listed + // with its size, and its shape once it has been downloaded and measured. + const sped = page.getByTestId('example-SPEDAg') + await expect(sped).toBeVisible() + await expect(sped).toContainText(/\d+(\.\d+)?\s*[kMG]B/) + + // Every row carries exactly one state marker. + const rows = await items.getByRole('button').all() + expect(rows.length).toBeGreaterThan(3) + let marked = 0 + for (const row of rows) { + const text = (await row.textContent()) ?? '' + if (text.includes('●') || text.includes('○')) marked++ + } + expect(marked, 'dataset rows carry no downloaded/not-downloaded marker') + .toBe(rows.length) + + await page.screenshot({ path: `${SHOTS}/02-4dstem-datasets.png` }) + ctx.assertNoJsErrors() +}) + +test('hovering a dataset shows a themed info card', async () => { + const { page } = ctx + await openExamples() + await page.getByTestId('examples-tech-4d-stem').hover() + await expect(page.getByTestId('examples-tech-4d-stem-items')).toBeVisible() + + await page.getByTestId('example-SPEDAg').hover() + const card = page.getByTestId('menu-hover-card') + await expect(card).toBeVisible({ timeout: 10_000 }) + await expect(card).toContainText('SPEDAg') + await expect(card).toContainText('4D-STEM') + await expect(card).toContainText(/Size/) + await expect(card).toContainText(/On disk|Not downloaded/) + await page.screenshot({ path: `${SHOTS}/04-hover-card.png` }) + + // It is OUR panel, not the OS bubble — so the row must not also carry a + // native title attribute racing it. + await expect(page.getByTestId('example-SPEDAg')).not.toHaveAttribute('title', /./) + + // An undownloaded set reads differently. + await page.getByTestId('example-FeAlStripes').hover() + await expect(card).toContainText('Not downloaded') + await page.screenshot({ path: `${SHOTS}/05-hover-card-undownloaded.png` }) + ctx.assertNoJsErrors() +}) + +test('Dummy Data is its own submenu and still loads', async () => { + const { page } = ctx + await openExamples() + await page.getByTestId('examples-dummy-data').hover() + const items = page.getByTestId('examples-dummy-data-items') + await expect(items).toBeVisible({ timeout: 15_000 }) + await expect(page.getByTestId('tutorial-navigation')).toBeVisible() + await page.screenshot({ path: `${SHOTS}/03-dummy-data.png` }) + + // It is instant + no-download, so actually clicking one is fair game. + await page.getByTestId('tutorial-navigation').click() + await expect(page.getByTestId('subwindow').first()) + .toBeVisible({ timeout: 60_000 }) + ctx.assertNoJsErrors() +}) + +test('Show Example Data Directory reports a real path', async () => { + const { page, backend } = ctx + await openExamples() + await expect(page.getByTestId('examples-show-dir')).toBeVisible() + await page.getByTestId('examples-show-dir').click() + // `open_path` is consumed by the renderer and never echoed back on the + // PLOTAPP channel, so waiting on a message would wait forever — the backend + // logs the reveal for exactly this reason. + await backend.waitForLog('revealing example data directory', 20_000) + const line = backend.logBuffer.find( + (l: string) => l.includes('revealing example data directory')) + expect(line, 'the backend never reported the directory').toBeTruthy() + expect(line).toContain('em_database') + ctx.assertNoJsErrors() +}) diff --git a/pyproject.toml b/pyproject.toml index ff0e705..d9b36e9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,6 +33,19 @@ dependencies = [ # [hdf5,image] extras cover .hspy / .tif reading (hyperspy pins the same). "rosettasciio[hdf5,image]>=0.14.0", "pyxem", + # The Examples menu's datasets. One curated, checksummed, cited collection + # (grouped by technique, with real acquisition metadata) instead of each + # analysis library's own fixtures — adding a dataset is a YAML file + # upstream rather than a loader here. See spyde/backend/example_catalogue. + # + # MARKER, not a plain dependency: em-database 0.1.0 declares + # requires-python >=3.12 while SpyDE supports >=3.10, so depending on it + # unconditionally makes the whole project unresolvable on 3.10/3.11 — the + # backend then fails to start at all, taking the app with it. On those + # versions the Examples menu falls back to Dummy Data only + # (example_catalogue.available() is False). Drop the marker once upstream + # relaxes its floor. + "em-database>=0.1.0; python_version >= '3.12'", "imageio-ffmpeg", "imagecodecs", "psygnal>=0.10.0", diff --git a/spyde/backend/_session_actions.py b/spyde/backend/_session_actions.py index 5264605..4f2bcba 100644 --- a/spyde/backend/_session_actions.py +++ b/spyde/backend/_session_actions.py @@ -223,6 +223,10 @@ def dispatch_action(self, msg: dict) -> None: self._set_signal_type(plot, payload.get("signal_type", "")) elif action == "load_example": self.load_example_data(payload["name"]) + elif action == "example_catalogue": + self.emit_example_catalogue(warm=bool(payload.get("warm", True))) + elif action == "show_example_dir": + self.show_example_dir() elif action == "set_active": wid = payload.get("window_id", window_id) if wid is not None: diff --git a/spyde/backend/_session_files.py b/spyde/backend/_session_files.py index 1c6617e..57e817e 100644 --- a/spyde/backend/_session_files.py +++ b/spyde/backend/_session_files.py @@ -21,7 +21,7 @@ from hyperspy.signal import BaseSignal from spyde.backend import ipc -from spyde.backend.ipc import emit_status, emit_error +from spyde.backend.ipc import emit, emit_status, emit_error log = logging.getLogger(__name__) @@ -617,9 +617,57 @@ def _carry_axes_from_reference(new: BaseSignal, ref: BaseSignal) -> None: except Exception: pass - def load_example_data(self, name: str) -> None: - import pyxem.data as _pxd + def emit_example_catalogue(self, warm: bool = False) -> None: + """Send the Examples menu its contents. + + Cheap by construction (no file is opened), so the renderer can ask for + it every time the menu opens and always show current download state. + With *warm*, any downloaded dataset whose shape we have not read yet is + read on a worker and the catalogue re-sent — that pass opens files, so + it must never be on the menu's own path. + """ + from spyde.backend import example_catalogue as catalogue + + emit({"type": "example_catalogue", **catalogue.catalogue()}) + if not warm: + return + + def _warm() -> None: + try: + if catalogue.warm_shapes(): + self._dispatch_to_main( + lambda: emit({"type": "example_catalogue", + **catalogue.catalogue()})) + except Exception as e: + log.debug("warming example shapes failed: %s", e) + + threading.Thread(target=_warm, daemon=True, + name="example-shapes").start() + + def show_example_dir(self) -> None: + """Reveal the example-data directory in the OS file manager. + The backend knows where it is (em-database owns the location and it is + overridable via ``EM_DATABASE_DATA_DIR``); only the main process can + open a folder, so it goes out as a message rather than being opened + here. The directory is created first — revealing a path that does not + exist yet does nothing on every platform, which reads as a dead menu + item before the first download. + """ + from spyde.backend import example_catalogue as catalogue + + path = catalogue.data_dir() + try: + os.makedirs(path, exist_ok=True) + except OSError as e: + log.debug("creating the example data dir failed: %s", e) + # Logged, not just emitted: `open_path` is consumed by the renderer and + # never echoed back on the PLOTAPP channel, so this line is the only + # thing an e2e (or a bug report) can observe. + log.info("[examples] revealing example data directory: %s", path) + emit({"type": "open_path", "path": path}) + + def load_example_data(self, name: str) -> None: emit_status(f"Loading example: {name}…") threading.Thread( target=self._load_example_thread, @@ -629,56 +677,62 @@ def load_example_data(self, name: str) -> None: ).start() def _load_example_thread(self, name: str) -> None: + """Download (if needed) and open one **em-database** dataset, lazily. + + em-database is the single source of example data (see + ``example_catalogue``): it downloads through pooch with a checksum, so + this only has to supply the progress/cancel monitor and then open the + file it hands back. + """ + from spyde.backend import example_catalogue as catalogue from spyde.backend.example_download import ( - DownloadCancelled, patched_example_downloader, + DownloadCancelled, example_progress, ) try: # Wait for the Dask cluster before registering the signal — _add_signal # builds the navigator compute, which needs the client. A load fired # during startup queues here instead of racing ahead with a None client. self._await_dask() - import pyxem.data as _pxd - loader = getattr(_pxd, name, None) - if loader is None: + ds = catalogue.resolve(name) + if ds is None: emit_error(f"Unknown example dataset: {name}") return - # Progress + cancel for the (possible) pooch download: the renderer - # shows a toast with a bar and a Cancel button (download_progress / - # download_done messages; download_cancel action). A cache hit never - # downloads → no toast. + # Progress + cancel for the (possible) download: the renderer shows + # a toast with a bar and a Cancel button (download_progress / + # download_done messages; download_cancel action). Already on disk + # → pooch never streams → no toast. try: - with patched_example_downloader(f"example:{name}", name): - sig = self._load_example_lazy(loader) + with example_progress(f"example:{name}", name) as monitor: + path = ds.download(progressbar=monitor) except DownloadCancelled: emit_status(f"Download cancelled: {name}") return + + sig = self._load_example_lazy(str(path)) + # Now that it is open, remember its shape so the Examples menu can + # show it without ever reopening the file. + catalogue.record_shape(str(path), sig) _apply_example_calibration(sig, name) self._maybe_set_insitu_signal_type(sig) - self._add_signal(sig, source_path=None) + self._add_signal(sig, source_path=str(path)) + self.emit_example_catalogue() # the entry is downloaded now except Exception as e: import traceback traceback.print_exc() emit_error(f"Failed to load example {name}: {e}") - def _load_example_lazy(self, loader) -> BaseSignal: - """Load a pyxem example as a LAZY (Dask-backed) signal. + def _load_example_lazy(self, path: str) -> BaseSignal: + """Open a downloaded example as a LAZY (Dask-backed) signal. - pyxem example loaders forward ``**kwargs`` to ``hs.load``, so passing - ``lazy=True`` reads the already-downloaded file straight off disk as a - dask array — no eager 668 MB materialise, no zspy re-save. Falls back to - an in-place ``as_lazy()`` wrap only if a loader doesn't honour ``lazy``. + Lazy so a multi-GB store is never materialised just to display it (the + navigator reads it a chunk at a time). Falls back to kikuchipy for the + EBSD files hyperspy's reader sniffing refuses to disambiguate — the + same fallback the catalogue uses to read their shape. """ - for kwargs in ({"allow_download": True, "lazy": True}, {"lazy": True}): - try: - sig = loader(**kwargs) - except TypeError: - continue - return sig if getattr(sig, "_lazy", False) else sig.as_lazy() - # Loader doesn't accept those kwargs at all → eager once, wrap lazy. - try: - sig = loader(allow_download=True) - except TypeError: - sig = loader() + from spyde.backend.example_catalogue import _open_lazily + sig = _open_lazily(path) + if isinstance(sig, (list, tuple)): + sig = sig[0] return sig if getattr(sig, "_lazy", False) else sig.as_lazy() def _to_lazy(self, sig: BaseSignal, name: str) -> BaseSignal: diff --git a/spyde/backend/example_catalogue.py b/spyde/backend/example_catalogue.py new file mode 100644 index 0000000..b270e48 --- /dev/null +++ b/spyde/backend/example_catalogue.py @@ -0,0 +1,369 @@ +""" +example_catalogue.py — the Examples menu's catalogue, from **em-database**. + +SpyDE's example datasets come from the `em-database +`_ package rather than from each +analysis library's own hand-rolled fixtures. One curated, checksummed, cited +collection beats a handful of loaders scattered across pyxem/exspy/kikuchipy: +the datasets carry real acquisition metadata, they are versioned independently +of SpyDE, and adding one is a YAML file upstream rather than a code change +here. + +What this module turns that into is the shape the menu needs: + +* **grouped by technique** — ``metadata["technique"]`` is already there + (4D-STEM, EELS, EBSD, Cryo-EM, In-situ TEM, STEM), so the menu gets one + submenu per technique for free; +* **sized** — ``data_size`` is a ready-made human string ("1.4 GB"); +* **shaped** — see :func:`_shape_of`, the one thing em-database does not + (yet) carry for every dataset; +* **marked downloaded or not** — ``DownloadableDataset.filepath()`` returns + the local path or ``None``, which is exactly that question. + +Nothing here downloads anything. Building the catalogue must stay cheap enough +to run every time the menu opens. +""" +from __future__ import annotations + +import inspect +import json +import logging +import os + +log = logging.getLogger(__name__) + +#: Techniques in the order the menu should show them — the modalities SpyDE is +#: built around first, then the rest alphabetically. Anything not listed still +#: appears, after these. +TECHNIQUE_ORDER = ("4D-STEM", "EELS", "EDS", "EBSD", "STEM", "In-situ TEM", + "Cryo-EM") + +#: Shapes read out of downloaded files, cached beside the data so the menu +#: doesn't reopen every file each time it opens. +_SHAPE_CACHE_NAME = ".spyde_shapes.json" + +_shape_cache: dict | None = None + + +def available() -> bool: + """Is em-database importable? The menu degrades to Dummy Data if not.""" + try: + import em_database # noqa: F401 + return True + except Exception as e: + log.debug("em-database unavailable: %s", e) + return False + + +def data_dir() -> str: + """Where em-database keeps downloaded datasets (``EM_DATABASE_DATA_DIR`` + or ``~/em_database``). Shown by the Examples menu's "Show Example Data + Directory".""" + try: + import em_database + return str(em_database.get_data_dir()) + except Exception: + return os.path.join(os.path.expanduser("~"), "em_database") + + +def datasets() -> list[tuple[str, object]]: + """``(key, dataset)`` for every dataset em-database exposes. + + The key is the class name — stable, and what ``load_example {name}`` + carries. Filtered by ``issubclass`` rather than by name so the base class + and incidental imports (``Path``) in the module namespace stay out. + """ + try: + import em_database + import em_database.data as data + except Exception as e: + log.debug("listing em-database datasets failed: %s", e) + return [] + + base = em_database.DownloadableDataset + out = [] + for name in dir(data): + if name.startswith("_"): + continue + obj = getattr(data, name, None) + if not inspect.isclass(obj) or obj is base or not issubclass(obj, base): + continue + try: + out.append((name, obj())) + except Exception as e: + log.debug("dataset %s could not be instantiated: %s", name, e) + return sorted(out, key=lambda kv: kv[0].lower()) + + +def resolve(key: str): + """The dataset object for a catalogue key, or None.""" + try: + import em_database.data as data + obj = getattr(data, str(key), None) + return obj() if inspect.isclass(obj) else None + except Exception as e: + log.debug("resolving example %r failed: %s", key, e) + return None + + +# ───────────────────────────────────────────────────────────────────────────── +# Shape +# ───────────────────────────────────────────────────────────────────────────── + +def _cache_path() -> str: + return os.path.join(data_dir(), _SHAPE_CACHE_NAME) + + +def _load_shape_cache() -> dict: + global _shape_cache + if _shape_cache is None: + try: + with open(_cache_path(), encoding="utf-8") as fh: + _shape_cache = json.load(fh) + except Exception: + _shape_cache = {} + return _shape_cache + + +def _save_shape_cache() -> None: + try: + os.makedirs(data_dir(), exist_ok=True) + with open(_cache_path(), "w", encoding="utf-8") as fh: + json.dump(_shape_cache or {}, fh) + except Exception as e: + log.debug("writing the shape cache failed: %s", e) + + +def _declared_shape(ds) -> str | None: + """The shape em-database declares, if it does. + + Upstream (``cssfrancis/em_data``) describes each dataset with a YAML file; + a ``shape`` there is authoritative and means the menu can show the shape + BEFORE anything is downloaded. Not every dataset carries one yet, hence + :func:`_read_shape`. + """ + for source in (getattr(ds, "metadata", None) or {}, ds): + for attr in ("shape", "data_shape"): + val = (source.get(attr) if isinstance(source, dict) + else getattr(source, attr, None)) + if val: + return _format_shape(val) + return None + + +def _format_shape(val) -> str: + """Normalise whatever shape upstream gives into ``nav | sig`` display form.""" + if isinstance(val, str): + return val.strip() + try: + if (isinstance(val, dict)): + nav, sig = val.get("navigation"), val.get("signal") + if nav and sig: + return f"{_tuple_str(nav)} | {_tuple_str(sig)}" + return _tuple_str(nav or sig) + return _tuple_str(val) + except Exception: + return str(val) + + +def _tuple_str(seq) -> str: + return "×".join(str(int(v)) for v in seq) + + +def _cache_key(path: str) -> str | None: + """Size+mtime keyed, so a re-download invalidates the cached shape.""" + try: + stat = os.stat(path) + except OSError: + return None + return f"{os.path.basename(path)}:{stat.st_size}:{int(stat.st_mtime)}" + + +def shape_of_signal(sig) -> str | None: + """``nav | sig`` for a loaded signal.""" + try: + am = sig.axes_manager + nav = tuple(int(a.size) for a in am.navigation_axes) + sgl = tuple(int(a.size) for a in am.signal_axes) + return (f"{_tuple_str(nav)} | {_tuple_str(sgl)}" if nav + else _tuple_str(sgl)) + except Exception as e: + log.debug("shape of signal failed: %s", e) + return None + + +def record_shape(path: str | None, sig) -> None: + """Remember a dataset's shape from a signal we already have in hand. + + Called when an example is loaded — free, exact, and it means the menu shows + the shape from then on without ever opening the file just to look. + """ + if not path: + return + key = _cache_key(path) + shape = shape_of_signal(sig) + if not key or not shape: + return + cache = _load_shape_cache() + if cache.get(key) != shape: + cache[key] = shape + _save_shape_cache() + + +def _open_lazily(path: str): + """Lazily open a downloaded example, metadata only. + + hyperspy's reader sniffing is ambiguous for some EBSD ``.h5`` files (it + offers Arina/Delmic/Tofwerk/USID and refuses to choose), and those are + exactly the files whose reader lives in kikuchipy. So: hyperspy first, + kikuchipy as the fallback — which also keeps kikuchipy optional, since a + dataset it cannot open simply has no shape shown. + """ + import hyperspy.api as hs + try: + return hs.load(path, lazy=True) + except Exception as e: + log.debug("hyperspy could not open %s (%s); trying kikuchipy", path, e) + import kikuchipy as kp + return kp.load(path, lazy=True) + + +def read_shape(path: str) -> str | None: + """Read a DOWNLOADED file's shape, without reading its data. + + A lazy load touches the store's metadata only, so this is cheap even for + the multi-GB sets — but it still opens a file, which is why the catalogue + never calls it inline (see :func:`warm_shapes`). + """ + key = _cache_key(path) + if key is None: + return None + cache = _load_shape_cache() + if key in cache: + return cache[key] or None + + shape = None + try: + sig = _open_lazily(path) + if isinstance(sig, (list, tuple)): + sig = sig[0] + shape = shape_of_signal(sig) + try: + sig.data = None # drop the dask graph promptly + except Exception: + pass + except Exception as e: + log.debug("reading the shape of %s failed: %s", path, e) + + cache[key] = shape or "" + _save_shape_cache() + return shape + + +def _cached_shape(path: str | None) -> str | None: + """The shape we already know for a downloaded file — never opens it.""" + if not path: + return None + key = _cache_key(path) + return (_load_shape_cache().get(key) or None) if key else None + + +def _shape_of(ds, path: str | None) -> str | None: + """Declared shape if em-database has one, else one we have already read. + + Deliberately does NOT open files: the catalogue is rebuilt every time the + Examples menu opens, and reading five multi-GB stores inline made that a + 7-second click. :func:`warm_shapes` fills the gap off the menu's path. + """ + return _declared_shape(ds) or _cached_shape(path) + + +def warm_shapes() -> bool: + """Read the shape of every downloaded dataset we don't know yet. + + Slow the first time (it opens each file), so run it on a worker and + re-emit the catalogue if it returns True. Cheap and a no-op after that. + """ + changed = False + for _key, ds in datasets(): + if _declared_shape(ds): + continue + try: + path = ds.filepath() + except Exception: + path = None + if path and _cached_shape(path) is None: + if read_shape(path): + changed = True + return changed + + +# ───────────────────────────────────────────────────────────────────────────── +# The catalogue +# ───────────────────────────────────────────────────────────────────────────── + +def _label(key: str) -> str: + """``LayeredCuNb4DSTEM`` -> ``Layered Cu Nb 4D STEM`` is worse than the + class name, so the class name IS the label. Kept as a hook for an upstream + ``title`` field.""" + return key + + +def _technique(ds) -> str: + md = getattr(ds, "metadata", None) or {} + return str(md.get("technique") or "Other").strip() or "Other" + + +def entry(key: str, ds) -> dict: + """One catalogue row — everything the menu draws for a dataset.""" + try: + path = ds.filepath() + except Exception as e: + log.debug("filepath() for %s failed: %s", key, e) + path = None + md = getattr(ds, "metadata", None) or {} + return { + "key": key, + "label": _label(key), + "technique": _technique(ds), + "size": str(getattr(ds, "data_size", "") or ""), + "shape": _shape_of(ds, path), + "downloaded": bool(path), + "path": path or None, + "file": str(getattr(ds, "file", "") or ""), + "description": str(getattr(ds, "description", "") or ""), + "tags": list(md.get("tags") or []), + "microscope": str(md.get("microscope") or ""), + "voltage": str(md.get("voltage") or ""), + "license": str(getattr(ds, "license", "") or ""), + "source": str(getattr(ds, "source", "") or ""), + } + + +def catalogue() -> dict: + """The whole Examples menu, grouped by technique. + + ``{"available": bool, "data_dir": str, "groups": [{"technique", "items"}]}`` + — one group per technique in :data:`TECHNIQUE_ORDER`, then any others + alphabetically. + """ + items = [entry(key, ds) for key, ds in datasets()] + by_tech: dict[str, list[dict]] = {} + for it in items: + by_tech.setdefault(it["technique"], []).append(it) + + def order(tech: str): + try: + return (0, TECHNIQUE_ORDER.index(tech)) + except ValueError: + return (1, tech.lower()) + + groups = [{"technique": t, "items": by_tech[t]} + for t in sorted(by_tech, key=order)] + return { + "available": available(), + "data_dir": data_dir(), + "groups": groups, + "n_downloaded": sum(1 for it in items if it["downloaded"]), + "n_total": len(items), + } diff --git a/spyde/backend/example_download.py b/spyde/backend/example_download.py index cd77cc6..719aefc 100644 --- a/spyde/backend/example_download.py +++ b/spyde/backend/example_download.py @@ -1,10 +1,8 @@ """example_download.py — IPC progress + cancel for the Examples-menu downloads. -pyxem's example loaders fetch their data through pooch -(``Dataset.fetch_file_path`` → ``pooch.HTTPDownloader(progressbar=…)`` → -``kipper.fetch``). pooch accepts an arbitrary tqdm-like object as the -``progressbar`` — that one hook gives us both a byte-level progress stream and -a cancellation point, without touching how any individual loader works: +em-database fetches through pooch, and pooch accepts an arbitrary tqdm-like +object as its ``progressbar`` — that one hook gives us both a byte-level +progress stream and a cancellation point: - :class:`_IpcProgress` implements pooch's progress protocol (assignable ``total``, ``update(n)``, ``reset()``, ``close()``) and emits throttled @@ -13,10 +11,7 @@ file and deletes it on error, so no partial file is left in the cache; the exception is not a ``ValueError``/requests error, so pooch's retry loop does NOT re-attempt a cancelled download). -- :func:`patched_example_downloader` scopes the hook: for the duration of one - example load it swaps ``pyxem.data._data``'s view of the ``pooch`` module for - a proxy whose ``HTTPDownloader`` pins ``progressbar=`` to our monitor. On - exit it restores the real module and emits ``download_done``. +- :func:`example_progress` scopes one download and emits ``download_done``. - ``download_cancel`` (a staged action, wired in ``actions/registry.py``) sets the cancel flag for a token — the renderer's toast Cancel button sends it. @@ -123,39 +118,29 @@ def _emit(self, force: bool = False) -> None: "total": int(self._total)}) -class _PoochProxy: - """A stand-in for the ``pooch`` module inside ``pyxem.data._data`` whose - ``HTTPDownloader`` pins ``progressbar=`` to our monitor. Everything else - delegates to the real module.""" - - def __init__(self, real, monitor: _IpcProgress): - self._real = real - self._monitor = monitor - - def __getattr__(self, name): - return getattr(self._real, name) - - def HTTPDownloader(self, *args, **kwargs): # noqa: N802 (pooch API name) - kwargs["progressbar"] = self._monitor - return self._real.HTTPDownloader(*args, **kwargs) - - @contextmanager -def patched_example_downloader(token: str, label: str): - """Scope the progress/cancel hook around one example load. - - Progress messages only flow when pooch actually downloads (a cache hit - never constructs a progress bar), so the toast simply never appears for - already-cached data. ``download_done`` is emitted on exit iff a download - started.""" - import pyxem.data._data as _pxdata - +def example_progress(token: str, label: str): + """Scope the progress/cancel monitor around one example download. + + Yields the monitor to hand straight to + ``em_database.DownloadableDataset.download(progressbar=…)``. That method + forwards whatever it is given to ``pooch.HTTPDownloader``, and pooch's + ``progressbar`` contract is "any tqdm-like object" — so :class:`_IpcProgress` + plugs in with no interception at all. + + (This used to swap out ``pyxem.data._data``'s view of the ``pooch`` module + to reach the same hook, because pyxem's loaders construct the downloader + themselves and take no progress argument. em-database exposing the + parameter is what let that monkey-patch go.) + + Progress messages only flow when pooch actually streams — a file already on + disk never constructs a progress bar — so the toast simply never appears + for cached data, and ``download_done`` is emitted iff a download started. + """ cancel = threading.Event() with _LOCK: _CANCELS[token] = cancel monitor = _IpcProgress(token, label, cancel) - real_pooch = _pxdata.pooch - _pxdata.pooch = _PoochProxy(real_pooch, monitor) ok = False cancelled = False try: @@ -165,7 +150,6 @@ def patched_example_downloader(token: str, label: str): cancelled = True raise finally: - _pxdata.pooch = real_pooch with _LOCK: _CANCELS.pop(token, None) if monitor.started: diff --git a/spyde/tests/migrated/test_example_catalogue.py b/spyde/tests/migrated/test_example_catalogue.py new file mode 100644 index 0000000..9c8d214 --- /dev/null +++ b/spyde/tests/migrated/test_example_catalogue.py @@ -0,0 +1,147 @@ +"""The Examples menu's catalogue (``spyde/backend/example_catalogue.py``). + +The catalogue is rebuilt every time the menu opens, so the property that +matters most is that building it is CHEAP — it must not open a data file to +answer "what is in here?". The rest pins the three things the menu draws that +were not simply handed over by em-database: the technique grouping order, the +downloaded marker, and where a shape comes from. +""" +from __future__ import annotations + +import os + +import pytest + +from spyde.backend import example_catalogue as ec + +pytestmark = pytest.mark.skipif(not ec.available(), + reason="em-database is not installed " + "(it requires Python >= 3.12)") + + +class TestCatalogue: + def test_lists_datasets_grouped_by_technique(self): + cat = ec.catalogue() + assert cat["available"] + assert cat["n_total"] > 5, "em-database returned almost nothing" + techs = [g["technique"] for g in cat["groups"]] + assert "4D-STEM" in techs + assert len(techs) == len(set(techs)), "a technique was split in two" + assert sum(len(g["items"]) for g in cat["groups"]) == cat["n_total"] + + def test_the_modalities_spyde_is_built_around_come_first(self): + """4D-STEM before Cryo-EM, whatever order em-database enumerates in — + the menu should open on the things this app is for.""" + techs = [g["technique"] for g in ec.catalogue()["groups"]] + known = [t for t in techs if t in ec.TECHNIQUE_ORDER] + assert known == sorted(known, key=ec.TECHNIQUE_ORDER.index) + + def test_every_entry_has_what_the_menu_draws(self): + for group in ec.catalogue()["groups"]: + for item in group["items"]: + assert item["key"] and item["label"] + assert item["technique"] == group["technique"] + assert isinstance(item["downloaded"], bool) + assert item["size"], f"{item['key']} has no size to show" + + def test_the_base_class_is_not_a_dataset(self): + """``em_database.data`` also holds the DownloadableDataset base and an + incidental ``Path`` import; neither is an example.""" + keys = {k for k, _ in ec.datasets()} + assert "DownloadableDataset" not in keys + assert "Path" not in keys + + def test_resolve_round_trips_a_key(self): + key = ec.catalogue()["groups"][0]["items"][0]["key"] + ds = ec.resolve(key) + assert ds is not None and getattr(ds, "file", None) + assert ec.resolve("NoSuchDataset") is None + + def test_downloaded_reflects_the_file_on_disk(self): + """``filepath()`` returns the path or None — that IS the question the + marker asks, so the two must not disagree.""" + for group in ec.catalogue()["groups"]: + for item in group["items"]: + assert item["downloaded"] == bool(item["path"]) + if item["downloaded"]: + assert os.path.exists(item["path"]) + + +class TestBuildingItIsCheap: + def test_the_catalogue_never_opens_a_data_file(self, monkeypatch): + """The menu asks for this on every open. Reading five multi-GB stores + inline made that a 7-second click, so the catalogue uses only declared + and already-cached shapes — warm_shapes() does the reading, off the + menu's path. + """ + opened = [] + monkeypatch.setattr(ec, "_open_lazily", + lambda path: opened.append(path) or (_ for _ in ()).throw( + AssertionError(f"catalogue opened {path}"))) + cat = ec.catalogue() + assert cat["n_total"] > 0 + assert not opened + + def test_data_dir_is_reported_even_when_nothing_is_downloaded(self): + path = ec.data_dir() + assert path and os.path.isabs(path) + + +class TestShape: + def test_declared_shape_wins_over_reading_the_file(self, monkeypatch): + """When em-database's YAML carries a shape, it is authoritative — the + menu can then show a shape BEFORE anything is downloaded, and no file + is opened to find one.""" + class _Declared: + metadata = {"technique": "4D-STEM", "shape": {"navigation": (32, 32), + "signal": (256, 256)}} + data_size = "1 GB" + file = "x.zspy" + description = "" + license = "" + source = "" + + def filepath(self): + return None + + monkeypatch.setattr(ec, "read_shape", + lambda p: pytest.fail("read a file despite a declared shape")) + entry = ec.entry("Declared", _Declared()) + assert entry["shape"] == "32×32 | 256×256" + assert entry["downloaded"] is False + + def test_an_undownloaded_dataset_without_a_declared_shape_shows_none(self): + class _Bare: + metadata = {"technique": "EELS"} + data_size = "1 kB" + file = "y.hspy" + description = "" + license = "" + source = "" + + def filepath(self): + return None + + assert ec.entry("Bare", _Bare())["shape"] is None + + def test_record_shape_stamps_what_a_load_already_knows(self, tmp_path, + monkeypatch): + """Loading an example gives us the signal — so the shape is free and + exact, and the menu shows it from then on without reopening anything.""" + import hyperspy.api as hs + import numpy as np + + monkeypatch.setattr(ec, "data_dir", lambda: str(tmp_path)) + monkeypatch.setattr(ec, "_shape_cache", None, raising=False) + target = tmp_path / "example.hspy" + target.write_bytes(b"not really a file") + + sig = hs.signals.Signal2D(np.zeros((4, 5, 6, 7), np.float32)) + ec.record_shape(str(target), sig) + assert ec._cached_shape(str(target)) == "5×4 | 7×6" + + def test_shape_of_signal_is_nav_then_signal(self): + import hyperspy.api as hs + import numpy as np + sig = hs.signals.Signal2D(np.zeros((3, 8, 9), np.float32)) + assert ec.shape_of_signal(sig) == "3 | 9×8" diff --git a/spyde/tests/migrated/test_example_download.py b/spyde/tests/migrated/test_example_download.py index 3448e27..6b0ec60 100644 --- a/spyde/tests/migrated/test_example_download.py +++ b/spyde/tests/migrated/test_example_download.py @@ -1,9 +1,8 @@ """Examples-menu download progress + cancel (spyde/backend/example_download.py). Covers the pooch progress-object protocol (total / update / reset / close), -message throttling, the cancel flag → DownloadCancelled abort, the scoped -pyxem pooch proxy patch, and the download_cancel staged action. No network — -pooch itself is never invoked. +message throttling, the cancel flag → DownloadCancelled abort, and the +download_cancel staged action. No network — pooch itself is never invoked. """ from __future__ import annotations @@ -60,29 +59,31 @@ def test_cancel_raises(self, monkeypatch): class TestPatchedDownloader: - def test_proxy_scoped_and_restored(self, monkeypatch): - """Inside the context pyxem's pooch.HTTPDownloader pins progressbar= to - the monitor; outside, the real module is back.""" - import pyxem.data._data as _pxdata + def test_the_monitor_is_a_pooch_progressbar(self, monkeypatch): + """The monitor is handed straight to + ``em_database…download(progressbar=…)``, which forwards it to + ``pooch.HTTPDownloader``. So it has to BE a pooch progressbar — that is + what replaced the old proxy that swapped out pyxem's view of the pooch + module to reach the same hook.""" + import pooch dl, _ = _capture(monkeypatch) - real = _pxdata.pooch - with dl.patched_example_downloader("example:t", "t") as monitor: - assert _pxdata.pooch is not real - d = _pxdata.pooch.HTTPDownloader(progressbar=True) - assert d.progressbar is monitor - # everything else delegates to the real pooch - assert _pxdata.pooch.os_cache is real.os_cache - assert _pxdata.pooch is real + with dl.example_progress("example:t", "t") as monitor: + downloader = pooch.HTTPDownloader(progressbar=monitor) + assert downloader.progressbar is monitor + # pooch's contract: assignable total, update/reset/close. + monitor.total = 10 + monitor.update(5) + monitor.reset() + monitor.close() def test_done_emitted_only_when_started(self, monkeypatch): - import pyxem.data._data # noqa: F401 — ensure importable dl, msgs = _capture(monkeypatch) # Cache hit: no bytes flowed → no download_done. - with dl.patched_example_downloader("example:hit", "hit"): + with dl.example_progress("example:hit", "hit"): pass assert not [m for m in msgs if m["type"] == "download_done"] # Download happened → ok:true done. - with dl.patched_example_downloader("example:miss", "miss") as mon: + with dl.example_progress("example:miss", "miss") as mon: mon.total = 10 mon.update(10) done = [m for m in msgs if m["type"] == "download_done"] @@ -91,10 +92,9 @@ def test_done_emitted_only_when_started(self, monkeypatch): def test_cancel_flow(self, monkeypatch): """download_cancel (the staged action) flags the token; the next chunk raises, the context re-raises and reports cancelled:true.""" - import pyxem.data._data # noqa: F401 dl, msgs = _capture(monkeypatch) with pytest.raises(dl.DownloadCancelled): - with dl.patched_example_downloader("example:c", "c") as mon: + with dl.example_progress("example:c", "c") as mon: mon.total = 100 mon.update(10) dl.download_cancel(None, None, {"token": "example:c"}) diff --git a/spyde/tests/migrated/test_lazy_loading.py b/spyde/tests/migrated/test_lazy_loading.py index bcee435..be1cf20 100644 --- a/spyde/tests/migrated/test_lazy_loading.py +++ b/spyde/tests/migrated/test_lazy_loading.py @@ -50,29 +50,43 @@ def test_lazy_signal_flows_through_add_signal_without_client(self, window): # Navigator + signal windows were created. assert len(session._plots) >= 2 - def test_load_example_requests_a_lazy_load(self, window): - """The example loader must be asked for a LAZY load (pyxem forwards - lazy=True to hs.load → reads the downloaded file as dask, no eager - materialise, no zspy re-save).""" + def test_load_example_opens_the_file_lazily(self, window, monkeypatch): + """A downloaded example is opened as dask, never materialised — some of + them are gigabytes, and the navigator reads a chunk at a time.""" + import spyde.backend.example_catalogue as ec session = window["window"] - calls = [] + opened = [] - def fake_loader(**kwargs): - calls.append(kwargs) - return _eager_4d().as_lazy() # pretend the loader honoured lazy + def fake_open(path): + opened.append(path) + return _eager_4d().as_lazy() - sig = session._load_example_lazy(fake_loader) + monkeypatch.setattr(ec, "_open_lazily", fake_open) + sig = session._load_example_lazy("/tmp/example.zspy") assert sig._lazy is True - assert calls and calls[0].get("lazy") is True + assert opened == ["/tmp/example.zspy"] - def test_load_example_wraps_an_eager_only_loader(self, window): - """If a loader ignores lazy and returns eager, wrap via as_lazy() — + def test_load_example_wraps_an_eager_reader(self, window, monkeypatch): + """A reader that hands back an eager signal is wrapped via as_lazy() — still Dask-backed, still no disk round-trip.""" + import spyde.backend.example_catalogue as ec session = window["window"] - sig = session._load_example_lazy(lambda **k: _eager_4d()) + monkeypatch.setattr(ec, "_open_lazily", lambda path: _eager_4d()) + sig = session._load_example_lazy("/tmp/example.hspy") assert sig._lazy is True assert isinstance(sig.data, da.Array) + def test_load_example_takes_the_first_of_a_multi_signal_file(self, window, + monkeypatch): + """Some readers return a LIST of signals for one file; loading one must + not hand a list to _add_signal.""" + import spyde.backend.example_catalogue as ec + session = window["window"] + monkeypatch.setattr(ec, "_open_lazily", + lambda path: [_eager_4d(), _eager_4d()]) + sig = session._load_example_lazy("/tmp/multi.h5") + assert sig._lazy is True and not isinstance(sig, list) + class TestNonBlockingNavigator: """The display must NOT wait on the navigator compute, and a tree built diff --git a/uv.lock b/uv.lock index 1ea8674..dd98377 100644 --- a/uv.lock +++ b/uv.lock @@ -823,6 +823,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/02/10/5da547df7a391dcde17f59520a231527b8571e6f46fc8efb02ccb370ab12/docutils-0.22.4-py3-none-any.whl", hash = "sha256:d0013f540772d1420576855455d050a2180186c91c15779301ac2ccb3eeb68de", size = 633196, upload-time = "2025-12-18T19:00:18.077Z" }, ] +[[package]] +name = "em-database" +version = "0.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pooch", marker = "python_full_version >= '3.12'" }, + { name = "pyyaml", marker = "python_full_version >= '3.12'" }, + { name = "tqdm", marker = "python_full_version >= '3.12'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c8/1f/c009752d22cf7274b193313ac9616087c4eaa1167ec9aba59766b14c7a01/em_database-0.1.0.tar.gz", hash = "sha256:1a6bff7e9e8df843e3df457070d068fb892cdfa76faca1f6c1c5a37d040ea945", size = 19675, upload-time = "2026-07-28T13:35:13.731Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1e/ab/a3c64005115d0ab50f09869ea66d1c9d1d1290db040946fbe67aba42d843/em_database-0.1.0-py3-none-any.whl", hash = "sha256:858d1235e471e1f6647d9fd5e3332c570b663b2418a9defa00c757ab29235212", size = 30632, upload-time = "2026-07-28T13:35:12.692Z" }, +] + [[package]] name = "exceptiongroup" version = "1.3.1" @@ -3543,6 +3557,7 @@ dependencies = [ { name = "bokeh" }, { name = "dask" }, { name = "distributed" }, + { name = "em-database", marker = "python_full_version >= '3.12'" }, { name = "huggingface-hub" }, { name = "hyperspy" }, { name = "imagecodecs", version = "2025.3.30", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, @@ -3612,6 +3627,7 @@ requires-dist = [ { name = "bokeh" }, { name = "dask" }, { name = "distributed" }, + { name = "em-database", marker = "python_full_version >= '3.12'", specifier = ">=0.1.0" }, { name = "exspy", marker = "extra == 'eels'", specifier = ">=0.3.2" }, { name = "huggingface-hub" }, { name = "hyperspy", git = "https://github.com/cssfrancis/hyperspy.git?rev=f314cf42aa16a06d82449cd8dc8c8f57df2e3e72" }, From b9ce4a24a6c7c2cddff321b507dff0bbe1e89de3 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Tue, 28 Jul 2026 11:14:18 -0500 Subject: [PATCH 43/60] feat(examples): show the camera on the hover card MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Which detector a dataset came off is most of what tells you what to expect from it — counted vs integrating, frame rate, pixel size — and it is what someone picking an example to compare against their own data looks for. `detector` / `detector_manufacturer` are top-level dataset fields rather than metadata keys, which is easy to miss; a test asserts they keep coming through. Rendered as "Merlin (Quantum Detectors)", falling back to either half alone — a HAADF dataset has a manufacturer but no named detector. --- .../src/renderer/src/components/MenuBar.tsx | 21 ++++++++++++++----- electron/tests/examples_menu.spec.ts | 3 +++ spyde/backend/example_catalogue.py | 7 +++++++ .../tests/migrated/test_example_catalogue.py | 11 ++++++++++ 4 files changed, 37 insertions(+), 5 deletions(-) diff --git a/electron/src/renderer/src/components/MenuBar.tsx b/electron/src/renderer/src/components/MenuBar.tsx index 3e553af..21859dc 100644 --- a/electron/src/renderer/src/components/MenuBar.tsx +++ b/electron/src/renderer/src/components/MenuBar.tsx @@ -40,6 +40,16 @@ export interface ExampleItem { file?: string microscope?: string voltage?: string + detector?: string + detector_manufacturer?: string +} + +/** "CeleritasXS (Direct Electron)" — either half alone when that is all there + * is (a HAADF dataset has a manufacturer but no named detector). */ +function cameraOf(it: ExampleItem): string { + const [name, make] = [it.detector?.trim(), it.detector_manufacturer?.trim()] + if (name && make) return `${name} (${make})` + return name || make || '' } interface ExampleGroup { technique: string; items: ExampleItem[] } @@ -160,11 +170,12 @@ export function MenuBar({ onStartGuide }: { onStartGuide: (g: Guide) => void }) tip: { title: it.label, meta: [it.technique, it.microscope, it.voltage].filter(Boolean).join(' · '), - facts: [ - ...(it.shape ? [['Shape', it.shape] as [string, string]] : []), - ['Size', it.size] as [string, string], - ...(it.file ? [['File', it.file] as [string, string]] : []), - ], + facts: ([ + it.shape ? ['Shape', it.shape] : null, + ['Size', it.size], + cameraOf(it) ? ['Camera', cameraOf(it)] : null, + it.file ? ['File', it.file] : null, + ].filter(Boolean)) as [string, string][], body: it.description, downloaded: it.downloaded, }, diff --git a/electron/tests/examples_menu.spec.ts b/electron/tests/examples_menu.spec.ts index 65eaf33..e6324f9 100644 --- a/electron/tests/examples_menu.spec.ts +++ b/electron/tests/examples_menu.spec.ts @@ -92,6 +92,9 @@ test('hovering a dataset shows a themed info card', async () => { await expect(card).toContainText('SPEDAg') await expect(card).toContainText('4D-STEM') await expect(card).toContainText(/Size/) + // The camera is most of what tells you what to expect from a dataset. + await expect(card).toContainText('Camera') + await expect(card).toContainText('Merlin (Quantum Detectors)') await expect(card).toContainText(/On disk|Not downloaded/) await page.screenshot({ path: `${SHOTS}/04-hover-card.png` }) diff --git a/spyde/backend/example_catalogue.py b/spyde/backend/example_catalogue.py index b270e48..11ac9c0 100644 --- a/spyde/backend/example_catalogue.py +++ b/spyde/backend/example_catalogue.py @@ -335,6 +335,13 @@ def entry(key: str, ds) -> dict: "tags": list(md.get("tags") or []), "microscope": str(md.get("microscope") or ""), "voltage": str(md.get("voltage") or ""), + # The camera, as top-level dataset fields rather than metadata keys. + # Worth surfacing: which detector a dataset came off is most of what + # tells you what to expect from it (counted vs integrating, frame rate, + # pixel size), and it is the field a user picking an example to test + # their own data against actually looks for. + "detector": str(getattr(ds, "detector", "") or ""), + "detector_manufacturer": str(getattr(ds, "detector_manufacturer", "") or ""), "license": str(getattr(ds, "license", "") or ""), "source": str(getattr(ds, "source", "") or ""), } diff --git a/spyde/tests/migrated/test_example_catalogue.py b/spyde/tests/migrated/test_example_catalogue.py index 9c8d214..05ed8fa 100644 --- a/spyde/tests/migrated/test_example_catalogue.py +++ b/spyde/tests/migrated/test_example_catalogue.py @@ -44,6 +44,17 @@ def test_every_entry_has_what_the_menu_draws(self): assert isinstance(item["downloaded"], bool) assert item["size"], f"{item['key']} has no size to show" + def test_the_camera_is_carried_through(self): + """`detector` / `detector_manufacturer` are top-level dataset fields, + not metadata keys — easy to miss, and the hover card shows them.""" + items = [it for g in ec.catalogue()["groups"] for it in g["items"]] + assert all("detector" in it and "detector_manufacturer" in it + for it in items) + with_camera = [it for it in items + if it["detector"] or it["detector_manufacturer"]] + assert len(with_camera) > len(items) // 2, \ + "almost nothing carried a camera — the field name probably moved" + def test_the_base_class_is_not_a_dataset(self): """``em_database.data`` also holds the DownloadableDataset base and an incidental ``Path`` import; neither is an example.""" From 21d74547cfaea71c6eb5123378443f9630ddd0b6 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Tue, 28 Jul 2026 22:51:05 -0500 Subject: [PATCH 44/60] ci: the matrix was timing out on a coverage report nobody reads MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Half of PR #89's matrix showed "fail" at exactly 30m0s. None of them failed a test — they were CANCELLED by timeout-minutes mid-suite. 0.3.0 grew the suite to 2372 tests and legs were landing at 28m30s, so the budget had simply run out. Coverage was most of it. `--cov=spyde --cov-report=xml` costs +45% wall clock (measured 171s -> 248s over the same 25 files), and the Codecov upload is gated to `github.repository_owner == 'directelectron'` — so on every fork all twelve legs paid ~9 minutes to write a coverage.xml that was then discarded. Run it on one leg, and gate the upload to that same leg so the two can't drift. Also here: - `extras` died in 42s before running a single test: atomap imports fine but exposes no `__version__`. Ask importlib.metadata instead, which every installed distribution answers. - The e2e installed `--extra tests` only, so `requires_package:` hid every EELS/EBSD/atoms action and nine 0.3.0 specs could not pass no matter what — each paying a second full app boot for its retry. Install `--extra all`; that alone takes ~10 dead minutes off shard 1. - Shard ×3 ran 8.1m / 14.5m / >30m, and the long one was cancelled on test 76 of 115 so it uploaded no report at all. Go to ×4. --- .github/workflows/build.yml | 57 ++++++++++++++++++++++++++----------- 1 file changed, 40 insertions(+), 17 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 669f385..55890b3 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -21,7 +21,12 @@ jobs: test: name: ${{ matrix.os }}-py${{ matrix.python-version }} runs-on: ${{ matrix.os }} - timeout-minutes: 30 + # 30 was too tight once 0.3.0 grew the suite to 2372 tests: legs landed at + # 28m30s and half the matrix was CANCELLED mid-run at exactly 30m, which + # reads as a test failure but is a clock. The run is much shorter now (see + # the coverage note below); this is headroom, and a timeout never reached + # costs nothing. + timeout-minutes: 40 strategy: fail-fast: false matrix: @@ -58,11 +63,18 @@ jobs: # The single canonical Qt-free suite (testpaths=spyde/tests/migrated). # GPU tests skip themselves without CUDA/MPS. No display server required. + # + # Coverage on ONE leg only. `--cov` costs +45% wall-clock (measured + # 171s -> 248s over the same 25 files) and the upload below only runs on + # directelectron/spyde, so on every fork the other 11 legs paid ~9 minutes + # to write a coverage.xml nobody reads — which is what pushed half the + # matrix past timeout-minutes and got it cancelled mid-suite. - name: Run tests - run: uv run pytest spyde/tests/migrated --cov=spyde --cov-report=xml + run: uv run pytest spyde/tests/migrated ${{ (matrix.os == 'ubuntu-latest' && matrix.python-version == '3.12') && '--cov=spyde --cov-report=xml' || '' }} - name: Upload coverage to Codecov - if: ${{ github.repository_owner == 'directelectron' }} + if: ${{ github.repository_owner == 'directelectron' + && matrix.os == 'ubuntu-latest' && matrix.python-version == '3.12' }} uses: codecov/codecov-action@v5 with: token: ${{ secrets.CODECOV_TOKEN }} @@ -104,10 +116,14 @@ jobs: - name: Versions run: | uv run python -V - uv run python -c "import exspy, kikuchipy, atomap; \ - print('exspy', exspy.__version__); \ - print('kikuchipy', kikuchipy.__version__); \ - print('atomap', atomap.__version__)" + # importlib.metadata, not `pkg.__version__` — atomap imports fine but + # exposes no `__version__`, which failed this job before a single test + # ran. The point of the step is "the extra installed and imports", so + # ask the installed distribution, which every package answers. + uv run python -c "import importlib.metadata as m; \ + import exspy, kikuchipy, atomap; \ + print('\n'.join(f'{p} {m.version(p)}' \ + for p in ('exspy', 'kikuchipy', 'atomap')))" - name: Run tests (extras present) run: uv run pytest spyde/tests/migrated @@ -183,24 +199,24 @@ jobs: # Best-effort for now (continue-on-error): the e2e launches a real Electron app # + Dask LocalCluster and has not yet been validated green on a hosted runner. # Drop continue-on-error once it's confirmed stable in CI. - # Sharded ×3: the suite is deliberately serial WITHIN a run (one Electron + + # Sharded ×4: the suite is deliberately serial WITHIN a run (one Electron + # Dask cluster at a time — see playwright.config.ts), so wall-clock scales # only by splitting the spec files across independent runners. Each shard is # still fully serial, so the cluster-contention flakiness that forced # workers:1 cannot recur across shards (separate VMs). Playwright balances - # shards by FILE COUNT, not duration — ×2 measured 22.7m vs 10.0m — so ×3 - # keeps the worst shard under ~15m even while the known-failing specs pay - # retry + app-reboot tax (2026-07-11: 16 failing tests account for ~23 of - # the 33 spec-minutes; fixing them is the real speedup). + # shards by FILE COUNT, not duration — ×2 measured 22.7m vs 10.0m, and ×3 + # still ran 8.1m / 14.5m / >30m (shard 1 was CANCELLED at the timeout on + # test 76 of 115, so it produced no report at all). ×4 plus the domain + # extras below brings the worst shard back under the budget. e2e: - name: electron-e2e (ubuntu, shard ${{ matrix.shard }}/3) + name: electron-e2e (ubuntu, shard ${{ matrix.shard }}/4) runs-on: ubuntu-latest - timeout-minutes: 30 + timeout-minutes: 45 continue-on-error: true strategy: fail-fast: false matrix: - shard: [1, 2, 3] + shard: [1, 2, 3, 4] steps: - uses: actions/checkout@v4 with: @@ -220,8 +236,15 @@ jobs: version: "0.10.11" enable-cache: true + # WITH the domain extras. Without them the `requires_package:` gate hides + # every EELS/EBSD/atoms action, so the 0.3.0 specs that click those + # buttons cannot pass — 9 of them failed for that reason alone, and with + # retries:1 each one paid a second full app boot. Installing the extras + # makes them testable AND takes ~10 minutes of dead retry off shard 1. + # Not --frozen: the extras resolve per-platform, so the lock can't satisfy + # them (same reason as the `extras` job above). - name: Pre-sync the Python backend env - run: uv sync --frozen --extra tests + run: uv sync --extra tests --extra all - name: Install Electron dependencies working-directory: electron @@ -237,7 +260,7 @@ jobs: # downloads + screenshot writes) that is meant to be opt-in. - name: Run Electron e2e working-directory: electron - run: xvfb-run --auto-servernum npm run test:build -- --project=electron --shard=${{ matrix.shard }}/3 + run: xvfb-run --auto-servernum npm run test:build -- --project=electron --shard=${{ matrix.shard }}/4 - name: Upload Playwright report if: ${{ always() }} From 56c40ff1a8c9053fd36cd07e902293792a7aa95f Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Tue, 28 Jul 2026 22:56:21 -0500 Subject: [PATCH 45/60] perf(tests): stop sleeping half a second per test to reap a cluster that was never started MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The suite was not slow because of the tests. `Session.shutdown()` -> `DaskManager.shutdown()` ran its settle-and-reap tail unconditionally: a flat `time.sleep(0.5)`, a psutil walk of the machine's entire process tree, `wait_procs(timeout=1.0)`, and a full `gc.collect()`. Under SPYDE_NO_DASK — which conftest sets for every pytest Session, and which headless scripts use too — `start()` is never called, so there is no client, no cluster, no nanny and no worker grandchild anywhere. The tail was reaping nothing, slowly. Measured 669ms per teardown, and the fixtures build a Session per test: 2372 x 0.67s is ~26 of the suite's ~28 minutes. Return early when this manager never brought anything up. Gated on `_thread` rather than on client/cluster being None, because `start()` builds the cluster on a background thread — a shutdown racing a startup can legitimately see both attributes unset while a LocalCluster spawns. If start() was ever called the reap runs exactly as before, so the real app is untouched. Session teardown: 669ms -> 0ms. Same 25 files that took 171s serial now take 75s, with the same 340 tests passing. --- spyde/dask_manager.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/spyde/dask_manager.py b/spyde/dask_manager.py index a0b4df2..7ab5068 100644 --- a/spyde/dask_manager.py +++ b/spyde/dask_manager.py @@ -480,6 +480,26 @@ def shutdown(self) -> None: finally: self._cluster = None + # Nothing was ever brought up -> nothing to settle, reap or collect. + # + # `start()` is what spawns the cluster, and it is the only thing that + # sets `_thread`. Under SPYDE_NO_DASK (every pytest Session, plus + # headless scripts) it is never called, so there is no client, no + # cluster, no nanny and no worker grandchild in existence — yet the + # tail below still slept 0.5s, walked the machine's whole process tree + # via psutil, and ran a full gc.collect(). That is ~670ms per Session + # teardown, and the test fixtures build one per test: 2372 tests x 0.67s + # was ~26 of the suite's ~28 minutes, and it is why the CI matrix + # started getting cancelled at its 30-minute timeout. + # + # Deliberately gated on `_thread`, not just on client/cluster being + # None: `start()` builds the cluster on a background thread, so a + # shutdown racing a startup can legitimately see both attributes still + # unset while a LocalCluster is spawning. If start() was ever called we + # always reap, exactly as before. + if self._thread is None and client is None and cluster is None: + return + time.sleep(0.5) # Reap multiprocessing children we own directly. try: From 736f77e7caffaacd152b0c1dfa0c7780d3c277cb Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Tue, 28 Jul 2026 23:09:45 -0500 Subject: [PATCH 46/60] test(dask): pin both halves of the shutdown reap contract MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Skip the settle/reap tail when nothing was ever started, and still reap when it was. Verified by reverting the guard: the three "nothing started" tests fail, the four reap-safety tests pass either way — which is the point, since those are the ones protecting users from a leaked worker outliving the app. --- .../tests/migrated/test_dask_shutdown_fast.py | 113 ++++++++++++++++++ 1 file changed, 113 insertions(+) create mode 100644 spyde/tests/migrated/test_dask_shutdown_fast.py diff --git a/spyde/tests/migrated/test_dask_shutdown_fast.py b/spyde/tests/migrated/test_dask_shutdown_fast.py new file mode 100644 index 0000000..e15b24e --- /dev/null +++ b/spyde/tests/migrated/test_dask_shutdown_fast.py @@ -0,0 +1,113 @@ +"""DaskManager.shutdown() must not reap a cluster that was never started. + +The reap tail — settle sleep, psutil walk of the machine's whole process +tree, wait_procs, gc.collect() — exists to make sure a nanny's worker +GRANDCHILD dies with the app. Under SPYDE_NO_DASK (every pytest Session, +and any headless script) nothing was ever spawned, so it was reaping +nothing at a flat ~670ms a go. The fixtures build a Session per test, so +that was ~26 minutes of a ~28-minute suite, and it is what pushed the CI +matrix past its job timeout. + +Both halves matter, so both are pinned here: skip when nothing ran, and +still reap when something did. The second is the one that protects users — +a leaked Dask worker outlives the app and holds its memory. +""" +from __future__ import annotations + +import time + +import pytest + +from spyde.dask_manager import DaskManager + + +class TestNothingStarted: + def test_shutdown_is_immediate_when_start_was_never_called(self): + dm = DaskManager(n_workers=1, threads_per_worker=1) + t0 = time.perf_counter() + dm.shutdown() + elapsed = time.perf_counter() - t0 + # The old path slept 0.5s flat before it even began walking processes. + assert elapsed < 0.2, ( + f"shutdown() took {elapsed:.3f}s with no cluster ever started — the " + "settle/reap tail is running again; that is ~0.67s on every one of " + "the suite's Session teardowns" + ) + + def test_it_does_not_touch_the_process_tree(self, monkeypatch): + """The psutil walk is the expensive half on Windows (~30ms for a + ppid_map of the whole box) and is meaningless with no children.""" + import psutil + + called: list[str] = [] + monkeypatch.setattr( + psutil.Process, "children", + lambda self, recursive=False: called.append("children") or [], + ) + DaskManager(n_workers=1, threads_per_worker=1).shutdown() + assert called == [], "shutdown() walked the process tree with no cluster" + + def test_repeated_shutdown_stays_cheap(self): + """Session.shutdown() is called on every fixture teardown, and some + paths call it twice; neither may pay the tail.""" + dm = DaskManager(n_workers=1, threads_per_worker=1) + t0 = time.perf_counter() + for _ in range(5): + dm.shutdown() + assert time.perf_counter() - t0 < 0.5 + + +class TestSomethingStarted: + """The half that protects users: if start() ever ran, reap unconditionally. + + Gated on `_thread` rather than on client/cluster, because start() builds + the cluster on a BACKGROUND thread — a shutdown racing a startup can + legitimately see both attributes still unset while a LocalCluster spawns. + Skipping there would leak the workers it was about to create. + """ + + def _instrument(self, monkeypatch): + import psutil + + seen: list[str] = [] + monkeypatch.setattr( + psutil.Process, "children", + lambda self, recursive=False: seen.append("children") or [], + ) + monkeypatch.setattr(time, "sleep", lambda s: seen.append(f"sleep{s}")) + return seen + + def test_a_started_manager_still_reaps(self, monkeypatch): + dm = DaskManager(n_workers=1, threads_per_worker=1) + # What start() leaves behind, without actually spawning a cluster. + dm._thread = object() + seen = self._instrument(monkeypatch) + dm.shutdown() + assert "children" in seen, ( + "a manager that called start() skipped the reap — a nanny's worker " + "grandchild would outlive the app" + ) + + def test_a_live_cluster_still_reaps_even_without_the_thread(self, monkeypatch): + """restart() and the test helpers set _cluster directly.""" + class _Cluster: + def scale(self, n): pass + def close(self, timeout=None): pass + + dm = DaskManager(n_workers=1, threads_per_worker=1) + dm._cluster = _Cluster() + seen = self._instrument(monkeypatch) + dm.shutdown() + assert "children" in seen + + @pytest.mark.parametrize("attr", ["_client", "_cluster"]) + def test_either_handle_alone_is_enough_to_reap(self, monkeypatch, attr): + class _Handle: + def scale(self, n): pass + def close(self, timeout=None): pass + + dm = DaskManager(n_workers=1, threads_per_worker=1) + setattr(dm, attr, _Handle()) + seen = self._instrument(monkeypatch) + dm.shutdown() + assert "children" in seen From 3370c5d86f1fa8f85c0fe74f9098b5949d1dc8c6 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Tue, 28 Jul 2026 23:19:38 -0500 Subject: [PATCH 47/60] fix(tests): two races the suite's own slowness was hiding MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With the matrix now finishing in ~11-13m instead of being cancelled at 30, two tests started failing on 6 of 12 legs — inconsistently, and on no particular OS or Python. Neither is caused by the speedup; both were always racing and the old pacing hid it. PlotUpdateWorker.stop() only set a flag. The poll thread lived on for up to another tick plus whatever _check() was mid-flight, so a shut-down Session's poller kept walking its plots and painting. test_nav_progressive patches Plot.set_data on the CLASS, so it saw one: it measured a leftover 16x16 navigator instead of its own 12x12 and asserted 256 == 12 * 12. stop() now joins, bounded and never from inside the poll thread, and the test only counts plots belonging to its own session. Teardown goes from 0ms to 5ms — 0.3 minutes across the suite, for a stop() that stops. test_fit_wizard asserted on fit_catalogue the instant fit_open returned, but _send_catalogue is deliberately asynchronous: the first call builds nine hyperspy components (~680ms of sympy lambdify) and used to stall the caret opening. A real Session marshals that off-thread, so the assert only passed when an earlier test happened to warm the catalogue first. Wait for the message. --- spyde/tests/migrated/test_fit_wizard.py | 23 +++++++++++++++++++- spyde/tests/migrated/test_nav_progressive.py | 19 ++++++++++++++-- spyde/workers/plot_update_worker.py | 20 +++++++++++++++-- 3 files changed, 57 insertions(+), 5 deletions(-) diff --git a/spyde/tests/migrated/test_fit_wizard.py b/spyde/tests/migrated/test_fit_wizard.py index 26f02ad..a0825a0 100644 --- a/spyde/tests/migrated/test_fit_wizard.py +++ b/spyde/tests/migrated/test_fit_wizard.py @@ -65,6 +65,27 @@ def _messages_of(window, kind): return [m for m in window["messages"] if m.get("type") == kind] +def _await_messages(window, kind, timeout=30.0): + """Same, but for a message the handler emits ASYNCHRONOUSLY. + + `fit_open` returns before the palette exists on purpose: `_send_catalogue` + samples it via `run_on_worker` because the first call in a process builds + nine hyperspy components (~680 ms of sympy lambdify) and used to stall the + caret. A real Session has `_dispatch_to_main`, so that genuinely lands on + another thread — asserting straight after the call only passed when an + earlier test happened to warm the catalogue first. It was the flakiest + test in CI (5 of 12 legs). + """ + import time + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + got = _messages_of(window, kind) + if got: + return got + time.sleep(0.01) + return [] + + class TestOpenClose: def test_open_creates_a_wizard_and_sends_state(self, window, fitted): session, plot, tree, _ = fitted @@ -76,7 +97,7 @@ def test_open_sends_the_component_palette(self, window, fitted): """#56 — the picker needs SHAPES, not just names.""" session, plot, tree, _ = fitted fit_open(session, plot, {}) - cat = _messages_of(window, "fit_catalogue") + cat = _await_messages(window, "fit_catalogue") assert cat, "no catalogue was sent" kinds = {c["kind"] for c in cat[-1]["components"]} assert {"Gaussian", "PowerLaw", "Offset"} <= kinds diff --git a/spyde/tests/migrated/test_nav_progressive.py b/spyde/tests/migrated/test_nav_progressive.py index 628e0f8..b6fd774 100644 --- a/spyde/tests/migrated/test_nav_progressive.py +++ b/spyde/tests/migrated/test_nav_progressive.py @@ -229,9 +229,25 @@ def test_navigator_has_no_holes_when_complete(self, monkeypatch): last_size = [0] orig = Plot.set_data + # The session first, so the tracker can tell OUR navigator from anyone + # else's. The patch is on the CLASS, so it sees every Plot in the + # process — including one belonging to an earlier test whose poller + # thread has not wound down yet. That is not hypothetical: this test + # failed in CI with `assert 256 == (12 * 12)`, i.e. it measured a + # leftover 16x16 navigator and never looked at its own. + session = _make_session() + + def own(p): + if any(p is q for q in session._plots): + return True + # Registration in _plots can trail the first paint, so also accept + # a plot whose tree is one of ours. + tree = getattr(p, "signal_tree", None) + return tree is not None and any(tree is t for t in session.signal_trees) + def _track(self, data, *a, **k): try: - if getattr(self, "is_navigator", False): + if getattr(self, "is_navigator", False) and own(self): arr = np.asarray(data) last_finite[0] = int(np.isfinite(arr).sum()) last_size[0] = int(arr.size) @@ -241,7 +257,6 @@ def _track(self, data, *a, **k): monkeypatch.setattr(Plot, "set_data", _track) - session = _make_session() try: ny = nx = 12 base = np.arange(ny * nx * 8 * 8, dtype=np.float32).reshape(ny, nx, 8, 8) diff --git a/spyde/workers/plot_update_worker.py b/spyde/workers/plot_update_worker.py index e1f598d..39af80b 100644 --- a/spyde/workers/plot_update_worker.py +++ b/spyde/workers/plot_update_worker.py @@ -60,10 +60,26 @@ def start(self) -> None: ) self._thread.start() - def stop(self) -> None: - """Request the worker to stop. Returns immediately; thread exits on next tick.""" + def stop(self, timeout: float = 1.0) -> None: + """Stop the worker and wait for its thread to actually be gone. + + This used to set the flag and return, leaving the poll thread alive for + up to one more tick — plus however long a `_check()` already in flight + took, which can emit. A poller still walking a shut-down Session's + plots is a small leak in the app and outright confusing in tests: with + `Session.shutdown()` now fast, a previous test's worker could still be + painting while the next one measured, and one test read a leftover + 16x16 navigator as its own 12x12 (`assert 256 == (12 * 12)`). + + Bounded, and never joins from inside the poll thread, so it cannot + wedge — the thread is a daemon either way. + """ self._running = False self._seen.clear() + t = self._thread + if t is not None and t.is_alive() and t is not threading.current_thread(): + t.join(timeout=timeout) + self._thread = None def _loop(self) -> None: while self._running: From b7bab96089892c5d9eba27c50cdfdd835be062dd Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Tue, 28 Jul 2026 23:27:28 -0500 Subject: [PATCH 48/60] fix(e2e): the welcome tour was eating pointer events on every CI runner MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit action_scoping.spec.ts failed with "
from
subtree intercepts pointer events" on a hover that works on any dev box. The difference is not the app: a machine that has actually RUN SpyDE has tutorial_seen persisted in ~/.spyde, and a fresh runner does not — so FirstRunGate auto-opens the welcome tour and its full-screen overlay swallows the click. launchApp now always points SPYDE_SETTINGS_DIR at a fresh scratch dir seeded with tutorial_seen, so the outcome no longer depends on whose machine it is (and no spec can leak a persisted setting into the next). A spec's own SPYDE_SETTINGS_DIR still wins, which is how first_run.spec.ts gets a genuine first launch — both of its cases still pass. Pinned by a third case there: a default launch shows no overlay and a menu-bar hover lands. Verified in the app, screenshot read. --- electron/tests/_harness.cjs | 29 ++++++++++++++++++++++++++++- electron/tests/first_run.spec.ts | 29 +++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 1 deletion(-) diff --git a/electron/tests/_harness.cjs b/electron/tests/_harness.cjs index 110bada..5e9f364 100644 --- a/electron/tests/_harness.cjs +++ b/electron/tests/_harness.cjs @@ -18,12 +18,38 @@ */ const { _electron: electron } = require('@playwright/test') const { join } = require('path') +const { mkdtempSync, mkdirSync, writeFileSync } = require('fs') +const { tmpdir } = require('os') + +/** A scratch ~/.spyde with the welcome tour already dismissed. + * + * On a machine that has never RUN SpyDE — i.e. every CI runner — settings.json + * is absent, `tutorial_seen` is unset, and FirstRunGate auto-opens the welcome + * tour. Its full-screen `tour-overlay` then swallows pointer events, so a spec + * that hovers a titlebar fails with "
from
+ * subtree intercepts pointer events" — while passing on any dev box, because + * there the real ~/.spyde has tutorial_seen persisted from actual use. + * + * SPYDE_SETTINGS_DIR redirects only settings.json, never Electron's own + * profile (Chromium refuses to launch without a real one — see the note in + * first_run.spec.ts). A fresh dir per launch also means no spec can leak a + * persisted setting into the next. + */ +function _seenSettingsDir() { + const dir = join(mkdtempSync(join(tmpdir(), 'spyde-e2e-')), '.spyde') + mkdirSync(dir, { recursive: true }) + writeFileSync(join(dir, 'settings.json'), + JSON.stringify({ tutorial_seen: true })) + return dir +} /** * Launch the app. Returns { app, page, backend, jsErrors, assertNoJsErrors }. * @param {{dask?: boolean, env?: Record}} opts * dask=false (default) → SPYDE_NO_DASK=1 (renderer-only / eager, fast). * dask=true → real LocalCluster + client. + * env.SPYDE_SETTINGS_DIR → pass your own to opt OUT of the tour suppression + * above (first_run.spec.ts does exactly this to get a genuine first launch). */ async function launchApp(opts = {}) { const { dask = false, env = {} } = opts @@ -32,7 +58,8 @@ async function launchApp(opts = {}) { env: { ...process.env, ...(dask ? {} : { SPYDE_NO_DASK: '1' }), - ...env, + SPYDE_SETTINGS_DIR: _seenSettingsDir(), + ...env, // a spec's own SPYDE_SETTINGS_DIR wins }, }) diff --git a/electron/tests/first_run.spec.ts b/electron/tests/first_run.spec.ts index 471f82a..787dccf 100644 --- a/electron/tests/first_run.spec.ts +++ b/electron/tests/first_run.spec.ts @@ -103,4 +103,33 @@ test.describe('first-run welcome walkthrough', () => { rmSync(home, { recursive: true, force: true }) } }) + + test('a default harness launch never auto-opens the tour', async () => { + // Every OTHER spec launches without its own SPYDE_SETTINGS_DIR, and none of + // them wants the welcome tour: its full-screen overlay swallows pointer + // events, so a titlebar hover fails with "
from
subtree intercepts pointer events". That is + // exactly how action_scoping.spec.ts failed on CI while passing on every + // dev box — a machine that has actually RUN SpyDE has tutorial_seen + // persisted in the real ~/.spyde, and a fresh runner does not. + // + // launchApp now always points SPYDE_SETTINGS_DIR at a scratch dir seeded + // with tutorial_seen, so the outcome no longer depends on whose machine it + // is. The two tests above prove the opt-out still works: pass your own + // settings dir and you get a genuine first launch. + const ctx = await launchApp({ dask: false }) + try { + await ctx.page.waitForTimeout(3000) // past FirstRunGate's get_first_run + await expect(ctx.page.getByTestId('tour-overlay')).toHaveCount(0) + await ctx.page.screenshot({ path: join(SHOTS, '05-default-launch-no-tour.png') }) + + // And the thing the overlay was blocking — a hover — actually lands. + const bar = ctx.page.getByTestId('menu-bar') + await expect(bar).toBeVisible() + await bar.hover() + ctx.assertNoJsErrors() + } finally { + await ctx.app.close() + } + }) }) From 067e9f4f45ad6f4d5c1d54abef744965476005cd Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Tue, 28 Jul 2026 23:32:28 -0500 Subject: [PATCH 49/60] test(e2e): only mint a scratch settings dir when the spec has not brought one Otherwise every first_run launch also left an unused one behind. Its three cases still pass. --- electron/tests/_harness.cjs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/electron/tests/_harness.cjs b/electron/tests/_harness.cjs index 5e9f364..2cb8c64 100644 --- a/electron/tests/_harness.cjs +++ b/electron/tests/_harness.cjs @@ -58,8 +58,10 @@ async function launchApp(opts = {}) { env: { ...process.env, ...(dask ? {} : { SPYDE_NO_DASK: '1' }), - SPYDE_SETTINGS_DIR: _seenSettingsDir(), - ...env, // a spec's own SPYDE_SETTINGS_DIR wins + // Only mint a scratch dir if the spec didn't bring its own — otherwise + // every first_run launch would also leave an unused one behind. + ...(env.SPYDE_SETTINGS_DIR ? {} : { SPYDE_SETTINGS_DIR: _seenSettingsDir() }), + ...env, }, }) From d3b34ad6692196631ad33f973654b13b08a2a4dd Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Tue, 28 Jul 2026 23:36:15 -0500 Subject: [PATCH 50/60] =?UTF-8?q?docs(ci):=20correct=20the=20coverage=20co?= =?UTF-8?q?st=20=E2=80=94=20it=20is=20not=20why=20the=20matrix=20timed=20o?= =?UTF-8?q?ut?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The +45% figure came from this Windows dev box. On a hosted Linux runner it is ~15%: same commit, ubuntu-py3.12 with coverage 12.5m vs the uncovered legs at ~10.8m. Gating it to one leg is still right (the upload only runs on directelectron, so every other leg wrote a report nobody read) but it buys ~1.5m a leg, not the ~9m the comment claimed. The timeouts were DaskManager.shutdown(). --- .github/workflows/build.yml | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 55890b3..03f560d 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -64,11 +64,14 @@ jobs: # The single canonical Qt-free suite (testpaths=spyde/tests/migrated). # GPU tests skip themselves without CUDA/MPS. No display server required. # - # Coverage on ONE leg only. `--cov` costs +45% wall-clock (measured - # 171s -> 248s over the same 25 files) and the upload below only runs on - # directelectron/spyde, so on every fork the other 11 legs paid ~9 minutes - # to write a coverage.xml nobody reads — which is what pushed half the - # matrix past timeout-minutes and got it cancelled mid-suite. + # Coverage on ONE leg only: the upload below runs solely on + # directelectron/spyde, so on every fork the other 11 legs wrote a + # coverage.xml nobody reads. Cost is +45% on a Windows dev box (measured + # 171s -> 248s over the same 25 files) but only ~15% on a hosted Linux + # runner (12.5m with vs ~10.8m without, same commit) — so this is tidiness + # and ~1.5m a leg, NOT the reason the matrix used to time out. That was + # DaskManager.shutdown() sleeping 0.5s per Session teardown; see + # test_dask_shutdown_fast.py. - name: Run tests run: uv run pytest spyde/tests/migrated ${{ (matrix.os == 'ubuntu-latest' && matrix.python-version == '3.12') && '--cov=spyde --cov-report=xml' || '' }} From 477be25e0149efcae17ae0f01c5e780059538c2b Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Tue, 28 Jul 2026 23:43:36 -0500 Subject: [PATCH 51/60] fix(session): shutdown never stopped the trees' in-flight compute MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Found while reviewing the teardown change: running test_console_preview alone now printed a background traceback and a faulthandler module dump during interpreter teardown — "threaded navigator compute failed", signal_tree._bg_nav submitting into an executor that is already going away. Reverting the two teardown commits made it disappear, so this is mine. The cause is not new, only newly visible. BaseSignalTree.close() exists for exactly this: it flips every registered stopped_flag, cancels every registered future, and sets _nav_stop so the progressive-navigator thread bails on its next poll. Session.shutdown() simply never called it. That went unnoticed while shutdown() slept half a second on its way out — the fill usually finished inside the sleep. Close the trees first, in the same cancel-then-teardown order close() itself uses. Costs ~10ms a test (340 tests: 75.0s -> 78.6s) and buys background compute that actually stops rather than one that races the executors. Three runs of the file: 36 passed, zero teardown noise. --- spyde/backend/session.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/spyde/backend/session.py b/spyde/backend/session.py index 2d642d8..93d6b75 100644 --- a/spyde/backend/session.py +++ b/spyde/backend/session.py @@ -674,6 +674,24 @@ def shutdown(self) -> None: except Exception as e: log.debug("console shutdown failed: %s", e) self._closed = True # block compute_backend from recreating _nav_executor + # Stop the trees' in-flight compute BEFORE the executors go away. + # + # BaseSignalTree.close() exists for exactly this — it flips every + # registered stopped_flag, cancels every registered future, and sets + # _nav_stop so the progressive-navigator thread bails on its next poll — + # but shutdown() never called it. That did not show while shutdown() + # slept half a second on its way out: the background fill usually + # finished inside the sleep. With that gone, a live `_bg_nav` submits + # into an executor that is already shutting down and logs "threaded + # navigator compute failed" during interpreter teardown. + # + # Cancel-then-teardown is the same order close() itself uses, and it is + # the order that matters: reversing it is what produces the raise. + for tree in list(getattr(self, "signal_trees", []) or []): + try: + tree.close() + except Exception as e: + log.debug("closing signal tree during shutdown failed: %s", e) stats = getattr(self, "_dask_stats", None) if stats is not None: try: From 87ddf24325e56dabc3df6381549d7c9d9e676e9f Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Wed, 29 Jul 2026 00:06:49 -0500 Subject: [PATCH 52/60] docs(test): note how the real-cluster reap was verified MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit These stub the cluster because a real LocalCluster costs ~15s — exactly what this change exists to remove from the suite. Verified out-of-band instead: 2 workers, shutdown() 1.99s (settle + psutil walk + wait_procs all ran), zero survivors. --- spyde/tests/migrated/test_dask_shutdown_fast.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/spyde/tests/migrated/test_dask_shutdown_fast.py b/spyde/tests/migrated/test_dask_shutdown_fast.py index e15b24e..bcbb04e 100644 --- a/spyde/tests/migrated/test_dask_shutdown_fast.py +++ b/spyde/tests/migrated/test_dask_shutdown_fast.py @@ -11,6 +11,13 @@ Both halves matter, so both are pinned here: skip when nothing ran, and still reap when something did. The second is the one that protects users — a leaked Dask worker outlives the app and holds its memory. + +These stub the cluster rather than spinning one: a real ``LocalCluster`` +costs ~15 s, which is precisely the kind of thing this change exists to +remove from the suite. The real path was verified out-of-band instead — +a 2-worker cluster, ``shutdown()`` taking 1.99 s (so the settle + psutil +walk + ``wait_procs`` all ran) and zero surviving children. Re-run that +by hand if you touch the reap; do not add it here. """ from __future__ import annotations From dd342a05d1e4a618d9bd403fd2999b9d2e8020dd Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Wed, 29 Jul 2026 06:27:12 -0500 Subject: [PATCH 53/60] fix: distinct names for the two Fit maps windows, and tour/shell isolation for every spec MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three e2e failures, three different causes, all now fixed. fit_wizard could not target the committed maps window because there are TWO by design — the caret's live one, refreshed as positions fit, and the snapshot commit() keeps — and _open_maps titled both "Fit components". show_maps()'s own docstring warns about exactly this ("stacked identical-looking windows with no way to tell which was current") one paragraph above the call that did it. The live one is now "Fit components (live)"; the snapshot keeps the plain name, since it is the one that persists. The spec matches the breadcrumb anchored with $ and asserts both windows exist, so a user can tell the moving one from the kept one. examples_menu's assertion passed and then its afterAll timed out for 120s on app.close(): shell.openPath shells out to xdg-open, which on a headless runner has no file manager to reach. SPYDE_NO_SHELL_OPEN makes main log the resolved path instead — same test-seam shape as SPYDE_SETTINGS_DIR. The spec now also asserts the path reached the MAIN process, which the backend log cannot see, so this covers more than before rather than less. center_zero_beam failed on the tour overlay eating a hover — the bug b7bab96 supposedly fixed. It did not reach it: ~30 specs call _electron.launch() directly instead of using _harness.cjs. They all spread ...process.env, so a globalSetup that sets SPYDE_SETTINGS_DIR (a scratch dir with tutorial_seen) and SPYDE_NO_SHELL_OPEN is what actually reaches all of them. first_run.spec.ts passes its own settings dir and still gets a genuine first launch; tour.spec.ts still opens the tour by hand. Verified: fit_wizard, examples_menu, center_zero_beam, tour, action_scoping and all three first_run cases pass. --- electron/playwright.config.ts | 5 ++++ electron/src/main/index.ts | 12 ++++++++++ electron/tests/_harness.cjs | 5 ++++ electron/tests/examples_menu.spec.ts | 11 +++++++++ electron/tests/fit_wizard.spec.ts | 18 ++++++++++++-- electron/tests/global-setup.cjs | 36 ++++++++++++++++++++++++++++ spyde/actions/fit_action.py | 18 +++++++++++--- 7 files changed, 100 insertions(+), 5 deletions(-) create mode 100644 electron/tests/global-setup.cjs diff --git a/electron/playwright.config.ts b/electron/playwright.config.ts index 03e475f..a469e63 100644 --- a/electron/playwright.config.ts +++ b/electron/playwright.config.ts @@ -3,6 +3,11 @@ import { join } from 'path' export default defineConfig({ testDir: './tests', + // Settings isolation for EVERY spec — ~30 of them call _electron.launch() + // directly rather than going through _harness.cjs, and on a fresh CI runner + // they get the first-run welcome tour whose overlay eats pointer events. + // They all spread ...process.env, so setting it there is what reaches them. + globalSetup: require.resolve('./tests/global-setup.cjs'), timeout: 120_000, expect: { timeout: 15_000 }, retries: 1, diff --git a/electron/src/main/index.ts b/electron/src/main/index.ts index 81025e2..f967d08 100644 --- a/electron/src/main/index.ts +++ b/electron/src/main/index.ts @@ -868,6 +868,18 @@ ipcMain.on('open-path', async (_, target: string) => { console.warn(`[spyde] open-path rejected non-directory: ${resolved}`) return } + // Test seam, same shape as SPYDE_SETTINGS_DIR: handle the request and log + // it, but do not hand the path to the desktop. On a headless CI runner + // `shell.openPath` shells out to xdg-open, which has no file manager to + // reach and leaves something behind that stops the app exiting — the + // examples_menu spec's own assertion passed and then its afterAll timed + // out for 120s on app.close(). There is nothing to verify past this point + // in CI anyway ("did a file manager window appear" is not observable), so + // the spec still covers backend -> renderer -> main and stops here. + if (process.env.SPYDE_NO_SHELL_OPEN === '1') { + console.log(`[spyde] open-path (suppressed): ${resolved}`) + return + } const err = await shell.openPath(resolved) if (err) console.warn(`[spyde] open-path failed: ${err}`) } catch (e) { diff --git a/electron/tests/_harness.cjs b/electron/tests/_harness.cjs index 2cb8c64..d8da88f 100644 --- a/electron/tests/_harness.cjs +++ b/electron/tests/_harness.cjs @@ -61,6 +61,11 @@ async function launchApp(opts = {}) { // Only mint a scratch dir if the spec didn't bring its own — otherwise // every first_run launch would also leave an unused one behind. ...(env.SPYDE_SETTINGS_DIR ? {} : { SPYDE_SETTINGS_DIR: _seenSettingsDir() }), + // Never hand a path to the desktop from a test. On a headless runner + // xdg-open has no file manager to reach and leaves the app unable to + // exit — examples_menu's afterAll timed out for 120s on app.close(). + // See the open-path handler in src/main/index.ts. + SPYDE_NO_SHELL_OPEN: '1', ...env, }, }) diff --git a/electron/tests/examples_menu.spec.ts b/electron/tests/examples_menu.spec.ts index e6324f9..98a9301 100644 --- a/electron/tests/examples_menu.spec.ts +++ b/electron/tests/examples_menu.spec.ts @@ -138,5 +138,16 @@ test('Show Example Data Directory reports a real path', async () => { (l: string) => l.includes('revealing example data directory')) expect(line, 'the backend never reported the directory').toBeTruthy() expect(line).toContain('em_database') + + // …and that the path actually reached the MAIN process, which is the half + // the backend log cannot see. The harness sets SPYDE_NO_SHELL_OPEN so main + // logs the resolved directory instead of handing it to the desktop — on a + // headless runner xdg-open has no file manager to reach and left the app + // unable to exit, so this test's afterAll timed out for 120s on app.close() + // even though the assertion above had already passed. + await backend.waitForLog('open-path (suppressed)', 20_000) + const mainLine = backend.logBuffer.find( + (l: string) => l.includes('open-path (suppressed)')) + expect(mainLine, 'the main process never received the path').toContain('em_database') ctx.assertNoJsErrors() }) diff --git a/electron/tests/fit_wizard.spec.ts b/electron/tests/fit_wizard.spec.ts index 8d05e3f..5d43701 100644 --- a/electron/tests/fit_wizard.spec.ts +++ b/electron/tests/fit_wizard.spec.ts @@ -155,8 +155,22 @@ test('Fit caret: build a model, run it, commit component maps', async () => { await page.waitForTimeout(2_500) // The committed window carries one chip per component — the same // click-one / cmd-click-to-tile toggle the strain components use (#58). - await expect(page.getByTestId('subwindow').filter({ hasText: 'Fit components' })) - .toBeVisible({ timeout: 20_000 }) + // + // Match the BREADCRUMB, anchored. Two maps windows exist by design — the + // caret's live one (refreshed as positions fit) and this snapshot — and + // `filter({hasText})` is a substring match over the whole subwindow, so + // the old locator matched both and failed strict mode. `$` is what + // separates "Fit components" from "Fit components (live)". + const committed = page.getByTestId('subwindow').filter({ + has: page.getByTestId('window-breadcrumb').filter({ hasText: /Fit components$/ }), + }) + await expect(committed).toBeVisible({ timeout: 20_000 }) + await expect(committed, 'the committed snapshot must be exactly one window') + .toHaveCount(1) + // And the live one is still there under its own name, so a user can tell + // the moving window from the kept one. + await expect(page.getByTestId('window-breadcrumb') + .filter({ hasText: 'Fit components (live)' })).toHaveCount(1) await page.screenshot({ path: `${SHOTS}/09-committed.png`, fullPage: true }) const errs = backend.logBuffer.filter((l: string) => /Traceback|CRITICAL/.test(l)) diff --git a/electron/tests/global-setup.cjs b/electron/tests/global-setup.cjs new file mode 100644 index 0000000..dc62f99 --- /dev/null +++ b/electron/tests/global-setup.cjs @@ -0,0 +1,36 @@ +/** + * global-setup.cjs — settings isolation for EVERY spec, not just harness users. + * + * Two things bite only on a machine that has never RUN SpyDE, i.e. every CI + * runner, and neither is visible on a dev box: + * + * 1. `tutorial_seen` is unset, so FirstRunGate auto-opens the welcome tour and + * its full-screen `tour-overlay` swallows pointer events. A hover or click + * on a titlebar then fails with "
from
+ * subtree intercepts pointer events" — which is how action_scoping and + * center_zero_beam failed on CI while passing locally. + * + * 2. `shell.openPath` shells out to xdg-open, which has no file manager to + * reach and leaves the app unable to exit (examples_menu's afterAll timed + * out for 120s on app.close()). + * + * _harness.cjs already handles both per launch, but only ~half the suite uses + * it — 30 specs call `_electron.launch()` directly. They all spread + * `...process.env`, so setting these here is what actually reaches all of them. + * A spec that wants the real behaviour (first_run.spec.ts wants a genuine first + * launch) passes its own SPYDE_SETTINGS_DIR, which still wins. + * + * SPYDE_SETTINGS_DIR redirects settings.json ONLY — never Electron's own + * profile, which Chromium refuses to launch without. + */ +const { mkdtempSync, mkdirSync, writeFileSync } = require('fs') +const { tmpdir } = require('os') +const { join } = require('path') + +module.exports = async () => { + const dir = join(mkdtempSync(join(tmpdir(), 'spyde-e2e-global-')), '.spyde') + mkdirSync(dir, { recursive: true }) + writeFileSync(join(dir, 'settings.json'), JSON.stringify({ tutorial_seen: true })) + process.env.SPYDE_SETTINGS_DIR = dir + process.env.SPYDE_NO_SHELL_OPEN = '1' +} diff --git a/spyde/actions/fit_action.py b/spyde/actions/fit_action.py index 58398be..1ff2ddf 100644 --- a/spyde/actions/fit_action.py +++ b/spyde/actions/fit_action.py @@ -946,6 +946,13 @@ def show_maps(self): one per fit stacked identical-looking windows with no way to tell which was current; and creating it only when a fit finished meant there was nothing to watch fill in. + + It is titled "(live)" to separate it from the snapshot commit() keeps. + Both used to be "Fit components", which is the same + no-way-to-tell-which-is-which problem one paragraph up — after a + commit the user had two identical titles side by side, and + fit_wizard.spec.ts could not target either (its locator matched both + and failed strict mode). """ maps = self.result_maps() if not maps: @@ -965,14 +972,19 @@ def show_maps(self): self.session._close_tree(tree) except Exception as e: log.debug("closing the previous maps window failed: %s", e) - self._maps_tree = self._open_maps(maps) + self._maps_tree = self._open_maps(maps, self.LIVE_MAPS_TITLE) self._maps_labels = labels return self._maps_tree - def _open_maps(self, maps): + #: The caret's live maps window, refreshed in place as positions fit. + LIVE_MAPS_TITLE = "Fit components (live)" + #: The snapshot commit() keeps. Plain name: it is the one that persists. + COMMITTED_MAPS_TITLE = "Fit components" + + def _open_maps(self, maps, title: str = COMMITTED_MAPS_TITLE): names = list(maps) return commit_result_tree( - self.session, title="Fit components", + self.session, title=title, primary=maps[names[0]], primary_label=names[0], views=[(n, maps[n]) for n in names[1:]], levels=None, cmap="viridis", From 3edf9aea5d1c49a811a3da24ed0331d8631509d3 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Wed, 29 Jul 2026 06:38:19 -0500 Subject: [PATCH 54/60] perf(tests): wait for the selector debounce instead of sleeping 0.8s per fixture MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The dataset fixtures loaded a signal and then slept a flat 0.8 seconds to "let selector debounce timers fire". They are used ~930 times across the suite — roughly 12 minutes of the run spent sleeping. The timer being waited on is armed for max(0.12, live_delay + 0.1) (base_selector._arm_settle), so 0.8s was about 6x the margin needed. Wait for the actual state instead: nothing pending on the nav dispatcher, no selector holding a live settle timer, and every signal plot painted. The floor matters as much as the condition. At the instant _add_signal returns, the timer may not be armed yet and no plot has painted, so a pure condition poll would return immediately and hand the test a session that had not settled at all — never return before one full settle period has passed. It falls through after a 3s timeout rather than hanging, so a fixture that genuinely never settles still fails loudly in its test. Measured: 213 ms for the 4D-STEM fixture and 157 ms for the 2D one, the 0.12s timer now being the floor. Same 25 files that took 78.6s take 59.6s, with the same 340 tests passing. No xdist — 4x box load on a suite this thread-sensitive traded flakes for speed, and three races surfaced this week already. --- spyde/tests/migrated/conftest.py | 45 +++++++++++++++++++++++++++++++- 1 file changed, 44 insertions(+), 1 deletion(-) diff --git a/spyde/tests/migrated/conftest.py b/spyde/tests/migrated/conftest.py index 5c7f20d..5bca18f 100644 --- a/spyde/tests/migrated/conftest.py +++ b/spyde/tests/migrated/conftest.py @@ -57,9 +57,52 @@ def _make_session(): return Session(n_workers=1, threads_per_worker=1) +def _settled(session) -> bool: + """Have the selector debounce/settle timers fired and their updates landed?""" + from spyde.drawing.selectors import base_selector as bs + + # A queued nav update is still to run (the dispatcher is latest-wins, so a + # non-empty pending map means work is outstanding). + if bs._nav_dispatcher._pending: + return False + # Any selector with a live settle timer is still going to re-fire. + for _wid, sel in getattr(session, "_nav_selectors", {}).items(): + for s in (sel, getattr(sel, "selector", None)): + if s is not None and getattr(s, "_settle_timer", None) is not None: + return False + # And every signal plot has actually painted something. + plots = [p for p in session._plots if not getattr(p, "is_navigator", False)] + return bool(plots) and all( + getattr(p, "current_data", None) is not None for p in plots) + + +def _settle(session, timeout: float = 3.0) -> None: + """Wait for the selector debounce instead of sleeping a flat 0.8s. + + This was `time.sleep(0.8)`, and the fixtures below are used ~930 times + across the suite — about 12 minutes of the run spent sleeping. The timer + it was waiting on is armed for `max(0.12, live_delay + 0.1)` + (base_selector._arm_settle), so 0.8s was ~6x the margin actually needed. + + The floor matters as much as the condition: at the instant _add_signal + returns, the timer may not be armed yet and no plot has painted, so a + pure condition poll could return immediately and hand the test a session + that has not settled at all. Never return before one full settle period + has passed, then wait for the real state. Falls through after `timeout` + rather than hanging, so a fixture can still fail loudly in its test. + """ + start = time.monotonic() + floor = start + 0.13 # one settle period; the timer cannot beat it + deadline = start + timeout + while time.monotonic() < deadline: + if time.monotonic() >= floor and _settled(session): + return + time.sleep(0.005) + + def _load(session, signal): session._add_signal(signal, source_path=None) - time.sleep(0.8) # let selector debounce timers fire + _settle(session) def _bright_disk_4d(nav, sig=(16, 16)): From 89ecec9758eb812cc8146dfb52397ff1378dad61 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Wed, 29 Jul 2026 06:51:09 -0500 Subject: [PATCH 55/60] fix(tests): settle the load before console-preview tests drive the selector MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit test_run_update_fires_hook_only_on_change failed on ubuntu-py3.10 with `assert 2 == 1`: it installs a NAV_CHANGE_HOOKS counter immediately after _add_signal, while the settle timer that load armed is still pending. The timer then fires on the dispatcher thread, commits a position and trips the hook — so the test's own single _run_update() looks like two. This file is the empirically hot one: the same race produced a flake under xdist, another on macos-py3.11 (test_no_selector_falls_back_to_frame_zero fighting a selector that kept re-arming), and now this. Every test in it drives the nav cursor, so settle after all seven loads rather than only the one that happened to fail. Costs ~1s across the file. The fixtures already do this; these tests call _add_signal themselves and so never got it. --- spyde/tests/migrated/test_console_preview.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/spyde/tests/migrated/test_console_preview.py b/spyde/tests/migrated/test_console_preview.py index 4f3a68a..050aa36 100644 --- a/spyde/tests/migrated/test_console_preview.py +++ b/spyde/tests/migrated/test_console_preview.py @@ -42,6 +42,7 @@ # initialized-module circular import). Warming it here removes the race (the # established testharness pattern: ensure_heavy_imports BEFORE set_signal_type). from spyde.backend.heavy_imports import ensure_heavy_imports +from spyde.tests.migrated.conftest import _settle ensure_heavy_imports() @@ -182,6 +183,7 @@ def test_lazy_comparison_previews_one_frame(self, window): sig = _lazy_4d() full_shape = sig.data.shape session._add_signal(sig, source_path=None) + _settle(session) # the load's own settle timer must fire BEFORE the test drives the selector _wait_vars(session, msgs, lambda c: any(v["name"] == "s1" for v in c["vars"])) computed = {"full": False} @@ -214,6 +216,7 @@ def test_nav_sum_is_too_expensive(self, window): # 1 nav position per chunk → 3*4 = 12 source chunks for a full-nav sum. sig = _lazy_4d(nav=(3, 4)) session._add_signal(sig, source_path=None) + _settle(session) # the load's own settle timer must fire BEFORE the test drives the selector _wait_vars(session, msgs, lambda c: any(v["name"] == "s1" for v in c["vars"])) with patch.object(da.Array, "compute", @@ -268,6 +271,7 @@ def test_image_thumbnail_matches_pipeline(self, window): _ = session.console sig = _lazy_4d() session._add_signal(sig, source_path=None) + _settle(session) # the load's own settle timer must fire BEFORE the test drives the selector _wait_vars(session, msgs, lambda c: any(v["name"] == "s1" for v in c["vars"])) res = _preview(session, msgs, "s1") assert res["kind"] == "image" @@ -338,6 +342,7 @@ def test_thumbnail_reflects_cursor(self, window): _ = session.console sig, arr = _stamped_4d() session._add_signal(sig, source_path=None) + _settle(session) # the load's own settle timer must fire BEFORE the test drives the selector _wait_vars(session, msgs, lambda c: any(v["name"] == "s1" for v in c["vars"])) tree = session.signal_trees[0] @@ -359,6 +364,7 @@ def test_no_selector_falls_back_to_frame_zero(self, window): _ = session.console sig, arr = _stamped_4d() session._add_signal(sig, source_path=None) + _settle(session) # the load's own settle timer must fire BEFORE the test drives the selector _wait_vars(session, msgs, lambda c: any(v["name"] == "s1" for v in c["vars"])) tree = session.signal_trees[0] # Clear any resolved indices so _nav_indices_for returns (None, None) — @@ -438,6 +444,7 @@ def test_nav_move_reemits_at_new_cursor(self, window): _ = session.console sig, arr = _stamped_4d() session._add_signal(sig, source_path=None) + _settle(session) # the load's own settle timer must fire BEFORE the test drives the selector _wait_vars(session, msgs, lambda c: any(v["name"] == "s1" for v in c["vars"])) tree = session.signal_trees[0] nav = TestNavPositionResolution() @@ -526,6 +533,7 @@ def test_run_update_fires_hook_only_on_change(self, window, monkeypatch): session, msgs = window["window"], window["messages"] sig = _lazy_4d() session._add_signal(sig, source_path=None) + _settle(session) # the load's own settle timer must fire BEFORE the test drives the selector tree = session.signal_trees[0] sel = TestNavPositionResolution()._signal_selector(session, tree) assert sel is not None From d61587c8bd7a8df462fe494f1d2eaf28ee7ca8b3 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Wed, 29 Jul 2026 07:11:57 -0500 Subject: [PATCH 56/60] fix(e2e): the two "flaky" fit specs were measuring a stale overlay MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both read the model curve straight after a navigator move. goTo ends in a flat waitForTimeout (2.5s / 1.2s) and repainting the overlay is a backend round trip, so on a loaded runner the read scores the PREVIOUS position's curve. fit_navigate saw it directly: peak 404.84 / argmax 486 at both stops, so "the overlaid model curve did not change between positions" fired on two reads of the same stale curve. It now polls until the curve changes, which is the property the test is named for — and is faster than the fixed sleep whenever the repaint lands early. fit_quality hid it behind a tolerance. Its sweep logged worst misfits of 47%, 76%, 83% and 191% across CI runs, while re-measuring that very position moments later gave 2-12% — so "worst" was really "where the read was most stale", and comparing that noise to a dedicated refit is what tripped `worstBefore < worstAfter * 1.5`. Loosening the ratio would have buried a real measurement bug. Instead every misfit now samples until two consecutive reads agree; unlike waiting for a CHANGE it assumes nothing about neighbouring positions differing. The measurement is what improved, not the threshold: the sweep's worst goes 47-191% -> 3.9%, median 3.6-3.8% -> 2.9%, and the worst position now reads 3.9% -> 3.9% after a dedicated refit — i.e. the whole-scan fit did leave it at its own answer, which is exactly what the assertion claims. --- electron/tests/fit_navigate.spec.ts | 45 +++++++++++++++++++---------- electron/tests/fit_quality.spec.ts | 38 ++++++++++++++++++++---- 2 files changed, 63 insertions(+), 20 deletions(-) diff --git a/electron/tests/fit_navigate.spec.ts b/electron/tests/fit_navigate.spec.ts index 073c5a9..fd50579 100644 --- a/electron/tests/fit_navigate.spec.ts +++ b/electron/tests/fit_navigate.spec.ts @@ -111,21 +111,28 @@ test('the overlaid model follows the navigator after a scan fit', async () => { await page.screenshot({ path: `${SHOTS}/02-at-6-6.png`, fullPage: true }) await goTo(26, 26) - const at2626 = await readModel() + // Poll rather than read once, for the same reason as the curve read at the + // bottom: goTo's flat 2.5s can be shorter than the caret's refresh on a + // loaded runner, and a stale read here fails as "the model is IDENTICAL". + // The assertion keeps its teeth — if the fields never diverge, the poll + // times out and reports it. + let at2626: Record = {} + await expect + .poll(async () => { + at2626 = await readModel() + return Object.keys(at66) + .filter((k) => Math.abs(at66[k] - at2626[k]) > 1e-9).length + }, { + timeout: 60_000, + message: 'the model stayed IDENTICAL at two different navigator ' + + 'positions — the caret is showing a stale fit, not this position\'s', + }) + .toBeGreaterThan(0) await page.screenshot({ path: `${SHOTS}/03-at-26-26.png`, fullPage: true }) console.log('model at (6,6) :', JSON.stringify(at66)) console.log('model at (26,26):', JSON.stringify(at2626)) - const changed = Object.keys(at66).filter( - (k) => Math.abs(at66[k] - at2626[k]) > 1e-9) - expect( - changed.length, - `the model is IDENTICAL at two different navigator positions ` + - `(${JSON.stringify(at66)}) — the caret is showing a stale fit, not ` + - `this position's`, - ).toBeGreaterThan(0) - // ── and each named component must be the SAME PEAK at both stops ───── // Two gaussians are exchangeable, so an unconstrained fit puts the broad // one in whichever slot each position happens to pick. Measured before @@ -165,11 +172,19 @@ test('the overlaid model follows the navigator after a scan fit', async () => { const curveHere = await modelCurve() expect(curveHere, 'no overlaid model curve on the signal plot').toBeTruthy() await goTo(6, 6) - const curveThere = await modelCurve() - expect( - JSON.stringify(curveThere), - 'the overlaid model curve did not change between positions', - ).not.toEqual(JSON.stringify(curveHere)) + // POLL, don't read once. goTo ends in a flat waitForTimeout(2_500), and + // repainting the overlay is a backend round trip — on a loaded CI runner + // it can outlast that, so a single read returns the PREVIOUS position's + // curve and the assertion below fires on two identical values. That is + // exactly how this went flaky (peak 404.84/argmax 486 at both stops). + // Waiting for the change is also faster than the fixed sleep in the + // common case, and it is the property the test is named for. + await expect + .poll(async () => JSON.stringify(await modelCurve()), { + timeout: 60_000, + message: 'the overlaid model curve did not change between positions', + }) + .not.toEqual(JSON.stringify(curveHere)) await assertNoJsErrors() } finally { diff --git a/electron/tests/fit_quality.spec.ts b/electron/tests/fit_quality.spec.ts index 20afe3a..55f9c3c 100644 --- a/electron/tests/fit_quality.spec.ts +++ b/electron/tests/fit_quality.spec.ts @@ -75,6 +75,34 @@ test('the fitted model matches the spectrum on screen', async () => { return rms / (span || 1) } + /** misfit(), but only once the overlay has stopped changing. + * + * Every goTo below ends in a flat waitForTimeout, and repainting the model + * overlay is a backend round trip that can outlast it on a loaded runner. + * A single read then scores the PREVIOUS position's curve against THIS + * position's data, which is not a small error: the sweep logged worst + * misfits of 47%, 76%, 83% and 191% across CI runs, while re-measuring + * that very position a moment later gave 2-12%. So "worst" was really + * "where the read was most stale", and the assertion that compares it to a + * dedicated refit went flaky on the noise. + * + * Sampling until two consecutive reads agree fixes the measurement rather + * than loosening the tolerance around it — and unlike waiting for the + * curve to CHANGE, it makes no assumption that neighbouring positions have + * visibly different fits. + */ + const settledMisfit = async (timeout = 30_000) => { + const deadline = Date.now() + timeout + let prev = misfit((await curves())!) + while (Date.now() < deadline) { + await page.waitForTimeout(150) + const now = misfit((await curves())!) + if (Math.abs(now - prev) < 1e-9) return now + prev = now + } + return prev + } + const sig = sigWindow(page) await sig.getByTestId('subwindow-titlebar').hover() await sig.getByTestId('action-btn-Fit').click() @@ -115,7 +143,7 @@ test('the fitted model matches the spectrum on screen', async () => { await page.waitForTimeout(2_000) await page.screenshot({ path: `${SHOTS}/01-after-fit-all.png`, fullPage: true }) - const afterAll = misfit((await curves())!) + const afterAll = await settledMisfit() console.log(`misfit right after Fit all Spectra: ${(afterAll * 100).toFixed(1)}% of range`) // ── sweep the navigator: ONE position proves nothing ───────────────── @@ -146,7 +174,7 @@ test('the fitted model matches the spectrum on screen', async () => { for (const cy of [2, 10, 16, 24, 30]) { for (const cx of [2, 10, 16, 24, 30]) { await goTo(cx, cy) - const m = misfit((await curves())!) + const m = await settledMisfit() swept.push(m) if (m > worst.m) worst = { m, at: [cx, cy] } if (swept.length <= 3) { @@ -163,12 +191,12 @@ test('the fitted model matches the spectrum on screen', async () => { await goTo(worst.at[0], worst.at[1]) await page.screenshot({ path: `${SHOTS}/03-worst-position.png`, fullPage: true }) - const worstBefore = misfit((await curves())!) + const worstBefore = await settledMisfit() await page.locator('[data-testid="fit-spectrum"]').click() await expect(page.locator('[data-testid="fit-status"]')) .toContainText(/chi2/i, { timeout: 60_000 }) await page.waitForTimeout(1_500) - const worstAfter = misfit((await curves())!) + const worstAfter = await settledMisfit() await page.screenshot({ path: `${SHOTS}/04-worst-refitted.png`, fullPage: true }) console.log(`worst position: ${(worstBefore * 100).toFixed(1)}% -> ` + `${(worstAfter * 100).toFixed(1)}% after Fit spectrum`) @@ -191,7 +219,7 @@ test('the fitted model matches the spectrum on screen', async () => { .toContainText(/chi2/i, { timeout: 60_000 }) await page.waitForTimeout(1_500) await page.screenshot({ path: `${SHOTS}/02-after-fit-spectrum.png`, fullPage: true }) - const afterOne = misfit((await curves())!) + const afterOne = await settledMisfit() console.log(`misfit after Fit spectrum here : ${(afterOne * 100).toFixed(1)}% of range`) expect( From e96ddc8fc05efa86f173e00bb671c3e9032be43a Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Wed, 29 Jul 2026 07:17:49 -0500 Subject: [PATCH 57/60] fix(tests): pin the widget's answer too, so the cursor cannot be reverted MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit test_thumbnail_reflects_cursor failed on windows-py3.11 with "cursor never settled at (1, 2)" — _settle_cursor burning its entire 30s budget. _set_indices assigned current_indices only. These tests never move a real widget, and _run_update does `self.current_indices = self.get_selected_indices()` on the dispatcher thread, so the committed value was reverted to the widget's own unmoved position every time a settle timer fired. _settle_cursor re-applied, got reverted, re-applied — a coin flip every 50ms that it can lose for 30s straight under load. Pin get_selected_indices as well as current_indices: the state is then self-consistent and nothing has anything to revert. That is the technique the NAV_CHANGE_HOOKS test in this same file already uses, and _run_update is safe with it (array_equal short-circuits, and the settle's force=True path just re-commits the same value). 36 passed three times running; 83 passed across the neighbouring selector suites. --- spyde/tests/migrated/test_console_preview.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/spyde/tests/migrated/test_console_preview.py b/spyde/tests/migrated/test_console_preview.py index 050aa36..09c2edd 100644 --- a/spyde/tests/migrated/test_console_preview.py +++ b/spyde/tests/migrated/test_console_preview.py @@ -306,11 +306,24 @@ def _signal_selector(self, session, tree): @staticmethod def _set_indices(sel, value): - """Set current_indices on *sel*, going through the inner active selector - for a composite (IntegratingSSelector2D exposes current_indices as a - read-only property that delegates to ``.selector``).""" + """Park *sel* at *value*, going through the inner active selector for a + composite (IntegratingSSelector2D exposes current_indices as a + read-only property that delegates to ``.selector``). + + Pin what the WIDGET reports as well as the committed value. These tests + never move a real widget, and ``_run_update`` does + ``self.current_indices = self.get_selected_indices()`` on the dispatcher + thread — so setting ``current_indices`` alone is reverted to the + widget's own (unmoved) position the moment a settle timer or any other + update fires. ``_settle_cursor`` then re-applies, gets reverted again, + and is really flipping a coin every 50 ms: on windows-py3.11 it lost for + the whole 30 s budget and failed as "cursor never settled at (1, 2)". + Pinning both makes the state self-consistent, so nothing has anything to + revert — the same technique the NAV_CHANGE_HOOKS test below already uses. + """ inner = getattr(sel, "selector", None) target = inner if inner is not None else sel + target.get_selected_indices = lambda: value target.current_indices = value def _move_all(self, session, tree, value): From cdea03bad52e2dbbc71268f6db17d7d9a365d3c3 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Wed, 29 Jul 2026 07:45:30 -0500 Subject: [PATCH 58/60] fix(e2e): confirm the panel MOVED before judging the model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit fit_quality still flaked once after the settled-read fix. The retry's own log gave it away: [2,2] — the first sweep stop after the tab switch — still read 191.6% while re-measuring that position later gave 3.7%. "Two consecutive reads agree" can agree on the OLD value when the move has not been applied yet, so the stability check settled on the stale curve. Wait for the SPECTRUM to become this position's before judging the model against it, and require three consecutive agreements rather than two so one repeat cannot land in a lull while the overlay catches up. The sweep is now deterministic: three runs in a row give median 2.9%, worst 3.9% at [30,10], worst position 3.9% -> 3.9% after a dedicated refit. Before, the worst position and value moved every run — 191.6% at [2,2], 82.6% at [30,2], 76.3% at [24,16], 46.9% at [30,10]. That variance WAS the staleness. Also faster: goTo waits for the real change instead of a flat 1.2s, so the spec runs in 1.0-1.1m rather than 1.3m. --- electron/tests/fit_quality.spec.ts | 29 +++++++++++++++++++++++++++-- 1 file changed, 27 insertions(+), 2 deletions(-) diff --git a/electron/tests/fit_quality.spec.ts b/electron/tests/fit_quality.spec.ts index 55f9c3c..a127942 100644 --- a/electron/tests/fit_quality.spec.ts +++ b/electron/tests/fit_quality.spec.ts @@ -94,11 +94,15 @@ test('the fitted model matches the spectrum on screen', async () => { const settledMisfit = async (timeout = 30_000) => { const deadline = Date.now() + timeout let prev = misfit((await curves())!) + // THREE consecutive agreements, not two: one repeat can land in a lull + // while the model overlay is still catching up with the data. + let stable = 0 while (Date.now() < deadline) { await page.waitForTimeout(150) const now = misfit((await curves())!) - if (Math.abs(now - prev) < 1e-9) return now + stable = Math.abs(now - prev) < 1e-9 ? stable + 1 : 0 prev = now + if (stable >= 2) return now } return prev } @@ -156,7 +160,17 @@ test('the fitted model matches the spectrum on screen', async () => { }, figIds.nav) expect(cross, 'the navigator has no crosshair').toBeTruthy() + /** Cheap signature of the DATA curve — changes when the panel actually + * moves to a different navigator position. */ + const dataSig = async () => { + const c = await curves() + if (!c) return '' + const d = c.data + return `${d.length}:${d[0]}:${d[d.length >> 1]}:${d[d.length - 1]}` + } + const goTo = async (cx: number, cy: number) => { + const before = await dataSig() await page.evaluate(({ f, panel, id, x, y }) => { window.postMessage({ type: 'awi_event', figId: f, @@ -166,7 +180,18 @@ test('the fitted model matches the spectrum on screen', async () => { }), }, '*') }, { f: figIds.nav, panel: cross.panel_id, id: cross.id, x: cx, y: cy }) - await page.waitForTimeout(1_200) + // Wait for the SPECTRUM to become this position's, not a flat 1.2s. + // settledMisfit alone was not enough: "two consecutive reads agree" can + // agree on the OLD value when the move has not been applied yet, which + // is why [2,2] — the first sweep stop after the tab switch — still read + // 191.6% while re-measuring it later gave 3.7%. Confirming the panel + // moved before judging the model is what closes that window. + const deadline = Date.now() + 30_000 + while (Date.now() < deadline) { + await page.waitForTimeout(100) + if ((await dataSig()) !== before) break + } + await page.waitForTimeout(400) // let the model overlay follow the data } let worst = { m: 0, at: [0, 0] as number[] } From b38ef3aad1ca9a44bac3eaac9d481660ee43c588 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Wed, 29 Jul 2026 08:11:11 -0500 Subject: [PATCH 59/60] fix(e2e): wait for the MODEL to follow, not just the spectrum MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Third pass at this one, each driven by what CI actually logged. A flat 1.2s: [2,2] read 191.6% where re-measuring gave 3.7%. Sampling until two reads agree: a stale overlay is perfectly stable, so stability cannot tell "finished updating" from "has not started" — [2,2] still read 191.6%. Waiting for the DATA to change: the spectrum arrives first and the model follows as a separate push, so on a loaded runner [16,2] still scored 28.3% mid-sweep and the run went median 24.8%. The MODEL changing is the signal that the overlay belongs to this position — the same thing that made fit_navigate stable. goTo now waits for both curves, and falls through on timeout rather than hanging, with settledMisfit still the final guard. Two runs back to back, identical: [2,2] 3.7%, median 2.9%, worst 3.9% at [30,10], worst position 3.9% -> 3.9%. And faster again — 52.5s against 1.3m for the original flat-sleep version, because it waits for the condition instead of guessing at it. --- electron/tests/fit_quality.spec.ts | 40 +++++++++++++++++++----------- 1 file changed, 25 insertions(+), 15 deletions(-) diff --git a/electron/tests/fit_quality.spec.ts b/electron/tests/fit_quality.spec.ts index a127942..9d23b15 100644 --- a/electron/tests/fit_quality.spec.ts +++ b/electron/tests/fit_quality.spec.ts @@ -160,17 +160,19 @@ test('the fitted model matches the spectrum on screen', async () => { }, figIds.nav) expect(cross, 'the navigator has no crosshair').toBeTruthy() - /** Cheap signature of the DATA curve — changes when the panel actually - * moves to a different navigator position. */ - const dataSig = async () => { + const sigOf = (a: number[] | null) => + a ? `${a.length}:${a[0]}:${a[a.length >> 1]}:${a[a.length - 1]}` : '' + + /** Signatures of BOTH curves. They update independently, which is the + * whole problem: the spectrum can be this position's while the model + * overlay is still the previous one's. */ + const sigs = async () => { const c = await curves() - if (!c) return '' - const d = c.data - return `${d.length}:${d[0]}:${d[d.length >> 1]}:${d[d.length - 1]}` + return { data: sigOf(c?.data ?? null), model: sigOf(c?.model ?? null) } } const goTo = async (cx: number, cy: number) => { - const before = await dataSig() + const before = await sigs() await page.evaluate(({ f, panel, id, x, y }) => { window.postMessage({ type: 'awi_event', figId: f, @@ -180,18 +182,26 @@ test('the fitted model matches the spectrum on screen', async () => { }), }, '*') }, { f: figIds.nav, panel: cross.panel_id, id: cross.id, x: cx, y: cy }) - // Wait for the SPECTRUM to become this position's, not a flat 1.2s. - // settledMisfit alone was not enough: "two consecutive reads agree" can - // agree on the OLD value when the move has not been applied yet, which - // is why [2,2] — the first sweep stop after the tab switch — still read - // 191.6% while re-measuring it later gave 3.7%. Confirming the panel - // moved before judging the model is what closes that window. + + // Wait for BOTH curves to become this position's. + // + // Neither weaker version held on CI. A flat 1.2s: [2,2] read 191.6% + // where re-measuring gave 3.7%. "Sample until two reads agree": a stale + // overlay is perfectly stable, so stability cannot tell "finished + // updating" from "has not started". Waiting only for the DATA: the + // spectrum arrives first and the model follows separately, so [16,2] + // still scored 28.3% mid-sweep on a loaded runner. + // + // The MODEL changing is the signal that the overlay is this position's, + // and it is what made fit_navigate stable. Falls through on timeout + // rather than hanging — two positions could in principle fit alike, and + // settledMisfit below is still the final guard. const deadline = Date.now() + 30_000 while (Date.now() < deadline) { await page.waitForTimeout(100) - if ((await dataSig()) !== before) break + const now = await sigs() + if (now.data !== before.data && now.model !== before.model) break } - await page.waitForTimeout(400) // let the model overlay follow the data } let worst = { m: 0, at: [0, 0] as number[] } From c7ae61eea6f4ac87b350dfdebd0d9f0c949c9b1b Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Wed, 29 Jul 2026 08:37:02 -0500 Subject: [PATCH 60/60] fix(e2e): confirm the navigator LANDED before reading the caret MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit fit_quality came back clean this run; fit_navigate took its place, and the message was mine: "the model stayed IDENTICAL at two different navigator positions", the 60s poll expiring because the caret never showed the new position. That is a moved-too-early failure, not a stale-fit one. goTo posted a pointer_up and slept a flat 2.5s; when that was not enough the navigator had not landed, so both reads described the same position and no amount of polling afterwards could separate them. Wait for the SPECTRUM to change — the same gate that fixed fit_quality's sweep. Two runs back to back give identical parameters at both stops, and it is quicker too: 28.8s / 31.6s against 38.9s for the flat-sleep version. Note the sweep's per-position log in fit_quality can still print a stale figure on a loaded runner. It no longer decides anything: the worst position is re-measured with settledMisfit before the assertion, which is why that run passed with worst position 3.8% -> 3.8% despite the sweep printing 82.6%. --- electron/tests/fit_navigate.spec.ts | 28 +++++++++++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/electron/tests/fit_navigate.spec.ts b/electron/tests/fit_navigate.spec.ts index fd50579..b6c99d9 100644 --- a/electron/tests/fit_navigate.spec.ts +++ b/electron/tests/fit_navigate.spec.ts @@ -77,7 +77,24 @@ test('the overlaid model follows the navigator after a scan fit', async () => { } const cross = await navCrosshair() + /** Signature of the SPECTRUM on the signal plot — changes when the + * navigator actually lands on a different position. */ + const dataSig = async () => page.evaluate((f) => { + const hook = (window as any)._spyde_test_panel_json + for (const raw of hook ? hook(f) : []) { + const d = JSON.parse(raw) + if (!d.data_b64) continue + const bin = atob(d.data_b64) + const bytes = new Uint8Array(bin.length) + for (let i = 0; i < bin.length; i++) bytes[i] = bin.charCodeAt(i) + const v = Array.from(new Float64Array(bytes.buffer)) + return `${v.length}:${v[0]}:${v[v.length >> 1]}:${v[v.length - 1]}` + } + return '' + }, figIds.sig) + const goTo = async (cx: number, cy: number) => { + const before = await dataSig() await page.evaluate(({ f, panel, id, x, y }) => { window.postMessage({ type: 'awi_event', figId: f, @@ -87,7 +104,16 @@ test('the overlaid model follows the navigator after a scan fit', async () => { }), }, '*') }, { f: figIds.nav, panel: cross.panel_id, id: cross.id, x: cx, y: cy }) - await page.waitForTimeout(2_500) + // Confirm the navigator LANDED before reading anything, instead of + // assuming 2.5s was enough. When it was not, the caret never showed the + // new position and the reads below compared a position against itself — + // "the model stayed IDENTICAL at two different navigator positions", + // which is a moved-too-early failure, not a stale-fit one. + const deadline = Date.now() + 60_000 + while (Date.now() < deadline) { + await page.waitForTimeout(100) + if ((await dataSig()) !== before) break + } } const readModel = async () => {