diff --git a/CHANGELOG.md b/CHANGELOG.md index 688f1d9..83bc5c4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,21 @@ and this project adheres to [Semantic Versioning](https://semver.org/). This file is maintained using the shared changelog workflow in [`docs/ai/changelog.md`](docs/ai/changelog.md). +## [Unreleased] + +### Added + +- `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 + +- **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)`. + +### 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. + ## [0.8.0] - 2026-04-20 ### Added diff --git a/TODO.md b/TODO.md index 7bd0f72..269ff9f 100644 --- a/TODO.md +++ b/TODO.md @@ -9,6 +9,7 @@ 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. - [ ] **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. diff --git a/pyproject.toml b/pyproject.toml index 2b0282a..444eb18 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "trspecfit" -version = "0.8.0" +version = "0.8.1" authors = [ {name = "Johannes Mahl", email = "johannes.a.mahl@gmail.com"}, ] diff --git a/src/trspecfit/fitlib.py b/src/trspecfit/fitlib.py index 1eb2e19..2913e16 100644 --- a/src/trspecfit/fitlib.py +++ b/src/trspecfit/fitlib.py @@ -35,6 +35,7 @@ from lmfit.minimizer import MinimizerResult from numpy.typing import ArrayLike +from trspecfit import spectra from trspecfit.config.plot import PlotConfig from trspecfit.utils import lmfit as ulmfit from trspecfit.utils import plot as uplt @@ -67,7 +68,6 @@ def residual_fun( par: Any, x: ArrayLike, data: np.ndarray, - package: Any, fit_fun_str: str, unpack: int = 0, e_lim: list[int] | None = None, @@ -97,10 +97,9 @@ def residual_fun( - 1D: [n_energy] for energy-resolved fits - 2D: [n_time, n_energy] for time- and energy-resolved fits - package : module - Python module containing the fit function (typically trspecfit.spectra) fit_fun_str : str - Name of fit function within package (e.g., 'fit_model_mcp') + Name of fit function in ``trspecfit.spectra`` + (e.g., ``'fit_model_mcp'``, ``'fit_model_gir'``) unpack : {0, 1}, default=0 Parameter passing mode: @@ -145,7 +144,7 @@ def residual_fun( args = () # define the fit function - fit_fun = getattr(package, fit_fun_str) + fit_fun = getattr(spectra, fit_fun_str) # if the minimizer calling this is from the lmfit package, then # extract the value from their lmfit.Parameter() (dictionary) @@ -390,7 +389,7 @@ def fit_wrapper( ---------- const : tuple Constants for residual_fun: - (x, data, package, function_str, unpack, e_lim, t_lim) + (x, data, function_str, unpack, e_lim, t_lim) args : tuple Arguments for fit function (passed to residual_fun): Typically (model, dim) for MCP models @@ -947,7 +946,7 @@ def results_to_fit_2d( const : tuple Constants for residual_fun: - (x, data, package, function_str, unpack, e_lim, t_lim) + (x, data, function_str, unpack, e_lim, t_lim) Used to evaluate fit function at each time point. args : tuple Arguments for fit function (model, dim). @@ -977,7 +976,6 @@ def results_to_fit_2d( ( x_const, data_const, - package_const, fit_fun_const, unpack_const, e_lim_const, @@ -992,7 +990,6 @@ def results_to_fit_2d( results[i][1].params, x_const, np.asarray(data_const), - package_const, fit_fun_const, unpack=cast("int", unpack_const), e_lim=cast("list[int]", e_lim_const), @@ -1008,7 +1005,6 @@ def results_to_fit_2d( results.iloc[i].values, x_const, np.asarray(data_const), - package_const, fit_fun_const, unpack=cast("int", unpack_const), e_lim=cast("list[int]", e_lim_const), @@ -1037,7 +1033,6 @@ def plt_fit_res_1d( x: ArrayLike, y: ArrayLike, fit_fun_str: str, - package: Any, par_init: Any, par_fin: Any, args: tuple[Any, ...] | None = None, @@ -1064,9 +1059,8 @@ def plt_fit_res_1d( y : array Y-axis data (spectrum to be fitted) fit_fun_str : str - Name of fitting function in package (e.g., 'fit_model_mcp') - package : module - Python module containing fit_fun_str (typically trspecfit.spectra) + Name of fitting function in ``trspecfit.spectra`` + (e.g., ``'fit_model_mcp'``, ``'fit_model_gir'``) par_init : list or lmfit.Parameters Initial parameter guess. Can be empty list [] if show_init=False. par_fin : lmfit.MinimizerResult or lmfit.Parameters or list @@ -1133,7 +1127,7 @@ def plt_fit_res_1d( save_path = kwargs.get("save_path", "") # Get fit function - fit_fun = getattr(package, fit_fun_str) + fit_fun = getattr(spectra, fit_fun_str) x_arr = np.asarray(x, dtype=float) y_arr = np.asarray(y, dtype=float) diff --git a/src/trspecfit/mcp.py b/src/trspecfit/mcp.py index 2d00e2c..e7a017f 100644 --- a/src/trspecfit/mcp.py +++ b/src/trspecfit/mcp.py @@ -216,6 +216,25 @@ def __repr__(self) -> str: cls = type(self).__name__ return f"{cls}('{self.name}', {n_comp} comp, {n_par} pars, dim={dim})" + # + def __getstate__(self) -> dict[str, Any]: + """Pickle protocol: strip back-refs and transient state. + + Pickled Models are for short-lived process-boundary transfer + (multiprocessing, ``copy.deepcopy``), **not** for persistence. + ``parent_file`` (and, on Dynamics/Profile subclasses, + ``parent_model``) is nulled to keep pickles bounded in size — + otherwise the whole ``Project`` graph would be dragged in. + Transient fit state (``const``, ``args``) is stripped because it + is meant to be re-populated by the caller before the next fit. + """ + + state = self.__dict__.copy() + for key in ("parent_file", "parent_model", "const", "args"): + if key in state: + state[key] = None + return state + @property def plot_config(self) -> PlotConfig: """ @@ -1238,6 +1257,30 @@ def __repr__(self) -> str: n = len(self.pars) return f"Component('{self.comp_name}', type='{self.comp_type}', {n} pars)" + # + def __getstate__(self) -> dict[str, Any]: + """Pickle protocol: strip module and back-refs. + + ``package`` is a live module reference (e.g. + ``trspecfit.functions.energy``), which is not picklable. Replace + it with the dotted module name; ``__setstate__`` restores the + module via ``importlib.import_module``. ``parent_model`` is + nulled — pickled Components are detached from their parent Model. + """ + + state = self.__dict__.copy() + state["package"] = self.package.__name__ + state["parent_model"] = None + return state + + # + def __setstate__(self, state: dict[str, Any]) -> None: + """Pickle protocol: restore the ``package`` module from its name.""" + + package_name = state.pop("package") + self.__dict__.update(state) + self.package = importlib.import_module(package_name) + # [automatic] create self.fct attribute that will update if either # self.package or self.fct_str changes [attribute is read only] @property @@ -1973,6 +2016,20 @@ def __repr__(self) -> str: extra = f" [{', '.join(flags)}]" if flags else "" return f"Par('{self.name}'{extra})" + # + def __getstate__(self) -> dict[str, Any]: + """Pickle protocol: null the ``parent_model`` back-reference. + + Pickled ``Par`` instances are detached from their parent Model; + anything that relied on ``parent_model`` (e.g. expression + resolution that needs the full parameter set) must re-attach + after unpickling. + """ + + state = self.__dict__.copy() + state["parent_model"] = None + return state + # def describe(self, detail: int = 0) -> None: """ diff --git a/src/trspecfit/trspecfit.py b/src/trspecfit/trspecfit.py index b329bae..c94e0a4 100644 --- a/src/trspecfit/trspecfit.py +++ b/src/trspecfit/trspecfit.py @@ -57,7 +57,7 @@ import time import types import warnings -from collections.abc import Callable, Sequence +from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Literal, cast # TYPE_CHECKING is False at runtime so this import is skipped during execution. @@ -71,7 +71,7 @@ from IPython.display import display from ruamel.yaml import YAML -from trspecfit import fitlib, mcp, spectra +from trspecfit import fitlib, mcp # standardized plotting configuration from trspecfit.config.plot import PlotConfig @@ -135,10 +135,9 @@ class Project: displayed or saved - 1: Interactive / notebook / UI mode -- show timing, fit results, save plots and data - spec_lib : module - Module containing spectrum fitting functions (default: spectra) spec_fun_str : str - Name of fitting function in spec_lib + 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) @@ -226,19 +225,10 @@ def _set_defaults(self) -> None: self.da_fmt = "%04d" self.da_slices_fmt = "%06d" # Advanced settings - self.spec_lib = spectra self.spec_fun_str = "fit_model_gir" self.skip_first_n_spec = -1 self.first_n_spec_only = -1 - @property - def spec_fun(self) -> Callable: - """ - Dynamically get the spectrum fitting function. - """ - - return cast("Callable", getattr(self.spec_lib, self.spec_fun_str)) - # def __repr__(self) -> str: return f"Project(path='{self.path}', name='{self.name}')" @@ -905,7 +895,6 @@ def fit_2d( const: tuple[Any, ...] = ( np.array([]), # x — unused by fit_project_mcp concat_data, - spectra, "fit_project_mcp", 0, [], # no additional e_lim slicing @@ -1592,7 +1581,6 @@ def describe_model( x=self.energy, y=self.data_base, fit_fun_str=self.p.spec_fun_str, - package=self.p.spec_lib, par_init=[], par_fin=mod.lmfit_pars, args=(mod, 1), @@ -2033,7 +2021,6 @@ def fit_baseline( self.model_base.const = ( self.energy, self.data_base, - self.p.spec_lib, _fun_str, 0, self.e_lim, @@ -2075,7 +2062,6 @@ def fit_baseline( x=self.energy, y=self.data_base, fit_fun_str=self.p.spec_fun_str, - package=self.p.spec_lib, par_init=initial_guess, par_fin=self.model_base.result[1], args=self.model_base.args, @@ -2206,12 +2192,11 @@ def fit_spectrum( # define (and create) path where spectrum fit results will be saved to path_spec_results = self.create_model_path(model_name) - # const = (x, data, package, fnctn string, unpack, energy limits, time limits) + # const = (x, data, fnctn string, unpack, energy limits, time limits) _fun_str = self.p.spec_fun_str self.model_spec.const = ( self.energy, self.data_spec, - self.p.spec_lib, _fun_str, 0, self.e_lim, @@ -2257,7 +2242,6 @@ def fit_spectrum( x=self.energy, y=self.data_spec, fit_fun_str=self.p.spec_fun_str, - package=self.p.spec_lib, par_init=initial_guess, par_fin=self.model_spec.result[1], args=self.model_spec.args, @@ -2390,11 +2374,10 @@ def fit_slice_by_slice( self.model_sbs.lmfit_pars, return_type="list" ) - # const = (x, data, package, fnctn str, unpack, energy limits, time limits) + # const = (x, data, fnctn str, unpack, energy limits, time limits) self.model_sbs.const = ( self.energy, s, - self.p.spec_lib, _fun_str, 0, self.e_lim, @@ -2424,7 +2407,6 @@ def fit_slice_by_slice( x=self.model_sbs.const[0], y=self.model_sbs.const[1], fit_fun_str=self.p.spec_fun_str, - package=self.p.spec_lib, par_init=initial_guess, par_fin=result_sbs[1], args=self.model_sbs.args, @@ -2830,11 +2812,10 @@ def fit_2d(self, model_name: str, stages: int = 1, **fit_wrapper_kwargs) -> None else: _args = (self.model_2d, 2) - # const [x, data, package, function string, unpack, energy limits, time limits] + # const [x, data, function string, unpack, energy limits, time limits] self.model_2d.const = ( self.energy, self.data, - self.p.spec_lib, _fun_str, 0, self.e_lim, diff --git a/tests/test_file.py b/tests/test_file.py index 574c4c1..71fa24f 100644 --- a/tests/test_file.py +++ b/tests/test_file.py @@ -687,7 +687,6 @@ def _call_residual(self, file, *, res_type="res"): const = ( file.energy, file.data, - file.p.spec_lib, file.p.spec_fun_str, 0, file.e_lim, diff --git a/tests/test_gir_integration.py b/tests/test_gir_integration.py index 4732518..3b9e258 100644 --- a/tests/test_gir_integration.py +++ b/tests/test_gir_integration.py @@ -369,7 +369,6 @@ def test_residual_same_gir_vs_mcp(self): par=par, x=file.energy, data=data, - package=spectra, fit_fun_str="fit_model_gir", args=(plan, theta_indices, model, 2), ) @@ -379,7 +378,6 @@ def test_residual_same_gir_vs_mcp(self): par=par, x=file.energy, data=data, - package=spectra, fit_fun_str="fit_model_mcp", args=(model, 2), ) @@ -442,7 +440,6 @@ def test_residual_with_slicing(self): par=par, x=file.energy, data=data, - package=spectra, fit_fun_str="fit_model_gir", e_lim=e_lim, t_lim=t_lim, @@ -453,7 +450,6 @@ def test_residual_with_slicing(self): par=par, x=file.energy, data=data, - package=spectra, fit_fun_str="fit_model_mcp", e_lim=e_lim, t_lim=t_lim, @@ -505,7 +501,6 @@ def test_residual_same_gir_vs_mcp_irf(self, dyn_model): par=par, x=file.energy, data=data, - package=spectra, fit_fun_str="fit_model_gir", args=(plan, theta_indices, model, 2), ) @@ -514,7 +509,6 @@ def test_residual_same_gir_vs_mcp_irf(self, dyn_model): par=par, x=file.energy, data=data, - package=spectra, fit_fun_str="fit_model_mcp", args=(model, 2), ) @@ -576,7 +570,6 @@ def test_subcycle_residual_gir_vs_mcp(self): par=par, x=file.energy, data=data, - package=spectra, fit_fun_str="fit_model_gir", args=(plan, theta_indices, model, 2), ) @@ -584,7 +577,6 @@ def test_subcycle_residual_gir_vs_mcp(self): par=par, x=file.energy, data=data, - package=spectra, fit_fun_str="fit_model_mcp", args=(model, 2), ) @@ -790,7 +782,6 @@ def test_residual_same_gir_vs_mcp_1d(self): par=par, x=file.energy, data=data, - package=spectra, fit_fun_str="fit_model_gir", args=(plan, theta_indices, model, 1), ) @@ -800,7 +791,6 @@ def test_residual_same_gir_vs_mcp_1d(self): par=par, x=file.energy, data=data, - package=spectra, fit_fun_str="fit_model_mcp", args=(model, 1), ) @@ -834,7 +824,6 @@ def test_residual_with_e_lim_1d(self): par=par, x=file.energy, data=data, - package=spectra, fit_fun_str="fit_model_gir", e_lim=e_lim, args=(plan, theta_indices, model, 1), @@ -844,7 +833,6 @@ def test_residual_with_e_lim_1d(self): par=par, x=file.energy, data=data, - package=spectra, fit_fun_str="fit_model_mcp", e_lim=e_lim, args=(model, 1), @@ -880,7 +868,6 @@ def test_residual_same_gir_vs_mcp_profile_1d(self): par=par, x=file.energy, data=data, - package=spectra, fit_fun_str="fit_model_gir", args=(plan, theta_indices, model, 1), ) @@ -888,7 +875,6 @@ def test_residual_same_gir_vs_mcp_profile_1d(self): par=par, x=file.energy, data=data, - package=spectra, fit_fun_str="fit_model_mcp", args=(model, 1), ) @@ -973,7 +959,6 @@ def test_residual_same_gir_vs_mcp_profile_2d(self): par=par, x=file.energy, data=data, - package=spectra, fit_fun_str="fit_model_gir", args=(plan, theta_indices, model, 2), ) @@ -981,7 +966,6 @@ def test_residual_same_gir_vs_mcp_profile_2d(self): par=par, x=file.energy, data=data, - package=spectra, fit_fun_str="fit_model_mcp", args=(model, 2), ) diff --git a/tests/test_mcp_library.py b/tests/test_mcp_library.py index 0e137af..46fe974 100644 --- a/tests/test_mcp_library.py +++ b/tests/test_mcp_library.py @@ -844,5 +844,137 @@ def test_profile_with_dynamics(self): assert not np.allclose(val_early, val_late) +# +# +class TestMCPPickling: + """Pickle / deepcopy contract for Model, Component, and Par. + + These exist to (a) enable multiprocessing dispatch of live models + (e.g. lmfit.emcee ``workers > 1``, future slice-by-slice parallelism) + and (b) make ``copy.deepcopy`` work for safe comparison workflows. + Pickled instances are for short-lived process-boundary transfer, + NOT for persistence — parent back-refs are nulled. + """ + + # + def _make_fittable_file(self): + """File with a 1D energy model loaded and data populated.""" + + from trspecfit import File, Project + + project = Project(path="tests") + project.show_output = 0 + file = File(parent_project=project, energy=np.linspace(80, 90, 101)) + file.load_model( + model_yaml="models/file_energy.yaml", + model_info="single_glp", + ) + # synthesize clean 1D data from the model + assert file.model_active is not None # type guard + file.model_active.create_value_1d() + assert file.model_active.value_1d is not None # type guard + file.data_base = file.model_active.value_1d.copy() + file.e_lim = [0, len(file.energy)] + return file + + # + def test_pickle_component_roundtrip(self): + """Component roundtrips through pickle with package module restored.""" + + import pickle + + comp = Component("GLP_01") + comp.add_pars( + { + "A": [10, True, 0, 20], + "x0": [85, True, 80, 90], + "F": [1.5, True, 1, 2], + "m": [0.3, False, 0, 1], + } + ) + + blob = pickle.dumps(comp) + restored = pickle.loads(blob) + + # package reloaded as module (not a string) + import types + + assert isinstance(restored.package, types.ModuleType) + assert restored.package.__name__ == comp.package.__name__ + assert callable(restored.fct) + # par_dict values survive roundtrip + assert restored.par_dict == comp.par_dict + # parent_model is nulled on pickle (detached state) + assert restored.parent_model is None + + # + def test_pickle_model_roundtrip_and_evaluate(self): + """Pickled Model can still evaluate create_value_1d after roundtrip.""" + + import pickle + + file = self._make_fittable_file() + model = file.model_active + expected = model.value_1d.copy() + + restored = pickle.loads(pickle.dumps(model)) + # parent_file stripped on pickle (detached state) + assert restored.parent_file is None + # transient fit state stripped + assert restored.const is None + assert restored.args is None + # model can still evaluate after roundtrip + restored.create_value_1d() + np.testing.assert_allclose(restored.value_1d, expected, rtol=1e-12) + + # + def test_deepcopy_model_is_independent(self): + """copy.deepcopy produces a fresh Model that can be mutated safely.""" + + import copy + + file = self._make_fittable_file() + model = file.model_active + clone = copy.deepcopy(model) + + assert clone is not model + assert clone.components[0] is not model.components[0] + + # Mutating the clone must not affect the original + original_par_name = clone.parameter_names[0] + clone.lmfit_pars[original_par_name].value = 42.0 + assert model.lmfit_pars[original_par_name].value != 42.0 + + # + @pytest.mark.slow + def test_mcmc_workers_2_does_not_pickle_error(self): + """lmfit.emcee(workers=2) must not fail with TypeError: cannot pickle module. + + This was a latent bug before the pickle hooks landed: MCMC + parallelism was parameterized via ``MC(workers=N)`` but silently + broken because the residual closure could not cross the + multiprocessing boundary. + """ + + from trspecfit.utils.lmfit import MC + + file = self._make_fittable_file() + # Tiny MCMC settings so the test runs quickly. nwalkers must exceed + # 2 * n_params for emcee's red-blue move, hence 32 for a 4-param GLP. + mc = MC(use_mc=1, steps=20, nwalkers=32, burn=5, thin=1, workers=2) + + # fit_baseline raises TypeError before the pickle hooks were added; + # after the hooks, it completes (even if MCMC itself converges poorly + # on 20 steps, the point here is that it doesn't crash). + try: + file.fit_baseline( + model_name="single_glp", stages=1, try_ci=0, mc_settings=mc + ) + except TypeError as e: + if "pickle" in str(e) or "module" in str(e): + pytest.fail(f"MCMC workers=2 hit a pickle error: {e}") + raise + + if __name__ == "__main__": pytest.main([__file__])