Skip to content

issues from the PR addressed. Bigger rewrite, so separate branch #260

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

Draft
rozyczko wants to merge 2 commits into
base: bayesian_extend
Choose a base branch
from
Draft

issues from the PR addressed. Bigger rewrite, so separate branch #260

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e9d9130
issues from the PR addressed. Big rewrite, so separate branch
rozyczko Jun 11, 2026
f8e3929
actualy refactor out sampling methods
rozyczko Jun 11, 2026
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
164 changes: 31 additions & 133 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 `Fitter.mcmc_sample` method. 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. Under the hood this uses BUMPS' DREAM sampler, so the underlying minimizer must be switched to BUMPS.\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 All @@ -45,7 +45,7 @@
"\n",
"The trade-off is computational cost: MCMC requires thousands of model evaluations, whereas an MLE fit may converge in dozens. For the simple 4-parameter model used here the difference is negligible, but for expensive models it is worth starting with MLE and only switching to MCMC when you need the richer output.\n",
"\n",
"In this tutorial we re-use the QENS dataset and the Lorentzian-with-resolution model from the [Fitting QENS](fitting-qens.ipynb) tutorial, but instead of returning a single best-fit value with a symmetric error bar we will draw thousands of samples from the posterior.\n"
"In this tutorial we re-use the QENS dataset and the Lorentzian-with-resolution model from the [Fitting QENS](fitting-qens.ipynb) tutorial, but instead of returning a single best-fit value with a symmetric error bar we will draw thousands of samples from the posterior."
]
},
{
Expand Down Expand Up @@ -250,20 +250,20 @@
"\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, switch it to BUMPS, and call `Fitter.mcmc_sample`, which returns a dictionary with the following keys:\n",
"DREAM only works with the BUMPS minimizer. We reuse the ``mle_fitter`` created above, 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",
"\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;\n",
"- `param_names`: the unique names of the parameters, in the same column order as `draws`;\n",
"- `internal_bumps_object`: the underlying BUMPS `MCMCDraw` object, useful for advanced diagnostics;\n",
"- `state`: the underlying BUMPS `MCMCDraw` object, useful for advanced diagnostics;\n",
"- `logp`: the log-posterior of each retained sample.\n",
"\n",
"The key sampling parameters are:\n",
"\n",
"- `samples` (4000): the total number of posterior draws to generate across all chains;\n",
"- `samples` (10000): the total number of posterior draws to generate across all chains;\n",
"- `burn` (500): the number of initial *burn-in* iterations to discard — the sampler needs time to find the typical set of the posterior, and early samples are not representative;\n",
"- `thin` (2): the *thinning* interval — only every second sample is kept, which reduces autocorrelation between consecutive draws;\n",
"\n",
"First, we switch to the BUMPS minimizer:\n"
"First, we switch to the BUMPS minimizer:"
]
},
{
Expand All @@ -285,17 +285,13 @@
"metadata": {},
"outputs": [],
"source": [
"result = mle_fitter.mcmc_sample(\n",
" x=omega,\n",
" y=intensity_obs,\n",
" weights=1 / intensity_error,\n",
" samples=10000,\n",
" burn=500,\n",
" thin=2,\n",
")\n",
"from easyscience.fitting import Sampler\n",
"\n",
"sampler = Sampler(mle_fitter, omega, intensity_obs, weights=1 / intensity_error)\n",
"results = sampler.sample(samples=10000, burn=500, thin=2)\n",
"\n",
"print(f'Drew {result[\"draws\"].shape[0]} samples for {result[\"draws\"].shape[1]} parameters.')\n",
"print('parameters:', result['param_names'])"
"print(f'Drew {results.draws.shape[0]} samples for {results.draws.shape[1]} parameters.')\n",
"print('parameters:', results.param_names)"
]
},
{
Expand All @@ -320,12 +316,9 @@
"metadata": {},
"outputs": [],
"source": [
"draws = result['draws']\n",
"logp = result['logp']\n",
"if callable(logp):\n",
" _, logp = logp()\n",
" logp = logp.flatten()\n",
"name_to_col = {name: idx for idx, name in enumerate(result['param_names'])}\n",
"draws = results.draws\n",
"logp = results.logp\n",
"name_to_col = {name: idx for idx, name in enumerate(results.param_names)}\n",
"\n",
"\n",
"def column_for(parameter):\n",
Expand Down Expand Up @@ -367,7 +360,7 @@
"The MLE finds the single point that maximises the likelihood, while the Bayesian median is the central value of the *posterior* — which also accounts for the prior $p(\\theta)$. If a parameter's posterior is skewed (asymmetric), the median and the mode (which approximates the MLE) will not coincide. The table below shows both so you can spot any such differences.\n",
"\n",
"\n",
"Note that the columns of `result['draws']` are ordered by `result['param_names']` (which use the parameters' `unique_name`), so it is worth building a small helper to look up a column by friendly name.\n"
"Note that the columns of `results.draws` are ordered by `results.param_names` (which use the parameters' `unique_name`), so it is worth building a small helper to look up a column by friendly name."
]
},
{
Expand Down Expand Up @@ -493,100 +486,16 @@
"plt.show()"
]
},
{
"cell_type": "markdown",
"id": "bcadc094",
"metadata": {},
"source": [
"## Persist and reload a chain\n",
"\n",
"In a real analysis you may want to **save a running chain to disk** for later inspection, or to resume it across compute sessions.\n",
"\n",
"`easyscience` provides robust I/O via ``save_sampler_state(state, prefix)`` and ``load_sampler_state(prefix)``. Additionally, ``Fitter.mcmc_sample()`` accepts a ``resume_state`` parameter — pass the ``MCMCDraw`` object from a previous run (or one reloaded from disk via ``load_sampler_state``) and DREAM **continues** the saved chain instead of starting cold.\n",
"\n",
"**Ring-buffer contract (important!)** DREAM stores draws in a fixed-size ring buffer sized to *samples*. To add N new draws on top of an existing chain of M draws without losing the old ones, pass ``samples = M + N`` and ``burn = 0``:\n",
"\n",
"```python\n",
"extended = fitter.mcmc_sample(..., samples=M + N, burn=0,\n",
" resume_state=previous_state)\n",
"```\n",
"\n",
"Here we save the chain from the previous exercise to disk, reload it, and verify the round-trip preserved the draws."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b7ef160c",
"metadata": {},
"outputs": [],
"source": [
"import glob\n",
"import os\n",
"\n",
"from easyscience.fitting.minimizers import save_sampler_state\n",
"\n",
"# Save the chain from the previous exercise to disk.\n",
"# save_sampler_state writes several files using 'prefix' as a filename stem.\n",
"prefix = os.path.abspath('dream_chain_example')\n",
"save_sampler_state(result['internal_bumps_object'], prefix)\n",
"\n",
"saved_files = sorted(glob.glob(prefix + '-*'))\n",
"print(f'Saved {len(saved_files)} file(s):')\n",
"for f in saved_files:\n",
" print(f' {os.path.basename(f)}')"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3acaa5c6",
"metadata": {},
"outputs": [],
"source": [
"from easyscience.fitting.minimizers import load_sampler_state\n",
"\n",
"# Reload the saved state from disk.\n",
"reloaded = load_sampler_state(prefix)\n",
"\n",
"# Compare the draws before and after the round-trip. We draw with\n",
"# portion=1.0 on both so the comparison uses the full chain: an in-memory\n",
"# state may carry a convergence-trimmed `portion` (so draw() returns only a\n",
"# tail of the samples), whereas a state reloaded from disk always reports\n",
"# portion=1.0. The underlying chain data is identical either way.\n",
"_draw_orig = result['internal_bumps_object'].draw(portion=1.0)\n",
"_draw_reloaded = reloaded.draw(portion=1.0)\n",
"\n",
"print('Draws match:', np.allclose(_draw_orig.points, _draw_reloaded.points))\n",
"print('logp match:', np.allclose(_draw_orig.logp, _draw_reloaded.logp))\n",
"print('Nvar (params):', reloaded.Nvar)\n",
"print('Npop (pop. size):', reloaded.Npop)\n",
"print('Ncr:', reloaded.Ncr)\n",
"print('Labels:', reloaded.labels)\n",
"\n",
"# The reloaded state is a fully functional MCMCDraw that can be passed\n",
"# as resume_state= to MultiFitter.sample() to continue the chain.\n",
"# (Note: BUMPS' DREAM resume has a known limitation with thin>1 and\n",
"# outlier removal. load_sampler_state also works around a BUMPS >=1.0.4\n",
"# regression where a short chain's single-row buffers are read back as 1-D\n",
"# arrays; prefer it over calling bumps.dream.state.load_state directly.)"
]
},
{
"cell_type": "markdown",
"id": "03339658",
"metadata": {},
"source": [
"## Extend the chain and check convergence\n",
"\n",
"The original run used ``samples=10000`` (5000 retained after thinning). Here we **extend the chain by 5000 more raw samples** using ``resume_state`` — DREAM continues from the saved population instead of starting cold.\n",
"\n",
"**Ring-buffer contract:** DREAM stores draws in a fixed-size ring buffer sized to *samples* (raw samples before thinning). To keep everything, pass ``samples = old_raw + new_raw`` and ``burn=0``:\n",
"The original run used ``samples=10000`` (5000 retained after thinning). Here we **extend the chain by 5000 more raw samples** using ``sampler.extend()`` — DREAM continues from the previous in-memory state instead of starting cold.\n",
"\n",
"```python\n",
"extended = fitter.mcmc_sample(..., samples=old_raw + new_raw, burn=0,\n",
" resume_state=previous_state)\n",
"```\n",
"``sampler.extend(additional_samples=5000)`` does the ring-buffer arithmetic for you automatically: it reads the number of stored generations from the saved state and sizes the new buffer to ``old_generations + additional_samples`` so no existing draws are lost.\n",
"\n",
"After the extension we compare the posterior summaries and check convergence with Gelman-Rubin R-hat."
]
Expand All @@ -599,25 +508,14 @@
"outputs": [],
"source": [
"# Extend the existing chain by 5000 more raw samples.\n",
"# samples = original_raw + new_raw = 10000 + 5000 = 15000 raw\n",
"# burn=0 because the chain is already warm.\n",
"total_raw = 10000 + 5000\n",
"\n",
"extended = mle_fitter.mcmc_sample(\n",
" x=omega,\n",
" y=intensity_obs,\n",
" weights=1 / intensity_error,\n",
" samples=total_raw,\n",
" burn=0,\n",
" thin=2,\n",
" resume_state=result['internal_bumps_object'], # continue from where we left off\n",
")\n",
"# extend() handles the ring-buffer arithmetic automatically.\n",
"extended_results = sampler.extend(additional_samples=5000, thin=2)\n",
"\n",
"short_draws = result['draws']\n",
"extended_draws = extended['draws']\n",
"short_draws = results.draws\n",
"extended_draws = extended_results.draws\n",
"\n",
"print(f'Short chain: {short_draws.shape[0]} retained draws')\n",
"print(f'Extended chain: {extended_draws.shape[0]} retained draws (requested {total_raw} raw)')"
"print(f'Extended chain: {extended_draws.shape[0]} retained draws')"
]
},
{
Expand All @@ -628,16 +526,16 @@
"outputs": [],
"source": [
"# Build helpers to look up parameter columns.\n",
"name_to_col_ext = {name: idx for idx, name in enumerate(extended['param_names'])}\n",
"name_to_col_orig = {name: idx for idx, name in enumerate(result['param_names'])}\n",
"name_to_col_ext = {name: idx for idx, name in enumerate(extended_results.param_names)}\n",
"name_to_col_orig = {name: idx for idx, name in enumerate(results.param_names)}\n",
"\n",
"\n",
"def col_orig(par):\n",
" return result['draws'][:, name_to_col_orig[par.unique_name]]\n",
" return results.draws[:, name_to_col_orig[par.unique_name]]\n",
"\n",
"\n",
"def col_ext(par):\n",
" return extended['draws'][:, name_to_col_ext[par.unique_name]]\n",
" return extended_results.draws[:, name_to_col_ext[par.unique_name]]\n",
"\n",
"\n",
"# Side-by-side posterior summary: compare the first 5000 draws (before extension)\n",
Expand Down Expand Up @@ -696,16 +594,16 @@
"# Convergence diagnostic: Gelman-Rubin R-hat from the extended chain state.\n",
"# Values close to 1.0 indicate good convergence.\n",
"print('Gelman-Rubin R-hat (extended chain) — values < 1.05 indicate convergence:')\n",
"rhat = extended['internal_bumps_object'].gelman()\n",
"for name, val in zip(extended['param_names'], rhat):\n",
"rhat = extended_results.state.gelman()\n",
"for name, val in zip(extended_results.param_names, rhat):\n",
" status = '✓' if val < 1.05 else '?' if val < 1.1 else '✗'\n",
" print(f' {name:<20s} {val:.4f} {status}')"
]
}
],
"metadata": {
"kernelspec": {
"display_name": ".venv (3.11.9)",
"display_name": "era",
"language": "python",
"name": "python3"
},
Expand All @@ -719,7 +617,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.9"
"version": "3.12.11"
}
},
"nbformat": 4,
Expand Down
4 changes: 3 additions & 1 deletion src/easyscience/fitting/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,10 @@
from .available_minimizers import AvailableMinimizers
from .fitter import Fitter
from .minimizers.utils import FitResults
from .sampler import Sampler
from .sampler import SamplingResults

# Causes circular import
# from .multi_fitter import MultiFitter # noqa: F401, E402

all = [AvailableMinimizers, Fitter, FitResults]
all = [AvailableMinimizers, Fitter, FitResults, Sampler, SamplingResults]
Loading
Loading