From a8946cea84f192f79365fc5c398e12bcfc03d058 Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Thu, 13 Aug 2026 10:18:29 +0200 Subject: [PATCH 1/3] Add Bayesian posterior sampling to Analysis1d Expose the EasyScience Fitter on Analysis1d and add MCMC posterior sampling on top of it, using the BUMPS DREAM sampler introduced in easyscience 2.5.1 (easyscience.fitting.Sampler). Least-squares fitting reports a single point with a curvature-derived uncertainty, which is only trustworthy when parameters are uncorrelated and roughly Gaussian. Sampling maps the whole posterior instead, so correlated and skewed parameters get honest credible intervals. The sampling machinery lives in a mixin with three hooks (build the fitter, bind the data, list the chain parameters) so that Analysis and ParameterAnalysis can reuse it. ParameterAnalysis is not an AnalysisBase and builds a MultiFitter over binding models rather than over itself, so a shared base class would not have worked. Notable details: - fit() now uses a cached Fitter instead of building one per call, and the cache is invalidated through the existing dirty-flag pattern. - Bounds are the prior in DREAM, so sampling refuses to run with any infinite bound. suggest_bounds() proposes finite ones from the fitted values and uncertainties; it is advisory until .apply() is called and never loosens a bound that is already finite, so physical limits survive. A zero-width suggestion is flagged rather than invented. - Sampling restores parameter values afterwards, since BUMPS leaves them wherever the last likelihood evaluation put them. - Chains are reported under Parameter.name, not the internal unique_name. Those names are per-session, so save_chain() writes a sidecar mapping them to stable names and load_chain() uses it; loading without one warns rather than mislabelling the columns. - After sampling, a warning fires when the posterior has piled up against a bound, which catches both bounds that are too tight and degenerate parameters that drift until a bound stops them. - BUMPS crashes with a bare IndexError inside its own outlier removal when chains scatter, which in practice means a degenerate model. That is re-raised with the likely cause and a workaround. Co-Authored-By: Claude Opus 5 (1M context) --- docs/docs/tutorials/bayesian.ipynb | 306 ++++++ docs/docs/tutorials/index.md | 3 + docs/mkdocs.yml | 1 + pixi.lock | 3 +- pyproject.toml | 23 +- src/easydynamics/analysis/__init__.py | 10 + src/easydynamics/analysis/analysis1d.py | 80 +- .../analysis/bayesian_sampling.py | 993 ++++++++++++++++++ src/easydynamics/analysis/posterior.py | 608 +++++++++++ src/easydynamics/utils/__init__.py | 11 +- src/easydynamics/utils/posterior_plotting.py | 256 +++++ .../fitting/test_bayesian_sampling.py | 207 ++++ .../analysis/test_analysis1d_bayesian.py | 486 +++++++++ .../easydynamics/analysis/test_posterior.py | 296 ++++++ .../utils/test_posterior_plotting.py | 148 +++ 15 files changed, 3407 insertions(+), 24 deletions(-) create mode 100644 docs/docs/tutorials/bayesian.ipynb create mode 100644 src/easydynamics/analysis/bayesian_sampling.py create mode 100644 src/easydynamics/analysis/posterior.py create mode 100644 src/easydynamics/utils/posterior_plotting.py create mode 100644 tests/integration/fitting/test_bayesian_sampling.py create mode 100644 tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py create mode 100644 tests/unit/easydynamics/analysis/test_posterior.py create mode 100644 tests/unit/easydynamics/utils/test_posterior_plotting.py diff --git a/docs/docs/tutorials/bayesian.ipynb b/docs/docs/tutorials/bayesian.ipynb new file mode 100644 index 000000000..44a415147 --- /dev/null +++ b/docs/docs/tutorials/bayesian.ipynb @@ -0,0 +1,306 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "eac0b8bb", + "metadata": {}, + "source": [ + "# Bayesian analysis\n", + "\n", + "Fitting with `fit()` finds the single set of parameter values that best matches the data, and reports an uncertainty derived from the curvature of $\\chi^2$ at that point. That uncertainty is only trustworthy when the parameters are uncorrelated and their uncertainties are close to Gaussian, which in QENS is often not the case.\n", + "\n", + "A **Bayesian** analysis answers a different question: instead of one best point, it maps out the whole *posterior distribution* over the parameters. From that you can read off credible intervals that stay honest when parameters are correlated or their distributions are skewed, and you can see the correlations directly.\n", + "\n", + "EasyDynamics does this with the DREAM sampler from [BUMPS](https://bumps.readthedocs.io/), through `sample_posterior()`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "02cb7aec", + "metadata": {}, + "outputs": [], + "source": [ + "import pooch\n", + "\n", + "import easydynamics as edyn\n", + "import easydynamics.sample_model as sm\n", + "from easydynamics.analysis.analysis1d import Analysis1d\n", + "\n", + "%matplotlib inline" + ] + }, + { + "cell_type": "markdown", + "id": "0499fea7", + "metadata": {}, + "source": [ + "## Load the data\n", + "\n", + "We use the same artificial vanadium measurement as the [Analysis 1D](analysis1d.ipynb) tutorial, and analyse a single Q slice." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fb407621", + "metadata": {}, + "outputs": [], + "source": [ + "vanadium_experiment = edyn.Experiment('Vanadium')\n", + "\n", + "file_path = pooch.retrieve(\n", + " url='https://github.com/easyscience/dynamics-lib/raw/refs/heads/master/docs/docs/tutorials/data/vanadium_data_example.h5',\n", + " known_hash='16cc1b327c303feeb88fb9dda5390dc4880b62396b1793f98c6fef0b27c7b873',\n", + ")\n", + "\n", + "vanadium_experiment.load_hdf5(filename=file_path)" + ] + }, + { + "cell_type": "markdown", + "id": "fcdfd395", + "metadata": {}, + "source": [ + "## Build the model and fit it\n", + "\n", + "As in [Tutorial 1](tutorial1_brownian.ipynb), a vanadium measurement is modelled with the Gaussian as the *sample*: what is being measured is the resolution function itself, so there is nothing to convolve it with.\n", + "\n", + "Sampling does not require a fit first, but it benefits from one: DREAM starts its chains in a small ball around the parameters' current values, so beginning from fitted values means less burn-in is needed before the chains reach the interesting region." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "de3297cf", + "metadata": {}, + "outputs": [], + "source": [ + "vanadium_components = sm.ComponentCollection()\n", + "vanadium_components.append_component(sm.Gaussian(width=0.1, area=1, name='Res. Gauss'))\n", + "\n", + "instrument_model = sm.InstrumentModel(\n", + " background_model=sm.BackgroundModel(components=sm.Polynomial(coefficients=[0.001])),\n", + ")\n", + "\n", + "analysis = Analysis1d(\n", + " display_name='Vanadium Analysis',\n", + " experiment=vanadium_experiment,\n", + " sample_model=sm.SampleModel(components=vanadium_components),\n", + " instrument_model=instrument_model,\n", + " Q_index=5,\n", + ")\n", + "\n", + "fit_result = analysis.fit()\n", + "print(f'reduced chi-squared = {fit_result.reduced_chi2:.4f}')" + ] + }, + { + "cell_type": "markdown", + "id": "b0a709f7", + "metadata": {}, + "source": [ + "## Bounds are the prior\n", + "\n", + "In DREAM, each parameter's `min` and `max` define a uniform prior, so **every free parameter must have finite bounds** before sampling. Most parameters start with at least one infinite bound, so `sample_posterior()` would refuse to run.\n", + "\n", + "`suggest_bounds()` proposes bounds from the fitted values and uncertainties. It is advisory: it changes nothing until you call `.apply()`, and it only ever fills in an *infinite* bound, so physical limits you have already set (an area that cannot go below zero, say) are left alone." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a004cd83", + "metadata": {}, + "outputs": [], + "source": [ + "suggestions = analysis.suggest_bounds()\n", + "print(suggestions)" + ] + }, + { + "cell_type": "markdown", + "id": "ab263f41", + "metadata": {}, + "source": [ + "The defaults are deliberately generous — 10 standard deviations plus 20% of the value. Because the bounds are a uniform prior, being too *narrow* is the dangerous mistake: it truncates the posterior and makes the uncertainty look smaller than it is. The 20% term is there for parameters whose fitted uncertainty comes back as zero. All three settings (`n_sigma`, `relative_pad`, `absolute_floor`) can be adjusted, and you can always set `min` and `max` by hand.\n", + "\n", + "It is worth reading the table before applying it. A suggestion many orders of magnitude larger than the parameter itself is a useful warning sign: it means the fit returned a huge uncertainty, which usually happens because two parameters are **degenerate** — the data determines only some combination of them, so one can grow while the other shrinks with no effect on the fit. That is a problem to fix in the model, not with the sampler." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1a92a899", + "metadata": {}, + "outputs": [], + "source": [ + "changed = suggestions.apply()\n", + "print(f'Applied bounds to: {[parameter.name for parameter in changed]}')" + ] + }, + { + "cell_type": "markdown", + "id": "8cf71e53", + "metadata": {}, + "source": [ + "## Sample the posterior\n", + "\n", + "`sample_posterior()` runs the chains. The three numbers that matter are:\n", + "\n", + "- `samples` — how many draws to collect in total. More is better, at linear cost.\n", + "- `burn` — generations discarded at the start, while the chains are still travelling towards the bulk of the posterior.\n", + "- `thin` — keep only every n-th generation, which reduces the correlation between neighbouring draws.\n", + "\n", + "Sampling never moves your parameters: their values are restored afterwards, so the model is left exactly as the fit left it." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "861fa264", + "metadata": {}, + "outputs": [], + "source": [ + "results = analysis.sample_posterior(samples=4000, burn=300, thin=2)\n", + "\n", + "print(f'Collected {results.draws.shape[0]} draws for {results.draws.shape[1]} parameters.')" + ] + }, + { + "cell_type": "markdown", + "id": "f2ce5232", + "metadata": {}, + "source": [ + "## Did the chains converge?\n", + "\n", + "Always look at the traces before trusting the numbers. A converged chain looks like a \"hairy caterpillar\": noisy, but flat and stationary. A visible drift or slow wander means the chain has not settled and needs a longer burn-in or more samples." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b2fe638c", + "metadata": {}, + "outputs": [], + "source": [ + "analysis.plot_trace()" + ] + }, + { + "cell_type": "markdown", + "id": "2b3ea7b4", + "metadata": {}, + "source": [ + "## Summarize the posterior\n", + "\n", + "`posterior_summary()` reports the median and the 68% credible interval of each parameter, under the parameter's own name and unit. The interval is asymmetric in general, which is precisely the information a single symmetric error bar throws away." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8470375e", + "metadata": {}, + "outputs": [], + "source": [ + "analysis.posterior_summary()" + ] + }, + { + "cell_type": "markdown", + "id": "60b4a448", + "metadata": {}, + "source": [ + "## Correlations between parameters\n", + "\n", + "The corner plot is the part least available from a least-squares fit. The diagonal shows each parameter's own distribution; each off-diagonal panel shows a pair. A round blob means the two are independent, while a tilted, narrow ridge means they are correlated and the data constrains only a combination of them." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c96509cc", + "metadata": {}, + "outputs": [], + "source": [ + "analysis.plot_corner()" + ] + }, + { + "cell_type": "markdown", + "id": "8f819c75", + "metadata": {}, + "source": [ + "## Does the model actually describe the data?\n", + "\n", + "The posterior predictive plot re-evaluates the model for a sample of posterior draws and shades the region they cover. If the data wanders outside the band in a systematic way, the model is missing a feature, and no amount of parameter tuning will fix it." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cce95199", + "metadata": {}, + "outputs": [], + "source": [ + "analysis.plot_posterior_predictive(n_draws=100)" + ] + }, + { + "cell_type": "markdown", + "id": "5947da0f", + "metadata": {}, + "source": [ + "## Continuing and storing a chain\n", + "\n", + "If the traces suggest the chain needs to run longer, `extend_sampling()` continues the existing chain rather than starting over, so nothing already computed is thrown away." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "974e486b", + "metadata": {}, + "outputs": [], + "source": [ + "extended = analysis.extend_sampling(additional_samples=1000, thin=2)\n", + "print(f'Chain now holds {extended.draws.shape[0]} draws.')" + ] + }, + { + "cell_type": "markdown", + "id": "69793d31", + "metadata": {}, + "source": [ + "Chains are expensive, so they can be saved and reloaded with `analysis.save_chain(path)` and `analysis.load_chain(path)`. A reloaded chain can be summarized, plotted, or extended further, exactly like a fresh one." + ] + }, + { + "cell_type": "markdown", + "id": "a3448aee", + "metadata": {}, + "source": [ + "## Things to watch out for\n", + "\n", + "**Data without uncertainties.** If your data carries no variances, the weights fall back to 1, which means the sampler assumes a noise level of 1 in whatever units the intensity happens to be. Least-squares does not care, since that scale cancels out of the best-fit position, but a posterior *does*: its width scales directly with the assumed noise, so the credible intervals will be wrong by whatever factor the true noise differs from 1. Bayesian analysis is not a way to avoid needing uncertainties on your data.\n", + "\n", + "**Sampling only some parameters.** `sample_posterior(parameters=[...])` restricts the chain to a subset, which is faster because the number of chains scales with the number of parameters. Be careful with the result: the other parameters are held *fixed*, which is not the same as averaging over them. The intervals you get are conditional on those fixed values, and will be too narrow whenever the parameters are correlated.\n", + "\n", + "**Degenerate parameters.** As seen above, they are a modelling problem rather than a sampling one. `suggest_bounds()` returning absurd values, or a warning that the posterior has piled up against its bounds, are both signs to go back and look at the model." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "default", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/docs/tutorials/index.md b/docs/docs/tutorials/index.md index 0617edb2a..23a36adc5 100644 --- a/docs/docs/tutorials/index.md +++ b/docs/docs/tutorials/index.md @@ -62,3 +62,6 @@ tutorials. - [Analysis](analysis.ipynb) - Learn how to fit a model to your data. - [Analysis 1D](analysis1d.ipynb) - Learn how to fit a model to your data at a particular Q. +- [Bayesian analysis](bayesian.ipynb) - Learn how to map out the full + posterior distribution of your parameters, including their + correlations, instead of a single best-fit point. diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index 64e94f967..3db04ace0 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -217,6 +217,7 @@ nav: - Experiment: tutorials/experiment.ipynb - Analysis: tutorials/analysis.ipynb - Analysis 1D: tutorials/analysis1d.ipynb + - Bayesian analysis: tutorials/bayesian.ipynb - API Reference: - API Reference: api-reference/index.md - analysis: api-reference/analysis.md diff --git a/pixi.lock b/pixi.lock index d1ebb0ed7..f07f02cf8 100644 --- a/pixi.lock +++ b/pixi.lock @@ -7535,12 +7535,13 @@ packages: name: easydynamics requires_dist: - darkdetect - - easyscience + - easyscience>=2.5.1 - ipykernel - ipympl - ipython - ipywidgets - jupyterlab + - matplotlib - pixi-kernel - plopp - pooch diff --git a/pyproject.toml b/pyproject.toml index b39261f87..3330304ea 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,17 +23,18 @@ classifiers = [ ] requires-python = '>=3.12' dependencies = [ - 'easyscience', # The base library of the EasyScience framework - 'pooch', # Data downloader - 'darkdetect', # Detecting dark mode (system-level) - 'plopp', # Plotting library - 'jupyterlab', # Jupyter notebooks - 'pixi-kernel', # Pixi Jupyter kernel - 'ipykernel', # Jupyter kernel (required for running notebooks) - 'ipywidgets', # Widgets (needed for interactive matplotlib backends) - 'ipympl', # Matplotlib Jupyter widget backend (%matplotlib widget) - 'IPython', # Interactive Python shell - 'sympy', # Symbolic mathematics (used for expression components) + 'easyscience>=2.5.1', # The base library of the EasyScience framework. 2.5.1 adds fitting.Sampler + 'matplotlib', # Plotting (posterior trace, corner, and predictive plots) + 'pooch', # Data downloader + 'darkdetect', # Detecting dark mode (system-level) + 'plopp', # Plotting library + 'jupyterlab', # Jupyter notebooks + 'pixi-kernel', # Pixi Jupyter kernel + 'ipykernel', # Jupyter kernel (required for running notebooks) + 'ipywidgets', # Widgets (needed for interactive matplotlib backends) + 'ipympl', # Matplotlib Jupyter widget backend (%matplotlib widget) + 'IPython', # Interactive Python shell + 'sympy', # Symbolic mathematics (used for expression components) ] [project.optional-dependencies] diff --git a/src/easydynamics/analysis/__init__.py b/src/easydynamics/analysis/__init__.py index 289ec02f5..89a45cea3 100644 --- a/src/easydynamics/analysis/__init__.py +++ b/src/easydynamics/analysis/__init__.py @@ -2,9 +2,19 @@ # SPDX-License-Identifier: BSD-3-Clause from easydynamics.analysis.analysis import Analysis +from easydynamics.analysis.bayesian_sampling import BayesianSamplingMixin from easydynamics.analysis.parameter_analysis import ParameterAnalysis +from easydynamics.analysis.posterior import BoundsSuggestion +from easydynamics.analysis.posterior import BoundsSuggestions +from easydynamics.analysis.posterior import ParameterPosterior +from easydynamics.analysis.posterior import PosteriorSummary __all__ = [ 'Analysis', + 'BayesianSamplingMixin', + 'BoundsSuggestion', + 'BoundsSuggestions', 'ParameterAnalysis', + 'ParameterPosterior', + 'PosteriorSummary', ] diff --git a/src/easydynamics/analysis/analysis1d.py b/src/easydynamics/analysis/analysis1d.py index e50b09737..eaec732b1 100644 --- a/src/easydynamics/analysis/analysis1d.py +++ b/src/easydynamics/analysis/analysis1d.py @@ -12,6 +12,7 @@ from plopp.backends.matplotlib.figure import InteractiveFigure from easydynamics.analysis.analysis_base import AnalysisBase +from easydynamics.analysis.bayesian_sampling import BayesianSamplingMixin from easydynamics.convolution.convolution import Convolution from easydynamics.experiment import Experiment from easydynamics.sample_model import InstrumentModel @@ -25,12 +26,16 @@ from easydynamics.utils.utils import verify_Q_index -class Analysis1d(AnalysisBase): +class Analysis1d(BayesianSamplingMixin, AnalysisBase): """ For analysing one-dimensional data, i.e. intensity as function of energy for a single Q index. Is used primarily in the Analysis class, but can also be used on its own for simpler analyses. + In addition to least-squares fitting with :meth:`fit`, the posterior distribution of the free + parameters can be explored with :meth:`sample_posterior`; see + :class:`~easydynamics.analysis.bayesian_sampling.BayesianSamplingMixin`. + Examples -------- **Fitting a single Q slice** @@ -116,6 +121,7 @@ def __init__( self._fit_result = None self._convolver = None self._convolver_is_dirty = True + self._init_bayesian_state() super().__init__( display_name=display_name, @@ -245,27 +251,72 @@ def fit(self) -> FitResults: if self._experiment is None: raise ValueError('No experiment is associated with this Analysis.') - if ( - self.sample_model.component_collections_is_dirty - or self.instrument_model.resolution_model.component_collections_is_dirty - ): - self._convolver_is_dirty = True + self._prepare_for_sampling() - self._ensure_convolver_current() + x, y, weights = self._get_sampling_data() + fit_result = self.fitter.fit(x=x, y=y, weights=weights) + + self._fit_result = fit_result + + return fit_result - fitter = EasyScienceFitter( + ############# + # Hooks for BayesianSamplingMixin + ############# + + def _build_bayesian_fitter(self) -> EasyScienceFitter: + """ + Build the EasyScience Fitter for this Analysis. + + Returns + ------- + EasyScienceFitter + A Fitter bound to this Analysis and its fit function. + """ + return EasyScienceFitter( fit_object=self, fit_function=self.as_fit_function(), ) + def _get_sampling_data(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Get the finite data for the chosen Q index, as used by both fitting and sampling. + + Returns + ------- + tuple[np.ndarray, np.ndarray, np.ndarray] + The ``(x, y, weights)`` triple. + """ x, y, weights, _ = self.experiment.extract_x_y_weights_only_finite( Q_index=self._require_Q_index() ) - fit_result = fitter.fit(x=x, y=y, weights=weights) + return x, y, weights - self._fit_result = fit_result + def _get_chain_parameters(self) -> list[Parameter]: + """ + Get the free parameters of this Analysis. - return fit_result + Returns + ------- + list[Parameter] + The parameters that are free to vary, which are the ones the sampler explores. + """ + return self.get_free_parameters() + + def _prepare_for_sampling(self) -> None: + """ + Rebuild the convolver if anything it depends on has changed. + + The energy grid is fixed for the duration of a fit or a sampling run, so the convolution + objects are built once here and reused for every model evaluation. + """ + if ( + self.sample_model.component_collections_is_dirty + or self.instrument_model.resolution_model.component_collections_is_dirty + ): + self._convolver_is_dirty = True + + self._ensure_convolver_current() def as_fit_function( self, @@ -483,6 +534,7 @@ def rebin(self, dimensions: dict[str, int | sc.Variable]) -> None: if self._Q_index is not None and self.experiment is not None: self._masked_energy = self.experiment.get_masked_energy(Q_index=self._Q_index) self._convolver_is_dirty = True + self._invalidate_bayesian_sampler() def refresh_convolver(self, energy: sc.Variable | None = None) -> None: """Refresh the pre-built Convolution object for the current Q index.""" @@ -523,10 +575,13 @@ def _on_Q_index_changed(self) -> None: if self._Q_index is None: self._masked_energy = None self._convolver_is_dirty = True + self._invalidate_bayesian_sampler() return masked_energy = self.experiment.get_masked_energy(Q_index=self._Q_index) self._masked_energy = masked_energy self._convolver_is_dirty = True + # A different Q index means different data, and the Sampler binds its data at construction. + self._invalidate_bayesian_sampler() def _on_experiment_changed(self) -> None: """Mark the convolver as dirty when the experiment changes.""" @@ -535,16 +590,19 @@ def _on_experiment_changed(self) -> None: if self._Q_index is not None and self.experiment is not None: self._masked_energy = self.experiment.get_masked_energy(Q_index=self._Q_index) self._convolver_is_dirty = True + self._invalidate_bayesian_sampler() def _on_sample_model_changed(self) -> None: """Mark the convolver as dirty when the sample model changes.""" super()._on_sample_model_changed() self._convolver_is_dirty = True + self._invalidate_fitter() def _on_instrument_model_changed(self) -> None: """Mark the convolver as dirty when the instrument model changes.""" super()._on_instrument_model_changed() self._convolver_is_dirty = True + self._invalidate_fitter() def _on_convolution_settings_changed(self) -> None: """Mark the convolver as dirty when the convolution settings change.""" diff --git a/src/easydynamics/analysis/bayesian_sampling.py b/src/easydynamics/analysis/bayesian_sampling.py new file mode 100644 index 000000000..0ea38d8ef --- /dev/null +++ b/src/easydynamics/analysis/bayesian_sampling.py @@ -0,0 +1,993 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +""" +Shared Bayesian MCMC sampling machinery for the Analysis classes. + +Everything that does not depend on how a particular Analysis is wired up lives here: caching the +Fitter and the Sampler, guarding the parameter bounds, restoring parameter values afterwards, and +turning raw draws into a readable summary. A concrete Analysis supplies the three things that do +differ, via :meth:`BayesianSamplingMixin._build_bayesian_fitter`, +:meth:`BayesianSamplingMixin._get_sampling_data`, and +:meth:`BayesianSamplingMixin._get_chain_parameters`. +""" + +from __future__ import annotations + +import json +import warnings +from pathlib import Path +from typing import TYPE_CHECKING +from typing import Any + +import numpy as np +from easyscience.fitting import AvailableMinimizers +from easyscience.fitting import Sampler + +from easydynamics.analysis.posterior import PosteriorSummary +from easydynamics.analysis.posterior import parameters_at_bounds +from easydynamics.analysis.posterior import suggest_bounds_for_parameters +from easydynamics.analysis.posterior import summarize_draws +from easydynamics.analysis.posterior import unbounded_parameters + +if TYPE_CHECKING: + import os + from collections.abc import Callable + + from easyscience.fitting.fitter import Fitter + from easyscience.fitting.sampler import SamplingResults + from easyscience.variable import Parameter + from matplotlib.figure import Figure + + from easydynamics.analysis.posterior import BoundsSuggestions + +# Suffix of the sidecar mapping chain columns to stable parameter names, written next to the BUMPS +# chain files by save_chain(). +_NAME_MAP_SUFFIX = '.parameter-names.json' + + +class BayesianSamplingMixin: + """ + Bayesian MCMC sampling on top of an Analysis, backed by the BUMPS DREAM sampler. + + Sampling explores the full posterior distribution of the free parameters rather than reporting + a single best-fit point, which is worth doing when parameters are correlated or their + uncertainties are strongly non-Gaussian -- both common in QENS. + + Running :meth:`fit` first is not required, but it helps: DREAM seeds its population in a small + ball around the parameters' current values, so starting from fitted values shortens the burn-in + needed to reach the typical set. + + Notes + ----- + All free parameters must have finite bounds before sampling, because in DREAM the bounds are + the prior. :meth:`suggest_bounds` proposes bounds for any parameter still missing one. + + Examples + -------- + ```python + analysis.fit() + analysis.suggest_bounds().apply() + results = analysis.sample_posterior(samples=10000, burn=2000, thin=10) + analysis.posterior_summary() + ``` + """ + + ############# + # Setup + ############# + + def _init_bayesian_state(self) -> None: + """ + Initialize the cached sampling state. + + Must be called by the concrete Analysis before any observer callback can fire, in the same + way as the other cached objects on the class. + """ + self._fitter = None + self._fitter_is_dirty = True + self._bayesian_sampler = None + self._bayesian_sampler_is_dirty = True + self._posterior_result = None + # Maps a chain column's unique_name to the parameter name it had when saved. Only populated + # by load_chain, because unique_names are per-session and do not survive a round trip. + self._chain_name_map = {} + + def _invalidate_fitter(self) -> None: + """ + Mark the cached Fitter and Sampler as needing a rebuild. + + The Sampler binds its data at construction, so anything that invalidates the Fitter + invalidates the Sampler too. + """ + self._fitter_is_dirty = True + self._bayesian_sampler_is_dirty = True + + def _invalidate_bayesian_sampler(self) -> None: + """ + Mark only the cached Sampler as needing a rebuild. + + Used when the data changed but the model did not. + """ + self._bayesian_sampler_is_dirty = True + + ############# + # Hooks for concrete Analysis classes + ############# + + def _build_bayesian_fitter(self) -> Fitter: + """ + Build the EasyScience Fitter (or MultiFitter) for this Analysis. + + Returns + ------- + Fitter + A configured Fitter or MultiFitter. + + Raises + ------ + NotImplementedError + If the concrete Analysis does not implement it. + """ + raise NotImplementedError('Subclasses must implement _build_bayesian_fitter.') + + def _get_sampling_data(self) -> tuple: + """ + Get the ``(x, y, weights)`` to bind to the Sampler. + + Each element is either an array (single dataset) or a list of arrays (MultiFitter). + + Returns + ------- + tuple + The ``(x, y, weights)`` triple. + + Raises + ------ + NotImplementedError + If the concrete Analysis does not implement it. + """ + raise NotImplementedError('Subclasses must implement _get_sampling_data.') + + def _get_chain_parameters(self) -> list[Parameter]: + """ + Get the free parameters that will appear as columns of the chain. + + Returns + ------- + list[Parameter] + The free parameters of the underlying model(s). + + Raises + ------ + NotImplementedError + If the concrete Analysis does not implement it. + """ + raise NotImplementedError('Subclasses must implement _get_chain_parameters.') + + def _prepare_for_sampling(self) -> None: + """ + Bring any cached computation up to date before a sampling run. + + The default does nothing; Analysis classes that cache a convolver override it. + """ + + ############# + # Properties + ############# + + @property + def fitter(self) -> Fitter: + """ + The EasyScience Fitter used for fitting and sampling, built on first use. + + Exposed so the minimizer, tolerance, and maximum evaluation count can be configured + directly, e.g. ``analysis.fitter.switch_minimizer(AvailableMinimizers.Bumps)``. + + Returns + ------- + Fitter + The cached Fitter or MultiFitter. + """ + if self._fitter_is_dirty or self._fitter is None: + self._fitter = self._build_bayesian_fitter() + self._fitter_is_dirty = False + return self._fitter + + @property + def bayesian_sampler(self) -> Sampler | None: + """ + The EasyScience Sampler holding the MCMC chain, or None before the first run. + + Named to avoid confusion with the SampleModel: this samples the posterior, not the sample. + + Returns + ------- + Sampler | None + The cached Sampler, or None if no chain has been started. + """ + return self._bayesian_sampler + + @property + def posterior_result(self) -> SamplingResults | None: + """ + The results of the most recent sampling run, or None if there has not been one. + + Returns + ------- + SamplingResults | None + The most recent sampling results. + """ + return self._posterior_result + + ############# + # Bounds + ############# + + def suggest_bounds( + self, + n_sigma: float = 10.0, + relative_pad: float = 0.2, + absolute_floor: float | None = None, + ) -> BoundsSuggestions: + """ + Propose finite bounds for free parameters that still have an infinite one. + + Nothing is changed until :meth:`BoundsSuggestions.apply` is called, so the proposal can be + reviewed first. Bounds that are already finite are never widened or narrowed, so physical + limits such as a non-negative area are left alone. + + Because the bounds act as a uniform prior in DREAM, a generous width is the safe choice: + too tight a bound truncates the posterior and understates the uncertainty. + + Parameters + ---------- + n_sigma : float, default=10.0 + How many standard deviations of the fitted uncertainty to allow on each side. + relative_pad : float, default=0.2 + Extra half-width as a fraction of the absolute parameter value. This guards against + minimizers that report a zero or absurdly small uncertainty. + absolute_floor : float | None, default=None + A minimum half-width in the parameter's own units, for when neither the uncertainty nor + the value carries the natural scale. + + Returns + ------- + BoundsSuggestions + The proposed bounds, which must be applied explicitly. + """ + return suggest_bounds_for_parameters( + self._get_chain_parameters(), + n_sigma=n_sigma, + relative_pad=relative_pad, + absolute_floor=absolute_floor, + ) + + def check_bounds_for_sampling(self) -> None: + """ + Verify that every free parameter has finite bounds. + + Raises + ------ + ValueError + If any free parameter has an infinite lower or upper bound. + """ + unbounded = unbounded_parameters(self._get_chain_parameters()) + if not unbounded: + return + names = ', '.join(parameter.name for parameter in unbounded) + raise ValueError( + f'Bayesian sampling requires finite bounds on every free parameter, because the ' + f'bounds act as the prior. These parameters are unbounded: {names}. ' + f'Set their min and max, or call suggest_bounds() to propose values.' + ) + + ############# + # Sampling + ############# + + def sample_posterior( + self, + samples: int = 10000, + burn: int = 2000, + thin: int = 10, + population: int | None = None, + parameters: list[Parameter] | list[str] | None = None, + **sampler_options: dict[str, Any], + ) -> SamplingResults: + """ + Draw samples from the posterior distribution of the free parameters. + + This starts a fresh chain, replacing any existing one; use :meth:`extend_sampling` to + continue a chain instead. Parameter values are restored to what they were beforehand, so + sampling never silently moves the model off its fitted values; use + :meth:`set_parameters_to_posterior_median` to adopt the posterior. + + Parameters + ---------- + samples : int, default=10000 + Number of raw samples to draw across all chains, before thinning. This is a guaranteed + minimum rather than an exact count. + burn : int, default=2000 + Burn-in generations to discard before collecting samples. + thin : int, default=10 + Thinning interval, which reduces autocorrelation between retained draws. + population : int | None, default=None + DREAM population scale factor: BUMPS runs ``ceil(population * n_parameters)`` chains. + parameters : list[Parameter] | list[str] | None, default=None + Restrict the chain to these parameters, given as Parameter objects or names. All other + free parameters are held fixed for the duration of the run. Note that holding a + parameter fixed is not the same as marginalizing over it: the resulting credible + intervals are conditional on the fixed values and will be too narrow if the parameters + are correlated. The default samples every free parameter. + **sampler_options : dict[str, Any] + Forwarded to the EasyScience Sampler, e.g. ``sampler_kwargs``, ``progress_callback``, + or ``abort_test``. + + Returns + ------- + SamplingResults + The sampling results, also stored on :attr:`posterior_result`. + """ + return self._run_sampling( + parameters=parameters, + run=lambda sampler: sampler.sample( + samples=samples, + burn=burn, + thin=thin, + population=population, + **sampler_options, + ), + ) + + def extend_sampling( + self, + additional_samples: int = 5000, + thin: int = 10, + parameters: list[Parameter] | list[str] | None = None, + **sampler_options: dict[str, Any], + ) -> SamplingResults: + """ + Continue the existing chain with additional samples. + + Parameters + ---------- + additional_samples : int, default=5000 + Number of additional samples to draw, in the same units as ``samples``. + thin : int, default=10 + Thinning interval for the retained draws. + parameters : list[Parameter] | list[str] | None, default=None + The same restriction as in :meth:`sample_posterior`. Pass the same value that started + the chain, since the chain's columns cannot change on extension. + **sampler_options : dict[str, Any] + Forwarded to the EasyScience Sampler. + + Returns + ------- + SamplingResults + The sampling results for the full extended chain. + + Raises + ------ + RuntimeError + If there is no chain to extend. + """ + if self._bayesian_sampler is None: + raise RuntimeError( + 'No chain to extend. Call sample_posterior() or load_chain() first.' + ) + return self._run_sampling( + parameters=parameters, + run=lambda sampler: sampler.extend( + additional_samples=additional_samples, + thin=thin, + **sampler_options, + ), + reuse_sampler=True, + ) + + def _run_sampling( + self, + parameters: list[Parameter] | list[str] | None, + run: Callable[[Sampler], SamplingResults], + reuse_sampler: bool = False, + ) -> SamplingResults: + """ + Run a sampling operation with all the surrounding guards in place. + + Checks the bounds, switches the minimizer to BUMPS, optionally holds parameters fixed, + runs, and then restores the parameter values, fixed flags, and minimizer. + + Parameters + ---------- + parameters : list[Parameter] | list[str] | None + Parameters to restrict the chain to, or None for all free parameters. + run : Callable[[Sampler], SamplingResults] + The operation to perform on the prepared Sampler. + reuse_sampler : bool, default=False + Whether to reuse the cached Sampler rather than rebuilding it. Required when extending + a chain, since the chain lives on the Sampler. + + Returns + ------- + SamplingResults + The results of the run. + + Raises + ------ + RuntimeError + If the BUMPS sampler fails while removing outlier chains, which points at degenerate + parameters. + """ + held_fixed = self._resolve_parameters_to_hold_fixed(parameters) + self._warn_about_held_parameters(held_fixed) + + with _FixedParameters(held_fixed): + self.check_bounds_for_sampling() + self._prepare_for_sampling() + + chain_parameters = self._get_chain_parameters() + saved_values = [(p, p.value) for p in chain_parameters] + + fitter = self.fitter + original_minimizer = fitter.minimizer.enum + fitter.switch_minimizer(AvailableMinimizers.Bumps) + try: + sampler = self._get_or_build_sampler(reuse_sampler=reuse_sampler) + results = run(sampler) + except IndexError as error: + # BUMPS' own outlier removal indexes past the end of its buffer when chains + # scatter wildly, which in practice means the model is not identifiable. The bare + # IndexError says nothing useful, so point at the likely cause instead. + raise RuntimeError( + 'The BUMPS sampler failed while removing outlier chains. This usually means ' + 'the chains scattered because two or more free parameters are degenerate, so ' + 'the data cannot determine them separately. Check for degenerate parameters ' + "and fix one of them, or retry with sampler_kwargs={'outliers': 'none'}." + ) from error + finally: + fitter.switch_minimizer(original_minimizer) + for parameter, value in saved_values: + parameter.value = value + + # A fresh chain is labelled with this session's unique names, so any mapping left over + # from a loaded chain no longer applies. + self._chain_name_map = { + parameter.unique_name: parameter.name for parameter in chain_parameters + } + + self._posterior_result = results + self._warn_about_bounds_occupancy(results, self._resolve_chain_parameters(results)) + return results + + def _get_or_build_sampler(self, reuse_sampler: bool) -> Sampler: + """ + Get the cached Sampler, rebuilding it if the data or model changed. + + Parameters + ---------- + reuse_sampler : bool + Whether to reuse the cached Sampler even if it is marked dirty. + + Returns + ------- + Sampler + The Sampler to run. + """ + needs_rebuild = self._bayesian_sampler is None or ( + self._bayesian_sampler_is_dirty and not reuse_sampler + ) + if needs_rebuild: + x, y, weights = self._get_sampling_data() + self._bayesian_sampler = Sampler(self.fitter, x, y, weights=weights) + self._bayesian_sampler_is_dirty = False + return self._bayesian_sampler + + def _resolve_parameters_to_hold_fixed( + self, + parameters: list[Parameter] | list[str] | None, + ) -> list[Parameter]: + """ + Work out which free parameters must be held fixed to honour a subset request. + + Parameters + ---------- + parameters : list[Parameter] | list[str] | None + The requested subset, as Parameter objects or names, or None for all free parameters. + + Returns + ------- + list[Parameter] + The free parameters that are not in the requested subset. + + Raises + ------ + TypeError + If parameters is not a list of Parameters or strings, or None. + ValueError + If a requested name does not match any free parameter, or the subset is empty. + """ + if parameters is None: + return [] + if not isinstance(parameters, (list, tuple)): + raise TypeError('parameters must be a list of Parameters, a list of names, or None.') + + free = self._get_chain_parameters() + by_name = {parameter.name: parameter for parameter in free} + requested = [] + for entry in parameters: + if isinstance(entry, str): + if entry not in by_name: + available = ', '.join(sorted(by_name)) + raise ValueError(f'No free parameter named {entry!r}. Available: {available}.') + requested.append(by_name[entry]) + elif hasattr(entry, 'unique_name'): + requested.append(entry) + else: + raise TypeError( + 'parameters must contain Parameter objects or parameter names (strings).' + ) + + requested_unique_names = {parameter.unique_name for parameter in requested} + if not requested_unique_names: + raise ValueError('parameters must name at least one parameter to sample.') + return [ + parameter for parameter in free if parameter.unique_name not in requested_unique_names + ] + + @staticmethod + def _warn_about_held_parameters(held_fixed: list[Parameter]) -> None: + """ + Warn that holding parameters fixed makes the credible intervals conditional. + + Parameters + ---------- + held_fixed : list[Parameter] + The parameters being held fixed for the run. + """ + if not held_fixed: + return + names = ', '.join(parameter.name for parameter in held_fixed) + warnings.warn( + ( + f'Holding these parameters fixed while sampling: {names}. ' + f'Fixing a parameter is not the same as marginalizing over it, so the resulting ' + f'credible intervals are conditional on these values and will be too narrow if ' + f'the parameters are correlated.' + ), + UserWarning, + stacklevel=4, + ) + + @staticmethod + def _warn_about_bounds_occupancy( + results: SamplingResults, + parameters_by_column: list[Parameter | None], + ) -> None: + """ + Warn when the posterior has piled up against a bound. + + Parameters + ---------- + results : SamplingResults + The sampling results to inspect. + parameters_by_column : list[Parameter | None] + The parameter for each column of the chain, or None where none could be matched. + """ + piled_up = parameters_at_bounds(results.draws, parameters_by_column) + if not piled_up: + return + details = ', '.join( + f'{name} ({fraction:.0%} of draws)' for name, fraction in piled_up.items() + ) + warnings.warn( + ( + f'The posterior is piled up against the bounds for: {details}. ' + f'The bounds, rather than the data, are setting these credible intervals. ' + f'Widen the bounds, or check whether these parameters are degenerate with others.' + ), + UserWarning, + stacklevel=4, + ) + + ############# + # Results + ############# + + def posterior_summary(self) -> PosteriorSummary: + """ + Summarize the marginal posterior of each sampled parameter. + + Reports the median and the 68% credible interval under the parameter's own name and unit, + rather than the opaque unique name the sampler uses internally. Requires a completed + sampling run. + + Returns + ------- + PosteriorSummary + One entry per sampled parameter. + """ + results = self._require_posterior_result() + return summarize_draws( + draws=results.draws, + fallback_names=self._chain_display_names(results), + parameters_by_column=self._resolve_chain_parameters(results), + ) + + def set_parameters_to_posterior_median(self) -> list[Parameter]: + """ + Set every sampled parameter to the median of its marginal posterior. + + Note that the vector of marginal medians is not in general the same as the + highest-posterior point, and for strongly correlated parameters it need not even be a good + fit. Requires a completed sampling run. + + Returns + ------- + list[Parameter] + The parameters that were changed. + """ + results = self._require_posterior_result() + changed = [] + for column, parameter in enumerate(self._resolve_chain_parameters(results)): + if parameter is None: + continue + parameter.value = float(np.median(results.draws[:, column])) + changed.append(parameter) + return changed + + def _require_posterior_result(self) -> SamplingResults: + """ + Get the stored sampling results, raising if there are none. + + Returns + ------- + SamplingResults + The most recent sampling results. + + Raises + ------ + RuntimeError + If no sampling has been run yet. + """ + if self._posterior_result is None: + raise RuntimeError( + 'No posterior samples yet. Call sample_posterior() or load_chain() first.' + ) + return self._posterior_result + + ############# + # Persistence + ############# + + def save_chain(self, path: str | os.PathLike) -> None: + """ + Save the MCMC chain to disk. + + Writes the BUMPS chain files alongside a sidecar recording the parameter names and a + fingerprint of the data that was sampled. + + Parameters + ---------- + path : str | os.PathLike + Path prefix for the chain files. + + Raises + ------ + RuntimeError + If there is no chain to save. + """ + if self._bayesian_sampler is None: + raise RuntimeError('No chain to save. Call sample_posterior() first.') + self._bayesian_sampler.save(path) + # The BUMPS sidecar records unique names, which are handed out per session and so mean + # nothing on reload. Record the parameter names alongside them, which are stable. + Path(f'{path}{_NAME_MAP_SUFFIX}').write_text( + json.dumps(self._chain_name_map, indent=2), + encoding='utf-8', + ) + + def load_chain(self, path: str | os.PathLike, skip: int = 0) -> SamplingResults: + """ + Load a previously saved MCMC chain. + + The loaded chain can be inspected, summarized, or continued with :meth:`extend_sampling`. A + chain saved from different data loads with a warning. + + Parameters + ---------- + path : str | os.PathLike + The path prefix the chain was saved under. + skip : int, default=0 + Number of initial samples to skip when reading the chain. + + Returns + ------- + SamplingResults + The loaded sampling results, also stored on :attr:`posterior_result`. + """ + self._prepare_for_sampling() + name_map_path = Path(f'{path}{_NAME_MAP_SUFFIX}') + if name_map_path.is_file(): + self._chain_name_map = json.loads(name_map_path.read_text(encoding='utf-8')) + else: + self._chain_name_map = {} + warnings.warn( + ( + f'No parameter-name sidecar found at {name_map_path}. The chain will be ' + f'reported under the internal names it was saved with, because those cannot ' + f'be matched to this Analysis.' + ), + UserWarning, + stacklevel=2, + ) + + fitter = self.fitter + original_minimizer = fitter.minimizer.enum + fitter.switch_minimizer(AvailableMinimizers.Bumps) + try: + sampler = self._get_or_build_sampler(reuse_sampler=False) + results = sampler.load_state(path, skip=skip) + finally: + fitter.switch_minimizer(original_minimizer) + self._posterior_result = results + return results + + ############# + # Plotting + ############# + + def plot_trace(self, **kwargs: dict[str, Any]) -> Figure: + """ + Plot the chain trace of each sampled parameter. + + A well-mixed chain looks like a "hairy caterpillar" with no drift; visible trends mean the + chain has not converged and needs a longer burn-in. Requires a completed sampling run. + + Parameters + ---------- + **kwargs : dict[str, Any] + Forwarded to :func:`easydynamics.utils.posterior_plotting.plot_trace`. + + Returns + ------- + Figure + The matplotlib Figure. + """ + from easydynamics.utils.posterior_plotting import plot_trace + + results = self._require_posterior_result() + return plot_trace( + draws=results.draws, + logp=results.logp, + names=self._chain_display_names(results), + title=self.display_name, + **kwargs, + ) + + def plot_corner(self, **kwargs: dict[str, Any]) -> Figure: + """ + Plot the marginal and pairwise posterior distributions. + + Diagonal panels show each parameter's marginal distribution; off-diagonal panels show the + joint distribution of a pair, where a strong diagonal ridge means the two are correlated. + Requires a completed sampling run. + + Parameters + ---------- + **kwargs : dict[str, Any] + Forwarded to :func:`easydynamics.utils.posterior_plotting.plot_corner`. + + Returns + ------- + Figure + The matplotlib Figure. + """ + from easydynamics.utils.posterior_plotting import plot_corner + + results = self._require_posterior_result() + return plot_corner( + draws=results.draws, + names=self._chain_display_names(results), + title=self.display_name, + **kwargs, + ) + + def plot_posterior_predictive( + self, + n_draws: int = 200, + credible_interval: float = 68.0, + **kwargs: dict[str, Any], + ) -> Figure: + """ + Plot the data against the credible band implied by the posterior. + + The model is re-evaluated for a random subset of the posterior draws, and the spread of + those curves becomes the band. Data straying outside the band systematically points at a + model that is missing something, rather than at parameters that need tuning. Requires a + completed sampling run. + + Parameters + ---------- + n_draws : int, default=200 + How many posterior draws to evaluate the model for. Each draw costs one full model + evaluation, so this trades smoothness of the band against time. + credible_interval : float, default=68.0 + Width of the credible band, as a percentage. + **kwargs : dict[str, Any] + Forwarded to :func:`easydynamics.utils.posterior_plotting.plot_posterior_predictive`. + + Returns + ------- + Figure + The matplotlib Figure. + + Raises + ------ + NotImplementedError + If this Analysis binds a list of datasets rather than a single one. + ValueError + If n_draws is not a positive integer. + """ + from easydynamics.utils.posterior_plotting import plot_posterior_predictive + + if not isinstance(n_draws, int) or isinstance(n_draws, bool) or n_draws < 1: + raise ValueError(f'n_draws must be a positive integer. Got {n_draws}.') + + results = self._require_posterior_result() + x, y, weights = self._get_sampling_data() + if isinstance(x, (list, tuple)): + raise NotImplementedError( + 'plot_posterior_predictive supports a single dataset only. Plot each dataset ' + 'from its own Analysis1d instead.' + ) + + predictions = self._evaluate_over_draws(results, x, n_draws) + y_err = None if weights is None else 1.0 / np.asarray(weights) + return plot_posterior_predictive( + x=np.asarray(x), + y=np.asarray(y), + predictions=predictions, + y_err=y_err, + title=self.display_name, + credible_interval=credible_interval, + **kwargs, + ) + + def _evaluate_over_draws( + self, + results: SamplingResults, + x: np.ndarray, + n_draws: int, + ) -> np.ndarray: + """ + Evaluate the model once per posterior draw, restoring the parameters afterwards. + + Parameters + ---------- + results : SamplingResults + The sampling results supplying the draws. + x : np.ndarray + The independent variable to evaluate the model on. + n_draws : int + How many draws to evaluate. Draws are taken evenly across the chain. + + Returns + ------- + np.ndarray + Model evaluations, shape ``(n_selected, len(x))``. + """ + self._prepare_for_sampling() + + columns = [ + (parameter, column) + for column, parameter in enumerate(self._resolve_chain_parameters(results)) + if parameter is not None + ] + saved_values = [(parameter, parameter.value) for parameter, _ in columns] + + total = results.draws.shape[0] + indices = np.unique(np.linspace(0, total - 1, min(n_draws, total)).astype(int)) + + fit_function = self.fitter.fit_function + predictions = [] + try: + for index in indices: + for parameter, column in columns: + parameter.value = float(results.draws[index, column]) + predictions.append(np.asarray(fit_function(x))) + finally: + for parameter, value in saved_values: + parameter.value = value + + return np.vstack(predictions) + + def _resolve_chain_parameters(self, results: SamplingResults) -> list[Parameter | None]: + """ + Match each column of the chain to one of this Analysis's parameters. + + Columns are matched on ``unique_name`` first. That fails for a chain loaded from disk, + because unique names are handed out per session, so a saved chain also records the + parameter names and those are used as a fallback. + + Parameters + ---------- + results : SamplingResults + The sampling results whose columns should be matched. + + Returns + ------- + list[Parameter | None] + The parameter for each column, or None where no match could be made. + """ + parameters = self._get_chain_parameters() + by_unique_name = {p.unique_name: p for p in parameters} + by_name = {p.name: p for p in parameters} + resolved = [] + for unique_name in results.param_names: + parameter = by_unique_name.get(unique_name) + if parameter is None: + saved_name = self._chain_name_map.get(unique_name) + parameter = None if saved_name is None else by_name.get(saved_name) + resolved.append(parameter) + return resolved + + def _chain_display_names(self, results: SamplingResults) -> list[str]: + """ + Translate the chain's column names into parameter names. + + Parameters + ---------- + results : SamplingResults + The sampling results whose columns should be named. + + Returns + ------- + list[str] + One name per column of the chain. + """ + resolved = self._resolve_chain_parameters(results) + return [ + self._chain_name_map.get(unique_name, unique_name) + if parameter is None + else parameter.name + for unique_name, parameter in zip(results.param_names, resolved, strict=True) + ] + + +class _FixedParameters: + """ + Context manager that temporarily fixes parameters and restores their flags on exit. + """ + + def __init__(self, parameters: list[Parameter]) -> None: + """ + Initialize the context manager. + + Parameters + ---------- + parameters : list[Parameter] + The parameters to hold fixed for the duration of the block. + """ + self._parameters = list(parameters) + self._saved: list[tuple[Parameter, bool]] = [] + + def __enter__(self) -> None: + """ + Fix the parameters, remembering their previous state. + """ + self._saved = [(parameter, parameter.fixed) for parameter in self._parameters] + for parameter in self._parameters: + parameter.fixed = True + + def __exit__(self, *_exc_info: object) -> None: + """ + Restore the previous fixed state of every parameter. + + Parameters + ---------- + *_exc_info : object + Exception information, ignored. + """ + for parameter, was_fixed in self._saved: + parameter.fixed = was_fixed diff --git a/src/easydynamics/analysis/posterior.py b/src/easydynamics/analysis/posterior.py new file mode 100644 index 000000000..e22ae87d9 --- /dev/null +++ b/src/easydynamics/analysis/posterior.py @@ -0,0 +1,608 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +""" +Bounds suggestions and posterior summaries for Bayesian sampling. + +The helpers here are deliberately free of any Analysis or Fitter machinery: they operate on plain +``Parameter`` objects and on the ``(n_draws, n_parameters)`` array produced by the sampler, so they +can be unit-tested on their own and reused by every Analysis class. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import numpy as np + +if TYPE_CHECKING: + from easyscience.variable import Parameter + +# Fraction of the allowed range at each end that counts as "at the bound" when checking whether +# the posterior has piled up against a bound. +BOUND_EDGE_FRACTION = 0.05 + +# Fraction of draws inside those edge bands above which a pile-up is reported. A posterior spread +# uniformly across its bounds -- the signature of a bound, rather than the data, setting the +# credible interval -- puts 2 * BOUND_EDGE_FRACTION of its draws there. A posterior comfortably +# inside its bounds puts essentially none there, so the threshold sits well below the uniform value +# to stay sensitive to partly-clipped posteriors without risking false positives. +BOUND_OCCUPANCY_THRESHOLD = 0.05 + + +@dataclass(frozen=True) +class BoundsSuggestion: + """ + A proposed pair of bounds for a single parameter. + + Attributes + ---------- + parameter : Parameter + The parameter the suggestion applies to. + suggested_min : float + The proposed lower bound. Equal to the parameter's current lower bound when that is already + finite. + suggested_max : float + The proposed upper bound. Equal to the parameter's current upper bound when that is already + finite. + reason : str + Empty when the suggestion is usable. Otherwise, why the parameter needs manual attention. + """ + + parameter: Parameter + suggested_min: float + suggested_max: float + reason: str + + @property + def needs_attention(self) -> bool: + """ + Whether this parameter could not be given a usable suggestion. + + Returns + ------- + bool + True when no usable bounds could be derived and the user must set them by hand. + """ + return bool(self.reason) + + @property + def changes_bounds(self) -> bool: + """ + Whether applying this suggestion would actually change the parameter. + + Returns + ------- + bool + True when either bound differs from the parameter's current bound. + """ + return self.suggested_min != self.parameter.min or self.suggested_max != self.parameter.max + + +class BoundsSuggestions: + """ + The result of :func:`suggest_bounds_for_parameters`, rendered as a table. + + This is advisory: nothing is changed until :meth:`apply` is called. Suggestions only ever fill + in an infinite bound; a bound that is already finite is never widened or narrowed, so physical + limits such as a non-negative area survive untouched. + """ + + def __init__(self, suggestions: list[BoundsSuggestion]) -> None: + """ + Initialize the collection. + + Parameters + ---------- + suggestions : list[BoundsSuggestion] + The per-parameter suggestions. + """ + self._suggestions = list(suggestions) + + @property + def suggestions(self) -> list[BoundsSuggestion]: + """ + All suggestions, including those needing manual attention. + + Returns + ------- + list[BoundsSuggestion] + The per-parameter suggestions. + """ + return list(self._suggestions) + + @property + def needing_attention(self) -> list[BoundsSuggestion]: + """ + The suggestions for which no usable bounds could be derived. + + Returns + ------- + list[BoundsSuggestion] + Suggestions whose parameters must be bounded by hand. + """ + return [s for s in self._suggestions if s.needs_attention] + + def apply(self) -> list[Parameter]: + """ + Set the suggested bounds on every parameter that has a usable suggestion. + + Parameters needing manual attention are skipped rather than guessed at. + + Returns + ------- + list[Parameter] + The parameters whose bounds were changed. + """ + changed = [] + for suggestion in self._suggestions: + if suggestion.needs_attention or not suggestion.changes_bounds: + continue + suggestion.parameter.min = suggestion.suggested_min + suggestion.parameter.max = suggestion.suggested_max + changed.append(suggestion.parameter) + return changed + + def __len__(self) -> int: + """ + Return the number of suggestions. + + Returns + ------- + int + The number of suggestions. + """ + return len(self._suggestions) + + def __iter__(self) -> iter: + """ + Iterate over the suggestions. + + Returns + ------- + iter + An iterator over the suggestions. + """ + return iter(self._suggestions) + + def __repr__(self) -> str: + """ + Render the suggestions as a table. + + Returns + ------- + str + A table of current and suggested bounds, one row per parameter. + """ + if not self._suggestions: + return 'BoundsSuggestions(no free parameters)' + + header = f'{"parameter":<28s} {"current":>26s} {"suggested":>26s}' + lines = ['BoundsSuggestions', header, '-' * len(header)] + for s in self._suggestions: + current = f'({s.parameter.min:.4g}, {s.parameter.max:.4g})' + if s.needs_attention: + suggested = f'-- {s.reason}' + else: + suggested = f'({s.suggested_min:.4g}, {s.suggested_max:.4g})' + lines.append(f'{s.parameter.name:<28s} {current:>26s} {suggested:>26s}') + + attention = self.needing_attention + if attention: + lines.append('') + lines.append( + f'{len(attention)} parameter(s) need bounds set by hand; .apply() will skip them.' + ) + return '\n'.join(lines) + + +def suggest_bounds_for_parameters( + parameters: list[Parameter], + n_sigma: float = 10.0, + relative_pad: float = 0.2, + absolute_floor: float | None = None, +) -> BoundsSuggestions: + """ + Propose finite bounds for parameters that currently have an infinite one. + + The half-width of a proposed bound is ``n_sigma * error + relative_pad * abs(value)``, floored + at ``absolute_floor`` when one is given. The ``relative_pad`` term matters because + least-squares minimizers sometimes report a zero or absurdly small uncertainty; without it such + a parameter would be given a zero-width bound. When the half-width still comes out as zero or + non-finite, the parameter is flagged for manual attention rather than given an invented scale. + + In BUMPS' DREAM sampler the bounds act as a uniform prior, so a generous width is the safe + choice: too narrow a bound truncates the posterior and understates the uncertainty. Hence the + deliberately loose ``n_sigma`` default. + + A ``TypeError`` is raised if any of the three settings is not a number, and a ``ValueError`` if + any is negative. + + Parameters + ---------- + parameters : list[Parameter] + The parameters to propose bounds for. + n_sigma : float, default=10.0 + How many standard deviations of the parameter's fitted uncertainty to allow on each side. + relative_pad : float, default=0.2 + Extra half-width as a fraction of the absolute parameter value, guarding against + artificially small uncertainties. + absolute_floor : float | None, default=None + A minimum half-width, in the parameter's own units. Use it when the natural scale is known + but neither the uncertainty nor the value carries it. + + Returns + ------- + BoundsSuggestions + The proposed bounds, which must be applied explicitly. + """ + _verify_nonneg_number(n_sigma, 'n_sigma') + _verify_nonneg_number(relative_pad, 'relative_pad') + if absolute_floor is not None: + _verify_nonneg_number(absolute_floor, 'absolute_floor') + + suggestions = [ + _suggest_bounds_for_parameter( + parameter=parameter, + n_sigma=n_sigma, + relative_pad=relative_pad, + absolute_floor=absolute_floor, + ) + for parameter in parameters + ] + return BoundsSuggestions(suggestions) + + +def _suggest_bounds_for_parameter( + parameter: Parameter, + n_sigma: float, + relative_pad: float, + absolute_floor: float | None, +) -> BoundsSuggestion: + """ + Propose bounds for a single parameter. + + Parameters + ---------- + parameter : Parameter + The parameter to propose bounds for. + n_sigma : float + How many standard deviations to allow on each side. + relative_pad : float + Extra half-width as a fraction of the absolute parameter value. + absolute_floor : float | None + A minimum half-width, or None. + + Returns + ------- + BoundsSuggestion + The proposal for this parameter. + """ + current_min = float(parameter.min) + current_max = float(parameter.max) + min_is_finite = np.isfinite(current_min) + max_is_finite = np.isfinite(current_max) + + # Nothing to fill in: a bound that is already finite is never touched. + if min_is_finite and max_is_finite: + return BoundsSuggestion( + parameter=parameter, + suggested_min=current_min, + suggested_max=current_max, + reason='', + ) + + value = float(parameter.value) + error = float(parameter.error) + if not np.isfinite(value): + return BoundsSuggestion( + parameter=parameter, + suggested_min=current_min, + suggested_max=current_max, + reason='value is not finite', + ) + + half_width = relative_pad * abs(value) + if np.isfinite(error): + half_width += n_sigma * error + if absolute_floor is not None: + half_width = max(half_width, absolute_floor) + + if not np.isfinite(half_width) or half_width <= 0: + return BoundsSuggestion( + parameter=parameter, + suggested_min=current_min, + suggested_max=current_max, + reason='no scale information (zero value and uncertainty)', + ) + + return BoundsSuggestion( + parameter=parameter, + suggested_min=current_min if min_is_finite else value - half_width, + suggested_max=current_max if max_is_finite else value + half_width, + reason='', + ) + + +def unbounded_parameters(parameters: list[Parameter]) -> list[Parameter]: + """ + Find parameters with a non-finite lower or upper bound. + + Parameters + ---------- + parameters : list[Parameter] + The parameters to check. + + Returns + ------- + list[Parameter] + Those parameters that have at least one infinite bound. + """ + return [ + parameter + for parameter in parameters + if not (np.isfinite(parameter.min) and np.isfinite(parameter.max)) + ] + + +def parameters_at_bounds( + draws: np.ndarray, + parameters_by_column: list[Parameter | None], +) -> dict[str, float]: + """ + Find parameters whose posterior has piled up against one of its bounds. + + A chain that spends much of its time hard against a bound is a sign that the bound, rather than + the data, is setting the credible interval. That happens when a bound is too tight, and also + when two parameters are degenerate and the pair drifts until it is stopped by a bound. + + Parameters + ---------- + draws : np.ndarray + Posterior draws, shape ``(n_draws, n_parameters)``. + parameters_by_column : list[Parameter | None] + The parameter for each column of ``draws``, or None where no parameter could be matched. + + Returns + ------- + dict[str, float] + Mapping of parameter name to the fraction of draws sitting in the outer + ``BOUND_EDGE_FRACTION`` of its allowed range, for those parameters where that fraction + exceeds ``BOUND_OCCUPANCY_THRESHOLD``. + """ + piled_up = {} + for column, parameter in enumerate(parameters_by_column): + if parameter is None: + continue + low = float(parameter.min) + high = float(parameter.max) + if not (np.isfinite(low) and np.isfinite(high)) or high <= low: + continue + edge = BOUND_EDGE_FRACTION * (high - low) + values = draws[:, column] + at_edge = (values <= low + edge) | (values >= high - edge) + fraction = float(np.count_nonzero(at_edge)) / len(values) + if fraction > BOUND_OCCUPANCY_THRESHOLD: + piled_up[parameter.name] = fraction + return piled_up + + +@dataclass(frozen=True) +class ParameterPosterior: + """ + The marginal posterior of a single parameter. + + Attributes + ---------- + name : str + The parameter's name. + unit : str + The parameter's unit, as a string. + median : float + The 50th percentile of the marginal posterior. + lower : float + The 16th percentile. + upper : float + The 84th percentile. + value : float + The parameter's current value, for comparison with the median. + """ + + name: str + unit: str + median: float + lower: float + upper: float + value: float + + @property + def minus(self) -> float: + """ + Distance from the median down to the 16th percentile. + + Returns + ------- + float + The lower half of the 68% credible interval. + """ + return self.median - self.lower + + @property + def plus(self) -> float: + """ + Distance from the median up to the 84th percentile. + + Returns + ------- + float + The upper half of the 68% credible interval. + """ + return self.upper - self.median + + +class PosteriorSummary: + """ + Marginal posterior summaries for every sampled parameter, rendered as a table. + """ + + def __init__(self, entries: list[ParameterPosterior]) -> None: + """ + Initialize the summary. + + Parameters + ---------- + entries : list[ParameterPosterior] + One entry per sampled parameter. + """ + self._entries = list(entries) + + @property + def entries(self) -> list[ParameterPosterior]: + """ + The per-parameter summaries. + + Returns + ------- + list[ParameterPosterior] + One entry per sampled parameter. + """ + return list(self._entries) + + def __len__(self) -> int: + """ + Return the number of summarized parameters. + + Returns + ------- + int + The number of entries. + """ + return len(self._entries) + + def __iter__(self) -> iter: + """ + Iterate over the entries. + + Returns + ------- + iter + An iterator over the entries. + """ + return iter(self._entries) + + def __getitem__(self, name: str) -> ParameterPosterior: + """ + Look up a parameter's summary by name. + + Parameters + ---------- + name : str + The parameter name. + + Returns + ------- + ParameterPosterior + The summary for that parameter. + + Raises + ------ + KeyError + If no sampled parameter has that name. + """ + for entry in self._entries: + if entry.name == name: + return entry + raise KeyError(f'No sampled parameter named {name!r}.') + + def __repr__(self) -> str: + """ + Render the summary as a table. + + Returns + ------- + str + A table with the median and 68% credible interval of each parameter. + """ + if not self._entries: + return 'PosteriorSummary(no parameters)' + + header = ( + f'{"parameter":<28s} {"unit":>10s} {"median":>14s} ' + f'{"-":>12s} {"+":>12s} {"current":>14s}' + ) + lines = ['PosteriorSummary', header, '-' * len(header)] + lines.extend( + f'{e.name:<28s} {e.unit:>10s} {e.median:>14.5g} ' + f'{e.minus:>12.4g} {e.plus:>12.4g} {e.value:>14.5g}' + for e in self._entries + ) + return '\n'.join(lines) + + +def summarize_draws( + draws: np.ndarray, + fallback_names: list[str], + parameters_by_column: list[Parameter | None], +) -> PosteriorSummary: + """ + Summarize posterior draws under the parameters' own names and units. + + The sampler labels its columns with each parameter's ``unique_name`` (``Parameter_4`` and the + like), which is not what a user recognises, so columns are reported under ``Parameter.name`` + wherever a parameter could be matched. + + Parameters + ---------- + draws : np.ndarray + Posterior draws, shape ``(n_draws, n_parameters)``. + fallback_names : list[str] + Label to use for any column with no matching parameter, one per column. + parameters_by_column : list[Parameter | None] + The parameter for each column of ``draws``, or None where none could be matched. + + Returns + ------- + PosteriorSummary + One entry per column of ``draws``, in column order. + """ + entries = [] + for column, parameter in enumerate(parameters_by_column): + lower, median, upper = ( + float(percentile) for percentile in np.percentile(draws[:, column], [16, 50, 84]) + ) + entries.append( + ParameterPosterior( + name=fallback_names[column] if parameter is None else parameter.name, + unit='' if parameter is None else str(parameter.unit), + median=median, + lower=lower, + upper=upper, + value=float('nan') if parameter is None else float(parameter.value), + ) + ) + return PosteriorSummary(entries) + + +def _verify_nonneg_number(value: object, name: str) -> None: + """ + Raise if a value is not a non-negative number. + + Parameters + ---------- + value : object + The object to verify. + name : str + The name of the object, for the error message. + + Raises + ------ + TypeError + If value is not an int or float. + ValueError + If value is negative. + """ + if not isinstance(value, (int, float)) or isinstance(value, bool): + raise TypeError(f'{name} must be a number. Got {type(value)}.') + if value < 0: + raise ValueError(f'{name} must be non-negative. Got {value}.') diff --git a/src/easydynamics/utils/__init__.py b/src/easydynamics/utils/__init__.py index 5e644a06b..1c3402ced 100644 --- a/src/easydynamics/utils/__init__.py +++ b/src/easydynamics/utils/__init__.py @@ -3,5 +3,14 @@ from easydynamics.utils.detailed_balance import detailed_balance_factor from easydynamics.utils.plotting import slicerplot_with_residuals +from easydynamics.utils.posterior_plotting import plot_corner +from easydynamics.utils.posterior_plotting import plot_posterior_predictive +from easydynamics.utils.posterior_plotting import plot_trace -__all__ = ['detailed_balance_factor', 'slicerplot_with_residuals'] +__all__ = [ + 'detailed_balance_factor', + 'plot_corner', + 'plot_posterior_predictive', + 'plot_trace', + 'slicerplot_with_residuals', +] diff --git a/src/easydynamics/utils/posterior_plotting.py b/src/easydynamics/utils/posterior_plotting.py new file mode 100644 index 000000000..bfb1a1320 --- /dev/null +++ b/src/easydynamics/utils/posterior_plotting.py @@ -0,0 +1,256 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +""" +Diagnostic plots for Bayesian posterior samples. + +These take plain arrays rather than an Analysis, so they can be used on any chain, including one +loaded from disk. The Analysis classes wrap them in convenience methods. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import matplotlib.pyplot as plt +import numpy as np + +if TYPE_CHECKING: + from matplotlib.figure import Figure + + +def plot_trace( + draws: np.ndarray, + names: list[str], + logp: np.ndarray | None = None, + title: str | None = None, + figsize: tuple[float, float] | None = None, +) -> Figure: + """ + Plot the chain trace of every sampled parameter. + + A converged chain looks like a "hairy caterpillar": noisy but stationary, with no drift or long + excursions. A visible trend means the chain has not reached the typical set and needs a longer + burn-in. + + A ``ValueError`` is raised if ``draws`` is not two-dimensional, or if ``names`` does not have + one entry per column. + + Parameters + ---------- + draws : np.ndarray + Posterior draws, shape ``(n_draws, n_parameters)``. + names : list[str] + One label per column of ``draws``. + logp : np.ndarray | None, default=None + Log-posterior values, plotted in an extra panel when given. + title : str | None, default=None + Figure title. + figsize : tuple[float, float] | None, default=None + Figure size in inches. Defaults to a height that scales with the number of panels. + + Returns + ------- + Figure + The matplotlib Figure. + """ + draws = np.asarray(draws) + _verify_draws(draws, names) + + n_panels = draws.shape[1] + (1 if logp is not None else 0) + if figsize is None: + figsize = (10.0, max(2.0, 1.6 * n_panels)) + + fig, axes = plt.subplots(n_panels, 1, figsize=figsize, sharex=True, squeeze=False) + axes = axes[:, 0] + + for axis, column, name in zip(axes, range(draws.shape[1]), names, strict=False): + axis.plot(draws[:, column], lw=0.5) + axis.set_ylabel(name, fontsize=8) + axis.set_xlim(0, len(draws) - 1) + + if logp is not None: + axes[-1].plot(np.asarray(logp), lw=0.5, color='C4') + axes[-1].set_ylabel('log-posterior', fontsize=8) + + axes[-1].set_xlabel('sample index') + if title is not None: + fig.suptitle(title) + fig.tight_layout() + return fig + + +def plot_corner( + draws: np.ndarray, + names: list[str], + title: str | None = None, + bins: int = 40, + figsize: tuple[float, float] | None = None, +) -> Figure: + """ + Plot marginal and pairwise posterior distributions. + + Diagonal panels show each parameter's marginal distribution. Off-diagonal panels show the joint + distribution of a pair: a compact blob means the two are independent, while a narrow diagonal + ridge means they are correlated and cannot be determined separately from this data. + + A ``ValueError`` is raised if ``draws`` is not two-dimensional, or if ``names`` does not have + one entry per column. + + Parameters + ---------- + draws : np.ndarray + Posterior draws, shape ``(n_draws, n_parameters)``. + names : list[str] + One label per column of ``draws``. + title : str | None, default=None + Figure title. + bins : int, default=40 + Number of bins for the marginal histograms. + figsize : tuple[float, float] | None, default=None + Figure size in inches. Defaults to a square that scales with the parameter count. + + Returns + ------- + Figure + The matplotlib Figure. + """ + draws = np.asarray(draws) + _verify_draws(draws, names) + + n = draws.shape[1] + if figsize is None: + side = max(4.0, 2.0 * n) + figsize = (side, side) + + fig, axes = plt.subplots(n, n, figsize=figsize, squeeze=False) + for row in range(n): + for col in range(n): + axis = axes[row, col] + if col > row: + axis.set_visible(False) + continue + if row == col: + axis.hist(draws[:, row], bins=bins, color='C0', histtype='stepfilled', alpha=0.7) + axis.set_yticks([]) + else: + axis.hexbin(draws[:, col], draws[:, row], gridsize=30, cmap='Blues', mincnt=1) + if row == n - 1: + axis.set_xlabel(names[col], fontsize=8) + else: + axis.set_xticklabels([]) + if col == 0 and row != 0: + axis.set_ylabel(names[row], fontsize=8) + else: + axis.set_yticklabels([]) + axis.tick_params(labelsize=7) + + if title is not None: + fig.suptitle(title) + fig.tight_layout() + return fig + + +def plot_posterior_predictive( + x: np.ndarray, + y: np.ndarray, + predictions: np.ndarray, + y_err: np.ndarray | None = None, + title: str | None = None, + credible_interval: float = 68.0, + figsize: tuple[float, float] = (8.0, 5.0), +) -> Figure: + """ + Plot the data against the credible band implied by the posterior. + + The band shows where the model says the data should lie, given the posterior. If the data + strays outside it systematically, the model is missing something that no amount of parameter + tuning will fix. + + Parameters + ---------- + x : np.ndarray + Independent variable of the data. + y : np.ndarray + Observed values. + predictions : np.ndarray + Model evaluations, shape ``(n_draws, len(x))``, one row per posterior draw. + y_err : np.ndarray | None, default=None + Standard deviation of the observed values, drawn as error bars when given. + title : str | None, default=None + Figure title. + credible_interval : float, default=68.0 + Width of the credible band, as a percentage. + figsize : tuple[float, float], default=(8.0, 5.0) + Figure size in inches. + + Returns + ------- + Figure + The matplotlib Figure. + + Raises + ------ + ValueError + If ``predictions`` is not two-dimensional with one column per point in ``x``, or if + ``credible_interval`` is not between 0 and 100. + """ + x = np.asarray(x) + y = np.asarray(y) + predictions = np.asarray(predictions) + if predictions.ndim != 2 or predictions.shape[1] != len(x): + raise ValueError( + f'predictions must have shape (n_draws, {len(x)}). Got {predictions.shape}.' + ) + if not 0 < credible_interval < 100: + raise ValueError(f'credible_interval must be between 0 and 100. Got {credible_interval}.') + + tail = (100.0 - credible_interval) / 2.0 + lower, median, upper = np.percentile(predictions, [tail, 50.0, 100.0 - tail], axis=0) + + fig, axis = plt.subplots(figsize=figsize) + if y_err is None: + axis.plot(x, y, 'o', mfc='none', color='black', label='Data', markersize=4) + else: + axis.errorbar( + x, y, np.asarray(y_err), fmt='o', mfc='none', color='black', label='Data', markersize=4 + ) + axis.fill_between( + x, + lower, + upper, + color='C3', + alpha=0.3, + label=f'{credible_interval:.0f}% credible band', + ) + axis.plot(x, median, '-', color='C3', label='Posterior median') + axis.legend() + if title is not None: + axis.set_title(title) + fig.tight_layout() + return fig + + +def _verify_draws(draws: np.ndarray, names: list[str]) -> None: + """ + Verify that a draws array is two-dimensional and matches its labels. + + Parameters + ---------- + draws : np.ndarray + The posterior draws to check. + names : list[str] + The labels to check against. + + Raises + ------ + ValueError + If ``draws`` is not two-dimensional or its column count differs from ``len(names)``. + """ + if draws.ndim != 2: + raise ValueError(f'draws must be two-dimensional. Got shape {draws.shape}.') + if draws.shape[1] != len(names): + raise ValueError( + f'names must have one entry per column of draws. ' + f'Got {len(names)} names for {draws.shape[1]} columns.' + ) diff --git a/tests/integration/fitting/test_bayesian_sampling.py b/tests/integration/fitting/test_bayesian_sampling.py new file mode 100644 index 000000000..704a37ea2 --- /dev/null +++ b/tests/integration/fitting/test_bayesian_sampling.py @@ -0,0 +1,207 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +""" +Integration tests running real BUMPS DREAM chains through Analysis1d. + +These are slow by nature. They deliberately run with ``sampler_kwargs={'trim': False}``: BUMPS' +automatic burn-point trimming re-runs a convergence detector on every call and can crash inside its +own outlier removal on the very short chains used here. +""" + +import warnings + +import matplotlib as mpl +import numpy as np +import pytest +import scipp as sc + +mpl.use('Agg') + +from easydynamics.analysis.analysis1d import Analysis1d +from easydynamics.experiment import Experiment +from easydynamics.sample_model import InstrumentModel +from easydynamics.sample_model import SampleModel +from easydynamics.sample_model.components.gaussian import Gaussian + +TRUE_AREA = 9.0 +TRUE_WIDTH = 1.2 +NOISE = 0.05 + +# Keep the chains short enough to stay usable in CI; long enough to locate the peak. +SAMPLE_KWARGS = { + 'samples': 2000, + 'burn': 100, + 'thin': 2, + 'sampler_kwargs': {'trim': False}, +} + + +def build_analysis(): + energy_values = np.linspace(-5.0, 5.0, 60) + truth = TRUE_AREA / (TRUE_WIDTH * np.sqrt(2 * np.pi)) + truth = truth * np.exp(-0.5 * (energy_values / TRUE_WIDTH) ** 2) + observed = truth + np.random.default_rng(0).normal(0.0, NOISE, size=truth.shape) + + data = sc.array( + dims=['Q', 'energy'], + values=observed[None, :], + variances=np.full_like(observed, NOISE**2)[None, :], + ) + experiment = Experiment( + data=sc.DataArray( + data=data, + coords={ + 'Q': sc.array(dims=['Q'], values=[1.0], unit='1/Angstrom'), + 'energy': sc.array(dims=['energy'], values=energy_values, unit='meV'), + }, + ) + ) + analysis = Analysis1d( + display_name='BayesianIntegration', + experiment=experiment, + sample_model=SampleModel( + components=Gaussian(area=TRUE_AREA, width=TRUE_WIDTH, center=0.0) + ), + instrument_model=InstrumentModel(), + Q_index=0, + ) + # The energy offset shifts the spectrum exactly as the Gaussian centre does. Leaving both free + # makes the model unidentifiable, which no amount of sampling can repair. + analysis.instrument_model.fix_energy_offset(Q_index=0) + return analysis + + +@pytest.fixture(scope='module') +def sampled_analysis(): + analysis = build_analysis() + analysis.fit() + analysis.suggest_bounds().apply() + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + analysis.sample_posterior(**SAMPLE_KWARGS) + return analysis + + +class TestRealChain: + def test_chain_has_one_column_per_free_parameter(self, sampled_analysis): + # EXPECT + results = sampled_analysis.posterior_result + assert results.draws.shape[1] == len(sampled_analysis.get_free_parameters()) + assert results.draws.shape[0] > 0 + + @pytest.mark.parametrize( + ('name', 'truth'), + [('Gaussian area', TRUE_AREA), ('Gaussian width', TRUE_WIDTH)], + ) + def test_posterior_recovers_the_true_parameters(self, sampled_analysis, name, truth): + # WHEN + entry = sampled_analysis.posterior_summary()[name] + + # EXPECT the truth sits within a few posterior standard deviations of the median. A 68% + # interval is deliberately not used: it excludes the truth about a third of the time for + # any single noise realization, which would make this test flaky rather than strict. + spread = max(entry.minus, entry.plus) + assert abs(entry.median - truth) < 4 * spread + + def test_summary_is_reported_under_parameter_names_and_units(self, sampled_analysis): + # WHEN + summary = sampled_analysis.posterior_summary() + + # EXPECT + assert {entry.name for entry in summary} == { + p.name for p in sampled_analysis.get_free_parameters() + } + assert all(entry.unit == 'meV' for entry in summary) + + def test_sampling_leaves_the_fitted_values_untouched(self): + # WHEN + analysis = build_analysis() + analysis.fit() + analysis.suggest_bounds().apply() + before = [float(p.value) for p in analysis.get_free_parameters()] + + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + analysis.sample_posterior(**SAMPLE_KWARGS) + + # EXPECT + after = [float(p.value) for p in analysis.get_free_parameters()] + assert after == pytest.approx(before) + + def test_extend_grows_the_chain(self, sampled_analysis): + # WHEN + before = int(sampled_analysis.posterior_result.state.Ngen) + + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + extended = sampled_analysis.extend_sampling( + additional_samples=500, thin=2, sampler_kwargs={'trim': False} + ) + + # EXPECT + assert int(extended.state.Ngen) > before + + def test_save_and_load_round_trip_keeps_parameter_identity(self, sampled_analysis, tmp_path): + # WHEN + prefix = str(tmp_path / 'chain') + sampled_analysis.save_chain(prefix) + + fresh = build_analysis() + fresh.fit() + fresh.suggest_bounds().apply() + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + fresh.load_chain(prefix) + + # EXPECT the reloaded chain is reported under real names, not internal unique names + summary = fresh.posterior_summary() + assert {entry.name for entry in summary} == {p.name for p in fresh.get_free_parameters()} + assert all(np.isfinite(entry.value) for entry in summary) + + def test_subset_sampling_produces_a_single_column(self): + # WHEN + analysis = build_analysis() + analysis.fit() + analysis.suggest_bounds().apply() + + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + results = analysis.sample_posterior(parameters=['Gaussian width'], **SAMPLE_KWARGS) + + # EXPECT + assert results.draws.shape[1] == 1 + assert analysis.posterior_summary().entries[0].name == 'Gaussian width' + + def test_plots_render(self, sampled_analysis): + # WHEN + import matplotlib.pyplot as plt + + n_parameters = len(sampled_analysis.get_free_parameters()) + trace = sampled_analysis.plot_trace() + corner = sampled_analysis.plot_corner() + predictive = sampled_analysis.plot_posterior_predictive(n_draws=20) + + # EXPECT + assert len(trace.axes) == n_parameters + 1 + assert len(corner.axes) == n_parameters**2 + assert len(predictive.axes) == 1 + plt.close('all') + + def test_posterior_median_is_close_to_the_least_squares_fit(self): + # WHEN + analysis = build_analysis() + analysis.fit() + analysis.suggest_bounds().apply() + fitted = {p.name: float(p.value) for p in analysis.get_free_parameters()} + + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + analysis.sample_posterior(**SAMPLE_KWARGS) + summary = analysis.posterior_summary() + + # EXPECT the two agree within the posterior's own uncertainty, since with flat priors the + # maximum-likelihood point sits inside the bulk of the posterior + for entry in summary: + spread = max(entry.minus, entry.plus) + assert abs(entry.median - fitted[entry.name]) < 5 * spread diff --git a/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py b/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py new file mode 100644 index 000000000..b29a2b54a --- /dev/null +++ b/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py @@ -0,0 +1,486 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +"""Unit tests for Bayesian sampling on Analysis1d, with the EasyScience Sampler mocked out.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock +from unittest.mock import patch + +import numpy as np +import pytest +import scipp as sc +from easyscience.fitting import AvailableMinimizers + +from easydynamics.analysis.analysis1d import Analysis1d +from easydynamics.experiment import Experiment +from easydynamics.sample_model import InstrumentModel +from easydynamics.sample_model import SampleModel +from easydynamics.sample_model.components.gaussian import Gaussian + +SAMPLER_PATH = 'easydynamics.analysis.bayesian_sampling.Sampler' + + +def make_analysis(): + energy_values = np.linspace(-5.0, 5.0, 20) + intensity = 3.0 * np.exp(-0.5 * (energy_values / 1.2) ** 2) + data = sc.array( + dims=['Q', 'energy'], + values=intensity[None, :], + variances=np.full_like(intensity, 0.01)[None, :], + ) + experiment = Experiment( + data=sc.DataArray( + data=data, + coords={ + 'Q': sc.array(dims=['Q'], values=[1.0], unit='1/Angstrom'), + 'energy': sc.array(dims=['energy'], values=energy_values, unit='meV'), + }, + ) + ) + analysis = Analysis1d( + display_name='TestBayesian', + experiment=experiment, + sample_model=SampleModel(components=Gaussian(area=3.0, width=1.2, center=0.0)), + instrument_model=InstrumentModel(), + Q_index=0, + ) + analysis.instrument_model.fix_energy_offset(Q_index=0) + return analysis + + +def bound_all(analysis, half_width=5.0): + """Give every free parameter finite bounds so the pre-flight passes.""" + for parameter in analysis.get_free_parameters(): + parameter.min = float(parameter.value) - half_width + parameter.max = float(parameter.value) + half_width + + +def fake_results(analysis, n_draws=100, values=None): + """Build a SamplingResults-shaped object for the free parameters of an analysis.""" + parameters = analysis.get_free_parameters() + if values is None: + draws = np.tile([float(p.value) for p in parameters], (n_draws, 1)) + else: + draws = np.asarray(values, dtype=float) + return SimpleNamespace( + draws=draws, + param_names=[p.unique_name for p in parameters], + logp=np.zeros(draws.shape[0]), + state=MagicMock(Ngen=10, Npop=4), + ) + + +@pytest.fixture +def analysis(): + return make_analysis() + + +class TestFitterExposure: + def test_fitter_is_built_lazily_and_cached(self, analysis): + # WHEN + fitter = analysis.fitter + + # EXPECT + assert fitter is analysis.fitter + assert fitter.fit_object is analysis + + def test_fitter_is_rebuilt_when_the_sample_model_changes(self, analysis): + # WHEN + original = analysis.fitter + analysis.sample_model = SampleModel(components=Gaussian(area=1.0)) + + # EXPECT + assert analysis.fitter is not original + + def test_minimizer_can_be_switched_through_the_fitter(self, analysis): + # WHEN + analysis.fitter.switch_minimizer(AvailableMinimizers.Bumps) + + # EXPECT + assert analysis.fitter.minimizer.enum == AvailableMinimizers.Bumps + + def test_fit_uses_the_persistent_fitter(self, analysis): + # WHEN + result = analysis.fit() + + # EXPECT + assert result is analysis._fit_result + assert np.isfinite(result.reduced_chi2) + + +class TestBoundsPreflight: + def test_sampling_refuses_unbounded_parameters(self, analysis): + # EXPECT + with pytest.raises(ValueError, match='finite bounds'): + analysis.sample_posterior(samples=10) + + def test_error_names_the_offending_parameters(self, analysis): + # EXPECT + with pytest.raises(ValueError, match='Gaussian area'): + analysis.check_bounds_for_sampling() + + def test_bounded_parameters_pass(self, analysis): + # WHEN + bound_all(analysis) + + # EXPECT: does not raise + analysis.check_bounds_for_sampling() + + def test_suggest_bounds_covers_the_free_parameters(self, analysis): + # WHEN + suggestions = analysis.suggest_bounds() + + # EXPECT + assert len(suggestions) == len(analysis.get_free_parameters()) + + +class TestSamplePosterior: + def test_restores_parameter_values_and_minimizer(self, analysis): + # WHEN + bound_all(analysis) + before = [(p.unique_name, p.value) for p in analysis.get_free_parameters()] + + with patch(SAMPLER_PATH) as sampler_class: + + def mutate_then_return(**_kwargs): + # The real sampler leaves the parameters wherever the last evaluation put them. + for parameter in analysis.get_free_parameters(): + parameter.value = float(parameter.value) + 1.0 + return fake_results(analysis) + + sampler_class.return_value.sample.side_effect = mutate_then_return + analysis.sample_posterior(samples=10, burn=1, thin=1) + + # EXPECT + after = [(p.unique_name, p.value) for p in analysis.get_free_parameters()] + assert after == before + assert analysis.fitter.minimizer.enum == AvailableMinimizers.LMFit_leastsq + + def test_switches_to_bumps_for_the_run(self, analysis): + # WHEN + bound_all(analysis) + seen = [] + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: ( + seen.append(analysis.fitter.minimizer.enum), + fake_results(analysis), + )[1] + analysis.sample_posterior(samples=10) + + # EXPECT + assert seen == [AvailableMinimizers.Bumps] + + def test_restores_the_minimizer_even_when_sampling_raises(self, analysis): + # WHEN + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = RuntimeError('boom') + with pytest.raises(RuntimeError, match='boom'): + analysis.sample_posterior(samples=10) + + # EXPECT + assert analysis.fitter.minimizer.enum == AvailableMinimizers.LMFit_leastsq + + def test_forwards_sampling_arguments(self, analysis): + # WHEN + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + analysis.sample_posterior(samples=123, burn=7, thin=3, population=5) + + # EXPECT + kwargs = sampler_class.return_value.sample.call_args.kwargs + assert kwargs['samples'] == 123 + assert kwargs['burn'] == 7 + assert kwargs['thin'] == 3 + assert kwargs['population'] == 5 + + def test_stores_the_result(self, analysis): + # WHEN + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + expected = fake_results(analysis) + sampler_class.return_value.sample.return_value = expected + returned = analysis.sample_posterior(samples=10) + + # EXPECT + assert returned is expected + assert analysis.posterior_result is expected + + def test_warns_when_the_posterior_piles_up_against_a_bound(self, analysis): + # WHEN a parameter's draws span its whole allowed range + bound_all(analysis) + parameters = analysis.get_free_parameters() + draws = np.tile([float(p.value) for p in parameters], (500, 1)) + draws[:, 0] = np.linspace(parameters[0].min, parameters[0].max, 500) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_results(analysis, values=draws) + + # EXPECT + with pytest.warns(UserWarning, match='piled up'): + analysis.sample_posterior(samples=10) + + def test_does_not_warn_when_the_posterior_is_well_inside(self, analysis): + # WHEN + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_results(analysis) + + # EXPECT + with warnings_as_errors(): + analysis.sample_posterior(samples=10) + + +class TestParameterSubset: + def test_holds_other_parameters_fixed_during_the_run(self, analysis): + # WHEN + bound_all(analysis) + target = analysis.get_free_parameters()[0] + seen = {} + + with patch(SAMPLER_PATH) as sampler_class: + + def record(**_kwargs): + seen['free'] = [p.unique_name for p in analysis.get_free_parameters()] + return fake_results(analysis) + + sampler_class.return_value.sample.side_effect = record + with pytest.warns(UserWarning, match='Holding these parameters fixed'): + analysis.sample_posterior(samples=10, parameters=[target.name]) + + # EXPECT + assert seen['free'] == [target.unique_name] + + def test_restores_the_fixed_flags_afterwards(self, analysis): + # WHEN + bound_all(analysis) + before = [(p.unique_name, p.fixed) for p in analysis.get_all_parameters()] + target = analysis.get_free_parameters()[0] + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + with pytest.warns(UserWarning): + analysis.sample_posterior(samples=10, parameters=[target]) + + # EXPECT + assert [(p.unique_name, p.fixed) for p in analysis.get_all_parameters()] == before + + def test_unknown_parameter_name_raises(self, analysis): + # WHEN + bound_all(analysis) + + # EXPECT + with pytest.raises(ValueError, match='No free parameter named'): + analysis.sample_posterior(samples=10, parameters=['not a parameter']) + + def test_non_list_parameters_raises(self, analysis): + # EXPECT + with pytest.raises(TypeError, match='must be a list'): + analysis.sample_posterior(samples=10, parameters='Gaussian area') + + def test_empty_parameter_list_raises(self, analysis): + # EXPECT + with pytest.raises(ValueError, match='at least one parameter'): + analysis.sample_posterior(samples=10, parameters=[]) + + +class TestSamplerCaching: + def test_sampler_is_reused_between_runs(self, analysis): + # WHEN + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + analysis.sample_posterior(samples=10) + analysis.sample_posterior(samples=10) + + # EXPECT the data is bound once, not per run + assert sampler_class.call_count == 1 + + def test_changing_the_q_index_rebuilds_the_sampler(self, analysis): + # WHEN + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + analysis.sample_posterior(samples=10) + analysis.Q_index = 0 + analysis.sample_posterior(samples=10) + + # EXPECT the Sampler binds its data at construction, so it must be rebuilt + assert sampler_class.call_count == 2 + + def test_binds_the_same_data_the_fit_uses(self, analysis): + # WHEN + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + analysis.sample_posterior(samples=10) + + # EXPECT + expected_x, expected_y, expected_w = analysis._get_sampling_data() + args, kwargs = sampler_class.call_args + assert np.array_equal(args[1], expected_x) + assert np.array_equal(args[2], expected_y) + assert np.array_equal(kwargs['weights'], expected_w) + + +class TestExtendAndPersistence: + def test_extend_without_a_chain_raises(self, analysis): + # EXPECT + with pytest.raises(RuntimeError, match='No chain to extend'): + analysis.extend_sampling() + + def test_extend_delegates_to_the_sampler(self, analysis): + # WHEN + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + sampler_class.return_value.extend.side_effect = lambda **_k: fake_results(analysis) + analysis.sample_posterior(samples=10) + analysis.extend_sampling(additional_samples=42, thin=2) + + # EXPECT + kwargs = sampler_class.return_value.extend.call_args.kwargs + assert kwargs['additional_samples'] == 42 + assert kwargs['thin'] == 2 + + def test_save_without_a_chain_raises(self, analysis): + # EXPECT + with pytest.raises(RuntimeError, match='No chain to save'): + analysis.save_chain('somewhere') + + def test_save_writes_the_parameter_name_sidecar(self, analysis, tmp_path): + # WHEN + import json + + bound_all(analysis) + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + analysis.sample_posterior(samples=10) + analysis.save_chain(str(tmp_path / 'chain')) + + # EXPECT the unique names are recorded against the stable parameter names + sidecar = tmp_path / 'chain.parameter-names.json' + assert sidecar.is_file() + mapping = json.loads(sidecar.read_text(encoding='utf-8')) + assert set(mapping.values()) == {p.name for p in analysis.get_free_parameters()} + + def test_load_without_a_sidecar_warns(self, analysis, tmp_path): + # WHEN + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.load_state.return_value = fake_results(analysis) + + # EXPECT + with pytest.warns(UserWarning, match='No parameter-name sidecar'): + analysis.load_chain(str(tmp_path / 'missing')) + + +class TestResults: + def test_summary_without_sampling_raises(self, analysis): + # EXPECT + with pytest.raises(RuntimeError, match='No posterior samples yet'): + analysis.posterior_summary() + + def test_summary_uses_parameter_names_and_units(self, analysis): + # WHEN + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + analysis.sample_posterior(samples=10) + + # EXPECT + summary = analysis.posterior_summary() + names = {entry.name for entry in summary} + assert names == {p.name for p in analysis.get_free_parameters()} + assert all(entry.unit == 'meV' for entry in summary) + + def test_set_parameters_to_posterior_median(self, analysis): + # WHEN + bound_all(analysis) + parameters = analysis.get_free_parameters() + draws = np.tile([float(p.value) + 2.0 for p in parameters], (50, 1)) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_results(analysis, values=draws) + expected = [float(p.value) + 2.0 for p in parameters] + analysis.sample_posterior(samples=10) + + changed = analysis.set_parameters_to_posterior_median() + + # EXPECT + assert len(changed) == len(parameters) + assert [float(p.value) for p in parameters] == pytest.approx(expected) + + def test_median_without_sampling_raises(self, analysis): + # EXPECT + with pytest.raises(RuntimeError, match='No posterior samples yet'): + analysis.set_parameters_to_posterior_median() + + +class TestPlots: + def test_predictive_rejects_a_bad_draw_count(self, analysis): + # WHEN + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + analysis.sample_posterior(samples=10) + + # EXPECT + with pytest.raises(ValueError, match='positive integer'): + analysis.plot_posterior_predictive(n_draws=0) + + def test_predictive_restores_parameter_values(self, analysis): + # WHEN + bound_all(analysis) + parameters = analysis.get_free_parameters() + draws = np.tile([float(p.value) + 0.5 for p in parameters], (20, 1)) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_results(analysis, values=draws) + analysis.sample_posterior(samples=10) + + before = [float(p.value) for p in parameters] + analysis.plot_posterior_predictive(n_draws=5) + + # EXPECT + assert [float(p.value) for p in parameters] == pytest.approx(before) + + def test_plots_without_sampling_raise(self, analysis): + # EXPECT + with pytest.raises(RuntimeError): + analysis.plot_trace() + with pytest.raises(RuntimeError): + analysis.plot_corner() + + +class warnings_as_errors: + """Context manager asserting that no UserWarning is emitted inside the block.""" + + def __enter__(self): + import warnings + + self._ctx = warnings.catch_warnings(record=True) + self._caught = self._ctx.__enter__() + warnings.simplefilter('always') + return self + + def __exit__(self, *exc_info): + caught = [w for w in self._caught if issubclass(w.category, UserWarning)] + self._ctx.__exit__(*exc_info) + if exc_info[0] is None: + assert not caught, f'unexpected warnings: {[str(w.message) for w in caught]}' + return False diff --git a/tests/unit/easydynamics/analysis/test_posterior.py b/tests/unit/easydynamics/analysis/test_posterior.py new file mode 100644 index 000000000..88dc29bf6 --- /dev/null +++ b/tests/unit/easydynamics/analysis/test_posterior.py @@ -0,0 +1,296 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +import numpy as np +import pytest +from easyscience.variable import Parameter + +from easydynamics.analysis.posterior import BoundsSuggestion +from easydynamics.analysis.posterior import BoundsSuggestions +from easydynamics.analysis.posterior import parameters_at_bounds +from easydynamics.analysis.posterior import suggest_bounds_for_parameters +from easydynamics.analysis.posterior import summarize_draws +from easydynamics.analysis.posterior import unbounded_parameters + + +def make_parameter(name='p', value=1.0, error=0.0, minimum=-np.inf, maximum=np.inf, unit='meV'): + parameter = Parameter(name=name, value=value, unit=unit) + parameter.min = minimum + parameter.max = maximum + if error: + parameter.variance = error**2 + return parameter + + +class TestSuggestBounds: + def test_fills_in_both_infinite_sides(self): + # WHEN + parameter = make_parameter(value=10.0, error=0.5) + + # THEN + suggestions = suggest_bounds_for_parameters([parameter], n_sigma=10.0, relative_pad=0.2) + + # EXPECT: 10 * 0.5 + 0.2 * 10 = 7 + suggestion = suggestions.suggestions[0] + assert suggestion.suggested_min == pytest.approx(3.0) + assert suggestion.suggested_max == pytest.approx(17.0) + assert not suggestion.needs_attention + + def test_never_loosens_an_existing_finite_bound(self): + # WHEN a physical lower bound is already set + parameter = make_parameter(value=1.2, error=1.5, minimum=1e-10) + + # THEN + suggestion = suggest_bounds_for_parameters([parameter]).suggestions[0] + + # EXPECT the finite side survives untouched, even though the sigma rule would go negative + assert suggestion.suggested_min == pytest.approx(1e-10) + assert suggestion.suggested_max > 1.2 + + def test_fully_bounded_parameter_is_left_alone(self): + # WHEN + parameter = make_parameter(value=1.0, error=0.1, minimum=0.0, maximum=2.0) + + # THEN + suggestion = suggest_bounds_for_parameters([parameter]).suggestions[0] + + # EXPECT + assert suggestion.suggested_min == pytest.approx(0.0) + assert suggestion.suggested_max == pytest.approx(2.0) + assert not suggestion.changes_bounds + + def test_zero_error_falls_back_to_the_relative_pad(self): + # WHEN a minimizer reports no uncertainty at all + parameter = make_parameter(value=4.0, error=0.0) + + # THEN + suggestion = suggest_bounds_for_parameters([parameter], relative_pad=0.25).suggestions[0] + + # EXPECT the pad still yields a usable width + assert suggestion.suggested_min == pytest.approx(3.0) + assert suggestion.suggested_max == pytest.approx(5.0) + assert not suggestion.needs_attention + + def test_zero_value_and_zero_error_is_flagged_not_guessed(self): + # WHEN there is no scale information anywhere + parameter = make_parameter(value=0.0, error=0.0) + + # THEN + suggestion = suggest_bounds_for_parameters([parameter]).suggestions[0] + + # EXPECT + assert suggestion.needs_attention + assert 'no scale information' in suggestion.reason + assert not np.isfinite(suggestion.suggested_min) + + def test_absolute_floor_rescues_a_scaleless_parameter(self): + # WHEN + parameter = make_parameter(value=0.0, error=0.0) + + # THEN + suggestion = suggest_bounds_for_parameters([parameter], absolute_floor=0.5).suggestions[0] + + # EXPECT + assert not suggestion.needs_attention + assert suggestion.suggested_min == pytest.approx(-0.5) + assert suggestion.suggested_max == pytest.approx(0.5) + + def test_non_finite_value_is_flagged(self): + # WHEN + parameter = make_parameter(value=1.0) + parameter.value = np.inf + + # THEN + suggestion = suggest_bounds_for_parameters([parameter]).suggestions[0] + + # EXPECT + assert suggestion.needs_attention + assert 'not finite' in suggestion.reason + + @pytest.mark.parametrize('kwargs', [{'n_sigma': -1.0}, {'relative_pad': -0.1}]) + def test_negative_settings_raise(self, kwargs): + # EXPECT + with pytest.raises(ValueError): + suggest_bounds_for_parameters([make_parameter()], **kwargs) + + def test_non_numeric_setting_raises(self): + # EXPECT + with pytest.raises(TypeError): + suggest_bounds_for_parameters([make_parameter()], n_sigma='wide') + + +class TestBoundsSuggestionsApply: + def test_apply_sets_bounds_and_reports_changes(self): + # WHEN + parameter = make_parameter(value=10.0, error=0.5) + suggestions = suggest_bounds_for_parameters([parameter]) + + # THEN nothing has changed until apply is called + assert parameter.max == np.inf + changed = suggestions.apply() + + # EXPECT + assert changed == [parameter] + assert parameter.min == pytest.approx(3.0) + assert parameter.max == pytest.approx(17.0) + + def test_apply_skips_parameters_needing_attention(self): + # WHEN + parameter = make_parameter(value=0.0, error=0.0) + suggestions = suggest_bounds_for_parameters([parameter]) + + # THEN + changed = suggestions.apply() + + # EXPECT the unusable suggestion is skipped rather than written + assert changed == [] + assert parameter.min == -np.inf + + def test_repr_lists_parameters_and_flags_attention(self): + # WHEN + good = make_parameter(name='good', value=10.0, error=0.5) + bad = make_parameter(name='bad', value=0.0, error=0.0) + + # THEN + text = repr(suggest_bounds_for_parameters([good, bad])) + + # EXPECT + assert 'good' in text + assert 'bad' in text + assert 'need bounds set by hand' in text + + def test_repr_with_no_parameters(self): + # EXPECT + assert 'no free parameters' in repr(BoundsSuggestions([])) + + def test_len_and_iteration(self): + # WHEN + suggestions = suggest_bounds_for_parameters([make_parameter(), make_parameter()]) + + # EXPECT + assert len(suggestions) == 2 + assert all(isinstance(s, BoundsSuggestion) for s in suggestions) + + +class TestUnboundedParameters: + def test_finds_parameters_with_an_infinite_side(self): + # WHEN + bounded = make_parameter(name='bounded', minimum=0.0, maximum=1.0) + half_open = make_parameter(name='half_open', minimum=0.0) + + # THEN + result = unbounded_parameters([bounded, half_open]) + + # EXPECT + assert result == [half_open] + + +class TestParametersAtBounds: + def test_uniform_posterior_across_the_bounds_is_reported(self): + # WHEN a posterior fills its whole allowed range, the bound is setting the interval + parameter = make_parameter(minimum=0.0, maximum=1.0) + draws = np.linspace(0.0, 1.0, 1000).reshape(-1, 1) + + # THEN + result = parameters_at_bounds(draws, [parameter]) + + # EXPECT + assert parameter.name in result + assert result[parameter.name] == pytest.approx(0.1, abs=0.01) + + def test_posterior_well_inside_its_bounds_is_not_reported(self): + # WHEN + parameter = make_parameter(minimum=0.0, maximum=1.0) + draws = np.random.default_rng(0).normal(0.5, 0.02, size=1000).reshape(-1, 1) + + # THEN + result = parameters_at_bounds(draws, [parameter]) + + # EXPECT + assert result == {} + + def test_partly_clipped_posterior_is_reported(self): + # WHEN a posterior fills most, but not all, of its allowed range. A real bound-limited + # chain looks like this rather than perfectly uniform, so the threshold has to catch it. + parameter = make_parameter(minimum=0.0, maximum=1.0) + draws = np.linspace(0.02, 0.98, 1000).reshape(-1, 1) + + # THEN + result = parameters_at_bounds(draws, [parameter]) + + # EXPECT + assert parameter.name in result + + def test_posterior_pinned_at_one_bound_is_reported(self): + # WHEN + parameter = make_parameter(minimum=0.0, maximum=1.0) + draws = np.abs(np.random.default_rng(0).normal(0.0, 0.02, size=1000)).reshape(-1, 1) + + # THEN + result = parameters_at_bounds(draws, [parameter]) + + # EXPECT + assert result[parameter.name] > 0.9 + + def test_unmatched_and_unbounded_columns_are_skipped(self): + # WHEN + unbounded = make_parameter() + draws = np.zeros((10, 2)) + + # THEN + result = parameters_at_bounds(draws, [None, unbounded]) + + # EXPECT + assert result == {} + + +class TestSummarizeDraws: + def test_reports_parameter_names_units_and_percentiles(self): + # WHEN + parameter = make_parameter(name='Gaussian width', value=1.5, unit='meV') + draws = np.linspace(0.0, 100.0, 101).reshape(-1, 1) + + # THEN + summary = summarize_draws(draws, ['Parameter_0'], [parameter]) + + # EXPECT + entry = summary['Gaussian width'] + assert entry.unit == 'meV' + assert entry.median == pytest.approx(50.0) + assert entry.lower == pytest.approx(16.0) + assert entry.upper == pytest.approx(84.0) + assert entry.minus == pytest.approx(34.0) + assert entry.plus == pytest.approx(34.0) + assert entry.value == pytest.approx(1.5) + + def test_unmatched_column_falls_back_to_the_supplied_name(self): + # WHEN + draws = np.zeros((10, 1)) + + # THEN + summary = summarize_draws(draws, ['Parameter_7'], [None]) + + # EXPECT + entry = summary.entries[0] + assert entry.name == 'Parameter_7' + assert entry.unit == '' + assert np.isnan(entry.value) + + def test_lookup_of_missing_name_raises(self): + # WHEN + summary = summarize_draws(np.zeros((5, 1)), ['x'], [None]) + + # EXPECT + with pytest.raises(KeyError): + summary['not a parameter'] + + def test_repr_contains_the_parameter_name(self): + # WHEN + parameter = make_parameter(name='Gaussian area') + + # THEN + text = repr(summarize_draws(np.zeros((5, 1)), ['x'], [parameter])) + + # EXPECT + assert 'Gaussian area' in text + assert 'median' in text diff --git a/tests/unit/easydynamics/utils/test_posterior_plotting.py b/tests/unit/easydynamics/utils/test_posterior_plotting.py new file mode 100644 index 000000000..64d0129e9 --- /dev/null +++ b/tests/unit/easydynamics/utils/test_posterior_plotting.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +import matplotlib as mpl +import numpy as np +import pytest + +mpl.use('Agg') + +import matplotlib.pyplot as plt + +from easydynamics.utils.posterior_plotting import plot_corner +from easydynamics.utils.posterior_plotting import plot_posterior_predictive +from easydynamics.utils.posterior_plotting import plot_trace + + +@pytest.fixture(autouse=True) +def close_figures(): + yield + plt.close('all') + + +@pytest.fixture +def draws(): + return np.random.default_rng(0).normal(size=(200, 3)) + + +class TestPlotTrace: + def test_one_panel_per_parameter(self, draws): + # WHEN + fig = plot_trace(draws=draws, names=['a', 'b', 'c']) + + # EXPECT + assert len(fig.axes) == 3 + + def test_logp_adds_a_panel(self, draws): + # WHEN + fig = plot_trace(draws=draws, names=['a', 'b', 'c'], logp=np.zeros(len(draws))) + + # EXPECT + assert len(fig.axes) == 4 + assert fig.axes[-1].get_ylabel() == 'log-posterior' + + def test_names_label_the_panels(self, draws): + # WHEN + fig = plot_trace(draws=draws, names=['alpha', 'beta', 'gamma']) + + # EXPECT + assert [axis.get_ylabel() for axis in fig.axes] == ['alpha', 'beta', 'gamma'] + + def test_single_parameter_works(self): + # WHEN + fig = plot_trace(draws=np.zeros((10, 1)), names=['only']) + + # EXPECT + assert len(fig.axes) == 1 + + def test_mismatched_names_raise(self, draws): + # EXPECT + with pytest.raises(ValueError, match='one entry per column'): + plot_trace(draws=draws, names=['a', 'b']) + + def test_one_dimensional_draws_raise(self): + # EXPECT + with pytest.raises(ValueError, match='two-dimensional'): + plot_trace(draws=np.zeros(10), names=['a']) + + +class TestPlotCorner: + def test_grid_is_square_in_the_parameter_count(self, draws): + # WHEN + fig = plot_corner(draws=draws, names=['a', 'b', 'c']) + + # EXPECT + assert len(fig.axes) == 9 + + def test_upper_triangle_is_hidden(self, draws): + # WHEN + fig = plot_corner(draws=draws, names=['a', 'b', 'c']) + + # EXPECT: 3 hidden panels above the diagonal of a 3x3 grid + assert sum(not axis.get_visible() for axis in fig.axes) == 3 + + def test_mismatched_names_raise(self, draws): + # EXPECT + with pytest.raises(ValueError, match='one entry per column'): + plot_corner(draws=draws, names=['a']) + + +class TestPlotPosteriorPredictive: + def test_returns_a_figure_with_data_and_band(self): + # WHEN + x = np.linspace(0.0, 1.0, 25) + predictions = np.random.default_rng(0).normal(size=(50, 25)) + + fig = plot_posterior_predictive(x=x, y=np.zeros(25), predictions=predictions) + + # EXPECT + labels = [text.get_text() for text in fig.axes[0].get_legend().get_texts()] + assert 'Data' in labels + assert any('credible band' in label for label in labels) + + def test_error_bars_are_drawn_when_given(self): + # WHEN + x = np.linspace(0.0, 1.0, 10) + + fig = plot_posterior_predictive( + x=x, + y=np.zeros(10), + predictions=np.zeros((5, 10)), + y_err=np.full(10, 0.1), + ) + + # EXPECT + assert len(fig.axes[0].containers) == 1 + + def test_wrong_prediction_shape_raises(self): + # EXPECT + with pytest.raises(ValueError, match='predictions must have shape'): + plot_posterior_predictive(x=np.zeros(10), y=np.zeros(10), predictions=np.zeros((5, 3))) + + @pytest.mark.parametrize('interval', [0.0, 100.0, -5.0]) + def test_invalid_credible_interval_raises(self, interval): + # EXPECT + with pytest.raises(ValueError, match='credible_interval'): + plot_posterior_predictive( + x=np.zeros(4), + y=np.zeros(4), + predictions=np.zeros((5, 4)), + credible_interval=interval, + ) + + def test_band_widens_with_the_credible_interval(self): + # WHEN + x = np.linspace(0.0, 1.0, 8) + predictions = np.random.default_rng(0).normal(size=(400, 8)) + + narrow = plot_posterior_predictive( + x=x, y=np.zeros(8), predictions=predictions, credible_interval=50.0 + ) + wide = plot_posterior_predictive( + x=x, y=np.zeros(8), predictions=predictions, credible_interval=95.0 + ) + + # EXPECT + narrow_span = narrow.axes[0].collections[0].get_paths()[0].get_extents().height + wide_span = wide.axes[0].collections[0].get_paths()[0].get_extents().height + assert wide_span > narrow_span From bb1062308945fb165518e4cb13248eeb13011d5f Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Fri, 14 Aug 2026 13:07:35 +0200 Subject: [PATCH 2/3] Compose the posterior sampler instead of mixing it in Review feedback: bayesian_sampling.py had a lot in it that belonged elsewhere, and it was unclear why it was a mixin at all. It was a mixin because ParameterAnalysis is not an AnalysisBase and fits its binding models rather than itself, so a shared base class does not work. That was a reason, not a good one: it injected some forty methods into every Analysis class. The sampler is now composed. An Analysis exposes one `bayesian` property, and hands the sampler the few things that differ between the Analysis classes -- the data, the free parameters, their labels, and a hook to refresh cached computation -- so PosteriorSampler needs no knowledge of how any Analysis is built, and no Analysis inherits sampling machinery it does not use. Labelling moves to posterior_labels.py. Building it once for a fixed set of parameters also removes the quadratic cost the old code needed a scoped cache to avoid: the counts and lookups are computed in the constructor rather than per column. Plotting stays in posterior_plotting.py, where it already lived. The sampler keeps three short delegates so a chain can still be plotted from the object holding it, but none of the drawing happens there. The public API becomes analysis.bayesian.sample() and friends, and the explicit suggest_bounds().apply() step stays: in DREAM the bounds are the prior, and an unbounded parameter gives a confident-looking interval set by nothing. Co-Authored-By: Claude Opus 5 (1M context) --- docs/docs/tutorials/bayesian.ipynb | 30 +- src/easydynamics/analysis/__init__.py | 6 +- src/easydynamics/analysis/analysis1d.py | 89 +- .../analysis/bayesian_sampling.py | 993 ------------------ src/easydynamics/analysis/posterior.py | 37 +- src/easydynamics/analysis/posterior_labels.py | 182 ++++ .../analysis/posterior_sampling.py | 933 ++++++++++++++++ src/easydynamics/utils/posterior_plotting.py | 118 ++- .../fitting/test_bayesian_sampling.py | 44 +- .../analysis/test_analysis1d_bayesian.py | 86 +- .../easydynamics/analysis/test_posterior.py | 15 +- .../analysis/test_posterior_labels.py | 109 ++ 12 files changed, 1537 insertions(+), 1105 deletions(-) delete mode 100644 src/easydynamics/analysis/bayesian_sampling.py create mode 100644 src/easydynamics/analysis/posterior_labels.py create mode 100644 src/easydynamics/analysis/posterior_sampling.py create mode 100644 tests/unit/easydynamics/analysis/test_posterior_labels.py diff --git a/docs/docs/tutorials/bayesian.ipynb b/docs/docs/tutorials/bayesian.ipynb index 44a415147..51a68acd2 100644 --- a/docs/docs/tutorials/bayesian.ipynb +++ b/docs/docs/tutorials/bayesian.ipynb @@ -11,7 +11,7 @@ "\n", "A **Bayesian** analysis answers a different question: instead of one best point, it maps out the whole *posterior distribution* over the parameters. From that you can read off credible intervals that stay honest when parameters are correlated or their distributions are skewed, and you can see the correlations directly.\n", "\n", - "EasyDynamics does this with the DREAM sampler from [BUMPS](https://bumps.readthedocs.io/), through `sample_posterior()`." + "EasyDynamics does this with the DREAM sampler from [BUMPS](https://bumps.readthedocs.io/), through `bayesian.sample()`." ] }, { @@ -102,9 +102,9 @@ "source": [ "## Bounds are the prior\n", "\n", - "In DREAM, each parameter's `min` and `max` define a uniform prior, so **every free parameter must have finite bounds** before sampling. Most parameters start with at least one infinite bound, so `sample_posterior()` would refuse to run.\n", + "In DREAM, each parameter's `min` and `max` define a uniform prior, so **every free parameter must have finite bounds** before sampling. Most parameters start with at least one infinite bound, so `bayesian.sample()` would refuse to run.\n", "\n", - "`suggest_bounds()` proposes bounds from the fitted values and uncertainties. It is advisory: it changes nothing until you call `.apply()`, and it only ever fills in an *infinite* bound, so physical limits you have already set (an area that cannot go below zero, say) are left alone." + "`bayesian.suggest_bounds()` proposes bounds from the fitted values and uncertainties. It is advisory: it changes nothing until you call `.apply()`, and it only ever fills in an *infinite* bound, so physical limits you have already set (an area that cannot go below zero, say) are left alone." ] }, { @@ -114,7 +114,7 @@ "metadata": {}, "outputs": [], "source": [ - "suggestions = analysis.suggest_bounds()\n", + "suggestions = analysis.bayesian.suggest_bounds()\n", "print(suggestions)" ] }, @@ -146,7 +146,7 @@ "source": [ "## Sample the posterior\n", "\n", - "`sample_posterior()` runs the chains. The three numbers that matter are:\n", + "`bayesian.sample()` runs the chains. The three numbers that matter are:\n", "\n", "- `samples` — how many draws to collect in total. More is better, at linear cost.\n", "- `burn` — generations discarded at the start, while the chains are still travelling towards the bulk of the posterior.\n", @@ -162,7 +162,7 @@ "metadata": {}, "outputs": [], "source": [ - "results = analysis.sample_posterior(samples=4000, burn=300, thin=2)\n", + "results = analysis.bayesian.sample(samples=4000, burn=300, thin=2)\n", "\n", "print(f'Collected {results.draws.shape[0]} draws for {results.draws.shape[1]} parameters.')" ] @@ -184,7 +184,7 @@ "metadata": {}, "outputs": [], "source": [ - "analysis.plot_trace()" + "analysis.bayesian.plot_trace()" ] }, { @@ -194,7 +194,7 @@ "source": [ "## Summarize the posterior\n", "\n", - "`posterior_summary()` reports the median and the 68% credible interval of each parameter, under the parameter's own name and unit. The interval is asymmetric in general, which is precisely the information a single symmetric error bar throws away." + "`bayesian.summary()` reports the median and the 68% credible interval of each parameter, under the parameter's own name and unit. The interval is asymmetric in general, which is precisely the information a single symmetric error bar throws away." ] }, { @@ -204,7 +204,7 @@ "metadata": {}, "outputs": [], "source": [ - "analysis.posterior_summary()" + "analysis.bayesian.summary()" ] }, { @@ -224,7 +224,7 @@ "metadata": {}, "outputs": [], "source": [ - "analysis.plot_corner()" + "analysis.bayesian.plot_corner()" ] }, { @@ -244,7 +244,7 @@ "metadata": {}, "outputs": [], "source": [ - "analysis.plot_posterior_predictive(n_draws=100)" + "analysis.bayesian.plot_posterior_predictive(n_draws=100)" ] }, { @@ -264,7 +264,7 @@ "metadata": {}, "outputs": [], "source": [ - "extended = analysis.extend_sampling(additional_samples=1000, thin=2)\n", + "extended = analysis.bayesian.extend(additional_samples=1000, thin=2)\n", "print(f'Chain now holds {extended.draws.shape[0]} draws.')" ] }, @@ -273,7 +273,7 @@ "id": "69793d31", "metadata": {}, "source": [ - "Chains are expensive, so they can be saved and reloaded with `analysis.save_chain(path)` and `analysis.load_chain(path)`. A reloaded chain can be summarized, plotted, or extended further, exactly like a fresh one." + "Chains are expensive, so they can be saved and reloaded with `analysis.bayesian.save(path)` and `analysis.bayesian.load(path)`. A reloaded chain can be summarized, plotted, or extended further, exactly like a fresh one." ] }, { @@ -285,9 +285,9 @@ "\n", "**Data without uncertainties.** If your data carries no variances, the weights fall back to 1, which means the sampler assumes a noise level of 1 in whatever units the intensity happens to be. Least-squares does not care, since that scale cancels out of the best-fit position, but a posterior *does*: its width scales directly with the assumed noise, so the credible intervals will be wrong by whatever factor the true noise differs from 1. Bayesian analysis is not a way to avoid needing uncertainties on your data.\n", "\n", - "**Sampling only some parameters.** `sample_posterior(parameters=[...])` restricts the chain to a subset, which is faster because the number of chains scales with the number of parameters. Be careful with the result: the other parameters are held *fixed*, which is not the same as averaging over them. The intervals you get are conditional on those fixed values, and will be too narrow whenever the parameters are correlated.\n", + "**Sampling only some parameters.** `bayesian.sample(parameters=[...])` restricts the chain to a subset, which is faster because the number of chains scales with the number of parameters. Be careful with the result: the other parameters are held *fixed*, which is not the same as averaging over them. The intervals you get are conditional on those fixed values, and will be too narrow whenever the parameters are correlated.\n", "\n", - "**Degenerate parameters.** As seen above, they are a modelling problem rather than a sampling one. `suggest_bounds()` returning absurd values, or a warning that the posterior has piled up against its bounds, are both signs to go back and look at the model." + "**Degenerate parameters.** As seen above, they are a modelling problem rather than a sampling one. `bayesian.suggest_bounds()` returning absurd values, or a warning that the posterior has piled up against its bounds, are both signs to go back and look at the model." ] } ], diff --git a/src/easydynamics/analysis/__init__.py b/src/easydynamics/analysis/__init__.py index 89a45cea3..89126ecdf 100644 --- a/src/easydynamics/analysis/__init__.py +++ b/src/easydynamics/analysis/__init__.py @@ -2,19 +2,21 @@ # SPDX-License-Identifier: BSD-3-Clause from easydynamics.analysis.analysis import Analysis -from easydynamics.analysis.bayesian_sampling import BayesianSamplingMixin from easydynamics.analysis.parameter_analysis import ParameterAnalysis from easydynamics.analysis.posterior import BoundsSuggestion from easydynamics.analysis.posterior import BoundsSuggestions from easydynamics.analysis.posterior import ParameterPosterior from easydynamics.analysis.posterior import PosteriorSummary +from easydynamics.analysis.posterior_labels import ParameterLabels +from easydynamics.analysis.posterior_sampling import PosteriorSampler __all__ = [ 'Analysis', - 'BayesianSamplingMixin', 'BoundsSuggestion', 'BoundsSuggestions', 'ParameterAnalysis', + 'ParameterLabels', 'ParameterPosterior', + 'PosteriorSampler', 'PosteriorSummary', ] diff --git a/src/easydynamics/analysis/analysis1d.py b/src/easydynamics/analysis/analysis1d.py index eaec732b1..5c8a68be3 100644 --- a/src/easydynamics/analysis/analysis1d.py +++ b/src/easydynamics/analysis/analysis1d.py @@ -12,7 +12,8 @@ from plopp.backends.matplotlib.figure import InteractiveFigure from easydynamics.analysis.analysis_base import AnalysisBase -from easydynamics.analysis.bayesian_sampling import BayesianSamplingMixin +from easydynamics.analysis.posterior_labels import ParameterLabels +from easydynamics.analysis.posterior_sampling import PosteriorSampler from easydynamics.convolution.convolution import Convolution from easydynamics.experiment import Experiment from easydynamics.sample_model import InstrumentModel @@ -26,15 +27,15 @@ from easydynamics.utils.utils import verify_Q_index -class Analysis1d(BayesianSamplingMixin, AnalysisBase): +class Analysis1d(AnalysisBase): """ For analysing one-dimensional data, i.e. intensity as function of energy for a single Q index. Is used primarily in the Analysis class, but can also be used on its own for simpler analyses. - In addition to least-squares fitting with :meth:`fit`, the posterior distribution of the free - parameters can be explored with :meth:`sample_posterior`; see - :class:`~easydynamics.analysis.bayesian_sampling.BayesianSamplingMixin`. + Besides least-squares fitting with :meth:`fit`, the posterior distribution of the free + parameters can be explored through :attr:`bayesian`; see + :class:`~easydynamics.analysis.posterior_sampling.PosteriorSampler`. Examples -------- @@ -121,7 +122,9 @@ def __init__( self._fit_result = None self._convolver = None self._convolver_is_dirty = True - self._init_bayesian_state() + self._fitter = None + self._fitter_is_dirty = True + self._bayesian = None super().__init__( display_name=display_name, @@ -253,32 +256,82 @@ def fit(self) -> FitResults: self._prepare_for_sampling() - x, y, weights = self._get_sampling_data() + x, y, weights = self._sampling_data() fit_result = self.fitter.fit(x=x, y=y, weights=weights) self._fit_result = fit_result return fit_result + @property + def fitter(self) -> EasyScienceFitter: + """ + The EasyScience Fitter used for fitting and sampling, built on first use. + + Exposed so the minimizer, tolerance, and maximum evaluation count can be configured + directly, e.g. ``analysis.fitter.switch_minimizer(AvailableMinimizers.Bumps)``. + + Returns + ------- + EasyScienceFitter + The cached Fitter. + """ + if self._fitter_is_dirty or self._fitter is None: + self._fitter = EasyScienceFitter( + fit_object=self, + fit_function=self.as_fit_function(), + ) + self._fitter_is_dirty = False + return self._fitter + + @property + def bayesian(self) -> PosteriorSampler: + """ + Bayesian posterior sampling for this Analysis, created on first use. + + Returns + ------- + PosteriorSampler + The sampler, which holds any chain that has been run. + """ + if self._bayesian is None: + self._bayesian = PosteriorSampler( + analysis=self, + sampling_data=self._sampling_data, + chain_parameters=self._chain_parameters, + parameter_labels=self._parameter_labels, + prepare=self._prepare_for_sampling, + ) + return self._bayesian + + def _invalidate_fitter(self) -> None: + """Mark the Fitter, and the Sampler built from it, as needing a rebuild.""" + self._fitter_is_dirty = True + self._invalidate_bayesian_sampler() + + def _invalidate_bayesian_sampler(self) -> None: + """Mark the Sampler as needing a rebuild, the data having changed.""" + if self._bayesian is not None: + self._bayesian.invalidate() + ############# - # Hooks for BayesianSamplingMixin + # The contract PosteriorSampler relies on ############# - def _build_bayesian_fitter(self) -> EasyScienceFitter: + def _parameter_labels(self) -> ParameterLabels: """ - Build the EasyScience Fitter for this Analysis. + Get labels for the chain's parameters. + + A single Q index holds one copy of each parameter, so nothing needs qualifying. Returns ------- - EasyScienceFitter - A Fitter bound to this Analysis and its fit function. + ParameterLabels + Labels over the current free parameters. """ - return EasyScienceFitter( - fit_object=self, - fit_function=self.as_fit_function(), - ) + return ParameterLabels(self._chain_parameters()) - def _get_sampling_data(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + def _sampling_data(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """ Get the finite data for the chosen Q index, as used by both fitting and sampling. @@ -292,7 +345,7 @@ def _get_sampling_data(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: ) return x, y, weights - def _get_chain_parameters(self) -> list[Parameter]: + def _chain_parameters(self) -> list[Parameter]: """ Get the free parameters of this Analysis. diff --git a/src/easydynamics/analysis/bayesian_sampling.py b/src/easydynamics/analysis/bayesian_sampling.py deleted file mode 100644 index 0ea38d8ef..000000000 --- a/src/easydynamics/analysis/bayesian_sampling.py +++ /dev/null @@ -1,993 +0,0 @@ -# SPDX-FileCopyrightText: 2026 EasyScience contributors -# SPDX-License-Identifier: BSD-3-Clause - -""" -Shared Bayesian MCMC sampling machinery for the Analysis classes. - -Everything that does not depend on how a particular Analysis is wired up lives here: caching the -Fitter and the Sampler, guarding the parameter bounds, restoring parameter values afterwards, and -turning raw draws into a readable summary. A concrete Analysis supplies the three things that do -differ, via :meth:`BayesianSamplingMixin._build_bayesian_fitter`, -:meth:`BayesianSamplingMixin._get_sampling_data`, and -:meth:`BayesianSamplingMixin._get_chain_parameters`. -""" - -from __future__ import annotations - -import json -import warnings -from pathlib import Path -from typing import TYPE_CHECKING -from typing import Any - -import numpy as np -from easyscience.fitting import AvailableMinimizers -from easyscience.fitting import Sampler - -from easydynamics.analysis.posterior import PosteriorSummary -from easydynamics.analysis.posterior import parameters_at_bounds -from easydynamics.analysis.posterior import suggest_bounds_for_parameters -from easydynamics.analysis.posterior import summarize_draws -from easydynamics.analysis.posterior import unbounded_parameters - -if TYPE_CHECKING: - import os - from collections.abc import Callable - - from easyscience.fitting.fitter import Fitter - from easyscience.fitting.sampler import SamplingResults - from easyscience.variable import Parameter - from matplotlib.figure import Figure - - from easydynamics.analysis.posterior import BoundsSuggestions - -# Suffix of the sidecar mapping chain columns to stable parameter names, written next to the BUMPS -# chain files by save_chain(). -_NAME_MAP_SUFFIX = '.parameter-names.json' - - -class BayesianSamplingMixin: - """ - Bayesian MCMC sampling on top of an Analysis, backed by the BUMPS DREAM sampler. - - Sampling explores the full posterior distribution of the free parameters rather than reporting - a single best-fit point, which is worth doing when parameters are correlated or their - uncertainties are strongly non-Gaussian -- both common in QENS. - - Running :meth:`fit` first is not required, but it helps: DREAM seeds its population in a small - ball around the parameters' current values, so starting from fitted values shortens the burn-in - needed to reach the typical set. - - Notes - ----- - All free parameters must have finite bounds before sampling, because in DREAM the bounds are - the prior. :meth:`suggest_bounds` proposes bounds for any parameter still missing one. - - Examples - -------- - ```python - analysis.fit() - analysis.suggest_bounds().apply() - results = analysis.sample_posterior(samples=10000, burn=2000, thin=10) - analysis.posterior_summary() - ``` - """ - - ############# - # Setup - ############# - - def _init_bayesian_state(self) -> None: - """ - Initialize the cached sampling state. - - Must be called by the concrete Analysis before any observer callback can fire, in the same - way as the other cached objects on the class. - """ - self._fitter = None - self._fitter_is_dirty = True - self._bayesian_sampler = None - self._bayesian_sampler_is_dirty = True - self._posterior_result = None - # Maps a chain column's unique_name to the parameter name it had when saved. Only populated - # by load_chain, because unique_names are per-session and do not survive a round trip. - self._chain_name_map = {} - - def _invalidate_fitter(self) -> None: - """ - Mark the cached Fitter and Sampler as needing a rebuild. - - The Sampler binds its data at construction, so anything that invalidates the Fitter - invalidates the Sampler too. - """ - self._fitter_is_dirty = True - self._bayesian_sampler_is_dirty = True - - def _invalidate_bayesian_sampler(self) -> None: - """ - Mark only the cached Sampler as needing a rebuild. - - Used when the data changed but the model did not. - """ - self._bayesian_sampler_is_dirty = True - - ############# - # Hooks for concrete Analysis classes - ############# - - def _build_bayesian_fitter(self) -> Fitter: - """ - Build the EasyScience Fitter (or MultiFitter) for this Analysis. - - Returns - ------- - Fitter - A configured Fitter or MultiFitter. - - Raises - ------ - NotImplementedError - If the concrete Analysis does not implement it. - """ - raise NotImplementedError('Subclasses must implement _build_bayesian_fitter.') - - def _get_sampling_data(self) -> tuple: - """ - Get the ``(x, y, weights)`` to bind to the Sampler. - - Each element is either an array (single dataset) or a list of arrays (MultiFitter). - - Returns - ------- - tuple - The ``(x, y, weights)`` triple. - - Raises - ------ - NotImplementedError - If the concrete Analysis does not implement it. - """ - raise NotImplementedError('Subclasses must implement _get_sampling_data.') - - def _get_chain_parameters(self) -> list[Parameter]: - """ - Get the free parameters that will appear as columns of the chain. - - Returns - ------- - list[Parameter] - The free parameters of the underlying model(s). - - Raises - ------ - NotImplementedError - If the concrete Analysis does not implement it. - """ - raise NotImplementedError('Subclasses must implement _get_chain_parameters.') - - def _prepare_for_sampling(self) -> None: - """ - Bring any cached computation up to date before a sampling run. - - The default does nothing; Analysis classes that cache a convolver override it. - """ - - ############# - # Properties - ############# - - @property - def fitter(self) -> Fitter: - """ - The EasyScience Fitter used for fitting and sampling, built on first use. - - Exposed so the minimizer, tolerance, and maximum evaluation count can be configured - directly, e.g. ``analysis.fitter.switch_minimizer(AvailableMinimizers.Bumps)``. - - Returns - ------- - Fitter - The cached Fitter or MultiFitter. - """ - if self._fitter_is_dirty or self._fitter is None: - self._fitter = self._build_bayesian_fitter() - self._fitter_is_dirty = False - return self._fitter - - @property - def bayesian_sampler(self) -> Sampler | None: - """ - The EasyScience Sampler holding the MCMC chain, or None before the first run. - - Named to avoid confusion with the SampleModel: this samples the posterior, not the sample. - - Returns - ------- - Sampler | None - The cached Sampler, or None if no chain has been started. - """ - return self._bayesian_sampler - - @property - def posterior_result(self) -> SamplingResults | None: - """ - The results of the most recent sampling run, or None if there has not been one. - - Returns - ------- - SamplingResults | None - The most recent sampling results. - """ - return self._posterior_result - - ############# - # Bounds - ############# - - def suggest_bounds( - self, - n_sigma: float = 10.0, - relative_pad: float = 0.2, - absolute_floor: float | None = None, - ) -> BoundsSuggestions: - """ - Propose finite bounds for free parameters that still have an infinite one. - - Nothing is changed until :meth:`BoundsSuggestions.apply` is called, so the proposal can be - reviewed first. Bounds that are already finite are never widened or narrowed, so physical - limits such as a non-negative area are left alone. - - Because the bounds act as a uniform prior in DREAM, a generous width is the safe choice: - too tight a bound truncates the posterior and understates the uncertainty. - - Parameters - ---------- - n_sigma : float, default=10.0 - How many standard deviations of the fitted uncertainty to allow on each side. - relative_pad : float, default=0.2 - Extra half-width as a fraction of the absolute parameter value. This guards against - minimizers that report a zero or absurdly small uncertainty. - absolute_floor : float | None, default=None - A minimum half-width in the parameter's own units, for when neither the uncertainty nor - the value carries the natural scale. - - Returns - ------- - BoundsSuggestions - The proposed bounds, which must be applied explicitly. - """ - return suggest_bounds_for_parameters( - self._get_chain_parameters(), - n_sigma=n_sigma, - relative_pad=relative_pad, - absolute_floor=absolute_floor, - ) - - def check_bounds_for_sampling(self) -> None: - """ - Verify that every free parameter has finite bounds. - - Raises - ------ - ValueError - If any free parameter has an infinite lower or upper bound. - """ - unbounded = unbounded_parameters(self._get_chain_parameters()) - if not unbounded: - return - names = ', '.join(parameter.name for parameter in unbounded) - raise ValueError( - f'Bayesian sampling requires finite bounds on every free parameter, because the ' - f'bounds act as the prior. These parameters are unbounded: {names}. ' - f'Set their min and max, or call suggest_bounds() to propose values.' - ) - - ############# - # Sampling - ############# - - def sample_posterior( - self, - samples: int = 10000, - burn: int = 2000, - thin: int = 10, - population: int | None = None, - parameters: list[Parameter] | list[str] | None = None, - **sampler_options: dict[str, Any], - ) -> SamplingResults: - """ - Draw samples from the posterior distribution of the free parameters. - - This starts a fresh chain, replacing any existing one; use :meth:`extend_sampling` to - continue a chain instead. Parameter values are restored to what they were beforehand, so - sampling never silently moves the model off its fitted values; use - :meth:`set_parameters_to_posterior_median` to adopt the posterior. - - Parameters - ---------- - samples : int, default=10000 - Number of raw samples to draw across all chains, before thinning. This is a guaranteed - minimum rather than an exact count. - burn : int, default=2000 - Burn-in generations to discard before collecting samples. - thin : int, default=10 - Thinning interval, which reduces autocorrelation between retained draws. - population : int | None, default=None - DREAM population scale factor: BUMPS runs ``ceil(population * n_parameters)`` chains. - parameters : list[Parameter] | list[str] | None, default=None - Restrict the chain to these parameters, given as Parameter objects or names. All other - free parameters are held fixed for the duration of the run. Note that holding a - parameter fixed is not the same as marginalizing over it: the resulting credible - intervals are conditional on the fixed values and will be too narrow if the parameters - are correlated. The default samples every free parameter. - **sampler_options : dict[str, Any] - Forwarded to the EasyScience Sampler, e.g. ``sampler_kwargs``, ``progress_callback``, - or ``abort_test``. - - Returns - ------- - SamplingResults - The sampling results, also stored on :attr:`posterior_result`. - """ - return self._run_sampling( - parameters=parameters, - run=lambda sampler: sampler.sample( - samples=samples, - burn=burn, - thin=thin, - population=population, - **sampler_options, - ), - ) - - def extend_sampling( - self, - additional_samples: int = 5000, - thin: int = 10, - parameters: list[Parameter] | list[str] | None = None, - **sampler_options: dict[str, Any], - ) -> SamplingResults: - """ - Continue the existing chain with additional samples. - - Parameters - ---------- - additional_samples : int, default=5000 - Number of additional samples to draw, in the same units as ``samples``. - thin : int, default=10 - Thinning interval for the retained draws. - parameters : list[Parameter] | list[str] | None, default=None - The same restriction as in :meth:`sample_posterior`. Pass the same value that started - the chain, since the chain's columns cannot change on extension. - **sampler_options : dict[str, Any] - Forwarded to the EasyScience Sampler. - - Returns - ------- - SamplingResults - The sampling results for the full extended chain. - - Raises - ------ - RuntimeError - If there is no chain to extend. - """ - if self._bayesian_sampler is None: - raise RuntimeError( - 'No chain to extend. Call sample_posterior() or load_chain() first.' - ) - return self._run_sampling( - parameters=parameters, - run=lambda sampler: sampler.extend( - additional_samples=additional_samples, - thin=thin, - **sampler_options, - ), - reuse_sampler=True, - ) - - def _run_sampling( - self, - parameters: list[Parameter] | list[str] | None, - run: Callable[[Sampler], SamplingResults], - reuse_sampler: bool = False, - ) -> SamplingResults: - """ - Run a sampling operation with all the surrounding guards in place. - - Checks the bounds, switches the minimizer to BUMPS, optionally holds parameters fixed, - runs, and then restores the parameter values, fixed flags, and minimizer. - - Parameters - ---------- - parameters : list[Parameter] | list[str] | None - Parameters to restrict the chain to, or None for all free parameters. - run : Callable[[Sampler], SamplingResults] - The operation to perform on the prepared Sampler. - reuse_sampler : bool, default=False - Whether to reuse the cached Sampler rather than rebuilding it. Required when extending - a chain, since the chain lives on the Sampler. - - Returns - ------- - SamplingResults - The results of the run. - - Raises - ------ - RuntimeError - If the BUMPS sampler fails while removing outlier chains, which points at degenerate - parameters. - """ - held_fixed = self._resolve_parameters_to_hold_fixed(parameters) - self._warn_about_held_parameters(held_fixed) - - with _FixedParameters(held_fixed): - self.check_bounds_for_sampling() - self._prepare_for_sampling() - - chain_parameters = self._get_chain_parameters() - saved_values = [(p, p.value) for p in chain_parameters] - - fitter = self.fitter - original_minimizer = fitter.minimizer.enum - fitter.switch_minimizer(AvailableMinimizers.Bumps) - try: - sampler = self._get_or_build_sampler(reuse_sampler=reuse_sampler) - results = run(sampler) - except IndexError as error: - # BUMPS' own outlier removal indexes past the end of its buffer when chains - # scatter wildly, which in practice means the model is not identifiable. The bare - # IndexError says nothing useful, so point at the likely cause instead. - raise RuntimeError( - 'The BUMPS sampler failed while removing outlier chains. This usually means ' - 'the chains scattered because two or more free parameters are degenerate, so ' - 'the data cannot determine them separately. Check for degenerate parameters ' - "and fix one of them, or retry with sampler_kwargs={'outliers': 'none'}." - ) from error - finally: - fitter.switch_minimizer(original_minimizer) - for parameter, value in saved_values: - parameter.value = value - - # A fresh chain is labelled with this session's unique names, so any mapping left over - # from a loaded chain no longer applies. - self._chain_name_map = { - parameter.unique_name: parameter.name for parameter in chain_parameters - } - - self._posterior_result = results - self._warn_about_bounds_occupancy(results, self._resolve_chain_parameters(results)) - return results - - def _get_or_build_sampler(self, reuse_sampler: bool) -> Sampler: - """ - Get the cached Sampler, rebuilding it if the data or model changed. - - Parameters - ---------- - reuse_sampler : bool - Whether to reuse the cached Sampler even if it is marked dirty. - - Returns - ------- - Sampler - The Sampler to run. - """ - needs_rebuild = self._bayesian_sampler is None or ( - self._bayesian_sampler_is_dirty and not reuse_sampler - ) - if needs_rebuild: - x, y, weights = self._get_sampling_data() - self._bayesian_sampler = Sampler(self.fitter, x, y, weights=weights) - self._bayesian_sampler_is_dirty = False - return self._bayesian_sampler - - def _resolve_parameters_to_hold_fixed( - self, - parameters: list[Parameter] | list[str] | None, - ) -> list[Parameter]: - """ - Work out which free parameters must be held fixed to honour a subset request. - - Parameters - ---------- - parameters : list[Parameter] | list[str] | None - The requested subset, as Parameter objects or names, or None for all free parameters. - - Returns - ------- - list[Parameter] - The free parameters that are not in the requested subset. - - Raises - ------ - TypeError - If parameters is not a list of Parameters or strings, or None. - ValueError - If a requested name does not match any free parameter, or the subset is empty. - """ - if parameters is None: - return [] - if not isinstance(parameters, (list, tuple)): - raise TypeError('parameters must be a list of Parameters, a list of names, or None.') - - free = self._get_chain_parameters() - by_name = {parameter.name: parameter for parameter in free} - requested = [] - for entry in parameters: - if isinstance(entry, str): - if entry not in by_name: - available = ', '.join(sorted(by_name)) - raise ValueError(f'No free parameter named {entry!r}. Available: {available}.') - requested.append(by_name[entry]) - elif hasattr(entry, 'unique_name'): - requested.append(entry) - else: - raise TypeError( - 'parameters must contain Parameter objects or parameter names (strings).' - ) - - requested_unique_names = {parameter.unique_name for parameter in requested} - if not requested_unique_names: - raise ValueError('parameters must name at least one parameter to sample.') - return [ - parameter for parameter in free if parameter.unique_name not in requested_unique_names - ] - - @staticmethod - def _warn_about_held_parameters(held_fixed: list[Parameter]) -> None: - """ - Warn that holding parameters fixed makes the credible intervals conditional. - - Parameters - ---------- - held_fixed : list[Parameter] - The parameters being held fixed for the run. - """ - if not held_fixed: - return - names = ', '.join(parameter.name for parameter in held_fixed) - warnings.warn( - ( - f'Holding these parameters fixed while sampling: {names}. ' - f'Fixing a parameter is not the same as marginalizing over it, so the resulting ' - f'credible intervals are conditional on these values and will be too narrow if ' - f'the parameters are correlated.' - ), - UserWarning, - stacklevel=4, - ) - - @staticmethod - def _warn_about_bounds_occupancy( - results: SamplingResults, - parameters_by_column: list[Parameter | None], - ) -> None: - """ - Warn when the posterior has piled up against a bound. - - Parameters - ---------- - results : SamplingResults - The sampling results to inspect. - parameters_by_column : list[Parameter | None] - The parameter for each column of the chain, or None where none could be matched. - """ - piled_up = parameters_at_bounds(results.draws, parameters_by_column) - if not piled_up: - return - details = ', '.join( - f'{name} ({fraction:.0%} of draws)' for name, fraction in piled_up.items() - ) - warnings.warn( - ( - f'The posterior is piled up against the bounds for: {details}. ' - f'The bounds, rather than the data, are setting these credible intervals. ' - f'Widen the bounds, or check whether these parameters are degenerate with others.' - ), - UserWarning, - stacklevel=4, - ) - - ############# - # Results - ############# - - def posterior_summary(self) -> PosteriorSummary: - """ - Summarize the marginal posterior of each sampled parameter. - - Reports the median and the 68% credible interval under the parameter's own name and unit, - rather than the opaque unique name the sampler uses internally. Requires a completed - sampling run. - - Returns - ------- - PosteriorSummary - One entry per sampled parameter. - """ - results = self._require_posterior_result() - return summarize_draws( - draws=results.draws, - fallback_names=self._chain_display_names(results), - parameters_by_column=self._resolve_chain_parameters(results), - ) - - def set_parameters_to_posterior_median(self) -> list[Parameter]: - """ - Set every sampled parameter to the median of its marginal posterior. - - Note that the vector of marginal medians is not in general the same as the - highest-posterior point, and for strongly correlated parameters it need not even be a good - fit. Requires a completed sampling run. - - Returns - ------- - list[Parameter] - The parameters that were changed. - """ - results = self._require_posterior_result() - changed = [] - for column, parameter in enumerate(self._resolve_chain_parameters(results)): - if parameter is None: - continue - parameter.value = float(np.median(results.draws[:, column])) - changed.append(parameter) - return changed - - def _require_posterior_result(self) -> SamplingResults: - """ - Get the stored sampling results, raising if there are none. - - Returns - ------- - SamplingResults - The most recent sampling results. - - Raises - ------ - RuntimeError - If no sampling has been run yet. - """ - if self._posterior_result is None: - raise RuntimeError( - 'No posterior samples yet. Call sample_posterior() or load_chain() first.' - ) - return self._posterior_result - - ############# - # Persistence - ############# - - def save_chain(self, path: str | os.PathLike) -> None: - """ - Save the MCMC chain to disk. - - Writes the BUMPS chain files alongside a sidecar recording the parameter names and a - fingerprint of the data that was sampled. - - Parameters - ---------- - path : str | os.PathLike - Path prefix for the chain files. - - Raises - ------ - RuntimeError - If there is no chain to save. - """ - if self._bayesian_sampler is None: - raise RuntimeError('No chain to save. Call sample_posterior() first.') - self._bayesian_sampler.save(path) - # The BUMPS sidecar records unique names, which are handed out per session and so mean - # nothing on reload. Record the parameter names alongside them, which are stable. - Path(f'{path}{_NAME_MAP_SUFFIX}').write_text( - json.dumps(self._chain_name_map, indent=2), - encoding='utf-8', - ) - - def load_chain(self, path: str | os.PathLike, skip: int = 0) -> SamplingResults: - """ - Load a previously saved MCMC chain. - - The loaded chain can be inspected, summarized, or continued with :meth:`extend_sampling`. A - chain saved from different data loads with a warning. - - Parameters - ---------- - path : str | os.PathLike - The path prefix the chain was saved under. - skip : int, default=0 - Number of initial samples to skip when reading the chain. - - Returns - ------- - SamplingResults - The loaded sampling results, also stored on :attr:`posterior_result`. - """ - self._prepare_for_sampling() - name_map_path = Path(f'{path}{_NAME_MAP_SUFFIX}') - if name_map_path.is_file(): - self._chain_name_map = json.loads(name_map_path.read_text(encoding='utf-8')) - else: - self._chain_name_map = {} - warnings.warn( - ( - f'No parameter-name sidecar found at {name_map_path}. The chain will be ' - f'reported under the internal names it was saved with, because those cannot ' - f'be matched to this Analysis.' - ), - UserWarning, - stacklevel=2, - ) - - fitter = self.fitter - original_minimizer = fitter.minimizer.enum - fitter.switch_minimizer(AvailableMinimizers.Bumps) - try: - sampler = self._get_or_build_sampler(reuse_sampler=False) - results = sampler.load_state(path, skip=skip) - finally: - fitter.switch_minimizer(original_minimizer) - self._posterior_result = results - return results - - ############# - # Plotting - ############# - - def plot_trace(self, **kwargs: dict[str, Any]) -> Figure: - """ - Plot the chain trace of each sampled parameter. - - A well-mixed chain looks like a "hairy caterpillar" with no drift; visible trends mean the - chain has not converged and needs a longer burn-in. Requires a completed sampling run. - - Parameters - ---------- - **kwargs : dict[str, Any] - Forwarded to :func:`easydynamics.utils.posterior_plotting.plot_trace`. - - Returns - ------- - Figure - The matplotlib Figure. - """ - from easydynamics.utils.posterior_plotting import plot_trace - - results = self._require_posterior_result() - return plot_trace( - draws=results.draws, - logp=results.logp, - names=self._chain_display_names(results), - title=self.display_name, - **kwargs, - ) - - def plot_corner(self, **kwargs: dict[str, Any]) -> Figure: - """ - Plot the marginal and pairwise posterior distributions. - - Diagonal panels show each parameter's marginal distribution; off-diagonal panels show the - joint distribution of a pair, where a strong diagonal ridge means the two are correlated. - Requires a completed sampling run. - - Parameters - ---------- - **kwargs : dict[str, Any] - Forwarded to :func:`easydynamics.utils.posterior_plotting.plot_corner`. - - Returns - ------- - Figure - The matplotlib Figure. - """ - from easydynamics.utils.posterior_plotting import plot_corner - - results = self._require_posterior_result() - return plot_corner( - draws=results.draws, - names=self._chain_display_names(results), - title=self.display_name, - **kwargs, - ) - - def plot_posterior_predictive( - self, - n_draws: int = 200, - credible_interval: float = 68.0, - **kwargs: dict[str, Any], - ) -> Figure: - """ - Plot the data against the credible band implied by the posterior. - - The model is re-evaluated for a random subset of the posterior draws, and the spread of - those curves becomes the band. Data straying outside the band systematically points at a - model that is missing something, rather than at parameters that need tuning. Requires a - completed sampling run. - - Parameters - ---------- - n_draws : int, default=200 - How many posterior draws to evaluate the model for. Each draw costs one full model - evaluation, so this trades smoothness of the band against time. - credible_interval : float, default=68.0 - Width of the credible band, as a percentage. - **kwargs : dict[str, Any] - Forwarded to :func:`easydynamics.utils.posterior_plotting.plot_posterior_predictive`. - - Returns - ------- - Figure - The matplotlib Figure. - - Raises - ------ - NotImplementedError - If this Analysis binds a list of datasets rather than a single one. - ValueError - If n_draws is not a positive integer. - """ - from easydynamics.utils.posterior_plotting import plot_posterior_predictive - - if not isinstance(n_draws, int) or isinstance(n_draws, bool) or n_draws < 1: - raise ValueError(f'n_draws must be a positive integer. Got {n_draws}.') - - results = self._require_posterior_result() - x, y, weights = self._get_sampling_data() - if isinstance(x, (list, tuple)): - raise NotImplementedError( - 'plot_posterior_predictive supports a single dataset only. Plot each dataset ' - 'from its own Analysis1d instead.' - ) - - predictions = self._evaluate_over_draws(results, x, n_draws) - y_err = None if weights is None else 1.0 / np.asarray(weights) - return plot_posterior_predictive( - x=np.asarray(x), - y=np.asarray(y), - predictions=predictions, - y_err=y_err, - title=self.display_name, - credible_interval=credible_interval, - **kwargs, - ) - - def _evaluate_over_draws( - self, - results: SamplingResults, - x: np.ndarray, - n_draws: int, - ) -> np.ndarray: - """ - Evaluate the model once per posterior draw, restoring the parameters afterwards. - - Parameters - ---------- - results : SamplingResults - The sampling results supplying the draws. - x : np.ndarray - The independent variable to evaluate the model on. - n_draws : int - How many draws to evaluate. Draws are taken evenly across the chain. - - Returns - ------- - np.ndarray - Model evaluations, shape ``(n_selected, len(x))``. - """ - self._prepare_for_sampling() - - columns = [ - (parameter, column) - for column, parameter in enumerate(self._resolve_chain_parameters(results)) - if parameter is not None - ] - saved_values = [(parameter, parameter.value) for parameter, _ in columns] - - total = results.draws.shape[0] - indices = np.unique(np.linspace(0, total - 1, min(n_draws, total)).astype(int)) - - fit_function = self.fitter.fit_function - predictions = [] - try: - for index in indices: - for parameter, column in columns: - parameter.value = float(results.draws[index, column]) - predictions.append(np.asarray(fit_function(x))) - finally: - for parameter, value in saved_values: - parameter.value = value - - return np.vstack(predictions) - - def _resolve_chain_parameters(self, results: SamplingResults) -> list[Parameter | None]: - """ - Match each column of the chain to one of this Analysis's parameters. - - Columns are matched on ``unique_name`` first. That fails for a chain loaded from disk, - because unique names are handed out per session, so a saved chain also records the - parameter names and those are used as a fallback. - - Parameters - ---------- - results : SamplingResults - The sampling results whose columns should be matched. - - Returns - ------- - list[Parameter | None] - The parameter for each column, or None where no match could be made. - """ - parameters = self._get_chain_parameters() - by_unique_name = {p.unique_name: p for p in parameters} - by_name = {p.name: p for p in parameters} - resolved = [] - for unique_name in results.param_names: - parameter = by_unique_name.get(unique_name) - if parameter is None: - saved_name = self._chain_name_map.get(unique_name) - parameter = None if saved_name is None else by_name.get(saved_name) - resolved.append(parameter) - return resolved - - def _chain_display_names(self, results: SamplingResults) -> list[str]: - """ - Translate the chain's column names into parameter names. - - Parameters - ---------- - results : SamplingResults - The sampling results whose columns should be named. - - Returns - ------- - list[str] - One name per column of the chain. - """ - resolved = self._resolve_chain_parameters(results) - return [ - self._chain_name_map.get(unique_name, unique_name) - if parameter is None - else parameter.name - for unique_name, parameter in zip(results.param_names, resolved, strict=True) - ] - - -class _FixedParameters: - """ - Context manager that temporarily fixes parameters and restores their flags on exit. - """ - - def __init__(self, parameters: list[Parameter]) -> None: - """ - Initialize the context manager. - - Parameters - ---------- - parameters : list[Parameter] - The parameters to hold fixed for the duration of the block. - """ - self._parameters = list(parameters) - self._saved: list[tuple[Parameter, bool]] = [] - - def __enter__(self) -> None: - """ - Fix the parameters, remembering their previous state. - """ - self._saved = [(parameter, parameter.fixed) for parameter in self._parameters] - for parameter in self._parameters: - parameter.fixed = True - - def __exit__(self, *_exc_info: object) -> None: - """ - Restore the previous fixed state of every parameter. - - Parameters - ---------- - *_exc_info : object - Exception information, ignored. - """ - for parameter, was_fixed in self._saved: - parameter.fixed = was_fixed diff --git a/src/easydynamics/analysis/posterior.py b/src/easydynamics/analysis/posterior.py index e22ae87d9..4853bb1a5 100644 --- a/src/easydynamics/analysis/posterior.py +++ b/src/easydynamics/analysis/posterior.py @@ -40,6 +40,8 @@ class BoundsSuggestion: ---------- parameter : Parameter The parameter the suggestion applies to. + label : str + The name the parameter is reported under, qualified where several share a name. suggested_min : float The proposed lower bound. Equal to the parameter's current lower bound when that is already finite. @@ -51,6 +53,7 @@ class BoundsSuggestion: """ parameter: Parameter + label: str suggested_min: float suggested_max: float reason: str @@ -178,7 +181,8 @@ def __repr__(self) -> str: if not self._suggestions: return 'BoundsSuggestions(no free parameters)' - header = f'{"parameter":<28s} {"current":>26s} {"suggested":>26s}' + width = max(len('parameter'), *(len(s.label) for s in self._suggestions)) + header = f'{"parameter":<{width}s} {"current":>26s} {"suggested":>26s}' lines = ['BoundsSuggestions', header, '-' * len(header)] for s in self._suggestions: current = f'({s.parameter.min:.4g}, {s.parameter.max:.4g})' @@ -186,7 +190,7 @@ def __repr__(self) -> str: suggested = f'-- {s.reason}' else: suggested = f'({s.suggested_min:.4g}, {s.suggested_max:.4g})' - lines.append(f'{s.parameter.name:<28s} {current:>26s} {suggested:>26s}') + lines.append(f'{s.label:<{width}s} {current:>26s} {suggested:>26s}') attention = self.needing_attention if attention: @@ -199,6 +203,7 @@ def __repr__(self) -> str: def suggest_bounds_for_parameters( parameters: list[Parameter], + labels: list[str] | None = None, n_sigma: float = 10.0, relative_pad: float = 0.2, absolute_floor: float | None = None, @@ -223,6 +228,9 @@ def suggest_bounds_for_parameters( ---------- parameters : list[Parameter] The parameters to propose bounds for. + labels : list[str] | None, default=None + The name to report each parameter under, one per parameter. Defaults to the parameters' own + names. n_sigma : float, default=10.0 How many standard deviations of the parameter's fitted uncertainty to allow on each side. relative_pad : float, default=0.2 @@ -242,20 +250,24 @@ def suggest_bounds_for_parameters( if absolute_floor is not None: _verify_nonneg_number(absolute_floor, 'absolute_floor') + if labels is None: + labels = [parameter.name for parameter in parameters] suggestions = [ _suggest_bounds_for_parameter( parameter=parameter, + label=label, n_sigma=n_sigma, relative_pad=relative_pad, absolute_floor=absolute_floor, ) - for parameter in parameters + for parameter, label in zip(parameters, labels, strict=True) ] return BoundsSuggestions(suggestions) def _suggest_bounds_for_parameter( parameter: Parameter, + label: str, n_sigma: float, relative_pad: float, absolute_floor: float | None, @@ -267,6 +279,8 @@ def _suggest_bounds_for_parameter( ---------- parameter : Parameter The parameter to propose bounds for. + label : str + The name to report the parameter under. n_sigma : float How many standard deviations to allow on each side. relative_pad : float @@ -288,6 +302,7 @@ def _suggest_bounds_for_parameter( if min_is_finite and max_is_finite: return BoundsSuggestion( parameter=parameter, + label=label, suggested_min=current_min, suggested_max=current_max, reason='', @@ -298,6 +313,7 @@ def _suggest_bounds_for_parameter( if not np.isfinite(value): return BoundsSuggestion( parameter=parameter, + label=label, suggested_min=current_min, suggested_max=current_max, reason='value is not finite', @@ -312,6 +328,7 @@ def _suggest_bounds_for_parameter( if not np.isfinite(half_width) or half_width <= 0: return BoundsSuggestion( parameter=parameter, + label=label, suggested_min=current_min, suggested_max=current_max, reason='no scale information (zero value and uncertainty)', @@ -319,6 +336,7 @@ def _suggest_bounds_for_parameter( return BoundsSuggestion( parameter=parameter, + label=label, suggested_min=current_min if min_is_finite else value - half_width, suggested_max=current_max if max_is_finite else value + half_width, reason='', @@ -527,13 +545,14 @@ def __repr__(self) -> str: if not self._entries: return 'PosteriorSummary(no parameters)' + width = max(len('parameter'), *(len(e.name) for e in self._entries)) header = ( - f'{"parameter":<28s} {"unit":>10s} {"median":>14s} ' + f'{"parameter":<{width}s} {"unit":>10s} {"median":>14s} ' f'{"-":>12s} {"+":>12s} {"current":>14s}' ) lines = ['PosteriorSummary', header, '-' * len(header)] lines.extend( - f'{e.name:<28s} {e.unit:>10s} {e.median:>14.5g} ' + f'{e.name:<{width}s} {e.unit:>10s} {e.median:>14.5g} ' f'{e.minus:>12.4g} {e.plus:>12.4g} {e.value:>14.5g}' for e in self._entries ) @@ -542,7 +561,7 @@ def __repr__(self) -> str: def summarize_draws( draws: np.ndarray, - fallback_names: list[str], + labels: list[str], parameters_by_column: list[Parameter | None], ) -> PosteriorSummary: """ @@ -556,8 +575,8 @@ def summarize_draws( ---------- draws : np.ndarray Posterior draws, shape ``(n_draws, n_parameters)``. - fallback_names : list[str] - Label to use for any column with no matching parameter, one per column. + labels : list[str] + The label to report each column under, one per column. parameters_by_column : list[Parameter | None] The parameter for each column of ``draws``, or None where none could be matched. @@ -573,7 +592,7 @@ def summarize_draws( ) entries.append( ParameterPosterior( - name=fallback_names[column] if parameter is None else parameter.name, + name=labels[column], unit='' if parameter is None else str(parameter.unit), median=median, lower=lower, diff --git a/src/easydynamics/analysis/posterior_labels.py b/src/easydynamics/analysis/posterior_labels.py new file mode 100644 index 000000000..84ff9212e --- /dev/null +++ b/src/easydynamics/analysis/posterior_labels.py @@ -0,0 +1,182 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +""" +Naming the columns of an MCMC chain. + +The sampler labels its columns with each parameter's ``unique_name`` -- ``Parameter_4`` and the +like -- which is not what a user recognises, and which is handed out per session so it does not +survive a saved chain either. This turns those columns back into readable labels. +""" + +from __future__ import annotations + +from collections import Counter +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Callable + + from easyscience.variable import Parameter + + +class ParameterLabels: + """ + Readable labels and units for the columns of a chain. + + Built once for a fixed set of parameters, so the name counts and lookups are computed a single + time. Doing this per column instead is quadratic in the parameter count, which is seconds of + work for an analysis with many Q values. + + Parameters + ---------- + parameters : list[Parameter] + The parameters that can appear as columns. + qualify : Callable[[Parameter], str | None] | None, default=None + Returns a qualifier for a parameter whose name is shared with another, for example its Q + index. Only consulted when the bare name really is ambiguous, so an analysis with nothing + to disambiguate keeps its short names. Returning None leaves the name unqualified. + """ + + def __init__( + self, + parameters: list[Parameter], + qualify: Callable[[Parameter], str | None] | None = None, + ) -> None: + self._parameters = list(parameters) + self._qualify = qualify + self._counts = Counter(parameter.name for parameter in self._parameters) + self._by_unique_name = {p.unique_name: p for p in self._parameters} + self._by_label = {self.label(p): p for p in self._parameters} + + @property + def parameters(self) -> list[Parameter]: + """ + The parameters these labels describe. + + Returns + ------- + list[Parameter] + The parameters given at construction. + """ + return list(self._parameters) + + def label(self, parameter: Parameter) -> str: + """ + Get the label a parameter is reported under. + + Parameters + ---------- + parameter : Parameter + The parameter to label. + + Returns + ------- + str + The parameter's name, qualified only where that name is shared with another parameter. + """ + if self._counts[parameter.name] <= 1 or self._qualify is None: + return parameter.name + qualifier = self._qualify(parameter) + return parameter.name if qualifier is None else f'{parameter.name} ({qualifier})' + + def name_map(self) -> dict[str, str]: + """ + Map each parameter's ``unique_name`` to its label. + + Saved alongside a chain, because unique names are per-session: without this a reloaded + chain cannot be matched back to any parameter. + + Returns + ------- + dict[str, str] + Mapping of unique name to label. + """ + return {p.unique_name: self.label(p) for p in self._parameters} + + def resolve( + self, + column_names: list[str], + saved_labels: dict[str, str] | None = None, + ) -> list[Parameter | None]: + """ + Match each column of a chain to a parameter. + + Columns are matched on ``unique_name`` first. That fails for a chain loaded from disk, + where the saved labels are used instead. + + Parameters + ---------- + column_names : list[str] + The sampler's name for each column. + saved_labels : dict[str, str] | None, default=None + Mapping of unique name to label, as recorded when a chain was saved. + + Returns + ------- + list[Parameter | None] + The parameter for each column, or None where no match could be made. + """ + saved_labels = saved_labels or {} + resolved = [] + for unique_name in column_names: + parameter = self._by_unique_name.get(unique_name) + if parameter is None: + parameter = self._by_label.get(saved_labels.get(unique_name, '')) + resolved.append(parameter) + return resolved + + def display_names( + self, + column_names: list[str], + saved_labels: dict[str, str] | None = None, + ) -> list[str]: + """ + Get a readable label for each column of a chain. + + Parameters + ---------- + column_names : list[str] + The sampler's name for each column. + saved_labels : dict[str, str] | None, default=None + Mapping of unique name to label, as recorded when a chain was saved. + + Returns + ------- + list[str] + One label per column, falling back to the saved label and then to the raw column name. + """ + saved_labels = saved_labels or {} + return [ + saved_labels.get(unique_name, unique_name) + if parameter is None + else self.label(parameter) + for unique_name, parameter in zip( + column_names, self.resolve(column_names, saved_labels), strict=True + ) + ] + + def units( + self, + column_names: list[str], + saved_labels: dict[str, str] | None = None, + ) -> list[str]: + """ + Get the unit of each column of a chain. + + Parameters + ---------- + column_names : list[str] + The sampler's name for each column. + saved_labels : dict[str, str] | None, default=None + Mapping of unique name to label, as recorded when a chain was saved. + + Returns + ------- + list[str] + One unit per column, empty where no parameter could be matched. + """ + return [ + '' if parameter is None else str(parameter.unit) + for parameter in self.resolve(column_names, saved_labels) + ] diff --git a/src/easydynamics/analysis/posterior_sampling.py b/src/easydynamics/analysis/posterior_sampling.py new file mode 100644 index 000000000..f782b2d00 --- /dev/null +++ b/src/easydynamics/analysis/posterior_sampling.py @@ -0,0 +1,933 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +""" +Bayesian MCMC sampling for the Analysis classes, backed by the BUMPS DREAM sampler. + +The sampler is composed into an Analysis rather than inherited by it: an Analysis exposes one +``bayesian`` property, and everything to do with sampling lives here instead of being mixed into +three classes. Labelling lives in :mod:`easydynamics.analysis.posterior_labels` and the figures in +:mod:`easydynamics.utils.posterior_plotting`; this module only runs chains. +""" + +from __future__ import annotations + +import json +import warnings +from pathlib import Path +from typing import TYPE_CHECKING +from typing import Any + +import numpy as np +from easyscience.fitting import AvailableMinimizers +from easyscience.fitting import Sampler + +from easydynamics.analysis.posterior import parameters_at_bounds +from easydynamics.analysis.posterior import suggest_bounds_for_parameters +from easydynamics.analysis.posterior import summarize_draws +from easydynamics.analysis.posterior import unbounded_parameters + +if TYPE_CHECKING: + import os + from collections.abc import Callable + + from easyscience.fitting.sampler import SamplingResults + from easyscience.variable import Parameter + from matplotlib.figure import Figure + + from easydynamics.analysis.posterior import BoundsSuggestions + from easydynamics.analysis.posterior import PosteriorSummary + from easydynamics.analysis.posterior_labels import ParameterLabels + +# Suffix of the sidecar mapping chain columns to stable labels, written next to the BUMPS chain +# files by save(). +_LABEL_MAP_SUFFIX = '.parameter-names.json' + + +class PosteriorSampler: + """ + Draws samples from the posterior distribution of an Analysis' free parameters. + + Reached as ``analysis.bayesian``. Sampling explores the whole posterior rather than reporting a + single best-fit point, which is worth doing when parameters are correlated or their + uncertainties are strongly non-Gaussian, both common in QENS. + + Running a fit first is not required, but it helps: DREAM seeds its population in a small ball + around the parameters' current values, so starting from fitted values shortens the burn-in. + + The Analysis passes in everything that differs between the Analysis classes, so this class + needs no knowledge of how any of them is built. + + Parameters + ---------- + analysis : object + The Analysis being sampled, used for its ``display_name`` and its ``fitter``. + sampling_data : Callable[[], tuple] + Returns the ``(x, y, weights)`` to bind to the sampler. Each is an array, or a list of + arrays for a multi-dataset fit. + chain_parameters : Callable[[], list[Parameter]] + Returns the free parameters that will form the chain's columns. + parameter_labels : Callable[[], ParameterLabels] + Returns labels for those parameters. + prepare : Callable[[], None] | None, default=None + Brings any cached computation on the Analysis up to date before a run. + + Notes + ----- + Every free parameter must have finite bounds before sampling, because in DREAM the bounds are + the prior. :meth:`suggest_bounds` proposes bounds for any parameter still missing one. + + Examples + -------- + ```python + analysis.fit() + analysis.bayesian.suggest_bounds().apply() + analysis.bayesian.sample(samples=10000, burn=2000, thin=10) + analysis.bayesian.summary() + ``` + """ + + def __init__( + self, + analysis: object, + sampling_data: Callable[[], tuple], + chain_parameters: Callable[[], list[Parameter]], + parameter_labels: Callable[[], ParameterLabels], + prepare: Callable[[], None] | None = None, + ) -> None: + self._analysis = analysis + self._sampling_data = sampling_data + self._chain_parameters = chain_parameters + self._parameter_labels = parameter_labels + self._prepare_hook = prepare + self._sampler: Sampler | None = None + self._sampler_is_dirty = True + self._results: SamplingResults | None = None + # Maps a chain column's unique_name to the label it had when saved. Only populated by + # load(), because unique names are per-session and do not survive a round trip. + self._saved_labels: dict[str, str] = {} + + ############# + # State + ############# + + def invalidate(self) -> None: + """ + Mark the underlying Sampler as needing a rebuild. + + Called by the Analysis when its data changes, since the Sampler binds its data at + construction. + """ + self._sampler_is_dirty = True + + @property + def sampler(self) -> Sampler | None: + """ + The EasyScience Sampler holding the chain, or None before the first run. + + Returns + ------- + Sampler | None + The cached Sampler. + """ + return self._sampler + + @property + def results(self) -> SamplingResults | None: + """ + The results of the most recent run, or None if there has not been one. + + Returns + ------- + SamplingResults | None + The most recent sampling results. + """ + return self._results + + ############# + # Bounds + ############# + + def suggest_bounds( + self, + n_sigma: float = 10.0, + relative_pad: float = 0.2, + absolute_floor: float | None = None, + ) -> BoundsSuggestions: + """ + Propose finite bounds for free parameters that still have an infinite one. + + Nothing changes until :meth:`BoundsSuggestions.apply` is called, so the proposal can be + reviewed first. Bounds that are already finite are never widened or narrowed, so physical + limits such as a non-negative area are left alone. + + Because the bounds act as a uniform prior in DREAM, a generous width is the safe choice: + too tight a bound truncates the posterior and understates the uncertainty. + + Parameters + ---------- + n_sigma : float, default=10.0 + How many standard deviations of the fitted uncertainty to allow on each side. + relative_pad : float, default=0.2 + Extra half-width as a fraction of the absolute parameter value, guarding against + minimizers that report a zero or absurdly small uncertainty. + absolute_floor : float | None, default=None + A minimum half-width in the parameter's own units, for when neither the uncertainty nor + the value carries the natural scale. + + Returns + ------- + BoundsSuggestions + The proposed bounds, which must be applied explicitly. + """ + labels = self._labels() + return suggest_bounds_for_parameters( + labels.parameters, + labels=[labels.label(parameter) for parameter in labels.parameters], + n_sigma=n_sigma, + relative_pad=relative_pad, + absolute_floor=absolute_floor, + ) + + def check_bounds(self) -> None: + """ + Verify that every free parameter has finite bounds. + + Raises + ------ + ValueError + If any free parameter has an infinite lower or upper bound. + """ + labels = self._labels() + unbounded = unbounded_parameters(labels.parameters) + if not unbounded: + return + names = ', '.join(labels.label(parameter) for parameter in unbounded) + raise ValueError( + f'Bayesian sampling requires finite bounds on every free parameter, because the ' + f'bounds act as the prior. These parameters are unbounded: {names}. ' + f'Set their min and max, or call suggest_bounds() to propose values.' + ) + + ############# + # Sampling + ############# + + def sample( + self, + samples: int = 10000, + burn: int = 2000, + thin: int = 10, + population: int | None = None, + parameters: list[Parameter] | list[str] | None = None, + **sampler_options: dict[str, Any], + ) -> SamplingResults: + """ + Draw samples from the posterior distribution of the free parameters. + + Starts a fresh chain, replacing any existing one; use :meth:`extend` to continue one. + Parameter values are restored afterwards, so sampling never silently moves the model off + its fitted values; use :meth:`set_parameters_to_median` to adopt the posterior. + + Parameters + ---------- + samples : int, default=10000 + Number of raw samples to draw across all chains, before thinning. A guaranteed minimum + rather than an exact count. + burn : int, default=2000 + Burn-in generations to discard before collecting samples. + thin : int, default=10 + Thinning interval, which reduces autocorrelation between retained draws. + population : int | None, default=None + DREAM population scale factor: BUMPS runs ``ceil(population * n_parameters)`` chains. + parameters : list[Parameter] | list[str] | None, default=None + Restrict the chain to these parameters, given as Parameter objects or labels. All other + free parameters are held fixed for the run. Holding a parameter fixed is not the same + as marginalizing over it: the resulting intervals are conditional on those values and + will be too narrow if the parameters are correlated. The default samples everything. + **sampler_options : dict[str, Any] + Forwarded to the EasyScience Sampler, e.g. ``sampler_kwargs`` or ``progress_callback``. + + Returns + ------- + SamplingResults + The sampling results, also stored on :attr:`results`. + """ + return self._run( + parameters=parameters, + run=lambda sampler: sampler.sample( + samples=samples, burn=burn, thin=thin, population=population, **sampler_options + ), + ) + + def extend( + self, + additional_samples: int = 5000, + thin: int = 10, + parameters: list[Parameter] | list[str] | None = None, + **sampler_options: dict[str, Any], + ) -> SamplingResults: + """ + Continue the existing chain with additional samples. + + Parameters + ---------- + additional_samples : int, default=5000 + Number of additional samples to draw, in the same units as ``samples``. + thin : int, default=10 + Thinning interval for the retained draws. + parameters : list[Parameter] | list[str] | None, default=None + The same restriction as in :meth:`sample`. It must leave the chain the same width, + since BUMPS resumes from a stored chain whose columns are fixed. + **sampler_options : dict[str, Any] + Forwarded to the EasyScience Sampler. + + Returns + ------- + SamplingResults + The sampling results for the full extended chain. + + Raises + ------ + RuntimeError + If there is no chain to extend. + """ + if self._sampler is None: + raise RuntimeError('No chain to extend. Call sample() or load() first.') + return self._run( + parameters=parameters, + run=lambda sampler: sampler.extend( + additional_samples=additional_samples, thin=thin, **sampler_options + ), + reuse_sampler=True, + ) + + def _run( + self, + parameters: list[Parameter] | list[str] | None, + run: Callable[[Sampler], SamplingResults], + reuse_sampler: bool = False, + ) -> SamplingResults: + """ + Run a sampling operation with the surrounding guards in place. + + Checks the bounds, switches the minimizer to BUMPS, optionally holds parameters fixed, and + restores the parameter values, fixed flags and minimizer afterwards. + + Parameters + ---------- + parameters : list[Parameter] | list[str] | None + Parameters to restrict the chain to, or None for all free parameters. + run : Callable[[Sampler], SamplingResults] + The operation to perform on the prepared Sampler. + reuse_sampler : bool, default=False + Whether to reuse the cached Sampler, as an extension must. + + Returns + ------- + SamplingResults + The results of the run. + + Raises + ------ + IndexError + Re-raised untouched when it did not come from BUMPS, since that is a bug here rather + than a modelling problem. + RuntimeError + If the BUMPS sampler fails while removing outlier chains. + """ + held_fixed = self._resolve_parameters_to_hold_fixed(parameters) + _warn_about_held_parameters(self._labels(), held_fixed) + + with _FixedParameters(held_fixed): + self.check_bounds() + self._prepare() + + chain_parameters = self._chain_parameters() + saved_values = [(p, p.value) for p in chain_parameters] + + if reuse_sampler: + self._verify_chain_shape_unchanged(chain_parameters) + + fitter = self._analysis.fitter + original_minimizer = fitter.minimizer.enum + fitter.switch_minimizer(AvailableMinimizers.Bumps) + try: + results = run(self._get_or_build_sampler(reuse_sampler=reuse_sampler)) + except IndexError as error: + if not _raised_inside_bumps(error): + raise + raise RuntimeError( + 'The BUMPS sampler failed while removing outlier chains. This happens when ' + 'the chains scatter because two or more free parameters are degenerate, and ' + 'also on short chains, where BUMPS has too few generations to work with. ' + 'Check for degenerate parameters, raise samples, or switch the outlier ' + "removal off with sampler_kwargs={'outliers': 'none'}." + ) from error + finally: + fitter.switch_minimizer(original_minimizer) + for parameter, value in saved_values: + parameter.value = value + + # Labelled outside the block above, so a subset run records the labels a full run would. + # Inside it the other parameters are fixed, nothing looks ambiguous, and the sidecar would + # be written with unqualified names that no longer match on reload. + self._saved_labels = self._labels().name_map() + self._results = results + self._warn_about_bounds_occupancy(results) + return results + + def _get_or_build_sampler(self, reuse_sampler: bool) -> Sampler: + """ + Get the cached Sampler, rebuilding it if the data changed. + + Parameters + ---------- + reuse_sampler : bool + Whether to reuse the cached Sampler even if it is marked dirty. + + Returns + ------- + Sampler + The Sampler to run. + """ + if self._sampler is None or (self._sampler_is_dirty and not reuse_sampler): + x, y, weights = self._sampling_data() + self._sampler = Sampler(self._analysis.fitter, x, y, weights=weights) + self._sampler_is_dirty = False + return self._sampler + + def _verify_chain_shape_unchanged(self, chain_parameters: list[Parameter]) -> None: + """ + Check that an extension keeps the chain's columns. + + Parameters + ---------- + chain_parameters : list[Parameter] + The parameters that would form the chain for this run. + + Raises + ------ + ValueError + If the number of parameters differs from the existing chain's. + """ + if self._results is None: + return + existing = self._results.draws.shape[1] + if len(chain_parameters) != existing: + raise ValueError( + f'Cannot extend a chain of {existing} parameters with a run of ' + f'{len(chain_parameters)}. An extension continues the stored chain, whose columns ' + f'are fixed, so it needs the same parameters the chain was started with. Start a ' + f'fresh chain with sample() instead.' + ) + + def _resolve_parameters_to_hold_fixed( + self, + parameters: list[Parameter] | list[str] | None, + ) -> list[Parameter]: + """ + Work out which free parameters must be held fixed to honour a subset request. + + Parameters + ---------- + parameters : list[Parameter] | list[str] | None + The requested subset, as Parameter objects or labels, or None for everything. + + Returns + ------- + list[Parameter] + The free parameters that are not in the requested subset. + + Raises + ------ + TypeError + If parameters is not a list of Parameters or strings, or None. + ValueError + If a requested label matches no free parameter, or the subset is empty. + """ + if parameters is None: + return [] + if not isinstance(parameters, (list, tuple)): + raise TypeError('parameters must be a list of Parameters, a list of labels, or None.') + + labels = self._labels() + by_label = {labels.label(parameter): parameter for parameter in labels.parameters} + requested = [] + for entry in parameters: + if isinstance(entry, str): + if entry not in by_label: + raise ValueError( + f'No free parameter named {entry!r}. ' + f'Available: {", ".join(sorted(by_label))}.' + ) + requested.append(by_label[entry]) + elif hasattr(entry, 'unique_name'): + requested.append(entry) + else: + raise TypeError('parameters must contain Parameter objects or labels (strings).') + + wanted = {parameter.unique_name for parameter in requested} + if not wanted: + raise ValueError('parameters must name at least one parameter to sample.') + return [p for p in labels.parameters if p.unique_name not in wanted] + + def _warn_about_bounds_occupancy(self, results: SamplingResults) -> None: + """ + Warn when the posterior has piled up against a bound. + + Parameters + ---------- + results : SamplingResults + The sampling results to inspect. + """ + piled_up = parameters_at_bounds(results.draws, self._resolve(results)) + if not piled_up: + return + details = ', '.join( + f'{name} ({fraction:.0%} of draws)' for name, fraction in piled_up.items() + ) + warnings.warn( + ( + f'The posterior is piled up against the bounds for: {details}. ' + f'The bounds, rather than the data, are setting these credible intervals. ' + f'Widen the bounds, or check whether these parameters are degenerate with others.' + ), + UserWarning, + stacklevel=4, + ) + + ############# + # Results + ############# + + def summary(self) -> PosteriorSummary: + """ + Summarize the marginal posterior of each sampled parameter. + + Reports the median and the 68% credible interval under the parameter's own label and unit. + + Returns + ------- + PosteriorSummary + One entry per sampled parameter. + """ + results = self._require_results() + labels = self._labels() + return summarize_draws( + draws=results.draws, + labels=labels.display_names(results.param_names, self._saved_labels), + parameters_by_column=self._resolve(results), + ) + + def set_parameters_to_median(self) -> list[Parameter]: + """ + Set every sampled parameter to the median of its marginal posterior. + + The vector of marginal medians is not in general the highest-posterior point, and for + strongly correlated parameters need not even be a good fit. + + Returns + ------- + list[Parameter] + The parameters that were changed. + """ + results = self._require_results() + changed = [] + for column, parameter in enumerate(self._resolve(results)): + if parameter is None: + continue + parameter.value = float(np.median(results.draws[:, column])) + changed.append(parameter) + return changed + + ############# + # Persistence + ############# + + def save(self, path: str | os.PathLike) -> None: + """ + Save the MCMC chain to disk. + + Writes the BUMPS chain files plus a sidecar recording the column labels, because the unique + names BUMPS stores are per-session and cannot be matched up again on their own. + + Parameters + ---------- + path : str | os.PathLike + Path prefix for the chain files. + + Raises + ------ + RuntimeError + If there is no chain to save. + """ + if self._sampler is None: + raise RuntimeError('No chain to save. Call sample() first.') + self._sampler.save(path) + Path(f'{path}{_LABEL_MAP_SUFFIX}').write_text( + json.dumps(self._saved_labels, indent=2), encoding='utf-8' + ) + + def load(self, path: str | os.PathLike, skip: int = 0) -> SamplingResults: + """ + Load a previously saved MCMC chain. + + The loaded chain can be summarized, plotted, or continued with :meth:`extend`. + + Parameters + ---------- + path : str | os.PathLike + The path prefix the chain was saved under. + skip : int, default=0 + Number of initial samples to skip when reading the chain. + + Returns + ------- + SamplingResults + The loaded results, also stored on :attr:`results`. + """ + self._prepare() + sidecar = Path(f'{path}{_LABEL_MAP_SUFFIX}') + if sidecar.is_file(): + self._saved_labels = json.loads(sidecar.read_text(encoding='utf-8')) + else: + self._saved_labels = {} + warnings.warn( + ( + f'No parameter-name sidecar found at {sidecar}. The chain will be reported ' + f'under the internal names it was saved with, because those cannot be matched ' + f'to this Analysis.' + ), + UserWarning, + stacklevel=2, + ) + + fitter = self._analysis.fitter + original_minimizer = fitter.minimizer.enum + fitter.switch_minimizer(AvailableMinimizers.Bumps) + try: + self._results = self._get_or_build_sampler(reuse_sampler=False).load_state( + path, skip=skip + ) + finally: + fitter.switch_minimizer(original_minimizer) + return self._results + + ############# + # Figures, each one a call into posterior_plotting + ############# + + def plot_trace(self, **kwargs: dict[str, Any]) -> Figure: + """ + Plot the chain trace of each sampled parameter. + + Parameters + ---------- + **kwargs : dict[str, Any] + Forwarded to :func:`easydynamics.utils.posterior_plotting.plot_trace`. + + Returns + ------- + Figure + The matplotlib Figure. + """ + from easydynamics.utils.posterior_plotting import plot_trace + + results = self._require_results() + return plot_trace( + draws=results.draws, + names=self._display_names(results), + logp=results.logp, + units=self._units(results), + title=self._analysis.display_name, + **kwargs, + ) + + def plot_corner(self, **kwargs: dict[str, Any]) -> Figure: + """ + Plot the marginal and pairwise posterior distributions. + + Parameters + ---------- + **kwargs : dict[str, Any] + Forwarded to :func:`easydynamics.utils.posterior_plotting.plot_corner`. + + Returns + ------- + Figure + The matplotlib Figure. + """ + from easydynamics.utils.posterior_plotting import plot_corner + + results = self._require_results() + return plot_corner( + draws=results.draws, + names=self._display_names(results), + units=self._units(results), + title=self._analysis.display_name, + **kwargs, + ) + + def plot_posterior_predictive( + self, + n_draws: int = 200, + credible_interval: float = 68.0, + **kwargs: dict[str, Any], + ) -> Figure: + """ + Plot the data against the credible band implied by the posterior. + + Parameters + ---------- + n_draws : int, default=200 + How many posterior draws to evaluate the model for. Each costs a full model evaluation. + credible_interval : float, default=68.0 + Width of the credible band, as a percentage. + **kwargs : dict[str, Any] + Forwarded to :func:`easydynamics.utils.posterior_plotting.plot_posterior_predictive`. + + Returns + ------- + Figure + The matplotlib Figure. + + Raises + ------ + NotImplementedError + If this Analysis binds a list of datasets rather than a single one. + ValueError + If n_draws is not a positive integer. + """ + from easydynamics.utils.posterior_plotting import plot_posterior_predictive + + if not isinstance(n_draws, int) or isinstance(n_draws, bool) or n_draws < 1: + raise ValueError(f'n_draws must be a positive integer. Got {n_draws}.') + + self._require_results() + x, y, weights = self._sampling_data() + if isinstance(x, (list, tuple)): + raise NotImplementedError( + 'plot_posterior_predictive supports a single dataset only. Plot each dataset ' + 'from its own Analysis1d instead.' + ) + + energy = getattr(self._analysis, 'energy', None) + sample_model = getattr(self._analysis, 'sample_model', None) + y_unit = None if sample_model is None else getattr(sample_model, 'y_unit', None) + kwargs.setdefault('xlabel', None if energy is None else f'Energy ({energy.unit})') + kwargs.setdefault('ylabel', 'Intensity' if y_unit is None else f'Intensity ({y_unit})') + + return plot_posterior_predictive( + x=np.asarray(x), + y=np.asarray(y), + predictions=self.predictions(n_draws), + y_err=None if weights is None else 1.0 / np.asarray(weights), + title=self._analysis.display_name, + credible_interval=credible_interval, + **kwargs, + ) + + def predictions(self, n_draws: int = 200) -> np.ndarray: + """ + Evaluate the model once per posterior draw, restoring the parameters afterwards. + + Parameters + ---------- + n_draws : int, default=200 + How many draws to evaluate, taken evenly across the chain. + + Returns + ------- + np.ndarray + Model evaluations, shape ``(n_selected, len(x))``. + """ + results = self._require_results() + self._prepare() + + x, _, _ = self._sampling_data() + columns = [ + (parameter, column) + for column, parameter in enumerate(self._resolve(results)) + if parameter is not None + ] + saved_values = [(parameter, parameter.value) for parameter, _ in columns] + + total = results.draws.shape[0] + indices = np.unique(np.linspace(0, total - 1, min(n_draws, total)).astype(int)) + + fit_function = self._analysis.fitter.fit_function + predictions = [] + try: + for index in indices: + for parameter, column in columns: + parameter.value = float(results.draws[index, column]) + predictions.append(np.asarray(fit_function(x))) + finally: + for parameter, value in saved_values: + parameter.value = value + return np.vstack(predictions) + + ############# + # Talking to the Analysis + ############# + + def _labels(self) -> ParameterLabels: + """ + Get the label helper for the current free parameters. + + Returns + ------- + ParameterLabels + Built fresh, because which parameters are free can change between calls. + """ + return self._parameter_labels() + + def _prepare(self) -> None: + """Bring any cached computation on the Analysis up to date before a run.""" + if self._prepare_hook is not None: + self._prepare_hook() + + def _resolve(self, results: SamplingResults) -> list[Parameter | None]: + """ + Match each column of a chain to a parameter. + + Parameters + ---------- + results : SamplingResults + The results whose columns should be matched. + + Returns + ------- + list[Parameter | None] + The parameter for each column, or None where none could be matched. + """ + return self._labels().resolve(results.param_names, self._saved_labels) + + def _display_names(self, results: SamplingResults) -> list[str]: + """ + Get a readable label for each column of a chain. + + Parameters + ---------- + results : SamplingResults + The results whose columns should be named. + + Returns + ------- + list[str] + One label per column. + """ + return self._labels().display_names(results.param_names, self._saved_labels) + + def _units(self, results: SamplingResults) -> list[str]: + """ + Get the unit of each column of a chain. + + Parameters + ---------- + results : SamplingResults + The results whose columns should be described. + + Returns + ------- + list[str] + One unit per column. + """ + return self._labels().units(results.param_names, self._saved_labels) + + def _require_results(self) -> SamplingResults: + """ + Get the stored results, raising if there are none. + + Returns + ------- + SamplingResults + The most recent sampling results. + + Raises + ------ + RuntimeError + If no sampling has been run yet. + """ + if self._results is None: + raise RuntimeError('No posterior samples yet. Call sample() or load() first.') + return self._results + + +def _warn_about_held_parameters(labels: object, held_fixed: list[Parameter]) -> None: + """ + Warn that holding parameters fixed makes the credible intervals conditional. + + Parameters + ---------- + labels : object + The ParameterLabels used to name them. + held_fixed : list[Parameter] + The parameters being held fixed for the run. + """ + if not held_fixed: + return + names = ', '.join(labels.label(parameter) for parameter in held_fixed) + warnings.warn( + ( + f'Holding these parameters fixed while sampling: {names}. ' + f'Fixing a parameter is not the same as marginalizing over it, so the resulting ' + f'credible intervals are conditional on these values and will be too narrow if the ' + f'parameters are correlated.' + ), + UserWarning, + stacklevel=4, + ) + + +def _raised_inside_bumps(error: BaseException) -> bool: + """ + Check whether an exception came from inside BUMPS. + + Used so only BUMPS' own failures are relabelled, and a bug in this package is not reported as a + modelling problem. + + Parameters + ---------- + error : BaseException + The exception to inspect. + + Returns + ------- + bool + True when any frame of the traceback lies in the bumps package. + """ + traceback = error.__traceback__ + while traceback is not None: + module = traceback.tb_frame.f_globals.get('__name__', '') + if module == 'bumps' or module.startswith('bumps.'): + return True + traceback = traceback.tb_next + return False + + +class _FixedParameters: + """Context manager that temporarily fixes parameters and restores their flags on exit.""" + + def __init__(self, parameters: list[Parameter]) -> None: + self._parameters = list(parameters) + self._saved: list[tuple[Parameter, bool]] = [] + + def __enter__(self) -> None: + """Fix the parameters, remembering their previous state.""" + self._saved = [(parameter, parameter.fixed) for parameter in self._parameters] + for parameter in self._parameters: + parameter.fixed = True + + def __exit__(self, *_exc_info: object) -> None: + """ + Restore the previous fixed state of every parameter. + + Parameters + ---------- + *_exc_info : object + Exception information, ignored. + """ + for parameter, was_fixed in self._saved: + parameter.fixed = was_fixed diff --git a/src/easydynamics/utils/posterior_plotting.py b/src/easydynamics/utils/posterior_plotting.py index bfb1a1320..8e75cf425 100644 --- a/src/easydynamics/utils/posterior_plotting.py +++ b/src/easydynamics/utils/posterior_plotting.py @@ -14,6 +14,7 @@ import matplotlib.pyplot as plt import numpy as np +from matplotlib.ticker import MaxNLocator if TYPE_CHECKING: from matplotlib.figure import Figure @@ -23,6 +24,7 @@ def plot_trace( draws: np.ndarray, names: list[str], logp: np.ndarray | None = None, + units: list[str] | None = None, title: str | None = None, figsize: tuple[float, float] | None = None, ) -> Figure: @@ -44,6 +46,9 @@ def plot_trace( One label per column of ``draws``. logp : np.ndarray | None, default=None Log-posterior values, plotted in an extra panel when given. + units : list[str] | None, default=None + Unit of each column, appended to its label. Entries that are empty or dimensionless are + skipped, since a bare "dimensionless" only adds clutter. title : str | None, default=None Figure title. figsize : tuple[float, float] | None, default=None @@ -66,7 +71,7 @@ def plot_trace( for axis, column, name in zip(axes, range(draws.shape[1]), names, strict=False): axis.plot(draws[:, column], lw=0.5) - axis.set_ylabel(name, fontsize=8) + axis.set_ylabel(_with_unit(name, units, column), fontsize=8) axis.set_xlim(0, len(draws) - 1) if logp is not None: @@ -83,6 +88,7 @@ def plot_trace( def plot_corner( draws: np.ndarray, names: list[str], + units: list[str] | None = None, title: str | None = None, bins: int = 40, figsize: tuple[float, float] | None = None, @@ -103,6 +109,9 @@ def plot_corner( Posterior draws, shape ``(n_draws, n_parameters)``. names : list[str] One label per column of ``draws``. + units : list[str] | None, default=None + Unit of each column, appended to its label. Entries that are empty or dimensionless are + skipped, since a bare "dimensionless" only adds clutter. title : str | None, default=None Figure title. bins : int, default=40 @@ -143,7 +152,25 @@ def plot_corner( axis.set_ylabel(names[row], fontsize=8) else: axis.set_yticklabels([]) + if row == 0 and col == 0: + # The top-left panel is a histogram, so its vertical axis counts draws rather than + # carrying a parameter. Say so, instead of leaving it blank as if by omission. + axis.set_ylabel('counts', fontsize=8) axis.tick_params(labelsize=7) + axis.xaxis.set_major_locator(MaxNLocator(nbins=4)) + if row != col: + axis.yaxis.set_major_locator(MaxNLocator(nbins=4)) + + # Matplotlib parks the shared exponent ("1e-8") at the end of the axis, where it lands on top + # of the axis label. Fold it into the label instead. + fig.canvas.draw() + for row in range(n): + for col in range(row + 1): + axis = axes[row, col] + if row == n - 1: + _absorb_offset(axis.xaxis, axis.set_xlabel, names[col], units, col) + if col == 0 and row != 0: + _absorb_offset(axis.yaxis, axis.set_ylabel, names[row], units, row) if title is not None: fig.suptitle(title) @@ -158,6 +185,8 @@ def plot_posterior_predictive( y_err: np.ndarray | None = None, title: str | None = None, credible_interval: float = 68.0, + xlabel: str | None = None, + ylabel: str | None = None, figsize: tuple[float, float] = (8.0, 5.0), ) -> Figure: """ @@ -181,6 +210,10 @@ def plot_posterior_predictive( Figure title. credible_interval : float, default=68.0 Width of the credible band, as a percentage. + xlabel : str | None, default=None + Label for the independent axis. + ylabel : str | None, default=None + Label for the dependent axis. figsize : tuple[float, float], default=(8.0, 5.0) Figure size in inches. @@ -224,6 +257,10 @@ def plot_posterior_predictive( label=f'{credible_interval:.0f}% credible band', ) axis.plot(x, median, '-', color='C3', label='Posterior median') + if xlabel is not None: + axis.set_xlabel(xlabel) + if ylabel is not None: + axis.set_ylabel(ylabel) axis.legend() if title is not None: axis.set_title(title) @@ -231,6 +268,85 @@ def plot_posterior_predictive( return fig +def _unit_for(units: list[str] | None, column: int) -> str: + """ + Get the unit to show for a column, if it is worth showing. + + Parameters + ---------- + units : list[str] | None + The units, one per column, or None. + column : int + The column to look up. + + Returns + ------- + str + The unit, or an empty string when there is none worth printing. + """ + if units is None or column >= len(units): + return '' + unit = (units[column] or '').strip() + return '' if unit.lower() in ('', 'dimensionless', 'none') else unit + + +def _with_unit(name: str, units: list[str] | None, column: int) -> str: + """ + Append a column's unit to its label. + + Parameters + ---------- + name : str + The label to extend. + units : list[str] | None + The units, one per column, or None. + column : int + The column the label belongs to. + + Returns + ------- + str + The label, with the unit in parentheses when there is one. + """ + unit = _unit_for(units, column) + return f'{name} ({unit})' if unit else name + + +def _absorb_offset( + axis_object: object, + set_label: object, + name: str, + units: list[str] | None = None, + column: int = 0, +) -> None: + """ + Move an axis' shared exponent into its label, so the two stop overlapping. + + The exponent and the unit share one set of parentheses, since two adjacent parentheticals read + badly: ``D (1e-8 m^2/s)`` rather than ``D (1e-8) (m^2/s)``. + + Parameters + ---------- + axis_object : object + The matplotlib ``XAxis`` or ``YAxis`` carrying the offset text. + set_label : object + The corresponding ``set_xlabel`` or ``set_ylabel`` callable. + name : str + The label the axis should carry, before the exponent and unit are appended. + units : list[str] | None, default=None + The units, one per column, or None. + column : int, default=0 + The column the axis belongs to. + """ + offset_text = axis_object.get_offset_text() + offset = offset_text.get_text() + unit = _unit_for(units, column) + suffix = ' '.join(part for part in (offset, unit) if part) + set_label(f'{name} ({suffix})' if suffix else name, fontsize=8) + if offset: + offset_text.set_visible(False) + + def _verify_draws(draws: np.ndarray, names: list[str]) -> None: """ Verify that a draws array is two-dimensional and matches its labels. diff --git a/tests/integration/fitting/test_bayesian_sampling.py b/tests/integration/fitting/test_bayesian_sampling.py index 704a37ea2..b6adb3bd2 100644 --- a/tests/integration/fitting/test_bayesian_sampling.py +++ b/tests/integration/fitting/test_bayesian_sampling.py @@ -76,17 +76,17 @@ def build_analysis(): def sampled_analysis(): analysis = build_analysis() analysis.fit() - analysis.suggest_bounds().apply() + analysis.bayesian.suggest_bounds().apply() with warnings.catch_warnings(): warnings.simplefilter('ignore') - analysis.sample_posterior(**SAMPLE_KWARGS) + analysis.bayesian.sample(**SAMPLE_KWARGS) return analysis class TestRealChain: def test_chain_has_one_column_per_free_parameter(self, sampled_analysis): # EXPECT - results = sampled_analysis.posterior_result + results = sampled_analysis.bayesian.results assert results.draws.shape[1] == len(sampled_analysis.get_free_parameters()) assert results.draws.shape[0] > 0 @@ -96,7 +96,7 @@ def test_chain_has_one_column_per_free_parameter(self, sampled_analysis): ) def test_posterior_recovers_the_true_parameters(self, sampled_analysis, name, truth): # WHEN - entry = sampled_analysis.posterior_summary()[name] + entry = sampled_analysis.bayesian.summary()[name] # EXPECT the truth sits within a few posterior standard deviations of the median. A 68% # interval is deliberately not used: it excludes the truth about a third of the time for @@ -106,7 +106,7 @@ def test_posterior_recovers_the_true_parameters(self, sampled_analysis, name, tr def test_summary_is_reported_under_parameter_names_and_units(self, sampled_analysis): # WHEN - summary = sampled_analysis.posterior_summary() + summary = sampled_analysis.bayesian.summary() # EXPECT assert {entry.name for entry in summary} == { @@ -118,12 +118,12 @@ def test_sampling_leaves_the_fitted_values_untouched(self): # WHEN analysis = build_analysis() analysis.fit() - analysis.suggest_bounds().apply() + analysis.bayesian.suggest_bounds().apply() before = [float(p.value) for p in analysis.get_free_parameters()] with warnings.catch_warnings(): warnings.simplefilter('ignore') - analysis.sample_posterior(**SAMPLE_KWARGS) + analysis.bayesian.sample(**SAMPLE_KWARGS) # EXPECT after = [float(p.value) for p in analysis.get_free_parameters()] @@ -131,11 +131,11 @@ def test_sampling_leaves_the_fitted_values_untouched(self): def test_extend_grows_the_chain(self, sampled_analysis): # WHEN - before = int(sampled_analysis.posterior_result.state.Ngen) + before = int(sampled_analysis.bayesian.results.state.Ngen) with warnings.catch_warnings(): warnings.simplefilter('ignore') - extended = sampled_analysis.extend_sampling( + extended = sampled_analysis.bayesian.extend( additional_samples=500, thin=2, sampler_kwargs={'trim': False} ) @@ -145,17 +145,17 @@ def test_extend_grows_the_chain(self, sampled_analysis): def test_save_and_load_round_trip_keeps_parameter_identity(self, sampled_analysis, tmp_path): # WHEN prefix = str(tmp_path / 'chain') - sampled_analysis.save_chain(prefix) + sampled_analysis.bayesian.save(prefix) fresh = build_analysis() fresh.fit() - fresh.suggest_bounds().apply() + fresh.bayesian.suggest_bounds().apply() with warnings.catch_warnings(): warnings.simplefilter('ignore') - fresh.load_chain(prefix) + fresh.bayesian.load(prefix) # EXPECT the reloaded chain is reported under real names, not internal unique names - summary = fresh.posterior_summary() + summary = fresh.bayesian.summary() assert {entry.name for entry in summary} == {p.name for p in fresh.get_free_parameters()} assert all(np.isfinite(entry.value) for entry in summary) @@ -163,24 +163,24 @@ def test_subset_sampling_produces_a_single_column(self): # WHEN analysis = build_analysis() analysis.fit() - analysis.suggest_bounds().apply() + analysis.bayesian.suggest_bounds().apply() with warnings.catch_warnings(): warnings.simplefilter('ignore') - results = analysis.sample_posterior(parameters=['Gaussian width'], **SAMPLE_KWARGS) + results = analysis.bayesian.sample(parameters=['Gaussian width'], **SAMPLE_KWARGS) # EXPECT assert results.draws.shape[1] == 1 - assert analysis.posterior_summary().entries[0].name == 'Gaussian width' + assert analysis.bayesian.summary().entries[0].name == 'Gaussian width' def test_plots_render(self, sampled_analysis): # WHEN import matplotlib.pyplot as plt n_parameters = len(sampled_analysis.get_free_parameters()) - trace = sampled_analysis.plot_trace() - corner = sampled_analysis.plot_corner() - predictive = sampled_analysis.plot_posterior_predictive(n_draws=20) + trace = sampled_analysis.bayesian.plot_trace() + corner = sampled_analysis.bayesian.plot_corner() + predictive = sampled_analysis.bayesian.plot_posterior_predictive(n_draws=20) # EXPECT assert len(trace.axes) == n_parameters + 1 @@ -192,13 +192,13 @@ def test_posterior_median_is_close_to_the_least_squares_fit(self): # WHEN analysis = build_analysis() analysis.fit() - analysis.suggest_bounds().apply() + analysis.bayesian.suggest_bounds().apply() fitted = {p.name: float(p.value) for p in analysis.get_free_parameters()} with warnings.catch_warnings(): warnings.simplefilter('ignore') - analysis.sample_posterior(**SAMPLE_KWARGS) - summary = analysis.posterior_summary() + analysis.bayesian.sample(**SAMPLE_KWARGS) + summary = analysis.bayesian.summary() # EXPECT the two agree within the posterior's own uncertainty, since with flat priors the # maximum-likelihood point sits inside the bulk of the posterior diff --git a/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py b/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py index b29a2b54a..345bc2724 100644 --- a/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py +++ b/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py @@ -18,7 +18,7 @@ from easydynamics.sample_model import SampleModel from easydynamics.sample_model.components.gaussian import Gaussian -SAMPLER_PATH = 'easydynamics.analysis.bayesian_sampling.Sampler' +SAMPLER_PATH = 'easydynamics.analysis.posterior_sampling.Sampler' def make_analysis(): @@ -113,23 +113,23 @@ class TestBoundsPreflight: def test_sampling_refuses_unbounded_parameters(self, analysis): # EXPECT with pytest.raises(ValueError, match='finite bounds'): - analysis.sample_posterior(samples=10) + analysis.bayesian.sample(samples=10) def test_error_names_the_offending_parameters(self, analysis): # EXPECT with pytest.raises(ValueError, match='Gaussian area'): - analysis.check_bounds_for_sampling() + analysis.bayesian.check_bounds() def test_bounded_parameters_pass(self, analysis): # WHEN bound_all(analysis) # EXPECT: does not raise - analysis.check_bounds_for_sampling() + analysis.bayesian.check_bounds() def test_suggest_bounds_covers_the_free_parameters(self, analysis): # WHEN - suggestions = analysis.suggest_bounds() + suggestions = analysis.bayesian.suggest_bounds() # EXPECT assert len(suggestions) == len(analysis.get_free_parameters()) @@ -150,7 +150,7 @@ def mutate_then_return(**_kwargs): return fake_results(analysis) sampler_class.return_value.sample.side_effect = mutate_then_return - analysis.sample_posterior(samples=10, burn=1, thin=1) + analysis.bayesian.sample(samples=10, burn=1, thin=1) # EXPECT after = [(p.unique_name, p.value) for p in analysis.get_free_parameters()] @@ -167,7 +167,7 @@ def test_switches_to_bumps_for_the_run(self, analysis): seen.append(analysis.fitter.minimizer.enum), fake_results(analysis), )[1] - analysis.sample_posterior(samples=10) + analysis.bayesian.sample(samples=10) # EXPECT assert seen == [AvailableMinimizers.Bumps] @@ -179,7 +179,7 @@ def test_restores_the_minimizer_even_when_sampling_raises(self, analysis): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.side_effect = RuntimeError('boom') with pytest.raises(RuntimeError, match='boom'): - analysis.sample_posterior(samples=10) + analysis.bayesian.sample(samples=10) # EXPECT assert analysis.fitter.minimizer.enum == AvailableMinimizers.LMFit_leastsq @@ -190,7 +190,7 @@ def test_forwards_sampling_arguments(self, analysis): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) - analysis.sample_posterior(samples=123, burn=7, thin=3, population=5) + analysis.bayesian.sample(samples=123, burn=7, thin=3, population=5) # EXPECT kwargs = sampler_class.return_value.sample.call_args.kwargs @@ -206,11 +206,11 @@ def test_stores_the_result(self, analysis): with patch(SAMPLER_PATH) as sampler_class: expected = fake_results(analysis) sampler_class.return_value.sample.return_value = expected - returned = analysis.sample_posterior(samples=10) + returned = analysis.bayesian.sample(samples=10) # EXPECT assert returned is expected - assert analysis.posterior_result is expected + assert analysis.bayesian.results is expected def test_warns_when_the_posterior_piles_up_against_a_bound(self, analysis): # WHEN a parameter's draws span its whole allowed range @@ -224,7 +224,7 @@ def test_warns_when_the_posterior_piles_up_against_a_bound(self, analysis): # EXPECT with pytest.warns(UserWarning, match='piled up'): - analysis.sample_posterior(samples=10) + analysis.bayesian.sample(samples=10) def test_does_not_warn_when_the_posterior_is_well_inside(self, analysis): # WHEN @@ -235,7 +235,7 @@ def test_does_not_warn_when_the_posterior_is_well_inside(self, analysis): # EXPECT with warnings_as_errors(): - analysis.sample_posterior(samples=10) + analysis.bayesian.sample(samples=10) class TestParameterSubset: @@ -253,7 +253,7 @@ def record(**_kwargs): sampler_class.return_value.sample.side_effect = record with pytest.warns(UserWarning, match='Holding these parameters fixed'): - analysis.sample_posterior(samples=10, parameters=[target.name]) + analysis.bayesian.sample(samples=10, parameters=[target.name]) # EXPECT assert seen['free'] == [target.unique_name] @@ -267,7 +267,7 @@ def test_restores_the_fixed_flags_afterwards(self, analysis): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) with pytest.warns(UserWarning): - analysis.sample_posterior(samples=10, parameters=[target]) + analysis.bayesian.sample(samples=10, parameters=[target]) # EXPECT assert [(p.unique_name, p.fixed) for p in analysis.get_all_parameters()] == before @@ -278,17 +278,17 @@ def test_unknown_parameter_name_raises(self, analysis): # EXPECT with pytest.raises(ValueError, match='No free parameter named'): - analysis.sample_posterior(samples=10, parameters=['not a parameter']) + analysis.bayesian.sample(samples=10, parameters=['not a parameter']) def test_non_list_parameters_raises(self, analysis): # EXPECT with pytest.raises(TypeError, match='must be a list'): - analysis.sample_posterior(samples=10, parameters='Gaussian area') + analysis.bayesian.sample(samples=10, parameters='Gaussian area') def test_empty_parameter_list_raises(self, analysis): # EXPECT with pytest.raises(ValueError, match='at least one parameter'): - analysis.sample_posterior(samples=10, parameters=[]) + analysis.bayesian.sample(samples=10, parameters=[]) class TestSamplerCaching: @@ -298,8 +298,8 @@ def test_sampler_is_reused_between_runs(self, analysis): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) - analysis.sample_posterior(samples=10) - analysis.sample_posterior(samples=10) + analysis.bayesian.sample(samples=10) + analysis.bayesian.sample(samples=10) # EXPECT the data is bound once, not per run assert sampler_class.call_count == 1 @@ -310,9 +310,9 @@ def test_changing_the_q_index_rebuilds_the_sampler(self, analysis): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) - analysis.sample_posterior(samples=10) + analysis.bayesian.sample(samples=10) analysis.Q_index = 0 - analysis.sample_posterior(samples=10) + analysis.bayesian.sample(samples=10) # EXPECT the Sampler binds its data at construction, so it must be rebuilt assert sampler_class.call_count == 2 @@ -323,10 +323,10 @@ def test_binds_the_same_data_the_fit_uses(self, analysis): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) - analysis.sample_posterior(samples=10) + analysis.bayesian.sample(samples=10) # EXPECT - expected_x, expected_y, expected_w = analysis._get_sampling_data() + expected_x, expected_y, expected_w = analysis._sampling_data() args, kwargs = sampler_class.call_args assert np.array_equal(args[1], expected_x) assert np.array_equal(args[2], expected_y) @@ -337,7 +337,7 @@ class TestExtendAndPersistence: def test_extend_without_a_chain_raises(self, analysis): # EXPECT with pytest.raises(RuntimeError, match='No chain to extend'): - analysis.extend_sampling() + analysis.bayesian.extend() def test_extend_delegates_to_the_sampler(self, analysis): # WHEN @@ -346,8 +346,8 @@ def test_extend_delegates_to_the_sampler(self, analysis): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) sampler_class.return_value.extend.side_effect = lambda **_k: fake_results(analysis) - analysis.sample_posterior(samples=10) - analysis.extend_sampling(additional_samples=42, thin=2) + analysis.bayesian.sample(samples=10) + analysis.bayesian.extend(additional_samples=42, thin=2) # EXPECT kwargs = sampler_class.return_value.extend.call_args.kwargs @@ -357,7 +357,7 @@ def test_extend_delegates_to_the_sampler(self, analysis): def test_save_without_a_chain_raises(self, analysis): # EXPECT with pytest.raises(RuntimeError, match='No chain to save'): - analysis.save_chain('somewhere') + analysis.bayesian.save('somewhere') def test_save_writes_the_parameter_name_sidecar(self, analysis, tmp_path): # WHEN @@ -366,8 +366,8 @@ def test_save_writes_the_parameter_name_sidecar(self, analysis, tmp_path): bound_all(analysis) with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) - analysis.sample_posterior(samples=10) - analysis.save_chain(str(tmp_path / 'chain')) + analysis.bayesian.sample(samples=10) + analysis.bayesian.save(str(tmp_path / 'chain')) # EXPECT the unique names are recorded against the stable parameter names sidecar = tmp_path / 'chain.parameter-names.json' @@ -384,14 +384,14 @@ def test_load_without_a_sidecar_warns(self, analysis, tmp_path): # EXPECT with pytest.warns(UserWarning, match='No parameter-name sidecar'): - analysis.load_chain(str(tmp_path / 'missing')) + analysis.bayesian.load(str(tmp_path / 'missing')) class TestResults: def test_summary_without_sampling_raises(self, analysis): # EXPECT with pytest.raises(RuntimeError, match='No posterior samples yet'): - analysis.posterior_summary() + analysis.bayesian.summary() def test_summary_uses_parameter_names_and_units(self, analysis): # WHEN @@ -399,10 +399,10 @@ def test_summary_uses_parameter_names_and_units(self, analysis): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) - analysis.sample_posterior(samples=10) + analysis.bayesian.sample(samples=10) # EXPECT - summary = analysis.posterior_summary() + summary = analysis.bayesian.summary() names = {entry.name for entry in summary} assert names == {p.name for p in analysis.get_free_parameters()} assert all(entry.unit == 'meV' for entry in summary) @@ -416,9 +416,9 @@ def test_set_parameters_to_posterior_median(self, analysis): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.return_value = fake_results(analysis, values=draws) expected = [float(p.value) + 2.0 for p in parameters] - analysis.sample_posterior(samples=10) + analysis.bayesian.sample(samples=10) - changed = analysis.set_parameters_to_posterior_median() + changed = analysis.bayesian.set_parameters_to_median() # EXPECT assert len(changed) == len(parameters) @@ -427,7 +427,7 @@ def test_set_parameters_to_posterior_median(self, analysis): def test_median_without_sampling_raises(self, analysis): # EXPECT with pytest.raises(RuntimeError, match='No posterior samples yet'): - analysis.set_parameters_to_posterior_median() + analysis.bayesian.set_parameters_to_median() class TestPlots: @@ -437,11 +437,11 @@ def test_predictive_rejects_a_bad_draw_count(self, analysis): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) - analysis.sample_posterior(samples=10) + analysis.bayesian.sample(samples=10) # EXPECT with pytest.raises(ValueError, match='positive integer'): - analysis.plot_posterior_predictive(n_draws=0) + analysis.bayesian.plot_posterior_predictive(n_draws=0) def test_predictive_restores_parameter_values(self, analysis): # WHEN @@ -451,10 +451,10 @@ def test_predictive_restores_parameter_values(self, analysis): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.return_value = fake_results(analysis, values=draws) - analysis.sample_posterior(samples=10) + analysis.bayesian.sample(samples=10) before = [float(p.value) for p in parameters] - analysis.plot_posterior_predictive(n_draws=5) + analysis.bayesian.plot_posterior_predictive(n_draws=5) # EXPECT assert [float(p.value) for p in parameters] == pytest.approx(before) @@ -462,9 +462,9 @@ def test_predictive_restores_parameter_values(self, analysis): def test_plots_without_sampling_raise(self, analysis): # EXPECT with pytest.raises(RuntimeError): - analysis.plot_trace() + analysis.bayesian.plot_trace() with pytest.raises(RuntimeError): - analysis.plot_corner() + analysis.bayesian.plot_corner() class warnings_as_errors: diff --git a/tests/unit/easydynamics/analysis/test_posterior.py b/tests/unit/easydynamics/analysis/test_posterior.py index 88dc29bf6..659412774 100644 --- a/tests/unit/easydynamics/analysis/test_posterior.py +++ b/tests/unit/easydynamics/analysis/test_posterior.py @@ -251,7 +251,7 @@ def test_reports_parameter_names_units_and_percentiles(self): draws = np.linspace(0.0, 100.0, 101).reshape(-1, 1) # THEN - summary = summarize_draws(draws, ['Parameter_0'], [parameter]) + summary = summarize_draws(draws, ['Gaussian width'], [parameter]) # EXPECT entry = summary['Gaussian width'] @@ -263,6 +263,17 @@ def test_reports_parameter_names_units_and_percentiles(self): assert entry.plus == pytest.approx(34.0) assert entry.value == pytest.approx(1.5) + def test_labels_are_reported_verbatim(self): + # WHEN a caller supplies a qualified label, as a multi-Q analysis does + parameter = make_parameter(name='Gaussian width') + + # THEN + summary = summarize_draws(np.zeros((5, 1)), ['Gaussian width (Q_index=2)'], [parameter]) + + # EXPECT + assert summary.entries[0].name == 'Gaussian width (Q_index=2)' + assert summary.entries[0].unit == 'meV' + def test_unmatched_column_falls_back_to_the_supplied_name(self): # WHEN draws = np.zeros((10, 1)) @@ -289,7 +300,7 @@ def test_repr_contains_the_parameter_name(self): parameter = make_parameter(name='Gaussian area') # THEN - text = repr(summarize_draws(np.zeros((5, 1)), ['x'], [parameter])) + text = repr(summarize_draws(np.zeros((5, 1)), ['Gaussian area'], [parameter])) # EXPECT assert 'Gaussian area' in text diff --git a/tests/unit/easydynamics/analysis/test_posterior_labels.py b/tests/unit/easydynamics/analysis/test_posterior_labels.py new file mode 100644 index 000000000..b4b08f86c --- /dev/null +++ b/tests/unit/easydynamics/analysis/test_posterior_labels.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +import numpy as np +from easyscience.variable import Parameter + +from easydynamics.analysis.posterior_labels import ParameterLabels + + +def make_parameter(name, unit='meV'): + return Parameter(name=name, value=1.0, unit=unit) + + +class TestLabelling: + def test_unique_names_are_left_alone(self): + # WHEN nothing is ambiguous, a qualifier would only cost width + parameters = [make_parameter('area'), make_parameter('width')] + labels = ParameterLabels(parameters, qualify=lambda _p: 'Q_index=0') + + # EXPECT + assert [labels.label(p) for p in parameters] == ['area', 'width'] + + def test_shared_names_are_qualified(self): + # WHEN two parameters share a name + first, second = make_parameter('width'), make_parameter('width') + owners = {first.unique_name: 'Q_index=0', second.unique_name: 'Q_index=1'} + labels = ParameterLabels([first, second], qualify=lambda p: owners[p.unique_name]) + + # EXPECT + assert labels.label(first) == 'width (Q_index=0)' + assert labels.label(second) == 'width (Q_index=1)' + + def test_a_qualifier_that_declines_leaves_the_name_alone(self): + # WHEN the qualifier cannot identify an owner, as for a parameter shared across Q + first, second = make_parameter('width'), make_parameter('width') + labels = ParameterLabels([first, second], qualify=lambda _p: None) + + # EXPECT the plain name rather than an invented qualifier + assert labels.label(first) == 'width' + + def test_without_a_qualifier_names_stay_bare(self): + # WHEN + first, second = make_parameter('width'), make_parameter('width') + labels = ParameterLabels([first, second]) + + # EXPECT + assert labels.label(first) == 'width' + + +class TestChainColumns: + def test_columns_resolve_by_unique_name(self): + # WHEN + parameters = [make_parameter('area'), make_parameter('width')] + labels = ParameterLabels(parameters) + columns = [p.unique_name for p in reversed(parameters)] + + # EXPECT resolution follows the chain's order, not the parameter list's + assert labels.resolve(columns) == list(reversed(parameters)) + assert labels.display_names(columns) == ['width', 'area'] + assert labels.units(columns) == ['meV', 'meV'] + + def test_a_saved_chain_resolves_through_its_labels(self): + # WHEN a chain was saved in another session, so its unique names mean nothing here + original = make_parameter('width') + saved = {original.unique_name: 'width'} + current = make_parameter('width') + labels = ParameterLabels([current]) + + # EXPECT the saved label finds the parameter this session has + assert labels.resolve([original.unique_name], saved) == [current] + assert labels.display_names([original.unique_name], saved) == ['width'] + + def test_an_unknown_column_is_reported_not_guessed(self): + # WHEN + labels = ParameterLabels([make_parameter('area')]) + + # EXPECT None rather than a wrong parameter, and the raw name to show something + assert labels.resolve(['Parameter_999']) == [None] + assert labels.display_names(['Parameter_999']) == ['Parameter_999'] + assert labels.units(['Parameter_999']) == [''] + + def test_name_map_records_labels_against_unique_names(self): + # WHEN + first, second = make_parameter('width'), make_parameter('width') + owners = {first.unique_name: 'Q_index=0', second.unique_name: 'Q_index=1'} + labels = ParameterLabels([first, second], qualify=lambda p: owners[p.unique_name]) + + # EXPECT what save() writes alongside a chain + assert labels.name_map() == { + first.unique_name: 'width (Q_index=0)', + second.unique_name: 'width (Q_index=1)', + } + + +class TestCost: + def test_labelling_does_not_rescan_per_parameter(self): + # WHEN there are many parameters. Computing the name counts per parameter is quadratic, + # which was seconds of work for an analysis with many Q values. + parameters = [make_parameter(f'p{i // 2}') for i in range(400)] + labels = ParameterLabels(parameters, qualify=lambda _p: 'q') + + # EXPECT labelling all of them stays cheap + import time + + start = time.perf_counter() + names = [labels.label(p) for p in parameters] + assert time.perf_counter() - start < 0.5 + assert len(names) == len(parameters) + assert np.all([n.endswith('(q)') for n in names]) From 85f8a714df32ce1eafc6f2e2653388df030eb04e Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Thu, 13 Aug 2026 14:26:55 +0200 Subject: [PATCH 3/3] Warm the tutorial data cache before running notebooks in parallel The notebook tests run with '-n auto', and five of the notebooks fetch vanadium_data_example.h5 through pooch. On a cold cache the workers race: one is still writing the file into the cache while another opens it, which fails on Windows with "PermissionError: Permission denied". This failed twice in a row on windows-latest, always on that file, always with the other sixteen notebooks passing. The race is pre-existing, but adding a fifth notebook that wants the same file, and lengthening tutorial 1, made it reliable rather than rare. Fetching every tutorial data file once, before the parallel run starts, leaves the workers with nothing to do but read, which is safe. The prefetch reads the URLs and hashes out of the notebooks themselves, so it cannot drift from what they actually download, and it never fails the run: a file it cannot fetch is left to the notebook that needs it, which reports the problem with far more context. Co-Authored-By: Claude Opus 5 (1M context) (cherry picked from commit 46d745a4a73e8025c37e02e490fab67cfdb5ff22) --- pixi.toml | 8 ++- tools/prefetch_tutorial_data.py | 92 +++++++++++++++++++++++++++++++++ 2 files changed, 99 insertions(+), 1 deletion(-) create mode 100644 tools/prefetch_tutorial_data.py diff --git a/pixi.toml b/pixi.toml index f26b4fb7e..db46523e1 100644 --- a/pixi.toml +++ b/pixi.toml @@ -100,7 +100,13 @@ user = { features = ['py-max', 'user'] } unit-tests = 'python -m pytest tests/unit/ --color=yes -v' functional-tests = 'python -m pytest tests/functional/ --color=yes -v' integration-tests = 'python -m pytest tests/integration/ --color=yes -n auto -v' -notebook-tests = 'python -m pytest --nbmake docs/docs/tutorials/ --nbmake-timeout=1200 --color=yes -n auto -v' +# Warm the pooch cache first. Several notebooks fetch the same file, and running them with +# '-n auto' has the workers race: one writes the file while another opens it, which fails on +# Windows. Fetching up front leaves the parallel run with nothing to do but read. +prefetch-tutorial-data = 'python tools/prefetch_tutorial_data.py' +notebook-tests = { cmd = 'python -m pytest --nbmake docs/docs/tutorials/ --nbmake-timeout=1200 --color=yes -n auto -v', depends-on = [ + 'prefetch-tutorial-data', +] } test = { depends-on = ['unit-tests'] } diff --git a/tools/prefetch_tutorial_data.py b/tools/prefetch_tutorial_data.py new file mode 100644 index 000000000..839b1897a --- /dev/null +++ b/tools/prefetch_tutorial_data.py @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +""" +Download every data file the tutorial notebooks fetch, once, before they are run. + +The notebooks are executed in parallel with ``pytest -n auto``, and several of them fetch the same +file through ``pooch``. On a cold cache the workers race: one is still writing the file into the +cache while another tries to open it, which fails on Windows with a permission error. Fetching +everything up front leaves the parallel run with nothing to do but read. + +Run as ``python tools/prefetch_tutorial_data.py``; it is wired into the ``notebook-tests`` task. +""" + +from __future__ import annotations + +import json +import re +import sys +from pathlib import Path + +import pooch + +TUTORIALS = Path(__file__).resolve().parent.parent / 'docs' / 'docs' / 'tutorials' + +# Matches the pooch.retrieve(url=..., known_hash=...) calls the notebooks use, in either order. +URL_PATTERN = re.compile(r"url\s*=\s*f?['\"]([^'\"]+)['\"]") +HASH_PATTERN = re.compile(r"known_hash\s*=\s*['\"]([^'\"]+)['\"]") + + +def find_downloads() -> dict[str, str]: + """ + Collect the ``(url, known_hash)`` pairs the notebooks fetch. + + Returns + ------- + dict[str, str] + Mapping of URL to expected hash, deduplicated across notebooks. + """ + downloads: dict[str, str] = {} + for notebook in sorted(TUTORIALS.glob('*.ipynb')): + cells = json.loads(notebook.read_text(encoding='utf-8'))['cells'] + for cell in cells: + if cell['cell_type'] != 'code': + continue + source = ''.join(cell['source']) + if 'pooch.retrieve' not in source: + continue + urls = URL_PATTERN.findall(source) + hashes = HASH_PATTERN.findall(source) + # Only pairs are usable; a templated URL without a literal hash is skipped rather than + # guessed at, and the notebook will simply fetch it itself. + for url, known_hash in zip(urls, hashes, strict=False): + downloads[url] = known_hash + return downloads + + +def main() -> int: + """ + Fetch every tutorial data file into the pooch cache. + + Deliberately never fails: this only warms a cache. A file that cannot be fetched here is left + to the notebook that needs it, which reports the problem with far more context than this script + could, and which is where the failure belongs. + + Returns + ------- + int + Always zero. + """ + downloads = find_downloads() + if not downloads: + sys.stdout.write('No tutorial downloads found.\n') + return 0 + + failures = 0 + for url, known_hash in downloads.items(): + name = url.rsplit('/', 1)[-1] + try: + pooch.retrieve(url=url, known_hash=known_hash) + except Exception as error: # noqa: BLE001 - report and continue, the notebook will retry + failures += 1 + sys.stdout.write(f'could not prefetch {name}, leaving it to the notebook: {error}\n') + else: + sys.stdout.write(f'cached {name}\n') + + sys.stdout.write(f'{len(downloads) - failures}/{len(downloads)} tutorial data files ready.\n') + return 0 + + +if __name__ == '__main__': + sys.exit(main())