Pickleable models

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

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.

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"},
 ]

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)

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:
         """

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,

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,