Sbs cleanup

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)`.
 

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.

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.

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,