diff --git a/docs/docs/tutorials/bayesian.ipynb b/docs/docs/tutorials/bayesian.ipynb new file mode 100644 index 000000000..51a68acd2 --- /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 `bayesian.sample()`." + ] + }, + { + "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 `bayesian.sample()` would refuse to run.\n", + "\n", + "`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." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a004cd83", + "metadata": {}, + "outputs": [], + "source": [ + "suggestions = analysis.bayesian.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", + "`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", + "- `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.bayesian.sample(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.bayesian.plot_trace()" + ] + }, + { + "cell_type": "markdown", + "id": "2b3ea7b4", + "metadata": {}, + "source": [ + "## Summarize the posterior\n", + "\n", + "`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." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8470375e", + "metadata": {}, + "outputs": [], + "source": [ + "analysis.bayesian.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.bayesian.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.bayesian.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.bayesian.extend(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.bayesian.save(path)` and `analysis.bayesian.load(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.** `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. `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." + ] + } + ], + "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/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/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..89126ecdf 100644 --- a/src/easydynamics/analysis/__init__.py +++ b/src/easydynamics/analysis/__init__.py @@ -3,8 +3,20 @@ from easydynamics.analysis.analysis import Analysis 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', + 'BoundsSuggestion', + 'BoundsSuggestions', 'ParameterAnalysis', + 'ParameterLabels', + 'ParameterPosterior', + 'PosteriorSampler', + 'PosteriorSummary', ] diff --git a/src/easydynamics/analysis/analysis1d.py b/src/easydynamics/analysis/analysis1d.py index e50b09737..5c8a68be3 100644 --- a/src/easydynamics/analysis/analysis1d.py +++ b/src/easydynamics/analysis/analysis1d.py @@ -12,6 +12,8 @@ from plopp.backends.matplotlib.figure import InteractiveFigure from easydynamics.analysis.analysis_base import AnalysisBase +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 @@ -31,6 +33,10 @@ class Analysis1d(AnalysisBase): Is used primarily in the Analysis class, but can also be used on its own for simpler analyses. + 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 -------- **Fitting a single Q slice** @@ -116,6 +122,9 @@ def __init__( self._fit_result = None self._convolver = None self._convolver_is_dirty = True + self._fitter = None + self._fitter_is_dirty = True + self._bayesian = None super().__init__( display_name=display_name, @@ -245,27 +254,122 @@ 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._sampling_data() + fit_result = self.fitter.fit(x=x, y=y, weights=weights) - fitter = EasyScienceFitter( - fit_object=self, - fit_function=self.as_fit_function(), - ) + 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() + + ############# + # The contract PosteriorSampler relies on + ############# + + def _parameter_labels(self) -> ParameterLabels: + """ + Get labels for the chain's parameters. + + A single Q index holds one copy of each parameter, so nothing needs qualifying. + + Returns + ------- + ParameterLabels + Labels over the current free parameters. + """ + return ParameterLabels(self._chain_parameters()) + + 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. + + 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 _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 +587,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 +628,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 +643,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/posterior.py b/src/easydynamics/analysis/posterior.py new file mode 100644 index 000000000..d8d65af77 --- /dev/null +++ b/src/easydynamics/analysis/posterior.py @@ -0,0 +1,668 @@ +# 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. + 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. + 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 + label: str + 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)' + + 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})' + if s.needs_attention: + suggested = f'-- {s.reason}' + else: + suggested = f'({s.suggested_min:.4g}, {s.suggested_max:.4g})' + lines.append(f'{s.label:<{width}s} {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], + labels: list[str] | None = None, + 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. + 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 + 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') + + 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, 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, +) -> BoundsSuggestion: + """ + Propose bounds for a single parameter. + + Parameters + ---------- + 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 + 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, + label=label, + 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, + label=label, + suggested_min=current_min, + suggested_max=current_max, + reason='value is not finite', + ) + + # A NaN uncertainty is exactly the degenerate fit this helper exists to guard against, so it + # is flagged rather than silently treated like a zero error, which would yield deceptively + # tight bounds of value +/- relative_pad * |value|. + if not np.isfinite(error): + return BoundsSuggestion( + parameter=parameter, + label=label, + suggested_min=current_min, + suggested_max=current_max, + reason='fitted uncertainty is not finite', + ) + + half_width = relative_pad * abs(value) + 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, + label=label, + suggested_min=current_min, + suggested_max=current_max, + reason='no scale information (zero value and uncertainty)', + ) + + 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='', + ) + + +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 degenerate_parameters(parameters: list[Parameter]) -> list[Parameter]: + """ + Find parameters whose finite bounds enclose no range at all. + + A zero-width range (``min >= max``) gives DREAM nothing to explore: as the prior it has zero + volume, and letting it through surfaces only as NaNs deep inside the sampler, far from the + cause. + + Parameters + ---------- + parameters : list[Parameter] + The parameters to check. + + Returns + ------- + list[Parameter] + Those parameters whose bounds are both finite with ``min >= max``. + """ + return [ + parameter + for parameter in parameters + if np.isfinite(parameter.min) + and np.isfinite(parameter.max) + and float(parameter.min) >= float(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 the parameter's ``unique_name`` -- ``name`` is not used as the key because two + same-named parameters would collide -- 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``. The caller resolves the unique names back to + readable labels where the result is reported. + """ + if draws.shape[0] == 0: + return {} + 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.unique_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)' + + width = max(len('parameter'), *(len(e.name) for e in self._entries)) + header = ( + f'{"parameter":<{width}s} {"unit":>10s} {"median":>14s} ' + f'{"-":>12s} {"+":>12s} {"current":>14s}' + ) + lines = ['PosteriorSummary', header, '-' * len(header)] + lines.extend( + 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 + ) + return '\n'.join(lines) + + +def summarize_draws( + draws: np.ndarray, + labels: 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)``. + 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. + + 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=labels[column], + 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/analysis/posterior_labels.py b/src/easydynamics/analysis/posterior_labels.py new file mode 100644 index 000000000..fe406f706 --- /dev/null +++ b/src/easydynamics/analysis/posterior_labels.py @@ -0,0 +1,200 @@ +# 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} + # Display labels can still collide after qualification -- two same-named parameters with + # no qualifier, or one that declines. A colliding label cannot round-trip through a saved + # chain: both columns would silently resolve to whichever parameter was registered last. + # So the labels used as lookup keys, and written to the sidecar by name_map(), carry a + # deterministic positional suffix wherever they collide, while label() keeps the bare + # display name. + label_counts = Counter(self.label(p) for p in self._parameters) + occurrence: Counter = Counter() + self._storage_labels: dict[str, str] = {} + for p in self._parameters: + base = self.label(p) + if label_counts[base] > 1: + occurrence[base] += 1 + self._storage_labels[p.unique_name] = f'{base} [{occurrence[base]}]' + else: + self._storage_labels[p.unique_name] = base + self._by_label = {self._storage_labels[p.unique_name]: 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. Where two parameters share a display label, + the recorded labels carry a deterministic positional suffix (``width [1]``, ``width [2]``) + so each column can be matched back to exactly one parameter. + + Returns + ------- + dict[str, str] + Mapping of unique name to label, collision-free. + """ + return dict(self._storage_labels) + + 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..2c23a9235 --- /dev/null +++ b/src/easydynamics/analysis/posterior_sampling.py @@ -0,0 +1,1278 @@ +# 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 sys +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 degenerate_parameters +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, or finite bounds that + enclose no range (``min >= max``). + """ + labels = self._labels() + unbounded = unbounded_parameters(labels.parameters) + if unbounded: + 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.' + ) + degenerate = degenerate_parameters(labels.parameters) + if degenerate: + names = ', '.join(labels.label(parameter) for parameter in degenerate) + raise ValueError( + f'Bayesian sampling requires min < max on every free parameter, because the ' + f'bounds act as the prior and a zero-width range leaves the sampler nothing to ' + f'explore. These parameters have degenerate bounds: {names}. ' + f'Widen their min and max, or fix them instead of sampling them.' + ) + + ############# + # Sampling + ############# + + def sample( + self, + samples: int = 10000, + burn: int = 2000, + thin: int = 10, + population: int | None = None, + parameters: list[Parameter] | list[str] | None = None, + progress: bool = False, + **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. + progress : bool, default=False + Print a progress line, redrawn in place as the sampler advances and closed with a done + marker when the run finishes. Off by default so scripted runs stay quiet; a + ``progress_callback`` given in ``sampler_options`` takes precedence over it. + **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`. + + Notes + ----- + Runs are not reproducible. BUMPS' DREAM sampler draws from NumPy's global random state and + the underlying EasyScience Sampler exposes no seed control, so two identical calls return + two different chains. Their summaries should nevertheless agree to well within the reported + credible intervals; if they do not, the chain is too short to have converged. + """ + reporter = _install_progress_reporter(progress, sampler_options) + completed = False + try: + results = self._run( + parameters=parameters, + run=lambda sampler: sampler.sample( + samples=samples, burn=burn, thin=thin, population=population, **sampler_options + ), + ) + completed = True + finally: + if reporter is not None: + reporter.close(completed=completed) + return results + + def extend( + self, + additional_samples: int = 5000, + thin: int = 10, + parameters: list[Parameter] | list[str] | None = None, + progress: bool = False, + **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. + progress : bool, default=False + Print a progress line, redrawn in place as the sampler advances, as in :meth:`sample`. + **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, or the previous run failed and left no results. + + Notes + ----- + A ``ValueError`` propagates from the run guards if the model or data changed since the + chain was started, or if this run's parameters differ from the ones the chain holds. + + Like :meth:`sample`, extensions are not reproducible: the sampler draws from NumPy's global + random state and exposes no seed control. + """ + if self._sampler is None: + raise RuntimeError('No chain to extend. Call sample() or load() first.') + reporter = _install_progress_reporter(progress, sampler_options) + completed = False + try: + results = self._run( + parameters=parameters, + run=lambda sampler: sampler.extend( + additional_samples=additional_samples, thin=thin, **sampler_options + ), + reuse_sampler=True, + ) + completed = True + finally: + if reporter is not None: + reporter.close(completed=completed) + return results + + 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. + ValueError + If there are no free parameters to sample. + """ + 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() + if not chain_parameters: + raise ValueError( + 'There are no free parameters to sample: every parameter is fixed. ' + 'Free at least one parameter before sampling.' + ) + 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, as an extension must. + + Returns + ------- + Sampler + The Sampler to run. + + Raises + ------ + ValueError + If the cached Sampler must be reused but the model or data has changed since it was + built, so continuing its chain would silently mix draws against different data. + """ + if reuse_sampler and self._sampler is not None and self._sampler_is_dirty: + raise ValueError( + 'Cannot extend the chain: the model or data has changed since the chain was ' + 'started, and an extension would mix draws taken against different data. ' + 'Start a fresh chain with sample() instead.' + ) + 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, both in count and in identity. + + Parameters + ---------- + chain_parameters : list[Parameter] + The parameters that would form the chain for this run. + + Raises + ------ + RuntimeError + If there are no stored results to continue from, as after a failed run. + ValueError + If the number or the identity of the parameters differs from the existing chain's. + """ + if self._results is None: + raise RuntimeError( + 'Cannot extend: the previous run failed and left no results to continue from. ' + 'Start a fresh chain with sample() instead.' + ) + 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.' + ) + + # An equal count is not enough: the columns must be draws of the same parameters. For a + # chain from this session the stored column names are current unique names; for a loaded + # chain they are foreign, so they are resolved through the saved labels instead. + requested = {parameter.unique_name for parameter in chain_parameters} + if set(self._results.param_names) == requested: + return + resolved = self._resolve(self._results) + if ( + all(parameter is not None for parameter in resolved) + and {parameter.unique_name for parameter in resolved} == requested + ): + return + labels = self._labels() + chain_names = ', '.join(self._display_names(self._results)) + run_names = ', '.join(labels.label(parameter) for parameter in chain_parameters) + raise ValueError( + f'Cannot extend the chain: it holds draws of [{chain_names}], but this run would ' + f'sample [{run_names}]. An extension continues the stored chain, whose columns are ' + f'fixed, so it needs the same parameters the chain was started with. Start a fresh ' + f'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 or Parameter matches no free parameter of this analysis, 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} + by_unique_name = {parameter.unique_name: 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'): + # A Parameter object gets the same membership check a label does. Without it a + # fixed or foreign parameter slips through, every free parameter ends up held + # fixed, and the run dies with a cryptic zero-parameter failure deep in BUMPS. + if entry.unique_name not in by_unique_name: + name = getattr(entry, 'name', entry.unique_name) + raise ValueError( + f'Parameter {name!r} is not a free parameter of this analysis, so it ' + f'cannot be sampled. It is either fixed or not part of this analysis. ' + f'Available: {", ".join(sorted(by_label))}.' + ) + requested.append(by_unique_name[entry.unique_name]) + 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 + labels = self._labels() + by_unique_name = {parameter.unique_name: parameter for parameter in labels.parameters} + details = ', '.join( + f'{labels.label(by_unique_name[unique_name])} ({fraction:.0%} of draws)' + if unique_name in by_unique_name + else f'{unique_name} ({fraction:.0%} of draws)' + for unique_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) + if not self._saved_labels: + # A chain loaded without a sidecar has no labels to record. Writing an empty sidecar + # would be worse than none: the next load() would find a "valid" file, warn about + # nothing, and report every column under its raw internal name. + warnings.warn( + ( + f'No parameter labels are recorded for this chain, so no parameter-name ' + f'sidecar was written next to {path}; the chain was probably loaded without ' + f'one. A future load() will report the columns under their internal names.' + ), + UserWarning, + stacklevel=2, + ) + return + 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}') + self._saved_labels = ( + json.loads(sidecar.read_text(encoding='utf-8')) if sidecar.is_file() else {} + ) + if not self._saved_labels: + # An empty sidecar is as unusable as a missing one, so both warn the same way. + warnings.warn( + ( + f'No parameter-name sidecar with usable content found at {sidecar}. The ' + f'chain will be reported under the internal names it was saved with, because ' + f'those cannot be matched 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_marginal(self, parameter: Parameter | str, **kwargs: dict[str, Any]) -> Figure: + """ + Plot the marginal posterior distribution of a single sampled parameter. + + Shows a density-normalized histogram of the parameter's draws, with the median and the 16th + and 84th percentiles marked -- the same 68% credible interval :meth:`summary` reports. + + Parameters + ---------- + parameter : Parameter | str + The parameter to plot, as a Parameter object or its label. + **kwargs : dict[str, Any] + Forwarded to :func:`easydynamics.utils.posterior_plotting.plot_marginal`. + + Returns + ------- + Figure + The matplotlib Figure. + """ + from easydynamics.utils.posterior_plotting import plot_marginal + + results = self._require_results() + column = self._resolve_column(results, parameter) + return plot_marginal( + values=results.draws[:, column], + name=self._display_names(results)[column], + unit=self._units(results)[column], + title=self._analysis.display_name, + **kwargs, + ) + + def plot_correlations(self, **kwargs: dict[str, Any]) -> Figure: + """ + Plot the Pearson correlation matrix of the sampled parameters. + + A strongly correlated pair cannot be determined separately from this data. The matrix + condenses what the off-diagonal panels of :meth:`plot_corner` show, one number per pair, + which scales better to many parameters. + + Parameters + ---------- + **kwargs : dict[str, Any] + Forwarded to :func:`easydynamics.utils.posterior_plotting.plot_correlations`. + + Returns + ------- + Figure + The matplotlib Figure. + """ + from easydynamics.utils.posterior_plotting import plot_correlations + + results = self._require_results() + return plot_correlations( + draws=results.draws, + names=self._display_names(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})') + + # When the data carries no variances the weights are all-ones placeholders, and inverting + # them would fabricate error bars of 1.0 that the data never had. + experiment = getattr(self._analysis, 'experiment', None) + has_variances = experiment is None or getattr(experiment, 'has_variances', True) + + return plot_posterior_predictive( + x=np.asarray(x), + y=np.asarray(y), + predictions=self.predictions(n_draws), + y_err=1.0 / np.asarray(weights) if weights is not None and has_variances else None, + 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 _resolve_column(self, results: SamplingResults, parameter: Parameter | str) -> int: + """ + Find the chain column holding a parameter's draws. + + Labels are matched against the columns' display names, so the same names the summary and + the plots report under are the ones accepted here. Parameter objects are matched through + the resolved columns, so a parameter reloaded from a saved chain is found too. + + Parameters + ---------- + results : SamplingResults + The results whose columns should be searched. + parameter : Parameter | str + The parameter to look for, as a Parameter object or its label. + + Returns + ------- + int + The index of the column holding the parameter's draws. + + Raises + ------ + TypeError + If parameter is neither a Parameter object nor a string. + ValueError + If the parameter matches no column of the chain. + """ + names = self._display_names(results) + if isinstance(parameter, str): + matches = [column for column, name in enumerate(names) if name == parameter] + elif hasattr(parameter, 'unique_name'): + matches = [ + column + for column, candidate in enumerate(self._resolve(results)) + if candidate is not None and candidate.unique_name == parameter.unique_name + ] + else: + raise TypeError('parameter must be a Parameter object or a label (string).') + if not matches: + requested = ( + parameter if isinstance(parameter, str) else getattr(parameter, 'name', '?') + ) + raise ValueError( + f'No sampled parameter named {requested!r}. Available: {", ".join(sorted(names))}.' + ) + return matches[0] + + 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 _install_progress_reporter( + progress: bool, + sampler_options: dict[str, Any], +) -> _SamplingProgress | None: + """ + Put a progress reporter into the sampler options when one is asked for. + + A ``progress_callback`` the caller supplied themselves is left untouched, since an explicit + callback is more specific than the boolean convenience flag. + + Parameters + ---------- + progress : bool + Whether a progress line was requested. + sampler_options : dict[str, Any] + The options about to be forwarded to the EasyScience Sampler, modified in place. + + Returns + ------- + _SamplingProgress | None + The installed reporter, which the caller must close after the run, or None when nothing was + installed. + """ + if not progress or 'progress_callback' in sampler_options: + return None + reporter = _SamplingProgress() + sampler_options['progress_callback'] = reporter + return reporter + + +class _SamplingProgress: + """ + Renders the sampler's per-generation callbacks as a single self-overwriting progress line. + + BUMPS invokes the callback once per DREAM generation, which for a long run is far too often to + print, so the line is only redrawn when the percentage changes. Carriage-return output works in + terminals and notebooks alike, and needs no extra dependency. + + The generation total in the payload is the backend's own estimate, and it overestimates when + DREAM runs more chains than the estimate assumes, so a finished run can stop short of 100%. The + line is therefore closed with an explicit done marker rather than trusting the estimate. + """ + + def __init__(self) -> None: + self._last_percent = -1 + self._line_length = 0 + self._printed = False + + def __call__(self, payload: dict[str, Any]) -> None: + """ + Handle one progress callback from the sampler. + + Parameters + ---------- + payload : dict[str, Any] + The sampler's progress payload. ``iteration`` carries the DREAM generation and + ``total_steps``, when present, the estimated total number of generations. + """ + iteration = payload.get('iteration') + if iteration is None: + return + total = payload.get('total_steps') + if total: + # Clamped, so the line never reports more than 100% when the run outlives the + # backend's estimate of its own length. + percent = min(100, int(100 * iteration / total)) + if percent == self._last_percent: + return + self._last_percent = percent + line = f'Sampling: {percent:3d}% ({iteration}/{total} generations)' + else: + line = f'Sampling: generation {iteration}' + self._write(line) + + def close(self, completed: bool) -> None: + """ + End the progress line, so any later output starts on a line of its own. + + Parameters + ---------- + completed : bool + Whether the run finished. A finished run gets a done marker; a failed one only has its + line terminated, so the exception is not decorated with a claim of success. + """ + if not self._printed: + return + if completed: + self._write('Sampling: done') + sys.stdout.write('\n') + sys.stdout.flush() + + def _write(self, line: str) -> None: + """ + Redraw the progress line in place. + + Parameters + ---------- + line : str + The text to show, padded so it fully overwrites a longer previous line. + """ + sys.stdout.write(f'\r{line.ljust(self._line_length)}') + sys.stdout.flush() + self._line_length = max(self._line_length, len(line)) + self._printed = True + + +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/experiment/experiment.py b/src/easydynamics/experiment/experiment.py index d064326f5..5be71f109 100644 --- a/src/easydynamics/experiment/experiment.py +++ b/src/easydynamics/experiment/experiment.py @@ -585,6 +585,21 @@ def _extract_x_y_var(self, Q_index: int) -> tuple[np.ndarray, np.ndarray, np.nda var = data.variances return x, y, var + @property + def has_variances(self) -> bool: + """ + Whether the data carries variances. + + When it does not, :meth:`extract_x_y_weights_only_finite` falls back to all-ones weights, + which are placeholders for the fit rather than measured uncertainties. + + Returns + ------- + bool + True when there is data and it has variances. + """ + return self._binned_data is not None and self._binned_data.variances is not None + def extract_x_y_weights_only_finite( self, Q_index: int ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: 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..c3576d164 --- /dev/null +++ b/src/easydynamics/utils/posterior_plotting.py @@ -0,0 +1,592 @@ +# 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 + +import warnings +from typing import TYPE_CHECKING + +import matplotlib.pyplot as plt +import numpy as np +from matplotlib import colormaps +from matplotlib.ticker import MaxNLocator + +if TYPE_CHECKING: + from matplotlib.figure import Figure + + +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: + """ + 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. + + 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, one per draw, 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 + Figure size in inches. Defaults to a height that scales with the number of panels. + + Returns + ------- + Figure + The matplotlib Figure. + + Raises + ------ + ValueError + If ``draws`` is not two-dimensional or is empty, if ``names`` does not have one entry per + column, or if ``logp`` does not have one entry per draw. + """ + draws = np.asarray(draws) + _verify_draws(draws, names) + if logp is not None: + logp = np.asarray(logp) + if logp.ndim != 1 or logp.shape[0] != draws.shape[0]: + raise ValueError( + f'logp must have one entry per draw. ' + f'Got shape {logp.shape} for {draws.shape[0]} draws.' + ) + + 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(_with_unit(name, units, column), fontsize=8) + # A single draw would make (0, len - 1) a zero-width range; matplotlib's autoscaling + # handles that case better than an explicit degenerate limit would. + if len(draws) > 1: + axis.set_xlim(0, len(draws) - 1) + + if logp is not None: + axes[-1].plot(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], + units: list[str] | None = None, + 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. + + Parameters + ---------- + draws : np.ndarray + 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 + 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. + + Raises + ------ + ValueError + If ``draws`` is not two-dimensional or is empty, if ``names`` does not have one entry per + column, or if any column contains non-finite values. + """ + draws = np.asarray(draws) + _verify_draws(draws, names) + + # Caught up front, because numpy would otherwise report it as an obscure + # "range [nan, nan]" error from inside the histogram. + finite_columns = np.isfinite(draws).all(axis=0) + if not finite_columns.all(): + bad = ', '.join(name for name, ok in zip(names, finite_columns, strict=True) if not ok) + raise ValueError(f'draws contain non-finite values (NaN or infinity) in: {bad}.') + + n = draws.shape[1] + if figsize is None: + side = max(4.0, 2.0 * n) + figsize = (side, side) + + # One shared limit per column, applied to the diagonal histogram and every hexbin panel below + # it, so the ticks of a column line up instead of each panel autoscaling on its own. + limits = _column_limits(draws) + + 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) + axis.set_ylim(limits[row]) + axis.set_xlim(limits[col]) + 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([]) + 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) + fig.tight_layout() + return fig + + +def plot_marginal( + values: np.ndarray, + name: str, + unit: str | None = None, + title: str | None = None, + bins: int = 40, + figsize: tuple[float, float] = (8.0, 5.0), +) -> Figure: + """ + Plot the marginal posterior distribution of a single parameter. + + Shows a density-normalized histogram of the parameter's draws, with the median and the 16th and + 84th percentiles marked -- the same 68% credible interval the posterior summary reports. + + Parameters + ---------- + values : np.ndarray + The parameter's posterior draws, one-dimensional. + name : str + The label the parameter is reported under. + unit : str | None, default=None + The parameter's unit, appended to the axis label. Empty or dimensionless units are skipped, + since a bare "dimensionless" only adds clutter. + title : str | None, default=None + Figure title. + bins : int, default=40 + Number of histogram bins. + figsize : tuple[float, float], default=(8.0, 5.0) + Figure size in inches. + + Returns + ------- + Figure + The matplotlib Figure. + + Raises + ------ + ValueError + If ``values`` is not one-dimensional, is empty, or contains non-finite entries. + """ + values = np.asarray(values) + if values.ndim != 1: + raise ValueError(f'values must be one-dimensional. Got shape {values.shape}.') + if values.size == 0: + raise ValueError('values is empty: there are no samples to plot.') + # Caught up front, because numpy would otherwise report it as an obscure + # "range [nan, nan]" error from inside the histogram. + if not np.isfinite(values).all(): + raise ValueError(f'values contain non-finite entries (NaN or infinity) for {name}.') + + lower, median, upper = np.percentile(values, [16.0, 50.0, 84.0]) + + fig, axis = plt.subplots(figsize=figsize) + axis.hist(values, bins=bins, density=True, color='C0', histtype='stepfilled', alpha=0.7) + axis.axvline(median, color='C3', lw=1.5, label='Median') + axis.axvline(lower, color='C3', lw=1.0, ls='--', label='68% credible interval') + axis.axvline(upper, color='C3', lw=1.0, ls='--') + axis.set_xlabel(_with_unit(name, [unit] if unit is not None else None, 0)) + axis.set_ylabel('Probability density') + axis.legend() + if title is not None: + axis.set_title(title) + fig.tight_layout() + return fig + + +def plot_correlations( + draws: np.ndarray, + names: list[str], + title: str | None = None, + figsize: tuple[float, float] | None = None, +) -> Figure: + """ + Plot the Pearson correlation matrix of the sampled parameters. + + A strongly correlated pair (an entry near +1 or -1) cannot be determined separately from this + data: the chain trades one off against the other. The matrix condenses what the off-diagonal + panels of the corner plot show, one number per pair, which scales better to many parameters. + + Correlations are dimensionless, so the labels carry no units. A constant column has no defined + correlation with anything; its cells are shown greyed out and marked "n/a" rather than failing. + A ``ValueError`` propagates from the input validation if ``draws`` is not two-dimensional or + is empty, 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. + figsize : tuple[float, float] | None, default=None + Figure size in inches. Defaults to a square that scales with the parameter count, plus room + for the colorbar. + + Returns + ------- + Figure + The matplotlib Figure. + """ + draws = np.asarray(draws) + _verify_draws(draws, names) + + matrix = _correlation_matrix(draws) + n = draws.shape[1] + if figsize is None: + side = max(4.0, 0.9 * n + 2.0) + figsize = (side + 1.5, side) + + # A diverging map centred on zero, so positive and negative correlations read as two hues + # around a neutral midpoint. Cells with no defined correlation are greyed out. + colormap = colormaps['RdBu_r'].with_extremes(bad='0.85') + + fig, axis = plt.subplots(figsize=figsize) + image = axis.imshow(np.ma.masked_invalid(matrix), cmap=colormap, vmin=-1.0, vmax=1.0) + axis.set_xticks(range(n), labels=names, rotation=45, ha='right', fontsize=8) + axis.set_yticks(range(n), labels=names, fontsize=8) + for row in range(n): + for col in range(n): + value = matrix[row, col] + defined = bool(np.isfinite(value)) + axis.text( + col, + row, + f'{value:.2f}' if defined else 'n/a', + ha='center', + va='center', + fontsize=8, + # Saturated cells at the ends of the map are too dark for black text. + color='white' if defined and abs(value) > 0.6 else 'black', + ) + fig.colorbar(image, ax=axis, label='Pearson correlation') + if title is not None: + axis.set_title(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, + xlabel: str | None = None, + ylabel: str | None = None, + 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. + 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. + + 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') + 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) + fig.tight_layout() + return fig + + +def _column_limits(draws: np.ndarray) -> list[tuple[float, float]]: + """ + Compute one shared axis range per column of a corner plot. + + Parameters + ---------- + draws : np.ndarray + Posterior draws, shape ``(n_draws, n_parameters)``, all finite. + + Returns + ------- + list[tuple[float, float]] + A padded ``(low, high)`` range per column, widened to a usable span when a column is + constant. + """ + lows = draws.min(axis=0) + highs = draws.max(axis=0) + spans = highs - lows + pads = np.where(spans > 0, 0.05 * spans, 0.05 * np.maximum(np.abs(highs), 1.0)) + return [(float(low), float(high)) for low, high in zip(lows - pads, highs + pads, strict=True)] + + +def _correlation_matrix(draws: np.ndarray) -> np.ndarray: + """ + Compute the Pearson correlation matrix of a chain's columns. + + Parameters + ---------- + draws : np.ndarray + Posterior draws, shape ``(n_draws, n_parameters)``. + + Returns + ------- + np.ndarray + The ``(n_parameters, n_parameters)`` correlation matrix, two-dimensional even for a + single-parameter chain, with NaN wherever a column has zero variance. Numpy's + division-by-zero warnings for those columns are suppressed, since the NaNs are handled by + the caller rather than being a numerical accident. + """ + with np.errstate(invalid='ignore', divide='ignore'), warnings.catch_warnings(): + warnings.simplefilter('ignore', RuntimeWarning) + matrix = np.corrcoef(draws, rowvar=False) + # np.corrcoef collapses a single-column input to a 0-d scalar; restore the 1x1 matrix. + return np.atleast_2d(np.asarray(matrix, dtype=float)) + + +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. + + 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[0] == 0: + raise ValueError('draws is empty: there are no samples to plot.') + if draws.shape[1] == 0: + raise ValueError('draws has no columns: there are no parameters to plot.') + 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..f341a1c84 --- /dev/null +++ b/tests/integration/fitting/test_bayesian_sampling.py @@ -0,0 +1,234 @@ +# 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.bayesian.suggest_bounds().apply() + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + analysis.bayesian.sample(**SAMPLE_KWARGS) + return analysis + + +class TestRealChain: + def test_chain_has_one_column_per_free_parameter(self, sampled_analysis): + # THEN + results = sampled_analysis.bayesian.results + + # EXPECT + 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): + # THEN + 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 + # 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): + # THEN + summary = sampled_analysis.bayesian.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.bayesian.suggest_bounds().apply() + before = [float(p.value) for p in analysis.get_free_parameters()] + + # THEN + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + analysis.bayesian.sample(**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.bayesian.results.state.Ngen) + + # THEN + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + extended = sampled_analysis.bayesian.extend( + 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.bayesian.save(prefix) + + fresh = build_analysis() + fresh.fit() + fresh.bayesian.suggest_bounds().apply() + + # THEN + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + fresh.bayesian.load(prefix) + + # EXPECT the reloaded chain is reported under real names, not internal unique names + 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) + + def test_subset_sampling_produces_a_single_column(self): + # WHEN + analysis = build_analysis() + analysis.fit() + analysis.bayesian.suggest_bounds().apply() + + # THEN + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + results = analysis.bayesian.sample(parameters=['Gaussian width'], **SAMPLE_KWARGS) + + # EXPECT + assert results.draws.shape[1] == 1 + 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()) + + # THEN + 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 + assert len(corner.axes) == n_parameters**2 + assert len(predictive.axes) == 1 + plt.close('all') + + def test_marginal_and_correlation_figures_render(self, sampled_analysis): + # WHEN + import matplotlib.pyplot as plt + + n_parameters = len(sampled_analysis.get_free_parameters()) + + # THEN + marginal = sampled_analysis.bayesian.plot_marginal('Gaussian width') + correlations = sampled_analysis.bayesian.plot_correlations() + + # EXPECT a real chain renders both figures + assert len(marginal.axes) == 1 + matrix = correlations.axes[0].images[0].get_array() + assert matrix.shape == (n_parameters, n_parameters) + assert np.asarray(np.diag(matrix)) == pytest.approx(np.ones(n_parameters)) + plt.close('all') + + def test_posterior_median_is_close_to_the_least_squares_fit(self): + # WHEN + analysis = build_analysis() + analysis.fit() + analysis.bayesian.suggest_bounds().apply() + fitted = {p.name: float(p.value) for p in analysis.get_free_parameters()} + + # THEN + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + 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 + 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.py b/tests/unit/easydynamics/analysis/test_analysis1d.py index 85804738f..99ab53ab5 100644 --- a/tests/unit/easydynamics/analysis/test_analysis1d.py +++ b/tests/unit/easydynamics/analysis/test_analysis1d.py @@ -8,6 +8,7 @@ import numpy as np import pytest import scipp as sc +from easyscience.fitting import AvailableMinimizers from easyscience.variable import Parameter from easydynamics.analysis.analysis1d import Analysis1d @@ -132,6 +133,77 @@ def test__calculate_adds_sample_and_background(self, analysis1d): analysis1d._evaluate_with_convolution.assert_called_once() analysis1d._evaluate_direct.assert_called_once() + ############# + # The cached fitter + ############# + + @pytest.fixture + def fittable_analysis1d(self): + # The analysis1d fixture holds three points against three free parameters, which leaves + # a fit with no degrees of freedom. This one has a curve to land on. + 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='TestFittable', + 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 test_fitter_is_built_lazily_and_cached(self, analysis1d): + # THEN + fitter = analysis1d.fitter + + # EXPECT + assert fitter is analysis1d.fitter + assert fitter.fit_object is analysis1d + + def test_fitter_is_rebuilt_when_the_sample_model_changes(self, analysis1d): + # WHEN + original = analysis1d.fitter + + # THEN + analysis1d.sample_model = SampleModel(components=Gaussian(area=1.0)) + + # EXPECT + assert analysis1d.fitter is not original + + def test_minimizer_can_be_switched_through_the_fitter(self, analysis1d): + # THEN + analysis1d.fitter.switch_minimizer(AvailableMinimizers.Bumps) + + # EXPECT + assert analysis1d.fitter.minimizer.enum == AvailableMinimizers.Bumps + + def test_fit_uses_the_persistent_fitter(self, fittable_analysis1d): + # THEN + result = fittable_analysis1d.fit() + + # EXPECT + assert result is fittable_analysis1d._fit_result + assert np.isfinite(result.reduced_chi2) + + ############# + # Fitting + ############# + def test_fit_raises_if_no_experiment(self, analysis1d): # WHEN THEN analysis1d._experiment = None diff --git a/tests/unit/easydynamics/analysis/test_posterior.py b/tests/unit/easydynamics/analysis/test_posterior.py new file mode 100644 index 000000000..28817e340 --- /dev/null +++ b/tests/unit/easydynamics/analysis/test_posterior.py @@ -0,0 +1,366 @@ +# 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 degenerate_parameters +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_error_is_flagged_not_silently_narrowed(self): + # WHEN a degenerate fit reports a NaN uncertainty + parameter = make_parameter(value=5.0) + parameter.variance = np.nan + + # THEN + suggestion = suggest_bounds_for_parameters([parameter]).suggestions[0] + + # EXPECT a flag, rather than deceptively tight bounds from the relative pad alone + assert suggestion.needs_attention + assert 'uncertainty is not finite' in suggestion.reason + + 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): + # THEN EXPECT + with pytest.raises(ValueError): + suggest_bounds_for_parameters([make_parameter()], **kwargs) + + def test_non_numeric_setting_raises(self): + # THEN 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]) + + # WHEN nothing has changed until apply is called + assert parameter.max == np.inf + + # THEN + 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): + # WHEN THEN EXPECT + assert 'no free parameters' in repr(BoundsSuggestions([])) + + def test_len_and_iteration(self): + # WHEN + suggestions = suggest_bounds_for_parameters([make_parameter(), make_parameter()]) + + # THEN 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 TestDegenerateParameters: + def test_finds_zero_width_ranges(self): + # WHEN one parameter's finite bounds enclose no range at all. The setters refuse identical + # bounds, but a deserialized or hand-built parameter can still carry them, so the internal + # state is written directly. + healthy = make_parameter(name='healthy', minimum=0.0, maximum=2.0) + degenerate = make_parameter(name='degenerate', value=1.0, minimum=0.0, maximum=1.0) + degenerate._min.value = 1.0 + + # THEN + result = degenerate_parameters([healthy, degenerate]) + + # EXPECT + assert result == [degenerate] + + def test_infinite_bounds_are_not_reported_as_degenerate(self): + # WHEN a bound is infinite, that is unboundedness rather than degeneracy + parameter = make_parameter() + + # THEN EXPECT + assert degenerate_parameters([parameter]) == [] + + +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.unique_name in result + assert result[parameter.unique_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.unique_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.unique_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 == {} + + def test_same_named_parameters_do_not_collide(self): + # WHEN two parameters share a name and both posteriors are pinned at a bound + first = make_parameter(name='width', minimum=0.0, maximum=1.0) + second = make_parameter(name='width', minimum=0.0, maximum=1.0) + draws = np.zeros((100, 2)) + + # THEN + result = parameters_at_bounds(draws, [first, second]) + + # EXPECT one entry per parameter, keyed so they cannot overwrite each other + assert len(result) == 2 + assert set(result) == {first.unique_name, second.unique_name} + + def test_zero_row_draws_return_nothing_rather_than_dividing_by_zero(self): + # WHEN + parameter = make_parameter(minimum=0.0, maximum=1.0) + draws = np.zeros((0, 1)) + + # THEN EXPECT + assert parameters_at_bounds(draws, [parameter]) == {} + + +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, ['Gaussian width'], [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_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)) + + # 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]) + + # THEN 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)), ['Gaussian area'], [parameter])) + + # EXPECT + assert 'Gaussian area' in text + assert 'median' 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..618c8284c --- /dev/null +++ b/tests/unit/easydynamics/analysis/test_posterior_labels.py @@ -0,0 +1,149 @@ +# 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')] + + # THEN + 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'} + + # THEN + 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') + + # THEN + 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') + + # THEN + 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)] + + # THEN 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') + + # THEN + 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): + # THEN + 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_colliding_labels_get_distinct_sidecar_entries(self): + # WHEN two parameters end up with the same display label + first, second = make_parameter('width'), make_parameter('width') + + # THEN + labels = ParameterLabels([first, second]) + + # EXPECT deterministic positional suffixes in the sidecar mapping, rather than a silent + # last-write-wins, while the display label stays bare + assert labels.name_map() == { + first.unique_name: 'width [1]', + second.unique_name: 'width [2]', + } + assert labels.label(first) == 'width' + + def test_colliding_labels_round_trip_to_their_own_parameters(self): + # WHEN a chain of two same-labelled parameters was saved in another session + old_first, old_second = make_parameter('width'), make_parameter('width') + saved = ParameterLabels([old_first, old_second]).name_map() + new_first, new_second = make_parameter('width'), make_parameter('width') + + # THEN + labels = ParameterLabels([new_first, new_second]) + resolved = labels.resolve([old_first.unique_name, old_second.unique_name], saved) + + # EXPECT each column finds its own parameter, not both the same one + assert resolved == [new_first, new_second] + + 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'} + + # THEN + 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') + + # THEN 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]) diff --git a/tests/unit/easydynamics/analysis/test_posterior_sampling.py b/tests/unit/easydynamics/analysis/test_posterior_sampling.py new file mode 100644 index 000000000..5fb5cd24e --- /dev/null +++ b/tests/unit/easydynamics/analysis/test_posterior_sampling.py @@ -0,0 +1,907 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +""" +Unit tests for the posterior sampler, driven through an 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 easyscience.variable import Parameter + +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.posterior_sampling.Sampler' + + +def make_analysis(with_variances=True): + 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, :] if with_variances else 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 TestPosteriorSampler: + ############# + # Bounds pre-flight + ############# + + def test_sampling_refuses_unbounded_parameters(self, analysis): + # THEN EXPECT + with pytest.raises(ValueError, match='finite bounds'): + analysis.bayesian.sample(samples=10) + + def test_error_names_the_offending_parameters(self, analysis): + # THEN EXPECT + with pytest.raises(ValueError, match='Gaussian area'): + analysis.bayesian.check_bounds() + + def test_bounded_parameters_pass(self, analysis): + # WHEN + bound_all(analysis) + + # THEN EXPECT: does not raise + analysis.bayesian.check_bounds() + + def test_suggest_bounds_covers_the_free_parameters(self, analysis): + # THEN + suggestions = analysis.bayesian.suggest_bounds() + + # EXPECT + assert len(suggestions) == len(analysis.get_free_parameters()) + + def test_degenerate_bounds_are_rejected(self, analysis): + # WHEN one parameter's bounds collapse to a zero-width range, which internal state can + # carry even though the setters refuse it + bound_all(analysis) + parameter = analysis.get_free_parameters()[0] + parameter._min.value = float(parameter.max) + + # THEN EXPECT + with pytest.raises(ValueError, match='degenerate bounds'): + analysis.bayesian.check_bounds() + + ############# + # Sampling + ############# + + 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()] + + # THEN + 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.bayesian.sample(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 = [] + + # THEN + 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.bayesian.sample(samples=10) + + # EXPECT + assert seen == [AvailableMinimizers.Bumps] + + def test_restores_the_minimizer_even_when_sampling_raises(self, analysis): + # WHEN + bound_all(analysis) + + # THEN + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = RuntimeError('boom') + with pytest.raises(RuntimeError, match='boom'): + analysis.bayesian.sample(samples=10) + + # EXPECT + assert analysis.fitter.minimizer.enum == AvailableMinimizers.LMFit_leastsq + + def test_forwards_sampling_arguments(self, analysis): + # WHEN + bound_all(analysis) + + # THEN + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + analysis.bayesian.sample(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) + + # THEN + with patch(SAMPLER_PATH) as sampler_class: + expected = fake_results(analysis) + sampler_class.return_value.sample.return_value = expected + returned = analysis.bayesian.sample(samples=10) + + # EXPECT + assert returned 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 + 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) + + # THEN EXPECT + with pytest.warns(UserWarning, match='piled up'): + analysis.bayesian.sample(samples=10) + + def test_sampling_with_no_free_parameters_raises(self, analysis): + # WHEN every parameter is fixed + for parameter in analysis.get_free_parameters(): + parameter.fixed = True + + # THEN EXPECT a clear refusal, rather than a zero-parameter failure deep in BUMPS + with pytest.raises(ValueError, match='no free parameters to sample'): + analysis.bayesian.sample(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) + + # THEN EXPECT + with warnings_as_errors(): + analysis.bayesian.sample(samples=10) + + ############# + # Parameter subsets + ############# + + def test_holds_other_parameters_fixed_during_the_run(self, analysis): + # WHEN + bound_all(analysis) + target = analysis.get_free_parameters()[0] + seen = {} + + # THEN + 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.bayesian.sample(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] + + # THEN + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + with pytest.warns(UserWarning): + analysis.bayesian.sample(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) + + # THEN EXPECT + with pytest.raises(ValueError, match='No free parameter named'): + analysis.bayesian.sample(samples=10, parameters=['not a parameter']) + + def test_fixed_parameter_object_is_rejected(self, analysis): + # WHEN a Parameter object that is currently fixed is requested + bound_all(analysis) + target = analysis.get_free_parameters()[0] + target.fixed = True + + # THEN EXPECT the same membership check a label gets, instead of every free parameter + # ending up held fixed and BUMPS failing with zero parameters + with pytest.raises(ValueError, match='not a free parameter'): + analysis.bayesian.sample(samples=10, parameters=[target]) + + def test_parameter_from_another_model_is_rejected(self, analysis): + # WHEN + bound_all(analysis) + foreign = Parameter(name='foreign', value=1.0, unit='meV') + + # THEN EXPECT + with pytest.raises(ValueError, match='not a free parameter'): + analysis.bayesian.sample(samples=10, parameters=[foreign]) + + def test_non_list_parameters_raises(self, analysis): + # THEN EXPECT + with pytest.raises(TypeError, match='must be a list'): + analysis.bayesian.sample(samples=10, parameters='Gaussian area') + + def test_empty_parameter_list_raises(self, analysis): + # THEN EXPECT + with pytest.raises(ValueError, match='at least one parameter'): + analysis.bayesian.sample(samples=10, parameters=[]) + + ############# + # Sampler caching + ############# + + def test_sampler_is_reused_between_runs(self, analysis): + # WHEN + bound_all(analysis) + + # THEN + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + 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 + + def test_changing_the_q_index_rebuilds_the_sampler(self, analysis): + # WHEN + bound_all(analysis) + + # THEN + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + analysis.bayesian.sample(samples=10) + analysis.Q_index = 0 + analysis.bayesian.sample(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) + + # THEN + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + analysis.bayesian.sample(samples=10) + + # EXPECT + 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) + assert np.array_equal(kwargs['weights'], expected_w) + + ############# + # Extending and persistence + ############# + + def test_extend_without_a_chain_raises(self, analysis): + # THEN EXPECT + with pytest.raises(RuntimeError, match='No chain to extend'): + analysis.bayesian.extend() + + def test_extend_delegates_to_the_sampler(self, analysis): + # WHEN + bound_all(analysis) + + # THEN + 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.bayesian.sample(samples=10) + analysis.bayesian.extend(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_extend_with_different_parameters_raises(self, analysis): + # WHEN a chain was sampled over one parameter + bound_all(analysis) + first, second = analysis.get_free_parameters()[:2] + + 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) + with pytest.warns(UserWarning, match='Holding these parameters fixed'): + analysis.bayesian.sample(samples=10, parameters=[first.name]) + + # THEN EXPECT extending with a different parameter, even at the same chain width, + # is refused rather than silently merging draws of different quantities + with ( + pytest.warns(UserWarning, match='Holding these parameters fixed'), + pytest.raises(ValueError, match='holds draws of'), + ): + analysis.bayesian.extend(parameters=[second.name]) + + def test_extend_after_a_data_change_raises(self, analysis): + # WHEN the data changed after the chain was started + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + analysis.bayesian.sample(samples=10) + analysis.bayesian.invalidate() + + # THEN EXPECT the stale chain is refused rather than silently continued + with pytest.raises(ValueError, match='model or data has changed'): + analysis.bayesian.extend() + + def test_extend_after_a_failed_run_raises(self, analysis): + # WHEN the previous run failed after building the sampler, leaving no results + 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.bayesian.sample(samples=10) + + # THEN EXPECT + with pytest.raises(RuntimeError, match='left no results'): + analysis.bayesian.extend() + + def test_save_without_a_chain_raises(self, analysis): + # THEN EXPECT + with pytest.raises(RuntimeError, match='No chain to save'): + analysis.bayesian.save('somewhere') + + def test_save_writes_the_parameter_name_sidecar(self, analysis, tmp_path): + # WHEN + import json + + bound_all(analysis) + + # THEN + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + 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' + 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) + + # THEN EXPECT + with pytest.warns(UserWarning, match='No parameter-name sidecar'): + analysis.bayesian.load(str(tmp_path / 'missing')) + + def test_load_with_an_empty_sidecar_warns_like_a_missing_one(self, analysis, tmp_path): + # WHEN a sidecar file exists but records no labels + bound_all(analysis) + (tmp_path / 'chain.parameter-names.json').write_text('{}', encoding='utf-8') + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.load_state.return_value = fake_results(analysis) + + # THEN EXPECT + with pytest.warns(UserWarning, match='No parameter-name sidecar'): + analysis.bayesian.load(str(tmp_path / 'chain')) + + def test_load_passes_skip_through(self, analysis, tmp_path): + # WHEN + bound_all(analysis) + + # THEN + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.load_state.return_value = fake_results(analysis) + with pytest.warns(UserWarning, match='No parameter-name sidecar'): + analysis.bayesian.load(str(tmp_path / 'chain'), skip=7) + + # EXPECT + assert sampler_class.return_value.load_state.call_args.kwargs['skip'] == 7 + + def test_save_after_a_sidecarless_load_writes_no_empty_sidecar(self, analysis, tmp_path): + # WHEN a chain was loaded without a sidecar, so there are no labels to record + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.load_state.return_value = fake_results(analysis) + with pytest.warns(UserWarning, match='No parameter-name sidecar'): + analysis.bayesian.load(str(tmp_path / 'original')) + + # THEN EXPECT saving warns instead of writing an empty sidecar, which a later load() + # would mistake for a valid one and resolve every column to raw names + with pytest.warns(UserWarning, match='no parameter-name sidecar was written'): + analysis.bayesian.save(str(tmp_path / 'resaved')) + + # EXPECT + assert not (tmp_path / 'resaved.parameter-names.json').exists() + + ############# + # Results + ############# + + def test_summary_without_sampling_raises(self, analysis): + # THEN EXPECT + with pytest.raises(RuntimeError, match='No posterior samples yet'): + analysis.bayesian.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.bayesian.sample(samples=10) + + # THEN + summary = analysis.bayesian.summary() + + # EXPECT + 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.bayesian.sample(samples=10) + + # THEN + changed = analysis.bayesian.set_parameters_to_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): + # THEN EXPECT + with pytest.raises(RuntimeError, match='No posterior samples yet'): + analysis.bayesian.set_parameters_to_median() + + ############# + # Plots + ############# + + 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.bayesian.sample(samples=10) + + # THEN EXPECT + with pytest.raises(ValueError, match='positive integer'): + analysis.bayesian.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.bayesian.sample(samples=10) + + before = [float(p.value) for p in parameters] + + # THEN + analysis.bayesian.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): + # THEN EXPECT + with pytest.raises(RuntimeError): + analysis.bayesian.plot_trace() + with pytest.raises(RuntimeError): + analysis.bayesian.plot_corner() + + def test_predictive_forwards_the_measured_error_bars(self, analysis): + # WHEN the data carries variances of 0.01, i.e. an uncertainty of 0.1 + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + analysis.bayesian.sample(samples=10) + + # THEN + with patch('easydynamics.utils.posterior_plotting.plot_posterior_predictive') as plot: + analysis.bayesian.plot_posterior_predictive(n_draws=2) + + # EXPECT + assert plot.call_args.kwargs['y_err'] == pytest.approx(np.full(20, 0.1)) + + def test_predictive_omits_error_bars_when_the_data_has_no_variances(self): + # WHEN the data has no variances, so the weights are all-ones placeholders + analysis = make_analysis(with_variances=False) + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + analysis.bayesian.sample(samples=10) + + # THEN + with patch('easydynamics.utils.posterior_plotting.plot_posterior_predictive') as plot: + analysis.bayesian.plot_posterior_predictive(n_draws=2) + + # EXPECT no error bars fabricated from the placeholder weights + assert plot.call_args.kwargs['y_err'] is None + + def test_marginal_forwards_the_resolved_column(self, analysis): + # WHEN the width column carries distinctive draws + bound_all(analysis) + parameters = analysis.get_free_parameters() + column = [p.name for p in parameters].index('Gaussian width') + draws = np.tile([float(p.value) for p in parameters], (30, 1)) + draws[:, column] += np.random.default_rng(0).normal(scale=0.01, size=30) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_results(analysis, values=draws) + analysis.bayesian.sample(samples=10) + + # THEN + with patch('easydynamics.utils.posterior_plotting.plot_marginal') as plot: + analysis.bayesian.plot_marginal('Gaussian width', bins=13) + + # EXPECT the label resolved to that column's draws, name and unit + kwargs = plot.call_args.kwargs + assert np.array_equal(kwargs['values'], draws[:, column]) + assert kwargs['name'] == 'Gaussian width' + assert kwargs['unit'] == 'meV' + assert kwargs['bins'] == 13 + + def test_marginal_accepts_a_parameter_object(self, analysis): + # WHEN + bound_all(analysis) + 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) + analysis.bayesian.sample(samples=10) + + # THEN + with patch('easydynamics.utils.posterior_plotting.plot_marginal') as plot: + analysis.bayesian.plot_marginal(target) + + # EXPECT + assert plot.call_args.kwargs['name'] == target.name + + def test_marginal_unknown_label_raises_naming_the_available_ones(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.bayesian.sample(samples=10) + + # THEN EXPECT + with pytest.raises(ValueError, match='No sampled parameter named') as excinfo: + analysis.bayesian.plot_marginal('not a parameter') + assert 'Gaussian width' in str(excinfo.value) + + def test_marginal_foreign_parameter_raises(self, analysis): + # WHEN + bound_all(analysis) + foreign = Parameter(name='foreign', value=1.0, unit='meV') + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + analysis.bayesian.sample(samples=10) + + # THEN EXPECT + with pytest.raises(ValueError, match='No sampled parameter named'): + analysis.bayesian.plot_marginal(foreign) + + def test_marginal_without_sampling_raises(self, analysis): + # THEN EXPECT + with pytest.raises(RuntimeError, match='No posterior samples yet'): + analysis.bayesian.plot_marginal('Gaussian width') + + def test_correlations_use_the_display_names(self, analysis): + # WHEN + bound_all(analysis) + parameters = analysis.get_free_parameters() + draws = np.tile([float(p.value) for p in parameters], (30, 1)) + draws += np.random.default_rng(0).normal(scale=0.01, size=draws.shape) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_results(analysis, values=draws) + analysis.bayesian.sample(samples=10) + + # THEN + with patch('easydynamics.utils.posterior_plotting.plot_correlations') as plot: + analysis.bayesian.plot_correlations() + + # EXPECT the chain's draws under the parameters' own names + kwargs = plot.call_args.kwargs + assert np.array_equal(kwargs['draws'], draws) + assert kwargs['names'] == [p.name for p in parameters] + + def test_correlations_without_sampling_raise(self, analysis): + # THEN EXPECT + with pytest.raises(RuntimeError, match='No posterior samples yet'): + analysis.bayesian.plot_correlations() + + ############# + # Progress reporting + ############# + + def test_progress_is_off_by_default(self, analysis): + # WHEN + bound_all(analysis) + + # THEN + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + analysis.bayesian.sample(samples=10) + + # EXPECT + assert 'progress_callback' not in sampler_class.return_value.sample.call_args.kwargs + + def test_progress_installs_a_callback(self, analysis): + # WHEN + bound_all(analysis) + + # THEN + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + analysis.bayesian.sample(samples=10, progress=True) + + # EXPECT + kwargs = sampler_class.return_value.sample.call_args.kwargs + assert callable(kwargs['progress_callback']) + + def test_progress_reports_and_finishes_the_line(self, analysis, capsys): + # WHEN + bound_all(analysis) + + # THEN the sampler drives the installed callback, as BUMPS does per generation + with patch(SAMPLER_PATH) as sampler_class: + + def run_reporting_progress(**kwargs): + for iteration in (1, 5, 10): + kwargs['progress_callback']({ + 'iteration': iteration, + 'total_steps': 10, + 'sampling': True, + }) + return fake_results(analysis) + + sampler_class.return_value.sample.side_effect = run_reporting_progress + results = analysis.bayesian.sample(samples=10, progress=True) + + # EXPECT progress was printed and the line was finished, without breaking the results + out = capsys.readouterr().out + assert '100%' in out + assert 'Sampling: done' in out + assert out.endswith('\n') + assert analysis.bayesian.results is results + + def test_progress_line_is_not_marked_done_when_sampling_fails(self, analysis, capsys): + # WHEN + bound_all(analysis) + + # THEN + with patch(SAMPLER_PATH) as sampler_class: + + def fail_after_progress(**kwargs): + kwargs['progress_callback']({'iteration': 1, 'total_steps': 10}) + raise RuntimeError('boom') + + sampler_class.return_value.sample.side_effect = fail_after_progress + with pytest.raises(RuntimeError, match='boom'): + analysis.bayesian.sample(samples=10, progress=True) + + # EXPECT the line is terminated but not decorated with a claim of success + out = capsys.readouterr().out + assert 'done' not in out + assert out.endswith('\n') + + def test_progress_defers_to_an_explicit_callback(self, analysis): + # WHEN the caller supplies their own callback alongside progress=True + bound_all(analysis) + explicit = MagicMock() + + # THEN + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + analysis.bayesian.sample(samples=10, progress=True, progress_callback=explicit) + + # EXPECT the explicit callback is forwarded untouched + kwargs = sampler_class.return_value.sample.call_args.kwargs + assert kwargs['progress_callback'] is explicit + + def test_extend_supports_progress(self, analysis): + # WHEN + bound_all(analysis) + + # THEN + 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.bayesian.sample(samples=10) + analysis.bayesian.extend(additional_samples=10, progress=True) + + # EXPECT + kwargs = sampler_class.return_value.extend.call_args.kwargs + assert callable(kwargs['progress_callback']) + + ############# + # Predictions + ############# + + def test_predictions_have_one_row_per_selected_draw(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, n_draws=100 + ) + analysis.bayesian.sample(samples=10) + + # THEN + predictions = analysis.bayesian.predictions(n_draws=10) + + # EXPECT + x, _, _ = analysis._sampling_data() + assert predictions.shape == (10, len(x)) + + def test_predictions_clamp_to_the_chain_length(self, analysis): + # WHEN more draws are requested than the chain holds + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results( + analysis, n_draws=100 + ) + analysis.bayesian.sample(samples=10) + + # THEN + predictions = analysis.bayesian.predictions(n_draws=500) + + # EXPECT one row per available draw, not 500 + x, _, _ = analysis._sampling_data() + assert predictions.shape == (100, len(x)) + + def test_predictions_take_draws_evenly_across_the_chain(self, analysis): + # WHEN the area column identifies each draw, since the model scales linearly with it + bound_all(analysis, half_width=500.0) + parameters = analysis.get_free_parameters() + column = [p.name for p in parameters].index('Gaussian area') + draws = np.tile([float(p.value) for p in parameters], (100, 1)) + draws[:, column] = 1.0 + np.arange(100.0) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_results(analysis, values=draws) + analysis.bayesian.sample(samples=10) + + # THEN + predictions = analysis.bayesian.predictions(n_draws=5) + + # EXPECT rows for draws 0, 24, 49, 74 and 99, read back through the model's linear + # scaling with the area + amplitudes = predictions.max(axis=1) + expected = draws[[0, 24, 49, 74, 99], column] + assert amplitudes / amplitudes[0] == pytest.approx(expected / expected[0]) + + +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/experiment/test_experiment.py b/tests/unit/easydynamics/experiment/test_experiment.py index 2329e29a2..1aebc0436 100644 --- a/tests/unit/easydynamics/experiment/test_experiment.py +++ b/tests/unit/easydynamics/experiment/test_experiment.py @@ -637,6 +637,19 @@ def testextract_x_y_weights_only_finite_zero_variance(self, experiment_with_data assert np.array_equal(weights, np.ones_like(y)) assert np.array_equal(mask, np.isfinite(y) & np.isfinite(x)) + def test_has_variances_true_when_the_data_carries_them(self, experiment_with_data): + # WHEN THEN EXPECT + assert experiment_with_data.has_variances + + def test_has_variances_false_when_the_data_has_none(self, experiment): + # WHEN THEN EXPECT the fixture's data has no variances, so the all-ones weights that + # extract_x_y_weights_only_finite falls back to are recognisable as placeholders + assert not experiment.has_variances + + def test_has_variances_false_without_data(self): + # WHEN THEN EXPECT + assert not Experiment().has_variances + ############## # test dunder methods ############## 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..c470bdd20 --- /dev/null +++ b/tests/unit/easydynamics/utils/test_posterior_plotting.py @@ -0,0 +1,352 @@ +# 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_correlations +from easydynamics.utils.posterior_plotting import plot_marginal +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): + # THEN + fig = plot_trace(draws=draws, names=['a', 'b', 'c']) + + # EXPECT + assert len(fig.axes) == 3 + + def test_logp_adds_a_panel(self, draws): + # THEN + 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): + # THEN + 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): + # THEN + fig = plot_trace(draws=np.zeros((10, 1)), names=['only']) + + # EXPECT + assert len(fig.axes) == 1 + + def test_mismatched_names_raise(self, draws): + # THEN EXPECT + with pytest.raises(ValueError, match='one entry per column'): + plot_trace(draws=draws, names=['a', 'b']) + + def test_one_dimensional_draws_raise(self): + # THEN EXPECT + with pytest.raises(ValueError, match='two-dimensional'): + plot_trace(draws=np.zeros(10), names=['a']) + + def test_zero_row_draws_raise(self): + # THEN EXPECT + with pytest.raises(ValueError, match='no samples'): + plot_trace(draws=np.zeros((0, 2)), names=['a', 'b']) + + def test_zero_column_draws_raise(self): + # THEN EXPECT + with pytest.raises(ValueError, match='no parameters'): + plot_trace(draws=np.zeros((5, 0)), names=[]) + + def test_a_single_draw_keeps_a_usable_axis(self): + # THEN + fig = plot_trace(draws=np.ones((1, 2)), names=['a', 'b']) + + # EXPECT a non-inverted, non-degenerate x range + left, right = fig.axes[0].get_xlim() + assert left < right + + def test_mismatched_logp_length_raises(self, draws): + # THEN EXPECT + with pytest.raises(ValueError, match='one entry per draw'): + plot_trace(draws=draws, names=['a', 'b', 'c'], logp=np.zeros(len(draws) - 1)) + + +class TestPlotCorner: + def test_grid_is_square_in_the_parameter_count(self, draws): + # THEN + fig = plot_corner(draws=draws, names=['a', 'b', 'c']) + + # EXPECT + assert len(fig.axes) == 9 + + def test_upper_triangle_is_hidden(self, draws): + # THEN + 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): + # THEN EXPECT + with pytest.raises(ValueError, match='one entry per column'): + plot_corner(draws=draws, names=['a']) + + def test_non_finite_draws_raise_naming_the_column(self, draws): + # WHEN one column contains a NaN + draws[5, 1] = np.nan + + # THEN EXPECT a clear error naming that column, not numpy's "range [nan, nan]" + with pytest.raises(ValueError, match='non-finite') as excinfo: + plot_corner(draws=draws, names=['a', 'b', 'c']) + assert ': b.' in str(excinfo.value) + + def test_columns_share_limits_between_histogram_and_hexbin_panels(self, draws): + # THEN + fig = plot_corner(draws=draws, names=['a', 'b', 'c']) + + # EXPECT every panel of a column agrees with the diagonal histogram on x-limits, so the + # ticks line up down the column + grid = np.array(fig.axes, dtype=object).reshape(3, 3) + for col in range(3): + column_limits = [grid[row, col].get_xlim() for row in range(col, 3)] + assert all(limits == pytest.approx(column_limits[0]) for limits in column_limits) + + +class TestPlotMarginal: + @pytest.fixture + def values(self): + return np.random.default_rng(0).normal(size=5000) + + def test_returns_a_single_axis_figure(self, values): + # THEN + fig = plot_marginal(values=values, name='width') + + # EXPECT + assert len(fig.axes) == 1 + + def test_marks_the_median_and_the_credible_interval(self, values): + # THEN + fig = plot_marginal(values=values, name='width') + + # EXPECT three vertical lines at the 16th, 50th and 84th percentiles + positions = sorted(line.get_xdata()[0] for line in fig.axes[0].lines) + assert positions == pytest.approx(np.percentile(values, [16.0, 50.0, 84.0])) + + def test_legend_names_the_median_and_the_interval(self, values): + # THEN + fig = plot_marginal(values=values, name='width') + + # EXPECT + labels = [text.get_text() for text in fig.axes[0].get_legend().get_texts()] + assert 'Median' in labels + assert any('credible interval' in label for label in labels) + + def test_histogram_is_density_normalized(self, values): + # WHEN values are standard-normal draws, whose density peaks near 0.4 + + # THEN + fig = plot_marginal(values=values, name='width') + + # EXPECT the peak reads as a probability density, not a raw count of thousands + peak = fig.axes[0].dataLim.ymax + assert 0.2 < peak < 0.7 + + def test_unit_is_appended_to_the_label(self, values): + # THEN + fig = plot_marginal(values=values, name='width', unit='meV') + + # EXPECT + assert fig.axes[0].get_xlabel() == 'width (meV)' + + def test_dimensionless_unit_is_skipped(self, values): + # THEN + fig = plot_marginal(values=values, name='area', unit='dimensionless') + + # EXPECT + assert fig.axes[0].get_xlabel() == 'area' + + def test_two_dimensional_values_raise(self): + # THEN EXPECT + with pytest.raises(ValueError, match='one-dimensional'): + plot_marginal(values=np.zeros((10, 2)), name='width') + + def test_empty_values_raise(self): + # THEN EXPECT + with pytest.raises(ValueError, match='no samples'): + plot_marginal(values=np.zeros(0), name='width') + + def test_non_finite_values_raise_naming_the_parameter(self, values): + # WHEN + values[3] = np.nan + + # THEN EXPECT a clear error, not numpy's "range [nan, nan]" + with pytest.raises(ValueError, match='non-finite') as excinfo: + plot_marginal(values=values, name='width') + assert 'width' in str(excinfo.value) + + +class TestPlotCorrelations: + def test_labels_both_axes_with_the_names(self, draws): + # THEN + fig = plot_correlations(draws=draws, names=['a', 'b', 'c']) + + # EXPECT + axis = fig.axes[0] + assert [text.get_text() for text in axis.get_xticklabels()] == ['a', 'b', 'c'] + assert [text.get_text() for text in axis.get_yticklabels()] == ['a', 'b', 'c'] + + def test_diagonal_is_one(self, draws): + # THEN + fig = plot_correlations(draws=draws, names=['a', 'b', 'c']) + + # EXPECT + matrix = fig.axes[0].images[0].get_array() + assert np.asarray(np.diag(matrix)) == pytest.approx(np.ones(3)) + + def test_every_cell_is_annotated(self, draws): + # THEN + fig = plot_correlations(draws=draws, names=['a', 'b', 'c']) + + # EXPECT + assert len(fig.axes[0].texts) == 9 + + def test_color_limits_span_the_full_correlation_range(self, draws): + # THEN + fig = plot_correlations(draws=draws, names=['a', 'b', 'c']) + + # EXPECT the diverging map is centred on 0 regardless of the data + assert fig.axes[0].images[0].get_clim() == (-1.0, 1.0) + + def test_has_a_colorbar(self, draws): + # THEN + fig = plot_correlations(draws=draws, names=['a', 'b', 'c']) + + # EXPECT + assert len(fig.axes) == 2 + + def test_correlated_columns_read_near_one(self): + # WHEN two columns are almost the same draw + rng = np.random.default_rng(0) + base = rng.normal(size=500) + draws = np.column_stack([base, base + rng.normal(scale=1e-6, size=500)]) + + # THEN + fig = plot_correlations(draws=draws, names=['a', 'b']) + + # EXPECT + matrix = fig.axes[0].images[0].get_array() + assert matrix[0, 1] == pytest.approx(1.0, abs=1e-6) + + def test_single_parameter_chain_works(self): + # THEN + fig = plot_correlations(draws=np.random.default_rng(0).normal(size=(50, 1)), names=['a']) + + # EXPECT a 1x1 matrix whose only entry is 1 + matrix = fig.axes[0].images[0].get_array() + assert matrix.shape == (1, 1) + assert matrix[0, 0] == pytest.approx(1.0) + + def test_constant_column_is_masked_without_warnings(self, draws): + # WHEN one column has zero variance, so its correlations are undefined + import warnings + + draws[:, 1] = 2.5 + + # THEN numpy's zero-variance warnings are suppressed rather than leaking out + with warnings.catch_warnings(): + warnings.simplefilter('error') + fig = plot_correlations(draws=draws, names=['a', 'b', 'c']) + + # EXPECT the undefined cells are masked and annotated as unavailable + matrix = fig.axes[0].images[0].get_array() + assert matrix.mask[0, 1] + assert any(text.get_text() == 'n/a' for text in fig.axes[0].texts) + + def test_mismatched_names_raise(self, draws): + # THEN EXPECT + with pytest.raises(ValueError, match='one entry per column'): + plot_correlations(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)) + + # THEN + 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) + + # THEN + 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): + # THEN 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): + # THEN 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)) + + # THEN + 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 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())