diff --git a/CHANGELOG.md b/CHANGELOG.md index 83bc5c4..a942ac6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,7 @@ This file is maintained using the shared changelog workflow in ### Added +- **Slice-by-Slice parallelism**: `File.fit_slice_by_slice()` accepts an `n_workers` keyword argument that dispatches per-slice fits across a `ProcessPoolExecutor` using the `spawn` start method (only portable option — Windows lacks `fork`). Default is `os.cpu_count() - 1`, capped at the number of slices. Set `n_workers=1` to keep the original serial path as a debug escape hatch. Workers reuse one pickled model installed at startup, render plots with the non-interactive Agg backend, and report progress via `tqdm`. On Linux/macOS spawn startup is a few hundred ms per worker; on Windows ~1-2s per worker, so very small fits (~< 20 slices) usually want `n_workers=1`. SbS seeding is now explicit too: `seed_source` chooses the shared template (`"model"`, `"baseline"`, or `"explicit"`), and `seed_adapt` controls the optional per-slice x0 tweak (`None` or `"argmax_shift"`). - `Model`, `Component`, and `Par` are now pickleable (and therefore deep-copyable) via `__getstate__` / `__setstate__`. This enables `copy.deepcopy(model)` and lets live models cross process boundaries, which unblocks future multiprocessing workflows and fixes latent MCMC parallelism (see `Fixed`). Pickled instances are for short-lived transfer, not persistence — parent back-references (`parent_file`, `parent_model`) and transient fit state (`const`, `args`) are nulled. ### Changed @@ -18,6 +19,10 @@ This file is maintained using the shared changelog workflow in - **Breaking (internal):** removed `Project.spec_lib` attribute and `Project.spec_fun` property. Fitting functions always live in `trspecfit.spectra`, so the indirection is hardcoded. Only affects code reaching into `project.spec_lib` / `project.spec_fun`; no public fit-workflow API changes. - `fitlib.residual_fun()` and `fitlib.plt_fit_res_1d()` no longer accept a `package` argument. The constant tuple passed to `fit_wrapper()` drops its third entry (`package`) and now has shape `(x, data, function_str, unpack, e_lim, t_lim)`. +### Removed + +- **Breaking:** `Project.skip_first_n_spec` and `Project.first_n_spec_only` debug controls (and the `first_n_spec_only` / `skip_first_n_spec` parameters on `fitlib.results_to_df()`, plus the now-unused `fitlib.results_select()` helper). With Slice-by-Slice parallelism a 200-slice fit takes seconds, so the "fit only the first N slices" debug shortcut no longer earns its keep. Users who want to fit a sub-range can slice the input array directly: `file.data = file.data[start:stop]; file.time = file.time[start:stop]`. + ### Fixed - **MCMC `workers > 1`**: `lmfit.emcee(workers=N)` via `ulmfit.MC(workers=N)` previously failed with `TypeError: cannot pickle 'module' object` because the residual closure carried a live module reference. The pickleable-model work plus the `spec_lib` removal close both sources of the error; MCMC parallel sampling now works end-to-end. diff --git a/TODO.md b/TODO.md index 269ff9f..9aa4669 100644 --- a/TODO.md +++ b/TODO.md @@ -9,9 +9,9 @@ Note: `fitlib.py` hardcodes `__lnsigma` value/min/max for MCMC sampling — make ## Performance & architecture -- [ ] **Slice-by-slice parallelism**: `n_workers` kwarg on `File.fit_slice_by_slice()`, `ProcessPoolExecutor` dispatch with tqdm progress, Agg backend in workers. Precondition (pickleable Model) shipped. - [ ] **Project-level fit backend**: `Project.fit_2d()` already supports `Project`/`File`/`Static` vary levels, but it currently evaluates through `fit_project_mcp()` and `Model.create_value_2d()` rather than the GIR scheduler/evaluator path. Decide whether to lower the multi-file residual to GIR or explicitly prefer project-managed per-file loops when we want maximum graph-IR speedups. - [ ] **JAX backend / Jacobian follow-on**: if we revisit a JAX evaluator, analytic Jacobians, or optimizer replacement, use [docs/design/jax-planning.md](docs/design/jax-planning.md) as the roadmap for scope, sequencing, and open technical constraints. +- [ ] **MCMC multiprocessing context**: `lmfit.emcee(workers=N)` currently inherits Python's default multiprocessing start method, which triggers a Python 3.12 `fork()` deprecation warning in multithreaded test runs. Investigate whether we can supply a `spawn`-backed worker pool or otherwise steer emcee/lmfit away from raw `fork`. - [ ] **Evaluation order correctness**: component eval order depends on coincidental list position; make it explicit. One option: build a directed acyclic graph (DAG) at model construction and topological-sort. - [ ] **Freeze non-varying pars**: pars without time-dependence (or profile dependence) are re-evaluated at every aux-axis point; could evaluate once and reuse. diff --git a/docs/design/roundtrip_test_matrix.md b/docs/design/roundtrip_test_matrix.md new file mode 100644 index 0000000..0567180 --- /dev/null +++ b/docs/design/roundtrip_test_matrix.md @@ -0,0 +1,239 @@ +# Roundtrip Test Matrix + +This document tracks the intended roundtrip-test surface for single-file fits. +It is meant to answer three questions at a glance: + +1. Which workflow API is under test? +2. Which backend expectation is under test? +3. Which supported model family is under test? + +Project-level fitting and MCMC/parallel execution are important too, but they +are tracked as secondary dimensions so the main matrix stays readable. + +## Axes + +### Workflow axis + +- `B`: `File.fit_baseline()` +- `Sp`: `File.fit_spectrum()` +- `SbS`: `File.fit_slice_by_slice()` +- `2D`: `File.fit_2d()` + +### Scope axis + +- `SF`: single-file workflows +- `P`: project-level workflows via `Project.fit_2d()` + +### Backend axis + +- `M`: force MCP with `project.spec_fun_str = "fit_model_mcp"` and recover truth +- `G`: run default compiled dispatch and recover truth on the GIR path +- `C`: assert `delta(MCP, GIR) = 0` + +For `C`, prefer `fit_model_compare` when the workflow supports it. Otherwise use +direct parity checks on residuals or evaluated outputs. + +### Cell notation + +- `M/G/C`: all three test types should exist for that workflow/model-family pair +- `-`: not applicable for that pair + +### Execution-mode qualifiers + +These are not a full extra matrix axis; they are required focused variants on +top of the main matrix: + +- `Opt`: normal optimizer-based fit, no MCMC +- `MC1`: MCMC enabled with `workers=1` +- `MC2`: MCMC enabled with `workers=2` +- `W1`: explicit serial execution for a workflow with worker support +- `W2`: parallel execution with `workers=2` + +Use `2` as the standard parallel test setting. We usually care about +`1` versus `>1`, not about many-worker scaling in CI. + +## Canonical model families + +Each row below should have one canonical fixture. Closely related variants can +be parameterized under the same row instead of getting separate rows. + +| ID | Model family | Representative fixture(s) | +| --- | --- | --- | +| `F1` | Plain energy model | `single_glp`, `glp_only` | +| `F2` | Static expressions in energy model: direct refs, fan-out, forward refs, static chains | `two_glp_expr_amplitude`, `expression_fan_out`, `energy_expression_forward_reference`, `expression_chain`, `glp_expression` | +| `F3` | Top-level standard dynamics | `single_glp` + `MonoExpPos` | +| `F4` | Top-level dynamics with IRF / convolution | `single_glp` + `MonoExpPosIRF` and other lowerable IRF kernels | +| `F5` | Top-level subcycle / multi-cycle dynamics | `single_glp` + `["ModelNone", "MonoExpNeg", "MonoExpPosExpr"]`, `frequency=10` | +| `F6` | Top-level profile only | `single_gauss` + `roundtrip_pLinear_x0` / `roundtrip_pExpDecay_A` | +| `F7` | Top-level profile plus separate top-level dynamics on another parameter | `single_gauss` + profile on `Gauss_01_A` + dynamics on `Gauss_01_x0` | +| `F8` | Profile-internal dynamics | `single_gauss` + profile on `Gauss_01_A` + dynamics on `Gauss_01_A_pExpDecay_01_A` | +| `F9` | Expression parameter referencing a top-level time-dependent base parameter | `two_glp_expr_amplitude` + dynamics on `GLP_01_A` | +| `F10` | Expression parameter referencing a top-level profiled base parameter | `two_glp_expr_amplitude` + profile on `GLP_01_A` | +| `F11` | Expression parameter referencing a profiled parameter whose profile internals are time-dependent | `two_glp_expr_amplitude` + profile on `GLP_01_A` + dynamics on `GLP_01_A_pExpDecay_01_A` | +| `F12` | Mixed expression referencing both profiled and time-dependent base parameters | `two_glp_mixed_profile_dynamics` with profile on `GLP_01_A` and dynamics on `GLP_01_x0` | + +## Target matrix + +Scope: `SF` (single-file) + +| Family | B | Sp | SbS | 2D | Notes | +| --- | --- | --- | --- | --- | --- | +| `F1` Plain energy | `M/G/C` | `M/G/C` | `M/G/C` | `-` | Core 1D workflow coverage | +| `F2` Static expressions | `M/G/C` | `M/G/C` | `M/G/C` | `-` | Include at least one direct-ref case and one fan-out or forward-ref case | +| `F3` Standard dynamics | `-` | `-` | `-` | `M/G/C` | Core 2D dynamic family | +| `F4` IRF dynamics | `-` | `-` | `-` | `M/G/C` | One canonical roundtrip plus parametrized parity across kernels | +| `F5` Subcycle dynamics | `-` | `-` | `-` | `M/G/C` | Important for multi-cycle indexing and expression prefixing | +| `F6` Profile only | `M/G/C` | `M/G/C` | `M/G/C` | `-` | Covers aux-axis plumbing in 1D APIs | +| `F7` Profile + separate dynamics | `-` | `-` | `-` | `M/G/C` | Top-level mixed feature case | +| `F8` Profile-internal dynamics | `-` | `-` | `-` | `M/G/C` | Single-cycle only | +| `F9` Expr -> time-dependent base par | `-` | `-` | `-` | `M/G/C` | High-value bug class for update ordering and pickling | +| `F10` Expr -> profiled base par | `M/G/C` | `M/G/C` | `M/G/C` | `-` | Expression namespace must see profiled values | +| `F11` Expr -> profiled base par with profile-internal dynamics | `-` | `-` | `-` | `M/G/C` | Single-cycle only | +| `F12` Mixed expr(profile + dynamics refs) | `-` | `-` | `-` | `M/G/C` | Stress case for combined namespace resolution | + +## Project-level matrix + +Scope: `P` (`Project.fit_2d()`) + +Project-level fitting should be tracked separately because it is currently +wired through `fit_project_mcp`, so the single-file GIR/MCP expectations do not +apply cleanly yet. + +| Family | 2D | Notes | +| --- | --- | --- | +| `PF1` Shared plain dynamics across files | `M` | Current core project roundtrip surface | +| `PF2` Project-level expressions | `M` | Includes file/project prefix rewriting and shared refs | +| `PF3` Shared dynamics with IRF | `M` | Add once project fixtures exist | +| `PF4` Shared subcycle dynamics | `M` | Add once project fixtures exist | + +Future: + +- if project-level GIR lands, upgrade applicable cells from `M` to `M/G/C` +- until then, do not force fake GIR coverage into the project matrix + +## MCMC and worker policy + +Yes: MCMC should be tracked. + +Yes: worker mode should be tracked anywhere the code can execute differently +between serial and parallel paths. + +But neither should multiply every cell in the main matrix. Instead use focused +requirements: + +### MCMC requirements + +MCMC is a second-layer contract on top of the clean optimizer roundtrips. + +Minimum MCMC set: + +- `MC1`: one canonical `B` test on `F1` +- `MC2`: one canonical `B` test on `F1` +- `MC2`: one expression-sensitive case, preferably `F9` or `F10` +- `MC2`: one 2D varying case, preferably `F3` or `F8` + +Rationale: + +- `MC1` checks that MCMC itself still works +- `MC2` checks pickling / process-boundary behavior +- expression-heavy and nested-model cases are the highest-value bug surfaces + +### SbS worker requirements + +Yes, `SbS` should eventually distinguish `W1` and `W2`, but only after +parallel SbS exists as a real API. + +Current status: + +- today `File.fit_slice_by_slice()` does not expose a worker-count API, so only + serial `SbS` roundtrips are testable + +Future requirement after `n_workers` lands: + +- `W1`: one canonical `SbS` roundtrip on `F1` +- `W2`: the same canonical `SbS` roundtrip on `F1` +- `W2`: one expression/profile-sensitive `SbS` case, likely `F2` or `F6` + +### Project worker requirements + +For project-level fits, add worker variants only when project execution gains a +parallel path that is semantically different from serial execution. + +## Practical rule + +Use this rule to decide whether a new dimension deserves explicit tracking: + +- add it as a full matrix axis only if it changes almost every cell +- otherwise add it as a focused secondary requirement + +By that rule: + +- project-level fits: yes, but separate matrix +- MCMC: yes, as focused secondary coverage +- `workers=1` vs `workers=2`: yes, but only for APIs that actually expose + worker-dependent behavior + +## Minimum test shape per cell + +For each required cell, the minimum useful test is: + +- simulate noiseless data from a truth model +- fit through the target workflow API +- assert recovered non-expression parameters match truth +- for `C`, assert MCP and GIR agree exactly or within a tight tolerance + +For MCMC-focused cells, the minimum useful test is: + +- run the fit with `mc_settings.use_mc=1` +- assert no crash for `MC1` +- assert no crash and no pickling/serialization failure for `MC2` +- when runtime allows, also assert basic parameter recovery or constraint + preservation + +Noisy roundtrip tests are still valuable, but should be a second layer. The +clean matrix above is the baseline contract. + +## Current snapshot + +This is the current high-level state of the suite, not a substitute for the +table above. + +- Covered reasonably well today: + - `F1` on `B`, `Sp`, `SbS`, and `2D` for GIR-path or compare-mode smoke + - `F3` on `2D` for GIR roundtrip and compare-mode + - `F4` on `2D` for parity / compare-mode + - `F5` on `2D` for parity / compare-mode + - `F6` on `B` for GIR roundtrip + - `F8` on `2D` for GIR roundtrip + - project-level `M` roundtrips for plain shared-dynamics fits + +- Thin or missing today: + - forced `M` roundtrip coverage for almost every family + - full workflow roundtrips for `F2`, `F5`, `F7`, `F9`, `F10`, `F11`, `F12` + - explicit `SbS` roundtrips outside the plain-energy family + - expression-heavy roundtrips through serialization-sensitive paths + - MCMC coverage beyond a simple plain-model smoke case + - any worker-mode matrix for `SbS` because parallel `SbS` does not exist yet + - project-level coverage for expression/subcycle/IRF families + +## Suggested implementation order + +If we fill this incrementally, the highest-value order is: + +1. Add forced-`M` twins for the existing plain and profile roundtrips. +2. Add `F9` and `F10` because expression + varying-parameter interactions are a known bug surface. +3. Add `MC2` coverage for one expression-heavy case and one nested-model case. +4. Add one canonical `F5` subcycle roundtrip through `fit_2d`. +5. Add one canonical `F4` IRF roundtrip through `fit_2d`. +6. Add `F7`, `F11`, and `F12` as the mixed-feature stress cases. +7. Expand the separate project-level matrix, starting with project expressions. + +## Non-goals + +This matrix does not try to track: + +- invalid / explicitly unsupported model combinations +- low-level evaluator unit tests +- plotting-only behavior + +Those should stay in their existing focused tests. diff --git a/pyproject.toml b/pyproject.toml index 444eb18..2d49347 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "trspecfit" -version = "0.8.1" +version = "0.8.2" authors = [ {name = "Johannes Mahl", email = "johannes.a.mahl@gmail.com"}, ] @@ -107,6 +107,7 @@ markers = ["slow: long-running round-trip tests (skipped by default, use -m slow filterwarnings = [ "ignore:FigureCanvasAgg is non-interactive:UserWarning", "ignore:invalid value encountered in scalar divide:RuntimeWarning:lmfit", + "ignore:This process .* use of fork\\(\\) may lead to deadlocks in the child\\.:DeprecationWarning", ] [tool.ruff] diff --git a/src/trspecfit/fitlib.py b/src/trspecfit/fitlib.py index 2913e16..d23415c 100644 --- a/src/trspecfit/fitlib.py +++ b/src/trspecfit/fitlib.py @@ -771,66 +771,12 @@ def fit_wrapper( # -# -def results_select(data: Any, skip: int = -1, n: int = -1, dim: int = 1) -> Any: - """ - Select slice of results array for partial fitting analysis. - - Parameters - ---------- - data : array-like - Data array to slice - skip : int, default=-1 - Number of initial elements to skip. -1 means skip none. - n : int, default=-1 - Number of elements to include. -1 means include all (after skip). - dim : int, default=1 - Dimensionality (currently only dim=1 supported). - - Returns - ------- - array-like - Sliced data - - Examples - -------- - >>> # Get all data - >>> results_select(data) - - >>> # Skip first 10 points - >>> results_select(data, skip=10) - - >>> # Get first 50 points only - >>> results_select(data, n=50) - - >>> # Get points 10-60 - >>> results_select(data, skip=10, n=60) - """ - - if dim == 1: - if n == -1: - if skip == -1: - return data # full data set - out = data[skip:] - else: # "+1" accounts for Python's exclusive upper bound - if skip == -1: - out = data[: n + 1] - else: - out = data[skip : n + 1] - else: - raise ValueError(f"Unsupported dim={dim}; only dim=1 is implemented") - - return out - - # def results_to_df( results: list[Any], x: ArrayLike | None = None, index: ArrayLike | None = None, config: PlotConfig | None = None, - first_n_spec_only: int = -1, - skip_first_n_spec: int = -1, save_df: int = 0, save_path: PathLike = "", ) -> pd.DataFrame: @@ -852,10 +798,6 @@ def results_to_df( Index values (e.g., slice numbers). If provided, included as column. config : PlotConfig, optional Plot configuration. If None, uses defaults. - first_n_spec_only : int, default=-1 - If != -1, include only first N results (for partial fitting). - skip_first_n_spec : int, default=-1 - If != -1, skip first N results (for partial fitting). save_df : {-1, 0, 1}, default=0 Save outputs: @@ -886,22 +828,22 @@ def results_to_df( # get columns names for plot before adding x/index cols_plt = df.columns - # select x (time) and index data if passed + # insert x (time) and index data if passed if x is not None: - x_save = results_select(data=x, skip=skip_first_n_spec, n=first_n_spec_only) - df.insert(0, config.y_label, x_save) # and insert into dataframe + df.insert(0, config.y_label, x) if index is not None: - ind_save = results_select( - data=index, skip=skip_first_n_spec, n=first_n_spec_only - ) - df.insert(0, "index", ind_save) # and insert into dataframe + df.insert(0, "index", index) # get par_fin([1]) of first slice(index=0) # (their "vary" attribute is the same for all) df_par_fin_slice0 = ulmfit.par_to_df( lmfit_params=results[0][1].params, col_type="min" ) - save_array = [-1 if not vary else 1 for vary in df_par_fin_slice0["vary"]] + if save_df < 0: + # Silent/API mode should not display parameter-evolution figures. + save_array = len(df_par_fin_slice0["vary"]) * [-1] + else: + save_array = [-1 if not vary else 1 for vary in df_par_fin_slice0["vary"]] if save_df != 0: # save the dataframe (index, x axis, parameter1, parameter2, ... @@ -909,7 +851,7 @@ def results_to_df( # plot individual parameters as a function of time (s) plt_fit_res_pars( df=df.loc[:, list(cols_plt)], - x=x_save if x is not None else None, + x=x, config=config, save_img=save_array, save_path=save_path, diff --git a/src/trspecfit/graph_ir.py b/src/trspecfit/graph_ir.py index 7db2cdb..804f590 100644 --- a/src/trspecfit/graph_ir.py +++ b/src/trspecfit/graph_ir.py @@ -1933,19 +1933,12 @@ def _bind_expr_to_rows( # -def _resolve_convolution_target_row( +def _walk_convolution_to_param_plus_trace( conv_node: GraphNode, edges: list[GraphEdge], id_to_node: dict[int, GraphNode], - name_to_row: dict[str, int], -) -> int: - """Walk a conv chain back to the underlying PARAM_PLUS_TRACE row. - - CONVOLUTION nodes wrap a resolved trace: their TRACE_INPUT source is - either the PARAM_PLUS_TRACE (single conv) or an earlier CONVOLUTION - (chained conv). The lowered plan rewrites the PPT row in place, so - every conv in a chain shares the same target row. - """ +) -> GraphNode: + """Walk a conv chain back to its underlying PARAM_PLUS_TRACE node.""" current = conv_node # Bounded walk: guards against pathological cycles in malformed graphs. @@ -1962,7 +1955,7 @@ def _resolve_convolution_target_row( ) parent = id_to_node[trace_parents[0].source] if parent.kind == NodeKind.PARAM_PLUS_TRACE: - return name_to_row[parent.name] + return parent if parent.kind != NodeKind.CONVOLUTION: raise ValueError( f"CONVOLUTION chain for {conv_node.name!r} walks through" @@ -1972,6 +1965,25 @@ def _resolve_convolution_target_row( raise ValueError(f"CONVOLUTION chain for {conv_node.name!r} exceeds graph size") +# +def _resolve_convolution_target_row( + conv_node: GraphNode, + edges: list[GraphEdge], + id_to_node: dict[int, GraphNode], + name_to_row: dict[str, int], +) -> int: + """Walk a conv chain back to the underlying PARAM_PLUS_TRACE row. + + CONVOLUTION nodes wrap a resolved trace: their TRACE_INPUT source is + either the PARAM_PLUS_TRACE (single conv) or an earlier CONVOLUTION + (chained conv). The lowered plan rewrites the PPT row in place, so + every conv in a chain shares the same target row. + """ + + ppt_node = _walk_convolution_to_param_plus_trace(conv_node, edges, id_to_node) + return name_to_row[ppt_node.name] + + # def _topological_sort(graph: GraphIR) -> list[int]: """Topological sort of graph node IDs. @@ -2526,11 +2538,19 @@ def schedule_2d(graph: GraphIR) -> ScheduledPlan2D: for edge in rep_param_edges[1:]: src_node = id_to_node[edge.source] src_row = name_to_row[src_node.name] - # PARAM_PLUS_TRACE nodes have a "_resolved" suffix; strip - # it for profile component parsing but use the resolved row. + # PARAM_PLUS_TRACE nodes have a "_resolved" suffix; CONVOLUTION + # nodes wrap a PPT for profile-time-dynamics and carry their own + # "__dynamics" name. In both cases, walk back to + # the underlying PPT so profile component parsing sees the + # original profile parameter name. parse_name = src_node.name if src_node.kind == NodeKind.PARAM_PLUS_TRACE: parse_name = parse_name.removesuffix("_resolved") + elif src_node.kind == NodeKind.CONVOLUTION: + ppt_node = _walk_convolution_to_param_plus_trace( + src_node, graph.edges, id_to_node + ) + parse_name = ppt_node.name.removesuffix("_resolved") comp_name, func_name = _parse_profile_component_param_name( group_name, parse_name, diff --git a/src/trspecfit/trspecfit.py b/src/trspecfit/trspecfit.py index c94e0a4..014b78f 100644 --- a/src/trspecfit/trspecfit.py +++ b/src/trspecfit/trspecfit.py @@ -51,7 +51,10 @@ See examples/ directory for complete workflows. """ +import concurrent.futures import copy +import multiprocessing +import os import pathlib import re import time @@ -70,6 +73,7 @@ import pandas as pd from IPython.display import display from ruamel.yaml import YAML +from tqdm import tqdm from trspecfit import fitlib, mcp @@ -80,7 +84,6 @@ from trspecfit.functions import energy as fcts_energy from trspecfit.functions import profile as fcts_profile from trspecfit.functions import time as fcts_time -from trspecfit.utils import arrays as uarr from trspecfit.utils import lmfit as ulmfit from trspecfit.utils import parsing as uparsing from trspecfit.utils import plot as uplt @@ -94,6 +97,198 @@ # "conv" functions in individual subcycles are currently ignored +# --- Slice-by-Slice multiprocessing workers ---------------------------------- +# These globals are populated only inside ProcessPoolExecutor worker +# processes by ``_sbs_worker_init``. They let workers reuse a single +# pickled Model and dispatch_args across all slices they process, +# avoiding per-task pickle overhead. + +_WORKER_MODEL: "mcp.Model | None" = None +_WORKER_DISPATCH_ARGS: tuple[Any, ...] | None = None +_WORKER_SEED_TEMPLATE: list[float] | None = None + + +# +def _extract_sbs_seed_template( + seed_values: Any, parameter_names: Sequence[str] +) -> list[float]: + """Normalize explicit SbS seed values to a full ordered parameter list.""" + + if isinstance(seed_values, dict): + missing = [name for name in parameter_names if name not in seed_values] + extra = [name for name in seed_values if name not in parameter_names] + if missing or extra: + raise ValueError( + "Explicit SbS seed dict must define exactly the model parameters.\n" + f"Missing: {missing or 'none'}\n" + f"Extra: {extra or 'none'}" + ) + out: list[float] = [] + for name in parameter_names: + value = seed_values[name] + if isinstance(value, (list, tuple, np.ndarray)): + if len(value) == 0: + raise ValueError( + f"Explicit SbS seed for parameter '{name}' is empty." + ) + value = value[0] + out.append(float(value)) + return out + + out = ulmfit.par_extract(seed_values, return_type="list") + if len(out) != len(parameter_names): + raise ValueError( + "Explicit SbS seed must provide one value per model parameter.\n" + f"Expected {len(parameter_names)} values, got {len(out)}." + ) + return out + + +# +def _prepare_sbs_model_for_slice( + model: "mcp.Model", + dispatch_args: tuple[Any, ...], + seed_template: list[float], + *, + seed_adapt: Literal["argmax_shift"] | None, + s: np.ndarray, + energy: np.ndarray, + e_lim: list[int] | None, + e_pos_pars: list[str], + e_pos_vals: pd.Series | None, + data_base_argmax_energy: float | None, + fit_fun_str: str, +) -> list[float]: + """Reset the shared SbS model to the seed template for one slice.""" + + model.update_value(new_par_values=seed_template, par_select="all") + + if seed_adapt == "argmax_shift": + if e_pos_vals is None or data_base_argmax_energy is None: + raise ValueError( + "argmax_shift seed adaptation requires baseline-derived x0 values." + ) + delta_max = energy[np.argmax(s)] - data_base_argmax_energy + new_e_vals = list(e_pos_vals.add(delta_max)) + model.update_value(new_par_values=new_e_vals, par_select=e_pos_pars) + + initial_guess = ulmfit.par_extract(model.lmfit_pars, return_type="list") + model.const = (energy, s, fit_fun_str, 0, e_lim, []) + model.args = dispatch_args + return initial_guess + + +# +def _sbs_worker_init( + model: "mcp.Model", + dispatch_args: tuple[Any, ...], + seed_template: list[float], +) -> None: + """Executor initializer: install per-worker model and Agg backend. + + Runs once per worker process before any task. Stashes the deep-pickled + model and GIR/MCP dispatch args as worker-local globals so individual + slice tasks don't have to pay the pickle cost on every submission. + Forces matplotlib to the non-interactive Agg backend so the per-slice + plot calls inside workers don't try to open a display. + """ + + global _WORKER_MODEL, _WORKER_DISPATCH_ARGS, _WORKER_SEED_TEMPLATE + import matplotlib + + matplotlib.use("Agg", force=True) + _WORKER_MODEL = model + _WORKER_DISPATCH_ARGS = dispatch_args + _WORKER_SEED_TEMPLATE = seed_template + + +# +def _sbs_fit_one_slice( + s_i: int, + s: np.ndarray, + *, + energy: np.ndarray, + e_lim: list[int] | None, + seed_adapt: Literal["argmax_shift"] | None, + e_pos_pars: list[str], + e_pos_vals: pd.Series | None, + data_base_argmax_energy: float | None, + fit_fun_str: str, + stages: int, + path_slice: pathlib.Path, + plot_config: PlotConfig, + fit_wrapper_kwargs: dict[str, Any], +) -> tuple[int, list[Any]]: + """Fit one energy slice in a worker process. + + Uses worker-local ``_WORKER_MODEL`` and ``_WORKER_DISPATCH_ARGS`` + installed by :func:`_sbs_worker_init`. Mutates the worker's model + state (x0 params, ``const``, ``args``) — this is safe because tasks + within a single worker run sequentially and each slice overwrites + the relevant fields. + + Returns + ------- + tuple[int, list] + (slice_index, fit_wrapper_result) so the caller can reassemble + out-of-order completions back into slice order. + """ + + assert _WORKER_MODEL is not None, "_sbs_worker_init must run first" + assert _WORKER_DISPATCH_ARGS is not None + assert _WORKER_SEED_TEMPLATE is not None + model = _WORKER_MODEL + dispatch_args = _WORKER_DISPATCH_ARGS + seed_template = _WORKER_SEED_TEMPLATE + + initial_guess = _prepare_sbs_model_for_slice( + model, + dispatch_args, + seed_template, + seed_adapt=seed_adapt, + s=s, + energy=energy, + e_lim=e_lim, + e_pos_pars=e_pos_pars, + e_pos_vals=e_pos_vals, + data_base_argmax_energy=data_base_argmax_energy, + fit_fun_str=fit_fun_str, + ) + const = model.const + args = model.args + assert const is not None + assert args is not None + + result_sbs = fitlib.fit_wrapper( + const=const, + args=args, + par_names=model.parameter_names, + par=model.lmfit_pars, + stages=stages, + show_output=0, + save_output=1, + save_path=path_slice, + **fit_wrapper_kwargs, + ) + + fitlib.plt_fit_res_1d( + x=const[0], + y=const[1], + fit_fun_str=fit_fun_str, + par_init=initial_guess, + par_fin=result_sbs[1], + args=args, + plot_sum=False, + show_init=True, + fit_lim=e_lim, + config=plot_config, + save_img=-1, + save_path=path_slice.with_suffix(".png"), + ) + + return s_i, result_sbs + + # # class Project: @@ -138,8 +333,6 @@ class Project: spec_fun_str : str Name of fitting function in ``trspecfit.spectra`` (e.g. ``'fit_model_gir'``, ``'fit_model_mcp'``, ``'fit_model_compare'``) - skip_first_n_spec, first_n_spec_only : int - Slice selection for partial fitting (-1 = all slices) Notes ----- @@ -226,8 +419,6 @@ def _set_defaults(self) -> None: self.da_slices_fmt = "%06d" # Advanced settings self.spec_fun_str = "fit_model_gir" - self.skip_first_n_spec = -1 - self.first_n_spec_only = -1 # def __repr__(self) -> str: @@ -2270,13 +2461,23 @@ def load_fit(self) -> None: # def fit_slice_by_slice( - self, model_name: str, stages: int = 1, **fit_wrapper_kwargs + self, + model_name: str, + stages: int = 1, + n_workers: int | None = None, + *, + seed_source: Literal["model", "baseline", "explicit"] = "baseline", + seed_values: Any | None = None, + seed_adapt: Literal["argmax_shift"] | None = "argmax_shift", + **fit_wrapper_kwargs, ) -> None: """ Fit time- and energy-resolved spectrum Slice-by-Slice (SbS). - Treats every time step as independent from other times. Requires fitting - the baseline first using fit_baseline(). + Treats every time step as independent from other times. Each slice starts + from a shared seed template selected via ``seed_source``, optionally + adapted per slice via ``seed_adapt``. There is no cross-slice warm start, + which keeps the workflow embarrassingly parallel. Parameters ---------- @@ -2288,15 +2489,53 @@ def fit_slice_by_slice( - 1: Single optimization with ``fit_alg_1`` - 2: Two-stage fit (``fit_alg_1`` then ``fit_alg_2``) + n_workers : int or None, default=None + Number of parallel worker processes. Slices are independent + so SbS is embarrassingly parallel. + + - ``None``: auto, ``os.cpu_count() - 1`` (leaves one core free) + - ``1``: serial path (debug escape hatch — no pool overhead) + - ``N > 1``: ``ProcessPoolExecutor`` with the ``spawn`` start + method, the only portable option (Windows lacks ``fork``). + Workers reuse one pickled model installed at startup, so + per-slice pickle overhead is bounded. + + Notes for very small fits: spawn worker startup costs + ~200-500ms on Linux/macOS and ~1-2s on Windows per worker, + so for fits of fewer than ~20 slices, ``n_workers=1`` is + usually faster. + + seed_source : {'model', 'baseline', 'explicit'}, default='baseline' + Shared parameter template used to seed every slice before any + per-slice adaptation is applied. + + - ``'model'``: use the current ``model_sbs.lmfit_pars`` values + - ``'baseline'``: use the stored ``fit_baseline()`` result values + - ``'explicit'``: use ``seed_values`` after normalizing it to the + model parameter order + + seed_values : optional + Explicit seed template used when ``seed_source='explicit'``. + Accepts the same broad value shapes as ``ulmfit.par_extract()`` + plus dicts keyed by parameter name. + + seed_adapt : {None, 'argmax_shift'}, default='argmax_shift' + Optional per-slice tweak applied after resolving the shared seed + template and before fitting a slice. + + - ``None``: use the same initial guess for every slice + - ``'argmax_shift'``: shift all parameters ending in ``'_x0'`` by the + difference between the current slice's argmax energy and the + baseline spectrum's argmax energy + **fit_wrapper_kwargs Additional keyword arguments passed to fitlib.fit_wrapper Notes ----- - Note: - Currently the energy position guesses (x0) are shifted on a per slice basis - according to the position in energy (x) of the maximum value of the spectrum - (NOT always a good idea!) + The actual per-slice fit results live in ``self.results_sbs``. + ``self.model_sbs`` is retained as shared model/reconstruction context + and is restored to the unadapted seed template after fitting. """ t_sbs = time.time() # start timing for SbS fit @@ -2308,18 +2547,27 @@ def fit_slice_by_slice( "be used for Slice-by-Slice fitting. " "Use a model without dynamics, or use fit_2d() instead." ) - if self.model_base is None: - raise ValueError( - "Baseline model is not fitted yet; run fit_baseline() first." - ) - if ( - self.data is None - or self.time is None - or self.energy is None - or self.data_base is None + if self.data is None or self.time is None or self.energy is None: + raise ValueError("Data/axes missing; cannot run Slice-by-Slice fit.") + if seed_source not in ("model", "baseline", "explicit"): + raise ValueError("seed_source must be 'model', 'baseline', or 'explicit'.") + if seed_adapt not in (None, "argmax_shift"): + raise ValueError("seed_adapt must be None or 'argmax_shift'.") + if seed_source == "baseline" and ( + self.model_base is None or not self.model_base.result ): raise ValueError( - "Data/axes/baseline missing; cannot run Slice-by-Slice fit." + "Baseline seed requested but baseline model is not fitted yet; " + "run fit_baseline() first or use seed_source='model'/'explicit'." + ) + if seed_source != "explicit" and seed_values is not None: + raise ValueError("seed_values is only used when seed_source='explicit'.") + if seed_source == "explicit" and seed_values is None: + raise ValueError("seed_source='explicit' requires seed_values.") + if seed_adapt == "argmax_shift" and self.data_base is None: + raise ValueError( + "seed_adapt='argmax_shift' requires baseline data; " + "run define_baseline() first or use seed_adapt=None." ) # define (and create) path where SbS fit results will be saved to @@ -2330,99 +2578,170 @@ def fit_slice_by_slice( ], ) - # set all fixed SbS fit parameters equal to baseline model results - base_df = ulmfit.par_to_df(self.model_base.lmfit_pars, col_type="min") - self.model_sbs.update_value( - new_par_values=list(base_df["value"]), par_select="all" - ) + if seed_source == "model": + seed_template = ulmfit.par_extract( + self.model_sbs.lmfit_pars, return_type="list" + ) + elif seed_source == "baseline": + assert self.model_base is not None # type guard + seed_template = ulmfit.par_extract( + self.model_base.result[1], return_type="list" + ) + else: + seed_template = _extract_sbs_seed_template( + seed_values, + self.model_sbs.parameter_names, + ) + + self.model_sbs.update_value(new_par_values=seed_template, par_select="all") # find all parameters with names ending in "x0" # so they can be updated for every slice e_pos_pars = [ name for name in self.model_sbs.parameter_names if name.endswith("_x0") ] - # find their corresponding values - e_pos_vals = uarr.get_item( - base_df, row=["name", e_pos_pars], col="value", astype="series" - ) + e_pos_vals: pd.Series | None + data_base_argmax_energy: float | None + if seed_adapt == "argmax_shift": + e_pos_vals = pd.Series( + data=[self.model_sbs.lmfit_pars[name].value for name in e_pos_pars], + index=e_pos_pars, + dtype=float, + ) + assert self.data_base is not None # type guard + data_base_argmax_energy = float(self.energy[np.argmax(self.data_base)]) + else: + e_pos_vals = None + data_base_argmax_energy = None # --- dispatch: GIR fast path vs interpreter --- _fun_str = self.p.spec_fun_str _args_sbs = self._build_1d_dispatch_args(self.model_sbs, _fun_str) - # cycle through all spectra and fit them - self.results_sbs = [] # (re-)initialize placeholder for results - for s_i, s in enumerate(self.data): - print(f"Analyzing slice number {s_i + 1}/{len(self.time)}", end="\r") - if s_i < self.p.skip_first_n_spec: - continue # skip past baseline spectra for debugging - # define path for files saved for this slice - path_slice = path_sbs_results / "slices" / str(self.p.da_slices_fmt % s_i) - - # update the "x0" peak energy guess(es) using - # "max(baseline) -(max current slice)" [ in eV] - delta_max = ( - self.energy[np.argmax(s)] - self.energy[np.argmax(self.data_base)] - ) - # update all guesses for parameters with names ending in "x0" - new_e_vals = list(e_pos_vals.add(delta_max)) - self.model_sbs.update_value( - new_par_values=new_e_vals, par_select=e_pos_pars - ) - # get initial guess - initial_guess = ulmfit.par_extract( - self.model_sbs.lmfit_pars, return_type="list" - ) - - # const = (x, data, fnctn str, unpack, energy limits, time limits) + n_slices = len(self.data) + + def _slice_path(s_i: int) -> pathlib.Path: + return path_sbs_results / "slices" / str(self.p.da_slices_fmt % s_i) + + # resolve worker count: None -> auto, otherwise honour user. + if n_workers is None: + n_workers = max(1, (os.cpu_count() or 1) - 1) + # No point spawning more workers than slices. + n_workers = max(1, min(n_workers, n_slices)) + + if n_workers == 1: + # serial path (debug escape hatch). + self.results_sbs = [] + for s_i, s in enumerate(self.data): + print(f"Analyzing slice number {s_i + 1}/{len(self.time)}", end="\r") + path_slice = _slice_path(s_i) + + initial_guess = _prepare_sbs_model_for_slice( + self.model_sbs, + _args_sbs, + seed_template, + seed_adapt=seed_adapt, + s=s, + energy=self.energy, + e_lim=self.e_lim, + e_pos_pars=e_pos_pars, + e_pos_vals=e_pos_vals, + data_base_argmax_energy=data_base_argmax_energy, + fit_fun_str=_fun_str, + ) + const = self.model_sbs.const + args = self.model_sbs.args + assert const is not None + assert args is not None + + result_sbs = fitlib.fit_wrapper( + const=const, + args=args, + par_names=self.model_sbs.parameter_names, + par=self.model_sbs.lmfit_pars, + stages=stages, + show_output=0, + save_output=1, + save_path=path_slice, + **fit_wrapper_kwargs, + ) + self.results_sbs.append(result_sbs) + + fitlib.plt_fit_res_1d( + x=const[0], + y=const[1], + fit_fun_str=self.p.spec_fun_str, + par_init=initial_guess, + par_fin=result_sbs[1], + args=args, + plot_sum=False, + show_init=True, + fit_lim=self.e_lim, + config=self.plot_config, + save_img=-1, + save_path=path_slice.with_suffix(".png"), + ) + else: + # parallel path: spawn pool, install model once per worker. + ctx = multiprocessing.get_context("spawn") + by_id: dict[int, list[Any]] = {} + with concurrent.futures.ProcessPoolExecutor( + max_workers=n_workers, + mp_context=ctx, + initializer=_sbs_worker_init, + initargs=(self.model_sbs, _args_sbs, seed_template), + ) as executor: + futures = { + executor.submit( + _sbs_fit_one_slice, + s_i, + self.data[s_i], + energy=self.energy, + e_lim=self.e_lim, + seed_adapt=seed_adapt, + e_pos_pars=e_pos_pars, + e_pos_vals=e_pos_vals, + data_base_argmax_energy=data_base_argmax_energy, + fit_fun_str=_fun_str, + stages=stages, + path_slice=_slice_path(s_i), + plot_config=self.plot_config, + fit_wrapper_kwargs=fit_wrapper_kwargs, + ): s_i + for s_i in range(n_slices) + } + try: + for fut in tqdm( + concurrent.futures.as_completed(futures), + total=len(futures), + desc=f"SbS fit ({n_workers} workers)", + ): + slice_idx, result_sbs = fut.result() + by_id[slice_idx] = result_sbs + except BaseException: + # fail-fast: cancel remaining futures and re-raise + for f in futures: + f.cancel() + raise + self.results_sbs = [by_id[i] for i in sorted(by_id)] + # mirror the serial path's final model state so downstream + # consumers (save_sbs_fit, plot helpers) see identical + # const/args regardless of which path produced the results. self.model_sbs.const = ( self.energy, - s, + self.data[n_slices - 1], _fun_str, 0, self.e_lim, [], ) - # args [for fit function called in residual function] self.model_sbs.args = _args_sbs - # fit with confidence intervals - result_sbs = fitlib.fit_wrapper( - const=self.model_sbs.const, - args=self.model_sbs.args, - par_names=self.model_sbs.parameter_names, - par=self.model_sbs.lmfit_pars, - stages=stages, - show_output=0, - save_output=1, - save_path=path_slice, - **fit_wrapper_kwargs, - ) - - # add final fit parameters to list of fit parameters of all spectra - self.results_sbs.append(result_sbs) - - # (optionally) plot and (always) save fit summary for this slice - fitlib.plt_fit_res_1d( - x=self.model_sbs.const[0], - y=self.model_sbs.const[1], - fit_fun_str=self.p.spec_fun_str, - par_init=initial_guess, - par_fin=result_sbs[1], - args=self.model_sbs.args, - plot_sum=False, - show_init=True, - fit_lim=self.e_lim, - config=self.plot_config, - save_img=-1, - save_path=path_slice.with_suffix(".png"), - ) - # - if s_i == self.p.first_n_spec_only: - break # for debugging: only fit first N spectra - if stages >= 1: self.save_sbs_fit(save_path=path_sbs_results) + self.model_sbs.update_value(new_par_values=seed_template, par_select="all") + self.model_sbs.args = _args_sbs + if stages >= 1: fitlib.time_display( t_start=t_sbs, print_str="Time elapsed for Slice-by-Slice fit: " ) @@ -2458,8 +2777,6 @@ def save_sbs_fit(self, save_path: PathLike) -> None: x=self.time, index=np.arange(0, len(self.time)), config=self.plot_config, - skip_first_n_spec=self.p.skip_first_n_spec, - first_n_spec_only=self.p.first_n_spec_only, save_df=-1 if self.p.show_output == 0 else 1, save_path=save_path, ) @@ -2475,19 +2792,17 @@ def save_sbs_fit(self, save_path: PathLike) -> None: ) # plot data, fit, and residual 2D maps - # (works if full 2D map is fitted/ no slices skipped) - if self.p.first_n_spec_only == -1 and self.p.skip_first_n_spec == -1: - fitlib.plt_fit_res_2d( - data=self.data, - fit=fit_2d_sbs, - x=self.energy, - y=self.time, - config=self.plot_config, - x_lim=self.e_lim, - y_lim=self.t_lim, - save_img=-1 if self.p.show_output == 0 else 1, - save_path=save_path, - ) + fitlib.plt_fit_res_2d( + data=self.data, + fit=fit_2d_sbs, + x=self.energy, + y=self.time, + config=self.plot_config, + x_lim=self.e_lim, + y_lim=self.t_lim, + save_img=-1 if self.p.show_output == 0 else 1, + save_path=save_path, + ) # def _resolve_model(self, model_name: str | None) -> mcp.Model: diff --git a/tests/test_file.py b/tests/test_file.py index 71fa24f..cba2ef7 100644 --- a/tests/test_file.py +++ b/tests/test_file.py @@ -934,33 +934,85 @@ def test_fit_baseline_no_data_base_raises(self): # -- fit_slice_by_slice -- # - def test_fit_sbs_no_baseline_model_raises(self): - """fit_slice_by_slice raises ValueError when baseline not fitted.""" + def test_fit_sbs_default_seed_requires_fitted_baseline(self): + """Default SbS seeding requires a completed baseline fit.""" file = self._make_file_with_model() - file.model_base = None - with pytest.raises(ValueError, match="fit_baseline"): + file.model_base = file.model_active + with pytest.raises(ValueError, match="Baseline seed requested"): file.fit_slice_by_slice("simple_energy") + # + def test_fit_sbs_model_seed_allows_no_baseline_fit(self): + """seed_source='model' can run without a baseline fit or baseline data.""" + + file = self._make_file_with_model() + file.p.spec_fun_str = "fit_model_mcp" + + with ( + unittest.mock.patch( + "trspecfit.trspecfit.fitlib.fit_wrapper", + return_value=[None, object(), None, None, None], + ) as mock_fit, + unittest.mock.patch("trspecfit.trspecfit.fitlib.plt_fit_res_1d"), + unittest.mock.patch.object(file, "save_sbs_fit"), + unittest.mock.patch("trspecfit.trspecfit.fitlib.time_display"), + ): + file.fit_slice_by_slice( + "simple_energy", + n_workers=1, + seed_source="model", + seed_adapt=None, + ) + + assert mock_fit.call_count == len(file.time) + + # + def test_fit_sbs_explicit_seed_requires_values(self): + """seed_source='explicit' must be accompanied by seed_values.""" + + file = self._make_file_with_model() + with pytest.raises(ValueError, match="requires seed_values"): + file.fit_slice_by_slice( + "simple_energy", + seed_source="explicit", + seed_adapt=None, + ) + + # + def test_fit_sbs_nonexplicit_seed_rejects_seed_values(self): + """seed_values should not be accepted for non-explicit seed sources.""" + + file = self._make_file_with_model() + with pytest.raises(ValueError, match="only used when seed_source='explicit'"): + file.fit_slice_by_slice( + "simple_energy", + seed_source="model", + seed_values=[1.0, 2.0], + seed_adapt=None, + ) + # def test_fit_sbs_no_data_raises(self): """fit_slice_by_slice raises ValueError when data is missing.""" file = self._make_file_with_model() - file.model_base = file.model_active # satisfy baseline check file.data = None with pytest.raises(ValueError, match="missing"): - file.fit_slice_by_slice("simple_energy") + file.fit_slice_by_slice( + "simple_energy", seed_source="model", seed_adapt=None + ) # def test_fit_sbs_no_time_raises(self): """fit_slice_by_slice raises ValueError when time axis is missing.""" file = self._make_file_with_model() - file.model_base = file.model_active file.time = None with pytest.raises(ValueError, match="missing"): - file.fit_slice_by_slice("simple_energy") + file.fit_slice_by_slice( + "simple_energy", seed_source="model", seed_adapt=None + ) # -- fit_2d -- diff --git a/tests/test_gir_integration.py b/tests/test_gir_integration.py index 3b9e258..d472407 100644 --- a/tests/test_gir_integration.py +++ b/tests/test_gir_integration.py @@ -21,6 +21,7 @@ schedule_1d, schedule_2d, ) +from trspecfit.utils import lmfit as ulmfit _ENERGY_YAML = "models/eval_2d_energy.yaml" _TIME_YAML = "models/file_time.yaml" @@ -1172,15 +1173,19 @@ class TestFileFitSliceBySlice: # @pytest.mark.slow - def test_compare_mode_through_fit_slice_by_slice(self): - """fit_slice_by_slice uses the 1D GIR path when the model lowers.""" + @pytest.mark.parametrize("n_workers", [1, 2]) + def test_compare_mode_through_fit_slice_by_slice(self, n_workers): + """fit_slice_by_slice uses the 1D GIR path when the model lowers. + + Run under both the serial (n_workers=1) and parallel + (n_workers=2) dispatch paths to ensure GIR/MCP parity holds in + both, and that the parallel path leaves the same downstream + model state (const, args) the serial path does. + """ - project = Project(path="tests", name="gir_sbs_cmp") + project = Project(path="tests", name=f"gir_sbs_cmp_w{n_workers}") project.show_output = 0 project.spec_fun_str = "fit_model_compare" - # Fit 3 slices so we exercise the hoisted-args reuse across the loop - # and the multi-row save_sbs_fit reconstruction path. - project.first_n_spec_only = 2 truth = _make_1d_truth_file(project) truth.model_active.create_value_1d() @@ -1191,10 +1196,132 @@ def test_compare_mode_through_fit_slice_by_slice(self): fit_file.fit_baseline(model_name="single_glp", stages=1, try_ci=0) # If GIR and interpreter disagree, fit_model_compare raises. - fit_file.fit_slice_by_slice(model_name="single_glp", stages=1, try_ci=0) + fit_file.fit_slice_by_slice( + model_name="single_glp", stages=1, try_ci=0, n_workers=n_workers + ) assert fit_file.model_sbs is not None assert fit_file.model_sbs.args is not None assert len(fit_file.model_sbs.args) == 4 assert isinstance(fit_file.model_sbs.args[0], ScheduledPlan1D) - assert len(fit_file.results_sbs) == 3 + assert len(fit_file.results_sbs) == len(truth.time) + + # + @pytest.mark.slow + def test_serial_and_parallel_produce_same_fit(self): + """n_workers=1 and n_workers=2 must converge to identical params. + + Slices are independent, so each path optimizes the same residual + on the same data with the same initial guess. Any divergence + would indicate a state-leak bug in worker dispatch. + """ + + def _run_sbs(name: str, n_workers: int) -> tuple[list, int]: + project = Project(path="tests", name=name) + project.show_output = 0 + project.spec_fun_str = "fit_model_mcp" + + truth = _make_1d_truth_file(project) + truth.model_active.create_value_1d() + spectrum_1d = truth.model_active.value_1d.copy() + data_2d = np.tile(spectrum_1d, (len(truth.time), 1)) + + fit_file = _make_1d_fit_file(project, data_2d, truth.energy, truth.time) + fit_file.fit_baseline(model_name="single_glp", stages=1, try_ci=0) + fit_file.fit_slice_by_slice( + model_name="single_glp", + stages=1, + try_ci=0, + n_workers=n_workers, + ) + assert fit_file.results_sbs is not None # type guard + return fit_file.results_sbs, len(truth.time) + + serial, n_serial = _run_sbs("sbs_serial", n_workers=1) + parallel, n_parallel = _run_sbs("sbs_parallel", n_workers=2) + + assert len(serial) == len(parallel) == n_serial == n_parallel + for s_i, (r_serial, r_parallel) in enumerate( + zip(serial, parallel, strict=True) + ): + # result_sbs[1] is the lmfit MinimizerResult; compare its + # final parameter values across the two paths. + params_serial = r_serial[1].params + params_parallel = r_parallel[1].params + assert list(params_serial.keys()) == list(params_parallel.keys()) + for name in params_serial.keys(): + np.testing.assert_allclose( + params_serial[name].value, + params_parallel[name].value, + rtol=1e-10, + atol=1e-12, + err_msg=f"slice {s_i} param {name} diverged", + ) + + # + @pytest.mark.slow + @pytest.mark.parametrize("n_workers", [1, 2]) + @pytest.mark.parametrize( + ("seed_source", "seed_adapt"), + [ + ("baseline", "argmax_shift"), + ("model", None), + ("explicit", None), + ], + ) + def test_fit_slice_by_slice_restores_seed_template( + self, + n_workers, + seed_source, + seed_adapt, + ): + """SbS leaves model_sbs at the shared seed template in both paths.""" + + project = Project( + path="tests", name=f"sbs_seed_restore_{seed_source}_w{n_workers}" + ) + project.show_output = 0 + project.spec_fun_str = "fit_model_mcp" + + truth = _make_1d_truth_file(project) + truth.model_active.create_value_1d() + spectrum_1d = truth.model_active.value_1d.copy() + data_2d = np.tile(spectrum_1d, (len(truth.time), 1)) + + fit_file = _make_1d_fit_file(project, data_2d, truth.energy, truth.time) + fit_file.fit_baseline(model_name="single_glp", stages=1, try_ci=0) + + assert fit_file.model_base is not None + expected = ulmfit.par_extract(fit_file.model_base.result[1], return_type="list") + seed_values = None + + if seed_source == "model": + expected = expected.copy() + expected[0] = expected[0] + 1.234 if expected[0] == 0 else expected[0] * 0.8 + fit_file.model_active.update_value(expected) + elif seed_source == "explicit": + expected = expected.copy() + expected[0] = expected[0] + 0.456 if expected[0] == 0 else expected[0] * 1.1 + seed_values = expected + + fit_file.fit_slice_by_slice( + model_name="single_glp", + stages=1, + try_ci=0, + n_workers=n_workers, + seed_source=seed_source, + seed_values=seed_values, + seed_adapt=seed_adapt, + ) + + assert fit_file.model_sbs is not None + for name, expected_value in zip( + fit_file.model_sbs.parameter_names, expected, strict=True + ): + np.testing.assert_allclose( + fit_file.model_sbs.lmfit_pars[name].value, + expected_value, + rtol=1e-12, + atol=1e-12, + err_msg=(f"{seed_source=} {n_workers=} left {name} in the wrong state"), + ) diff --git a/tests/test_graph_ir.py b/tests/test_graph_ir.py index 34cdda5..8dcafe3 100644 --- a/tests/test_graph_ir.py +++ b/tests/test_graph_ir.py @@ -1238,7 +1238,7 @@ def test_non_profiled_expression_refs_base(self): # Time-dependent profile parameter tests # -def _make_time_dep_profile_model(): +def _make_time_dep_profile_model(dynamics_model=None): """Create model with a profiled par whose profile slope is time-dependent. single_glp with pLinear(m=-0.5, b=0) on GLP_01_A, @@ -1264,7 +1264,7 @@ def _make_time_dep_profile_model(): target_model="single_glp", target_parameter="GLP_01_A_pLinear_01_m", dynamics_yaml="models/file_time.yaml", - dynamics_model=["MonoExpPos"], + dynamics_model=dynamics_model or ["MonoExpPos"], ) model = file.model_active assert model is not None @@ -1346,6 +1346,18 @@ def test_resolved_profile_par_in_sample_edges(self): break assert found_resolved_in_sample + # + def test_schedule_2d_accepts_profile_par_dynamics_convolution(self): + """Profile-param IRF dynamics lower through schedule_2d.""" + + _file, model = _make_time_dep_profile_model(["MonoExpPosIRF"]) + graph = build_graph(model) + + assert can_lower_2d(graph) + plan = schedule_2d(graph) + + assert plan.n_profile_samples > 0 + # Dynamics convolution tests #