Skip to content

Make Sampler completely Fitter-agnostic #310

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rozyczko wants to merge 1 commit into
base: sampler-engine-structure-280
Choose a base branch
from
Open

Make Sampler completely Fitter-agnostic #310

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 34 additions & 42 deletions docs/docs/tutorials/fitting-bayesian.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand All @@ -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",
Expand All @@ -193,7 +196,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "36a0c4f4",
"id": "10",
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -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",
Expand All @@ -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",
Expand All @@ -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",
Expand All @@ -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",
"```"
]
},
{
Expand All @@ -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",
Expand All @@ -307,7 +299,7 @@
},
{
"cell_type": "markdown",
"id": "8766b170",
"id": "15",
"metadata": {},
"source": [
"## Convergence diagnostics\n",
Expand All @@ -323,7 +315,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "3c49ab6f",
"id": "16",
"metadata": {},
"outputs": [],
"source": [
Expand Down Expand Up @@ -359,7 +351,7 @@
},
{
"cell_type": "markdown",
"id": "15",
"id": "17",
"metadata": {},
"source": [
"## Posterior summaries\n",
Expand All @@ -377,15 +369,15 @@
{
"cell_type": "code",
"execution_count": null,
"id": "ce3e38a8",
"id": "18",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "16",
"id": "19",
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -407,7 +399,7 @@
},
{
"cell_type": "markdown",
"id": "17",
"id": "20",
"metadata": {},
"source": [
"## Visualise the joint posterior\n",
Expand All @@ -418,7 +410,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "18",
"id": "21",
"metadata": {},
"outputs": [],
"source": [
Expand Down Expand Up @@ -452,7 +444,7 @@
},
{
"cell_type": "markdown",
"id": "19",
"id": "22",
"metadata": {},
"source": [
"## Posterior-predictive band\n",
Expand All @@ -463,7 +455,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "20",
"id": "23",
"metadata": {},
"outputs": [],
"source": [
Expand Down Expand Up @@ -499,7 +491,7 @@
},
{
"cell_type": "markdown",
"id": "03339658",
"id": "24",
"metadata": {},
"source": [
"## Extend the chain and check convergence\n",
Expand Down Expand Up @@ -543,7 +535,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "293b140b",
"id": "25",
"metadata": {},
"outputs": [],
"source": [
Expand Down Expand Up @@ -573,7 +565,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "9ec4302c",
"id": "26",
"metadata": {},
"outputs": [],
"source": [
Expand Down Expand Up @@ -614,7 +606,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "b0f30be6",
"id": "27",
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -641,7 +633,7 @@
},
{
"cell_type": "markdown",
"id": "50b7213a",
"id": "28",
"metadata": {},
"source": [
"### What is Gelman-Rubin R-hat?\n",
Expand Down Expand Up @@ -675,7 +667,7 @@
{
"cell_type": "code",
"execution_count": null,
"id": "3449e0a7",
"id": "29",
"metadata": {},
"outputs": [],
"source": [
Expand Down
93 changes: 6 additions & 87 deletions src/easyscience/fitting/fitter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -231,42 +233,24 @@ 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
----------
real_x : np.ndarray | None, default=None
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:
Expand Down Expand Up @@ -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(
Expand Down
Loading