From b5f4382124980053e3f34d212d7bc60e66d8b616 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Sun, 26 Jul 2026 20:51:41 -0500 Subject: [PATCH 01/33] =?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 00000000..4d5c5a17 --- /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/33] 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 56b5963b..7b79f070 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 bde015c0..3c0ba052 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 00000000..dd63abcd --- /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 00000000..69d5e542 --- /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 00000000..e7b38abc --- /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/33] 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 ac6453f2..669f3852 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 4d5c5a17..699457d2 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 e2e1476d..ff0e705e 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 fad6d3da..7047f6fa 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 ac0c7252..4c3a80f9 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 edda5d10..7d6b5ffe 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 00000000..e49a9006 --- /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/33] 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 00000000..280381c6 --- /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 00000000..263eb6d3 --- /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 00000000..cf25a159 --- /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 00000000..aba17219 --- /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 00000000..723c40ba --- /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 00000000..bdc95afb --- /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 00000000..6c3e2a24 --- /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 00000000..76c9d1ed --- /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/33] 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 263eb6d3..a7d392ed 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 cf25a159..6a52925e 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 00000000..f523e609 --- /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 723c40ba..f6142e75 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 00000000..cbf5af98 --- /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 76c9d1ed..33056257 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 a3e5ae5baaa5b3920778f2df510a90f7e8a8ae42 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Sun, 26 Jul 2026 23:45:20 -0500 Subject: [PATCH 06/33] 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 a7d392ed..578447b6 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 00000000..8c80db3e --- /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 5d6d396ab53db1273bb0780895121bb276520267 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Sun, 26 Jul 2026 23:09:46 -0500 Subject: [PATCH 07/33] 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 578447b6..76d68c4b 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 33056257..8175a06d 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 08/33] 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 00000000..69e9dd12 --- /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 00000000..f8050af2 --- /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 00000000..6aeaefdb --- /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 09/33] 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 76d68c4b..1c5a53f6 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 aba17219..7456384b 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 00000000..8ff8bfa9 --- /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 10/33] 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 1c5a53f6..777aabcd 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 7456384b..316c5abd 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 69e9dd12..7502b353 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 00000000..e7a17993 --- /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 00000000..d6ce9241 --- /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 11/33] 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 7502b353..9cab2358 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 00000000..31a6012e --- /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 00000000..29eae591 --- /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 955a2f79092add640423cb31fdaa06be76021cd0 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Mon, 27 Jul 2026 08:09:47 -0500 Subject: [PATCH 12/33] =?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 00000000..9e544062 --- /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 acac6d04..3fdc2acb 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 && ( + <> +