diff --git a/CHANGELOG.md b/CHANGELOG.md index ef5f481..9cac71e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,7 @@ This file is maintained using the shared changelog workflow in ### Changed +- Slice-by-Slice multiprocessing plumbing (worker globals, `_sbs_worker_init`, `_sbs_fit_one_slice`, and the seed-handling helpers) moved out of `trspecfit.py` into `trspecfit.utils.sbs`. `trspecfit.py` now opens directly on `class Project` instead of ~190 lines of worker plumbing. No public API changes. - **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)`. diff --git a/TODO.md b/TODO.md index 12b6b56..de95b28 100644 --- a/TODO.md +++ b/TODO.md @@ -17,6 +17,17 @@ Note: `fitlib.py` hardcodes `__lnsigma` value/min/max for MCMC sampling — make - [ ] **Pylance/Pyright `| None` noise in tests**: ~280 Pyright errors across tests, all from accessing `File`/`Model` attributes typed as `ndarray | None`. Current `# type guard` asserts are inconsistent and don't propagate through helper methods. Find a cleaner pattern (e.g. `TypeGuard`, narrowing wrapper, or Pyright config) and apply consistently. +## User and AI ergonomics + +- [ ] **Curate the public API surface before v1.0.0**: audit user-facing classes (`File`, `Project`, `Simulator`, `Model`, etc.) and decide which methods/attributes should be discoverable in notebooks and docs. Add curated `__dir__()` output for autocomplete, keep `__all__`/API docs aligned, and gradually rename or deprecate internal helper methods that should not look like primary user workflows. This should improve both human notebook ergonomics and AI/LLM efficiency by making the intended workflow surface smaller, clearer, and easier to infer. +- [ ] **Document API tiers**: add a short guide that separates stable user API (`Project`, `File`, `Simulator`, `PlotConfig`), advanced public API (`mcp.Model`, `Component`, `Par`, `ParameterSweep`, `MC`), and internal implementation modules (`graph_ir`, `eval_1d`, `eval_2d`, low-level parsing/HDF5 helpers). Use this as the source of truth for docs, tests, examples, and AI-agent guidance. +- [ ] **Add tool-neutral agent orientation**: add `AGENTS.md` or `docs/ai/agent-orientation.md` pointing agents to `CLAUDE.md`, `TODO.md`, `PLAN.md`, `docs/design/repo_architecture.md`, supported-model docs, common commands, and API-change guardrails. Keep it concise so any LLM can quickly find the intended workflow and repo boundaries. +- [ ] **Add more AI-friendly task recipes**: extend `docs/ai/` with checklists for common repo changes, such as adding YAML syntax, adding plotting options, changing fitting workflows, modifying GIR/evaluator behavior, extending save/load fields, and preparing a release. +- [ ] **Add minimal runnable workflow examples**: supplement notebooks with small script-like examples or docs snippets for the canonical public workflows: load data, load a model, set limits, fit baseline, fit 2D, inspect results, simulate data, and run a parameter sweep. +- [ ] **Improve public validation errors**: make user-facing errors state what failed, where it failed (file/model/component/parameter when applicable), and what the user or agent should change next. Prioritize YAML parsing, model loading, fit setup, and unsupported-model fallback paths. +- [ ] **Tighten public type hints and aliases**: reduce ambiguous `Any` on public APIs, document key aliases such as `ModelRef`, and keep return types crisp for IDEs, Pyright, generated docs, and LLM code navigation. +- [ ] **Document fast verification slices**: add focused pytest commands for common edits (public workflow/API, YAML parser, functions, GIR/evaluator, plotting) so contributors and agents can validate changes quickly before running the full suite. + ## Build & release - [ ] **Automate tagging and pushing**: automate `git tag v1.2.3` + `git push v1.2.3` as part of the release workflow. diff --git a/docs/design/roundtrip_test_matrix.md b/docs/design/roundtrip_test_matrix.md index 528fbc3..a20d1ed 100644 --- a/docs/design/roundtrip_test_matrix.md +++ b/docs/design/roundtrip_test_matrix.md @@ -146,11 +146,12 @@ serial path and `n_workers>1` crosses a process boundary. Current status: - the main `SbS` matrix runs with `n_workers=1` -- one focused `W2` test covers `F1` with `n_workers=2` +- focused `W2` tests cover `F1` and profile-bearing `F6` with `n_workers=2` Future requirement if worker-specific risk grows: -- `W2`: one expression/profile-sensitive `SbS` case, likely `F2` or `F6` +- add a more expression-heavy `W2` `SbS` case, likely `F2`, if process-boundary + risk shows up beyond the existing plain/profile cases ### Project worker requirements @@ -201,13 +202,13 @@ table above. - `F2` variants for direct, fan-out, and forward-reference expressions - noisy second-layer checks for `F3`, `F6`, and `F8` on the GIR path - focused MCMC checks for `MC1`, `MC2`, expression-sensitive `MC2`, and 2D `MC2` - - focused `W2` coverage for `fit_slice_by_slice()` + - focused `W2` coverage for plain and profile-bearing `fit_slice_by_slice()` - project-level `M` roundtrips for `PF1`, `PF2`, and `PF3` - Thin or missing today: - project-level `PF4` shared subcycle dynamics - project-level `G/C` coverage, because project fitting is still MCP-only - - expression/profile-sensitive `W2` coverage for `fit_slice_by_slice()` + - expression-heavy `W2` coverage for `fit_slice_by_slice()` - MCMC assertions beyond no-crash / process-boundary coverage - exhaustive noisy coverage, intentionally kept out of the main matrix @@ -216,8 +217,8 @@ table above. The original single-file matrix is implemented. Highest-value next steps: 1. Add `PF4` once a shared project-subcycle fixture exists. -2. Add a focused expression/profile-sensitive `W2` `SbS` test if process-boundary - risk shows up beyond the plain `F1` case. +2. Add a focused expression-heavy `W2` `SbS` test if process-boundary risk shows + up beyond the existing plain/profile cases. 3. Add lightweight recovery or constraint-preservation assertions to focused MCMC tests when runtime allows. 4. Upgrade project-level cells from `M` to `M/G/C` if project-level GIR lands. diff --git a/src/trspecfit/trspecfit.py b/src/trspecfit/trspecfit.py index 014b78f..a3e92ad 100644 --- a/src/trspecfit/trspecfit.py +++ b/src/trspecfit/trspecfit.py @@ -87,6 +87,7 @@ from trspecfit.utils import lmfit as ulmfit from trspecfit.utils import parsing as uparsing from trspecfit.utils import plot as uplt +from trspecfit.utils import sbs as usbs PathLike = str | pathlib.Path ModelRef = str | int | list[str] @@ -97,198 +98,6 @@ # "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: @@ -2588,7 +2397,7 @@ def fit_slice_by_slice( self.model_base.result[1], return_type="list" ) else: - seed_template = _extract_sbs_seed_template( + seed_template = usbs.extract_sbs_seed_template( seed_values, self.model_sbs.parameter_names, ) @@ -2636,7 +2445,7 @@ def _slice_path(s_i: int) -> pathlib.Path: 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( + initial_guess = usbs.prepare_sbs_model_for_slice( self.model_sbs, _args_sbs, seed_template, @@ -2688,12 +2497,12 @@ def _slice_path(s_i: int) -> pathlib.Path: with concurrent.futures.ProcessPoolExecutor( max_workers=n_workers, mp_context=ctx, - initializer=_sbs_worker_init, + initializer=usbs.sbs_worker_init, initargs=(self.model_sbs, _args_sbs, seed_template), ) as executor: futures = { executor.submit( - _sbs_fit_one_slice, + usbs.sbs_fit_one_slice, s_i, self.data[s_i], energy=self.energy, diff --git a/src/trspecfit/utils/sbs.py b/src/trspecfit/utils/sbs.py new file mode 100644 index 0000000..268375f --- /dev/null +++ b/src/trspecfit/utils/sbs.py @@ -0,0 +1,216 @@ +""" +Slice-by-Slice (SbS) multiprocessing helpers. + +Worker-process plumbing for ``File.fit_slice_by_slice()``. Workers are +spawned by ``ProcessPoolExecutor``; ``sbs_worker_init`` runs once per +worker to install a shared model and dispatch args as worker-local +globals, and ``sbs_fit_one_slice`` consumes them to fit a single slice. + +The module also exposes the seed-handling helpers used in the serial +path (``extract_sbs_seed_template``, ``prepare_sbs_model_for_slice``). +""" + +from __future__ import annotations + +import pathlib +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any, Literal + +import numpy as np +import pandas as pd + +from trspecfit import fitlib +from trspecfit.config.plot import PlotConfig +from trspecfit.utils import lmfit as ulmfit + +if TYPE_CHECKING: + from trspecfit import mcp + +# 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