From 5f8bd2f911ca87ed66800a04b44ad3729b88747c Mon Sep 17 00:00:00 2001 From: rozyczko Date: Thu, 17 Sep 2026 11:46:47 +0200 Subject: [PATCH] Make Sampler completely Fitter-agnostic --- docs/docs/tutorials/fitting-bayesian.ipynb | 76 +++---- src/easyscience/fitting/fitter.py | 93 +------- src/easyscience/fitting/multi_fitter.py | 112 ++-------- src/easyscience/fitting/reshaping.py | 216 ++++++++++++++++++ src/easyscience/fitting/sampler.py | 196 +++++++++++------ tests/integration/fitting/test_sampler.py | 135 ++++++------ tests/unit/fitting/test_multi_fitter.py | 15 -- tests/unit/fitting/test_reshaping.py | 100 +++++++++ tests/unit/fitting/test_sampler.py | 241 +++++++++++++-------- 9 files changed, 715 insertions(+), 469 deletions(-) create mode 100644 src/easyscience/fitting/reshaping.py create mode 100644 tests/unit/fitting/test_reshaping.py diff --git a/docs/docs/tutorials/fitting-bayesian.ipynb b/docs/docs/tutorials/fitting-bayesian.ipynb index 302800c2..ed90222e 100644 --- a/docs/docs/tutorials/fitting-bayesian.ipynb +++ b/docs/docs/tutorials/fitting-bayesian.ipynb @@ -28,7 +28,7 @@ "\n", "where $\\theta$ are the model parameters, $d$ is the observed data, $p(d \\mid \\theta)$ is the likelihood, and $p(\\theta)$ is the prior. In `easyscience`, the `min`/`max` bounds of a `Parameter` are interpreted as a **uniform prior**, and a Gaussian likelihood is constructed from the data and supplied weights.\n", "\n", - "`easyscience` exposes a Bayesian Markov-chain Monte Carlo (MCMC) sampler through the `Sampler` class. Under the hood this uses BUMPS' DREAM sampler, so the underlying minimizer must be switched to BUMPS.\n", + "`easyscience` exposes a Bayesian Markov-chain Monte Carlo (MCMC) sampler through the `Sampler` class. It is a parallel entry point to `Fitter`: both take a model object and a model function, so you can sample without ever creating a `Fitter`. Under the hood `Sampler` uses BUMPS' DREAM sampler, so the `bumps` package must be installed.\n", "\n", "```{note}\n", "This tutorial focuses on Bayesian analysis with a simple QENS model for illustration. For dedicated QENS fitting with more sophisticated models, consider using [`EasyDynamics`](https://github.com/easyscience/easydynamics).\n", @@ -138,12 +138,12 @@ }, { "cell_type": "markdown", - "id": "9", + "id": "7", "metadata": {}, "source": [ "## Defining parameters with priors\n", "\n", - "Create four `Parameter` objects, for the area $A$, $\\gamma$, $\\omega_0$ and $\\sigma$. The `min` and `max` arguments define a **uniform prior** on each parameter — the sampler will only consider values inside this range and will treat every value inside the range as equally plausible *a priori*.\n", + "Create four `Parameter` objects, for the area $A$, $\\gamma$, $\\omega_0$ and $\\sigma$. The `min` and `max` arguments define a **uniform prior** on each parameter — the sampler will only consider values inside this range and will treat every value inside the range as equally plausible *a priori*. The parameters are then collected in an `ObjBase` container: this is the model object that both `Fitter` and `Sampler` take, so it is defined here, before the optional fit.\n", "\n", "| Parameter | Initial Value | Min | Max |\n", "| --- | --- | --- | --- |\n", @@ -156,21 +156,24 @@ { "cell_type": "code", "execution_count": null, - "id": "10", + "id": "8", "metadata": {}, "outputs": [], "source": [ + "from easyscience import ObjBase\n", "from easyscience import Parameter\n", "\n", "area = Parameter(name='area', value=10, fixed=False, min=1, max=100)\n", "gamma = Parameter(name='gamma', value=8e-3, fixed=False, min=1e-4, max=1e-2)\n", "omega_0 = Parameter(name='omega_0', value=1e-3, fixed=False, min=0, max=2e-3)\n", - "sigma = Parameter(name='sigma', value=1e-3, fixed=False, min=1e-5, max=1e-1)" + "sigma = Parameter(name='sigma', value=1e-3, fixed=False, min=1e-5, max=1e-1)\n", + "\n", + "parameter_container = ObjBase(name='params', A=area, gamma=gamma, omega_0=omega_0, sigma=sigma)" ] }, { "cell_type": "markdown", - "id": "77ce3f63", + "id": "9", "metadata": {}, "source": [ "## The model\n", @@ -193,7 +196,7 @@ { "cell_type": "code", "execution_count": null, - "id": "36a0c4f4", + "id": "10", "metadata": {}, "outputs": [], "source": [ @@ -218,8 +221,8 @@ "## Maximum-likelihood fit (optional, but recommended)\n", "\n", "Perform a quick maximum-likelihood fit. This is **not** a prerequisite for sampling: `Sampler`\n", - "needs a *configured* `Fitter`, not a *fitted* one, and you can sample straight from the initial\n", - "parameter values.\n", + "does not need a `Fitter` at all, and you can sample straight from the initial parameter\n", + "values.\n", "\n", "It is worth doing anyway, for two reasons:\n", "\n", @@ -239,9 +242,6 @@ "outputs": [], "source": [ "from easyscience import Fitter\n", - "from easyscience import ObjBase\n", - "\n", - "parameter_container = ObjBase(name='params', A=area, gamma=gamma, omega_0=omega_0, sigma=sigma)\n", "\n", "mle_fitter = Fitter(parameter_container, intensity_model)\n", "mle_result = mle_fitter.fit(x=omega, y=intensity_obs, weights=1 / intensity_error)\n", @@ -261,7 +261,7 @@ "\n", "We now draw samples from the posterior distribution $p(\\theta \\mid d)$ using the BUMPS DREAM (DiffeRential Evolution Adaptive Metropolis) algorithm. DREAM is an ensemble MCMC method that runs multiple chains in parallel and automatically tunes the proposal distribution.\n", "\n", - "DREAM only works with the BUMPS minimizer. We reuse the ``mle_fitter`` created above — any configured `Fitter` would do, and it does not have to have been fitted — switch it to BUMPS, and create a `Sampler` instance bound to the fitter and data. Calling `sampler.sample()` returns a `SamplingResults` object with the following attributes:\n", + "Create a `Sampler` from the same `parameter_container` and `intensity_model` we gave the `Fitter`, bound to the data. No `Fitter` is involved: the sampler only needs the model object, the model function, the data and the `bumps` package. Calling `sampler.sample()` returns a `SamplingResults` object with the following attributes:\n", "\n", "- `draws`: a `(n_samples, n_parameters)` array of posterior samples: each **row** is one complete draw from the joint posterior (one value for every parameter simultaneously), and each **column** holds all sampled values for a single parameter. Note this is a *trimmed* view of the chain rather than the raw buffer, so `n_samples` is smaller than `samples / thin` — see the note under [Extend the chain](#extend-the-chain-and-check-convergence);\n", "- `param_names`: the unique names of the parameters, in the same column order as `draws`;\n", @@ -274,7 +274,9 @@ "- `burn` (500): the number of initial *burn-in* generations to discard — the sampler needs time to find the typical set of the posterior, and early samples are not representative. Note this counts generations, not raw samples, so `burn=500` discards `500 × n_chains` raw samples;\n", "- `thin` (2): the *thinning* interval — only every second generation is kept, which reduces autocorrelation between consecutive draws;\n", "\n", - "First, we switch to the BUMPS minimizer:" + "```{note}\n", + "If you already have a `Fitter`, `Sampler.from_fitter(mle_fitter, omega, intensity_obs, weights=1 / intensity_error)` builds the same sampler from the model bound to it. The fitter is neither needed afterwards nor modified.\n", + "```" ] }, { @@ -283,22 +285,12 @@ "id": "14", "metadata": {}, "outputs": [], - "source": [ - "from easyscience import AvailableMinimizers\n", - "\n", - "mle_fitter.switch_minimizer(AvailableMinimizers.Bumps)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "5a219fcd", - "metadata": {}, - "outputs": [], "source": [ "from easyscience.fitting import Sampler\n", "\n", - "sampler = Sampler(mle_fitter, omega, intensity_obs, weights=1 / intensity_error)\n", + "sampler = Sampler(\n", + " parameter_container, intensity_model, omega, intensity_obs, weights=1 / intensity_error\n", + ")\n", "results = sampler.sample(samples=10000, burn=500, thin=2)\n", "\n", "print(f'Drew {results.draws.shape[0]} samples for {results.draws.shape[1]} parameters.')\n", @@ -307,7 +299,7 @@ }, { "cell_type": "markdown", - "id": "8766b170", + "id": "15", "metadata": {}, "source": [ "## Convergence diagnostics\n", @@ -323,7 +315,7 @@ { "cell_type": "code", "execution_count": null, - "id": "3c49ab6f", + "id": "16", "metadata": {}, "outputs": [], "source": [ @@ -359,7 +351,7 @@ }, { "cell_type": "markdown", - "id": "15", + "id": "17", "metadata": {}, "source": [ "## Posterior summaries\n", @@ -377,7 +369,7 @@ { "cell_type": "code", "execution_count": null, - "id": "ce3e38a8", + "id": "18", "metadata": {}, "outputs": [], "source": [] @@ -385,7 +377,7 @@ { "cell_type": "code", "execution_count": null, - "id": "16", + "id": "19", "metadata": {}, "outputs": [], "source": [ @@ -407,7 +399,7 @@ }, { "cell_type": "markdown", - "id": "17", + "id": "20", "metadata": {}, "source": [ "## Visualise the joint posterior\n", @@ -418,7 +410,7 @@ { "cell_type": "code", "execution_count": null, - "id": "18", + "id": "21", "metadata": {}, "outputs": [], "source": [ @@ -452,7 +444,7 @@ }, { "cell_type": "markdown", - "id": "19", + "id": "22", "metadata": {}, "source": [ "## Posterior-predictive band\n", @@ -463,7 +455,7 @@ { "cell_type": "code", "execution_count": null, - "id": "20", + "id": "23", "metadata": {}, "outputs": [], "source": [ @@ -499,7 +491,7 @@ }, { "cell_type": "markdown", - "id": "03339658", + "id": "24", "metadata": {}, "source": [ "## Extend the chain and check convergence\n", @@ -543,7 +535,7 @@ { "cell_type": "code", "execution_count": null, - "id": "293b140b", + "id": "25", "metadata": {}, "outputs": [], "source": [ @@ -573,7 +565,7 @@ { "cell_type": "code", "execution_count": null, - "id": "9ec4302c", + "id": "26", "metadata": {}, "outputs": [], "source": [ @@ -614,7 +606,7 @@ { "cell_type": "code", "execution_count": null, - "id": "b0f30be6", + "id": "27", "metadata": {}, "outputs": [], "source": [ @@ -641,7 +633,7 @@ }, { "cell_type": "markdown", - "id": "50b7213a", + "id": "28", "metadata": {}, "source": [ "### What is Gelman-Rubin R-hat?\n", @@ -675,7 +667,7 @@ { "cell_type": "code", "execution_count": null, - "id": "3449e0a7", + "id": "29", "metadata": {}, "outputs": [], "source": [ diff --git a/src/easyscience/fitting/fitter.py b/src/easyscience/fitting/fitter.py index d1f20080..41857098 100644 --- a/src/easyscience/fitting/fitter.py +++ b/src/easyscience/fitting/fitter.py @@ -15,6 +15,8 @@ from .minimizers import FitResults from .minimizers import MinimizerBase from .minimizers.factory import factory +from .reshaping import inject_x +from .reshaping import reshape_dataset DEFAULT_MINIMIZER = AvailableMinimizers.LMFit_leastsq @@ -231,13 +233,10 @@ def _fit_function_wrapper( self, real_x: np.ndarray | None = None, flatten: bool = True, - dependent_dims: list[tuple[int, ...]] | None = None, ) -> Callable: """ - Simple fit function which injects the real X (independent) - values into the optimizer function. - - This will also flatten the results if needed. + Wrap the fit function so it evaluates on the real X (independent) + values instead of the optimizer's dummy x, flattening if needed. Parameters ---------- @@ -245,28 +244,13 @@ def _fit_function_wrapper( Independent x parameters to be injected. By default, None. flatten : bool, default=True Should the result be a flat 1D array? By default, True. - dependent_dims : list[tuple[int, ...]] | None, default=None - Unused for a single dataset; accepted so that callers can - pass it uniformly to ``Fitter`` and ``MultiFitter``. By - default, None. Returns ------- Callable Wrapped optimizer function. """ - fun = self._fit_function - - @functools.wraps(fun) - def wrapped_fit_function(x, **kwargs): - if real_x is not None: - x = real_x - dependent = fun(x, **kwargs) - if flatten: - dependent = dependent.flatten() - return dependent - - return wrapped_fit_function + return inject_x(self._fit_function, real_x, flatten=flatten) @property def fit(self) -> Callable: @@ -328,72 +312,7 @@ def inner_fit_callable( return inner_fit_callable - @staticmethod - def _precompute_reshaping( - x: np.ndarray, - y: np.ndarray, - weights: np.ndarray | None, - vectorized: bool, - ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray | None, tuple[int, ...]]: - """ - Check the dimensions of the inputs and reshape if necessary. - - Parameters - ---------- - x : np.ndarray - ND matrix of dependent points. - y : np.ndarray - N-1D matrix of independent points. - weights : np.ndarray | None - Optional weights for the fit. - vectorized : bool - Whether ``x`` already stores vectorized coordinates. - - Returns - ------- - tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray | None, tuple[int, ...]] - Reshaped x values, reshaped input data, flattened y values, - flattened weights, and the original x shape. - - Raises - ------ - ValueError - If the shapes of ``x`` and ``y`` are incompatible. - """ - # Make sure that they are np arrays - x_new = np.array(x) - y_new = np.array(y) - # Get the shape - x_shape = x_new.shape - # Check if the x data is 1D - if len(x_shape) > 1: - # It is ND data - # Check if the data is vectorized. i.e. should x be [NxMx...x Ndims] - if vectorized: - # Assert that the shapes are the same - if np.all(x_shape[:-1] != y_new.shape): - raise ValueError('The shape of the x and y data must be the same') - # If so do nothing but note that the data is vectorized - # x_shape = (-1,) # Should this be done? - else: - # Assert that the shapes are the same - if np.prod(x_new.shape[:-1]) != y_new.size: - raise ValueError('The number of elements in x and y data must be the same') - # Reshape the data to be [len(NxMx..), Ndims] i.e. flatten to columns - x_new = x_new.reshape(-1, x_shape[-1], order='F') - else: - # Assert that the shapes are the same - if np.all(x_shape != y_new.shape): - raise ValueError('The shape of the x and y data must be the same') - # It is 1D data - x_new = x.flatten() - # The optimizer needs a 1D array, flatten the y data - y_new = y_new.flatten() - if weights is not None: - weights = np.array(weights).flatten() - # Make a 'dummy' x array for the fit function - x_for_fit = np.array(range(y_new.size)) - return x_for_fit, x_new, y_new, weights, x_shape + _precompute_reshaping = staticmethod(reshape_dataset) @staticmethod def _post_compute_reshaping( diff --git a/src/easyscience/fitting/multi_fitter.py b/src/easyscience/fitting/multi_fitter.py index 21b2767c..d328047b 100644 --- a/src/easyscience/fitting/multi_fitter.py +++ b/src/easyscience/fitting/multi_fitter.py @@ -8,6 +8,8 @@ from ..base_classes import CollectionBase from .fitter import Fitter from .minimizers import FitResults +from .reshaping import inject_x_multi +from .reshaping import reshape_datasets class MultiFitter(Fitter): @@ -34,17 +36,27 @@ def __init__( # not possible to change the fitting engine. super().__init__(self._fit_objects, self._fit_functions[0]) + @property + def fit_functions(self) -> list[Callable]: + """ + Get the per-dataset fit functions, in dataset order. + + Returns + ------- + list[Callable] + One fit function per dataset. + """ + return list(self._fit_functions) + def _fit_function_wrapper( self, real_x: list[np.ndarray] | None = None, flatten: bool = True, - dependent_dims: list[tuple[int, ...]] | None = None, ) -> Callable: """ - Simple fit function which injects the N real X (independent) - values into the optimizer function. - - This will also flatten the results if needed. + Wrap the per-dataset fit functions into one function that + evaluates each on its real X (independent) values and + concatenates the results, flattening if needed. Parameters ---------- @@ -53,101 +65,15 @@ def _fit_function_wrapper( None. flatten : bool, default=True Should the result be a flat 1D array? By default, True. - dependent_dims : list[tuple[int, ...]] | None, default=None - Per-dataset dependent shapes used to slice the combined - output. When ``None``, ``self._dependent_dims`` (set by - ``fit``) is read at call time. By default, None. Returns ------- Callable Wrapped optimizer function. """ - # Extract of a list of callable functions. - # ``Fitter._fit_function_wrapper`` reads ``self._fit_function``, so it - # is repointed per dataset inside the loop; the original must be - # restored afterwards or every caller (``Fitter.fit`` aside, which - # snapshots it itself, e.g. sampling) is left with the *last* - # dataset's function on the user-visible ``fit_function`` surface. - wrapped_fns = [] - original_fit_function = self._fit_function - try: - for this_x, this_fun in zip(real_x, self._fit_functions): - self._fit_function = this_fun - wrapped_fns.append(Fitter._fit_function_wrapper(self, this_x, flatten=flatten)) - finally: - self._fit_function = original_fit_function - - def wrapped_fun(x, **kwargs): - # Generate an empty Y based on x - y = np.zeros_like(x) - i = 0 - dims = self._dependent_dims if dependent_dims is None else dependent_dims - # Iterate through wrapped functions, passing the WRONG x, the correct - # x was injected in the step above. - for idx, dim in enumerate(dims): - ep = i + np.prod(dim) - y[i:ep] = wrapped_fns[idx](x, **kwargs) - i = ep - return y - - return wrapped_fun - - @staticmethod - def _precompute_reshaping( - x: list[np.ndarray], - y: list[np.ndarray], - weights: list[np.ndarray] | None, - vectorized: bool, - ) -> tuple[np.ndarray, list[np.ndarray], np.ndarray, np.ndarray | None, list[tuple[int, ...]]]: - """ - Convert an array of X's and Y's to an acceptable shape for - fitting. + return inject_x_multi(self._fit_functions, real_x, self._dependent_dims, flatten=flatten) - Parameters - ---------- - x : list[np.ndarray] - List of independent variables. - y : list[np.ndarray] - List of dependent variables. - weights : list[np.ndarray] | None - Optional weights for each dataset. - vectorized : bool - When ``True``, each x array may be multi-dimensional (e.g. - an ``(N, M, 2)`` grid for a 2D model) and is left as-is. - When ``False`` (default), each x array is expected to be - 1-D. - - Returns - ------- - tuple[np.ndarray, list[np.ndarray], np.ndarray, np.ndarray | None, list[tuple[int, ...]]] - Reshaped x values, reshaped input data, flattened y values, - flattened weights, and stored dependent dimensions. - """ - if weights is None: - weights = [None] * len(x) - _, _x_new, _y_new, _weights, _dims = Fitter._precompute_reshaping( - x[0], y[0], weights[0], vectorized - ) - x_new = [_x_new] - y_new = [_y_new] - w_new = [_weights] - dims = [_dims] - for _x, _y, _w in zip(x[1::], y[1::], weights[1::]): - _, _x_new, _y_new, _weights, _dims = Fitter._precompute_reshaping( - _x, _y, _w, vectorized - ) - x_new.append(_x_new) - y_new.append(_y_new) - w_new.append(_weights) - dims.append(_dims) - y_new = np.hstack(y_new) - if w_new[0] is None: - w_new = None - else: - w_new = np.hstack(w_new) - x_fit = np.linspace(0, y_new.size - 1, y_new.size) - return x_fit, x_new, y_new, w_new, dims + _precompute_reshaping = staticmethod(reshape_datasets) def _post_compute_reshaping( self, diff --git a/src/easyscience/fitting/reshaping.py b/src/easyscience/fitting/reshaping.py new file mode 100644 index 00000000..c68a7739 --- /dev/null +++ b/src/easyscience/fitting/reshaping.py @@ -0,0 +1,216 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause +"""Data reshaping and fit-function wrapping shared by ``Fitter`` and ``Sampler``. + +Both fitting and sampling provide the engine witha flat 1-D ``y`` array +and a dummy 1-D ``x`` index array, while the user's fit function is called +with the real ``x``. The helpers here do the translation in one place: +``reshape_dataset``/``inject_x`` do one dataset, +``reshape_datasets``/``inject_x_multi`` do a list of datasets. +""" + +from __future__ import annotations + +import functools +from typing import Callable + +import numpy as np + + +def reshape_dataset( + x: np.ndarray, + y: np.ndarray, + weights: np.ndarray | None, + vectorized: bool, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray | None, tuple[int, ...]]: + """ + Check the dimensions of the inputs and reshape if necessary. + + Parameters + ---------- + x : np.ndarray + Independent points; 1-D, or ND with the coordinate components on + the last axis. + y : np.ndarray + Dependent points, one per observation. + weights : np.ndarray | None + Optional weights for the fit. + vectorized : bool + Whether ``x`` already stores vectorized coordinates. + + Returns + ------- + tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray | None, tuple[int, ...]] + Dummy x index array for the engine, reshaped x values, flattened + y values, flattened weights, and the shape of ``y`` (the dependent + dimensions; its product is the number of observations, which the + multi-dataset helpers use to slice the combined output). + + Raises + ------ + ValueError + If the shapes of ``x`` and ``y`` are incompatible. + """ + # Make sure that they are np arrays + x_new = np.array(x) + y_new = np.array(y) + # Get the shapes + x_shape = x_new.shape + y_shape = y_new.shape + # Check if the x data is 1D + if len(x_shape) > 1: + # It is ND data + # Check if the data is vectorized. i.e. should x be [NxMx...x Ndims] + if vectorized: + # Assert that the shapes are the same + if np.all(x_shape[:-1] != y_new.shape): + raise ValueError('The shape of the x and y data must be the same') + # If so do nothing but note that the data is vectorized + # x_shape = (-1,) # Should this be done? + else: + # Assert that the shapes are the same + if np.prod(x_new.shape[:-1]) != y_new.size: + raise ValueError('The number of elements in x and y data must be the same') + # Reshape the data to be [len(NxMx..), Ndims] i.e. flatten to columns + x_new = x_new.reshape(-1, x_shape[-1], order='F') + else: + # Assert that the shapes are the same + if np.all(x_shape != y_new.shape): + raise ValueError('The shape of the x and y data must be the same') + # It is 1D data + x_new = x.flatten() + # The optimizer needs a 1D array, flatten the y data + y_new = y_new.flatten() + if weights is not None: + weights = np.array(weights).flatten() + # Make a 'dummy' x array for the fit function + x_for_fit = np.array(range(y_new.size)) + return x_for_fit, x_new, y_new, weights, y_shape + + +def reshape_datasets( + x: list[np.ndarray], + y: list[np.ndarray], + weights: list[np.ndarray] | None, + vectorized: bool, +) -> tuple[np.ndarray, list[np.ndarray], np.ndarray, np.ndarray | None, list[tuple[int, ...]]]: + """ + Convert a list of X's and Y's to an acceptable shape for fitting. + + Parameters + ---------- + x : list[np.ndarray] + List of independent variables. + y : list[np.ndarray] + List of dependent variables. + weights : list[np.ndarray] | None + Optional weights for each dataset. + vectorized : bool + When ``True``, each x array may be multi-dimensional (e.g. an + ``(N, M, 2)`` grid for a 2D model) and is left as-is. When + ``False`` (default), each x array is expected to be 1-D. + + Returns + ------- + tuple[np.ndarray, list[np.ndarray], np.ndarray, np.ndarray | None, list[tuple[int, ...]]] + Dummy x index array for the engine, per-dataset reshaped x values, + concatenated y values, concatenated weights, and the per-dataset + dependent dimensions. + """ + if weights is None: + weights = [None] * len(x) + x_new, y_new, w_new, dims = [], [], [], [] + for _x, _y, _w in zip(x, y, weights): + _, _x_new, _y_new, _w_new, _dims = reshape_dataset(_x, _y, _w, vectorized) + x_new.append(_x_new) + y_new.append(_y_new) + w_new.append(_w_new) + dims.append(_dims) + y_new = np.hstack(y_new) + w_new = None if w_new[0] is None else np.hstack(w_new) + x_fit = np.linspace(0, y_new.size - 1, y_new.size) + return x_fit, x_new, y_new, w_new, dims + + +def inject_x( + fit_function: Callable, + real_x: np.ndarray | None = None, + flatten: bool = True, +) -> Callable: + """ + Wrap a fit function so it ignores the engine's dummy x and evaluates + on ``real_x`` instead, flattening the result if needed. + + Parameters + ---------- + fit_function : Callable + The user's fit function. + real_x : np.ndarray | None, default=None + Independent x values to be injected. By default, None. + flatten : bool, default=True + Should the result be a flat 1D array? By default, True. + + Returns + ------- + Callable + Wrapped optimizer function. + """ + + @functools.wraps(fit_function) + def wrapped_fit_function(x, **kwargs): + if real_x is not None: + x = real_x + dependent = fit_function(x, **kwargs) + if flatten: + dependent = dependent.flatten() + return dependent + + return wrapped_fit_function + + +def inject_x_multi( + fit_functions: list[Callable], + real_x: list[np.ndarray], + dims: list[tuple[int, ...]], + flatten: bool = True, +) -> Callable: + """ + Wrap one fit function per dataset into a single function whose output + is the concatenation of the per-dataset outputs. + + Parameters + ---------- + fit_functions : list[Callable] + One fit function per dataset. + real_x : list[np.ndarray] + One independent x array per dataset, injected into the matching + fit function. + dims : list[tuple[int, ...]] + Per-dataset dependent (``y``) shapes used to slice the combined + output, as returned by ``reshape_datasets``. + flatten : bool, default=True + Should each per-dataset result be flattened? By default, True. + + Returns + ------- + Callable + Wrapped optimizer function. + """ + wrapped_fns = [ + inject_x(this_fun, this_x, flatten=flatten) + for this_x, this_fun in zip(real_x, fit_functions) + ] + + def wrapped_fun(x, **kwargs): + # Generate an empty Y based on x + y = np.zeros_like(x) + i = 0 + # Iterate through wrapped functions, passing the WRONG x, the correct + # x was injected in the step above. + for wrapped, dim in zip(wrapped_fns, dims): + ep = i + np.prod(dim) + y[i:ep] = wrapped(x, **kwargs) + i = ep + return y + + return wrapped_fun diff --git a/src/easyscience/fitting/sampler.py b/src/easyscience/fitting/sampler.py index 345fe97b..3d9ef589 100644 --- a/src/easyscience/fitting/sampler.py +++ b/src/easyscience/fitting/sampler.py @@ -16,6 +16,8 @@ from easyscience import global_object from .engine_base import PARAMETER_PREFIX +from .reshaping import inject_x_multi +from .reshaping import reshape_datasets if TYPE_CHECKING: # avoid import cycles; only needed for type hints from bumps.dream.state import MCMCDraw @@ -165,7 +167,7 @@ def load_chain(path: str | os.PathLike, skip: int = 0) -> tuple[MCMCDraw, list[s """Reload a DREAM chain state saved by ``Sampler.save``. This is the standalone reader: unlike ``Sampler.load_state`` it needs no - fitter, model or data, so a saved chain can be inspected or post-processed + model or data, so a saved chain can be inspected or post-processed on a machine that does not have the model. Parameter names are restored from the sidecar when available (schema versions 1 and 2), falling back to the state's labels with the minimizer prefix stripped. @@ -254,28 +256,25 @@ class Sampler: effect on the sampler, and there are deliberately no setters: to sample different data, create a new ``Sampler``. - Construct directly with a configured ``Fitter`` (or ``MultiFitter``). - The only requirement is an installed ``bumps`` package. **Running a fit - first is not required**; sampling from the initial parameter values - works fine. - - It is often worth fitting first anyway. DREAM seeds its whole starting - population inside a tiny ball around the parameters' *current* values - (BUMPS' default ``init='eps'``), so sampling from fitted values starts the - chain in the right region and shortens the burn-in needed to reach the - typical set. From a poor initial guess, expect to burn for longer. + ``Sampler`` is a parallel entry point to ``Fitter``: it takes the same + ``(fit_object, fit_function)`` pair plus the data. Parameters ---------- - fitter : Fitter - A configured ``Fitter`` (or ``MultiFitter``) supplying the model and - fit function. + fit_object : object + The EasyScience model object holding the ``Parameter`` instances to + sample. For multiple datasets this is one object exposing all the + parameters (for example an ``EasyList`` of models, or the + ``fit_object`` of a ``MultiFitter``). + fit_function : Callable | list[Callable] + The model function, or one per dataset when ``x``, ``y`` and + ``weights`` are lists of arrays. x : np.ndarray | list[np.ndarray] - Independent variable array (or list of arrays for ``MultiFitter``). + Independent variable array (or list of arrays, one per dataset). y : np.ndarray | list[np.ndarray] - Dependent variable array (or list of arrays for ``MultiFitter``). + Dependent variable array (or list of arrays, one per dataset). weights : np.ndarray | list[np.ndarray] - Weight array (or list of arrays for ``MultiFitter``). Required: + Weight array (or list of arrays, one per dataset). Required: sampling has no default weighting, so a missing weight array is rejected here rather than deep inside the sampling engine. vectorized : bool, default=False @@ -289,14 +288,14 @@ class Sampler: Raises ------ TypeError - If ``fitter`` is not Fitter-shaped (no ``fit_function``), - if any dataset in ``x``/``y``/``weights`` is not a numeric array - (e.g. a string), or ``vectorized``/``sampler_kwargs`` have the wrong - type. + If ``fit_object`` has no ``get_fit_parameters``, ``fit_function`` + is not callable, any dataset in ``x``/``y``/``weights`` is not a + numeric array (e.g. a string), or ``vectorized``/``sampler_kwargs`` + have the wrong type. ValueError - If ``x``, ``y`` and ``weights`` do not hold matching structures - (all arrays, or lists of the same length), or any dataset is a - scalar or empty array. + If ``fit_function``, ``x``, ``y`` and ``weights`` do not hold + matching structures (all single, or lists of the same length), or + any dataset is a scalar or empty array. Notes ----- @@ -329,7 +328,7 @@ class Sampler: the whole chain:: sampler = Sampler( - fitter, x, y, weights=w, sampler_kwargs={'trim': False} + model, model, x, y, weights=w, sampler_kwargs={'trim': False} ) Note also that trimming does not survive a ``save()``/``load_state()`` @@ -340,29 +339,26 @@ class Sampler: def __init__( self, - fitter: 'Fitter', + fit_object: object, + fit_function: Callable | list[Callable], x: np.ndarray | list[np.ndarray], y: np.ndarray | list[np.ndarray], weights: np.ndarray | list[np.ndarray], vectorized: bool = False, sampler_kwargs: dict | None = None, ): - if not hasattr(fitter, 'fit_function'): - raise TypeError( - f'fitter must be a configured Fitter or MultiFitter, got {type(fitter).__name__}.' - ) - x_is_multi = isinstance(x, (list, tuple)) - if x_is_multi != isinstance(y, (list, tuple)): + is_multi = isinstance(x, (list, tuple)) + if is_multi != isinstance(y, (list, tuple)): raise ValueError('x and y must either both be arrays or both be lists of arrays.') - if x_is_multi and len(x) != len(y): + if is_multi and len(x) != len(y): raise ValueError( f'x and y must hold the same number of datasets, got {len(x)} and {len(y)}.' ) - if isinstance(weights, (list, tuple)) != x_is_multi: + if isinstance(weights, (list, tuple)) != is_multi: raise ValueError( 'weights must match the structure of x and y (array or list of arrays).' ) - if x_is_multi and len(weights) != len(x): + if is_multi and len(weights) != len(x): raise ValueError( f'weights must hold the same number of datasets as x and y, ' f'got {len(weights)} and {len(x)}.' @@ -370,44 +366,122 @@ def __init__( _validate_dataset_arrays('x', x) _validate_dataset_arrays('y', y) _validate_dataset_arrays('weights', weights) + if isinstance(fit_function, (list, tuple)) != is_multi: + raise ValueError( + 'fit_function must be a list of callables when x, y and weights are ' + 'lists of arrays, and a single callable otherwise.' + ) + if is_multi and len(fit_function) != len(x): + raise ValueError( + f'fit_function must hold one callable per dataset, ' + f'got {len(fit_function)} for {len(x)} datasets.' + ) + fit_functions = list(fit_function) if is_multi else [fit_function] + if not all(callable(f) for f in fit_functions): + raise TypeError('fit_function must be callable (or a list of callables).') + if not hasattr(fit_object, 'get_fit_parameters'): + raise TypeError( + f'fit_object must be an EasyScience model object exposing the parameters ' + f'to sample, got {type(fit_object).__name__}.' + ) if not isinstance(vectorized, bool): raise TypeError(f'vectorized must be a bool, got {type(vectorized).__name__}.') if sampler_kwargs is not None and not isinstance(sampler_kwargs, dict): raise TypeError( f'sampler_kwargs must be a dict or None, got {type(sampler_kwargs).__name__}.' ) - self._fitter = fitter - # Defensive copies, exposed read-only: mutating the caller's arrays - # (or the properties) cannot desynchronise the chain and the save() - # fingerprint from the data actually sampled. To sample different - # data, create a new Sampler. - self._x = _copy_data(x) - self._y = _copy_data(y) - self._weights = _copy_data(weights) + self._is_multi = is_multi + self._fit_object = fit_object + self._fit_functions = fit_functions + + self._x = _copy_data(x if is_multi else [x]) + self._y = _copy_data(y if is_multi else [y]) + self._weights = _copy_data(weights if is_multi else [weights]) self._vectorized = vectorized self._default_sampler_kwargs = dict(sampler_kwargs or {}) self._state: MCMCDraw | None = None # current chain state self._results: SamplingResults | None = None + @classmethod + def from_fitter( + cls, + fitter: Fitter, + x: np.ndarray | list[np.ndarray], + y: np.ndarray | list[np.ndarray], + weights: np.ndarray | list[np.ndarray], + vectorized: bool = False, + sampler_kwargs: dict | None = None, + ) -> Sampler: + """Build a ``Sampler`` from the model bound to an existing ``Fitter``. + + A convenience for the common fit-then-sample workflow: the sampler + takes the fitter's ``fit_object`` and fit function(s) and is + otherwise identical to one constructed directly. The fitter is not + retained or modified. A ``MultiFitter`` yields a multi-dataset + sampler, so ``x``, ``y`` and ``weights`` must then be lists of arrays. + + Parameters + ---------- + fitter : Fitter + A configured ``Fitter`` or ``MultiFitter``. It does not need to + have been fitted. + x : np.ndarray | list[np.ndarray] + Independent variable array (or list of arrays, one per dataset). + y : np.ndarray | list[np.ndarray] + Dependent variable array (or list of arrays, one per dataset). + weights : np.ndarray | list[np.ndarray] + Weight array (or list of arrays, one per dataset). + vectorized : bool, default=False + See ``Sampler``. + sampler_kwargs : dict | None, default=None + See ``Sampler``. + + Returns + ------- + Sampler + A sampler bound to the fitter's model and the given data. + """ + # A MultiFitter exposes one function per dataset as ``fit_functions``; + # a plain Fitter has a single ``fit_function``. + fit_function = getattr(fitter, 'fit_functions', None) or fitter.fit_function + return cls( + fitter.fit_object, + fit_function, + x, + y, + weights, + vectorized=vectorized, + sampler_kwargs=sampler_kwargs, + ) + + def _single_or_list(self, data: list): + """Return bound data the way it was passed in: one item or a list copy.""" + return list(data) if self._is_multi else data[0] + @property - def fitter(self) -> Fitter: - """The Fitter supplying the model and minimizer (read-only).""" - return self._fitter + def fit_object(self) -> object: + """The EasyScience model object holding the sampled parameters (read-only).""" + return self._fit_object + + @property + def fit_function(self) -> Callable | list[Callable]: + """The model function, or list of them for multiple datasets (read-only).""" + return self._single_or_list(self._fit_functions) @property def x(self) -> np.ndarray | list[np.ndarray]: """The bound independent variable data (read-only copy).""" - return list(self._x) if isinstance(self._x, list) else self._x + return self._single_or_list(self._x) @property def y(self) -> np.ndarray | list[np.ndarray]: """The bound dependent variable data (read-only copy).""" - return list(self._y) if isinstance(self._y, list) else self._y + return self._single_or_list(self._y) @property def weights(self) -> np.ndarray | list[np.ndarray]: """The bound weight data (read-only copy).""" - return list(self._weights) if isinstance(self._weights, list) else self._weights + return self._single_or_list(self._weights) @property def state(self) -> MCMCDraw | None: @@ -436,12 +510,7 @@ def logp(self) -> np.ndarray | None: def _fingerprint(self) -> str | None: """SHA-256 fingerprint of the bound (x, y, weights) data, or None.""" - x_list = list(self._x) if isinstance(self._x, (list, tuple)) else [self._x] - y_list = list(self._y) if isinstance(self._y, (list, tuple)) else [self._y] - w_list = ( - list(self._weights) if isinstance(self._weights, (list, tuple)) else [self._weights] - ) - return _data_fingerprint(x_list, y_list, w_list) + return _data_fingerprint(self._x, self._y, self._weights) def _run( self, @@ -468,19 +537,16 @@ def _run( ) from .samplers.sampler_bumps import DreamSampler - x_fit, x_new, y_new, w_new, dims = self._fitter._precompute_reshaping( + x_fit, x_new, y_new, w_new, dims = reshape_datasets( self._x, self._y, self._weights, self._vectorized ) - # The dims are passed explicitly so the fitter itself is never mutated. - wrapped = self._fitter._fit_function_wrapper(x_new, flatten=True, dependent_dims=dims) + wrapped = inject_x_multi(self._fit_functions, x_new, dims) merged_kwargs = {**self._default_sampler_kwargs, **(sampler_kwargs or {})} - # A fresh engine per run keeps the chain on the fitter's current fit - # function and parameters; chain continuity lives in ``resume_state``. - # This is where a sampler factory would plug in once there is more - # than one backend. - engine = DreamSampler(obj=self._fitter.fit_object, fit_function=wrapped) + # This is where a sampler factory would plug in once there + # is more than one backend. + engine = DreamSampler(obj=self._fit_object, fit_function=wrapped) result = engine.run( x=x_fit, y=y_new, @@ -678,7 +744,7 @@ def save(self, path: str | os.PathLike) -> None: ``.params.json`` sidecar with the parameter names, the easyscience version, and a fingerprint of the bound data (verified with a warning on ``load_state()``). Use ``load_chain`` to read the - files back without a fitter. + files back without the model. Parameters ---------- @@ -721,7 +787,7 @@ def save(self, path: str | os.PathLike) -> None: def load_state(self, path: str | os.PathLike, skip: int = 0) -> SamplingResults: """Load a previously saved chain into this sampler. - The sampler must be constructed with the same fitter and data used to + The sampler must be constructed with the same model and data used to create the chain — ``extend()`` then continues the saved chain. If the sidecar carries a data fingerprint and it does not match this sampler's bound data, a warning is logged (extending a chain against diff --git a/tests/integration/fitting/test_sampler.py b/tests/integration/fitting/test_sampler.py index c588763b..2bb51941 100644 --- a/tests/integration/fitting/test_sampler.py +++ b/tests/integration/fitting/test_sampler.py @@ -14,6 +14,7 @@ from easyscience import ObjBase from easyscience import Parameter +from easyscience.fitting import Fitter from easyscience.fitting import Sampler from easyscience.fitting import SamplingResults from easyscience.fitting.multi_fitter import MultiFitter @@ -62,11 +63,11 @@ def __call__(self, x): ) -def _fitter_and_data(): - """Build a 2-parameter MultiFitter over a small sine model. +def _model_and_data(): + """Build a 2-parameter sine model and a small dataset to sample. - The fitter keeps its default (LMFit) minimizer: sampling no longer - requires switching to BUMPS, only an installed ``bumps`` package. + No ``Fitter`` is involved: sampling needs only the model, its fit + function, the data, and an installed ``bumps`` package. """ pytest.importorskip('bumps') ref_sin = AbsSin(0.2, np.pi) @@ -76,18 +77,17 @@ def _fitter_and_data(): x = np.linspace(0, 5, 50) y = ref_sin(x) weights = np.ones_like(x) - f = MultiFitter([sp], [sp]) - return f, sp, x, y, weights + return sp, x, y, weights class TestSampler: - """Integration tests for ``Sampler(f, ...)`` / ``Sampler``.""" + """Integration tests for ``Sampler``.""" @pytest.mark.filterwarnings('ignore::UserWarning') def test_sample_returns_results_object(self): """sample() returns a populated SamplingResults, cached on the sampler.""" - f, sp, x, y, weights = _fitter_and_data() - sampler = Sampler(f, [x], [y], [weights]) + sp, x, y, weights = _model_and_data() + sampler = Sampler(sp, sp, x, y, weights) results = sampler.sample(samples=100, burn=20, thin=2) @@ -109,7 +109,7 @@ def test_sample_returns_results_object(self): @pytest.mark.filterwarnings('ignore::UserWarning') def test_sample_multi_dataset(self): - """Multi-dataset sampling via Sampler(f, ...) has correct param_names.""" + """Multi-dataset sampling has correct param_names.""" ref_sin_1 = AbsSin(0.2, np.pi) sp_sin_1 = AbsSin(0.354, 3.05) sp_line = Line(0.43, 6.1) @@ -130,9 +130,13 @@ def test_sample_multi_dataset(self): sp_line.c.fixed = False pytest.importorskip('bumps') + # Direct multi-dataset construction: one container object exposing all + # parameters (here the one MultiFitter builds) and one function per + # dataset. f = MultiFitter([sp_sin_1, sp_line], [sp_sin_1, sp_line]) - - sampler = Sampler(f, [x1, x2], [y1, y2], [weights, weights]) + sampler = Sampler( + f.fit_object, [sp_sin_1, sp_line], [x1, x2], [y1, y2], [weights, weights] + ) results = sampler.sample(samples=100, burn=20, thin=2) # All parameters across both models should appear @@ -142,8 +146,8 @@ def test_sample_multi_dataset(self): def test_sample_population(self): """Passing population should succeed and produce valid draws.""" - f, _, x, y, weights = _fitter_and_data() - sampler = Sampler(f, [x], [y], [weights]) + sp, x, y, weights = _model_and_data() + sampler = Sampler(sp, sp, x, y, weights) results = sampler.sample(samples=100, burn=20, thin=2, population=5) assert results.draws.shape[0] > 0 @@ -163,9 +167,7 @@ def test_sample_vectorized_2d(self): sp.phase.fixed = False pytest.importorskip('bumps') - f = MultiFitter([sp], [sp]) - - sampler = Sampler(f, [x2D], [y2D], [weights], vectorized=True) + sampler = Sampler(sp, sp, x2D, y2D, weights, vectorized=True) results = sampler.sample(samples=100, burn=20, thin=2) assert results.draws.ndim == 2 @@ -173,21 +175,24 @@ def test_sample_vectorized_2d(self): assert results.draws.shape[1] == len(results.param_names) @pytest.mark.filterwarnings('ignore::UserWarning') - def test_fit_function_restored_on_success(self): - """fit_function must be restored after a successful sample().""" - f, _, x, y, weights = _fitter_and_data() - sampler = Sampler(f, [x], [y], [weights]) - original_func = f.fit_function + def test_from_fitter_samples_and_leaves_fitter_untouched(self): + """from_fitter() on a plain, unfitted Fitter with its default LMFit + minimizer samples fine and never touches the fitter.""" + sp, x, y, weights = _model_and_data() + f = Fitter(sp, sp) + assert f.minimizer.package == 'lmfit' - sampler.sample(samples=100, burn=20, thin=2) - assert f.fit_function is original_func + sampler = Sampler.from_fitter(f, x, y, weights) + results = sampler.sample(samples=100, burn=20, thin=2) + + assert results.draws.shape[0] > 0 + assert f.fit_function is sp + assert f.minimizer.package == 'lmfit' @pytest.mark.filterwarnings('ignore::UserWarning') - def test_fit_function_untouched_multi_dataset(self): - """With 2+ datasets the per-dataset wrapping in MultiFitter must not - leave fit_function pointing at the LAST dataset's function after - sampling (regression: the single-dataset variant above is vacuous for - this bug because last == first == original).""" + def test_from_multi_fitter_untouched_multi_dataset(self): + """from_fitter() on a 2-dataset MultiFitter samples each dataset with + its own function and leaves the fitter's fit_function alone.""" ref_sin = AbsSin(0.2, np.pi) sp_sin = AbsSin(0.354, 3.05) sp_line = Line(0.43, 6.1) @@ -205,16 +210,18 @@ def test_fit_function_untouched_multi_dataset(self): original = f.fit_function assert original is sp_sin # two distinct per-dataset functions - sampler = Sampler(f, [x1, x2], [y1, y2], [weights, weights]) - sampler.sample(samples=50, burn=5, thin=1) + sampler = Sampler.from_fitter(f, [x1, x2], [y1, y2], [weights, weights]) + results = sampler.sample(samples=50, burn=5, thin=1) assert f.fit_function is original + assert sampler.fit_function == [sp_sin, sp_line] + assert results.draws.shape[0] > 0 @pytest.mark.filterwarnings('ignore::UserWarning') def test_sampler_kwargs_forwarded(self): """Per-call sampler_kwargs dict is forwarded to the BUMPS DREAM sampler.""" - f, _, x, y, weights = _fitter_and_data() - sampler = Sampler(f, [x], [y], [weights]) + sp, x, y, weights = _model_and_data() + sampler = Sampler(sp, sp, x, y, weights) results = sampler.sample(samples=100, burn=20, thin=2, sampler_kwargs={'init': 'random'}) @@ -226,8 +233,8 @@ def test_default_sampler_kwargs_merged(self, monkeypatch): """Constructor-level sampler_kwargs defaults are used; per-call kwargs win.""" from easyscience.fitting.samplers.sampler_bumps import DreamSampler - f, _, x, y, weights = _fitter_and_data() - sampler = Sampler(f, [x], [y], [weights], sampler_kwargs={'init': 'random'}) + sp, x, y, weights = _model_and_data() + sampler = Sampler(sp, sp, x, y, weights, sampler_kwargs={'init': 'random'}) captured = {} original_run = DreamSampler.run @@ -245,25 +252,11 @@ def spy(self, **kwargs): sampler.sample(samples=100, burn=20, thin=2, sampler_kwargs={'init': 'lhs'}) assert captured == {'init': 'lhs'} # per-call overrides default - @pytest.mark.filterwarnings('ignore::UserWarning') - def test_sample_with_lmfit_minimizer_active(self): - """Sampling works without switching the fitter's minimizer to BUMPS — - the new capability enabled by the ``DreamSampler`` engine (#280).""" - f, _, x, y, weights = _fitter_and_data() - assert f.minimizer.package == 'lmfit' # the default LMFit minimizer - - sampler = Sampler(f, [x], [y], [weights]) - results = sampler.sample(samples=100, burn=20, thin=2) - - assert results.draws.shape[0] > 0 - # The active minimizer is untouched by sampling. - assert f.minimizer.package == 'lmfit' - @pytest.mark.filterwarnings('ignore::UserWarning') def test_extend_chain(self): """extend(additional_samples=) continues the chain; ring-buffer math is done for the user.""" - f, _, x, y, weights = _fitter_and_data() - sampler = Sampler(f, [x], [y], [weights]) + sp, x, y, weights = _model_and_data() + sampler = Sampler(sp, sp, x, y, weights) first = sampler.sample(samples=100, burn=20, thin=1) n_first = first.draws.shape[0] @@ -282,8 +275,8 @@ def test_extend_with_thinning_keeps_existing_draws(self): generations (``Ngen * Npop``), not from the retained-draw count, which BUMPS divides by the thinning interval. """ - f, _, x, y, weights = _fitter_and_data() - sampler = Sampler(f, [x], [y], [weights]) + sp, x, y, weights = _model_and_data() + sampler = Sampler(sp, sp, x, y, weights) first = sampler.sample(samples=1000, burn=20, thin=10) n_first = first.draws.shape[0] @@ -296,8 +289,8 @@ def test_extend_with_thinning_keeps_existing_draws(self): @pytest.mark.filterwarnings('ignore::UserWarning') def test_extend_total_samples_override(self): """extend(total_samples=) bypasses the additional_samples arithmetic.""" - f, _, x, y, weights = _fitter_and_data() - sampler = Sampler(f, [x], [y], [weights]) + sp, x, y, weights = _model_and_data() + sampler = Sampler(sp, sp, x, y, weights) sampler.sample(samples=100, burn=20, thin=1) extended = sampler.extend(total_samples=150, thin=1) @@ -314,14 +307,14 @@ def test_extend_after_save_load_roundtrip(self, tmp_path, caplog): """ import logging - f, _, x, y, weights = _fitter_and_data() - sampler = Sampler(f, [x], [y], [weights]) + sp, x, y, weights = _model_and_data() + sampler = Sampler(sp, sp, x, y, weights) first = sampler.sample(samples=100, burn=20, thin=1) prefix = str(tmp_path / 'chain') sampler.save(prefix) - sampler2 = Sampler(f, [x], [y], [weights]) + sampler2 = Sampler(sp, sp, x, y, weights) loaded = sampler2.load_state(prefix) assert loaded.draws.shape[1] == first.draws.shape[1] @@ -340,8 +333,8 @@ def test_extend_preserves_nondefault_population(self): saved state on resume, otherwise BUMPS regenerates the default population and raises ``Cannot change Nvar, Npop or Ncr on resize``. """ - f, _, x, y, weights = _fitter_and_data() - sampler = Sampler(f, [x], [y], [weights]) + sp, x, y, weights = _model_and_data() + sampler = Sampler(sp, sp, x, y, weights) first = sampler.sample(samples=100, burn=20, thin=1, population=5) first_npop = first.state.Npop @@ -359,8 +352,8 @@ def test_save_warns_when_fingerprint_unavailable(self, tmp_path, caplog, monkeyp logs a warning and records ``null`` in the sidecar.""" import logging - f, _, x, y, weights = _fitter_and_data() - sampler = Sampler(f, [x], [y], [weights]) + sp, x, y, weights = _model_and_data() + sampler = Sampler(sp, sp, x, y, weights) sampler.sample(samples=100, burn=20, thin=2) monkeypatch.setattr('easyscience.fitting.sampler._data_fingerprint', lambda *args: None) @@ -376,14 +369,14 @@ def test_save_warns_when_fingerprint_unavailable(self, tmp_path, caplog, monkeyp @pytest.mark.filterwarnings('ignore::UserWarning') def test_load_state_populates_results(self, tmp_path): """A freshly loaded sampler reports draws/logp/param_names without resampling.""" - f, sp, x, y, weights = _fitter_and_data() - sampler = Sampler(f, [x], [y], [weights]) + sp, x, y, weights = _model_and_data() + sampler = Sampler(sp, sp, x, y, weights) first = sampler.sample(samples=100, burn=20, thin=2) prefix = str(tmp_path / 'chain') sampler.save(prefix) - sampler2 = Sampler(f, [x], [y], [weights]) + sampler2 = Sampler(sp, sp, x, y, weights) assert sampler2.draws is None loaded = sampler2.load_state(prefix) @@ -409,22 +402,22 @@ def test_load_short_chain_regression(self, tmp_path): reader collapses it to a 1-D array and ``load_state`` raises ``IndexError`` without the 2-D coercion workaround. """ - f, _, x, y, weights = _fitter_and_data() - sampler = Sampler(f, [x], [y], [weights]) + sp, x, y, weights = _model_and_data() + sampler = Sampler(sp, sp, x, y, weights) sampler.sample(samples=20, burn=5, thin=1) prefix = str(tmp_path / 'short_chain') sampler.save(prefix) - sampler2 = Sampler(f, [x], [y], [weights]) + sampler2 = Sampler(sp, sp, x, y, weights) loaded = sampler2.load_state(prefix) assert loaded.draws.shape[0] > 0 @pytest.mark.filterwarnings('ignore::UserWarning') def test_load_fingerprint_mismatch_warns(self, tmp_path, caplog): """Loading a chain into a sampler bound to different data warns.""" - f, _, x, y, weights = _fitter_and_data() - sampler = Sampler(f, [x], [y], [weights]) + sp, x, y, weights = _model_and_data() + sampler = Sampler(sp, sp, x, y, weights) sampler.sample(samples=100, burn=20, thin=2) prefix = str(tmp_path / 'chain') @@ -433,7 +426,7 @@ def test_load_fingerprint_mismatch_warns(self, tmp_path, caplog): import logging other_y = y + 0.5 - sampler2 = Sampler(f, [x], [other_y], [weights]) + sampler2 = Sampler(sp, sp, x, other_y, weights) with caplog.at_level(logging.WARNING, logger='easyscience.fitting'): sampler2.load_state(prefix) assert 'does not match the data fingerprint' in caplog.text diff --git a/tests/unit/fitting/test_multi_fitter.py b/tests/unit/fitting/test_multi_fitter.py index 6c7401df..f47dfcb7 100644 --- a/tests/unit/fitting/test_multi_fitter.py +++ b/tests/unit/fitting/test_multi_fitter.py @@ -168,21 +168,6 @@ def test_fit_function_restored_with_multiple_datasets(self): expected = np.hstack([fit_objects[0](x[0]), fit_objects[1](x[1])]) assert np.allclose(y, expected) - def test_explicit_dependent_dims_do_not_touch_fitter(self): - """Passing ``dependent_dims`` (as the ``Sampler`` does) slices the - combined output without writing ``_dependent_dims`` onto the fitter.""" - fit_objects = [Line(1.0, 0.5), Line(2.0, 1.5)] - mf = MultiFitter(fit_objects, fit_objects) - assert mf._dependent_dims is None - - x = [np.array([0.0, 1.0, 2.0]), np.array([0.0, 1.0])] - wrapped = mf._fit_function_wrapper(x, flatten=True, dependent_dims=[(3,), (2,)]) - - y = wrapped(np.zeros(5)) - expected = np.hstack([fit_objects[0](x[0]), fit_objects[1](x[1])]) - assert np.allclose(y, expected) - assert mf._dependent_dims is None - # =================================================================== # MultiFitter._precompute_reshaping with weights=None diff --git a/tests/unit/fitting/test_reshaping.py b/tests/unit/fitting/test_reshaping.py new file mode 100644 index 00000000..aada7c91 --- /dev/null +++ b/tests/unit/fitting/test_reshaping.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause +"""Unit tests for ``reshaping.py`` — the data reshaping and fit-function +wrapping shared by ``Fitter``, ``MultiFitter`` and ``Sampler``.""" + +import numpy as np +import pytest + +from easyscience.fitting.reshaping import inject_x +from easyscience.fitting.reshaping import inject_x_multi +from easyscience.fitting.reshaping import reshape_dataset +from easyscience.fitting.reshaping import reshape_datasets + + +def _grid(n: int, m: int): + """A vectorized ``(n, m, 2)`` coordinate grid and its ``(n, m)`` dependent values.""" + X, Y = np.meshgrid(np.linspace(0, 1, m), np.linspace(0, 1, n)) + x = np.stack((X, Y), axis=2) + return x, X + 2 * Y + + +def _sum_xy(x): + return x[..., 0] + 2 * x[..., 1] + + +class TestReshapeDataset: + def test_1d_dims_are_y_shape(self): + x = np.linspace(0, 1, 5) + x_fit, x_new, y_new, w_new, dims = reshape_dataset(x, 2 * x, np.ones(5), vectorized=False) + assert dims == (5,) + assert x_fit.shape == y_new.shape == w_new.shape == (5,) + + def test_vectorized_dims_are_y_shape_not_x_shape(self): + """For multi-dimensional coordinates the dependent dims must exclude + the coordinate-component axis: ``(2, 3, 2)`` x holds 6 observations.""" + x, y = _grid(2, 3) + _, x_new, y_new, _, dims = reshape_dataset(x, y, None, vectorized=True) + assert dims == (2, 3) + assert y_new.shape == (6,) + assert x_new.shape == (2, 3, 2) + + def test_non_vectorized_nd_dims_are_y_shape(self): + x = np.random.default_rng(0).random((6, 2)) + y = np.arange(6.0) + _, x_new, y_new, _, dims = reshape_dataset(x, y, None, vectorized=False) + assert dims == (6,) + assert x_new.shape == (6, 2) + + +class TestInjectX: + def test_injects_real_x_and_flattens(self): + x, y = _grid(2, 3) + wrapped = inject_x(_sum_xy, x, flatten=True) + np.testing.assert_allclose(wrapped(np.zeros(6)), y.flatten()) + + +class TestInjectXMulti: + def test_two_multidimensional_datasets_are_sliced_by_observation_count(self): + """Regression: slicing by the product of the x shape allocated twice as + many output positions as observations for vectorized grids, which + raised on the first dataset and silently clipped on the last.""" + x1, y1 = _grid(2, 3) + x2, y2 = _grid(3, 4) + x_fit, x_new, y_new, _, dims = reshape_datasets([x1, x2], [y1, y2], None, vectorized=True) + assert dims == [(2, 3), (3, 4)] + assert y_new.shape == (18,) + + wrapped = inject_x_multi([_sum_xy, _sum_xy], x_new, dims) + np.testing.assert_allclose(wrapped(x_fit), np.hstack([y1.flatten(), y2.flatten()])) + + def test_multidimensional_dataset_first_then_1d(self): + """Dataset boundaries hold regardless of order; a 2D dataset that is + not last used to raise a broadcast error.""" + x1, y1 = _grid(2, 3) + x2 = np.linspace(0, 1, 4) + y2 = 3 * x2 + x_fit, x_new, y_new, _, dims = reshape_datasets([x1, x2], [y1, y2], None, vectorized=True) + assert dims == [(2, 3), (4,)] + + wrapped = inject_x_multi([_sum_xy, lambda x: 3 * x], x_new, dims) + np.testing.assert_allclose(wrapped(x_fit), np.hstack([y1.flatten(), y2])) + + def test_weights_none_for_all_datasets(self): + x = np.linspace(0, 1, 3) + _, _, _, w_new, _ = reshape_datasets([x, x], [x, x], [None, None], vectorized=False) + assert w_new is None + + def test_weights_are_concatenated(self): + x = np.linspace(0, 1, 3) + _, _, _, w_new, _ = reshape_datasets( + [x, x], [x, x], [np.ones(3), 2 * np.ones(3)], vectorized=False + ) + np.testing.assert_array_equal(w_new, [1, 1, 1, 2, 2, 2]) + + +@pytest.mark.parametrize('vectorized', [False, True]) +def test_shape_mismatch_raises(vectorized): + x = np.linspace(0, 1, 4) + with pytest.raises(ValueError, match='shape of the x and y data must be the same'): + reshape_dataset(x, np.zeros(3), None, vectorized=vectorized) diff --git a/tests/unit/fitting/test_sampler.py b/tests/unit/fitting/test_sampler.py index 47136b11..1e488945 100644 --- a/tests/unit/fitting/test_sampler.py +++ b/tests/unit/fitting/test_sampler.py @@ -15,6 +15,7 @@ from easyscience import ObjBase from easyscience import Parameter +from easyscience.fitting import Fitter from easyscience.fitting import Sampler from easyscience.fitting import SamplingResults from easyscience.fitting.engine_base import PARAMETER_PREFIX @@ -36,10 +37,15 @@ def __call__(self, x): return np.abs(np.sin(self.phase.value * x + self.offset.value)) -class _StubFitter: - """Duck-types the Fitter attributes checked by the Sampler constructor.""" +class _StubModel: + """Duck-types the model attribute checked by the Sampler constructor.""" - fit_function = None + def get_fit_parameters(self): + return [] + + +def _identity(x): + return x class _StubState: @@ -49,12 +55,8 @@ def __init__(self, labels): self.labels = list(labels) -def _fitter_and_data(): - """Build a 2-parameter MultiFitter over a small sine model. - - The fitter keeps its default (LMFit) minimizer: sampling no longer - requires switching to BUMPS, only an installed ``bumps`` package. - """ +def _model_and_data(): + """Build a 2-parameter sine model and a small dataset to sample.""" pytest.importorskip('bumps') ref_sin = AbsSin(0.2, np.pi) sp = AbsSin(0.354, 3.05) @@ -63,8 +65,7 @@ def _fitter_and_data(): x = np.linspace(0, 5, 50) y = ref_sin(x) weights = np.ones_like(x) - f = MultiFitter([sp], [sp]) - return f, sp, x, y, weights + return sp, x, y, weights def _xyw(): @@ -93,70 +94,109 @@ def _make_state(ngen=6, npop=5, nvar=2, seed=7): class TestSamplerConstructorValidation: - def test_rejects_fitter_without_fit_function(self): + def test_rejects_fit_object_without_fit_parameters(self): + x, y, w = _xyw() + with pytest.raises(TypeError, match='fit_object must be an EasyScience model'): + Sampler(object(), [_identity], [x], [y], [w]) + + def test_rejects_non_callable_fit_function(self): + x, y, w = _xyw() + with pytest.raises(TypeError, match='fit_function must be callable'): + Sampler(_StubModel(), 'not-callable', x, y, w) + with pytest.raises(TypeError, match='fit_function must be callable'): + Sampler(_StubModel(), [_identity, None], [x, x], [y, y], [w, w]) + + def test_rejects_fit_function_structure_mismatch(self): + """One callable for a list of datasets (or vice versa) is an error: + multi-dataset sampling takes one fit function per dataset.""" x, y, w = _xyw() - with pytest.raises(TypeError, match='fitter must be a configured Fitter'): - Sampler(object(), [x], [y], [w]) + with pytest.raises(ValueError, match='fit_function must be a list of callables'): + Sampler(_StubModel(), _identity, [x], [y], [w]) + with pytest.raises(ValueError, match='fit_function must be a list of callables'): + Sampler(_StubModel(), [_identity], x, y, w) + + def test_rejects_fit_function_count_mismatch(self): + x, y, w = _xyw() + with pytest.raises(ValueError, match='one callable per dataset'): + Sampler(_StubModel(), [_identity], [x, x], [y, y], [w, w]) + + def test_rejects_list_of_fit_objects(self): + """Multiple datasets take one container object exposing all the + parameters, not a bare list of models.""" + sp_1 = AbsSin(0.1, 1.0) + sp_2 = AbsSin(0.2, 2.0) + x, y, w = _xyw() + with pytest.raises(TypeError, match='fit_object must be an EasyScience model'): + Sampler([sp_1, sp_2], [sp_1, sp_2], [x, x], [y, y], [w, w]) def test_requires_weights(self): """Sampling has no default weighting, so weights are a required argument rather than a None that only blows up at sample().""" x, y, _ = _xyw() with pytest.raises(TypeError, match='weights'): - Sampler(_StubFitter(), [x], [y]) + Sampler(_StubModel(), [_identity], [x], [y]) def test_rejects_mixed_array_and_list(self): x, y, w = _xyw() with pytest.raises(ValueError, match='both be arrays or both be lists'): - Sampler(_StubFitter(), [x], y, [w]) + Sampler(_StubModel(), [_identity], [x], y, [w]) def test_rejects_dataset_count_mismatch(self): x, y, w = _xyw() with pytest.raises(ValueError, match='same number of datasets'): - Sampler(_StubFitter(), [x, x], [y], [w, w]) + Sampler(_StubModel(), [_identity, _identity], [x, x], [y], [w, w]) def test_rejects_weights_structure_mismatch(self): x, y, w = _xyw() with pytest.raises(ValueError, match='weights must match the structure'): - Sampler(_StubFitter(), [x], [y], w) + Sampler(_StubModel(), [_identity], [x], [y], w) def test_rejects_weights_count_mismatch(self): x, y, w = _xyw() with pytest.raises(ValueError, match='weights must hold the same number'): - Sampler(_StubFitter(), [x], [y], [w, w]) + Sampler(_StubModel(), [_identity], [x], [y], [w, w]) def test_rejects_non_bool_vectorized(self): x, y, w = _xyw() with pytest.raises(TypeError, match='vectorized must be a bool'): - Sampler(_StubFitter(), [x], [y], [w], vectorized=1) + Sampler(_StubModel(), [_identity], [x], [y], [w], vectorized=1) def test_rejects_non_dict_sampler_kwargs(self): x, y, w = _xyw() with pytest.raises(TypeError, match='sampler_kwargs must be a dict'): - Sampler(_StubFitter(), [x], [y], [w], sampler_kwargs=[('init', 'random')]) + Sampler(_StubModel(), [_identity], [x], [y], [w], sampler_kwargs=[('init', 'random')]) def test_accepts_single_arrays(self): x, y, w = _xyw() - sampler = Sampler(_StubFitter(), x, y, w) + sampler = Sampler(_StubModel(), _identity, x, y, w) assert sampler.results is None class TestSamplerDataBinding: def test_properties_expose_bound_data(self): x, y, w = _xyw() - f = _StubFitter() - sampler = Sampler(f, [x], [y], [w]) - assert sampler.fitter is f - np.testing.assert_array_equal(sampler.x[0], x) - np.testing.assert_array_equal(sampler.y[0], y) - np.testing.assert_array_equal(sampler.weights[0], w) + model = _StubModel() + sampler = Sampler(model, _identity, x, y, w) + assert sampler.fit_object is model + assert sampler.fit_function is _identity + np.testing.assert_array_equal(sampler.x, x) + np.testing.assert_array_equal(sampler.y, y) + np.testing.assert_array_equal(sampler.weights, w) + + def test_properties_keep_list_structure(self): + """Multi-dataset inputs come back as lists, in the order given.""" + x, y, w = _xyw() + sampler = Sampler(_StubModel(), [_identity, _identity], [x, 2 * x], [y, y], [w, w]) + assert sampler.fit_function == [_identity, _identity] + assert len(sampler.x) == 2 + np.testing.assert_array_equal(sampler.x[1], 2 * x) def test_inputs_are_copied(self): """Mutating the caller's arrays after construction must not change the bound data (nor the save() fingerprint derived from it).""" x, y, w = _xyw() y_original = y.copy() - sampler = Sampler(_StubFitter(), [x], [y], [w]) + sampler = Sampler(_StubModel(), [_identity], [x], [y], [w]) fingerprint_before = sampler._fingerprint() y[:] = 0.0 @@ -166,7 +206,7 @@ def test_inputs_are_copied(self): def test_bound_arrays_are_read_only(self): x, y, w = _xyw() - sampler = Sampler(_StubFitter(), x, y, w) + sampler = Sampler(_StubModel(), _identity, x, y, w) with pytest.raises(ValueError, match='read-only'): sampler.x[0] = 99.0 @@ -174,8 +214,8 @@ def test_data_properties_have_no_setters(self): """Bound data is deliberately immutable — sample new data with a new Sampler, so a chain can never be extended against different data.""" x, y, w = _xyw() - sampler = Sampler(_StubFitter(), [x], [y], [w]) - for name in ('fitter', 'x', 'y', 'weights'): + sampler = Sampler(_StubModel(), [_identity], [x], [y], [w]) + for name in ('fit_object', 'fit_function', 'x', 'y', 'weights'): with pytest.raises(AttributeError): setattr(sampler, name, None) @@ -183,20 +223,20 @@ def test_data_properties_have_no_setters(self): class TestSamplerPathValidation: def test_save_rejects_non_pathlike(self): x, y, w = _xyw() - sampler = Sampler(_StubFitter(), [x], [y], [w]) + sampler = Sampler(_StubModel(), [_identity], [x], [y], [w]) with pytest.raises(TypeError, match='path must be a str or os.PathLike'): sampler.save(123) def test_save_accepts_pathlike(self, tmp_path): """A Path object passes validation; the empty sampler then raises RuntimeError.""" x, y, w = _xyw() - sampler = Sampler(_StubFitter(), [x], [y], [w]) + sampler = Sampler(_StubModel(), [_identity], [x], [y], [w]) with pytest.raises(RuntimeError, match='No chain state to save'): sampler.save(tmp_path / 'chain') def test_load_state_rejects_non_pathlike(self): x, y, w = _xyw() - sampler = Sampler(_StubFitter(), [x], [y], [w]) + sampler = Sampler(_StubModel(), [_identity], [x], [y], [w]) with pytest.raises(TypeError, match='path must be a str or os.PathLike'): sampler.load_state(123) @@ -213,41 +253,28 @@ def test_load_chain_rejects_bad_skip(self, tmp_path, skip): class TestSamplerErrorPaths: def test_sample_requires_bumps_package(self, monkeypatch): """sample() must raise RuntimeError when the bumps package is not - installed — regardless of the active minimizer — and must not touch - the fitter.""" - sp = AbsSin(0.354, 3.05) - f = MultiFitter([sp], [sp]) - + installed.""" x, y, w = _xyw() - sampler = Sampler(f, [x], [y], [w]) - minimizer_before = f.minimizer + sampler = Sampler(_StubModel(), _identity, x, y, w) monkeypatch.setattr( 'easyscience.fitting.available_minimizers.bumps_engine_available', False ) with pytest.raises(RuntimeError, match='requires the bumps package'): sampler.sample(samples=10, burn=5, thin=1) - assert f.minimizer is minimizer_before - - def test_fitter_untouched_on_error(self): - """The fitter is never mutated by sampling, even when the engine - raises.""" - f, _, x, y, weights = _fitter_and_data() - sampler = Sampler(f, [x], [y], [weights]) - original_func = f.fit_function - minimizer_before = f.minimizer - - # Invalid `samples` is rejected by the engine (single source of - # validation). + + def test_engine_argument_errors_propagate(self): + """Invalid ``samples`` is rejected by the engine (single source of + validation) and surfaces unchanged.""" + sp, x, y, weights = _model_and_data() + sampler = Sampler(sp, sp, x, y, weights) + with pytest.raises(ValueError, match='samples must be a positive integer'): sampler.sample(samples=-1, burn=5, thin=1) - assert f.fit_function is original_func - assert f.minimizer is minimizer_before - def test_extend_requires_existing_state(self): """extend() before sample()/load_state() raises RuntimeError.""" x, y, w = _xyw() - sampler = Sampler(_StubFitter(), [x], [y], [w]) + sampler = Sampler(_StubModel(), [_identity], [x], [y], [w]) with pytest.raises(RuntimeError, match='No chain to extend'): sampler.extend(additional_samples=10) @@ -255,7 +282,7 @@ def test_extend_requires_existing_state(self): def test_save_raises_without_state(self, tmp_path): """save() before sample() raises RuntimeError.""" x, y, w = _xyw() - sampler = Sampler(_StubFitter(), [x], [y], [w]) + sampler = Sampler(_StubModel(), [_identity], [x], [y], [w]) with pytest.raises(RuntimeError, match='No chain state to save'): sampler.save(str(tmp_path / 'chain')) @@ -264,7 +291,7 @@ def test_sample_warns_when_replacing_existing_chain(self, monkeypatch, caplog): """sample() over an existing chain logs a replace warning; a fresh sampler does not.""" x, y, w = _xyw() - sampler = Sampler(_StubFitter(), [x], [y], [w]) + sampler = Sampler(_StubModel(), [_identity], [x], [y], [w]) dummy = SamplingResults( draws=np.zeros((1, 1)), param_names=['p'], logp=np.zeros(1), state=object() @@ -356,40 +383,40 @@ class TestSamplerConstructorDataValidation: def test_rejects_scalar_x(self): _, y, w = _xyw() with pytest.raises(ValueError, match='x must be an array of values, got a scalar'): - Sampler(_StubFitter(), 5.0, y, w) + Sampler(_StubModel(), _identity, 5.0, y, w) def test_rejects_scalar_dataset_in_list(self): x, y, w = _xyw() with pytest.raises(ValueError, match=r'y\[1\] must be an array of values'): - Sampler(_StubFitter(), [x, x], [y, 3.0], [w, w]) + Sampler(_StubModel(), [_identity, _identity], [x, x], [y, 3.0], [w, w]) def test_rejects_string_data(self): x, _, w = _xyw() with pytest.raises(TypeError, match='y must hold numeric values'): - Sampler(_StubFitter(), x, 'abc', w) + Sampler(_StubModel(), _identity, x, 'abc', w) def test_rejects_non_numeric_object_array(self): _, y, _ = _xyw() with pytest.raises(TypeError, match='x must hold numeric values'): - Sampler(_StubFitter(), np.array([{}, {}], dtype=object), y, np.ones(2)) + Sampler(_StubModel(), _identity, np.array([{}, {}], dtype=object), y, np.ones(2)) def test_rejects_empty_array(self): with pytest.raises(ValueError, match='x must not be empty'): - Sampler(_StubFitter(), np.array([]), np.array([]), np.array([])) + Sampler(_StubModel(), _identity, np.array([]), np.array([]), np.array([])) def test_rejects_ragged_dataset(self): with pytest.raises(TypeError, match=r'x\[0\] could not be converted'): - Sampler(_StubFitter(), [[1.0, [2.0, 3.0]]], [np.zeros(3)], [np.ones(3)]) + Sampler(_StubModel(), [_identity], [[1.0, [2.0, 3.0]]], [np.zeros(3)], [np.ones(3)]) def test_rejects_scalar_weights(self): x, y, _ = _xyw() with pytest.raises(ValueError, match='weights must be an array of values'): - Sampler(_StubFitter(), x, y, 2.0) + Sampler(_StubModel(), _identity, x, y, 2.0) def test_rejects_none_weight_entry(self): x, y, w = _xyw() with pytest.raises(TypeError, match=r'weights\[1\] must hold numeric values'): - Sampler(_StubFitter(), [x, x], [y, y], [w, None]) + Sampler(_StubModel(), [_identity, _identity], [x, x], [y, y], [w, None]) class TestDataFingerprint: @@ -398,7 +425,7 @@ def test_returns_none_on_unhashable_data(self): def test_fingerprint_of_single_arrays(self): x, y, w = _xyw() - sampler = Sampler(_StubFitter(), x, y, w) + sampler = Sampler(_StubModel(), _identity, x, y, w) assert isinstance(sampler._fingerprint(), str) @@ -407,7 +434,7 @@ class TestSamplerRunEngine: with the ``DreamSampler`` engine stubbed out.""" def test_run_stores_results_and_exposes_properties(self, monkeypatch): - f, _, x, y, weights = _fitter_and_data() + sp, x, y, weights = _model_and_data() from easyscience.fitting.samplers.sampler_bumps import DreamSampler canned = { @@ -424,10 +451,7 @@ def fake_run(self, **kwargs): monkeypatch.setattr(DreamSampler, 'run', fake_run) - sampler = Sampler(f, [x], [y], [weights], sampler_kwargs={'trim': False}) - original_func = f.fit_function - minimizer_before = f.minimizer - dims_before = f._dependent_dims + sampler = Sampler(sp, sp, x, y, weights, sampler_kwargs={'trim': False}) results = sampler.sample(samples=100, burn=10, thin=2, sampler_kwargs={'init': 'lhs'}) assert isinstance(results, SamplingResults) @@ -441,21 +465,12 @@ def fake_run(self, **kwargs): assert captured['samples'] == 100 assert captured['burn'] == 10 assert captured['resume_state'] is None - # The fitter is never mutated: a fresh engine gets the wrapped - # function directly, the active (LMFit) minimizer stays put, and the - # reshaping bookkeeping is passed to the wrapper rather than written - # onto the fitter. - assert f.fit_function is original_func - assert f.minimizer is minimizer_before - assert f._dependent_dims is dims_before - - def test_run_works_with_non_bumps_minimizer(self, monkeypatch): - """Sampling works with the default LMFit minimizer active — the - engine is constructed independently of the fitter's minimizer.""" - f, _, x, y, weights = _fitter_and_data() - from easyscience.fitting.samplers.sampler_bumps import DreamSampler - assert f.minimizer.package != 'bumps' # default is LMFit + def test_run_binds_engine_to_model_and_wrapped_function(self, monkeypatch): + """The engine is built from the sampler's own model and a wrapped + fit function; no Fitter or minimizer is involved.""" + sp, x, y, weights = _model_and_data() + from easyscience.fitting.samplers.sampler_bumps import DreamSampler constructed = {} original_init = DreamSampler.__init__ @@ -474,14 +489,48 @@ def spy_init(self, obj, fit_function): monkeypatch.setattr(DreamSampler, '__init__', spy_init) monkeypatch.setattr(DreamSampler, 'run', lambda self, **kwargs: dict(canned)) - sampler = Sampler(f, [x], [y], [weights]) + sampler = Sampler(sp, sp, x, y, weights) results = sampler.sample(samples=10, burn=0, thin=1) assert results.param_names == ['offset', 'phase'] - # The engine is bound to the fitter's model object and a wrapped - # fit function, not to the minimizer. - assert constructed['obj'] is f.fit_object + assert constructed['obj'] is sp assert callable(constructed['fit_function']) + assert constructed['fit_function'] is not sp + + +class TestSamplerFromFitter: + """``from_fitter`` mirrors direct construction and leaves the fitter alone.""" + + def test_from_plain_fitter(self): + sp, x, y, weights = _model_and_data() + f = Fitter(sp, sp) + + sampler = Sampler.from_fitter(f, x, y, weights, sampler_kwargs={'init': 'lhs'}) + + assert sampler.fit_object is sp + assert sampler.fit_function is sp + np.testing.assert_array_equal(sampler.x, x) + assert sampler._default_sampler_kwargs == {'init': 'lhs'} + + def test_from_multi_fitter(self): + sp_1 = AbsSin(0.1, 1.0) + sp_2 = AbsSin(0.2, 2.0) + x, y, w = _xyw() + f = MultiFitter([sp_1, sp_2], [sp_1, sp_2]) + + sampler = Sampler.from_fitter(f, [x, x], [y, y], [w, w]) + + # The fitter's container object exposes every model's parameters. + assert sampler.fit_object is f.fit_object + assert sampler.fit_function == [sp_1, sp_2] + assert f.fit_function is sp_1 # the fitter is untouched + + def test_multi_fitter_requires_list_data(self): + sp_1 = AbsSin(0.1, 1.0) + x, y, w = _xyw() + f = MultiFitter([sp_1], [sp_1]) + with pytest.raises(ValueError, match='fit_function must be a list of callables'): + Sampler.from_fitter(f, x, y, w) class TestSamplerExtendArithmetic: @@ -490,7 +539,7 @@ class TestSamplerExtendArithmetic: @staticmethod def _sampler_with_stub_state(monkeypatch, ngen=7, npop=5): x, y, w = _xyw() - sampler = Sampler(_StubFitter(), x, y, w) + sampler = Sampler(_StubModel(), _identity, x, y, w) sampler._state = SimpleNamespace(Ngen=ngen, Npop=npop) captured = {} dummy = SamplingResults( @@ -530,7 +579,7 @@ class TestSamplerPersistenceRoundTrip: @staticmethod def _sampler_with_state(): x, y, w = _xyw() - sampler = Sampler(_StubFitter(), x, y, w) + sampler = Sampler(_StubModel(), _identity, x, y, w) state = _make_state() _draw = state.draw() sampler._state = state @@ -595,7 +644,7 @@ def test_load_state_populates_results(self, tmp_path, caplog): sampler.save(prefix) x, y, w = _xyw() - fresh = Sampler(_StubFitter(), x, y, w) + fresh = Sampler(_StubModel(), _identity, x, y, w) with caplog.at_level(logging.WARNING, logger='easyscience.fitting'): results = fresh.load_state(prefix) @@ -613,7 +662,7 @@ def test_load_state_warns_on_different_data(self, tmp_path, caplog): sampler.save(prefix) x, y, w = _xyw() - other = Sampler(_StubFitter(), x, 2.0 * y, w) + other = Sampler(_StubModel(), _identity, x, 2.0 * y, w) with caplog.at_level(logging.WARNING, logger='easyscience.fitting'): other.load_state(prefix)