From a8946cea84f192f79365fc5c398e12bcfc03d058 Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Thu, 13 Aug 2026 10:18:29 +0200 Subject: [PATCH 01/29] Add Bayesian posterior sampling to Analysis1d Expose the EasyScience Fitter on Analysis1d and add MCMC posterior sampling on top of it, using the BUMPS DREAM sampler introduced in easyscience 2.5.1 (easyscience.fitting.Sampler). Least-squares fitting reports a single point with a curvature-derived uncertainty, which is only trustworthy when parameters are uncorrelated and roughly Gaussian. Sampling maps the whole posterior instead, so correlated and skewed parameters get honest credible intervals. The sampling machinery lives in a mixin with three hooks (build the fitter, bind the data, list the chain parameters) so that Analysis and ParameterAnalysis can reuse it. ParameterAnalysis is not an AnalysisBase and builds a MultiFitter over binding models rather than over itself, so a shared base class would not have worked. Notable details: - fit() now uses a cached Fitter instead of building one per call, and the cache is invalidated through the existing dirty-flag pattern. - Bounds are the prior in DREAM, so sampling refuses to run with any infinite bound. suggest_bounds() proposes finite ones from the fitted values and uncertainties; it is advisory until .apply() is called and never loosens a bound that is already finite, so physical limits survive. A zero-width suggestion is flagged rather than invented. - Sampling restores parameter values afterwards, since BUMPS leaves them wherever the last likelihood evaluation put them. - Chains are reported under Parameter.name, not the internal unique_name. Those names are per-session, so save_chain() writes a sidecar mapping them to stable names and load_chain() uses it; loading without one warns rather than mislabelling the columns. - After sampling, a warning fires when the posterior has piled up against a bound, which catches both bounds that are too tight and degenerate parameters that drift until a bound stops them. - BUMPS crashes with a bare IndexError inside its own outlier removal when chains scatter, which in practice means a degenerate model. That is re-raised with the likely cause and a workaround. Co-Authored-By: Claude Opus 5 (1M context) --- docs/docs/tutorials/bayesian.ipynb | 306 ++++++ docs/docs/tutorials/index.md | 3 + docs/mkdocs.yml | 1 + pixi.lock | 3 +- pyproject.toml | 23 +- src/easydynamics/analysis/__init__.py | 10 + src/easydynamics/analysis/analysis1d.py | 80 +- .../analysis/bayesian_sampling.py | 993 ++++++++++++++++++ src/easydynamics/analysis/posterior.py | 608 +++++++++++ src/easydynamics/utils/__init__.py | 11 +- src/easydynamics/utils/posterior_plotting.py | 256 +++++ .../fitting/test_bayesian_sampling.py | 207 ++++ .../analysis/test_analysis1d_bayesian.py | 486 +++++++++ .../easydynamics/analysis/test_posterior.py | 296 ++++++ .../utils/test_posterior_plotting.py | 148 +++ 15 files changed, 3407 insertions(+), 24 deletions(-) create mode 100644 docs/docs/tutorials/bayesian.ipynb create mode 100644 src/easydynamics/analysis/bayesian_sampling.py create mode 100644 src/easydynamics/analysis/posterior.py create mode 100644 src/easydynamics/utils/posterior_plotting.py create mode 100644 tests/integration/fitting/test_bayesian_sampling.py create mode 100644 tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py create mode 100644 tests/unit/easydynamics/analysis/test_posterior.py create mode 100644 tests/unit/easydynamics/utils/test_posterior_plotting.py diff --git a/docs/docs/tutorials/bayesian.ipynb b/docs/docs/tutorials/bayesian.ipynb new file mode 100644 index 000000000..44a415147 --- /dev/null +++ b/docs/docs/tutorials/bayesian.ipynb @@ -0,0 +1,306 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "eac0b8bb", + "metadata": {}, + "source": [ + "# Bayesian analysis\n", + "\n", + "Fitting with `fit()` finds the single set of parameter values that best matches the data, and reports an uncertainty derived from the curvature of $\\chi^2$ at that point. That uncertainty is only trustworthy when the parameters are uncorrelated and their uncertainties are close to Gaussian, which in QENS is often not the case.\n", + "\n", + "A **Bayesian** analysis answers a different question: instead of one best point, it maps out the whole *posterior distribution* over the parameters. From that you can read off credible intervals that stay honest when parameters are correlated or their distributions are skewed, and you can see the correlations directly.\n", + "\n", + "EasyDynamics does this with the DREAM sampler from [BUMPS](https://bumps.readthedocs.io/), through `sample_posterior()`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "02cb7aec", + "metadata": {}, + "outputs": [], + "source": [ + "import pooch\n", + "\n", + "import easydynamics as edyn\n", + "import easydynamics.sample_model as sm\n", + "from easydynamics.analysis.analysis1d import Analysis1d\n", + "\n", + "%matplotlib inline" + ] + }, + { + "cell_type": "markdown", + "id": "0499fea7", + "metadata": {}, + "source": [ + "## Load the data\n", + "\n", + "We use the same artificial vanadium measurement as the [Analysis 1D](analysis1d.ipynb) tutorial, and analyse a single Q slice." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fb407621", + "metadata": {}, + "outputs": [], + "source": [ + "vanadium_experiment = edyn.Experiment('Vanadium')\n", + "\n", + "file_path = pooch.retrieve(\n", + " url='https://github.com/easyscience/dynamics-lib/raw/refs/heads/master/docs/docs/tutorials/data/vanadium_data_example.h5',\n", + " known_hash='16cc1b327c303feeb88fb9dda5390dc4880b62396b1793f98c6fef0b27c7b873',\n", + ")\n", + "\n", + "vanadium_experiment.load_hdf5(filename=file_path)" + ] + }, + { + "cell_type": "markdown", + "id": "fcdfd395", + "metadata": {}, + "source": [ + "## Build the model and fit it\n", + "\n", + "As in [Tutorial 1](tutorial1_brownian.ipynb), a vanadium measurement is modelled with the Gaussian as the *sample*: what is being measured is the resolution function itself, so there is nothing to convolve it with.\n", + "\n", + "Sampling does not require a fit first, but it benefits from one: DREAM starts its chains in a small ball around the parameters' current values, so beginning from fitted values means less burn-in is needed before the chains reach the interesting region." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "de3297cf", + "metadata": {}, + "outputs": [], + "source": [ + "vanadium_components = sm.ComponentCollection()\n", + "vanadium_components.append_component(sm.Gaussian(width=0.1, area=1, name='Res. Gauss'))\n", + "\n", + "instrument_model = sm.InstrumentModel(\n", + " background_model=sm.BackgroundModel(components=sm.Polynomial(coefficients=[0.001])),\n", + ")\n", + "\n", + "analysis = Analysis1d(\n", + " display_name='Vanadium Analysis',\n", + " experiment=vanadium_experiment,\n", + " sample_model=sm.SampleModel(components=vanadium_components),\n", + " instrument_model=instrument_model,\n", + " Q_index=5,\n", + ")\n", + "\n", + "fit_result = analysis.fit()\n", + "print(f'reduced chi-squared = {fit_result.reduced_chi2:.4f}')" + ] + }, + { + "cell_type": "markdown", + "id": "b0a709f7", + "metadata": {}, + "source": [ + "## Bounds are the prior\n", + "\n", + "In DREAM, each parameter's `min` and `max` define a uniform prior, so **every free parameter must have finite bounds** before sampling. Most parameters start with at least one infinite bound, so `sample_posterior()` would refuse to run.\n", + "\n", + "`suggest_bounds()` proposes bounds from the fitted values and uncertainties. It is advisory: it changes nothing until you call `.apply()`, and it only ever fills in an *infinite* bound, so physical limits you have already set (an area that cannot go below zero, say) are left alone." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a004cd83", + "metadata": {}, + "outputs": [], + "source": [ + "suggestions = analysis.suggest_bounds()\n", + "print(suggestions)" + ] + }, + { + "cell_type": "markdown", + "id": "ab263f41", + "metadata": {}, + "source": [ + "The defaults are deliberately generous — 10 standard deviations plus 20% of the value. Because the bounds are a uniform prior, being too *narrow* is the dangerous mistake: it truncates the posterior and makes the uncertainty look smaller than it is. The 20% term is there for parameters whose fitted uncertainty comes back as zero. All three settings (`n_sigma`, `relative_pad`, `absolute_floor`) can be adjusted, and you can always set `min` and `max` by hand.\n", + "\n", + "It is worth reading the table before applying it. A suggestion many orders of magnitude larger than the parameter itself is a useful warning sign: it means the fit returned a huge uncertainty, which usually happens because two parameters are **degenerate** — the data determines only some combination of them, so one can grow while the other shrinks with no effect on the fit. That is a problem to fix in the model, not with the sampler." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1a92a899", + "metadata": {}, + "outputs": [], + "source": [ + "changed = suggestions.apply()\n", + "print(f'Applied bounds to: {[parameter.name for parameter in changed]}')" + ] + }, + { + "cell_type": "markdown", + "id": "8cf71e53", + "metadata": {}, + "source": [ + "## Sample the posterior\n", + "\n", + "`sample_posterior()` runs the chains. The three numbers that matter are:\n", + "\n", + "- `samples` — how many draws to collect in total. More is better, at linear cost.\n", + "- `burn` — generations discarded at the start, while the chains are still travelling towards the bulk of the posterior.\n", + "- `thin` — keep only every n-th generation, which reduces the correlation between neighbouring draws.\n", + "\n", + "Sampling never moves your parameters: their values are restored afterwards, so the model is left exactly as the fit left it." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "861fa264", + "metadata": {}, + "outputs": [], + "source": [ + "results = analysis.sample_posterior(samples=4000, burn=300, thin=2)\n", + "\n", + "print(f'Collected {results.draws.shape[0]} draws for {results.draws.shape[1]} parameters.')" + ] + }, + { + "cell_type": "markdown", + "id": "f2ce5232", + "metadata": {}, + "source": [ + "## Did the chains converge?\n", + "\n", + "Always look at the traces before trusting the numbers. A converged chain looks like a \"hairy caterpillar\": noisy, but flat and stationary. A visible drift or slow wander means the chain has not settled and needs a longer burn-in or more samples." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b2fe638c", + "metadata": {}, + "outputs": [], + "source": [ + "analysis.plot_trace()" + ] + }, + { + "cell_type": "markdown", + "id": "2b3ea7b4", + "metadata": {}, + "source": [ + "## Summarize the posterior\n", + "\n", + "`posterior_summary()` reports the median and the 68% credible interval of each parameter, under the parameter's own name and unit. The interval is asymmetric in general, which is precisely the information a single symmetric error bar throws away." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8470375e", + "metadata": {}, + "outputs": [], + "source": [ + "analysis.posterior_summary()" + ] + }, + { + "cell_type": "markdown", + "id": "60b4a448", + "metadata": {}, + "source": [ + "## Correlations between parameters\n", + "\n", + "The corner plot is the part least available from a least-squares fit. The diagonal shows each parameter's own distribution; each off-diagonal panel shows a pair. A round blob means the two are independent, while a tilted, narrow ridge means they are correlated and the data constrains only a combination of them." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c96509cc", + "metadata": {}, + "outputs": [], + "source": [ + "analysis.plot_corner()" + ] + }, + { + "cell_type": "markdown", + "id": "8f819c75", + "metadata": {}, + "source": [ + "## Does the model actually describe the data?\n", + "\n", + "The posterior predictive plot re-evaluates the model for a sample of posterior draws and shades the region they cover. If the data wanders outside the band in a systematic way, the model is missing a feature, and no amount of parameter tuning will fix it." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cce95199", + "metadata": {}, + "outputs": [], + "source": [ + "analysis.plot_posterior_predictive(n_draws=100)" + ] + }, + { + "cell_type": "markdown", + "id": "5947da0f", + "metadata": {}, + "source": [ + "## Continuing and storing a chain\n", + "\n", + "If the traces suggest the chain needs to run longer, `extend_sampling()` continues the existing chain rather than starting over, so nothing already computed is thrown away." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "974e486b", + "metadata": {}, + "outputs": [], + "source": [ + "extended = analysis.extend_sampling(additional_samples=1000, thin=2)\n", + "print(f'Chain now holds {extended.draws.shape[0]} draws.')" + ] + }, + { + "cell_type": "markdown", + "id": "69793d31", + "metadata": {}, + "source": [ + "Chains are expensive, so they can be saved and reloaded with `analysis.save_chain(path)` and `analysis.load_chain(path)`. A reloaded chain can be summarized, plotted, or extended further, exactly like a fresh one." + ] + }, + { + "cell_type": "markdown", + "id": "a3448aee", + "metadata": {}, + "source": [ + "## Things to watch out for\n", + "\n", + "**Data without uncertainties.** If your data carries no variances, the weights fall back to 1, which means the sampler assumes a noise level of 1 in whatever units the intensity happens to be. Least-squares does not care, since that scale cancels out of the best-fit position, but a posterior *does*: its width scales directly with the assumed noise, so the credible intervals will be wrong by whatever factor the true noise differs from 1. Bayesian analysis is not a way to avoid needing uncertainties on your data.\n", + "\n", + "**Sampling only some parameters.** `sample_posterior(parameters=[...])` restricts the chain to a subset, which is faster because the number of chains scales with the number of parameters. Be careful with the result: the other parameters are held *fixed*, which is not the same as averaging over them. The intervals you get are conditional on those fixed values, and will be too narrow whenever the parameters are correlated.\n", + "\n", + "**Degenerate parameters.** As seen above, they are a modelling problem rather than a sampling one. `suggest_bounds()` returning absurd values, or a warning that the posterior has piled up against its bounds, are both signs to go back and look at the model." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "default", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/docs/tutorials/index.md b/docs/docs/tutorials/index.md index 0617edb2a..23a36adc5 100644 --- a/docs/docs/tutorials/index.md +++ b/docs/docs/tutorials/index.md @@ -62,3 +62,6 @@ tutorials. - [Analysis](analysis.ipynb) - Learn how to fit a model to your data. - [Analysis 1D](analysis1d.ipynb) - Learn how to fit a model to your data at a particular Q. +- [Bayesian analysis](bayesian.ipynb) - Learn how to map out the full + posterior distribution of your parameters, including their + correlations, instead of a single best-fit point. diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index 64e94f967..3db04ace0 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -217,6 +217,7 @@ nav: - Experiment: tutorials/experiment.ipynb - Analysis: tutorials/analysis.ipynb - Analysis 1D: tutorials/analysis1d.ipynb + - Bayesian analysis: tutorials/bayesian.ipynb - API Reference: - API Reference: api-reference/index.md - analysis: api-reference/analysis.md diff --git a/pixi.lock b/pixi.lock index d1ebb0ed7..f07f02cf8 100644 --- a/pixi.lock +++ b/pixi.lock @@ -7535,12 +7535,13 @@ packages: name: easydynamics requires_dist: - darkdetect - - easyscience + - easyscience>=2.5.1 - ipykernel - ipympl - ipython - ipywidgets - jupyterlab + - matplotlib - pixi-kernel - plopp - pooch diff --git a/pyproject.toml b/pyproject.toml index b39261f87..3330304ea 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,17 +23,18 @@ classifiers = [ ] requires-python = '>=3.12' dependencies = [ - 'easyscience', # The base library of the EasyScience framework - 'pooch', # Data downloader - 'darkdetect', # Detecting dark mode (system-level) - 'plopp', # Plotting library - 'jupyterlab', # Jupyter notebooks - 'pixi-kernel', # Pixi Jupyter kernel - 'ipykernel', # Jupyter kernel (required for running notebooks) - 'ipywidgets', # Widgets (needed for interactive matplotlib backends) - 'ipympl', # Matplotlib Jupyter widget backend (%matplotlib widget) - 'IPython', # Interactive Python shell - 'sympy', # Symbolic mathematics (used for expression components) + 'easyscience>=2.5.1', # The base library of the EasyScience framework. 2.5.1 adds fitting.Sampler + 'matplotlib', # Plotting (posterior trace, corner, and predictive plots) + 'pooch', # Data downloader + 'darkdetect', # Detecting dark mode (system-level) + 'plopp', # Plotting library + 'jupyterlab', # Jupyter notebooks + 'pixi-kernel', # Pixi Jupyter kernel + 'ipykernel', # Jupyter kernel (required for running notebooks) + 'ipywidgets', # Widgets (needed for interactive matplotlib backends) + 'ipympl', # Matplotlib Jupyter widget backend (%matplotlib widget) + 'IPython', # Interactive Python shell + 'sympy', # Symbolic mathematics (used for expression components) ] [project.optional-dependencies] diff --git a/src/easydynamics/analysis/__init__.py b/src/easydynamics/analysis/__init__.py index 289ec02f5..89a45cea3 100644 --- a/src/easydynamics/analysis/__init__.py +++ b/src/easydynamics/analysis/__init__.py @@ -2,9 +2,19 @@ # SPDX-License-Identifier: BSD-3-Clause from easydynamics.analysis.analysis import Analysis +from easydynamics.analysis.bayesian_sampling import BayesianSamplingMixin from easydynamics.analysis.parameter_analysis import ParameterAnalysis +from easydynamics.analysis.posterior import BoundsSuggestion +from easydynamics.analysis.posterior import BoundsSuggestions +from easydynamics.analysis.posterior import ParameterPosterior +from easydynamics.analysis.posterior import PosteriorSummary __all__ = [ 'Analysis', + 'BayesianSamplingMixin', + 'BoundsSuggestion', + 'BoundsSuggestions', 'ParameterAnalysis', + 'ParameterPosterior', + 'PosteriorSummary', ] diff --git a/src/easydynamics/analysis/analysis1d.py b/src/easydynamics/analysis/analysis1d.py index e50b09737..eaec732b1 100644 --- a/src/easydynamics/analysis/analysis1d.py +++ b/src/easydynamics/analysis/analysis1d.py @@ -12,6 +12,7 @@ from plopp.backends.matplotlib.figure import InteractiveFigure from easydynamics.analysis.analysis_base import AnalysisBase +from easydynamics.analysis.bayesian_sampling import BayesianSamplingMixin from easydynamics.convolution.convolution import Convolution from easydynamics.experiment import Experiment from easydynamics.sample_model import InstrumentModel @@ -25,12 +26,16 @@ from easydynamics.utils.utils import verify_Q_index -class Analysis1d(AnalysisBase): +class Analysis1d(BayesianSamplingMixin, AnalysisBase): """ For analysing one-dimensional data, i.e. intensity as function of energy for a single Q index. Is used primarily in the Analysis class, but can also be used on its own for simpler analyses. + In addition to least-squares fitting with :meth:`fit`, the posterior distribution of the free + parameters can be explored with :meth:`sample_posterior`; see + :class:`~easydynamics.analysis.bayesian_sampling.BayesianSamplingMixin`. + Examples -------- **Fitting a single Q slice** @@ -116,6 +121,7 @@ def __init__( self._fit_result = None self._convolver = None self._convolver_is_dirty = True + self._init_bayesian_state() super().__init__( display_name=display_name, @@ -245,27 +251,72 @@ def fit(self) -> FitResults: if self._experiment is None: raise ValueError('No experiment is associated with this Analysis.') - if ( - self.sample_model.component_collections_is_dirty - or self.instrument_model.resolution_model.component_collections_is_dirty - ): - self._convolver_is_dirty = True + self._prepare_for_sampling() - self._ensure_convolver_current() + x, y, weights = self._get_sampling_data() + fit_result = self.fitter.fit(x=x, y=y, weights=weights) + + self._fit_result = fit_result + + return fit_result - fitter = EasyScienceFitter( + ############# + # Hooks for BayesianSamplingMixin + ############# + + def _build_bayesian_fitter(self) -> EasyScienceFitter: + """ + Build the EasyScience Fitter for this Analysis. + + Returns + ------- + EasyScienceFitter + A Fitter bound to this Analysis and its fit function. + """ + return EasyScienceFitter( fit_object=self, fit_function=self.as_fit_function(), ) + def _get_sampling_data(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Get the finite data for the chosen Q index, as used by both fitting and sampling. + + Returns + ------- + tuple[np.ndarray, np.ndarray, np.ndarray] + The ``(x, y, weights)`` triple. + """ x, y, weights, _ = self.experiment.extract_x_y_weights_only_finite( Q_index=self._require_Q_index() ) - fit_result = fitter.fit(x=x, y=y, weights=weights) + return x, y, weights - self._fit_result = fit_result + def _get_chain_parameters(self) -> list[Parameter]: + """ + Get the free parameters of this Analysis. - return fit_result + Returns + ------- + list[Parameter] + The parameters that are free to vary, which are the ones the sampler explores. + """ + return self.get_free_parameters() + + def _prepare_for_sampling(self) -> None: + """ + Rebuild the convolver if anything it depends on has changed. + + The energy grid is fixed for the duration of a fit or a sampling run, so the convolution + objects are built once here and reused for every model evaluation. + """ + if ( + self.sample_model.component_collections_is_dirty + or self.instrument_model.resolution_model.component_collections_is_dirty + ): + self._convolver_is_dirty = True + + self._ensure_convolver_current() def as_fit_function( self, @@ -483,6 +534,7 @@ def rebin(self, dimensions: dict[str, int | sc.Variable]) -> None: if self._Q_index is not None and self.experiment is not None: self._masked_energy = self.experiment.get_masked_energy(Q_index=self._Q_index) self._convolver_is_dirty = True + self._invalidate_bayesian_sampler() def refresh_convolver(self, energy: sc.Variable | None = None) -> None: """Refresh the pre-built Convolution object for the current Q index.""" @@ -523,10 +575,13 @@ def _on_Q_index_changed(self) -> None: if self._Q_index is None: self._masked_energy = None self._convolver_is_dirty = True + self._invalidate_bayesian_sampler() return masked_energy = self.experiment.get_masked_energy(Q_index=self._Q_index) self._masked_energy = masked_energy self._convolver_is_dirty = True + # A different Q index means different data, and the Sampler binds its data at construction. + self._invalidate_bayesian_sampler() def _on_experiment_changed(self) -> None: """Mark the convolver as dirty when the experiment changes.""" @@ -535,16 +590,19 @@ def _on_experiment_changed(self) -> None: if self._Q_index is not None and self.experiment is not None: self._masked_energy = self.experiment.get_masked_energy(Q_index=self._Q_index) self._convolver_is_dirty = True + self._invalidate_bayesian_sampler() def _on_sample_model_changed(self) -> None: """Mark the convolver as dirty when the sample model changes.""" super()._on_sample_model_changed() self._convolver_is_dirty = True + self._invalidate_fitter() def _on_instrument_model_changed(self) -> None: """Mark the convolver as dirty when the instrument model changes.""" super()._on_instrument_model_changed() self._convolver_is_dirty = True + self._invalidate_fitter() def _on_convolution_settings_changed(self) -> None: """Mark the convolver as dirty when the convolution settings change.""" diff --git a/src/easydynamics/analysis/bayesian_sampling.py b/src/easydynamics/analysis/bayesian_sampling.py new file mode 100644 index 000000000..0ea38d8ef --- /dev/null +++ b/src/easydynamics/analysis/bayesian_sampling.py @@ -0,0 +1,993 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +""" +Shared Bayesian MCMC sampling machinery for the Analysis classes. + +Everything that does not depend on how a particular Analysis is wired up lives here: caching the +Fitter and the Sampler, guarding the parameter bounds, restoring parameter values afterwards, and +turning raw draws into a readable summary. A concrete Analysis supplies the three things that do +differ, via :meth:`BayesianSamplingMixin._build_bayesian_fitter`, +:meth:`BayesianSamplingMixin._get_sampling_data`, and +:meth:`BayesianSamplingMixin._get_chain_parameters`. +""" + +from __future__ import annotations + +import json +import warnings +from pathlib import Path +from typing import TYPE_CHECKING +from typing import Any + +import numpy as np +from easyscience.fitting import AvailableMinimizers +from easyscience.fitting import Sampler + +from easydynamics.analysis.posterior import PosteriorSummary +from easydynamics.analysis.posterior import parameters_at_bounds +from easydynamics.analysis.posterior import suggest_bounds_for_parameters +from easydynamics.analysis.posterior import summarize_draws +from easydynamics.analysis.posterior import unbounded_parameters + +if TYPE_CHECKING: + import os + from collections.abc import Callable + + from easyscience.fitting.fitter import Fitter + from easyscience.fitting.sampler import SamplingResults + from easyscience.variable import Parameter + from matplotlib.figure import Figure + + from easydynamics.analysis.posterior import BoundsSuggestions + +# Suffix of the sidecar mapping chain columns to stable parameter names, written next to the BUMPS +# chain files by save_chain(). +_NAME_MAP_SUFFIX = '.parameter-names.json' + + +class BayesianSamplingMixin: + """ + Bayesian MCMC sampling on top of an Analysis, backed by the BUMPS DREAM sampler. + + Sampling explores the full posterior distribution of the free parameters rather than reporting + a single best-fit point, which is worth doing when parameters are correlated or their + uncertainties are strongly non-Gaussian -- both common in QENS. + + Running :meth:`fit` first is not required, but it helps: DREAM seeds its population in a small + ball around the parameters' current values, so starting from fitted values shortens the burn-in + needed to reach the typical set. + + Notes + ----- + All free parameters must have finite bounds before sampling, because in DREAM the bounds are + the prior. :meth:`suggest_bounds` proposes bounds for any parameter still missing one. + + Examples + -------- + ```python + analysis.fit() + analysis.suggest_bounds().apply() + results = analysis.sample_posterior(samples=10000, burn=2000, thin=10) + analysis.posterior_summary() + ``` + """ + + ############# + # Setup + ############# + + def _init_bayesian_state(self) -> None: + """ + Initialize the cached sampling state. + + Must be called by the concrete Analysis before any observer callback can fire, in the same + way as the other cached objects on the class. + """ + self._fitter = None + self._fitter_is_dirty = True + self._bayesian_sampler = None + self._bayesian_sampler_is_dirty = True + self._posterior_result = None + # Maps a chain column's unique_name to the parameter name it had when saved. Only populated + # by load_chain, because unique_names are per-session and do not survive a round trip. + self._chain_name_map = {} + + def _invalidate_fitter(self) -> None: + """ + Mark the cached Fitter and Sampler as needing a rebuild. + + The Sampler binds its data at construction, so anything that invalidates the Fitter + invalidates the Sampler too. + """ + self._fitter_is_dirty = True + self._bayesian_sampler_is_dirty = True + + def _invalidate_bayesian_sampler(self) -> None: + """ + Mark only the cached Sampler as needing a rebuild. + + Used when the data changed but the model did not. + """ + self._bayesian_sampler_is_dirty = True + + ############# + # Hooks for concrete Analysis classes + ############# + + def _build_bayesian_fitter(self) -> Fitter: + """ + Build the EasyScience Fitter (or MultiFitter) for this Analysis. + + Returns + ------- + Fitter + A configured Fitter or MultiFitter. + + Raises + ------ + NotImplementedError + If the concrete Analysis does not implement it. + """ + raise NotImplementedError('Subclasses must implement _build_bayesian_fitter.') + + def _get_sampling_data(self) -> tuple: + """ + Get the ``(x, y, weights)`` to bind to the Sampler. + + Each element is either an array (single dataset) or a list of arrays (MultiFitter). + + Returns + ------- + tuple + The ``(x, y, weights)`` triple. + + Raises + ------ + NotImplementedError + If the concrete Analysis does not implement it. + """ + raise NotImplementedError('Subclasses must implement _get_sampling_data.') + + def _get_chain_parameters(self) -> list[Parameter]: + """ + Get the free parameters that will appear as columns of the chain. + + Returns + ------- + list[Parameter] + The free parameters of the underlying model(s). + + Raises + ------ + NotImplementedError + If the concrete Analysis does not implement it. + """ + raise NotImplementedError('Subclasses must implement _get_chain_parameters.') + + def _prepare_for_sampling(self) -> None: + """ + Bring any cached computation up to date before a sampling run. + + The default does nothing; Analysis classes that cache a convolver override it. + """ + + ############# + # Properties + ############# + + @property + def fitter(self) -> Fitter: + """ + The EasyScience Fitter used for fitting and sampling, built on first use. + + Exposed so the minimizer, tolerance, and maximum evaluation count can be configured + directly, e.g. ``analysis.fitter.switch_minimizer(AvailableMinimizers.Bumps)``. + + Returns + ------- + Fitter + The cached Fitter or MultiFitter. + """ + if self._fitter_is_dirty or self._fitter is None: + self._fitter = self._build_bayesian_fitter() + self._fitter_is_dirty = False + return self._fitter + + @property + def bayesian_sampler(self) -> Sampler | None: + """ + The EasyScience Sampler holding the MCMC chain, or None before the first run. + + Named to avoid confusion with the SampleModel: this samples the posterior, not the sample. + + Returns + ------- + Sampler | None + The cached Sampler, or None if no chain has been started. + """ + return self._bayesian_sampler + + @property + def posterior_result(self) -> SamplingResults | None: + """ + The results of the most recent sampling run, or None if there has not been one. + + Returns + ------- + SamplingResults | None + The most recent sampling results. + """ + return self._posterior_result + + ############# + # Bounds + ############# + + def suggest_bounds( + self, + n_sigma: float = 10.0, + relative_pad: float = 0.2, + absolute_floor: float | None = None, + ) -> BoundsSuggestions: + """ + Propose finite bounds for free parameters that still have an infinite one. + + Nothing is changed until :meth:`BoundsSuggestions.apply` is called, so the proposal can be + reviewed first. Bounds that are already finite are never widened or narrowed, so physical + limits such as a non-negative area are left alone. + + Because the bounds act as a uniform prior in DREAM, a generous width is the safe choice: + too tight a bound truncates the posterior and understates the uncertainty. + + Parameters + ---------- + n_sigma : float, default=10.0 + How many standard deviations of the fitted uncertainty to allow on each side. + relative_pad : float, default=0.2 + Extra half-width as a fraction of the absolute parameter value. This guards against + minimizers that report a zero or absurdly small uncertainty. + absolute_floor : float | None, default=None + A minimum half-width in the parameter's own units, for when neither the uncertainty nor + the value carries the natural scale. + + Returns + ------- + BoundsSuggestions + The proposed bounds, which must be applied explicitly. + """ + return suggest_bounds_for_parameters( + self._get_chain_parameters(), + n_sigma=n_sigma, + relative_pad=relative_pad, + absolute_floor=absolute_floor, + ) + + def check_bounds_for_sampling(self) -> None: + """ + Verify that every free parameter has finite bounds. + + Raises + ------ + ValueError + If any free parameter has an infinite lower or upper bound. + """ + unbounded = unbounded_parameters(self._get_chain_parameters()) + if not unbounded: + return + names = ', '.join(parameter.name for parameter in unbounded) + raise ValueError( + f'Bayesian sampling requires finite bounds on every free parameter, because the ' + f'bounds act as the prior. These parameters are unbounded: {names}. ' + f'Set their min and max, or call suggest_bounds() to propose values.' + ) + + ############# + # Sampling + ############# + + def sample_posterior( + self, + samples: int = 10000, + burn: int = 2000, + thin: int = 10, + population: int | None = None, + parameters: list[Parameter] | list[str] | None = None, + **sampler_options: dict[str, Any], + ) -> SamplingResults: + """ + Draw samples from the posterior distribution of the free parameters. + + This starts a fresh chain, replacing any existing one; use :meth:`extend_sampling` to + continue a chain instead. Parameter values are restored to what they were beforehand, so + sampling never silently moves the model off its fitted values; use + :meth:`set_parameters_to_posterior_median` to adopt the posterior. + + Parameters + ---------- + samples : int, default=10000 + Number of raw samples to draw across all chains, before thinning. This is a guaranteed + minimum rather than an exact count. + burn : int, default=2000 + Burn-in generations to discard before collecting samples. + thin : int, default=10 + Thinning interval, which reduces autocorrelation between retained draws. + population : int | None, default=None + DREAM population scale factor: BUMPS runs ``ceil(population * n_parameters)`` chains. + parameters : list[Parameter] | list[str] | None, default=None + Restrict the chain to these parameters, given as Parameter objects or names. All other + free parameters are held fixed for the duration of the run. Note that holding a + parameter fixed is not the same as marginalizing over it: the resulting credible + intervals are conditional on the fixed values and will be too narrow if the parameters + are correlated. The default samples every free parameter. + **sampler_options : dict[str, Any] + Forwarded to the EasyScience Sampler, e.g. ``sampler_kwargs``, ``progress_callback``, + or ``abort_test``. + + Returns + ------- + SamplingResults + The sampling results, also stored on :attr:`posterior_result`. + """ + return self._run_sampling( + parameters=parameters, + run=lambda sampler: sampler.sample( + samples=samples, + burn=burn, + thin=thin, + population=population, + **sampler_options, + ), + ) + + def extend_sampling( + self, + additional_samples: int = 5000, + thin: int = 10, + parameters: list[Parameter] | list[str] | None = None, + **sampler_options: dict[str, Any], + ) -> SamplingResults: + """ + Continue the existing chain with additional samples. + + Parameters + ---------- + additional_samples : int, default=5000 + Number of additional samples to draw, in the same units as ``samples``. + thin : int, default=10 + Thinning interval for the retained draws. + parameters : list[Parameter] | list[str] | None, default=None + The same restriction as in :meth:`sample_posterior`. Pass the same value that started + the chain, since the chain's columns cannot change on extension. + **sampler_options : dict[str, Any] + Forwarded to the EasyScience Sampler. + + Returns + ------- + SamplingResults + The sampling results for the full extended chain. + + Raises + ------ + RuntimeError + If there is no chain to extend. + """ + if self._bayesian_sampler is None: + raise RuntimeError( + 'No chain to extend. Call sample_posterior() or load_chain() first.' + ) + return self._run_sampling( + parameters=parameters, + run=lambda sampler: sampler.extend( + additional_samples=additional_samples, + thin=thin, + **sampler_options, + ), + reuse_sampler=True, + ) + + def _run_sampling( + self, + parameters: list[Parameter] | list[str] | None, + run: Callable[[Sampler], SamplingResults], + reuse_sampler: bool = False, + ) -> SamplingResults: + """ + Run a sampling operation with all the surrounding guards in place. + + Checks the bounds, switches the minimizer to BUMPS, optionally holds parameters fixed, + runs, and then restores the parameter values, fixed flags, and minimizer. + + Parameters + ---------- + parameters : list[Parameter] | list[str] | None + Parameters to restrict the chain to, or None for all free parameters. + run : Callable[[Sampler], SamplingResults] + The operation to perform on the prepared Sampler. + reuse_sampler : bool, default=False + Whether to reuse the cached Sampler rather than rebuilding it. Required when extending + a chain, since the chain lives on the Sampler. + + Returns + ------- + SamplingResults + The results of the run. + + Raises + ------ + RuntimeError + If the BUMPS sampler fails while removing outlier chains, which points at degenerate + parameters. + """ + held_fixed = self._resolve_parameters_to_hold_fixed(parameters) + self._warn_about_held_parameters(held_fixed) + + with _FixedParameters(held_fixed): + self.check_bounds_for_sampling() + self._prepare_for_sampling() + + chain_parameters = self._get_chain_parameters() + saved_values = [(p, p.value) for p in chain_parameters] + + fitter = self.fitter + original_minimizer = fitter.minimizer.enum + fitter.switch_minimizer(AvailableMinimizers.Bumps) + try: + sampler = self._get_or_build_sampler(reuse_sampler=reuse_sampler) + results = run(sampler) + except IndexError as error: + # BUMPS' own outlier removal indexes past the end of its buffer when chains + # scatter wildly, which in practice means the model is not identifiable. The bare + # IndexError says nothing useful, so point at the likely cause instead. + raise RuntimeError( + 'The BUMPS sampler failed while removing outlier chains. This usually means ' + 'the chains scattered because two or more free parameters are degenerate, so ' + 'the data cannot determine them separately. Check for degenerate parameters ' + "and fix one of them, or retry with sampler_kwargs={'outliers': 'none'}." + ) from error + finally: + fitter.switch_minimizer(original_minimizer) + for parameter, value in saved_values: + parameter.value = value + + # A fresh chain is labelled with this session's unique names, so any mapping left over + # from a loaded chain no longer applies. + self._chain_name_map = { + parameter.unique_name: parameter.name for parameter in chain_parameters + } + + self._posterior_result = results + self._warn_about_bounds_occupancy(results, self._resolve_chain_parameters(results)) + return results + + def _get_or_build_sampler(self, reuse_sampler: bool) -> Sampler: + """ + Get the cached Sampler, rebuilding it if the data or model changed. + + Parameters + ---------- + reuse_sampler : bool + Whether to reuse the cached Sampler even if it is marked dirty. + + Returns + ------- + Sampler + The Sampler to run. + """ + needs_rebuild = self._bayesian_sampler is None or ( + self._bayesian_sampler_is_dirty and not reuse_sampler + ) + if needs_rebuild: + x, y, weights = self._get_sampling_data() + self._bayesian_sampler = Sampler(self.fitter, x, y, weights=weights) + self._bayesian_sampler_is_dirty = False + return self._bayesian_sampler + + def _resolve_parameters_to_hold_fixed( + self, + parameters: list[Parameter] | list[str] | None, + ) -> list[Parameter]: + """ + Work out which free parameters must be held fixed to honour a subset request. + + Parameters + ---------- + parameters : list[Parameter] | list[str] | None + The requested subset, as Parameter objects or names, or None for all free parameters. + + Returns + ------- + list[Parameter] + The free parameters that are not in the requested subset. + + Raises + ------ + TypeError + If parameters is not a list of Parameters or strings, or None. + ValueError + If a requested name does not match any free parameter, or the subset is empty. + """ + if parameters is None: + return [] + if not isinstance(parameters, (list, tuple)): + raise TypeError('parameters must be a list of Parameters, a list of names, or None.') + + free = self._get_chain_parameters() + by_name = {parameter.name: parameter for parameter in free} + requested = [] + for entry in parameters: + if isinstance(entry, str): + if entry not in by_name: + available = ', '.join(sorted(by_name)) + raise ValueError(f'No free parameter named {entry!r}. Available: {available}.') + requested.append(by_name[entry]) + elif hasattr(entry, 'unique_name'): + requested.append(entry) + else: + raise TypeError( + 'parameters must contain Parameter objects or parameter names (strings).' + ) + + requested_unique_names = {parameter.unique_name for parameter in requested} + if not requested_unique_names: + raise ValueError('parameters must name at least one parameter to sample.') + return [ + parameter for parameter in free if parameter.unique_name not in requested_unique_names + ] + + @staticmethod + def _warn_about_held_parameters(held_fixed: list[Parameter]) -> None: + """ + Warn that holding parameters fixed makes the credible intervals conditional. + + Parameters + ---------- + held_fixed : list[Parameter] + The parameters being held fixed for the run. + """ + if not held_fixed: + return + names = ', '.join(parameter.name for parameter in held_fixed) + warnings.warn( + ( + f'Holding these parameters fixed while sampling: {names}. ' + f'Fixing a parameter is not the same as marginalizing over it, so the resulting ' + f'credible intervals are conditional on these values and will be too narrow if ' + f'the parameters are correlated.' + ), + UserWarning, + stacklevel=4, + ) + + @staticmethod + def _warn_about_bounds_occupancy( + results: SamplingResults, + parameters_by_column: list[Parameter | None], + ) -> None: + """ + Warn when the posterior has piled up against a bound. + + Parameters + ---------- + results : SamplingResults + The sampling results to inspect. + parameters_by_column : list[Parameter | None] + The parameter for each column of the chain, or None where none could be matched. + """ + piled_up = parameters_at_bounds(results.draws, parameters_by_column) + if not piled_up: + return + details = ', '.join( + f'{name} ({fraction:.0%} of draws)' for name, fraction in piled_up.items() + ) + warnings.warn( + ( + f'The posterior is piled up against the bounds for: {details}. ' + f'The bounds, rather than the data, are setting these credible intervals. ' + f'Widen the bounds, or check whether these parameters are degenerate with others.' + ), + UserWarning, + stacklevel=4, + ) + + ############# + # Results + ############# + + def posterior_summary(self) -> PosteriorSummary: + """ + Summarize the marginal posterior of each sampled parameter. + + Reports the median and the 68% credible interval under the parameter's own name and unit, + rather than the opaque unique name the sampler uses internally. Requires a completed + sampling run. + + Returns + ------- + PosteriorSummary + One entry per sampled parameter. + """ + results = self._require_posterior_result() + return summarize_draws( + draws=results.draws, + fallback_names=self._chain_display_names(results), + parameters_by_column=self._resolve_chain_parameters(results), + ) + + def set_parameters_to_posterior_median(self) -> list[Parameter]: + """ + Set every sampled parameter to the median of its marginal posterior. + + Note that the vector of marginal medians is not in general the same as the + highest-posterior point, and for strongly correlated parameters it need not even be a good + fit. Requires a completed sampling run. + + Returns + ------- + list[Parameter] + The parameters that were changed. + """ + results = self._require_posterior_result() + changed = [] + for column, parameter in enumerate(self._resolve_chain_parameters(results)): + if parameter is None: + continue + parameter.value = float(np.median(results.draws[:, column])) + changed.append(parameter) + return changed + + def _require_posterior_result(self) -> SamplingResults: + """ + Get the stored sampling results, raising if there are none. + + Returns + ------- + SamplingResults + The most recent sampling results. + + Raises + ------ + RuntimeError + If no sampling has been run yet. + """ + if self._posterior_result is None: + raise RuntimeError( + 'No posterior samples yet. Call sample_posterior() or load_chain() first.' + ) + return self._posterior_result + + ############# + # Persistence + ############# + + def save_chain(self, path: str | os.PathLike) -> None: + """ + Save the MCMC chain to disk. + + Writes the BUMPS chain files alongside a sidecar recording the parameter names and a + fingerprint of the data that was sampled. + + Parameters + ---------- + path : str | os.PathLike + Path prefix for the chain files. + + Raises + ------ + RuntimeError + If there is no chain to save. + """ + if self._bayesian_sampler is None: + raise RuntimeError('No chain to save. Call sample_posterior() first.') + self._bayesian_sampler.save(path) + # The BUMPS sidecar records unique names, which are handed out per session and so mean + # nothing on reload. Record the parameter names alongside them, which are stable. + Path(f'{path}{_NAME_MAP_SUFFIX}').write_text( + json.dumps(self._chain_name_map, indent=2), + encoding='utf-8', + ) + + def load_chain(self, path: str | os.PathLike, skip: int = 0) -> SamplingResults: + """ + Load a previously saved MCMC chain. + + The loaded chain can be inspected, summarized, or continued with :meth:`extend_sampling`. A + chain saved from different data loads with a warning. + + Parameters + ---------- + path : str | os.PathLike + The path prefix the chain was saved under. + skip : int, default=0 + Number of initial samples to skip when reading the chain. + + Returns + ------- + SamplingResults + The loaded sampling results, also stored on :attr:`posterior_result`. + """ + self._prepare_for_sampling() + name_map_path = Path(f'{path}{_NAME_MAP_SUFFIX}') + if name_map_path.is_file(): + self._chain_name_map = json.loads(name_map_path.read_text(encoding='utf-8')) + else: + self._chain_name_map = {} + warnings.warn( + ( + f'No parameter-name sidecar found at {name_map_path}. The chain will be ' + f'reported under the internal names it was saved with, because those cannot ' + f'be matched to this Analysis.' + ), + UserWarning, + stacklevel=2, + ) + + fitter = self.fitter + original_minimizer = fitter.minimizer.enum + fitter.switch_minimizer(AvailableMinimizers.Bumps) + try: + sampler = self._get_or_build_sampler(reuse_sampler=False) + results = sampler.load_state(path, skip=skip) + finally: + fitter.switch_minimizer(original_minimizer) + self._posterior_result = results + return results + + ############# + # Plotting + ############# + + def plot_trace(self, **kwargs: dict[str, Any]) -> Figure: + """ + Plot the chain trace of each sampled parameter. + + A well-mixed chain looks like a "hairy caterpillar" with no drift; visible trends mean the + chain has not converged and needs a longer burn-in. Requires a completed sampling run. + + Parameters + ---------- + **kwargs : dict[str, Any] + Forwarded to :func:`easydynamics.utils.posterior_plotting.plot_trace`. + + Returns + ------- + Figure + The matplotlib Figure. + """ + from easydynamics.utils.posterior_plotting import plot_trace + + results = self._require_posterior_result() + return plot_trace( + draws=results.draws, + logp=results.logp, + names=self._chain_display_names(results), + title=self.display_name, + **kwargs, + ) + + def plot_corner(self, **kwargs: dict[str, Any]) -> Figure: + """ + Plot the marginal and pairwise posterior distributions. + + Diagonal panels show each parameter's marginal distribution; off-diagonal panels show the + joint distribution of a pair, where a strong diagonal ridge means the two are correlated. + Requires a completed sampling run. + + Parameters + ---------- + **kwargs : dict[str, Any] + Forwarded to :func:`easydynamics.utils.posterior_plotting.plot_corner`. + + Returns + ------- + Figure + The matplotlib Figure. + """ + from easydynamics.utils.posterior_plotting import plot_corner + + results = self._require_posterior_result() + return plot_corner( + draws=results.draws, + names=self._chain_display_names(results), + title=self.display_name, + **kwargs, + ) + + def plot_posterior_predictive( + self, + n_draws: int = 200, + credible_interval: float = 68.0, + **kwargs: dict[str, Any], + ) -> Figure: + """ + Plot the data against the credible band implied by the posterior. + + The model is re-evaluated for a random subset of the posterior draws, and the spread of + those curves becomes the band. Data straying outside the band systematically points at a + model that is missing something, rather than at parameters that need tuning. Requires a + completed sampling run. + + Parameters + ---------- + n_draws : int, default=200 + How many posterior draws to evaluate the model for. Each draw costs one full model + evaluation, so this trades smoothness of the band against time. + credible_interval : float, default=68.0 + Width of the credible band, as a percentage. + **kwargs : dict[str, Any] + Forwarded to :func:`easydynamics.utils.posterior_plotting.plot_posterior_predictive`. + + Returns + ------- + Figure + The matplotlib Figure. + + Raises + ------ + NotImplementedError + If this Analysis binds a list of datasets rather than a single one. + ValueError + If n_draws is not a positive integer. + """ + from easydynamics.utils.posterior_plotting import plot_posterior_predictive + + if not isinstance(n_draws, int) or isinstance(n_draws, bool) or n_draws < 1: + raise ValueError(f'n_draws must be a positive integer. Got {n_draws}.') + + results = self._require_posterior_result() + x, y, weights = self._get_sampling_data() + if isinstance(x, (list, tuple)): + raise NotImplementedError( + 'plot_posterior_predictive supports a single dataset only. Plot each dataset ' + 'from its own Analysis1d instead.' + ) + + predictions = self._evaluate_over_draws(results, x, n_draws) + y_err = None if weights is None else 1.0 / np.asarray(weights) + return plot_posterior_predictive( + x=np.asarray(x), + y=np.asarray(y), + predictions=predictions, + y_err=y_err, + title=self.display_name, + credible_interval=credible_interval, + **kwargs, + ) + + def _evaluate_over_draws( + self, + results: SamplingResults, + x: np.ndarray, + n_draws: int, + ) -> np.ndarray: + """ + Evaluate the model once per posterior draw, restoring the parameters afterwards. + + Parameters + ---------- + results : SamplingResults + The sampling results supplying the draws. + x : np.ndarray + The independent variable to evaluate the model on. + n_draws : int + How many draws to evaluate. Draws are taken evenly across the chain. + + Returns + ------- + np.ndarray + Model evaluations, shape ``(n_selected, len(x))``. + """ + self._prepare_for_sampling() + + columns = [ + (parameter, column) + for column, parameter in enumerate(self._resolve_chain_parameters(results)) + if parameter is not None + ] + saved_values = [(parameter, parameter.value) for parameter, _ in columns] + + total = results.draws.shape[0] + indices = np.unique(np.linspace(0, total - 1, min(n_draws, total)).astype(int)) + + fit_function = self.fitter.fit_function + predictions = [] + try: + for index in indices: + for parameter, column in columns: + parameter.value = float(results.draws[index, column]) + predictions.append(np.asarray(fit_function(x))) + finally: + for parameter, value in saved_values: + parameter.value = value + + return np.vstack(predictions) + + def _resolve_chain_parameters(self, results: SamplingResults) -> list[Parameter | None]: + """ + Match each column of the chain to one of this Analysis's parameters. + + Columns are matched on ``unique_name`` first. That fails for a chain loaded from disk, + because unique names are handed out per session, so a saved chain also records the + parameter names and those are used as a fallback. + + Parameters + ---------- + results : SamplingResults + The sampling results whose columns should be matched. + + Returns + ------- + list[Parameter | None] + The parameter for each column, or None where no match could be made. + """ + parameters = self._get_chain_parameters() + by_unique_name = {p.unique_name: p for p in parameters} + by_name = {p.name: p for p in parameters} + resolved = [] + for unique_name in results.param_names: + parameter = by_unique_name.get(unique_name) + if parameter is None: + saved_name = self._chain_name_map.get(unique_name) + parameter = None if saved_name is None else by_name.get(saved_name) + resolved.append(parameter) + return resolved + + def _chain_display_names(self, results: SamplingResults) -> list[str]: + """ + Translate the chain's column names into parameter names. + + Parameters + ---------- + results : SamplingResults + The sampling results whose columns should be named. + + Returns + ------- + list[str] + One name per column of the chain. + """ + resolved = self._resolve_chain_parameters(results) + return [ + self._chain_name_map.get(unique_name, unique_name) + if parameter is None + else parameter.name + for unique_name, parameter in zip(results.param_names, resolved, strict=True) + ] + + +class _FixedParameters: + """ + Context manager that temporarily fixes parameters and restores their flags on exit. + """ + + def __init__(self, parameters: list[Parameter]) -> None: + """ + Initialize the context manager. + + Parameters + ---------- + parameters : list[Parameter] + The parameters to hold fixed for the duration of the block. + """ + self._parameters = list(parameters) + self._saved: list[tuple[Parameter, bool]] = [] + + def __enter__(self) -> None: + """ + Fix the parameters, remembering their previous state. + """ + self._saved = [(parameter, parameter.fixed) for parameter in self._parameters] + for parameter in self._parameters: + parameter.fixed = True + + def __exit__(self, *_exc_info: object) -> None: + """ + Restore the previous fixed state of every parameter. + + Parameters + ---------- + *_exc_info : object + Exception information, ignored. + """ + for parameter, was_fixed in self._saved: + parameter.fixed = was_fixed diff --git a/src/easydynamics/analysis/posterior.py b/src/easydynamics/analysis/posterior.py new file mode 100644 index 000000000..e22ae87d9 --- /dev/null +++ b/src/easydynamics/analysis/posterior.py @@ -0,0 +1,608 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +""" +Bounds suggestions and posterior summaries for Bayesian sampling. + +The helpers here are deliberately free of any Analysis or Fitter machinery: they operate on plain +``Parameter`` objects and on the ``(n_draws, n_parameters)`` array produced by the sampler, so they +can be unit-tested on their own and reused by every Analysis class. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import numpy as np + +if TYPE_CHECKING: + from easyscience.variable import Parameter + +# Fraction of the allowed range at each end that counts as "at the bound" when checking whether +# the posterior has piled up against a bound. +BOUND_EDGE_FRACTION = 0.05 + +# Fraction of draws inside those edge bands above which a pile-up is reported. A posterior spread +# uniformly across its bounds -- the signature of a bound, rather than the data, setting the +# credible interval -- puts 2 * BOUND_EDGE_FRACTION of its draws there. A posterior comfortably +# inside its bounds puts essentially none there, so the threshold sits well below the uniform value +# to stay sensitive to partly-clipped posteriors without risking false positives. +BOUND_OCCUPANCY_THRESHOLD = 0.05 + + +@dataclass(frozen=True) +class BoundsSuggestion: + """ + A proposed pair of bounds for a single parameter. + + Attributes + ---------- + parameter : Parameter + The parameter the suggestion applies to. + suggested_min : float + The proposed lower bound. Equal to the parameter's current lower bound when that is already + finite. + suggested_max : float + The proposed upper bound. Equal to the parameter's current upper bound when that is already + finite. + reason : str + Empty when the suggestion is usable. Otherwise, why the parameter needs manual attention. + """ + + parameter: Parameter + suggested_min: float + suggested_max: float + reason: str + + @property + def needs_attention(self) -> bool: + """ + Whether this parameter could not be given a usable suggestion. + + Returns + ------- + bool + True when no usable bounds could be derived and the user must set them by hand. + """ + return bool(self.reason) + + @property + def changes_bounds(self) -> bool: + """ + Whether applying this suggestion would actually change the parameter. + + Returns + ------- + bool + True when either bound differs from the parameter's current bound. + """ + return self.suggested_min != self.parameter.min or self.suggested_max != self.parameter.max + + +class BoundsSuggestions: + """ + The result of :func:`suggest_bounds_for_parameters`, rendered as a table. + + This is advisory: nothing is changed until :meth:`apply` is called. Suggestions only ever fill + in an infinite bound; a bound that is already finite is never widened or narrowed, so physical + limits such as a non-negative area survive untouched. + """ + + def __init__(self, suggestions: list[BoundsSuggestion]) -> None: + """ + Initialize the collection. + + Parameters + ---------- + suggestions : list[BoundsSuggestion] + The per-parameter suggestions. + """ + self._suggestions = list(suggestions) + + @property + def suggestions(self) -> list[BoundsSuggestion]: + """ + All suggestions, including those needing manual attention. + + Returns + ------- + list[BoundsSuggestion] + The per-parameter suggestions. + """ + return list(self._suggestions) + + @property + def needing_attention(self) -> list[BoundsSuggestion]: + """ + The suggestions for which no usable bounds could be derived. + + Returns + ------- + list[BoundsSuggestion] + Suggestions whose parameters must be bounded by hand. + """ + return [s for s in self._suggestions if s.needs_attention] + + def apply(self) -> list[Parameter]: + """ + Set the suggested bounds on every parameter that has a usable suggestion. + + Parameters needing manual attention are skipped rather than guessed at. + + Returns + ------- + list[Parameter] + The parameters whose bounds were changed. + """ + changed = [] + for suggestion in self._suggestions: + if suggestion.needs_attention or not suggestion.changes_bounds: + continue + suggestion.parameter.min = suggestion.suggested_min + suggestion.parameter.max = suggestion.suggested_max + changed.append(suggestion.parameter) + return changed + + def __len__(self) -> int: + """ + Return the number of suggestions. + + Returns + ------- + int + The number of suggestions. + """ + return len(self._suggestions) + + def __iter__(self) -> iter: + """ + Iterate over the suggestions. + + Returns + ------- + iter + An iterator over the suggestions. + """ + return iter(self._suggestions) + + def __repr__(self) -> str: + """ + Render the suggestions as a table. + + Returns + ------- + str + A table of current and suggested bounds, one row per parameter. + """ + if not self._suggestions: + return 'BoundsSuggestions(no free parameters)' + + header = f'{"parameter":<28s} {"current":>26s} {"suggested":>26s}' + lines = ['BoundsSuggestions', header, '-' * len(header)] + for s in self._suggestions: + current = f'({s.parameter.min:.4g}, {s.parameter.max:.4g})' + if s.needs_attention: + suggested = f'-- {s.reason}' + else: + suggested = f'({s.suggested_min:.4g}, {s.suggested_max:.4g})' + lines.append(f'{s.parameter.name:<28s} {current:>26s} {suggested:>26s}') + + attention = self.needing_attention + if attention: + lines.append('') + lines.append( + f'{len(attention)} parameter(s) need bounds set by hand; .apply() will skip them.' + ) + return '\n'.join(lines) + + +def suggest_bounds_for_parameters( + parameters: list[Parameter], + n_sigma: float = 10.0, + relative_pad: float = 0.2, + absolute_floor: float | None = None, +) -> BoundsSuggestions: + """ + Propose finite bounds for parameters that currently have an infinite one. + + The half-width of a proposed bound is ``n_sigma * error + relative_pad * abs(value)``, floored + at ``absolute_floor`` when one is given. The ``relative_pad`` term matters because + least-squares minimizers sometimes report a zero or absurdly small uncertainty; without it such + a parameter would be given a zero-width bound. When the half-width still comes out as zero or + non-finite, the parameter is flagged for manual attention rather than given an invented scale. + + In BUMPS' DREAM sampler the bounds act as a uniform prior, so a generous width is the safe + choice: too narrow a bound truncates the posterior and understates the uncertainty. Hence the + deliberately loose ``n_sigma`` default. + + A ``TypeError`` is raised if any of the three settings is not a number, and a ``ValueError`` if + any is negative. + + Parameters + ---------- + parameters : list[Parameter] + The parameters to propose bounds for. + n_sigma : float, default=10.0 + How many standard deviations of the parameter's fitted uncertainty to allow on each side. + relative_pad : float, default=0.2 + Extra half-width as a fraction of the absolute parameter value, guarding against + artificially small uncertainties. + absolute_floor : float | None, default=None + A minimum half-width, in the parameter's own units. Use it when the natural scale is known + but neither the uncertainty nor the value carries it. + + Returns + ------- + BoundsSuggestions + The proposed bounds, which must be applied explicitly. + """ + _verify_nonneg_number(n_sigma, 'n_sigma') + _verify_nonneg_number(relative_pad, 'relative_pad') + if absolute_floor is not None: + _verify_nonneg_number(absolute_floor, 'absolute_floor') + + suggestions = [ + _suggest_bounds_for_parameter( + parameter=parameter, + n_sigma=n_sigma, + relative_pad=relative_pad, + absolute_floor=absolute_floor, + ) + for parameter in parameters + ] + return BoundsSuggestions(suggestions) + + +def _suggest_bounds_for_parameter( + parameter: Parameter, + n_sigma: float, + relative_pad: float, + absolute_floor: float | None, +) -> BoundsSuggestion: + """ + Propose bounds for a single parameter. + + Parameters + ---------- + parameter : Parameter + The parameter to propose bounds for. + n_sigma : float + How many standard deviations to allow on each side. + relative_pad : float + Extra half-width as a fraction of the absolute parameter value. + absolute_floor : float | None + A minimum half-width, or None. + + Returns + ------- + BoundsSuggestion + The proposal for this parameter. + """ + current_min = float(parameter.min) + current_max = float(parameter.max) + min_is_finite = np.isfinite(current_min) + max_is_finite = np.isfinite(current_max) + + # Nothing to fill in: a bound that is already finite is never touched. + if min_is_finite and max_is_finite: + return BoundsSuggestion( + parameter=parameter, + suggested_min=current_min, + suggested_max=current_max, + reason='', + ) + + value = float(parameter.value) + error = float(parameter.error) + if not np.isfinite(value): + return BoundsSuggestion( + parameter=parameter, + suggested_min=current_min, + suggested_max=current_max, + reason='value is not finite', + ) + + half_width = relative_pad * abs(value) + if np.isfinite(error): + half_width += n_sigma * error + if absolute_floor is not None: + half_width = max(half_width, absolute_floor) + + if not np.isfinite(half_width) or half_width <= 0: + return BoundsSuggestion( + parameter=parameter, + suggested_min=current_min, + suggested_max=current_max, + reason='no scale information (zero value and uncertainty)', + ) + + return BoundsSuggestion( + parameter=parameter, + suggested_min=current_min if min_is_finite else value - half_width, + suggested_max=current_max if max_is_finite else value + half_width, + reason='', + ) + + +def unbounded_parameters(parameters: list[Parameter]) -> list[Parameter]: + """ + Find parameters with a non-finite lower or upper bound. + + Parameters + ---------- + parameters : list[Parameter] + The parameters to check. + + Returns + ------- + list[Parameter] + Those parameters that have at least one infinite bound. + """ + return [ + parameter + for parameter in parameters + if not (np.isfinite(parameter.min) and np.isfinite(parameter.max)) + ] + + +def parameters_at_bounds( + draws: np.ndarray, + parameters_by_column: list[Parameter | None], +) -> dict[str, float]: + """ + Find parameters whose posterior has piled up against one of its bounds. + + A chain that spends much of its time hard against a bound is a sign that the bound, rather than + the data, is setting the credible interval. That happens when a bound is too tight, and also + when two parameters are degenerate and the pair drifts until it is stopped by a bound. + + Parameters + ---------- + draws : np.ndarray + Posterior draws, shape ``(n_draws, n_parameters)``. + parameters_by_column : list[Parameter | None] + The parameter for each column of ``draws``, or None where no parameter could be matched. + + Returns + ------- + dict[str, float] + Mapping of parameter name to the fraction of draws sitting in the outer + ``BOUND_EDGE_FRACTION`` of its allowed range, for those parameters where that fraction + exceeds ``BOUND_OCCUPANCY_THRESHOLD``. + """ + piled_up = {} + for column, parameter in enumerate(parameters_by_column): + if parameter is None: + continue + low = float(parameter.min) + high = float(parameter.max) + if not (np.isfinite(low) and np.isfinite(high)) or high <= low: + continue + edge = BOUND_EDGE_FRACTION * (high - low) + values = draws[:, column] + at_edge = (values <= low + edge) | (values >= high - edge) + fraction = float(np.count_nonzero(at_edge)) / len(values) + if fraction > BOUND_OCCUPANCY_THRESHOLD: + piled_up[parameter.name] = fraction + return piled_up + + +@dataclass(frozen=True) +class ParameterPosterior: + """ + The marginal posterior of a single parameter. + + Attributes + ---------- + name : str + The parameter's name. + unit : str + The parameter's unit, as a string. + median : float + The 50th percentile of the marginal posterior. + lower : float + The 16th percentile. + upper : float + The 84th percentile. + value : float + The parameter's current value, for comparison with the median. + """ + + name: str + unit: str + median: float + lower: float + upper: float + value: float + + @property + def minus(self) -> float: + """ + Distance from the median down to the 16th percentile. + + Returns + ------- + float + The lower half of the 68% credible interval. + """ + return self.median - self.lower + + @property + def plus(self) -> float: + """ + Distance from the median up to the 84th percentile. + + Returns + ------- + float + The upper half of the 68% credible interval. + """ + return self.upper - self.median + + +class PosteriorSummary: + """ + Marginal posterior summaries for every sampled parameter, rendered as a table. + """ + + def __init__(self, entries: list[ParameterPosterior]) -> None: + """ + Initialize the summary. + + Parameters + ---------- + entries : list[ParameterPosterior] + One entry per sampled parameter. + """ + self._entries = list(entries) + + @property + def entries(self) -> list[ParameterPosterior]: + """ + The per-parameter summaries. + + Returns + ------- + list[ParameterPosterior] + One entry per sampled parameter. + """ + return list(self._entries) + + def __len__(self) -> int: + """ + Return the number of summarized parameters. + + Returns + ------- + int + The number of entries. + """ + return len(self._entries) + + def __iter__(self) -> iter: + """ + Iterate over the entries. + + Returns + ------- + iter + An iterator over the entries. + """ + return iter(self._entries) + + def __getitem__(self, name: str) -> ParameterPosterior: + """ + Look up a parameter's summary by name. + + Parameters + ---------- + name : str + The parameter name. + + Returns + ------- + ParameterPosterior + The summary for that parameter. + + Raises + ------ + KeyError + If no sampled parameter has that name. + """ + for entry in self._entries: + if entry.name == name: + return entry + raise KeyError(f'No sampled parameter named {name!r}.') + + def __repr__(self) -> str: + """ + Render the summary as a table. + + Returns + ------- + str + A table with the median and 68% credible interval of each parameter. + """ + if not self._entries: + return 'PosteriorSummary(no parameters)' + + header = ( + f'{"parameter":<28s} {"unit":>10s} {"median":>14s} ' + f'{"-":>12s} {"+":>12s} {"current":>14s}' + ) + lines = ['PosteriorSummary', header, '-' * len(header)] + lines.extend( + f'{e.name:<28s} {e.unit:>10s} {e.median:>14.5g} ' + f'{e.minus:>12.4g} {e.plus:>12.4g} {e.value:>14.5g}' + for e in self._entries + ) + return '\n'.join(lines) + + +def summarize_draws( + draws: np.ndarray, + fallback_names: list[str], + parameters_by_column: list[Parameter | None], +) -> PosteriorSummary: + """ + Summarize posterior draws under the parameters' own names and units. + + The sampler labels its columns with each parameter's ``unique_name`` (``Parameter_4`` and the + like), which is not what a user recognises, so columns are reported under ``Parameter.name`` + wherever a parameter could be matched. + + Parameters + ---------- + draws : np.ndarray + Posterior draws, shape ``(n_draws, n_parameters)``. + fallback_names : list[str] + Label to use for any column with no matching parameter, one per column. + parameters_by_column : list[Parameter | None] + The parameter for each column of ``draws``, or None where none could be matched. + + Returns + ------- + PosteriorSummary + One entry per column of ``draws``, in column order. + """ + entries = [] + for column, parameter in enumerate(parameters_by_column): + lower, median, upper = ( + float(percentile) for percentile in np.percentile(draws[:, column], [16, 50, 84]) + ) + entries.append( + ParameterPosterior( + name=fallback_names[column] if parameter is None else parameter.name, + unit='' if parameter is None else str(parameter.unit), + median=median, + lower=lower, + upper=upper, + value=float('nan') if parameter is None else float(parameter.value), + ) + ) + return PosteriorSummary(entries) + + +def _verify_nonneg_number(value: object, name: str) -> None: + """ + Raise if a value is not a non-negative number. + + Parameters + ---------- + value : object + The object to verify. + name : str + The name of the object, for the error message. + + Raises + ------ + TypeError + If value is not an int or float. + ValueError + If value is negative. + """ + if not isinstance(value, (int, float)) or isinstance(value, bool): + raise TypeError(f'{name} must be a number. Got {type(value)}.') + if value < 0: + raise ValueError(f'{name} must be non-negative. Got {value}.') diff --git a/src/easydynamics/utils/__init__.py b/src/easydynamics/utils/__init__.py index 5e644a06b..1c3402ced 100644 --- a/src/easydynamics/utils/__init__.py +++ b/src/easydynamics/utils/__init__.py @@ -3,5 +3,14 @@ from easydynamics.utils.detailed_balance import detailed_balance_factor from easydynamics.utils.plotting import slicerplot_with_residuals +from easydynamics.utils.posterior_plotting import plot_corner +from easydynamics.utils.posterior_plotting import plot_posterior_predictive +from easydynamics.utils.posterior_plotting import plot_trace -__all__ = ['detailed_balance_factor', 'slicerplot_with_residuals'] +__all__ = [ + 'detailed_balance_factor', + 'plot_corner', + 'plot_posterior_predictive', + 'plot_trace', + 'slicerplot_with_residuals', +] diff --git a/src/easydynamics/utils/posterior_plotting.py b/src/easydynamics/utils/posterior_plotting.py new file mode 100644 index 000000000..bfb1a1320 --- /dev/null +++ b/src/easydynamics/utils/posterior_plotting.py @@ -0,0 +1,256 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +""" +Diagnostic plots for Bayesian posterior samples. + +These take plain arrays rather than an Analysis, so they can be used on any chain, including one +loaded from disk. The Analysis classes wrap them in convenience methods. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import matplotlib.pyplot as plt +import numpy as np + +if TYPE_CHECKING: + from matplotlib.figure import Figure + + +def plot_trace( + draws: np.ndarray, + names: list[str], + logp: np.ndarray | None = None, + title: str | None = None, + figsize: tuple[float, float] | None = None, +) -> Figure: + """ + Plot the chain trace of every sampled parameter. + + A converged chain looks like a "hairy caterpillar": noisy but stationary, with no drift or long + excursions. A visible trend means the chain has not reached the typical set and needs a longer + burn-in. + + A ``ValueError`` is raised if ``draws`` is not two-dimensional, or if ``names`` does not have + one entry per column. + + Parameters + ---------- + draws : np.ndarray + Posterior draws, shape ``(n_draws, n_parameters)``. + names : list[str] + One label per column of ``draws``. + logp : np.ndarray | None, default=None + Log-posterior values, plotted in an extra panel when given. + title : str | None, default=None + Figure title. + figsize : tuple[float, float] | None, default=None + Figure size in inches. Defaults to a height that scales with the number of panels. + + Returns + ------- + Figure + The matplotlib Figure. + """ + draws = np.asarray(draws) + _verify_draws(draws, names) + + n_panels = draws.shape[1] + (1 if logp is not None else 0) + if figsize is None: + figsize = (10.0, max(2.0, 1.6 * n_panels)) + + fig, axes = plt.subplots(n_panels, 1, figsize=figsize, sharex=True, squeeze=False) + axes = axes[:, 0] + + for axis, column, name in zip(axes, range(draws.shape[1]), names, strict=False): + axis.plot(draws[:, column], lw=0.5) + axis.set_ylabel(name, fontsize=8) + axis.set_xlim(0, len(draws) - 1) + + if logp is not None: + axes[-1].plot(np.asarray(logp), lw=0.5, color='C4') + axes[-1].set_ylabel('log-posterior', fontsize=8) + + axes[-1].set_xlabel('sample index') + if title is not None: + fig.suptitle(title) + fig.tight_layout() + return fig + + +def plot_corner( + draws: np.ndarray, + names: list[str], + title: str | None = None, + bins: int = 40, + figsize: tuple[float, float] | None = None, +) -> Figure: + """ + Plot marginal and pairwise posterior distributions. + + Diagonal panels show each parameter's marginal distribution. Off-diagonal panels show the joint + distribution of a pair: a compact blob means the two are independent, while a narrow diagonal + ridge means they are correlated and cannot be determined separately from this data. + + A ``ValueError`` is raised if ``draws`` is not two-dimensional, or if ``names`` does not have + one entry per column. + + Parameters + ---------- + draws : np.ndarray + Posterior draws, shape ``(n_draws, n_parameters)``. + names : list[str] + One label per column of ``draws``. + title : str | None, default=None + Figure title. + bins : int, default=40 + Number of bins for the marginal histograms. + figsize : tuple[float, float] | None, default=None + Figure size in inches. Defaults to a square that scales with the parameter count. + + Returns + ------- + Figure + The matplotlib Figure. + """ + draws = np.asarray(draws) + _verify_draws(draws, names) + + n = draws.shape[1] + if figsize is None: + side = max(4.0, 2.0 * n) + figsize = (side, side) + + fig, axes = plt.subplots(n, n, figsize=figsize, squeeze=False) + for row in range(n): + for col in range(n): + axis = axes[row, col] + if col > row: + axis.set_visible(False) + continue + if row == col: + axis.hist(draws[:, row], bins=bins, color='C0', histtype='stepfilled', alpha=0.7) + axis.set_yticks([]) + else: + axis.hexbin(draws[:, col], draws[:, row], gridsize=30, cmap='Blues', mincnt=1) + if row == n - 1: + axis.set_xlabel(names[col], fontsize=8) + else: + axis.set_xticklabels([]) + if col == 0 and row != 0: + axis.set_ylabel(names[row], fontsize=8) + else: + axis.set_yticklabels([]) + axis.tick_params(labelsize=7) + + if title is not None: + fig.suptitle(title) + fig.tight_layout() + return fig + + +def plot_posterior_predictive( + x: np.ndarray, + y: np.ndarray, + predictions: np.ndarray, + y_err: np.ndarray | None = None, + title: str | None = None, + credible_interval: float = 68.0, + figsize: tuple[float, float] = (8.0, 5.0), +) -> Figure: + """ + Plot the data against the credible band implied by the posterior. + + The band shows where the model says the data should lie, given the posterior. If the data + strays outside it systematically, the model is missing something that no amount of parameter + tuning will fix. + + Parameters + ---------- + x : np.ndarray + Independent variable of the data. + y : np.ndarray + Observed values. + predictions : np.ndarray + Model evaluations, shape ``(n_draws, len(x))``, one row per posterior draw. + y_err : np.ndarray | None, default=None + Standard deviation of the observed values, drawn as error bars when given. + title : str | None, default=None + Figure title. + credible_interval : float, default=68.0 + Width of the credible band, as a percentage. + figsize : tuple[float, float], default=(8.0, 5.0) + Figure size in inches. + + Returns + ------- + Figure + The matplotlib Figure. + + Raises + ------ + ValueError + If ``predictions`` is not two-dimensional with one column per point in ``x``, or if + ``credible_interval`` is not between 0 and 100. + """ + x = np.asarray(x) + y = np.asarray(y) + predictions = np.asarray(predictions) + if predictions.ndim != 2 or predictions.shape[1] != len(x): + raise ValueError( + f'predictions must have shape (n_draws, {len(x)}). Got {predictions.shape}.' + ) + if not 0 < credible_interval < 100: + raise ValueError(f'credible_interval must be between 0 and 100. Got {credible_interval}.') + + tail = (100.0 - credible_interval) / 2.0 + lower, median, upper = np.percentile(predictions, [tail, 50.0, 100.0 - tail], axis=0) + + fig, axis = plt.subplots(figsize=figsize) + if y_err is None: + axis.plot(x, y, 'o', mfc='none', color='black', label='Data', markersize=4) + else: + axis.errorbar( + x, y, np.asarray(y_err), fmt='o', mfc='none', color='black', label='Data', markersize=4 + ) + axis.fill_between( + x, + lower, + upper, + color='C3', + alpha=0.3, + label=f'{credible_interval:.0f}% credible band', + ) + axis.plot(x, median, '-', color='C3', label='Posterior median') + axis.legend() + if title is not None: + axis.set_title(title) + fig.tight_layout() + return fig + + +def _verify_draws(draws: np.ndarray, names: list[str]) -> None: + """ + Verify that a draws array is two-dimensional and matches its labels. + + Parameters + ---------- + draws : np.ndarray + The posterior draws to check. + names : list[str] + The labels to check against. + + Raises + ------ + ValueError + If ``draws`` is not two-dimensional or its column count differs from ``len(names)``. + """ + if draws.ndim != 2: + raise ValueError(f'draws must be two-dimensional. Got shape {draws.shape}.') + if draws.shape[1] != len(names): + raise ValueError( + f'names must have one entry per column of draws. ' + f'Got {len(names)} names for {draws.shape[1]} columns.' + ) diff --git a/tests/integration/fitting/test_bayesian_sampling.py b/tests/integration/fitting/test_bayesian_sampling.py new file mode 100644 index 000000000..704a37ea2 --- /dev/null +++ b/tests/integration/fitting/test_bayesian_sampling.py @@ -0,0 +1,207 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +""" +Integration tests running real BUMPS DREAM chains through Analysis1d. + +These are slow by nature. They deliberately run with ``sampler_kwargs={'trim': False}``: BUMPS' +automatic burn-point trimming re-runs a convergence detector on every call and can crash inside its +own outlier removal on the very short chains used here. +""" + +import warnings + +import matplotlib as mpl +import numpy as np +import pytest +import scipp as sc + +mpl.use('Agg') + +from easydynamics.analysis.analysis1d import Analysis1d +from easydynamics.experiment import Experiment +from easydynamics.sample_model import InstrumentModel +from easydynamics.sample_model import SampleModel +from easydynamics.sample_model.components.gaussian import Gaussian + +TRUE_AREA = 9.0 +TRUE_WIDTH = 1.2 +NOISE = 0.05 + +# Keep the chains short enough to stay usable in CI; long enough to locate the peak. +SAMPLE_KWARGS = { + 'samples': 2000, + 'burn': 100, + 'thin': 2, + 'sampler_kwargs': {'trim': False}, +} + + +def build_analysis(): + energy_values = np.linspace(-5.0, 5.0, 60) + truth = TRUE_AREA / (TRUE_WIDTH * np.sqrt(2 * np.pi)) + truth = truth * np.exp(-0.5 * (energy_values / TRUE_WIDTH) ** 2) + observed = truth + np.random.default_rng(0).normal(0.0, NOISE, size=truth.shape) + + data = sc.array( + dims=['Q', 'energy'], + values=observed[None, :], + variances=np.full_like(observed, NOISE**2)[None, :], + ) + experiment = Experiment( + data=sc.DataArray( + data=data, + coords={ + 'Q': sc.array(dims=['Q'], values=[1.0], unit='1/Angstrom'), + 'energy': sc.array(dims=['energy'], values=energy_values, unit='meV'), + }, + ) + ) + analysis = Analysis1d( + display_name='BayesianIntegration', + experiment=experiment, + sample_model=SampleModel( + components=Gaussian(area=TRUE_AREA, width=TRUE_WIDTH, center=0.0) + ), + instrument_model=InstrumentModel(), + Q_index=0, + ) + # The energy offset shifts the spectrum exactly as the Gaussian centre does. Leaving both free + # makes the model unidentifiable, which no amount of sampling can repair. + analysis.instrument_model.fix_energy_offset(Q_index=0) + return analysis + + +@pytest.fixture(scope='module') +def sampled_analysis(): + analysis = build_analysis() + analysis.fit() + analysis.suggest_bounds().apply() + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + analysis.sample_posterior(**SAMPLE_KWARGS) + return analysis + + +class TestRealChain: + def test_chain_has_one_column_per_free_parameter(self, sampled_analysis): + # EXPECT + results = sampled_analysis.posterior_result + assert results.draws.shape[1] == len(sampled_analysis.get_free_parameters()) + assert results.draws.shape[0] > 0 + + @pytest.mark.parametrize( + ('name', 'truth'), + [('Gaussian area', TRUE_AREA), ('Gaussian width', TRUE_WIDTH)], + ) + def test_posterior_recovers_the_true_parameters(self, sampled_analysis, name, truth): + # WHEN + entry = sampled_analysis.posterior_summary()[name] + + # EXPECT the truth sits within a few posterior standard deviations of the median. A 68% + # interval is deliberately not used: it excludes the truth about a third of the time for + # any single noise realization, which would make this test flaky rather than strict. + spread = max(entry.minus, entry.plus) + assert abs(entry.median - truth) < 4 * spread + + def test_summary_is_reported_under_parameter_names_and_units(self, sampled_analysis): + # WHEN + summary = sampled_analysis.posterior_summary() + + # EXPECT + assert {entry.name for entry in summary} == { + p.name for p in sampled_analysis.get_free_parameters() + } + assert all(entry.unit == 'meV' for entry in summary) + + def test_sampling_leaves_the_fitted_values_untouched(self): + # WHEN + analysis = build_analysis() + analysis.fit() + analysis.suggest_bounds().apply() + before = [float(p.value) for p in analysis.get_free_parameters()] + + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + analysis.sample_posterior(**SAMPLE_KWARGS) + + # EXPECT + after = [float(p.value) for p in analysis.get_free_parameters()] + assert after == pytest.approx(before) + + def test_extend_grows_the_chain(self, sampled_analysis): + # WHEN + before = int(sampled_analysis.posterior_result.state.Ngen) + + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + extended = sampled_analysis.extend_sampling( + additional_samples=500, thin=2, sampler_kwargs={'trim': False} + ) + + # EXPECT + assert int(extended.state.Ngen) > before + + def test_save_and_load_round_trip_keeps_parameter_identity(self, sampled_analysis, tmp_path): + # WHEN + prefix = str(tmp_path / 'chain') + sampled_analysis.save_chain(prefix) + + fresh = build_analysis() + fresh.fit() + fresh.suggest_bounds().apply() + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + fresh.load_chain(prefix) + + # EXPECT the reloaded chain is reported under real names, not internal unique names + summary = fresh.posterior_summary() + assert {entry.name for entry in summary} == {p.name for p in fresh.get_free_parameters()} + assert all(np.isfinite(entry.value) for entry in summary) + + def test_subset_sampling_produces_a_single_column(self): + # WHEN + analysis = build_analysis() + analysis.fit() + analysis.suggest_bounds().apply() + + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + results = analysis.sample_posterior(parameters=['Gaussian width'], **SAMPLE_KWARGS) + + # EXPECT + assert results.draws.shape[1] == 1 + assert analysis.posterior_summary().entries[0].name == 'Gaussian width' + + def test_plots_render(self, sampled_analysis): + # WHEN + import matplotlib.pyplot as plt + + n_parameters = len(sampled_analysis.get_free_parameters()) + trace = sampled_analysis.plot_trace() + corner = sampled_analysis.plot_corner() + predictive = sampled_analysis.plot_posterior_predictive(n_draws=20) + + # EXPECT + assert len(trace.axes) == n_parameters + 1 + assert len(corner.axes) == n_parameters**2 + assert len(predictive.axes) == 1 + plt.close('all') + + def test_posterior_median_is_close_to_the_least_squares_fit(self): + # WHEN + analysis = build_analysis() + analysis.fit() + analysis.suggest_bounds().apply() + fitted = {p.name: float(p.value) for p in analysis.get_free_parameters()} + + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + analysis.sample_posterior(**SAMPLE_KWARGS) + summary = analysis.posterior_summary() + + # EXPECT the two agree within the posterior's own uncertainty, since with flat priors the + # maximum-likelihood point sits inside the bulk of the posterior + for entry in summary: + spread = max(entry.minus, entry.plus) + assert abs(entry.median - fitted[entry.name]) < 5 * spread diff --git a/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py b/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py new file mode 100644 index 000000000..b29a2b54a --- /dev/null +++ b/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py @@ -0,0 +1,486 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +"""Unit tests for Bayesian sampling on Analysis1d, with the EasyScience Sampler mocked out.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock +from unittest.mock import patch + +import numpy as np +import pytest +import scipp as sc +from easyscience.fitting import AvailableMinimizers + +from easydynamics.analysis.analysis1d import Analysis1d +from easydynamics.experiment import Experiment +from easydynamics.sample_model import InstrumentModel +from easydynamics.sample_model import SampleModel +from easydynamics.sample_model.components.gaussian import Gaussian + +SAMPLER_PATH = 'easydynamics.analysis.bayesian_sampling.Sampler' + + +def make_analysis(): + energy_values = np.linspace(-5.0, 5.0, 20) + intensity = 3.0 * np.exp(-0.5 * (energy_values / 1.2) ** 2) + data = sc.array( + dims=['Q', 'energy'], + values=intensity[None, :], + variances=np.full_like(intensity, 0.01)[None, :], + ) + experiment = Experiment( + data=sc.DataArray( + data=data, + coords={ + 'Q': sc.array(dims=['Q'], values=[1.0], unit='1/Angstrom'), + 'energy': sc.array(dims=['energy'], values=energy_values, unit='meV'), + }, + ) + ) + analysis = Analysis1d( + display_name='TestBayesian', + experiment=experiment, + sample_model=SampleModel(components=Gaussian(area=3.0, width=1.2, center=0.0)), + instrument_model=InstrumentModel(), + Q_index=0, + ) + analysis.instrument_model.fix_energy_offset(Q_index=0) + return analysis + + +def bound_all(analysis, half_width=5.0): + """Give every free parameter finite bounds so the pre-flight passes.""" + for parameter in analysis.get_free_parameters(): + parameter.min = float(parameter.value) - half_width + parameter.max = float(parameter.value) + half_width + + +def fake_results(analysis, n_draws=100, values=None): + """Build a SamplingResults-shaped object for the free parameters of an analysis.""" + parameters = analysis.get_free_parameters() + if values is None: + draws = np.tile([float(p.value) for p in parameters], (n_draws, 1)) + else: + draws = np.asarray(values, dtype=float) + return SimpleNamespace( + draws=draws, + param_names=[p.unique_name for p in parameters], + logp=np.zeros(draws.shape[0]), + state=MagicMock(Ngen=10, Npop=4), + ) + + +@pytest.fixture +def analysis(): + return make_analysis() + + +class TestFitterExposure: + def test_fitter_is_built_lazily_and_cached(self, analysis): + # WHEN + fitter = analysis.fitter + + # EXPECT + assert fitter is analysis.fitter + assert fitter.fit_object is analysis + + def test_fitter_is_rebuilt_when_the_sample_model_changes(self, analysis): + # WHEN + original = analysis.fitter + analysis.sample_model = SampleModel(components=Gaussian(area=1.0)) + + # EXPECT + assert analysis.fitter is not original + + def test_minimizer_can_be_switched_through_the_fitter(self, analysis): + # WHEN + analysis.fitter.switch_minimizer(AvailableMinimizers.Bumps) + + # EXPECT + assert analysis.fitter.minimizer.enum == AvailableMinimizers.Bumps + + def test_fit_uses_the_persistent_fitter(self, analysis): + # WHEN + result = analysis.fit() + + # EXPECT + assert result is analysis._fit_result + assert np.isfinite(result.reduced_chi2) + + +class TestBoundsPreflight: + def test_sampling_refuses_unbounded_parameters(self, analysis): + # EXPECT + with pytest.raises(ValueError, match='finite bounds'): + analysis.sample_posterior(samples=10) + + def test_error_names_the_offending_parameters(self, analysis): + # EXPECT + with pytest.raises(ValueError, match='Gaussian area'): + analysis.check_bounds_for_sampling() + + def test_bounded_parameters_pass(self, analysis): + # WHEN + bound_all(analysis) + + # EXPECT: does not raise + analysis.check_bounds_for_sampling() + + def test_suggest_bounds_covers_the_free_parameters(self, analysis): + # WHEN + suggestions = analysis.suggest_bounds() + + # EXPECT + assert len(suggestions) == len(analysis.get_free_parameters()) + + +class TestSamplePosterior: + def test_restores_parameter_values_and_minimizer(self, analysis): + # WHEN + bound_all(analysis) + before = [(p.unique_name, p.value) for p in analysis.get_free_parameters()] + + with patch(SAMPLER_PATH) as sampler_class: + + def mutate_then_return(**_kwargs): + # The real sampler leaves the parameters wherever the last evaluation put them. + for parameter in analysis.get_free_parameters(): + parameter.value = float(parameter.value) + 1.0 + return fake_results(analysis) + + sampler_class.return_value.sample.side_effect = mutate_then_return + analysis.sample_posterior(samples=10, burn=1, thin=1) + + # EXPECT + after = [(p.unique_name, p.value) for p in analysis.get_free_parameters()] + assert after == before + assert analysis.fitter.minimizer.enum == AvailableMinimizers.LMFit_leastsq + + def test_switches_to_bumps_for_the_run(self, analysis): + # WHEN + bound_all(analysis) + seen = [] + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: ( + seen.append(analysis.fitter.minimizer.enum), + fake_results(analysis), + )[1] + analysis.sample_posterior(samples=10) + + # EXPECT + assert seen == [AvailableMinimizers.Bumps] + + def test_restores_the_minimizer_even_when_sampling_raises(self, analysis): + # WHEN + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = RuntimeError('boom') + with pytest.raises(RuntimeError, match='boom'): + analysis.sample_posterior(samples=10) + + # EXPECT + assert analysis.fitter.minimizer.enum == AvailableMinimizers.LMFit_leastsq + + def test_forwards_sampling_arguments(self, analysis): + # WHEN + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + analysis.sample_posterior(samples=123, burn=7, thin=3, population=5) + + # EXPECT + kwargs = sampler_class.return_value.sample.call_args.kwargs + assert kwargs['samples'] == 123 + assert kwargs['burn'] == 7 + assert kwargs['thin'] == 3 + assert kwargs['population'] == 5 + + def test_stores_the_result(self, analysis): + # WHEN + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + expected = fake_results(analysis) + sampler_class.return_value.sample.return_value = expected + returned = analysis.sample_posterior(samples=10) + + # EXPECT + assert returned is expected + assert analysis.posterior_result is expected + + def test_warns_when_the_posterior_piles_up_against_a_bound(self, analysis): + # WHEN a parameter's draws span its whole allowed range + bound_all(analysis) + parameters = analysis.get_free_parameters() + draws = np.tile([float(p.value) for p in parameters], (500, 1)) + draws[:, 0] = np.linspace(parameters[0].min, parameters[0].max, 500) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_results(analysis, values=draws) + + # EXPECT + with pytest.warns(UserWarning, match='piled up'): + analysis.sample_posterior(samples=10) + + def test_does_not_warn_when_the_posterior_is_well_inside(self, analysis): + # WHEN + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_results(analysis) + + # EXPECT + with warnings_as_errors(): + analysis.sample_posterior(samples=10) + + +class TestParameterSubset: + def test_holds_other_parameters_fixed_during_the_run(self, analysis): + # WHEN + bound_all(analysis) + target = analysis.get_free_parameters()[0] + seen = {} + + with patch(SAMPLER_PATH) as sampler_class: + + def record(**_kwargs): + seen['free'] = [p.unique_name for p in analysis.get_free_parameters()] + return fake_results(analysis) + + sampler_class.return_value.sample.side_effect = record + with pytest.warns(UserWarning, match='Holding these parameters fixed'): + analysis.sample_posterior(samples=10, parameters=[target.name]) + + # EXPECT + assert seen['free'] == [target.unique_name] + + def test_restores_the_fixed_flags_afterwards(self, analysis): + # WHEN + bound_all(analysis) + before = [(p.unique_name, p.fixed) for p in analysis.get_all_parameters()] + target = analysis.get_free_parameters()[0] + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + with pytest.warns(UserWarning): + analysis.sample_posterior(samples=10, parameters=[target]) + + # EXPECT + assert [(p.unique_name, p.fixed) for p in analysis.get_all_parameters()] == before + + def test_unknown_parameter_name_raises(self, analysis): + # WHEN + bound_all(analysis) + + # EXPECT + with pytest.raises(ValueError, match='No free parameter named'): + analysis.sample_posterior(samples=10, parameters=['not a parameter']) + + def test_non_list_parameters_raises(self, analysis): + # EXPECT + with pytest.raises(TypeError, match='must be a list'): + analysis.sample_posterior(samples=10, parameters='Gaussian area') + + def test_empty_parameter_list_raises(self, analysis): + # EXPECT + with pytest.raises(ValueError, match='at least one parameter'): + analysis.sample_posterior(samples=10, parameters=[]) + + +class TestSamplerCaching: + def test_sampler_is_reused_between_runs(self, analysis): + # WHEN + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + analysis.sample_posterior(samples=10) + analysis.sample_posterior(samples=10) + + # EXPECT the data is bound once, not per run + assert sampler_class.call_count == 1 + + def test_changing_the_q_index_rebuilds_the_sampler(self, analysis): + # WHEN + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + analysis.sample_posterior(samples=10) + analysis.Q_index = 0 + analysis.sample_posterior(samples=10) + + # EXPECT the Sampler binds its data at construction, so it must be rebuilt + assert sampler_class.call_count == 2 + + def test_binds_the_same_data_the_fit_uses(self, analysis): + # WHEN + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + analysis.sample_posterior(samples=10) + + # EXPECT + expected_x, expected_y, expected_w = analysis._get_sampling_data() + args, kwargs = sampler_class.call_args + assert np.array_equal(args[1], expected_x) + assert np.array_equal(args[2], expected_y) + assert np.array_equal(kwargs['weights'], expected_w) + + +class TestExtendAndPersistence: + def test_extend_without_a_chain_raises(self, analysis): + # EXPECT + with pytest.raises(RuntimeError, match='No chain to extend'): + analysis.extend_sampling() + + def test_extend_delegates_to_the_sampler(self, analysis): + # WHEN + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + sampler_class.return_value.extend.side_effect = lambda **_k: fake_results(analysis) + analysis.sample_posterior(samples=10) + analysis.extend_sampling(additional_samples=42, thin=2) + + # EXPECT + kwargs = sampler_class.return_value.extend.call_args.kwargs + assert kwargs['additional_samples'] == 42 + assert kwargs['thin'] == 2 + + def test_save_without_a_chain_raises(self, analysis): + # EXPECT + with pytest.raises(RuntimeError, match='No chain to save'): + analysis.save_chain('somewhere') + + def test_save_writes_the_parameter_name_sidecar(self, analysis, tmp_path): + # WHEN + import json + + bound_all(analysis) + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + analysis.sample_posterior(samples=10) + analysis.save_chain(str(tmp_path / 'chain')) + + # EXPECT the unique names are recorded against the stable parameter names + sidecar = tmp_path / 'chain.parameter-names.json' + assert sidecar.is_file() + mapping = json.loads(sidecar.read_text(encoding='utf-8')) + assert set(mapping.values()) == {p.name for p in analysis.get_free_parameters()} + + def test_load_without_a_sidecar_warns(self, analysis, tmp_path): + # WHEN + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.load_state.return_value = fake_results(analysis) + + # EXPECT + with pytest.warns(UserWarning, match='No parameter-name sidecar'): + analysis.load_chain(str(tmp_path / 'missing')) + + +class TestResults: + def test_summary_without_sampling_raises(self, analysis): + # EXPECT + with pytest.raises(RuntimeError, match='No posterior samples yet'): + analysis.posterior_summary() + + def test_summary_uses_parameter_names_and_units(self, analysis): + # WHEN + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + analysis.sample_posterior(samples=10) + + # EXPECT + summary = analysis.posterior_summary() + names = {entry.name for entry in summary} + assert names == {p.name for p in analysis.get_free_parameters()} + assert all(entry.unit == 'meV' for entry in summary) + + def test_set_parameters_to_posterior_median(self, analysis): + # WHEN + bound_all(analysis) + parameters = analysis.get_free_parameters() + draws = np.tile([float(p.value) + 2.0 for p in parameters], (50, 1)) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_results(analysis, values=draws) + expected = [float(p.value) + 2.0 for p in parameters] + analysis.sample_posterior(samples=10) + + changed = analysis.set_parameters_to_posterior_median() + + # EXPECT + assert len(changed) == len(parameters) + assert [float(p.value) for p in parameters] == pytest.approx(expected) + + def test_median_without_sampling_raises(self, analysis): + # EXPECT + with pytest.raises(RuntimeError, match='No posterior samples yet'): + analysis.set_parameters_to_posterior_median() + + +class TestPlots: + def test_predictive_rejects_a_bad_draw_count(self, analysis): + # WHEN + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + analysis.sample_posterior(samples=10) + + # EXPECT + with pytest.raises(ValueError, match='positive integer'): + analysis.plot_posterior_predictive(n_draws=0) + + def test_predictive_restores_parameter_values(self, analysis): + # WHEN + bound_all(analysis) + parameters = analysis.get_free_parameters() + draws = np.tile([float(p.value) + 0.5 for p in parameters], (20, 1)) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_results(analysis, values=draws) + analysis.sample_posterior(samples=10) + + before = [float(p.value) for p in parameters] + analysis.plot_posterior_predictive(n_draws=5) + + # EXPECT + assert [float(p.value) for p in parameters] == pytest.approx(before) + + def test_plots_without_sampling_raise(self, analysis): + # EXPECT + with pytest.raises(RuntimeError): + analysis.plot_trace() + with pytest.raises(RuntimeError): + analysis.plot_corner() + + +class warnings_as_errors: + """Context manager asserting that no UserWarning is emitted inside the block.""" + + def __enter__(self): + import warnings + + self._ctx = warnings.catch_warnings(record=True) + self._caught = self._ctx.__enter__() + warnings.simplefilter('always') + return self + + def __exit__(self, *exc_info): + caught = [w for w in self._caught if issubclass(w.category, UserWarning)] + self._ctx.__exit__(*exc_info) + if exc_info[0] is None: + assert not caught, f'unexpected warnings: {[str(w.message) for w in caught]}' + return False diff --git a/tests/unit/easydynamics/analysis/test_posterior.py b/tests/unit/easydynamics/analysis/test_posterior.py new file mode 100644 index 000000000..88dc29bf6 --- /dev/null +++ b/tests/unit/easydynamics/analysis/test_posterior.py @@ -0,0 +1,296 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +import numpy as np +import pytest +from easyscience.variable import Parameter + +from easydynamics.analysis.posterior import BoundsSuggestion +from easydynamics.analysis.posterior import BoundsSuggestions +from easydynamics.analysis.posterior import parameters_at_bounds +from easydynamics.analysis.posterior import suggest_bounds_for_parameters +from easydynamics.analysis.posterior import summarize_draws +from easydynamics.analysis.posterior import unbounded_parameters + + +def make_parameter(name='p', value=1.0, error=0.0, minimum=-np.inf, maximum=np.inf, unit='meV'): + parameter = Parameter(name=name, value=value, unit=unit) + parameter.min = minimum + parameter.max = maximum + if error: + parameter.variance = error**2 + return parameter + + +class TestSuggestBounds: + def test_fills_in_both_infinite_sides(self): + # WHEN + parameter = make_parameter(value=10.0, error=0.5) + + # THEN + suggestions = suggest_bounds_for_parameters([parameter], n_sigma=10.0, relative_pad=0.2) + + # EXPECT: 10 * 0.5 + 0.2 * 10 = 7 + suggestion = suggestions.suggestions[0] + assert suggestion.suggested_min == pytest.approx(3.0) + assert suggestion.suggested_max == pytest.approx(17.0) + assert not suggestion.needs_attention + + def test_never_loosens_an_existing_finite_bound(self): + # WHEN a physical lower bound is already set + parameter = make_parameter(value=1.2, error=1.5, minimum=1e-10) + + # THEN + suggestion = suggest_bounds_for_parameters([parameter]).suggestions[0] + + # EXPECT the finite side survives untouched, even though the sigma rule would go negative + assert suggestion.suggested_min == pytest.approx(1e-10) + assert suggestion.suggested_max > 1.2 + + def test_fully_bounded_parameter_is_left_alone(self): + # WHEN + parameter = make_parameter(value=1.0, error=0.1, minimum=0.0, maximum=2.0) + + # THEN + suggestion = suggest_bounds_for_parameters([parameter]).suggestions[0] + + # EXPECT + assert suggestion.suggested_min == pytest.approx(0.0) + assert suggestion.suggested_max == pytest.approx(2.0) + assert not suggestion.changes_bounds + + def test_zero_error_falls_back_to_the_relative_pad(self): + # WHEN a minimizer reports no uncertainty at all + parameter = make_parameter(value=4.0, error=0.0) + + # THEN + suggestion = suggest_bounds_for_parameters([parameter], relative_pad=0.25).suggestions[0] + + # EXPECT the pad still yields a usable width + assert suggestion.suggested_min == pytest.approx(3.0) + assert suggestion.suggested_max == pytest.approx(5.0) + assert not suggestion.needs_attention + + def test_zero_value_and_zero_error_is_flagged_not_guessed(self): + # WHEN there is no scale information anywhere + parameter = make_parameter(value=0.0, error=0.0) + + # THEN + suggestion = suggest_bounds_for_parameters([parameter]).suggestions[0] + + # EXPECT + assert suggestion.needs_attention + assert 'no scale information' in suggestion.reason + assert not np.isfinite(suggestion.suggested_min) + + def test_absolute_floor_rescues_a_scaleless_parameter(self): + # WHEN + parameter = make_parameter(value=0.0, error=0.0) + + # THEN + suggestion = suggest_bounds_for_parameters([parameter], absolute_floor=0.5).suggestions[0] + + # EXPECT + assert not suggestion.needs_attention + assert suggestion.suggested_min == pytest.approx(-0.5) + assert suggestion.suggested_max == pytest.approx(0.5) + + def test_non_finite_value_is_flagged(self): + # WHEN + parameter = make_parameter(value=1.0) + parameter.value = np.inf + + # THEN + suggestion = suggest_bounds_for_parameters([parameter]).suggestions[0] + + # EXPECT + assert suggestion.needs_attention + assert 'not finite' in suggestion.reason + + @pytest.mark.parametrize('kwargs', [{'n_sigma': -1.0}, {'relative_pad': -0.1}]) + def test_negative_settings_raise(self, kwargs): + # EXPECT + with pytest.raises(ValueError): + suggest_bounds_for_parameters([make_parameter()], **kwargs) + + def test_non_numeric_setting_raises(self): + # EXPECT + with pytest.raises(TypeError): + suggest_bounds_for_parameters([make_parameter()], n_sigma='wide') + + +class TestBoundsSuggestionsApply: + def test_apply_sets_bounds_and_reports_changes(self): + # WHEN + parameter = make_parameter(value=10.0, error=0.5) + suggestions = suggest_bounds_for_parameters([parameter]) + + # THEN nothing has changed until apply is called + assert parameter.max == np.inf + changed = suggestions.apply() + + # EXPECT + assert changed == [parameter] + assert parameter.min == pytest.approx(3.0) + assert parameter.max == pytest.approx(17.0) + + def test_apply_skips_parameters_needing_attention(self): + # WHEN + parameter = make_parameter(value=0.0, error=0.0) + suggestions = suggest_bounds_for_parameters([parameter]) + + # THEN + changed = suggestions.apply() + + # EXPECT the unusable suggestion is skipped rather than written + assert changed == [] + assert parameter.min == -np.inf + + def test_repr_lists_parameters_and_flags_attention(self): + # WHEN + good = make_parameter(name='good', value=10.0, error=0.5) + bad = make_parameter(name='bad', value=0.0, error=0.0) + + # THEN + text = repr(suggest_bounds_for_parameters([good, bad])) + + # EXPECT + assert 'good' in text + assert 'bad' in text + assert 'need bounds set by hand' in text + + def test_repr_with_no_parameters(self): + # EXPECT + assert 'no free parameters' in repr(BoundsSuggestions([])) + + def test_len_and_iteration(self): + # WHEN + suggestions = suggest_bounds_for_parameters([make_parameter(), make_parameter()]) + + # EXPECT + assert len(suggestions) == 2 + assert all(isinstance(s, BoundsSuggestion) for s in suggestions) + + +class TestUnboundedParameters: + def test_finds_parameters_with_an_infinite_side(self): + # WHEN + bounded = make_parameter(name='bounded', minimum=0.0, maximum=1.0) + half_open = make_parameter(name='half_open', minimum=0.0) + + # THEN + result = unbounded_parameters([bounded, half_open]) + + # EXPECT + assert result == [half_open] + + +class TestParametersAtBounds: + def test_uniform_posterior_across_the_bounds_is_reported(self): + # WHEN a posterior fills its whole allowed range, the bound is setting the interval + parameter = make_parameter(minimum=0.0, maximum=1.0) + draws = np.linspace(0.0, 1.0, 1000).reshape(-1, 1) + + # THEN + result = parameters_at_bounds(draws, [parameter]) + + # EXPECT + assert parameter.name in result + assert result[parameter.name] == pytest.approx(0.1, abs=0.01) + + def test_posterior_well_inside_its_bounds_is_not_reported(self): + # WHEN + parameter = make_parameter(minimum=0.0, maximum=1.0) + draws = np.random.default_rng(0).normal(0.5, 0.02, size=1000).reshape(-1, 1) + + # THEN + result = parameters_at_bounds(draws, [parameter]) + + # EXPECT + assert result == {} + + def test_partly_clipped_posterior_is_reported(self): + # WHEN a posterior fills most, but not all, of its allowed range. A real bound-limited + # chain looks like this rather than perfectly uniform, so the threshold has to catch it. + parameter = make_parameter(minimum=0.0, maximum=1.0) + draws = np.linspace(0.02, 0.98, 1000).reshape(-1, 1) + + # THEN + result = parameters_at_bounds(draws, [parameter]) + + # EXPECT + assert parameter.name in result + + def test_posterior_pinned_at_one_bound_is_reported(self): + # WHEN + parameter = make_parameter(minimum=0.0, maximum=1.0) + draws = np.abs(np.random.default_rng(0).normal(0.0, 0.02, size=1000)).reshape(-1, 1) + + # THEN + result = parameters_at_bounds(draws, [parameter]) + + # EXPECT + assert result[parameter.name] > 0.9 + + def test_unmatched_and_unbounded_columns_are_skipped(self): + # WHEN + unbounded = make_parameter() + draws = np.zeros((10, 2)) + + # THEN + result = parameters_at_bounds(draws, [None, unbounded]) + + # EXPECT + assert result == {} + + +class TestSummarizeDraws: + def test_reports_parameter_names_units_and_percentiles(self): + # WHEN + parameter = make_parameter(name='Gaussian width', value=1.5, unit='meV') + draws = np.linspace(0.0, 100.0, 101).reshape(-1, 1) + + # THEN + summary = summarize_draws(draws, ['Parameter_0'], [parameter]) + + # EXPECT + entry = summary['Gaussian width'] + assert entry.unit == 'meV' + assert entry.median == pytest.approx(50.0) + assert entry.lower == pytest.approx(16.0) + assert entry.upper == pytest.approx(84.0) + assert entry.minus == pytest.approx(34.0) + assert entry.plus == pytest.approx(34.0) + assert entry.value == pytest.approx(1.5) + + def test_unmatched_column_falls_back_to_the_supplied_name(self): + # WHEN + draws = np.zeros((10, 1)) + + # THEN + summary = summarize_draws(draws, ['Parameter_7'], [None]) + + # EXPECT + entry = summary.entries[0] + assert entry.name == 'Parameter_7' + assert entry.unit == '' + assert np.isnan(entry.value) + + def test_lookup_of_missing_name_raises(self): + # WHEN + summary = summarize_draws(np.zeros((5, 1)), ['x'], [None]) + + # EXPECT + with pytest.raises(KeyError): + summary['not a parameter'] + + def test_repr_contains_the_parameter_name(self): + # WHEN + parameter = make_parameter(name='Gaussian area') + + # THEN + text = repr(summarize_draws(np.zeros((5, 1)), ['x'], [parameter])) + + # EXPECT + assert 'Gaussian area' in text + assert 'median' in text diff --git a/tests/unit/easydynamics/utils/test_posterior_plotting.py b/tests/unit/easydynamics/utils/test_posterior_plotting.py new file mode 100644 index 000000000..64d0129e9 --- /dev/null +++ b/tests/unit/easydynamics/utils/test_posterior_plotting.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +import matplotlib as mpl +import numpy as np +import pytest + +mpl.use('Agg') + +import matplotlib.pyplot as plt + +from easydynamics.utils.posterior_plotting import plot_corner +from easydynamics.utils.posterior_plotting import plot_posterior_predictive +from easydynamics.utils.posterior_plotting import plot_trace + + +@pytest.fixture(autouse=True) +def close_figures(): + yield + plt.close('all') + + +@pytest.fixture +def draws(): + return np.random.default_rng(0).normal(size=(200, 3)) + + +class TestPlotTrace: + def test_one_panel_per_parameter(self, draws): + # WHEN + fig = plot_trace(draws=draws, names=['a', 'b', 'c']) + + # EXPECT + assert len(fig.axes) == 3 + + def test_logp_adds_a_panel(self, draws): + # WHEN + fig = plot_trace(draws=draws, names=['a', 'b', 'c'], logp=np.zeros(len(draws))) + + # EXPECT + assert len(fig.axes) == 4 + assert fig.axes[-1].get_ylabel() == 'log-posterior' + + def test_names_label_the_panels(self, draws): + # WHEN + fig = plot_trace(draws=draws, names=['alpha', 'beta', 'gamma']) + + # EXPECT + assert [axis.get_ylabel() for axis in fig.axes] == ['alpha', 'beta', 'gamma'] + + def test_single_parameter_works(self): + # WHEN + fig = plot_trace(draws=np.zeros((10, 1)), names=['only']) + + # EXPECT + assert len(fig.axes) == 1 + + def test_mismatched_names_raise(self, draws): + # EXPECT + with pytest.raises(ValueError, match='one entry per column'): + plot_trace(draws=draws, names=['a', 'b']) + + def test_one_dimensional_draws_raise(self): + # EXPECT + with pytest.raises(ValueError, match='two-dimensional'): + plot_trace(draws=np.zeros(10), names=['a']) + + +class TestPlotCorner: + def test_grid_is_square_in_the_parameter_count(self, draws): + # WHEN + fig = plot_corner(draws=draws, names=['a', 'b', 'c']) + + # EXPECT + assert len(fig.axes) == 9 + + def test_upper_triangle_is_hidden(self, draws): + # WHEN + fig = plot_corner(draws=draws, names=['a', 'b', 'c']) + + # EXPECT: 3 hidden panels above the diagonal of a 3x3 grid + assert sum(not axis.get_visible() for axis in fig.axes) == 3 + + def test_mismatched_names_raise(self, draws): + # EXPECT + with pytest.raises(ValueError, match='one entry per column'): + plot_corner(draws=draws, names=['a']) + + +class TestPlotPosteriorPredictive: + def test_returns_a_figure_with_data_and_band(self): + # WHEN + x = np.linspace(0.0, 1.0, 25) + predictions = np.random.default_rng(0).normal(size=(50, 25)) + + fig = plot_posterior_predictive(x=x, y=np.zeros(25), predictions=predictions) + + # EXPECT + labels = [text.get_text() for text in fig.axes[0].get_legend().get_texts()] + assert 'Data' in labels + assert any('credible band' in label for label in labels) + + def test_error_bars_are_drawn_when_given(self): + # WHEN + x = np.linspace(0.0, 1.0, 10) + + fig = plot_posterior_predictive( + x=x, + y=np.zeros(10), + predictions=np.zeros((5, 10)), + y_err=np.full(10, 0.1), + ) + + # EXPECT + assert len(fig.axes[0].containers) == 1 + + def test_wrong_prediction_shape_raises(self): + # EXPECT + with pytest.raises(ValueError, match='predictions must have shape'): + plot_posterior_predictive(x=np.zeros(10), y=np.zeros(10), predictions=np.zeros((5, 3))) + + @pytest.mark.parametrize('interval', [0.0, 100.0, -5.0]) + def test_invalid_credible_interval_raises(self, interval): + # EXPECT + with pytest.raises(ValueError, match='credible_interval'): + plot_posterior_predictive( + x=np.zeros(4), + y=np.zeros(4), + predictions=np.zeros((5, 4)), + credible_interval=interval, + ) + + def test_band_widens_with_the_credible_interval(self): + # WHEN + x = np.linspace(0.0, 1.0, 8) + predictions = np.random.default_rng(0).normal(size=(400, 8)) + + narrow = plot_posterior_predictive( + x=x, y=np.zeros(8), predictions=predictions, credible_interval=50.0 + ) + wide = plot_posterior_predictive( + x=x, y=np.zeros(8), predictions=predictions, credible_interval=95.0 + ) + + # EXPECT + narrow_span = narrow.axes[0].collections[0].get_paths()[0].get_extents().height + wide_span = wide.axes[0].collections[0].get_paths()[0].get_extents().height + assert wide_span > narrow_span From 7db52f9876eec66882690768c913f2822574447e Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Thu, 13 Aug 2026 12:28:44 +0200 Subject: [PATCH 02/29] Add Bayesian posterior sampling to Analysis and ParameterAnalysis Extends the sampling introduced for Analysis1d to the remaining two Analysis classes, using the mixin hooks added with it. No new sampling machinery: each class supplies its fitter, its data, and its chain parameters, and everything else is shared. Analysis gains sample_posterior(fit_method=...), mirroring fit(): - 'independent' gives each Q index its own chain, delegating to the Analysis1d objects, and returns one result per Q (or a single result when a Q_index is given). - 'simultaneous' runs one chain over every Q at once through a MultiFitter, refreshing each per-Q convolver against its masked energy grid first, exactly as the simultaneous fit does. ParameterAnalysis samples the binding models. Its fit() built the MultiFitter inline, so the per-target data, functions, and models are now resolved by a shared _build_fit_inputs() that both paths use, which also guarantees fitting and sampling see the same targets in the same order with the same unit conversions. Parameter labels needed rethinking. A multi-Q analysis holds one copy of each parameter per Q, all sharing a name, so a summary showed several identical rows and a name could not pick a parameter out. Labels are now produced by an overridable parameter_label(): Analysis qualifies by Q index, ParameterAnalysis by binding model, and both only when the bare name is actually ambiguous, so single-Q and single-binding cases keep their short names. The summary and bounds tables size themselves to the longest label rather than truncating. Also fixes Analysis.fit's docstring, which promised a single FitResults for a simultaneous fit. MultiFitter splits its combined result back up by dataset, so a list has always been returned. Tutorial 1 gains a Bayesian section on the two-step diffusion fit, where the posterior turns out to be about twelve times tighter than the reported least-squares uncertainties. That gap is real and worth explaining: the width fit has a reduced chi-squared near 150, so lmfit inflates its uncertainties by the square root of that, while the sampler takes the stated uncertainties at face value. Sampling the full simultaneous diffusion model was measured at over ten minutes, so the tutorial uses the ParameterAnalysis step instead. Co-Authored-By: Claude Opus 5 (1M context) --- docs/docs/tutorials/tutorial1_brownian.ipynb | 65 +++++ src/easydynamics/analysis/analysis.py | 188 ++++++++++++- .../analysis/bayesian_sampling.py | 63 ++++- .../analysis/parameter_analysis.py | 125 ++++++++- src/easydynamics/analysis/posterior.py | 46 +++- .../fitting/test_bayesian_sampling_multi_q.py | 216 +++++++++++++++ .../analysis/test_analysis_bayesian.py | 249 ++++++++++++++++++ .../test_parameter_analysis_bayesian.py | 206 +++++++++++++++ .../easydynamics/analysis/test_posterior.py | 16 +- 9 files changed, 1134 insertions(+), 40 deletions(-) create mode 100644 tests/integration/fitting/test_bayesian_sampling_multi_q.py create mode 100644 tests/unit/easydynamics/analysis/test_analysis_bayesian.py create mode 100644 tests/unit/easydynamics/analysis/test_parameter_analysis_bayesian.py diff --git a/docs/docs/tutorials/tutorial1_brownian.ipynb b/docs/docs/tutorials/tutorial1_brownian.ipynb index bb6403252..a3b17989c 100644 --- a/docs/docs/tutorials/tutorial1_brownian.ipynb +++ b/docs/docs/tutorials/tutorial1_brownian.ipynb @@ -531,6 +531,71 @@ "parameter_analysis.get_all_parameters()" ] }, + { + "cell_type": "markdown", + "id": "163c27bb", + "metadata": {}, + "source": [ + "### How certain are the diffusion parameters?\n", + "\n", + "The uncertainties printed above come from the curvature of $\\chi^2$ at the best fit. That is a good estimate when the parameters are uncorrelated and their uncertainties are roughly Gaussian, but $D$ and the scale are fitted to the same curve and need not be either. A Bayesian analysis maps the full posterior instead, so we can check.\n", + "\n", + "The bounds act as the prior, so every free parameter needs finite ones first. `suggest_bounds()` proposes them from the fit and is advisory until `.apply()` is called." + ] + }, + { + "cell_type": "code", + "id": "de797763", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "suggestions = parameter_analysis.suggest_bounds()\n", + "print(suggestions)\n", + "suggestions.apply()" + ] + }, + { + "cell_type": "code", + "id": "604cd4e9", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "parameter_analysis.sample_posterior(samples=4000, burn=200, thin=2)\n", + "parameter_analysis.posterior_summary()" + ] + }, + { + "cell_type": "markdown", + "id": "2e0cdab0", + "metadata": {}, + "source": [ + "The corner plot shows how the two parameters trade off against each other. A tilted, narrow ridge means the data pins down a combination of $D$ and the scale more tightly than either one separately." + ] + }, + { + "cell_type": "code", + "id": "6208979b", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "parameter_analysis.plot_corner()" + ] + }, + { + "cell_type": "markdown", + "id": "154d6137", + "metadata": {}, + "source": [ + "Notice that these credible intervals are **much narrower** than the uncertainties printed further up, and that difference is worth understanding rather than trusting.\n", + "\n", + "The least-squares fit of the widths has a reduced $\\chi^2$ of about 150: the Brownian model does not describe the fitted widths to within their error bars. `lmfit` responds by inflating its reported uncertainties by the square root of that, roughly a factor of 12, on the assumption that a poor fit means the input uncertainties were understated. The sampler makes no such adjustment — it takes the stated uncertainties at face value — so its intervals come out around twelve times tighter.\n", + "\n", + "Neither is simply right. The gap is a signal that the two-step model is not capturing the data, which is exactly what we address next by fitting the diffusion model to all the data at once." + ] + }, { "cell_type": "markdown", "id": "fc2f8434", diff --git a/src/easydynamics/analysis/analysis.py b/src/easydynamics/analysis/analysis.py index 8fb3d703d..4edc080cf 100644 --- a/src/easydynamics/analysis/analysis.py +++ b/src/easydynamics/analysis/analysis.py @@ -8,12 +8,14 @@ import scipp as sc from easyscience.fitting.minimizers.utils import FitResults from easyscience.fitting.multi_fitter import MultiFitter +from easyscience.fitting.sampler import SamplingResults from easyscience.variable import Parameter from plopp.backends.matplotlib.figure import InteractiveFigure from scipp import UnitError from easydynamics.analysis.analysis1d import Analysis1d from easydynamics.analysis.analysis_base import AnalysisBase +from easydynamics.analysis.bayesian_sampling import BayesianSamplingMixin from easydynamics.experiment import Experiment from easydynamics.sample_model import SampleModel from easydynamics.sample_model.instrument_model import InstrumentModel @@ -24,12 +26,16 @@ from easydynamics.utils.utils import verify_Q_index -class Analysis(AnalysisBase): +class Analysis(BayesianSamplingMixin, AnalysisBase): """ For analysing two-dimensional data, i.e. intensity as function of energy and Q. Supports independent fits of each Q value and simultaneous fits of all Q. + Besides least-squares fitting with :meth:`fit`, the posterior distribution of the free + parameters can be explored with :meth:`sample_posterior`; see + :class:`~easydynamics.analysis.bayesian_sampling.BayesianSamplingMixin`. + Examples -------- **Fitting vanadium data for instrument calibration** @@ -117,6 +123,7 @@ def __init__( self._analysis_list: list[Analysis1d] = [] self._analysis_list_is_dirty = True + self._init_bayesian_state() super().__init__( display_name=display_name, unique_name=unique_name, @@ -283,8 +290,9 @@ def fit( Returns ------- FitResults | list[FitResults] - A list of FitResults if fitting independently, or a single FitResults object if fitting - simultaneously. + A single FitResults when a specific Q index was fitted, and otherwise a list holding + one FitResults per Q index. A simultaneous fit also reports per-Q results, since the + underlying MultiFitter splits its combined result back up by dataset. """ if self.Q is None: @@ -302,6 +310,77 @@ def fit( return self._fit_all_Q_simultaneously() raise ValueError("Invalid fit method. Choose 'independent' or 'simultaneous'.") + def sample_posterior( + self, + samples: int = 10000, + burn: int = 2000, + thin: int = 10, + fit_method: str = 'independent', + Q_index: int | None = None, + **sampler_options: dict[str, Any], + ) -> SamplingResults | list[SamplingResults]: + """ + Draw samples from the posterior distribution, mirroring :meth:`fit`. + + With ``fit_method='independent'`` each Q index gets its own chain, which is the cheaper + option and keeps the Q values from influencing one another. With + ``fit_method='simultaneous'`` a single chain covers every Q at once, which is what you want + when parameters are shared across Q, but costs considerably more: DREAM runs + ``ceil(population * n_parameters)`` chains, and a simultaneous run has the parameters of + every Q in play at the same time. + + Parameters + ---------- + samples : int, default=10000 + Number of raw samples to draw across all chains, before thinning. + burn : int, default=2000 + Burn-in generations to discard before collecting samples. + thin : int, default=10 + Thinning interval, which reduces autocorrelation between retained draws. + fit_method : str, default='independent' + Either "independent" (a separate chain per Q index) or "simultaneous" (one chain over + all Q indices at once). + Q_index : int | None, default=None + With ``fit_method='independent'``, sample only this Q index. Ignored when sampling + simultaneously. + **sampler_options : dict[str, Any] + Forwarded to the sampler, e.g. ``population``, ``parameters``, ``sampler_kwargs``, + ``progress_callback``, or ``abort_test``. + + Returns + ------- + SamplingResults | list[SamplingResults] + A single SamplingResults when a specific Q index was sampled or when sampling + simultaneously, and otherwise a list holding one SamplingResults per Q index. + + Raises + ------ + ValueError + If there are no Q values available, or if fit_method is not "independent" or + "simultaneous". + """ + if self.Q is None: + raise ValueError( + 'No Q values available for sampling. Please check the experiment data.' + ) + + verify_Q_index(Q_index=Q_index, Q=self.Q, allow_none=True) + + if fit_method == 'independent': + if Q_index is not None: + return self.analysis_list[Q_index].sample_posterior( + samples=samples, burn=burn, thin=thin, **sampler_options + ) + return [ + analysis.sample_posterior(samples=samples, burn=burn, thin=thin, **sampler_options) + for analysis in self.analysis_list + ] + if fit_method == 'simultaneous': + return super().sample_posterior( + samples=samples, burn=burn, thin=thin, **sampler_options + ) + raise ValueError("Invalid fit method. Choose 'independent' or 'simultaneous'.") + def plot_data_and_model( self, Q_index: int | None = None, @@ -661,6 +740,7 @@ def _on_experiment_changed(self) -> None: """ super()._on_experiment_changed() self._analysis_list_is_dirty = True + self._invalidate_fitter() def _on_sample_model_changed(self) -> None: """ @@ -668,6 +748,7 @@ def _on_sample_model_changed(self) -> None: """ super()._on_sample_model_changed() self._analysis_list_is_dirty = True + self._invalidate_fitter() def _on_instrument_model_changed(self) -> None: """ @@ -675,6 +756,7 @@ def _on_instrument_model_changed(self) -> None: """ super()._on_instrument_model_changed() self._analysis_list_is_dirty = True + self._invalidate_fitter() def _on_convolution_settings_changed(self) -> None: """ @@ -682,6 +764,7 @@ def _on_convolution_settings_changed(self) -> None: """ super()._on_convolution_settings_changed() self._analysis_list_is_dirty = True + self._invalidate_fitter() def _ensure_analysis_list_current(self) -> None: """Rebuild the analysis list if any dependency has changed since it was last built.""" @@ -714,6 +797,105 @@ def _create_analysis_list(self) -> None: # Private methods ############# + ############# + # Hooks for BayesianSamplingMixin (simultaneous sampling over all Q) + ############# + + def _build_bayesian_fitter(self) -> MultiFitter: + """ + Build the MultiFitter covering every Q index. + + Returns + ------- + MultiFitter + A MultiFitter over the Analysis1d objects and their fit functions. + """ + return MultiFitter( + fit_objects=self.analysis_list, + fit_functions=self.get_fit_functions(), + ) + + def _get_sampling_data(self) -> tuple[list, list, list]: + """ + Get the per-Q data to bind to the Sampler, as lists of arrays. + + Returns + ------- + tuple[list, list, list] + The ``(x, y, weights)`` triple, one entry per Q index. + """ + xs, ys, ws = [], [], [] + for analysis1d in self.analysis_list: + x, y, weight, _ = self.experiment.extract_x_y_weights_only_finite(analysis1d.Q_index) + xs.append(x) + ys.append(y) + ws.append(weight) + return xs, ys, ws + + def _get_chain_parameters(self) -> list[Parameter]: + """ + Get the free parameters across every Q index. + + Each Q index holds its own copy of the model parameters, so the union is taken by + ``unique_name``. Parameters shared between Q indices therefore appear only once. + + Returns + ------- + list[Parameter] + The free parameters of the whole analysis, in Q order and without duplicates. + """ + parameters = {} + for analysis1d in self.analysis_list: + for parameter in analysis1d.get_free_parameters(): + parameters.setdefault(parameter.unique_name, parameter) + return list(parameters.values()) + + def parameter_label(self, parameter: Parameter) -> str: + """ + Label a parameter with the Q index it belongs to. + + Every Q index carries its own copy of each model parameter, all sharing a name, so the bare + name would produce several identical rows in a summary and could not be used to pick a + parameter out. Parameters shared across Q indices, and analyses holding a single Q, keep + their plain name. + + Parameters + ---------- + parameter : Parameter + The parameter to label. + + Returns + ------- + str + The parameter name, qualified by Q index where that is needed to tell copies apart. + """ + if not self._name_is_ambiguous(parameter): + return parameter.name + owners = [ + analysis1d.Q_index + for analysis1d in self.analysis_list + if any( + p.unique_name == parameter.unique_name for p in analysis1d.get_free_parameters() + ) + ] + if len(owners) != 1: + return parameter.name + return f'{parameter.name} (Q_index={owners[0]})' + + def _prepare_for_sampling(self) -> None: + """ + Rebuild every per-Q convolver against its masked energy grid. + + Mirrors what a simultaneous fit does, so that the model evaluations seen by the sampler + match the ones the fit would have made. + """ + for analysis1d in self.analysis_list: + _, _, _, mask = self.experiment.extract_x_y_weights_only_finite(analysis1d.Q_index) + mask_var = sc.array(dims=['energy'], values=mask) + analysis1d.refresh_convolver( + energy=self.experiment.get_masked_energy(Q_index=analysis1d.Q_index, mask=mask_var) + ) + def _fit_single_Q(self, Q_index: int) -> FitResults: """ Fit data for a single Q index. diff --git a/src/easydynamics/analysis/bayesian_sampling.py b/src/easydynamics/analysis/bayesian_sampling.py index 0ea38d8ef..28dfb43b4 100644 --- a/src/easydynamics/analysis/bayesian_sampling.py +++ b/src/easydynamics/analysis/bayesian_sampling.py @@ -256,8 +256,10 @@ def suggest_bounds( BoundsSuggestions The proposed bounds, which must be applied explicitly. """ + parameters = self._get_chain_parameters() return suggest_bounds_for_parameters( - self._get_chain_parameters(), + parameters, + labels=[self.parameter_label(parameter) for parameter in parameters], n_sigma=n_sigma, relative_pad=relative_pad, absolute_floor=absolute_floor, @@ -275,7 +277,7 @@ def check_bounds_for_sampling(self) -> None: unbounded = unbounded_parameters(self._get_chain_parameters()) if not unbounded: return - names = ', '.join(parameter.name for parameter in unbounded) + names = ', '.join(self.parameter_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}. ' @@ -453,7 +455,8 @@ def _run_sampling( # A fresh chain is labelled with this session's unique names, so any mapping left over # from a loaded chain no longer applies. self._chain_name_map = { - parameter.unique_name: parameter.name for parameter in chain_parameters + parameter.unique_name: self.parameter_label(parameter) + for parameter in chain_parameters } self._posterior_result = results @@ -610,7 +613,7 @@ def posterior_summary(self) -> PosteriorSummary: results = self._require_posterior_result() return summarize_draws( draws=results.draws, - fallback_names=self._chain_display_names(results), + labels=self._chain_display_names(results), parameters_by_column=self._resolve_chain_parameters(results), ) @@ -922,19 +925,59 @@ def _resolve_chain_parameters(self, results: SamplingResults) -> list[Parameter """ parameters = self._get_chain_parameters() by_unique_name = {p.unique_name: p for p in parameters} - by_name = {p.name: p for p in parameters} + by_label = {self.parameter_label(p): p for p in parameters} resolved = [] for unique_name in results.param_names: parameter = by_unique_name.get(unique_name) if parameter is None: - saved_name = self._chain_name_map.get(unique_name) - parameter = None if saved_name is None else by_name.get(saved_name) + saved_label = self._chain_name_map.get(unique_name) + parameter = None if saved_label is None else by_label.get(saved_label) resolved.append(parameter) return resolved + def parameter_label(self, parameter: Parameter) -> str: + """ + Get the label a parameter is reported under. + + The parameter's own name is enough when it identifies the parameter uniquely. Analysis + classes that can hold several identically named parameters override this to qualify the + name, but only when it is actually ambiguous -- see :meth:`_name_is_ambiguous`. + + Parameters + ---------- + parameter : Parameter + The parameter to label. + + Returns + ------- + str + The label to report the parameter under. + """ + return parameter.name + + def _name_is_ambiguous(self, parameter: Parameter) -> bool: + """ + Check whether another parameter in the chain shares this parameter's name. + + Qualifying a label is only worth the extra width when the bare name would be ambiguous, so + a single-Q analysis, or a single binding, keeps its short names. + + Parameters + ---------- + parameter : Parameter + The parameter to check. + + Returns + ------- + bool + True when at least one other chain parameter has the same name. + """ + names = [p.name for p in self._get_chain_parameters()] + return names.count(parameter.name) > 1 + def _chain_display_names(self, results: SamplingResults) -> list[str]: """ - Translate the chain's column names into parameter names. + Translate the chain's column names into readable labels. Parameters ---------- @@ -944,13 +987,13 @@ def _chain_display_names(self, results: SamplingResults) -> list[str]: Returns ------- list[str] - One name per column of the chain. + One label per column of the chain. """ resolved = self._resolve_chain_parameters(results) return [ self._chain_name_map.get(unique_name, unique_name) if parameter is None - else parameter.name + else self.parameter_label(parameter) for unique_name, parameter in zip(results.param_names, resolved, strict=True) ] diff --git a/src/easydynamics/analysis/parameter_analysis.py b/src/easydynamics/analysis/parameter_analysis.py index 7e24108e1..342f33c69 100644 --- a/src/easydynamics/analysis/parameter_analysis.py +++ b/src/easydynamics/analysis/parameter_analysis.py @@ -9,10 +9,12 @@ import scipp as sc from easyscience.fitting.minimizers.utils import FitResults from easyscience.fitting.multi_fitter import MultiFitter +from easyscience.variable import Parameter from matplotlib import rcParams from plopp.backends.matplotlib.figure import InteractiveFigure from easydynamics.analysis.analysis import Analysis +from easydynamics.analysis.bayesian_sampling import BayesianSamplingMixin from easydynamics.analysis.fit_binding import FitBinding from easydynamics.base_classes.easydynamics_modelbase import EasyDynamicsModelBase from easydynamics.utils.fit_target import FitTarget @@ -20,7 +22,7 @@ from easydynamics.utils.utils import convert_value_unit -class ParameterAnalysis(EasyDynamicsModelBase): +class ParameterAnalysis(BayesianSamplingMixin, EasyDynamicsModelBase): """ For analysing fitted parameters. @@ -98,6 +100,8 @@ def __init__( default, None. """ + self._init_bayesian_state() + super().__init__(display_name=display_name, unique_name=unique_name) self._parameters = self._verify_parameters(parameters) @@ -130,6 +134,7 @@ def parameters(self, value: sc.Dataset | Analysis | None) -> None: The new parameter dataset for the parameter analysis. """ self._parameters = self._verify_parameters(value) + self._invalidate_fitter() @property def bindings(self) -> list[FitBinding]: @@ -154,6 +159,7 @@ def bindings(self, value: FitBinding | list[FitBinding] | None) -> None: The new fit bindings for the parameter analysis. """ self._bindings = self._verify_bindings(value) + self._invalidate_fitter() ############# # Other methods @@ -163,18 +169,36 @@ def fit(self) -> FitResults: """ Fit the parameters using the specified fit functions and settings. + A ``ValueError`` is raised if no parameters Dataset is provided, if no fit bindings are + provided, or if a binding names a dataset key that is not in the parameters Dataset. + Returns ------- FitResults The results of the fit + """ + + xs, ys, ws, _, _ = self._build_fit_inputs() + return self.fitter.fit(x=xs, y=ys, weights=ws) + + def _build_fit_inputs(self) -> tuple[list, list, list, list, list]: + """ + Resolve every binding into the per-target data, fit functions, and models. + + Shared by fitting and sampling so that both see exactly the same targets, in the same + order, with the same unit conversions applied. + + Returns + ------- + tuple[list, list, list, list, list] + The ``(x, y, weights, functions, models)`` lists, one entry per fit target. Raises ------ ValueError - If no parameters Dataset is provided. If no fit functions are provided. If no parameter - names are found for the fit functions. + If no parameters Dataset is provided, if no fit bindings are provided, or if a binding + names a dataset key that is not in the parameters Dataset. """ - if self.parameters is None: raise ValueError('No parameters Dataset provided.') @@ -207,16 +231,91 @@ def fit(self) -> FitResults: funcs.append(target.function) models.append(binding.model) - mf = MultiFitter( - fit_objects=models, - fit_functions=funcs, - ) + return xs, ys, ws, funcs, models - return mf.fit( - x=xs, - y=ys, - weights=ws, - ) + ############# + # Hooks for BayesianSamplingMixin + ############# + + def _build_bayesian_fitter(self) -> MultiFitter: + """ + Build the MultiFitter over the binding models. + + Unlike the other Analysis classes, the objects being fitted are the binding models rather + than this object, so the parameters live on those models. + + Returns + ------- + MultiFitter + A MultiFitter over the per-target models and fit functions. + """ + _, _, _, funcs, models = self._build_fit_inputs() + return MultiFitter(fit_objects=models, fit_functions=funcs) + + def _get_sampling_data(self) -> tuple[list, list, list]: + """ + Get the per-target data to bind to the Sampler. + + Returns + ------- + tuple[list, list, list] + The ``(x, y, weights)`` triple, one entry per fit target. + """ + xs, ys, ws, _, _ = self._build_fit_inputs() + return xs, ys, ws + + def _get_chain_parameters(self) -> list[Parameter]: + """ + Get the free parameters across every binding model. + + A model appears once per target it is fitted against, so the union is taken by + ``unique_name`` to avoid counting its parameters more than once. + + Returns + ------- + list[Parameter] + The free parameters of the binding models, without duplicates. + """ + parameters = {} + for binding in self.bindings: + for parameter in binding.model.get_free_parameters(): + parameters.setdefault(parameter.unique_name, parameter) + return list(parameters.values()) + + def parameter_label(self, parameter: Parameter) -> str: + """ + Label a parameter with the binding model it belongs to. + + Two bindings can use models of the same kind, whose parameters would then share a name, so + the model's display name is prefixed when that is needed to tell them apart. A single + binding, or models that already name their parameters after themselves, keep their plain + names -- these get long quickly, and there is nothing to disambiguate. + + Parameters + ---------- + parameter : Parameter + The parameter to label. + + Returns + ------- + str + The parameter name, qualified by model where that is needed to tell copies apart. + """ + if not self._name_is_ambiguous(parameter): + return parameter.name + owners = [ + binding.model + for binding in self.bindings + if any( + p.unique_name == parameter.unique_name for p in binding.model.get_free_parameters() + ) + ] + if len(owners) != 1: + return parameter.name + model_name = owners[0].display_name + if model_name is None or parameter.name.startswith(model_name): + return parameter.name + return f'{model_name}: {parameter.name}' def plot( self, names: str | list[str] | None = None, **kwargs: dict[str, Any] diff --git a/src/easydynamics/analysis/posterior.py b/src/easydynamics/analysis/posterior.py index e22ae87d9..717e99e10 100644 --- a/src/easydynamics/analysis/posterior.py +++ b/src/easydynamics/analysis/posterior.py @@ -40,6 +40,9 @@ class BoundsSuggestion: ---------- parameter : Parameter The parameter the suggestion applies to. + label : str + The name the parameter is reported under. For a multi-Q analysis this is qualified by Q, + since every Q holds an identically named copy of each parameter. suggested_min : float The proposed lower bound. Equal to the parameter's current lower bound when that is already finite. @@ -51,6 +54,7 @@ class BoundsSuggestion: """ parameter: Parameter + label: str suggested_min: float suggested_max: float reason: str @@ -178,7 +182,8 @@ def __repr__(self) -> str: if not self._suggestions: return 'BoundsSuggestions(no free parameters)' - header = f'{"parameter":<28s} {"current":>26s} {"suggested":>26s}' + width = max(len('parameter'), *(len(s.label) for s in self._suggestions)) + header = f'{"parameter":<{width}s} {"current":>26s} {"suggested":>26s}' lines = ['BoundsSuggestions', header, '-' * len(header)] for s in self._suggestions: current = f'({s.parameter.min:.4g}, {s.parameter.max:.4g})' @@ -186,7 +191,7 @@ def __repr__(self) -> str: suggested = f'-- {s.reason}' else: suggested = f'({s.suggested_min:.4g}, {s.suggested_max:.4g})' - lines.append(f'{s.parameter.name:<28s} {current:>26s} {suggested:>26s}') + lines.append(f'{s.label:<{width}s} {current:>26s} {suggested:>26s}') attention = self.needing_attention if attention: @@ -199,6 +204,7 @@ def __repr__(self) -> str: def suggest_bounds_for_parameters( parameters: list[Parameter], + labels: list[str] | None = None, n_sigma: float = 10.0, relative_pad: float = 0.2, absolute_floor: float | None = None, @@ -223,6 +229,10 @@ def suggest_bounds_for_parameters( ---------- parameters : list[Parameter] The parameters to propose bounds for. + labels : list[str] | None, default=None + The name to report each parameter under, one per parameter. Defaults to the parameters' own + names, which is ambiguous when several share a name, as the per-Q copies of a multi-Q + analysis do. n_sigma : float, default=10.0 How many standard deviations of the parameter's fitted uncertainty to allow on each side. relative_pad : float, default=0.2 @@ -242,20 +252,24 @@ def suggest_bounds_for_parameters( if absolute_floor is not None: _verify_nonneg_number(absolute_floor, 'absolute_floor') + if labels is None: + labels = [parameter.name for parameter in parameters] suggestions = [ _suggest_bounds_for_parameter( parameter=parameter, + label=label, n_sigma=n_sigma, relative_pad=relative_pad, absolute_floor=absolute_floor, ) - for parameter in parameters + for parameter, label in zip(parameters, labels, strict=True) ] return BoundsSuggestions(suggestions) def _suggest_bounds_for_parameter( parameter: Parameter, + label: str, n_sigma: float, relative_pad: float, absolute_floor: float | None, @@ -267,6 +281,8 @@ def _suggest_bounds_for_parameter( ---------- parameter : Parameter The parameter to propose bounds for. + label : str + The name to report the parameter under. n_sigma : float How many standard deviations to allow on each side. relative_pad : float @@ -288,6 +304,7 @@ def _suggest_bounds_for_parameter( if min_is_finite and max_is_finite: return BoundsSuggestion( parameter=parameter, + label=label, suggested_min=current_min, suggested_max=current_max, reason='', @@ -298,6 +315,7 @@ def _suggest_bounds_for_parameter( if not np.isfinite(value): return BoundsSuggestion( parameter=parameter, + label=label, suggested_min=current_min, suggested_max=current_max, reason='value is not finite', @@ -312,6 +330,7 @@ def _suggest_bounds_for_parameter( if not np.isfinite(half_width) or half_width <= 0: return BoundsSuggestion( parameter=parameter, + label=label, suggested_min=current_min, suggested_max=current_max, reason='no scale information (zero value and uncertainty)', @@ -319,6 +338,7 @@ def _suggest_bounds_for_parameter( return BoundsSuggestion( parameter=parameter, + label=label, suggested_min=current_min if min_is_finite else value - half_width, suggested_max=current_max if max_is_finite else value + half_width, reason='', @@ -527,13 +547,14 @@ def __repr__(self) -> str: if not self._entries: return 'PosteriorSummary(no parameters)' + width = max(len('parameter'), *(len(e.name) for e in self._entries)) header = ( - f'{"parameter":<28s} {"unit":>10s} {"median":>14s} ' + f'{"parameter":<{width}s} {"unit":>10s} {"median":>14s} ' f'{"-":>12s} {"+":>12s} {"current":>14s}' ) lines = ['PosteriorSummary', header, '-' * len(header)] lines.extend( - f'{e.name:<28s} {e.unit:>10s} {e.median:>14.5g} ' + f'{e.name:<{width}s} {e.unit:>10s} {e.median:>14.5g} ' f'{e.minus:>12.4g} {e.plus:>12.4g} {e.value:>14.5g}' for e in self._entries ) @@ -542,22 +563,23 @@ def __repr__(self) -> str: def summarize_draws( draws: np.ndarray, - fallback_names: list[str], + labels: list[str], parameters_by_column: list[Parameter | None], ) -> PosteriorSummary: """ - Summarize posterior draws under the parameters' own names and units. + Summarize posterior draws under caller-supplied labels. 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. + like), which is not what a user recognises, so the caller supplies readable labels instead. A + plain parameter name is enough for a single dataset, but a multi-Q analysis holds one copy of + each parameter per Q, all sharing a name, so those labels have to be qualified by Q. Parameters ---------- draws : np.ndarray Posterior draws, shape ``(n_draws, n_parameters)``. - fallback_names : list[str] - Label to use for any column with no matching parameter, one per column. + labels : list[str] + The label to report each column under, one per column. parameters_by_column : list[Parameter | None] The parameter for each column of ``draws``, or None where none could be matched. @@ -573,7 +595,7 @@ def summarize_draws( ) entries.append( ParameterPosterior( - name=fallback_names[column] if parameter is None else parameter.name, + name=labels[column], unit='' if parameter is None else str(parameter.unit), median=median, lower=lower, diff --git a/tests/integration/fitting/test_bayesian_sampling_multi_q.py b/tests/integration/fitting/test_bayesian_sampling_multi_q.py new file mode 100644 index 000000000..5870aab49 --- /dev/null +++ b/tests/integration/fitting/test_bayesian_sampling_multi_q.py @@ -0,0 +1,216 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +""" +Integration tests running real BUMPS DREAM chains through Analysis and ParameterAnalysis. + +Slow by nature, and run with ``sampler_kwargs={'trim': False}`` for the same reason as the +single-Q integration tests: 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') + +import easydynamics as edyn +import easydynamics.sample_model as sm + +Q_VALUES = [0.5, 1.0, 1.5] +NOISE = 0.02 +TRUE_AREA = 2.0 + +SAMPLE_KWARGS = { + 'samples': 2000, + 'burn': 100, + 'thin': 2, + 'sampler_kwargs': {'trim': False}, +} + + +def true_width(q): + return 0.8 + 0.4 * q**2 + + +def build_analysis(): + energy_values = np.linspace(-5.0, 5.0, 40) + rng = np.random.default_rng(0) + rows = [] + for q in Q_VALUES: + width = true_width(q) + row = TRUE_AREA / (width * np.sqrt(2 * np.pi)) + row = row * np.exp(-0.5 * (energy_values / width) ** 2) + rows.append(row + rng.normal(0.0, NOISE, size=row.shape)) + observed = np.vstack(rows) + + experiment = edyn.Experiment( + data=sc.DataArray( + data=sc.array( + dims=['Q', 'energy'], + values=observed, + variances=np.full_like(observed, NOISE**2), + ), + coords={ + 'Q': sc.array(dims=['Q'], values=Q_VALUES, unit='1/Angstrom'), + 'energy': sc.array(dims=['energy'], values=energy_values, unit='meV'), + }, + ) + ) + return edyn.Analysis( + display_name='MultiQIntegration', + experiment=experiment, + sample_model=sm.SampleModel(components=sm.Gaussian(area=TRUE_AREA, width=1.0)), + instrument_model=sm.InstrumentModel(), + ) + + +@pytest.fixture(scope='module') +def simultaneously_sampled(): + analysis = build_analysis() + analysis.fit(fit_method='simultaneous') + analysis.suggest_bounds().apply() + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + analysis.sample_posterior(fit_method='simultaneous', **SAMPLE_KWARGS) + return analysis + + +class TestSimultaneousChain: + def test_chain_covers_every_q_index(self, simultaneously_sampled): + # EXPECT one column per free parameter across all Q, in one chain + results = simultaneously_sampled.posterior_result + assert results.draws.shape[1] == len(simultaneously_sampled._get_chain_parameters()) + assert results.draws.shape[1] == 3 * len(Q_VALUES) + + def test_summary_labels_are_unique_and_q_qualified(self, simultaneously_sampled): + # WHEN + names = [entry.name for entry in simultaneously_sampled.posterior_summary()] + + # EXPECT + assert len(set(names)) == len(names) + assert all('Q_index=' in name for name in names) + + @pytest.mark.parametrize('q_index', range(len(Q_VALUES))) + def test_posterior_recovers_the_true_width_at_each_q(self, simultaneously_sampled, q_index): + # WHEN + entry = simultaneously_sampled.posterior_summary()[f'Gaussian width (Q_index={q_index})'] + + # EXPECT the truth within a few posterior standard deviations. A 68% interval is not used + # here: it excludes the truth about a third of the time for a single noise realization. + spread = max(entry.minus, entry.plus) + assert abs(entry.median - true_width(Q_VALUES[q_index])) < 4 * spread + + def test_sampling_leaves_the_fitted_values_untouched(self): + # WHEN + analysis = build_analysis() + analysis.fit(fit_method='simultaneous') + analysis.suggest_bounds().apply() + before = [float(p.value) for p in analysis._get_chain_parameters()] + + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + analysis.sample_posterior(fit_method='simultaneous', **SAMPLE_KWARGS) + + # EXPECT + after = [float(p.value) for p in analysis._get_chain_parameters()] + assert after == pytest.approx(before) + + def test_plots_render(self, simultaneously_sampled): + # WHEN + import matplotlib.pyplot as plt + + n_parameters = len(simultaneously_sampled._get_chain_parameters()) + trace = simultaneously_sampled.plot_trace() + corner = simultaneously_sampled.plot_corner() + + # EXPECT + assert len(trace.axes) == n_parameters + 1 + assert len(corner.axes) == n_parameters**2 + plt.close('all') + + +class TestIndependentChains: + def test_one_chain_per_q_index(self): + # WHEN + analysis = build_analysis() + analysis.fit(fit_method='independent') + for analysis1d in analysis.analysis_list: + analysis1d.suggest_bounds().apply() + + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + results = analysis.sample_posterior(fit_method='independent', **SAMPLE_KWARGS) + + # EXPECT + assert len(results) == len(Q_VALUES) + for analysis1d, result in zip(analysis.analysis_list, results, strict=True): + assert result.draws.shape[1] == len(analysis1d.get_free_parameters()) + + def test_independent_and_simultaneous_agree_on_the_widths(self, simultaneously_sampled): + # WHEN the same data is sampled per-Q instead of all at once + analysis = build_analysis() + analysis.fit(fit_method='independent') + for analysis1d in analysis.analysis_list: + analysis1d.suggest_bounds().apply() + + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + analysis.sample_posterior(fit_method='independent', **SAMPLE_KWARGS) + + # EXPECT both routes land on the same widths, since nothing is shared across Q here + for q_index, analysis1d in enumerate(analysis.analysis_list): + independent = analysis1d.posterior_summary()['Gaussian width'] + simultaneous = simultaneously_sampled.posterior_summary()[ + f'Gaussian width (Q_index={q_index})' + ] + spread = max(independent.minus, independent.plus, simultaneous.plus) + assert abs(independent.median - simultaneous.median) < 4 * spread + + +class TestParameterAnalysisChain: + def test_recovers_a_straight_line_through_the_widths(self): + # WHEN the fitted widths are themselves fitted against a model of their Q dependence + q = np.array(Q_VALUES) + widths = true_width(q) + dataset = sc.Dataset({ + 'Gaussian width': sc.DataArray( + data=sc.array( + dims=['Q'], + values=widths, + variances=np.full_like(widths, 0.01**2), + unit='meV', + ), + coords={'Q': sc.array(dims=['Q'], values=q, unit='1/angstrom')}, + ) + }) + model = sm.Polynomial( + coefficients=[0.8, 0.0, 0.4], x_unit='1/angstrom', y_unit='meV', name='Width model' + ) + analysis = edyn.ParameterAnalysis( + parameters=dataset, + bindings=edyn.FitBinding(model=model, targets='Gaussian width'), + ) + analysis.fit() + + # The linear coefficient sits at exactly zero with a vanishing uncertainty, so the sigma + # rule has no scale to work from and flags it rather than inventing one. absolute_floor + # supplies the scale the data cannot. + flagged = analysis.suggest_bounds().needing_attention + assert [s.label for s in flagged] == ['Width model_c1'] + + analysis.suggest_bounds(absolute_floor=1.0).apply() + assert not analysis.suggest_bounds().needing_attention + + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + results = analysis.sample_posterior(**SAMPLE_KWARGS) + + # EXPECT a column per polynomial coefficient, and a readable summary + assert results.draws.shape[1] == len(analysis._get_chain_parameters()) + names = [entry.name for entry in analysis.posterior_summary()] + assert len(set(names)) == len(names) diff --git a/tests/unit/easydynamics/analysis/test_analysis_bayesian.py b/tests/unit/easydynamics/analysis/test_analysis_bayesian.py new file mode 100644 index 000000000..97e88a1d6 --- /dev/null +++ b/tests/unit/easydynamics/analysis/test_analysis_bayesian.py @@ -0,0 +1,249 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +"""Unit tests for Bayesian sampling on the 2D Analysis, 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 + +import easydynamics as edyn +import easydynamics.sample_model as sm + +SAMPLER_PATH = 'easydynamics.analysis.bayesian_sampling.Sampler' +Q_VALUES = [0.5, 1.0, 1.5] + + +def make_analysis(): + energy_values = np.linspace(-5.0, 5.0, 15) + rows = [2.0 * np.exp(-0.5 * (energy_values / (0.8 + 0.4 * q**2)) ** 2) for q in Q_VALUES] + observed = np.vstack(rows) + experiment = edyn.Experiment( + data=sc.DataArray( + data=sc.array( + dims=['Q', 'energy'], + values=observed, + variances=np.full_like(observed, 0.01), + ), + coords={ + 'Q': sc.array(dims=['Q'], values=Q_VALUES, unit='1/Angstrom'), + 'energy': sc.array(dims=['energy'], values=energy_values, unit='meV'), + }, + ) + ) + return edyn.Analysis( + display_name='TestMultiQ', + experiment=experiment, + sample_model=sm.SampleModel(components=sm.Gaussian(area=2.0, width=1.0)), + instrument_model=sm.InstrumentModel(), + ) + + +def bound_all(analysis, half_width=5.0): + for parameter in analysis._get_chain_parameters(): + parameter.min = float(parameter.value) - half_width + parameter.max = float(parameter.value) + half_width + + +def fake_results(parameters, n_draws=50): + draws = np.tile([float(p.value) for p in parameters], (n_draws, 1)) + return SimpleNamespace( + draws=draws, + param_names=[p.unique_name for p in parameters], + logp=np.zeros(n_draws), + state=MagicMock(Ngen=10, Npop=4), + ) + + +@pytest.fixture +def analysis(): + return make_analysis() + + +class TestChainParameters: + def test_union_covers_every_q_index(self, analysis): + # WHEN + parameters = analysis._get_chain_parameters() + + # EXPECT one copy of each per-Q parameter, with no duplicates + assert len(parameters) == sum(len(a.get_free_parameters()) for a in analysis.analysis_list) + assert len({p.unique_name for p in parameters}) == len(parameters) + + def test_labels_are_qualified_by_q_index(self, analysis): + # WHEN + labels = [analysis.parameter_label(p) for p in analysis._get_chain_parameters()] + + # EXPECT every per-Q copy is distinguishable, which the bare name would not be + assert len(set(labels)) == len(labels) + assert 'Gaussian width (Q_index=0)' in labels + assert 'Gaussian width (Q_index=2)' in labels + + def test_bare_names_would_collide(self, analysis): + # WHEN + names = [p.name for p in analysis._get_chain_parameters()] + + # EXPECT the collision the Q-qualified label exists to solve + assert len(set(names)) < len(names) + + +class TestBoundsPreflight: + def test_sampling_refuses_unbounded_parameters(self, analysis): + # EXPECT + with pytest.raises(ValueError, match='finite bounds'): + analysis.sample_posterior(fit_method='simultaneous', samples=10) + + def test_error_names_parameters_by_q_index(self, analysis): + # EXPECT + with pytest.raises(ValueError, match=r'Gaussian width \(Q_index=0\)'): + analysis.check_bounds_for_sampling() + + def test_suggest_bounds_labels_every_q(self, analysis): + # WHEN + suggestions = analysis.suggest_bounds() + + # EXPECT + labels = [s.label for s in suggestions] + assert len(set(labels)) == len(labels) + assert 'Gaussian area (Q_index=1)' in labels + + +class TestSimultaneousSampling: + def test_binds_one_dataset_per_q_index(self, analysis): + # WHEN + bound_all(analysis) + parameters = analysis._get_chain_parameters() + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_results(parameters) + analysis.sample_posterior(fit_method='simultaneous', samples=10) + + # EXPECT + args, kwargs = sampler_class.call_args + assert len(args[1]) == len(Q_VALUES) + assert len(args[2]) == len(Q_VALUES) + assert len(kwargs['weights']) == len(Q_VALUES) + + def test_returns_a_single_result(self, analysis): + # WHEN + bound_all(analysis) + parameters = analysis._get_chain_parameters() + + with patch(SAMPLER_PATH) as sampler_class: + expected = fake_results(parameters) + sampler_class.return_value.sample.return_value = expected + returned = analysis.sample_posterior(fit_method='simultaneous', samples=10) + + # EXPECT + assert returned is expected + assert analysis.posterior_result is expected + + def test_summary_is_labelled_by_q_index(self, analysis): + # WHEN + bound_all(analysis) + parameters = analysis._get_chain_parameters() + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_results(parameters) + analysis.sample_posterior(fit_method='simultaneous', samples=10) + + # EXPECT + names = [entry.name for entry in analysis.posterior_summary()] + assert len(set(names)) == len(names) + assert all('Q_index=' in name for name in names) + + def test_refreshes_every_convolver_before_sampling(self, analysis): + # WHEN + bound_all(analysis) + parameters = analysis._get_chain_parameters() + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_results(parameters) + for analysis1d in analysis.analysis_list: + analysis1d._convolver_is_dirty = True + analysis.sample_posterior(fit_method='simultaneous', samples=10) + + # EXPECT the sampler sees the same prepared convolvers a simultaneous fit would + assert all(not a._convolver_is_dirty for a in analysis.analysis_list) + + def test_uses_a_multifitter(self, analysis): + # WHEN + from easyscience.fitting.multi_fitter import MultiFitter + + # EXPECT + assert isinstance(analysis.fitter, MultiFitter) + assert len(analysis.fitter.fit_object) == len(Q_VALUES) + + +class TestIndependentSampling: + def test_returns_one_result_per_q_index(self, analysis): + # WHEN + for analysis1d in analysis.analysis_list: + for parameter in analysis1d.get_free_parameters(): + parameter.min = float(parameter.value) - 5.0 + parameter.max = float(parameter.value) + 5.0 + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results( + analysis.analysis_list[0].get_free_parameters() + ) + results = analysis.sample_posterior(fit_method='independent', samples=10) + + # EXPECT + assert isinstance(results, list) + assert len(results) == len(Q_VALUES) + + def test_single_q_index_returns_one_result(self, analysis): + # WHEN + target = analysis.analysis_list[1] + for parameter in target.get_free_parameters(): + parameter.min = float(parameter.value) - 5.0 + parameter.max = float(parameter.value) + 5.0 + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results( + target.get_free_parameters() + ) + result = analysis.sample_posterior(fit_method='independent', Q_index=1, samples=10) + + # EXPECT + assert not isinstance(result, list) + assert result is target.posterior_result + + def test_invalid_q_index_raises(self, analysis): + # EXPECT + with pytest.raises((ValueError, IndexError)): + analysis.sample_posterior(fit_method='independent', Q_index=99, samples=10) + + +class TestSamplePosteriorValidation: + def test_invalid_fit_method_raises(self, analysis): + # EXPECT + with pytest.raises(ValueError, match='Invalid fit method'): + analysis.sample_posterior(fit_method='nonsense') + + def test_missing_q_values_raises(self): + # WHEN + analysis = edyn.Analysis(display_name='Empty') + + # EXPECT + with pytest.raises(ValueError, match='No Q values available'): + analysis.sample_posterior() + + +class TestPredictivePlot: + def test_predictive_is_not_supported_for_multiple_datasets(self, analysis): + # WHEN + bound_all(analysis) + parameters = analysis._get_chain_parameters() + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_results(parameters) + analysis.sample_posterior(fit_method='simultaneous', samples=10) + + # EXPECT + with pytest.raises(NotImplementedError, match='single dataset only'): + analysis.plot_posterior_predictive() diff --git a/tests/unit/easydynamics/analysis/test_parameter_analysis_bayesian.py b/tests/unit/easydynamics/analysis/test_parameter_analysis_bayesian.py new file mode 100644 index 000000000..05c1858eb --- /dev/null +++ b/tests/unit/easydynamics/analysis/test_parameter_analysis_bayesian.py @@ -0,0 +1,206 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +"""Unit tests for Bayesian sampling on ParameterAnalysis, with the 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.multi_fitter import MultiFitter + +import easydynamics as edyn +import easydynamics.sample_model as sm + +SAMPLER_PATH = 'easydynamics.analysis.bayesian_sampling.Sampler' +Q = np.array([0.5, 0.8, 1.1, 1.4, 1.7, 2.0]) + + +def make_dataset(): + widths = 0.10 + 0.35 * Q + areas = 2.0 - 0.3 * Q + return sc.Dataset({ + 'Lorentzian width': sc.DataArray( + data=sc.array( + dims=['Q'], values=widths, variances=np.full_like(widths, 1e-4), unit='meV' + ), + coords={'Q': sc.array(dims=['Q'], values=Q, unit='1/angstrom')}, + ), + 'Lorentzian area': sc.DataArray( + data=sc.array( + dims=['Q'], values=areas, variances=np.full_like(areas, 4e-4), unit='meV' + ), + coords={'Q': sc.array(dims=['Q'], values=Q, unit='1/angstrom')}, + ), + }) + + +def make_analysis(two_bindings=True): + bindings = [ + edyn.FitBinding( + model=sm.Polynomial( + coefficients=[0.1, 0.35], x_unit='1/angstrom', y_unit='meV', name='Width line' + ), + targets='Lorentzian width', + ) + ] + if two_bindings: + bindings.append( + edyn.FitBinding( + model=sm.Polynomial( + coefficients=[2.0, -0.3], x_unit='1/angstrom', y_unit='meV', name='Area line' + ), + targets='Lorentzian area', + ) + ) + return edyn.ParameterAnalysis(parameters=make_dataset(), bindings=bindings) + + +def bound_all(analysis, half_width=5.0): + for parameter in analysis._get_chain_parameters(): + parameter.min = float(parameter.value) - half_width + parameter.max = float(parameter.value) + half_width + + +def fake_results(parameters, n_draws=50): + draws = np.tile([float(p.value) for p in parameters], (n_draws, 1)) + return SimpleNamespace( + draws=draws, + param_names=[p.unique_name for p in parameters], + logp=np.zeros(n_draws), + state=MagicMock(Ngen=10, Npop=4), + ) + + +@pytest.fixture +def analysis(): + return make_analysis() + + +class TestFitterExposure: + def test_fitter_is_a_cached_multifitter(self, analysis): + # EXPECT + assert isinstance(analysis.fitter, MultiFitter) + assert analysis.fitter is analysis.fitter + + def test_fit_still_returns_per_target_results(self, analysis): + # WHEN + results = analysis.fit() + + # EXPECT one result per fit target, as before + assert isinstance(results, list) + assert len(results) == 2 + + def test_changing_bindings_rebuilds_the_fitter(self, analysis): + # WHEN + original = analysis.fitter + analysis.bindings = analysis.bindings[:1] + + # EXPECT + assert analysis.fitter is not original + + def test_changing_parameters_rebuilds_the_fitter(self, analysis): + # WHEN + original = analysis.fitter + analysis.parameters = make_dataset() + + # EXPECT + assert analysis.fitter is not original + + +class TestChainParameters: + def test_covers_every_binding_model(self, analysis): + # WHEN + parameters = analysis._get_chain_parameters() + + # EXPECT both Polynomials contribute their two coefficients + assert len(parameters) == 4 + assert len({p.unique_name for p in parameters}) == 4 + + def test_labels_are_unique(self, analysis): + # WHEN + labels = [analysis.parameter_label(p) for p in analysis._get_chain_parameters()] + + # EXPECT + assert len(set(labels)) == len(labels) + + def test_model_name_is_not_repeated_in_the_label(self, analysis): + # WHEN a model already names its parameters after itself + labels = [analysis.parameter_label(p) for p in analysis._get_chain_parameters()] + + # EXPECT no 'Width line: Width line_c0' + assert 'Width line_c0' in labels + assert not any(label.count('Width line') > 1 for label in labels) + + +class TestSampling: + def test_refuses_unbounded_parameters(self, analysis): + # EXPECT + with pytest.raises(ValueError, match='finite bounds'): + analysis.sample_posterior(samples=10) + + def test_binds_one_dataset_per_target(self, analysis): + # WHEN + bound_all(analysis) + parameters = analysis._get_chain_parameters() + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_results(parameters) + analysis.sample_posterior(samples=10) + + # EXPECT + args, kwargs = sampler_class.call_args + assert len(args[1]) == 2 + assert len(kwargs['weights']) == 2 + + def test_summary_uses_model_qualified_labels(self, analysis): + # WHEN + bound_all(analysis) + parameters = analysis._get_chain_parameters() + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_results(parameters) + analysis.sample_posterior(samples=10) + + # EXPECT + names = [entry.name for entry in analysis.posterior_summary()] + assert len(set(names)) == len(names) + assert 'Width line_c0' in names + + def test_restores_parameter_values(self, analysis): + # WHEN + bound_all(analysis) + parameters = analysis._get_chain_parameters() + before = [float(p.value) for p in parameters] + + with patch(SAMPLER_PATH) as sampler_class: + + def mutate(**_kwargs): + for parameter in parameters: + parameter.value = float(parameter.value) + 1.0 + return fake_results(parameters) + + sampler_class.return_value.sample.side_effect = mutate + analysis.sample_posterior(samples=10) + + # EXPECT + assert [float(p.value) for p in parameters] == pytest.approx(before) + + def test_missing_parameters_dataset_raises(self): + # WHEN + analysis = edyn.ParameterAnalysis() + + # EXPECT + with pytest.raises(ValueError, match='No parameters Dataset'): + analysis.sample_posterior(samples=10) + + def test_missing_bindings_raises(self): + # WHEN + analysis = edyn.ParameterAnalysis(parameters=make_dataset()) + + # EXPECT + with pytest.raises(ValueError, match='No fit bindings'): + analysis.sample_posterior(samples=10) diff --git a/tests/unit/easydynamics/analysis/test_posterior.py b/tests/unit/easydynamics/analysis/test_posterior.py index 88dc29bf6..dc9a5e4a1 100644 --- a/tests/unit/easydynamics/analysis/test_posterior.py +++ b/tests/unit/easydynamics/analysis/test_posterior.py @@ -251,7 +251,7 @@ def test_reports_parameter_names_units_and_percentiles(self): draws = np.linspace(0.0, 100.0, 101).reshape(-1, 1) # THEN - summary = summarize_draws(draws, ['Parameter_0'], [parameter]) + summary = summarize_draws(draws, ['Gaussian width'], [parameter]) # EXPECT entry = summary['Gaussian width'] @@ -263,6 +263,18 @@ def test_reports_parameter_names_units_and_percentiles(self): assert entry.plus == pytest.approx(34.0) assert entry.value == pytest.approx(1.5) + def test_labels_qualified_by_q_are_kept_verbatim(self): + # WHEN a multi-Q analysis supplies Q-qualified labels, since every Q holds a copy of the + # same parameter and the bare name would repeat + parameter = make_parameter(name='Gaussian width') + + # THEN + summary = summarize_draws(np.zeros((5, 1)), ['Gaussian width (Q_index=2)'], [parameter]) + + # EXPECT + assert summary.entries[0].name == 'Gaussian width (Q_index=2)' + assert summary.entries[0].unit == 'meV' + def test_unmatched_column_falls_back_to_the_supplied_name(self): # WHEN draws = np.zeros((10, 1)) @@ -289,7 +301,7 @@ def test_repr_contains_the_parameter_name(self): parameter = make_parameter(name='Gaussian area') # THEN - text = repr(summarize_draws(np.zeros((5, 1)), ['x'], [parameter])) + text = repr(summarize_draws(np.zeros((5, 1)), ['Gaussian area'], [parameter])) # EXPECT assert 'Gaussian area' in text From 7ad29003a2322222c5f68eb624bd7e9904543b87 Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Thu, 13 Aug 2026 13:48:13 +0200 Subject: [PATCH 03/29] Label the posterior plot axes with units and quantities The summary table already reported each parameter's unit, but the plots did not, so a diffusion coefficient came out as a bare number. Units are now threaded through to plot_trace and plot_corner, and the posterior predictive plot gets axis labels taken from the analysis' own energy and intensity units. Details that needed care: - Matplotlib parks a shared exponent at the end of the axis, on top of the axis label. It is now folded into the label, sharing one set of parentheses with the unit, so a diffusion coefficient reads "diffusion_coefficient (1e-8 m^2/s)" rather than stacking two parentheticals or overlapping. - Dimensionless and empty units are skipped. A polynomial coefficient labelled "dimensionless" is noise. - The top-left panel of a corner plot is a histogram, so its vertical axis counts draws rather than carrying a parameter. It is now labelled "counts" instead of being left blank, which read as an omission. - Corner tick counts are capped, since four labelled ticks per panel is as much as a small panel can carry legibly. Co-Authored-By: Claude Opus 5 (1M context) --- .../analysis/bayesian_sampling.py | 42 +++++++ .../analysis/parameter_analysis.py | 29 +++-- src/easydynamics/utils/posterior_plotting.py | 118 +++++++++++++++++- .../analysis/test_analysis1d_bayesian.py | 76 +++++++++++ .../analysis/test_analysis_bayesian.py | 42 +++++++ .../analysis/test_bayesian_sampling_mixin.py | 47 +++++++ .../test_parameter_analysis_bayesian.py | 62 +++++++++ .../easydynamics/analysis/test_posterior.py | 16 +++ .../utils/test_posterior_plotting.py | 39 ++++++ 9 files changed, 458 insertions(+), 13 deletions(-) create mode 100644 tests/unit/easydynamics/analysis/test_bayesian_sampling_mixin.py diff --git a/src/easydynamics/analysis/bayesian_sampling.py b/src/easydynamics/analysis/bayesian_sampling.py index 28dfb43b4..c9adb11f4 100644 --- a/src/easydynamics/analysis/bayesian_sampling.py +++ b/src/easydynamics/analysis/bayesian_sampling.py @@ -760,6 +760,7 @@ def plot_trace(self, **kwargs: dict[str, Any]) -> Figure: from easydynamics.utils.posterior_plotting import plot_trace results = self._require_posterior_result() + kwargs.setdefault('units', self._chain_units(results)) return plot_trace( draws=results.draws, logp=results.logp, @@ -789,6 +790,7 @@ def plot_corner(self, **kwargs: dict[str, Any]) -> Figure: from easydynamics.utils.posterior_plotting import plot_corner results = self._require_posterior_result() + kwargs.setdefault('units', self._chain_units(results)) return plot_corner( draws=results.draws, names=self._chain_display_names(results), @@ -847,6 +849,8 @@ def plot_posterior_predictive( predictions = self._evaluate_over_draws(results, x, n_draws) y_err = None if weights is None else 1.0 / np.asarray(weights) + kwargs.setdefault('xlabel', self._predictive_axis_labels()[0]) + kwargs.setdefault('ylabel', self._predictive_axis_labels()[1]) return plot_posterior_predictive( x=np.asarray(x), y=np.asarray(y), @@ -857,6 +861,25 @@ def plot_posterior_predictive( **kwargs, ) + def _predictive_axis_labels(self) -> tuple[str | None, str | None]: + """ + Get default axis labels for the posterior predictive plot. + + The base implementation reads the energy and intensity units off the analysis when they are + available, and falls back to no label rather than guessing. + + Returns + ------- + tuple[str | None, str | None] + The ``(xlabel, ylabel)`` pair. + """ + energy = getattr(self, 'energy', None) + xlabel = None if energy is None else f'Energy ({energy.unit})' + sample_model = getattr(self, 'sample_model', None) + y_unit = None if sample_model is None else getattr(sample_model, 'y_unit', None) + ylabel = 'Intensity' if y_unit is None else f'Intensity ({y_unit})' + return xlabel, ylabel + def _evaluate_over_draws( self, results: SamplingResults, @@ -975,6 +998,25 @@ def _name_is_ambiguous(self, parameter: Parameter) -> bool: names = [p.name for p in self._get_chain_parameters()] return names.count(parameter.name) > 1 + def _chain_units(self, results: SamplingResults) -> list[str]: + """ + Get the unit of each column of the chain. + + Parameters + ---------- + results : SamplingResults + The sampling results whose columns should be described. + + 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_chain_parameters(results) + ] + def _chain_display_names(self, results: SamplingResults) -> list[str]: """ Translate the chain's column names into readable labels. diff --git a/src/easydynamics/analysis/parameter_analysis.py b/src/easydynamics/analysis/parameter_analysis.py index 342f33c69..47020b1d8 100644 --- a/src/easydynamics/analysis/parameter_analysis.py +++ b/src/easydynamics/analysis/parameter_analysis.py @@ -287,9 +287,12 @@ def parameter_label(self, parameter: Parameter) -> str: Label a parameter with the binding model it belongs to. Two bindings can use models of the same kind, whose parameters would then share a name, so - the model's display name is prefixed when that is needed to tell them apart. A single - binding, or models that already name their parameters after themselves, keep their plain - names -- these get long quickly, and there is nothing to disambiguate. + the owning model is prefixed when that is needed to tell them apart. A single binding, or + models that already name their parameters after themselves, keep their plain names -- these + get long quickly, and there is nothing to disambiguate. + + The prefix is the model's display name, unless two models share that too, in which case the + unique name is used: a label that does not actually disambiguate is worse than a long one. Parameters ---------- @@ -303,19 +306,21 @@ def parameter_label(self, parameter: Parameter) -> str: """ if not self._name_is_ambiguous(parameter): return parameter.name + + models = {binding.model.unique_name: binding.model for binding in self.bindings} owners = [ - binding.model - for binding in self.bindings - if any( - p.unique_name == parameter.unique_name for p in binding.model.get_free_parameters() - ) + model + for model in models.values() + if any(p.unique_name == parameter.unique_name for p in model.get_free_parameters()) ] if len(owners) != 1: return parameter.name - model_name = owners[0].display_name - if model_name is None or parameter.name.startswith(model_name): - return parameter.name - return f'{model_name}: {parameter.name}' + + owner = owners[0] + display_names = [model.display_name for model in models.values()] + if owner.display_name is None or display_names.count(owner.display_name) > 1: + return f'{owner.unique_name}: {parameter.name}' + return f'{owner.display_name}: {parameter.name}' def plot( self, names: str | list[str] | None = None, **kwargs: dict[str, Any] diff --git a/src/easydynamics/utils/posterior_plotting.py b/src/easydynamics/utils/posterior_plotting.py index bfb1a1320..8e75cf425 100644 --- a/src/easydynamics/utils/posterior_plotting.py +++ b/src/easydynamics/utils/posterior_plotting.py @@ -14,6 +14,7 @@ import matplotlib.pyplot as plt import numpy as np +from matplotlib.ticker import MaxNLocator if TYPE_CHECKING: from matplotlib.figure import Figure @@ -23,6 +24,7 @@ def plot_trace( draws: np.ndarray, names: list[str], logp: np.ndarray | None = None, + units: list[str] | None = None, title: str | None = None, figsize: tuple[float, float] | None = None, ) -> Figure: @@ -44,6 +46,9 @@ def plot_trace( One label per column of ``draws``. logp : np.ndarray | None, default=None Log-posterior values, plotted in an extra panel when given. + units : list[str] | None, default=None + Unit of each column, appended to its label. Entries that are empty or dimensionless are + skipped, since a bare "dimensionless" only adds clutter. title : str | None, default=None Figure title. figsize : tuple[float, float] | None, default=None @@ -66,7 +71,7 @@ def plot_trace( for axis, column, name in zip(axes, range(draws.shape[1]), names, strict=False): axis.plot(draws[:, column], lw=0.5) - axis.set_ylabel(name, fontsize=8) + axis.set_ylabel(_with_unit(name, units, column), fontsize=8) axis.set_xlim(0, len(draws) - 1) if logp is not None: @@ -83,6 +88,7 @@ def plot_trace( def plot_corner( draws: np.ndarray, names: list[str], + units: list[str] | None = None, title: str | None = None, bins: int = 40, figsize: tuple[float, float] | None = None, @@ -103,6 +109,9 @@ def plot_corner( Posterior draws, shape ``(n_draws, n_parameters)``. names : list[str] One label per column of ``draws``. + units : list[str] | None, default=None + Unit of each column, appended to its label. Entries that are empty or dimensionless are + skipped, since a bare "dimensionless" only adds clutter. title : str | None, default=None Figure title. bins : int, default=40 @@ -143,7 +152,25 @@ def plot_corner( axis.set_ylabel(names[row], fontsize=8) else: axis.set_yticklabels([]) + if row == 0 and col == 0: + # The top-left panel is a histogram, so its vertical axis counts draws rather than + # carrying a parameter. Say so, instead of leaving it blank as if by omission. + axis.set_ylabel('counts', fontsize=8) axis.tick_params(labelsize=7) + axis.xaxis.set_major_locator(MaxNLocator(nbins=4)) + if row != col: + axis.yaxis.set_major_locator(MaxNLocator(nbins=4)) + + # Matplotlib parks the shared exponent ("1e-8") at the end of the axis, where it lands on top + # of the axis label. Fold it into the label instead. + fig.canvas.draw() + for row in range(n): + for col in range(row + 1): + axis = axes[row, col] + if row == n - 1: + _absorb_offset(axis.xaxis, axis.set_xlabel, names[col], units, col) + if col == 0 and row != 0: + _absorb_offset(axis.yaxis, axis.set_ylabel, names[row], units, row) if title is not None: fig.suptitle(title) @@ -158,6 +185,8 @@ def plot_posterior_predictive( y_err: np.ndarray | None = None, title: str | None = None, credible_interval: float = 68.0, + xlabel: str | None = None, + ylabel: str | None = None, figsize: tuple[float, float] = (8.0, 5.0), ) -> Figure: """ @@ -181,6 +210,10 @@ def plot_posterior_predictive( Figure title. credible_interval : float, default=68.0 Width of the credible band, as a percentage. + xlabel : str | None, default=None + Label for the independent axis. + ylabel : str | None, default=None + Label for the dependent axis. figsize : tuple[float, float], default=(8.0, 5.0) Figure size in inches. @@ -224,6 +257,10 @@ def plot_posterior_predictive( label=f'{credible_interval:.0f}% credible band', ) axis.plot(x, median, '-', color='C3', label='Posterior median') + if xlabel is not None: + axis.set_xlabel(xlabel) + if ylabel is not None: + axis.set_ylabel(ylabel) axis.legend() if title is not None: axis.set_title(title) @@ -231,6 +268,85 @@ def plot_posterior_predictive( return fig +def _unit_for(units: list[str] | None, column: int) -> str: + """ + Get the unit to show for a column, if it is worth showing. + + Parameters + ---------- + units : list[str] | None + The units, one per column, or None. + column : int + The column to look up. + + Returns + ------- + str + The unit, or an empty string when there is none worth printing. + """ + if units is None or column >= len(units): + return '' + unit = (units[column] or '').strip() + return '' if unit.lower() in ('', 'dimensionless', 'none') else unit + + +def _with_unit(name: str, units: list[str] | None, column: int) -> str: + """ + Append a column's unit to its label. + + Parameters + ---------- + name : str + The label to extend. + units : list[str] | None + The units, one per column, or None. + column : int + The column the label belongs to. + + Returns + ------- + str + The label, with the unit in parentheses when there is one. + """ + unit = _unit_for(units, column) + return f'{name} ({unit})' if unit else name + + +def _absorb_offset( + axis_object: object, + set_label: object, + name: str, + units: list[str] | None = None, + column: int = 0, +) -> None: + """ + Move an axis' shared exponent into its label, so the two stop overlapping. + + The exponent and the unit share one set of parentheses, since two adjacent parentheticals read + badly: ``D (1e-8 m^2/s)`` rather than ``D (1e-8) (m^2/s)``. + + Parameters + ---------- + axis_object : object + The matplotlib ``XAxis`` or ``YAxis`` carrying the offset text. + set_label : object + The corresponding ``set_xlabel`` or ``set_ylabel`` callable. + name : str + The label the axis should carry, before the exponent and unit are appended. + units : list[str] | None, default=None + The units, one per column, or None. + column : int, default=0 + The column the axis belongs to. + """ + offset_text = axis_object.get_offset_text() + offset = offset_text.get_text() + unit = _unit_for(units, column) + suffix = ' '.join(part for part in (offset, unit) if part) + set_label(f'{name} ({suffix})' if suffix else name, fontsize=8) + if offset: + offset_text.set_visible(False) + + def _verify_draws(draws: np.ndarray, names: list[str]) -> None: """ Verify that a draws array is two-dimensional and matches its labels. diff --git a/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py b/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py index b29a2b54a..d3308aef5 100644 --- a/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py +++ b/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py @@ -484,3 +484,79 @@ def __exit__(self, *exc_info): if exc_info[0] is None: assert not caught, f'unexpected warnings: {[str(w.message) for w in caught]}' return False + + +class TestErrorPaths: + def test_bumps_outlier_crash_is_reported_as_a_degeneracy(self, analysis): + # WHEN BUMPS' own outlier removal indexes past the end of its buffer, which happens when + # chains scatter because the model is not identifiable + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = IndexError( + 'index 71 is out of bounds for axis 0 with size 40' + ) + + # EXPECT the bare IndexError is replaced by something actionable + with pytest.raises(RuntimeError, match='degenerate'): + analysis.sample_posterior(samples=10) + + def test_parameters_entry_of_the_wrong_type_raises(self, analysis): + # EXPECT + with pytest.raises(TypeError, match='Parameter objects or parameter names'): + analysis.sample_posterior(samples=10, parameters=[42]) + + def test_median_skips_columns_with_no_matching_parameter(self, analysis): + # WHEN a chain carries a column this analysis knows nothing about + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + results = fake_results(analysis) + results.param_names = [*results.param_names, 'Parameter_does_not_exist'] + results.draws = np.column_stack([results.draws, np.zeros(results.draws.shape[0])]) + sampler_class.return_value.sample.return_value = results + analysis.sample_posterior(samples=10) + + # EXPECT the unknown column is skipped rather than crashing + changed = analysis.set_parameters_to_posterior_median() + assert len(changed) == len(analysis.get_free_parameters()) + + def test_load_chain_uses_the_sidecar_when_present(self, analysis, tmp_path): + # WHEN a chain is saved and reloaded into a *different* analysis, whose unique names differ + bound_all(analysis) + with patch(SAMPLER_PATH) as sampler_class: + saved = fake_results(analysis) + sampler_class.return_value.sample.return_value = saved + analysis.sample_posterior(samples=10) + analysis.save_chain(str(tmp_path / 'chain')) + + fresh = make_analysis() + bound_all(fresh) + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.load_state.return_value = saved + fresh.load_chain(str(tmp_path / 'chain')) + + # EXPECT the sidecar maps the old unique names onto the new analysis's parameters + summary = fresh.posterior_summary() + assert {entry.name for entry in summary} == {p.name for p in fresh.get_free_parameters()} + assert all(np.isfinite(entry.value) for entry in summary) + + +class TestPlotRendering: + def test_trace_and_corner_render_from_a_chain(self, analysis): + # WHEN + import matplotlib as mpl + import matplotlib.pyplot as plt + + mpl.use('Agg') + bound_all(analysis) + n_parameters = len(analysis.get_free_parameters()) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_results(analysis) + analysis.sample_posterior(samples=10) + + # EXPECT + assert len(analysis.plot_trace().axes) == n_parameters + 1 + assert len(analysis.plot_corner().axes) == n_parameters**2 + plt.close('all') diff --git a/tests/unit/easydynamics/analysis/test_analysis_bayesian.py b/tests/unit/easydynamics/analysis/test_analysis_bayesian.py index 97e88a1d6..de82d93a9 100644 --- a/tests/unit/easydynamics/analysis/test_analysis_bayesian.py +++ b/tests/unit/easydynamics/analysis/test_analysis_bayesian.py @@ -247,3 +247,45 @@ def test_predictive_is_not_supported_for_multiple_datasets(self, analysis): # EXPECT with pytest.raises(NotImplementedError, match='single dataset only'): analysis.plot_posterior_predictive() + + +class TestParameterLabelEdgeCases: + def test_single_q_analysis_keeps_plain_names(self): + # WHEN there is only one Q index, nothing needs disambiguating + energy_values = np.linspace(-5.0, 5.0, 15) + intensity = 2.0 * np.exp(-0.5 * (energy_values / 1.2) ** 2) + experiment = edyn.Experiment( + data=sc.DataArray( + data=sc.array( + dims=['Q', 'energy'], + values=intensity[None, :], + variances=np.full_like(intensity, 0.01)[None, :], + ), + coords={ + 'Q': sc.array(dims=['Q'], values=[1.0], unit='1/Angstrom'), + 'energy': sc.array(dims=['energy'], values=energy_values, unit='meV'), + }, + ) + ) + analysis = edyn.Analysis( + display_name='SingleQ', + experiment=experiment, + sample_model=sm.SampleModel(components=sm.Gaussian(area=2.0, width=1.0)), + instrument_model=sm.InstrumentModel(), + ) + + # THEN + labels = [analysis.parameter_label(p) for p in analysis._get_chain_parameters()] + + # EXPECT the short form, not 'Gaussian width (Q_index=0)' + assert 'Gaussian width' in labels + assert not any('Q_index=' in label for label in labels) + + def test_parameter_from_outside_the_analysis_keeps_its_name(self, analysis): + # WHEN a parameter belongs to no Q index of this analysis + from easyscience.variable import Parameter + + stranger = Parameter(name='Gaussian width', value=1.0) + + # EXPECT it is returned unqualified rather than mislabelled + assert analysis.parameter_label(stranger) == 'Gaussian width' diff --git a/tests/unit/easydynamics/analysis/test_bayesian_sampling_mixin.py b/tests/unit/easydynamics/analysis/test_bayesian_sampling_mixin.py new file mode 100644 index 000000000..e524aee97 --- /dev/null +++ b/tests/unit/easydynamics/analysis/test_bayesian_sampling_mixin.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +"""Tests for the BayesianSamplingMixin contract itself, independent of any Analysis.""" + +import pytest + +from easydynamics.analysis.bayesian_sampling import BayesianSamplingMixin + + +class Incomplete(BayesianSamplingMixin): + """A subclass that implements none of the hooks.""" + + +@pytest.fixture +def incomplete(): + subject = Incomplete() + subject._init_bayesian_state() + return subject + + +class TestHookContract: + def test_building_a_fitter_must_be_implemented(self, incomplete): + # EXPECT + with pytest.raises(NotImplementedError, match='_build_bayesian_fitter'): + incomplete._build_bayesian_fitter() + + def test_getting_the_data_must_be_implemented(self, incomplete): + # EXPECT + with pytest.raises(NotImplementedError, match='_get_sampling_data'): + incomplete._get_sampling_data() + + def test_getting_the_chain_parameters_must_be_implemented(self, incomplete): + # EXPECT + with pytest.raises(NotImplementedError, match='_get_chain_parameters'): + incomplete._get_chain_parameters() + + def test_preparing_for_sampling_is_optional(self, incomplete): + # EXPECT the default hook is a no-op rather than a failure + assert incomplete._prepare_for_sampling() is None + + +class TestInitialState: + def test_nothing_is_cached_before_use(self, incomplete): + # EXPECT + assert incomplete.bayesian_sampler is None + assert incomplete.posterior_result is None diff --git a/tests/unit/easydynamics/analysis/test_parameter_analysis_bayesian.py b/tests/unit/easydynamics/analysis/test_parameter_analysis_bayesian.py index 05c1858eb..6dfc504e0 100644 --- a/tests/unit/easydynamics/analysis/test_parameter_analysis_bayesian.py +++ b/tests/unit/easydynamics/analysis/test_parameter_analysis_bayesian.py @@ -204,3 +204,65 @@ def test_missing_bindings_raises(self): # EXPECT with pytest.raises(ValueError, match='No fit bindings'): analysis.sample_posterior(samples=10) + + +class TestParameterLabelEdgeCases: + def test_colliding_names_are_qualified_by_model(self): + # WHEN two bindings use models whose parameters share a name + shared_name_model_a = sm.Polynomial( + coefficients=[0.1, 0.35], x_unit='1/angstrom', y_unit='meV', name='Line' + ) + shared_name_model_b = sm.Polynomial( + coefficients=[2.0, -0.3], x_unit='1/angstrom', y_unit='meV', name='Line' + ) + analysis = edyn.ParameterAnalysis( + parameters=make_dataset(), + bindings=[ + edyn.FitBinding(model=shared_name_model_a, targets='Lorentzian width'), + edyn.FitBinding(model=shared_name_model_b, targets='Lorentzian area'), + ], + ) + + # THEN + parameters = analysis._get_chain_parameters() + names = [p.name for p in parameters] + labels = [analysis.parameter_label(p) for p in parameters] + + # EXPECT the bare names collide, and the labels resolve it + assert len(set(names)) < len(names) + assert len(set(labels)) == len(labels) + + def test_single_binding_keeps_plain_names(self): + # WHEN + analysis = make_analysis(two_bindings=False) + + # EXPECT no model prefix, since there is nothing to disambiguate + labels = [analysis.parameter_label(p) for p in analysis._get_chain_parameters()] + assert labels == ['Width line_c0', 'Width line_c1'] + + def test_parameter_from_outside_the_analysis_keeps_its_name(self, analysis): + # WHEN a parameter belongs to none of the binding models + from easyscience.variable import Parameter + + stranger = Parameter(name='Width line_c0', value=1.0) + + # EXPECT it is returned unqualified rather than mislabelled + assert analysis.parameter_label(stranger) == 'Width line_c0' + + def test_models_without_a_display_name_fall_back_to_the_unique_name(self): + # WHEN two colliding models have no display name to tell them apart + model_a = sm.Polynomial(coefficients=[0.1, 0.35], x_unit='1/angstrom', y_unit='meV') + model_b = sm.Polynomial(coefficients=[2.0, -0.3], x_unit='1/angstrom', y_unit='meV') + analysis = edyn.ParameterAnalysis( + parameters=make_dataset(), + bindings=[ + edyn.FitBinding(model=model_a, targets='Lorentzian width'), + edyn.FitBinding(model=model_b, targets='Lorentzian area'), + ], + ) + + # THEN + labels = [analysis.parameter_label(p) for p in analysis._get_chain_parameters()] + + # EXPECT still unambiguous, which is what matters + assert len(set(labels)) == len(labels) diff --git a/tests/unit/easydynamics/analysis/test_posterior.py b/tests/unit/easydynamics/analysis/test_posterior.py index dc9a5e4a1..bd9ad1d8e 100644 --- a/tests/unit/easydynamics/analysis/test_posterior.py +++ b/tests/unit/easydynamics/analysis/test_posterior.py @@ -306,3 +306,19 @@ def test_repr_contains_the_parameter_name(self): # EXPECT assert 'Gaussian area' in text assert 'median' in text + + +class TestPosteriorSummaryContainer: + def test_len_and_iteration(self): + # WHEN + parameters = [make_parameter(name='a'), make_parameter(name='b')] + summary = summarize_draws(np.zeros((7, 2)), ['a', 'b'], parameters) + + # EXPECT + assert len(summary) == 2 + assert [entry.name for entry in summary] == ['a', 'b'] + assert len(summary.entries) == 2 + + def test_repr_with_no_entries(self): + # EXPECT + assert 'no parameters' in repr(summarize_draws(np.zeros((3, 0)), [], [])) diff --git a/tests/unit/easydynamics/utils/test_posterior_plotting.py b/tests/unit/easydynamics/utils/test_posterior_plotting.py index 64d0129e9..abc6caa69 100644 --- a/tests/unit/easydynamics/utils/test_posterior_plotting.py +++ b/tests/unit/easydynamics/utils/test_posterior_plotting.py @@ -65,6 +65,13 @@ def test_one_dimensional_draws_raise(self): with pytest.raises(ValueError, match='two-dimensional'): plot_trace(draws=np.zeros(10), names=['a']) + def test_labels_carry_units(self, draws): + # WHEN + fig = plot_trace(draws=draws, names=['a', 'b', 'c'], units=['meV', 'm^2/s', '']) + + # EXPECT the real units are shown, and an empty one is skipped + assert [axis.get_ylabel() for axis in fig.axes] == ['a (meV)', 'b (m^2/s)', 'c'] + class TestPlotCorner: def test_grid_is_square_in_the_parameter_count(self, draws): @@ -86,6 +93,24 @@ def test_mismatched_names_raise(self, draws): with pytest.raises(ValueError, match='one entry per column'): plot_corner(draws=draws, names=['a']) + def test_diagonal_panel_is_labelled_as_counts(self, draws): + # WHEN + fig = plot_corner(draws=draws, names=['a', 'b', 'c']) + + # EXPECT the top-left panel says what its vertical axis actually is. It is a histogram, so + # the parameter is on the x axis and labelling y with the parameter name would be wrong. + assert fig.axes[0].get_ylabel() == 'counts' + + def test_units_are_appended_to_labels(self, draws): + # WHEN + fig = plot_corner(draws=draws, names=['a', 'b', 'c'], units=['meV', '', 'dimensionless']) + + # EXPECT the real unit is shown, and empty or dimensionless ones are skipped + bottom_row = fig.axes[-3:] + assert bottom_row[0].get_xlabel() == 'a (meV)' + assert bottom_row[1].get_xlabel() == 'b' + assert bottom_row[2].get_xlabel() == 'c' + class TestPlotPosteriorPredictive: def test_returns_a_figure_with_data_and_band(self): @@ -146,3 +171,17 @@ def test_band_widens_with_the_credible_interval(self): 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 + + def test_axis_labels_are_set_when_given(self): + # WHEN + fig = plot_posterior_predictive( + x=np.zeros(4), + y=np.zeros(4), + predictions=np.zeros((5, 4)), + xlabel='Energy (meV)', + ylabel='Intensity', + ) + + # EXPECT + assert fig.axes[0].get_xlabel() == 'Energy (meV)' + assert fig.axes[0].get_ylabel() == 'Intensity' From 0c3c689ca1a13cf383f52d74ce7b502e7925edf5 Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Thu, 13 Aug 2026 14:14:05 +0200 Subject: [PATCH 04/29] Qualify parameter labels by model name, and cover the remaining branches Two fixes found by writing the tests codecov asked for. ParameterAnalysis qualified an ambiguous parameter with the owning model's display_name, but for several models -- the diffusion models among them -- display_name is the class name, so two models constructed as name='Diffusion A' and name='Diffusion B' both came back as "BrownianTranslationalDiffusion" and the label did not disambiguate anything. It now uses the model's name, matching the choice to report parameters under their name rather than their display name, and falls back to the unique name only when the names collide too. The rest is test coverage for branches that were reachable but untested: the label fallbacks, the BUMPS outlier crash being re-raised as a degeneracy hint, a chain column that matches no parameter, loading a chain through its sidecar, the mixin's unimplemented hooks, and the scientific-notation exponent being folded into an axis label. Co-Authored-By: Claude Opus 5 (1M context) --- .../analysis/parameter_analysis.py | 11 ++-- .../test_parameter_analysis_bayesian.py | 57 +++++++++++++++++++ .../utils/test_posterior_plotting.py | 16 ++++++ 3 files changed, 80 insertions(+), 4 deletions(-) diff --git a/src/easydynamics/analysis/parameter_analysis.py b/src/easydynamics/analysis/parameter_analysis.py index 47020b1d8..dee5563b3 100644 --- a/src/easydynamics/analysis/parameter_analysis.py +++ b/src/easydynamics/analysis/parameter_analysis.py @@ -291,7 +291,9 @@ def parameter_label(self, parameter: Parameter) -> str: models that already name their parameters after themselves, keep their plain names -- these get long quickly, and there is nothing to disambiguate. - The prefix is the model's display name, unless two models share that too, in which case the + The prefix is the model's name, matching the choice to report parameters under their name + rather than their display name -- for several models the display name is just the class + name, which would not tell two of them apart. If two models share a name as well, the unique name is used: a label that does not actually disambiguate is worse than a long one. Parameters @@ -317,10 +319,11 @@ def parameter_label(self, parameter: Parameter) -> str: return parameter.name owner = owners[0] - display_names = [model.display_name for model in models.values()] - if owner.display_name is None or display_names.count(owner.display_name) > 1: + owner_name = getattr(owner, 'name', None) or owner.display_name + model_names = [getattr(m, 'name', None) or m.display_name for m in models.values()] + if owner_name is None or model_names.count(owner_name) > 1: return f'{owner.unique_name}: {parameter.name}' - return f'{owner.display_name}: {parameter.name}' + return f'{owner_name}: {parameter.name}' def plot( self, names: str | list[str] | None = None, **kwargs: dict[str, Any] diff --git a/tests/unit/easydynamics/analysis/test_parameter_analysis_bayesian.py b/tests/unit/easydynamics/analysis/test_parameter_analysis_bayesian.py index 6dfc504e0..5dd1286bf 100644 --- a/tests/unit/easydynamics/analysis/test_parameter_analysis_bayesian.py +++ b/tests/unit/easydynamics/analysis/test_parameter_analysis_bayesian.py @@ -266,3 +266,60 @@ def test_models_without_a_display_name_fall_back_to_the_unique_name(self): # EXPECT still unambiguous, which is what matters assert len(set(labels)) == len(labels) + + def test_colliding_names_with_distinct_models_use_the_display_name(self): + # WHEN two diffusion models are bound to different targets. Their parameters are not named + # after the model, so the names collide while the model names do not. + analysis = edyn.ParameterAnalysis( + parameters=make_dataset(), + bindings=[ + edyn.FitBinding( + model=sm.BrownianTranslationalDiffusion( + name='Diffusion A', diffusion_coefficient=2.4e-9, scale=0.5 + ), + targets={'width': 'Lorentzian width'}, + ), + edyn.FitBinding( + model=sm.BrownianTranslationalDiffusion( + name='Diffusion B', diffusion_coefficient=2.4e-9, scale=0.5 + ), + targets={'area': 'Lorentzian area'}, + ), + ], + ) + + # THEN + parameters = analysis._get_chain_parameters() + labels = [analysis.parameter_label(p) for p in parameters] + + # EXPECT the model's name resolves the collision + assert len({p.name for p in parameters}) < len(parameters) + assert len(set(labels)) == len(labels) + assert any(label.startswith('Diffusion A: ') for label in labels) + assert any(label.startswith('Diffusion B: ') for label in labels) + + def test_ambiguous_name_owned_by_no_model_keeps_its_name(self): + # WHEN a parameter shares an ambiguous name but belongs to none of the models + from easyscience.variable import Parameter + + analysis = edyn.ParameterAnalysis( + parameters=make_dataset(), + bindings=[ + edyn.FitBinding( + model=sm.Polynomial( + coefficients=[0.1, 0.35], x_unit='1/angstrom', y_unit='meV', name='Line' + ), + targets='Lorentzian width', + ), + edyn.FitBinding( + model=sm.Polynomial( + coefficients=[2.0, -0.3], x_unit='1/angstrom', y_unit='meV', name='Line' + ), + targets='Lorentzian area', + ), + ], + ) + stranger = Parameter(name='Line_c0', value=1.0) + + # EXPECT it falls back to the plain name rather than claiming an owner + assert analysis.parameter_label(stranger) == 'Line_c0' diff --git a/tests/unit/easydynamics/utils/test_posterior_plotting.py b/tests/unit/easydynamics/utils/test_posterior_plotting.py index abc6caa69..b80f69d00 100644 --- a/tests/unit/easydynamics/utils/test_posterior_plotting.py +++ b/tests/unit/easydynamics/utils/test_posterior_plotting.py @@ -185,3 +185,19 @@ def test_axis_labels_are_set_when_given(self): # EXPECT assert fig.axes[0].get_xlabel() == 'Energy (meV)' assert fig.axes[0].get_ylabel() == 'Intensity' + + +class TestScientificNotation: + def test_shared_exponent_is_folded_into_the_label(self): + # WHEN the values are small enough that matplotlib factors out an exponent, which it parks + # on top of the axis label + draws = np.random.default_rng(0).normal(size=(200, 2)) * 1e-8 + 1.15e-8 + + fig = plot_corner(draws=draws, names=['D', 'scale'], units=['m^2/s', '']) + + # EXPECT the exponent and the unit share one parenthetical, and the overlapping offset + # text is hidden + xlabel = fig.axes[-2].get_xlabel() + assert xlabel.startswith('D (1e') + assert 'm^2/s' in xlabel + assert not fig.axes[-2].xaxis.get_offset_text().get_visible() From 46d745a4a73e8025c37e02e490fab67cfdb5ff22 Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Thu, 13 Aug 2026 14:26:55 +0200 Subject: [PATCH 05/29] Warm the tutorial data cache before running notebooks in parallel The notebook tests run with '-n auto', and five of the notebooks fetch vanadium_data_example.h5 through pooch. On a cold cache the workers race: one is still writing the file into the cache while another opens it, which fails on Windows with "PermissionError: Permission denied". This failed twice in a row on windows-latest, always on that file, always with the other sixteen notebooks passing. The race is pre-existing, but adding a fifth notebook that wants the same file, and lengthening tutorial 1, made it reliable rather than rare. Fetching every tutorial data file once, before the parallel run starts, leaves the workers with nothing to do but read, which is safe. The prefetch reads the URLs and hashes out of the notebooks themselves, so it cannot drift from what they actually download, and it never fails the run: a file it cannot fetch is left to the notebook that needs it, which reports the problem with far more context. Co-Authored-By: Claude Opus 5 (1M context) --- pixi.toml | 8 ++- tools/prefetch_tutorial_data.py | 92 +++++++++++++++++++++++++++++++++ 2 files changed, 99 insertions(+), 1 deletion(-) create mode 100644 tools/prefetch_tutorial_data.py diff --git a/pixi.toml b/pixi.toml index f26b4fb7e..db46523e1 100644 --- a/pixi.toml +++ b/pixi.toml @@ -100,7 +100,13 @@ user = { features = ['py-max', 'user'] } unit-tests = 'python -m pytest tests/unit/ --color=yes -v' functional-tests = 'python -m pytest tests/functional/ --color=yes -v' integration-tests = 'python -m pytest tests/integration/ --color=yes -n auto -v' -notebook-tests = 'python -m pytest --nbmake docs/docs/tutorials/ --nbmake-timeout=1200 --color=yes -n auto -v' +# Warm the pooch cache first. Several notebooks fetch the same file, and running them with +# '-n auto' has the workers race: one writes the file while another opens it, which fails on +# Windows. Fetching up front leaves the parallel run with nothing to do but read. +prefetch-tutorial-data = 'python tools/prefetch_tutorial_data.py' +notebook-tests = { cmd = 'python -m pytest --nbmake docs/docs/tutorials/ --nbmake-timeout=1200 --color=yes -n auto -v', depends-on = [ + 'prefetch-tutorial-data', +] } test = { depends-on = ['unit-tests'] } diff --git a/tools/prefetch_tutorial_data.py b/tools/prefetch_tutorial_data.py new file mode 100644 index 000000000..839b1897a --- /dev/null +++ b/tools/prefetch_tutorial_data.py @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +""" +Download every data file the tutorial notebooks fetch, once, before they are run. + +The notebooks are executed in parallel with ``pytest -n auto``, and several of them fetch the same +file through ``pooch``. On a cold cache the workers race: one is still writing the file into the +cache while another tries to open it, which fails on Windows with a permission error. Fetching +everything up front leaves the parallel run with nothing to do but read. + +Run as ``python tools/prefetch_tutorial_data.py``; it is wired into the ``notebook-tests`` task. +""" + +from __future__ import annotations + +import json +import re +import sys +from pathlib import Path + +import pooch + +TUTORIALS = Path(__file__).resolve().parent.parent / 'docs' / 'docs' / 'tutorials' + +# Matches the pooch.retrieve(url=..., known_hash=...) calls the notebooks use, in either order. +URL_PATTERN = re.compile(r"url\s*=\s*f?['\"]([^'\"]+)['\"]") +HASH_PATTERN = re.compile(r"known_hash\s*=\s*['\"]([^'\"]+)['\"]") + + +def find_downloads() -> dict[str, str]: + """ + Collect the ``(url, known_hash)`` pairs the notebooks fetch. + + Returns + ------- + dict[str, str] + Mapping of URL to expected hash, deduplicated across notebooks. + """ + downloads: dict[str, str] = {} + for notebook in sorted(TUTORIALS.glob('*.ipynb')): + cells = json.loads(notebook.read_text(encoding='utf-8'))['cells'] + for cell in cells: + if cell['cell_type'] != 'code': + continue + source = ''.join(cell['source']) + if 'pooch.retrieve' not in source: + continue + urls = URL_PATTERN.findall(source) + hashes = HASH_PATTERN.findall(source) + # Only pairs are usable; a templated URL without a literal hash is skipped rather than + # guessed at, and the notebook will simply fetch it itself. + for url, known_hash in zip(urls, hashes, strict=False): + downloads[url] = known_hash + return downloads + + +def main() -> int: + """ + Fetch every tutorial data file into the pooch cache. + + Deliberately never fails: this only warms a cache. A file that cannot be fetched here is left + to the notebook that needs it, which reports the problem with far more context than this script + could, and which is where the failure belongs. + + Returns + ------- + int + Always zero. + """ + downloads = find_downloads() + if not downloads: + sys.stdout.write('No tutorial downloads found.\n') + return 0 + + failures = 0 + for url, known_hash in downloads.items(): + name = url.rsplit('/', 1)[-1] + try: + pooch.retrieve(url=url, known_hash=known_hash) + except Exception as error: # noqa: BLE001 - report and continue, the notebook will retry + failures += 1 + sys.stdout.write(f'could not prefetch {name}, leaving it to the notebook: {error}\n') + else: + sys.stdout.write(f'cached {name}\n') + + sys.stdout.write(f'{len(downloads) - failures}/{len(downloads)} tutorial data files ready.\n') + return 0 + + +if __name__ == '__main__': + sys.exit(main()) From ff50a2c2146044daeadf6f47711c73034463fc40 Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Thu, 13 Aug 2026 15:09:04 +0200 Subject: [PATCH 06/29] Rebuild the fitter when a binding changes shape, and stabilise the integration tests Two problems found while reviewing the previous commits. Caching the MultiFitter on ParameterAnalysis introduced a regression. A FitBinding can be edited in place -- binding.targets = ... -- which ParameterAnalysis cannot observe. Changing the number of targets left the cached fitter holding one fit function against two datasets, and fit() died with "FitError: list index out of range". It rebuilt every call before, so this worked previously. The targets the fitter was built for are now recorded and compared, which is enough to catch an edit that cannot be observed directly. The integration tests then failed in CI on macOS, inside BUMPS' outlier removal, on an identifiable model. That matters beyond the test: the error message claimed the crash means degenerate parameters, and this shows short chains do it too. The message now names both causes, and the integration tests switch the outlier removal off, as they already do for the burn-point trimming. Co-Authored-By: Claude Opus 5 (1M context) --- .../analysis/bayesian_sampling.py | 16 ++++--- .../analysis/parameter_analysis.py | 46 ++++++++++++++++++- .../fitting/test_bayesian_sampling.py | 13 ++++-- .../fitting/test_bayesian_sampling_multi_q.py | 11 +++-- .../test_parameter_analysis_bayesian.py | 42 +++++++++++++++++ 5 files changed, 110 insertions(+), 18 deletions(-) diff --git a/src/easydynamics/analysis/bayesian_sampling.py b/src/easydynamics/analysis/bayesian_sampling.py index c9adb11f4..e874ec218 100644 --- a/src/easydynamics/analysis/bayesian_sampling.py +++ b/src/easydynamics/analysis/bayesian_sampling.py @@ -438,14 +438,16 @@ def _run_sampling( sampler = self._get_or_build_sampler(reuse_sampler=reuse_sampler) results = run(sampler) except IndexError as error: - # BUMPS' own outlier removal indexes past the end of its buffer when chains - # scatter wildly, which in practice means the model is not identifiable. The bare - # IndexError says nothing useful, so point at the likely cause instead. + # BUMPS' own outlier removal indexes past the end of its buffer. Seen both when + # chains scatter because the model is not identifiable, and on short chains where + # its buffer has too few generations to work with. The bare IndexError says + # nothing useful, so name both causes and the way out. raise RuntimeError( - 'The BUMPS sampler failed while removing outlier chains. This usually means ' - 'the chains scattered because two or more free parameters are degenerate, so ' - 'the data cannot determine them separately. Check for degenerate parameters ' - "and fix one of them, or retry with sampler_kwargs={'outliers': 'none'}." + '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) diff --git a/src/easydynamics/analysis/parameter_analysis.py b/src/easydynamics/analysis/parameter_analysis.py index dee5563b3..4a72d5b96 100644 --- a/src/easydynamics/analysis/parameter_analysis.py +++ b/src/easydynamics/analysis/parameter_analysis.py @@ -101,6 +101,9 @@ def __init__( """ self._init_bayesian_state() + # Which targets the cached fitter was built for, so an in-place edit of a FitBinding is + # noticed even though it cannot be observed directly. + self._fitter_targets = None super().__init__(display_name=display_name, unique_name=unique_name) @@ -178,7 +181,8 @@ def fit(self) -> FitResults: The results of the fit """ - xs, ys, ws, _, _ = self._build_fit_inputs() + xs, ys, ws, _, models = self._build_fit_inputs() + self._invalidate_fitter_if_targets_changed(models) return self.fitter.fit(x=xs, y=ys, weights=ws) def _build_fit_inputs(self) -> tuple[list, list, list, list, list]: @@ -250,8 +254,45 @@ def _build_bayesian_fitter(self) -> MultiFitter: A MultiFitter over the per-target models and fit functions. """ _, _, _, funcs, models = self._build_fit_inputs() + self._fitter_targets = self._target_signature(models) return MultiFitter(fit_objects=models, fit_functions=funcs) + @staticmethod + def _target_signature(models: list) -> tuple: + """ + Summarize which models the fitter was built for, in target order. + + Parameters + ---------- + models : list + The model behind each fit target. + + Returns + ------- + tuple + A comparable signature of the current targets. + """ + return tuple(model.unique_name for model in models) + + def _invalidate_fitter_if_targets_changed(self, models: list) -> None: + """ + Rebuild the cached fitter when the bindings no longer resolve to the same targets. + + A FitBinding can be edited in place -- ``binding.targets = ...`` -- which this object + cannot observe. Doing so changes how many datasets there are, while the cached MultiFitter + still holds the old fit functions, and the fit then dies deep inside the minimizer. Compare + the targets the fitter was built for against the current ones instead. + + Parameters + ---------- + models : list + The model behind each fit target, as currently resolved. + """ + if self._fitter is None: + return + if self._target_signature(models) != getattr(self, '_fitter_targets', None): + self._invalidate_fitter() + def _get_sampling_data(self) -> tuple[list, list, list]: """ Get the per-target data to bind to the Sampler. @@ -261,7 +302,8 @@ def _get_sampling_data(self) -> tuple[list, list, list]: tuple[list, list, list] The ``(x, y, weights)`` triple, one entry per fit target. """ - xs, ys, ws, _, _ = self._build_fit_inputs() + xs, ys, ws, _, models = self._build_fit_inputs() + self._invalidate_fitter_if_targets_changed(models) return xs, ys, ws def _get_chain_parameters(self) -> list[Parameter]: diff --git a/tests/integration/fitting/test_bayesian_sampling.py b/tests/integration/fitting/test_bayesian_sampling.py index 704a37ea2..0db56787e 100644 --- a/tests/integration/fitting/test_bayesian_sampling.py +++ b/tests/integration/fitting/test_bayesian_sampling.py @@ -4,9 +4,9 @@ """ 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. +These are slow by nature. Two BUMPS options are switched off deliberately: its burn-point trimming, +which re-runs a convergence detector on every call, and its outlier removal, which indexes past the +end of its own buffer on chains as short as these. Neither affects the sampling itself. """ import warnings @@ -33,7 +33,10 @@ 'samples': 2000, 'burn': 100, 'thin': 2, - 'sampler_kwargs': {'trim': False}, + # 'trim': BUMPS' burn-point detector re-runs on every call and is not worth paying + # for here. 'outliers': its outlier removal indexes past the end of its own buffer on + # chains this short, which has failed in CI; the sampling itself is unaffected. + 'sampler_kwargs': {'trim': False, 'outliers': 'none'}, } @@ -136,7 +139,7 @@ def test_extend_grows_the_chain(self, sampled_analysis): with warnings.catch_warnings(): warnings.simplefilter('ignore') extended = sampled_analysis.extend_sampling( - additional_samples=500, thin=2, sampler_kwargs={'trim': False} + additional_samples=500, thin=2, sampler_kwargs={'trim': False, 'outliers': 'none'} ) # EXPECT diff --git a/tests/integration/fitting/test_bayesian_sampling_multi_q.py b/tests/integration/fitting/test_bayesian_sampling_multi_q.py index 5870aab49..56a883f97 100644 --- a/tests/integration/fitting/test_bayesian_sampling_multi_q.py +++ b/tests/integration/fitting/test_bayesian_sampling_multi_q.py @@ -4,9 +4,9 @@ """ Integration tests running real BUMPS DREAM chains through Analysis and ParameterAnalysis. -Slow by nature, and run with ``sampler_kwargs={'trim': False}`` for the same reason as the -single-Q integration tests: 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. +Slow by nature, and with the same two BUMPS options switched off as the single-Q integration tests: +its burn-point trimming, which re-runs a convergence detector on every call, and its outlier +removal, which indexes past the end of its own buffer on chains as short as these. """ import warnings @@ -29,7 +29,10 @@ 'samples': 2000, 'burn': 100, 'thin': 2, - 'sampler_kwargs': {'trim': False}, + # 'trim': BUMPS' burn-point detector re-runs on every call and is not worth paying + # for here. 'outliers': its outlier removal indexes past the end of its own buffer on + # chains this short, which has failed in CI; the sampling itself is unaffected. + 'sampler_kwargs': {'trim': False, 'outliers': 'none'}, } diff --git a/tests/unit/easydynamics/analysis/test_parameter_analysis_bayesian.py b/tests/unit/easydynamics/analysis/test_parameter_analysis_bayesian.py index 5dd1286bf..1bb827aa5 100644 --- a/tests/unit/easydynamics/analysis/test_parameter_analysis_bayesian.py +++ b/tests/unit/easydynamics/analysis/test_parameter_analysis_bayesian.py @@ -323,3 +323,45 @@ def test_ambiguous_name_owned_by_no_model_keeps_its_name(self): # EXPECT it falls back to the plain name rather than claiming an owner assert analysis.parameter_label(stranger) == 'Line_c0' + + +class TestInPlaceBindingEdits: + def test_changing_the_number_of_targets_rebuilds_the_fitter(self): + # WHEN a binding is edited in place so that it resolves to two targets instead of one. + # ParameterAnalysis cannot observe this, and the cached fitter would otherwise still hold + # one fit function against two datasets, which dies inside the minimizer. + binding = edyn.FitBinding( + model=sm.BrownianTranslationalDiffusion( + name='Brownian', + lorentzian_name='Lorentzian', + diffusion_coefficient=2.4e-9, + scale=0.5, + ), + targets={'width': 'Lorentzian width'}, + ) + analysis = edyn.ParameterAnalysis(parameters=make_dataset(), bindings=[binding]) + assert len(analysis.fit()) == 1 + + binding.targets = {'width': 'Lorentzian width', 'area': 'Lorentzian area'} + + # EXPECT the fit follows the binding rather than failing on a stale fitter + assert len(analysis.fit()) == 2 + + def test_shrinking_the_targets_also_rebuilds(self): + # WHEN + binding = edyn.FitBinding( + model=sm.BrownianTranslationalDiffusion( + name='Brownian', + lorentzian_name='Lorentzian', + diffusion_coefficient=2.4e-9, + scale=0.5, + ), + targets={'width': 'Lorentzian width', 'area': 'Lorentzian area'}, + ) + analysis = edyn.ParameterAnalysis(parameters=make_dataset(), bindings=[binding]) + assert len(analysis.fit()) == 2 + + binding.targets = {'width': 'Lorentzian width'} + + # EXPECT + assert len(analysis.fit()) == 1 From b28b9e4dd5b4d82af4afc3ba70cb8eeb9c00bf81 Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Thu, 13 Aug 2026 15:51:14 +0200 Subject: [PATCH 07/29] Address the review findings on the sampling API Six issues found reviewing the previous commits. The sidecar could be written with the wrong labels. A subset run built the name map inside the block that holds the other parameters fixed, where nothing looks ambiguous, so a multi-Q chain recorded unqualified names that no longer matched on reload. The map is now built outside that block, where the free set is the user's real one. extend_sampling() accepted a different parameter subset. BUMPS resumes from a stored chain whose width is fixed, so that could only fail deep inside the sampler; it is now refused up front. The IndexError relabelling was unconditional, so an IndexError from this package would have been reported as a BUMPS modelling problem. It now only applies when the traceback passes through bumps. Labelling a chain was quadratic in the parameter count: collecting the parameters and scanning for their owner both happened per parameter, and each walks every sub-model. 75 parameters took 0.39 s, and every summary and plot pays it. The parameters are now collected once per pass, and Analysis keeps an owner index alongside its analysis list. The same case now measures at 0.00 s. Asking an Analysis for a summary after sampling independently reported that nothing had been sampled, moments after it had. It now says where the chains actually are. Applying bounds many orders of magnitude wider than the parameter is still allowed -- it is what the fit implied -- but no longer silent, so a scripted apply() cannot hide a degeneracy the table would have shown. Co-Authored-By: Claude Opus 5 (1M context) --- src/easydynamics/analysis/analysis.py | 77 ++++++++-- .../analysis/bayesian_sampling.py | 136 ++++++++++++++++-- src/easydynamics/analysis/posterior.py | 48 ++++++- .../analysis/test_analysis1d_bayesian.py | 88 ++++++++++-- .../analysis/test_analysis_bayesian.py | 24 ++++ .../easydynamics/analysis/test_posterior.py | 37 +++++ 6 files changed, 377 insertions(+), 33 deletions(-) diff --git a/src/easydynamics/analysis/analysis.py b/src/easydynamics/analysis/analysis.py index 4edc080cf..25824d959 100644 --- a/src/easydynamics/analysis/analysis.py +++ b/src/easydynamics/analysis/analysis.py @@ -123,6 +123,8 @@ def __init__( self._analysis_list: list[Analysis1d] = [] self._analysis_list_is_dirty = True + # Rebuilt with the analysis list; see _parameter_owner_index. + self._owner_index = None self._init_bayesian_state() super().__init__( display_name=display_name, @@ -740,6 +742,7 @@ def _on_experiment_changed(self) -> None: """ super()._on_experiment_changed() self._analysis_list_is_dirty = True + self._owner_index = None self._invalidate_fitter() def _on_sample_model_changed(self) -> None: @@ -748,6 +751,7 @@ def _on_sample_model_changed(self) -> None: """ super()._on_sample_model_changed() self._analysis_list_is_dirty = True + self._owner_index = None self._invalidate_fitter() def _on_instrument_model_changed(self) -> None: @@ -756,6 +760,7 @@ def _on_instrument_model_changed(self) -> None: """ super()._on_instrument_model_changed() self._analysis_list_is_dirty = True + self._owner_index = None self._invalidate_fitter() def _on_convolution_settings_changed(self) -> None: @@ -764,6 +769,7 @@ def _on_convolution_settings_changed(self) -> None: """ super()._on_convolution_settings_changed() self._analysis_list_is_dirty = True + self._owner_index = None self._invalidate_fitter() def _ensure_analysis_list_current(self) -> None: @@ -778,6 +784,7 @@ def _create_analysis_list(self) -> None: experiment, sample model, and instrument model. """ self._analysis_list = [] + self._owner_index = None for Q_index in range(len(self.Q)): # The ConvolutionSettings object is shared so user changes reach every Q index; # plan validity is tracked per convolver, not on the settings object. @@ -871,16 +878,68 @@ def parameter_label(self, parameter: Parameter) -> str: """ if not self._name_is_ambiguous(parameter): return parameter.name - owners = [ - analysis1d.Q_index - for analysis1d in self.analysis_list - if any( - p.unique_name == parameter.unique_name for p in analysis1d.get_free_parameters() - ) - ] - if len(owners) != 1: + owner = self._parameter_owner_index().get(parameter.unique_name) + if owner is None: return parameter.name - return f'{parameter.name} (Q_index={owners[0]})' + return f'{parameter.name} (Q_index={owner})' + + def _parameter_owner_index(self) -> dict[str, int]: + """ + Map each parameter to the Q index that owns it. + + Built once per analysis list and reused, because scanning the list for every parameter + makes labelling a chain quadratic in the parameter count -- seconds, for a dataset with + many Q values. Built from all parameters rather than only the free ones, so that fixing a + parameter cannot leave the map stale. + + Returns + ------- + dict[str, int] + Mapping of parameter ``unique_name`` to owning Q index. Parameters shared by more than + one Q index are left out, since no single Q identifies them. + """ + self._ensure_analysis_list_current() + if self._owner_index is None: + owners: dict[str, int | None] = {} + for analysis1d in self._analysis_list: + for parameter in analysis1d.get_all_parameters(): + if parameter.unique_name in owners: + owners[parameter.unique_name] = None + else: + owners[parameter.unique_name] = analysis1d.Q_index + self._owner_index = { + name: q_index for name, q_index in owners.items() if q_index is not None + } + return self._owner_index + + def _require_posterior_result(self) -> SamplingResults: + """ + Get the stored sampling results, pointing at the per-Q chains when those are what exist. + + Sampling independently stores a chain on each Analysis1d rather than here, so asking this + object for a summary afterwards would otherwise report that no sampling has happened, right + after it has. + + Returns + ------- + SamplingResults + The results of the most recent simultaneous run. + + Raises + ------ + RuntimeError + If no simultaneous sampling has been run on this Analysis. + """ + if self._posterior_result is None and any( + analysis1d.posterior_result is not None for analysis1d in self.analysis_list + ): + raise RuntimeError( + 'This Analysis holds no chain of its own, but its Q indices do: sampling with ' + "fit_method='independent' gives each Q its own chain. Use " + 'analysis.analysis_list[Q_index] to summarize or plot one of them, or sample with ' + "fit_method='simultaneous' to get a single chain over all Q." + ) + return super()._require_posterior_result() def _prepare_for_sampling(self) -> None: """ diff --git a/src/easydynamics/analysis/bayesian_sampling.py b/src/easydynamics/analysis/bayesian_sampling.py index e874ec218..1db37202d 100644 --- a/src/easydynamics/analysis/bayesian_sampling.py +++ b/src/easydynamics/analysis/bayesian_sampling.py @@ -16,6 +16,7 @@ import json import warnings +from contextlib import contextmanager from pathlib import Path from typing import TYPE_CHECKING from typing import Any @@ -33,6 +34,7 @@ if TYPE_CHECKING: import os from collections.abc import Callable + from collections.abc import Iterator from easyscience.fitting.fitter import Fitter from easyscience.fitting.sampler import SamplingResults @@ -89,6 +91,8 @@ def _init_bayesian_state(self) -> None: self._bayesian_sampler = None self._bayesian_sampler_is_dirty = True self._posterior_result = None + # Set only while a bulk operation holds the parameter list, see _bulk_parameter_access. + self._chain_parameters_cache = None # Maps a chain column's unique_name to the parameter name it had when saved. Only populated # by load_chain, because unique_names are per-session and do not survive a round trip. self._chain_name_map = {} @@ -256,10 +260,11 @@ def suggest_bounds( BoundsSuggestions The proposed bounds, which must be applied explicitly. """ - parameters = self._get_chain_parameters() + with self._bulk_parameter_access() as parameters: + labels = [self.parameter_label(parameter) for parameter in parameters] return suggest_bounds_for_parameters( parameters, - labels=[self.parameter_label(parameter) for parameter in parameters], + labels=labels, n_sigma=n_sigma, relative_pad=relative_pad, absolute_floor=absolute_floor, @@ -417,9 +422,11 @@ def _run_sampling( 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, which points at degenerate - parameters. + If the BUMPS sampler fails while removing outlier chains. """ held_fixed = self._resolve_parameters_to_hold_fixed(parameters) self._warn_about_held_parameters(held_fixed) @@ -431,6 +438,9 @@ def _run_sampling( chain_parameters = self._get_chain_parameters() saved_values = [(p, p.value) for p in chain_parameters] + if reuse_sampler: + self._verify_chain_shape_unchanged(chain_parameters) + fitter = self.fitter original_minimizer = fitter.minimizer.enum fitter.switch_minimizer(AvailableMinimizers.Bumps) @@ -438,6 +448,8 @@ def _run_sampling( sampler = self._get_or_build_sampler(reuse_sampler=reuse_sampler) results = run(sampler) except IndexError as error: + if not _raised_inside_bumps(error): + raise # BUMPS' own outlier removal indexes past the end of its buffer. Seen both when # chains scatter because the model is not identifiable, and on short chains where # its buffer has too few generations to work with. The bare IndexError says @@ -454,8 +466,10 @@ def _run_sampling( for parameter, value in saved_values: parameter.value = value - # A fresh chain is labelled with this session's unique names, so any mapping left over - # from a loaded chain no longer applies. + # Labelled outside the block above, so that a subset run records the same 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. + with self._bulk_parameter_access(): self._chain_name_map = { parameter.unique_name: self.parameter_label(parameter) for parameter in chain_parameters @@ -465,6 +479,35 @@ def _run_sampling( self._warn_about_bounds_occupancy(results, self._resolve_chain_parameters(results)) return results + def _verify_chain_shape_unchanged(self, chain_parameters: list[Parameter]) -> None: + """ + Check that an extension keeps the chain's columns. + + BUMPS resumes from a stored state whose width is fixed, so a run that would add or drop a + parameter cannot continue that chain. Caught here rather than left to fail obscurely inside + the sampler. + + Parameters + ---------- + chain_parameters : list[Parameter] + The parameters that would form the chain for this run. + + Raises + ------ + ValueError + If the number of parameters differs from the existing chain's. + """ + if self._posterior_result is None: + return + existing = self._posterior_result.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_posterior() instead.' + ) + def _get_or_build_sampler(self, reuse_sampler: bool) -> Sampler: """ Get the cached Sampler, rebuilding it if the data or model changed. @@ -948,7 +991,7 @@ def _resolve_chain_parameters(self, results: SamplingResults) -> list[Parameter list[Parameter | None] The parameter for each column, or None where no match could be made. """ - parameters = self._get_chain_parameters() + parameters = self._chain_parameters() by_unique_name = {p.unique_name: p for p in parameters} by_label = {self.parameter_label(p): p for p in parameters} resolved = [] @@ -997,9 +1040,45 @@ def _name_is_ambiguous(self, parameter: Parameter) -> bool: bool True when at least one other chain parameter has the same name. """ - names = [p.name for p in self._get_chain_parameters()] + names = [p.name for p in self._chain_parameters()] return names.count(parameter.name) > 1 + def _chain_parameters(self) -> list[Parameter]: + """ + Get the chain parameters, reusing the list when a bulk operation is in progress. + + Collecting them walks every sub-model, so labelling a chain one parameter at a time is + quadratic in the parameter count -- seconds, for a dataset with many Q values. + + Returns + ------- + list[Parameter] + The free parameters of the underlying model(s). + """ + if self._chain_parameters_cache is not None: + return self._chain_parameters_cache + return self._get_chain_parameters() + + @contextmanager + def _bulk_parameter_access(self) -> Iterator[list[Parameter]]: + """ + Collect the chain parameters once for the duration of a block. + + Scoped rather than stored, so the cache cannot outlive the operation that wanted it and go + stale against a changed model. + + Yields + ------ + list[Parameter] + The chain parameters, also served to :meth:`_chain_parameters` inside the block. + """ + previous = self._chain_parameters_cache + self._chain_parameters_cache = self._get_chain_parameters() + try: + yield self._chain_parameters_cache + finally: + self._chain_parameters_cache = previous + def _chain_units(self, results: SamplingResults) -> list[str]: """ Get the unit of each column of the chain. @@ -1033,13 +1112,40 @@ def _chain_display_names(self, results: SamplingResults) -> list[str]: list[str] One label per column of the chain. """ - resolved = self._resolve_chain_parameters(results) - return [ - self._chain_name_map.get(unique_name, unique_name) - if parameter is None - else self.parameter_label(parameter) - for unique_name, parameter in zip(results.param_names, resolved, strict=True) - ] + with self._bulk_parameter_access(): + resolved = self._resolve_chain_parameters(results) + return [ + self._chain_name_map.get(unique_name, unique_name) + if parameter is None + else self.parameter_label(parameter) + for unique_name, parameter in zip(results.param_names, resolved, strict=True) + ] + + +def _raised_inside_bumps(error: BaseException) -> bool: + """ + Check whether an exception came from inside BUMPS. + + Used to make sure only BUMPS' own failures are relabelled, so 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: diff --git a/src/easydynamics/analysis/posterior.py b/src/easydynamics/analysis/posterior.py index 717e99e10..080a7b1ac 100644 --- a/src/easydynamics/analysis/posterior.py +++ b/src/easydynamics/analysis/posterior.py @@ -11,6 +11,7 @@ from __future__ import annotations +import warnings from dataclasses import dataclass from typing import TYPE_CHECKING @@ -19,6 +20,11 @@ if TYPE_CHECKING: from easyscience.variable import Parameter +# How many times wider than the parameter's own value a suggested range may be before it is +# reported as suspicious. A fit that returns an uncertainty this large is describing a flat +# direction rather than a measurement. +ABSURD_WIDTH_FACTOR = 1e4 + # 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 @@ -132,7 +138,9 @@ 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. + Parameters needing manual attention are skipped rather than guessed at. A suggestion that + is absurdly wide is still applied -- it is what the fit implied -- but warned about, since + reading the table first is easy to skip in a script. Returns ------- @@ -140,12 +148,27 @@ def apply(self) -> list[Parameter]: The parameters whose bounds were changed. """ changed = [] + absurd = [] 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) + if _is_absurdly_wide(suggestion): + absurd.append(suggestion.label) + + if absurd: + warnings.warn( + ( + f'Applied bounds far wider than the parameter itself for: ' + f'{", ".join(absurd)}. That width comes from a very large fitted uncertainty, ' + f'which usually means these parameters are degenerate with others, so the ' + f'data cannot determine them separately. Sampling explores that whole range.' + ), + UserWarning, + stacklevel=2, + ) return changed def __len__(self) -> int: @@ -202,6 +225,29 @@ def __repr__(self) -> str: return '\n'.join(lines) +def _is_absurdly_wide(suggestion: BoundsSuggestion) -> bool: + """ + Check whether a suggested range dwarfs the parameter it describes. + + Parameters + ---------- + suggestion : BoundsSuggestion + The suggestion to judge. + + Returns + ------- + bool + True when the range is more than ``ABSURD_WIDTH_FACTOR`` times the parameter's magnitude. + """ + width = suggestion.suggested_max - suggestion.suggested_min + if not np.isfinite(width): + return True + scale = abs(float(suggestion.parameter.value)) + if scale == 0: + return False + return width > ABSURD_WIDTH_FACTOR * scale + + def suggest_bounds_for_parameters( parameters: list[Parameter], labels: list[str] | None = None, diff --git a/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py b/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py index d3308aef5..33fcab061 100644 --- a/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py +++ b/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py @@ -3,6 +3,7 @@ """Unit tests for Bayesian sampling on Analysis1d, with the EasyScience Sampler mocked out.""" +import types from types import SimpleNamespace from unittest.mock import MagicMock from unittest.mock import patch @@ -71,6 +72,20 @@ def fake_results(analysis, n_draws=100, values=None): ) +def _bumps_style_index_error(): + """Build a callable that raises an IndexError from a frame that looks like it is in BUMPS.""" + + def raise_index_error(**_kwargs): + raise IndexError('index 71 is out of bounds for axis 0 with size 40') + + # The relabelling walks the traceback for a frame belonging to the bumps package, so the + # function has to appear to live there. + return types.FunctionType( + raise_index_error.__code__, + {'__name__': 'bumps.dream.state', '__builtins__': __builtins__}, + ) + + @pytest.fixture def analysis(): return make_analysis() @@ -487,18 +502,29 @@ def __exit__(self, *exc_info): class TestErrorPaths: - def test_bumps_outlier_crash_is_reported_as_a_degeneracy(self, analysis): - # WHEN BUMPS' own outlier removal indexes past the end of its buffer, which happens when - # chains scatter because the model is not identifiable + def test_bumps_outlier_crash_is_reported_helpfully(self, analysis): + # WHEN BUMPS' own outlier removal indexes past the end of its buffer + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = _bumps_style_index_error() + + # EXPECT the bare IndexError is replaced by something actionable, naming both causes + with pytest.raises(RuntimeError, match='degenerate') as raised: + analysis.sample_posterior(samples=10) + assert 'short chains' in str(raised.value) + assert isinstance(raised.value.__cause__, IndexError) + + def test_an_index_error_of_our_own_is_not_relabelled(self, analysis): + # WHEN the IndexError comes from anywhere but BUMPS, it is a bug here and must not be + # dressed up as a modelling problem bound_all(analysis) with patch(SAMPLER_PATH) as sampler_class: - sampler_class.return_value.sample.side_effect = IndexError( - 'index 71 is out of bounds for axis 0 with size 40' - ) + sampler_class.return_value.sample.side_effect = IndexError('list index out of range') - # EXPECT the bare IndexError is replaced by something actionable - with pytest.raises(RuntimeError, match='degenerate'): + # EXPECT it propagates untouched + with pytest.raises(IndexError, match='list index out of range'): analysis.sample_posterior(samples=10) def test_parameters_entry_of_the_wrong_type_raises(self, analysis): @@ -560,3 +586,49 @@ def test_trace_and_corner_render_from_a_chain(self, analysis): assert len(analysis.plot_trace().axes) == n_parameters + 1 assert len(analysis.plot_corner().axes) == n_parameters**2 plt.close('all') + + +class TestExtendGuards: + def test_extending_with_a_different_subset_is_refused(self, analysis): + # WHEN a chain is started over all parameters and then extended over one + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + analysis.sample_posterior(samples=10) + + target = analysis.get_free_parameters()[0] + + # EXPECT refused up front, rather than failing obscurely inside BUMPS, which resumes + # from a stored chain whose width is fixed + with pytest.warns(UserWarning), pytest.raises(ValueError, match='Cannot extend'): + analysis.extend_sampling(additional_samples=10, parameters=[target.name]) + + def test_extending_with_the_same_parameters_is_allowed(self, analysis): + # WHEN + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + sampler_class.return_value.extend.side_effect = lambda **_k: fake_results(analysis) + analysis.sample_posterior(samples=10) + + # EXPECT: does not raise + analysis.extend_sampling(additional_samples=10) + + +class TestSidecarLabels: + def test_a_subset_run_records_the_same_labels_a_full_run_would(self, analysis): + # WHEN only one parameter is sampled. Inside the run the others are fixed, so nothing looks + # ambiguous; the recorded labels must still match what a full run would have written, or + # the chain cannot be matched up again on reload. + 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) + with pytest.warns(UserWarning): + analysis.sample_posterior(samples=10, parameters=[target.name]) + + # EXPECT + assert analysis._chain_name_map[target.unique_name] == analysis.parameter_label(target) diff --git a/tests/unit/easydynamics/analysis/test_analysis_bayesian.py b/tests/unit/easydynamics/analysis/test_analysis_bayesian.py index de82d93a9..3ce5722b5 100644 --- a/tests/unit/easydynamics/analysis/test_analysis_bayesian.py +++ b/tests/unit/easydynamics/analysis/test_analysis_bayesian.py @@ -289,3 +289,27 @@ def test_parameter_from_outside_the_analysis_keeps_its_name(self, analysis): # EXPECT it is returned unqualified rather than mislabelled assert analysis.parameter_label(stranger) == 'Gaussian width' + + +class TestIndependentSamplingDiscoverability: + def test_summary_points_at_the_per_q_chains(self, analysis): + # WHEN sampling independently, the chains live on the Analysis1d objects, not here + for analysis1d in analysis.analysis_list: + for parameter in analysis1d.get_free_parameters(): + parameter.min = float(parameter.value) - 5.0 + parameter.max = float(parameter.value) + 5.0 + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results( + analysis.analysis_list[0].get_free_parameters() + ) + analysis.sample_posterior(fit_method='independent', samples=10) + + # EXPECT the error says where the chains actually are, rather than claiming none exist + with pytest.raises(RuntimeError, match='analysis_list'): + analysis.posterior_summary() + + def test_untouched_analysis_still_reports_no_samples(self, analysis): + # EXPECT the plain message when nothing has been sampled anywhere + with pytest.raises(RuntimeError, match='No posterior samples yet'): + analysis.posterior_summary() diff --git a/tests/unit/easydynamics/analysis/test_posterior.py b/tests/unit/easydynamics/analysis/test_posterior.py index bd9ad1d8e..ce7a6acbe 100644 --- a/tests/unit/easydynamics/analysis/test_posterior.py +++ b/tests/unit/easydynamics/analysis/test_posterior.py @@ -322,3 +322,40 @@ def test_len_and_iteration(self): def test_repr_with_no_entries(self): # EXPECT assert 'no parameters' in repr(summarize_draws(np.zeros((3, 0)), [], [])) + + +class TestAbsurdBoundsWarning: + def test_applying_a_wildly_wide_bound_warns(self): + # WHEN a fit returns an enormous uncertainty, which is what a degenerate parameter looks + # like coming out of least squares + parameter = make_parameter(name='Delta area', value=1.0, error=1e9) + suggestions = suggest_bounds_for_parameters([parameter]) + + # EXPECT it is still applied, since it is what the fit implied, but not silently + with pytest.warns(UserWarning, match='far wider than the parameter'): + changed = suggestions.apply() + assert changed == [parameter] + + def test_a_sane_bound_applies_without_warning(self): + # WHEN + parameter = make_parameter(name='sane', value=10.0, error=0.5) + suggestions = suggest_bounds_for_parameters([parameter]) + + # EXPECT + import warnings as warnings_module + + with warnings_module.catch_warnings(): + warnings_module.simplefilter('error') + suggestions.apply() + + def test_a_zero_valued_parameter_is_not_called_absurd(self): + # WHEN there is no magnitude to compare the width against + parameter = make_parameter(name='zero', value=0.0, error=1.0) + suggestions = suggest_bounds_for_parameters([parameter]) + + # EXPECT no warning, since the ratio is meaningless rather than alarming + import warnings as warnings_module + + with warnings_module.catch_warnings(): + warnings_module.simplefilter('error') + suggestions.apply() From 439135a756843553b4eafed00556f378116281e5 Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Thu, 13 Aug 2026 16:05:04 +0200 Subject: [PATCH 08/29] Cover the review fixes, and drop a redundant guard Three lines the review fixes added were not reachable from the unit tests. Two are now covered: extending after a run that died before storing results, where the chain-shape guard has nothing to compare against, and a parameter shared across every Q index, which is left out of the owner map because no single Q identifies it. The third was the non-finite check in the absurd-width test, and it was redundant rather than untested: an infinite width already compares greater than any threshold, and the zero-scale case returns before it. Removed, so the behaviour is unchanged and there is no dead branch. Co-Authored-By: Claude Opus 5 (1M context) --- src/easydynamics/analysis/posterior.py | 6 +-- .../analysis/test_analysis1d_bayesian.py | 17 ++++++++ .../analysis/test_analysis_bayesian.py | 42 +++++++++++++++++++ 3 files changed, 62 insertions(+), 3 deletions(-) diff --git a/src/easydynamics/analysis/posterior.py b/src/easydynamics/analysis/posterior.py index 080a7b1ac..6a1136c6c 100644 --- a/src/easydynamics/analysis/posterior.py +++ b/src/easydynamics/analysis/posterior.py @@ -239,12 +239,12 @@ def _is_absurdly_wide(suggestion: BoundsSuggestion) -> bool: bool True when the range is more than ``ABSURD_WIDTH_FACTOR`` times the parameter's magnitude. """ - width = suggestion.suggested_max - suggestion.suggested_min - if not np.isfinite(width): - return True scale = abs(float(suggestion.parameter.value)) if scale == 0: + # No magnitude to compare against, so the ratio would be meaningless rather than alarming. return False + width = suggestion.suggested_max - suggestion.suggested_min + # An infinite width compares greater than any threshold, so it needs no separate check. return width > ABSURD_WIDTH_FACTOR * scale diff --git a/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py b/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py index 33fcab061..20bed83b8 100644 --- a/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py +++ b/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py @@ -632,3 +632,20 @@ def test_a_subset_run_records_the_same_labels_a_full_run_would(self, analysis): # EXPECT assert analysis._chain_name_map[target.unique_name] == analysis.parameter_label(target) + + def test_extending_after_a_failed_run_is_allowed(self, analysis): + # WHEN a run built the sampler but died before storing results, so there is a sampler to + # extend but no chain shape to compare against + bound_all(analysis) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = RuntimeError('died mid-run') + with pytest.raises(RuntimeError, match='died mid-run'): + analysis.sample_posterior(samples=10) + + assert analysis.bayesian_sampler is not None + assert analysis.posterior_result is None + + # EXPECT the shape guard steps aside rather than comparing against nothing + sampler_class.return_value.extend.side_effect = lambda **_k: fake_results(analysis) + analysis.extend_sampling(additional_samples=10) diff --git a/tests/unit/easydynamics/analysis/test_analysis_bayesian.py b/tests/unit/easydynamics/analysis/test_analysis_bayesian.py index 3ce5722b5..ee9209ebb 100644 --- a/tests/unit/easydynamics/analysis/test_analysis_bayesian.py +++ b/tests/unit/easydynamics/analysis/test_analysis_bayesian.py @@ -313,3 +313,45 @@ def test_untouched_analysis_still_reports_no_samples(self, analysis): # EXPECT the plain message when nothing has been sampled anywhere with pytest.raises(RuntimeError, match='No posterior samples yet'): analysis.posterior_summary() + + +class TestSharedParameterLabels: + def test_a_parameter_shared_across_q_is_not_tied_to_one_index(self): + # WHEN a diffusion model contributes global parameters, the same objects appear at every Q + energy_values = np.linspace(-5.0, 5.0, 15) + rows = [2.0 * np.exp(-0.5 * (energy_values / 1.2) ** 2) for _ in Q_VALUES] + observed = np.vstack(rows) + experiment = edyn.Experiment( + data=sc.DataArray( + data=sc.array( + dims=['Q', 'energy'], + values=observed, + variances=np.full_like(observed, 0.01), + ), + coords={ + 'Q': sc.array(dims=['Q'], values=Q_VALUES, unit='1/Angstrom'), + 'energy': sc.array(dims=['energy'], values=energy_values, unit='meV'), + }, + ) + ) + analysis = edyn.Analysis( + display_name='Shared', + experiment=experiment, + sample_model=sm.SampleModel( + components=sm.ComponentCollection(components=[sm.DeltaFunction(area=0.2)]), + diffusion_models=sm.BrownianTranslationalDiffusion( + name='Brownian', diffusion_coefficient=2.4e-9, scale=0.5 + ), + ), + instrument_model=sm.InstrumentModel(), + ) + + # THEN + owners = analysis._parameter_owner_index() + shared = [p for p in analysis._get_chain_parameters() if p.unique_name not in owners] + + # EXPECT the shared parameters are left out of the owner map, since no single Q owns them, + # and so keep their plain names rather than being labelled with an arbitrary Q + assert shared, 'expected the diffusion model to contribute parameters shared across Q' + for parameter in shared: + assert analysis.parameter_label(parameter) == parameter.name From e6bd67c63235aa1e074f36f5a5f80e06f09b07ec Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Thu, 13 Aug 2026 16:33:32 +0200 Subject: [PATCH 09/29] Gather the per-Q chains on Analysis after independent sampling Sampling with fit_method='independent' left the results only on the Analysis1d objects, so the Analysis that produced them could not report on them. It now gathers them, but only where gathering is sound. posterior_summary() collects every Q into one table, labelled by Q index, and set_parameters_to_posterior_median() applies each chain to its own Q. Both are per-parameter marginal operations, and a marginal is well defined within its own chain, so combining them across separate chains says nothing that was not sampled. plot_corner() deliberately does not aggregate. Independent sampling draws each Q separately, so no draw pairs a parameter at one Q with a parameter at another, and a corner plot built from them would show correlations that are an artefact of how the sampling was run rather than anything measured. It says so and points at the per-Q corner plots, which are real. plot_trace() likewise, the chains being separate runs of different lengths rather than one trace. posterior_results exposes the per-Q chains directly, and a simultaneous chain still takes precedence over stale per-Q ones. Co-Authored-By: Claude Opus 5 (1M context) --- src/easydynamics/analysis/analysis.py | 165 +++++++++++++++++- .../fitting/test_bayesian_sampling_multi_q.py | 42 +++++ .../analysis/test_analysis_bayesian.py | 131 +++++++++++++- 3 files changed, 325 insertions(+), 13 deletions(-) diff --git a/src/easydynamics/analysis/analysis.py b/src/easydynamics/analysis/analysis.py index 25824d959..e6289472c 100644 --- a/src/easydynamics/analysis/analysis.py +++ b/src/easydynamics/analysis/analysis.py @@ -10,12 +10,15 @@ from easyscience.fitting.multi_fitter import MultiFitter from easyscience.fitting.sampler import SamplingResults from easyscience.variable import Parameter +from matplotlib.figure import Figure from plopp.backends.matplotlib.figure import InteractiveFigure from scipp import UnitError from easydynamics.analysis.analysis1d import Analysis1d from easydynamics.analysis.analysis_base import AnalysisBase from easydynamics.analysis.bayesian_sampling import BayesianSamplingMixin +from easydynamics.analysis.posterior import PosteriorSummary +from easydynamics.analysis.posterior import summarize_draws from easydynamics.experiment import Experiment from easydynamics.sample_model import SampleModel from easydynamics.sample_model.instrument_model import InstrumentModel @@ -373,16 +376,37 @@ def sample_posterior( return self.analysis_list[Q_index].sample_posterior( samples=samples, burn=burn, thin=thin, **sampler_options ) - return [ + results = [ analysis.sample_posterior(samples=samples, burn=burn, thin=thin, **sampler_options) for analysis in self.analysis_list ] + # The per-Q chains stay on their own Analysis1d; this only records that they are the + # current results, so this object can gather them up again. + self._posterior_result = None + return results if fit_method == 'simultaneous': return super().sample_posterior( samples=samples, burn=burn, thin=thin, **sampler_options ) raise ValueError("Invalid fit method. Choose 'independent' or 'simultaneous'.") + @property + def posterior_results(self) -> list[SamplingResults | None] | None: + """ + The per-Q chains from independent sampling, or None if there are none. + + A simultaneous run produces one chain covering every Q, which is on + :attr:`posterior_result` instead. + + Returns + ------- + list[SamplingResults | None] | None + One entry per Q index, None where that Q has not been sampled, or None if no Q index + has been sampled at all. + """ + results = [analysis1d.posterior_result for analysis1d in self.analysis_list] + return results if any(result is not None for result in results) else None + def plot_data_and_model( self, Q_index: int | None = None, @@ -912,13 +936,137 @@ def _parameter_owner_index(self) -> dict[str, int]: } return self._owner_index + def posterior_summary(self) -> PosteriorSummary: + """ + Summarize the posterior, gathering the per-Q chains when sampling was independent. + + Every entry is a marginal distribution of one parameter, and a marginal is well defined + within its own chain, so collecting them into one table is sound even though the chains are + separate. Labels are qualified by Q index, so the table reads the same either way. + + A ``RuntimeError`` is raised if no sampling has been run, on this Analysis or on any of its + Q indices. + + Returns + ------- + PosteriorSummary + One entry per sampled parameter, across every Q index that has been sampled. + """ + if self._posterior_result is not None: + return super().posterior_summary() + + per_q = self.posterior_results + if per_q is None: + return super().posterior_summary() + + entries = [] + with self._bulk_parameter_access(): + for analysis1d, result in zip(self.analysis_list, per_q, strict=True): + if result is None: + continue + by_unique_name = {p.unique_name: p for p in analysis1d.get_free_parameters()} + resolved = [by_unique_name.get(name) for name in result.param_names] + labels = [ + unique_name if parameter is None else self.parameter_label(parameter) + for unique_name, parameter in zip(result.param_names, resolved, strict=True) + ] + entries.extend( + summarize_draws(result.draws, labels, resolved).entries, + ) + return PosteriorSummary(entries) + + def set_parameters_to_posterior_median(self) -> list[Parameter]: + """ + Set every sampled parameter to the median of its marginal posterior. + + Applies the per-Q chains to their own Q when sampling was independent. + + A ``RuntimeError`` is raised if no sampling has been run, on this Analysis or on any of its + Q indices. + + Returns + ------- + list[Parameter] + The parameters that were changed. + """ + if self._posterior_result is not None or self.posterior_results is None: + return super().set_parameters_to_posterior_median() + + changed = [] + for analysis1d in self.analysis_list: + if analysis1d.posterior_result is not None: + changed.extend(analysis1d.set_parameters_to_posterior_median()) + return changed + + def plot_corner(self, **kwargs: dict[str, Any]) -> Figure: + """ + Plot marginal and pairwise posterior distributions. + + Only available for a simultaneous chain. Independent sampling draws each Q separately, so + there are no samples pairing a parameter at one Q with a parameter at another, and a corner + plot built from them would show correlations that are an artefact of how the sampling was + run rather than anything measured. + + Parameters + ---------- + **kwargs : dict[str, Any] + Forwarded to :func:`easydynamics.utils.posterior_plotting.plot_corner`. + + Returns + ------- + Figure + The matplotlib Figure. + + Raises + ------ + RuntimeError + If only independent per-Q chains exist. + """ + if self._posterior_result is None and self.posterior_results is not None: + raise RuntimeError( + 'A corner plot needs one chain covering the parameters it compares, and ' + "fit_method='independent' samples each Q on its own, so no draw pairs one Q with " + 'another. Use analysis.analysis_list[Q_index].plot_corner() for the correlations ' + "within a Q, or sample with fit_method='simultaneous' to compare across Q." + ) + return super().plot_corner(**kwargs) + + def plot_trace(self, **kwargs: dict[str, Any]) -> Figure: + """ + Plot the chain trace of each sampled parameter. + + Only available for a simultaneous chain, since the per-Q chains are separate runs of + different lengths rather than one trace. + + Parameters + ---------- + **kwargs : dict[str, Any] + Forwarded to :func:`easydynamics.utils.posterior_plotting.plot_trace`. + + Returns + ------- + Figure + The matplotlib Figure. + + Raises + ------ + RuntimeError + If only independent per-Q chains exist. + """ + if self._posterior_result is None and self.posterior_results is not None: + raise RuntimeError( + 'Each Q index has its own chain, so there is no single trace to draw. Use ' + 'analysis.analysis_list[Q_index].plot_trace() for one of them, or sample with ' + "fit_method='simultaneous'." + ) + return super().plot_trace(**kwargs) + def _require_posterior_result(self) -> SamplingResults: """ Get the stored sampling results, pointing at the per-Q chains when those are what exist. Sampling independently stores a chain on each Analysis1d rather than here, so asking this - object for a summary afterwards would otherwise report that no sampling has happened, right - after it has. + object for one would otherwise report that no sampling has happened, right after it has. Returns ------- @@ -930,14 +1078,13 @@ def _require_posterior_result(self) -> SamplingResults: RuntimeError If no simultaneous sampling has been run on this Analysis. """ - if self._posterior_result is None and any( - analysis1d.posterior_result is not None for analysis1d in self.analysis_list - ): + if self._posterior_result is None and self.posterior_results is not None: raise RuntimeError( 'This Analysis holds no chain of its own, but its Q indices do: sampling with ' - "fit_method='independent' gives each Q its own chain. Use " - 'analysis.analysis_list[Q_index] to summarize or plot one of them, or sample with ' - "fit_method='simultaneous' to get a single chain over all Q." + "fit_method='independent' gives each Q its own chain. posterior_summary() and " + 'set_parameters_to_posterior_median() gather those up; for anything needing a ' + 'single chain, use analysis.analysis_list[Q_index], or sample with ' + "fit_method='simultaneous'." ) return super()._require_posterior_result() diff --git a/tests/integration/fitting/test_bayesian_sampling_multi_q.py b/tests/integration/fitting/test_bayesian_sampling_multi_q.py index 56a883f97..ad87ea99f 100644 --- a/tests/integration/fitting/test_bayesian_sampling_multi_q.py +++ b/tests/integration/fitting/test_bayesian_sampling_multi_q.py @@ -217,3 +217,45 @@ def test_recovers_a_straight_line_through_the_widths(self): assert results.draws.shape[1] == len(analysis._get_chain_parameters()) names = [entry.name for entry in analysis.posterior_summary()] assert len(set(names)) == len(names) + + +class TestAggregatedIndependentChains: + def test_summary_gathers_the_real_per_q_chains(self): + # WHEN each Q is sampled on its own + analysis = build_analysis() + analysis.fit(fit_method='independent') + for analysis1d in analysis.analysis_list: + analysis1d.suggest_bounds().apply() + + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + analysis.sample_posterior(fit_method='independent', **SAMPLE_KWARGS) + + # THEN + summary = analysis.posterior_summary() + + # EXPECT one table covering every Q, and the widths still recovered + assert len(summary) == sum(len(a.get_free_parameters()) for a in analysis.analysis_list) + for q_index in range(len(Q_VALUES)): + entry = summary[f'Gaussian width (Q_index={q_index})'] + spread = max(entry.minus, entry.plus) + assert abs(entry.median - true_width(Q_VALUES[q_index])) < 4 * spread + + def test_median_applies_each_chain_to_its_own_q(self): + # WHEN + analysis = build_analysis() + analysis.fit(fit_method='independent') + for analysis1d in analysis.analysis_list: + analysis1d.suggest_bounds().apply() + + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + analysis.sample_posterior(fit_method='independent', **SAMPLE_KWARGS) + + changed = analysis.set_parameters_to_posterior_median() + + # EXPECT every Q's parameters land on that Q's own median + assert len(changed) == sum(len(a.get_free_parameters()) for a in analysis.analysis_list) + summary = analysis.posterior_summary() + for entry in summary: + assert entry.value == pytest.approx(entry.median, rel=1e-6) diff --git a/tests/unit/easydynamics/analysis/test_analysis_bayesian.py b/tests/unit/easydynamics/analysis/test_analysis_bayesian.py index ee9209ebb..c48fe9803 100644 --- a/tests/unit/easydynamics/analysis/test_analysis_bayesian.py +++ b/tests/unit/easydynamics/analysis/test_analysis_bayesian.py @@ -7,10 +7,13 @@ from unittest.mock import MagicMock from unittest.mock import patch +import matplotlib as mpl import numpy as np import pytest import scipp as sc +mpl.use('Agg') + import easydynamics as edyn import easydynamics.sample_model as sm @@ -292,8 +295,9 @@ def test_parameter_from_outside_the_analysis_keeps_its_name(self, analysis): class TestIndependentSamplingDiscoverability: - def test_summary_points_at_the_per_q_chains(self, analysis): + def test_operations_needing_one_chain_point_at_the_per_q_chains(self, analysis): # WHEN sampling independently, the chains live on the Analysis1d objects, not here + remaining = iter(analysis.analysis_list) for analysis1d in analysis.analysis_list: for parameter in analysis1d.get_free_parameters(): parameter.min = float(parameter.value) - 5.0 @@ -301,13 +305,14 @@ def test_summary_points_at_the_per_q_chains(self, analysis): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.side_effect = lambda **_k: fake_results( - analysis.analysis_list[0].get_free_parameters() + next(remaining).get_free_parameters() ) analysis.sample_posterior(fit_method='independent', samples=10) - # EXPECT the error says where the chains actually are, rather than claiming none exist + # EXPECT anything that genuinely needs a single chain says where the chains actually are, + # rather than claiming none exist with pytest.raises(RuntimeError, match='analysis_list'): - analysis.posterior_summary() + analysis.plot_posterior_predictive() def test_untouched_analysis_still_reports_no_samples(self, analysis): # EXPECT the plain message when nothing has been sampled anywhere @@ -355,3 +360,121 @@ def test_a_parameter_shared_across_q_is_not_tied_to_one_index(self): assert shared, 'expected the diffusion model to contribute parameters shared across Q' for parameter in shared: assert analysis.parameter_label(parameter) == parameter.name + + +class TestAggregatingPerQChains: + def _sample_independently(self, analysis): + for analysis1d in analysis.analysis_list: + for parameter in analysis1d.get_free_parameters(): + parameter.min = float(parameter.value) - 5.0 + parameter.max = float(parameter.value) + 5.0 + + # The Q indices sample in order, and each must get a chain over its own parameters. + remaining = iter(analysis.analysis_list) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results( + next(remaining).get_free_parameters() + ) + analysis.sample_posterior(fit_method='independent', samples=10) + + def test_posterior_results_holds_one_chain_per_q(self, analysis): + # WHEN + self._sample_independently(analysis) + + # EXPECT + assert len(analysis.posterior_results) == len(Q_VALUES) + assert all(result is not None for result in analysis.posterior_results) + + def test_posterior_results_is_none_before_sampling(self, analysis): + # EXPECT + assert analysis.posterior_results is None + + def test_summary_gathers_every_q(self, analysis): + # WHEN + self._sample_independently(analysis) + + # THEN + summary = analysis.posterior_summary() + + # EXPECT one entry per free parameter per Q, each labelled by its Q index + expected = sum(len(a.get_free_parameters()) for a in analysis.analysis_list) + names = [entry.name for entry in summary] + assert len(summary) == expected + assert len(set(names)) == len(names) + assert all('Q_index=' in name for name in names) + + def test_median_applies_each_chain_to_its_own_q(self, analysis): + # WHEN + self._sample_independently(analysis) + + # THEN + changed = analysis.set_parameters_to_posterior_median() + + # EXPECT every Q's parameters are set, from that Q's own chain + expected = sum(len(a.get_free_parameters()) for a in analysis.analysis_list) + assert len(changed) == expected + + def test_corner_refuses_to_invent_cross_q_correlations(self, analysis): + # WHEN each Q was sampled separately, no draw pairs one Q with another + self._sample_independently(analysis) + + # EXPECT it says so, rather than plotting correlations that are an artefact of the run + with pytest.raises(RuntimeError, match='no draw pairs one Q'): + analysis.plot_corner() + + def test_trace_points_at_the_individual_chains(self, analysis): + # WHEN + self._sample_independently(analysis) + + # EXPECT + with pytest.raises(RuntimeError, match='no single trace'): + analysis.plot_trace() + + def test_a_simultaneous_chain_still_takes_precedence(self, analysis): + # WHEN a simultaneous run follows an independent one + self._sample_independently(analysis) + bound_all(analysis) + parameters = analysis._get_chain_parameters() + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_results(parameters) + analysis.sample_posterior(fit_method='simultaneous', samples=10) + + # EXPECT the single chain is summarized, not the stale per-Q ones + assert len(analysis.posterior_summary()) == len(parameters) + analysis.plot_corner() + + def test_only_the_sampled_q_indices_are_gathered(self, analysis): + # WHEN just one Q index is sampled + target = analysis.analysis_list[1] + for parameter in target.get_free_parameters(): + parameter.min = float(parameter.value) - 5.0 + parameter.max = float(parameter.value) + 5.0 + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results( + target.get_free_parameters() + ) + analysis.sample_posterior(fit_method='independent', Q_index=1, samples=10) + + # EXPECT the unsampled Q indices are passed over rather than breaking the aggregation + summary = analysis.posterior_summary() + assert len(summary) == len(target.get_free_parameters()) + assert all('Q_index=1' in entry.name for entry in summary) + assert len(analysis.set_parameters_to_posterior_median()) == len( + target.get_free_parameters() + ) + + def test_a_simultaneous_chain_serves_the_median_and_the_trace(self, analysis): + # WHEN + bound_all(analysis) + parameters = analysis._get_chain_parameters() + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_results(parameters) + analysis.sample_posterior(fit_method='simultaneous', samples=10) + + # EXPECT both come from the single chain, with no per-Q gathering involved + assert len(analysis.set_parameters_to_posterior_median()) == len(parameters) + assert len(analysis.plot_trace().axes) == len(parameters) + 1 From 6170a4312c37d3b8cfe76063e80316d4a877b59c Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Thu, 13 Aug 2026 17:11:35 +0200 Subject: [PATCH 10/29] Step through the per-Q corner plots with a slider Independent chains share no draws, so there is no joint distribution across Q to plot, and combining them would show correlations that came from how the sampling was run rather than from the data. Refusing outright was correct but unhelpful: the correlations within each Q are real and worth looking at. Analysis.plot_corner() now shows one Q at a time. Pass Q_index for a particular one, or leave it out in a notebook for a slider across the Q values that were sampled. A simultaneous chain is unaffected; it already covers every Q in one figure. Outside a notebook the error names the sampled Q indices rather than only saying no. The slider is built with append_display_data rather than the Output widget's context manager. The context manager is the obvious choice and captures nothing under some kernels, which would have shipped a slider with a permanently blank panel beside it. Verified by executing a notebook against a real kernel, and the test asserts the panel actually holds a figure, since an empty panel is the regression that matters. Co-Authored-By: Claude Opus 5 (1M context) --- docs/docs/tutorials/bayesian.ipynb | 6 +- src/easydynamics/analysis/analysis.py | 74 ++++++++++++++---- src/easydynamics/utils/posterior_plotting.py | 78 +++++++++++++++++++ .../analysis/test_analysis_bayesian.py | 62 ++++++++++++++- 4 files changed, 201 insertions(+), 19 deletions(-) diff --git a/docs/docs/tutorials/bayesian.ipynb b/docs/docs/tutorials/bayesian.ipynb index 44a415147..1629d9058 100644 --- a/docs/docs/tutorials/bayesian.ipynb +++ b/docs/docs/tutorials/bayesian.ipynb @@ -287,7 +287,11 @@ "\n", "**Sampling only some parameters.** `sample_posterior(parameters=[...])` restricts the chain to a subset, which is faster because the number of chains scales with the number of parameters. Be careful with the result: the other parameters are held *fixed*, which is not the same as averaging over them. The intervals you get are conditional on those fixed values, and will be too narrow whenever the parameters are correlated.\n", "\n", - "**Degenerate parameters.** As seen above, they are a modelling problem rather than a sampling one. `suggest_bounds()` returning absurd values, or a warning that the posterior has piled up against its bounds, are both signs to go back and look at the model." + "**Degenerate parameters.** As seen above, they are a modelling problem rather than a sampling one. `suggest_bounds()` returning absurd values, or a warning that the posterior has piled up against its bounds, are both signs to go back and look at the model.\n", + "\n", + "**Several Q values at once.** An `Analysis` can sample either way. `fit_method='independent'` gives each Q its own chain, which is cheaper; `fit_method='simultaneous'` runs one chain over every Q, which is what you need when parameters are shared across Q. `posterior_summary()` gathers the per-Q chains into one table either way.\n", + "\n", + "Corner plots are the exception. Independent chains share no draws, so nothing pairs a parameter at one Q with a parameter at another, and combining them would show correlations that came from how the sampling was run rather than from the data. `analysis.plot_corner()` therefore shows one Q at a time: pass `Q_index`, or leave it out in a notebook to get a slider across the sampled Q values." ] } ], diff --git a/src/easydynamics/analysis/analysis.py b/src/easydynamics/analysis/analysis.py index e6289472c..7561904d3 100644 --- a/src/easydynamics/analysis/analysis.py +++ b/src/easydynamics/analysis/analysis.py @@ -10,6 +10,7 @@ from easyscience.fitting.multi_fitter import MultiFitter from easyscience.fitting.sampler import SamplingResults from easyscience.variable import Parameter +from ipywidgets import VBox from matplotlib.figure import Figure from plopp.backends.matplotlib.figure import InteractiveFigure from scipp import UnitError @@ -25,6 +26,7 @@ from easydynamics.settings.convolution_settings import ConvolutionSettings from easydynamics.settings.detailed_balance_settings import DetailedBalanceSettings from easydynamics.utils.plotting import slicerplot_with_residuals +from easydynamics.utils.posterior_plotting import corner_with_slider from easydynamics.utils.utils import _in_notebook from easydynamics.utils.utils import verify_Q_index @@ -998,38 +1000,82 @@ def set_parameters_to_posterior_median(self) -> list[Parameter]: changed.extend(analysis1d.set_parameters_to_posterior_median()) return changed - def plot_corner(self, **kwargs: dict[str, Any]) -> Figure: + def plot_corner(self, Q_index: int | None = None, **kwargs: dict[str, Any]) -> Figure | VBox: """ Plot marginal and pairwise posterior distributions. - Only available for a simultaneous chain. Independent sampling draws each Q separately, so - there are no samples pairing a parameter at one Q with a parameter at another, and a corner - plot built from them would show correlations that are an artefact of how the sampling was - run rather than anything measured. + After independent sampling each Q has its own chain, and no draw pairs a parameter at one Q + with a parameter at another, so there is no joint distribution across Q to plot. Rather + than combine them into a figure showing correlations that are an artefact of how the + sampling was run, this steps through the chains one at a time: pick one with ``Q_index``, + or leave it out in a notebook to get a slider. Parameters ---------- + Q_index : int | None, default=None + Which Q index to plot, when the chains are per-Q. If None, a slider is returned. Not + used for a simultaneous chain, which already covers every Q. **kwargs : dict[str, Any] Forwarded to :func:`easydynamics.utils.posterior_plotting.plot_corner`. Returns ------- - Figure - The matplotlib Figure. + Figure | VBox + The matplotlib Figure, or an ipywidgets box with a Q slider. Raises ------ RuntimeError - If only independent per-Q chains exist. + If a slider is asked for outside a notebook. """ - if self._posterior_result is None and self.posterior_results is not None: + per_q = self.posterior_results + if self._posterior_result is not None or per_q is None: + return super().plot_corner(**kwargs) + + verify_Q_index(Q_index=Q_index, Q=self.Q, allow_none=True) + if Q_index is not None: + return self.analysis_list[Q_index].plot_corner(**kwargs) + + if not _in_notebook(): + sampled = [index for index, result in enumerate(per_q) if result is not None] raise RuntimeError( - 'A corner plot needs one chain covering the parameters it compares, and ' - "fit_method='independent' samples each Q on its own, so no draw pairs one Q with " - 'another. Use analysis.analysis_list[Q_index].plot_corner() for the correlations ' - "within a Q, or sample with fit_method='simultaneous' to compare across Q." + f'Each Q index has its own chain, and the slider needs a Jupyter notebook. ' + f'Pass Q_index to plot one of them; sampled Q indices are {sampled}.' ) - return super().plot_corner(**kwargs) + return corner_with_slider(self._per_q_corner_chains(per_q), title=self.display_name) + + def _per_q_corner_chains(self, per_q: list[SamplingResults | None]) -> dict[int, dict]: + """ + Describe each per-Q chain in the form the slider wants. + + Parameters + ---------- + per_q : list[SamplingResults | None] + The chain for each Q index, None where that Q has not been sampled. + + Returns + ------- + dict[int, dict] + Mapping of Q index to its draws, labels, and units. Q indices without a chain are left + out, so the slider only offers what can actually be drawn. + """ + chains = {} + for analysis1d, result in zip(self.analysis_list, per_q, strict=True): + if result is None: + continue + by_unique_name = {p.unique_name: p for p in analysis1d.get_free_parameters()} + resolved = [by_unique_name.get(name) for name in result.param_names] + chains[analysis1d.Q_index] = { + 'draws': result.draws, + # Labelled with the plain parameter name: the Q index is already on the slider, so + # repeating it in every axis label would only cost width. + 'names': [ + unique_name if parameter is None else parameter.name + for unique_name, parameter in zip(result.param_names, resolved, strict=True) + ], + 'units': ['' if p is None else str(p.unit) for p in resolved], + } + return chains def plot_trace(self, **kwargs: dict[str, Any]) -> Figure: """ diff --git a/src/easydynamics/utils/posterior_plotting.py b/src/easydynamics/utils/posterior_plotting.py index 8e75cf425..eb9cda994 100644 --- a/src/easydynamics/utils/posterior_plotting.py +++ b/src/easydynamics/utils/posterior_plotting.py @@ -11,12 +11,14 @@ from __future__ import annotations from typing import TYPE_CHECKING +from typing import Any import matplotlib.pyplot as plt import numpy as np from matplotlib.ticker import MaxNLocator if TYPE_CHECKING: + from ipywidgets import VBox from matplotlib.figure import Figure @@ -370,3 +372,79 @@ def _verify_draws(draws: np.ndarray, names: list[str]) -> None: f'names must have one entry per column of draws. ' f'Got {len(names)} names for {draws.shape[1]} columns.' ) + + +def corner_with_slider( + chains: dict[int, dict], + title: str | None = None, + **kwargs: dict[str, Any], +) -> VBox: + """ + Show one corner plot at a time, with a slider choosing which chain to look at. + + Chains sampled separately share no draws, so there is no joint distribution across them to + plot. Stepping through them one at a time shows the correlations that were actually sampled, + which is what a single combined figure could not do honestly. + + Parameters + ---------- + chains : dict[int, dict] + Mapping of index to a ``{'draws': ..., 'names': ..., 'units': ...}`` description of one + chain. ``units`` is optional. + title : str | None, default=None + Title prefix, extended with the selected index. + **kwargs : dict[str, Any] + Forwarded to :func:`plot_corner`. + + Returns + ------- + VBox + An ipywidgets box holding the slider and the figure. + + Raises + ------ + ValueError + If no chains are given. + """ + import ipywidgets as widgets + + if not chains: + raise ValueError('No chains to plot.') + + indices = sorted(chains) + output = widgets.Output() + + def draw(index: int) -> None: + """ + Render the corner plot for one chain. + + Parameters + ---------- + index : int + The chain to draw. + """ + chain = chains[index] + figure = plot_corner( + draws=chain['draws'], + names=chain['names'], + units=chain.get('units'), + title=title if title is None else f'{title} (Q index {index})', + **kwargs, + ) + # append_display_data rather than the `with output:` context manager, which captures + # nothing under some kernels and would leave the slider with a blank panel beside it. + output.outputs = () + output.append_display_data(figure) + # Rendered into the widget already, so the figure is closed rather than left for a backend + # to draw a second time. + plt.close(figure) + + slider = widgets.SelectionSlider( + options=indices, + value=indices[0], + description='Q index', + continuous_update=False, + ) + slider.observe(lambda change: draw(change['new']), names='value') + draw(indices[0]) + return widgets.VBox([slider, output]) diff --git a/tests/unit/easydynamics/analysis/test_analysis_bayesian.py b/tests/unit/easydynamics/analysis/test_analysis_bayesian.py index c48fe9803..bff7efa91 100644 --- a/tests/unit/easydynamics/analysis/test_analysis_bayesian.py +++ b/tests/unit/easydynamics/analysis/test_analysis_bayesian.py @@ -415,14 +415,68 @@ def test_median_applies_each_chain_to_its_own_q(self, analysis): expected = sum(len(a.get_free_parameters()) for a in analysis.analysis_list) assert len(changed) == expected - def test_corner_refuses_to_invent_cross_q_correlations(self, analysis): - # WHEN each Q was sampled separately, no draw pairs one Q with another + def test_corner_plots_one_q_at_a_time(self, analysis): + # WHEN each Q was sampled separately, no draw pairs one Q with another, so a corner plot + # can only show one chain at a time self._sample_independently(analysis) - # EXPECT it says so, rather than plotting correlations that are an artefact of the run - with pytest.raises(RuntimeError, match='no draw pairs one Q'): + # THEN + figure = analysis.plot_corner(Q_index=1) + + # EXPECT that Q's own chain, not a combination across Q + n_parameters = len(analysis.analysis_list[1].get_free_parameters()) + assert len(figure.axes) == n_parameters**2 + + def test_corner_offers_a_slider_in_a_notebook(self, analysis): + # WHEN + self._sample_independently(analysis) + + with patch('easydynamics.analysis.analysis._in_notebook', return_value=True): + widget = analysis.plot_corner() + + # EXPECT a slider over the sampled Q indices, and a panel that actually holds a figure. + # The obvious way to build this captures nothing and leaves the panel blank beside the + # slider, so an empty panel is the regression worth guarding. Which mime type arrives + # depends on the environment: a live kernel renders a PNG, plain pytest only the repr. + slider, panel = widget.children + assert list(slider.options) == list(range(len(Q_VALUES))) + assert panel.outputs, 'the initial chain was not drawn' + assert 'Figure' in str(panel.outputs[0]['data']) + + slider.value = 2 + assert panel.outputs, 'changing Q did not redraw' + assert 'Figure' in str(panel.outputs[0]['data']) + + def test_corner_without_a_notebook_or_q_index_says_what_to_do(self, analysis): + # WHEN + self._sample_independently(analysis) + + # EXPECT it names the sampled Q indices rather than just refusing + with ( + patch('easydynamics.analysis.analysis._in_notebook', return_value=False), + pytest.raises(RuntimeError, match=r'sampled Q indices are \[0, 1, 2\]'), + ): analysis.plot_corner() + def test_the_slider_only_offers_q_indices_that_were_sampled(self, analysis): + # WHEN only one Q index is sampled + target = analysis.analysis_list[2] + for parameter in target.get_free_parameters(): + parameter.min = float(parameter.value) - 5.0 + parameter.max = float(parameter.value) + 5.0 + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results( + target.get_free_parameters() + ) + analysis.sample_posterior(fit_method='independent', Q_index=2, samples=10) + + with patch('easydynamics.analysis.analysis._in_notebook', return_value=True): + widget = analysis.plot_corner() + + # EXPECT the slider cannot land on a Q with nothing to draw + assert list(widget.children[0].options) == [2] + def test_trace_points_at_the_individual_chains(self, analysis): # WHEN self._sample_independently(analysis) From 1292f44ddefb913303a20b420a6f8009553d4104 Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Thu, 13 Aug 2026 17:32:25 +0200 Subject: [PATCH 11/29] Show the per-Q corner slider in the Bayesian tutorial The slider was described in the tutorial's caveats but never demonstrated: every notebook call to plot_corner() went through the single-chain path, because the Bayesian tutorial used Analysis1d and tutorial 1 used ParameterAnalysis, neither of which has a Q dimension. So the only things exercising it were the unit tests. The tutorial now builds the full multi-Q Analysis, samples a few Q values, gathers them with posterior_summary(), and shows the slider. It samples Q indices 4, 8 and 12 rather than all sixteen. Sampling every Q measured at 70 s against 16 s for three, and the subset also shows two things worth showing: that sampling is slow enough to be worth trying a few Q values first, and that the slider offers only the Q values that were actually sampled. Verified against a real kernel that the cell emits a widget view, rather than only that the notebook ran without raising. Co-Authored-By: Claude Opus 5 (1M context) --- docs/docs/tutorials/bayesian.ipynb | 82 ++++++++++++++++++++++++++++++ 1 file changed, 82 insertions(+) diff --git a/docs/docs/tutorials/bayesian.ipynb b/docs/docs/tutorials/bayesian.ipynb index 1629d9058..46858d732 100644 --- a/docs/docs/tutorials/bayesian.ipynb +++ b/docs/docs/tutorials/bayesian.ipynb @@ -276,6 +276,88 @@ "Chains are expensive, so they can be saved and reloaded with `analysis.save_chain(path)` and `analysis.load_chain(path)`. A reloaded chain can be summarized, plotted, or extended further, exactly like a fresh one." ] }, + { + "cell_type": "markdown", + "id": "ef0998e4", + "metadata": {}, + "source": [ + "## Several Q values at once\n", + "\n", + "Everything so far used `Analysis1d`, a single Q slice. A full `Analysis` can sample too, either way round:\n", + "\n", + "- `fit_method='independent'` gives each Q its own chain. Cheaper, and the Q values cannot influence one another.\n", + "- `fit_method='simultaneous'` runs a single chain over every Q at once, which is what you need when parameters are shared across Q. It costs considerably more, because DREAM runs a number of chains proportional to the parameter count and a simultaneous run has every Q's parameters in play together.\n", + "\n", + "Sampling is much slower than fitting, so it is worth trying a few Q values before committing to all of them. Passing `Q_index` samples just that one." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b9d3b565", + "metadata": {}, + "outputs": [], + "source": [ + "# Fresh models, so this analysis is independent of the single-Q one above rather than\n", + "# sharing its already-sampled components.\n", + "all_q_components = sm.ComponentCollection()\n", + "all_q_components.append_component(sm.Gaussian(width=0.1, area=1, name='Res. Gauss'))\n", + "\n", + "full_analysis = edyn.Analysis(\n", + " display_name='Vanadium, all Q',\n", + " experiment=vanadium_experiment,\n", + " sample_model=sm.SampleModel(components=all_q_components),\n", + " instrument_model=sm.InstrumentModel(\n", + " background_model=sm.BackgroundModel(components=sm.Polynomial(coefficients=[0.001])),\n", + " ),\n", + ")\n", + "full_analysis.fit(fit_method='independent')\n", + "\n", + "for Q_index in (4, 8, 12):\n", + " full_analysis.analysis_list[Q_index].suggest_bounds().apply()\n", + " full_analysis.sample_posterior(\n", + " fit_method='independent', Q_index=Q_index, samples=3000, burn=200, thin=2\n", + " )" + ] + }, + { + "cell_type": "markdown", + "id": "740fa625", + "metadata": {}, + "source": [ + "`posterior_summary()` gathers the per-Q chains into one table, labelled by Q index. Each row is a marginal distribution, and a marginal is well defined within its own chain, so collecting them says nothing that was not sampled." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "511ef922", + "metadata": {}, + "outputs": [], + "source": [ + "full_analysis.posterior_summary()" + ] + }, + { + "cell_type": "markdown", + "id": "12347890", + "metadata": {}, + "source": [ + "Corner plots are the one thing that cannot be gathered up. The chains were run separately, so no draw pairs a parameter at one Q with a parameter at another, and a combined figure would show correlations that came from how the sampling was run rather than from the data.\n", + "\n", + "So `plot_corner()` steps through them instead. The slider offers only the Q values that were actually sampled — 4, 8 and 12 here — and `plot_corner(Q_index=8)` goes straight to one of them." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "38d0b23c", + "metadata": {}, + "outputs": [], + "source": [ + "full_analysis.plot_corner()" + ] + }, { "cell_type": "markdown", "id": "a3448aee", From 442a5bde2336926286e1a4e73ec16207211472ed Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Thu, 13 Aug 2026 17:42:55 +0200 Subject: [PATCH 12/29] Put the corner slider under the figure Matches where plopp puts its slicer controls, which is also where the existing slicerplot_with_residuals puts them via the figure's bottom bar. Co-Authored-By: Claude Opus 5 (1M context) --- src/easydynamics/utils/posterior_plotting.py | 3 ++- tests/unit/easydynamics/analysis/test_analysis_bayesian.py | 5 +++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/src/easydynamics/utils/posterior_plotting.py b/src/easydynamics/utils/posterior_plotting.py index eb9cda994..eee1ce85d 100644 --- a/src/easydynamics/utils/posterior_plotting.py +++ b/src/easydynamics/utils/posterior_plotting.py @@ -447,4 +447,5 @@ def draw(index: int) -> None: ) slider.observe(lambda change: draw(change['new']), names='value') draw(indices[0]) - return widgets.VBox([slider, output]) + # Slider under the figure, matching where plopp puts its slicer controls. + return widgets.VBox([output, slider]) diff --git a/tests/unit/easydynamics/analysis/test_analysis_bayesian.py b/tests/unit/easydynamics/analysis/test_analysis_bayesian.py index bff7efa91..a2c451e60 100644 --- a/tests/unit/easydynamics/analysis/test_analysis_bayesian.py +++ b/tests/unit/easydynamics/analysis/test_analysis_bayesian.py @@ -438,7 +438,8 @@ def test_corner_offers_a_slider_in_a_notebook(self, analysis): # The obvious way to build this captures nothing and leaves the panel blank beside the # slider, so an empty panel is the regression worth guarding. Which mime type arrives # depends on the environment: a live kernel renders a PNG, plain pytest only the repr. - slider, panel = widget.children + # The figure comes first and the slider sits under it, where plopp puts its controls. + panel, slider = widget.children assert list(slider.options) == list(range(len(Q_VALUES))) assert panel.outputs, 'the initial chain was not drawn' assert 'Figure' in str(panel.outputs[0]['data']) @@ -475,7 +476,7 @@ def test_the_slider_only_offers_q_indices_that_were_sampled(self, analysis): widget = analysis.plot_corner() # EXPECT the slider cannot land on a Q with nothing to draw - assert list(widget.children[0].options) == [2] + assert list(widget.children[1].options) == [2] def test_trace_points_at_the_individual_chains(self, analysis): # WHEN From bb1062308945fb165518e4cb13248eeb13011d5f Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Fri, 14 Aug 2026 13:07:35 +0200 Subject: [PATCH 13/29] Compose the posterior sampler instead of mixing it in Review feedback: bayesian_sampling.py had a lot in it that belonged elsewhere, and it was unclear why it was a mixin at all. It was a mixin because ParameterAnalysis is not an AnalysisBase and fits its binding models rather than itself, so a shared base class does not work. That was a reason, not a good one: it injected some forty methods into every Analysis class. The sampler is now composed. An Analysis exposes one `bayesian` property, and hands the sampler the few things that differ between the Analysis classes -- the data, the free parameters, their labels, and a hook to refresh cached computation -- so PosteriorSampler needs no knowledge of how any Analysis is built, and no Analysis inherits sampling machinery it does not use. Labelling moves to posterior_labels.py. Building it once for a fixed set of parameters also removes the quadratic cost the old code needed a scoped cache to avoid: the counts and lookups are computed in the constructor rather than per column. Plotting stays in posterior_plotting.py, where it already lived. The sampler keeps three short delegates so a chain can still be plotted from the object holding it, but none of the drawing happens there. The public API becomes analysis.bayesian.sample() and friends, and the explicit suggest_bounds().apply() step stays: in DREAM the bounds are the prior, and an unbounded parameter gives a confident-looking interval set by nothing. Co-Authored-By: Claude Opus 5 (1M context) --- docs/docs/tutorials/bayesian.ipynb | 30 +- src/easydynamics/analysis/__init__.py | 6 +- src/easydynamics/analysis/analysis1d.py | 89 +- .../analysis/bayesian_sampling.py | 993 ------------------ src/easydynamics/analysis/posterior.py | 37 +- src/easydynamics/analysis/posterior_labels.py | 182 ++++ .../analysis/posterior_sampling.py | 933 ++++++++++++++++ src/easydynamics/utils/posterior_plotting.py | 118 ++- .../fitting/test_bayesian_sampling.py | 44 +- .../analysis/test_analysis1d_bayesian.py | 86 +- .../easydynamics/analysis/test_posterior.py | 15 +- .../analysis/test_posterior_labels.py | 109 ++ 12 files changed, 1537 insertions(+), 1105 deletions(-) delete mode 100644 src/easydynamics/analysis/bayesian_sampling.py create mode 100644 src/easydynamics/analysis/posterior_labels.py create mode 100644 src/easydynamics/analysis/posterior_sampling.py create mode 100644 tests/unit/easydynamics/analysis/test_posterior_labels.py diff --git a/docs/docs/tutorials/bayesian.ipynb b/docs/docs/tutorials/bayesian.ipynb index 44a415147..51a68acd2 100644 --- a/docs/docs/tutorials/bayesian.ipynb +++ b/docs/docs/tutorials/bayesian.ipynb @@ -11,7 +11,7 @@ "\n", "A **Bayesian** analysis answers a different question: instead of one best point, it maps out the whole *posterior distribution* over the parameters. From that you can read off credible intervals that stay honest when parameters are correlated or their distributions are skewed, and you can see the correlations directly.\n", "\n", - "EasyDynamics does this with the DREAM sampler from [BUMPS](https://bumps.readthedocs.io/), through `sample_posterior()`." + "EasyDynamics does this with the DREAM sampler from [BUMPS](https://bumps.readthedocs.io/), through `bayesian.sample()`." ] }, { @@ -102,9 +102,9 @@ "source": [ "## Bounds are the prior\n", "\n", - "In DREAM, each parameter's `min` and `max` define a uniform prior, so **every free parameter must have finite bounds** before sampling. Most parameters start with at least one infinite bound, so `sample_posterior()` would refuse to run.\n", + "In DREAM, each parameter's `min` and `max` define a uniform prior, so **every free parameter must have finite bounds** before sampling. Most parameters start with at least one infinite bound, so `bayesian.sample()` would refuse to run.\n", "\n", - "`suggest_bounds()` proposes bounds from the fitted values and uncertainties. It is advisory: it changes nothing until you call `.apply()`, and it only ever fills in an *infinite* bound, so physical limits you have already set (an area that cannot go below zero, say) are left alone." + "`bayesian.suggest_bounds()` proposes bounds from the fitted values and uncertainties. It is advisory: it changes nothing until you call `.apply()`, and it only ever fills in an *infinite* bound, so physical limits you have already set (an area that cannot go below zero, say) are left alone." ] }, { @@ -114,7 +114,7 @@ "metadata": {}, "outputs": [], "source": [ - "suggestions = analysis.suggest_bounds()\n", + "suggestions = analysis.bayesian.suggest_bounds()\n", "print(suggestions)" ] }, @@ -146,7 +146,7 @@ "source": [ "## Sample the posterior\n", "\n", - "`sample_posterior()` runs the chains. The three numbers that matter are:\n", + "`bayesian.sample()` runs the chains. The three numbers that matter are:\n", "\n", "- `samples` — how many draws to collect in total. More is better, at linear cost.\n", "- `burn` — generations discarded at the start, while the chains are still travelling towards the bulk of the posterior.\n", @@ -162,7 +162,7 @@ "metadata": {}, "outputs": [], "source": [ - "results = analysis.sample_posterior(samples=4000, burn=300, thin=2)\n", + "results = analysis.bayesian.sample(samples=4000, burn=300, thin=2)\n", "\n", "print(f'Collected {results.draws.shape[0]} draws for {results.draws.shape[1]} parameters.')" ] @@ -184,7 +184,7 @@ "metadata": {}, "outputs": [], "source": [ - "analysis.plot_trace()" + "analysis.bayesian.plot_trace()" ] }, { @@ -194,7 +194,7 @@ "source": [ "## Summarize the posterior\n", "\n", - "`posterior_summary()` reports the median and the 68% credible interval of each parameter, under the parameter's own name and unit. The interval is asymmetric in general, which is precisely the information a single symmetric error bar throws away." + "`bayesian.summary()` reports the median and the 68% credible interval of each parameter, under the parameter's own name and unit. The interval is asymmetric in general, which is precisely the information a single symmetric error bar throws away." ] }, { @@ -204,7 +204,7 @@ "metadata": {}, "outputs": [], "source": [ - "analysis.posterior_summary()" + "analysis.bayesian.summary()" ] }, { @@ -224,7 +224,7 @@ "metadata": {}, "outputs": [], "source": [ - "analysis.plot_corner()" + "analysis.bayesian.plot_corner()" ] }, { @@ -244,7 +244,7 @@ "metadata": {}, "outputs": [], "source": [ - "analysis.plot_posterior_predictive(n_draws=100)" + "analysis.bayesian.plot_posterior_predictive(n_draws=100)" ] }, { @@ -264,7 +264,7 @@ "metadata": {}, "outputs": [], "source": [ - "extended = analysis.extend_sampling(additional_samples=1000, thin=2)\n", + "extended = analysis.bayesian.extend(additional_samples=1000, thin=2)\n", "print(f'Chain now holds {extended.draws.shape[0]} draws.')" ] }, @@ -273,7 +273,7 @@ "id": "69793d31", "metadata": {}, "source": [ - "Chains are expensive, so they can be saved and reloaded with `analysis.save_chain(path)` and `analysis.load_chain(path)`. A reloaded chain can be summarized, plotted, or extended further, exactly like a fresh one." + "Chains are expensive, so they can be saved and reloaded with `analysis.bayesian.save(path)` and `analysis.bayesian.load(path)`. A reloaded chain can be summarized, plotted, or extended further, exactly like a fresh one." ] }, { @@ -285,9 +285,9 @@ "\n", "**Data without uncertainties.** If your data carries no variances, the weights fall back to 1, which means the sampler assumes a noise level of 1 in whatever units the intensity happens to be. Least-squares does not care, since that scale cancels out of the best-fit position, but a posterior *does*: its width scales directly with the assumed noise, so the credible intervals will be wrong by whatever factor the true noise differs from 1. Bayesian analysis is not a way to avoid needing uncertainties on your data.\n", "\n", - "**Sampling only some parameters.** `sample_posterior(parameters=[...])` restricts the chain to a subset, which is faster because the number of chains scales with the number of parameters. Be careful with the result: the other parameters are held *fixed*, which is not the same as averaging over them. The intervals you get are conditional on those fixed values, and will be too narrow whenever the parameters are correlated.\n", + "**Sampling only some parameters.** `bayesian.sample(parameters=[...])` restricts the chain to a subset, which is faster because the number of chains scales with the number of parameters. Be careful with the result: the other parameters are held *fixed*, which is not the same as averaging over them. The intervals you get are conditional on those fixed values, and will be too narrow whenever the parameters are correlated.\n", "\n", - "**Degenerate parameters.** As seen above, they are a modelling problem rather than a sampling one. `suggest_bounds()` returning absurd values, or a warning that the posterior has piled up against its bounds, are both signs to go back and look at the model." + "**Degenerate parameters.** As seen above, they are a modelling problem rather than a sampling one. `bayesian.suggest_bounds()` returning absurd values, or a warning that the posterior has piled up against its bounds, are both signs to go back and look at the model." ] } ], diff --git a/src/easydynamics/analysis/__init__.py b/src/easydynamics/analysis/__init__.py index 89a45cea3..89126ecdf 100644 --- a/src/easydynamics/analysis/__init__.py +++ b/src/easydynamics/analysis/__init__.py @@ -2,19 +2,21 @@ # SPDX-License-Identifier: BSD-3-Clause from easydynamics.analysis.analysis import Analysis -from easydynamics.analysis.bayesian_sampling import BayesianSamplingMixin from easydynamics.analysis.parameter_analysis import ParameterAnalysis from easydynamics.analysis.posterior import BoundsSuggestion from easydynamics.analysis.posterior import BoundsSuggestions from easydynamics.analysis.posterior import ParameterPosterior from easydynamics.analysis.posterior import PosteriorSummary +from easydynamics.analysis.posterior_labels import ParameterLabels +from easydynamics.analysis.posterior_sampling import PosteriorSampler __all__ = [ 'Analysis', - 'BayesianSamplingMixin', 'BoundsSuggestion', 'BoundsSuggestions', 'ParameterAnalysis', + 'ParameterLabels', 'ParameterPosterior', + 'PosteriorSampler', 'PosteriorSummary', ] diff --git a/src/easydynamics/analysis/analysis1d.py b/src/easydynamics/analysis/analysis1d.py index eaec732b1..5c8a68be3 100644 --- a/src/easydynamics/analysis/analysis1d.py +++ b/src/easydynamics/analysis/analysis1d.py @@ -12,7 +12,8 @@ from plopp.backends.matplotlib.figure import InteractiveFigure from easydynamics.analysis.analysis_base import AnalysisBase -from easydynamics.analysis.bayesian_sampling import BayesianSamplingMixin +from easydynamics.analysis.posterior_labels import ParameterLabels +from easydynamics.analysis.posterior_sampling import PosteriorSampler from easydynamics.convolution.convolution import Convolution from easydynamics.experiment import Experiment from easydynamics.sample_model import InstrumentModel @@ -26,15 +27,15 @@ from easydynamics.utils.utils import verify_Q_index -class Analysis1d(BayesianSamplingMixin, AnalysisBase): +class Analysis1d(AnalysisBase): """ For analysing one-dimensional data, i.e. intensity as function of energy for a single Q index. Is used primarily in the Analysis class, but can also be used on its own for simpler analyses. - In addition to least-squares fitting with :meth:`fit`, the posterior distribution of the free - parameters can be explored with :meth:`sample_posterior`; see - :class:`~easydynamics.analysis.bayesian_sampling.BayesianSamplingMixin`. + Besides least-squares fitting with :meth:`fit`, the posterior distribution of the free + parameters can be explored through :attr:`bayesian`; see + :class:`~easydynamics.analysis.posterior_sampling.PosteriorSampler`. Examples -------- @@ -121,7 +122,9 @@ def __init__( self._fit_result = None self._convolver = None self._convolver_is_dirty = True - self._init_bayesian_state() + self._fitter = None + self._fitter_is_dirty = True + self._bayesian = None super().__init__( display_name=display_name, @@ -253,32 +256,82 @@ def fit(self) -> FitResults: self._prepare_for_sampling() - x, y, weights = self._get_sampling_data() + x, y, weights = self._sampling_data() fit_result = self.fitter.fit(x=x, y=y, weights=weights) self._fit_result = fit_result return fit_result + @property + def fitter(self) -> EasyScienceFitter: + """ + The EasyScience Fitter used for fitting and sampling, built on first use. + + Exposed so the minimizer, tolerance, and maximum evaluation count can be configured + directly, e.g. ``analysis.fitter.switch_minimizer(AvailableMinimizers.Bumps)``. + + Returns + ------- + EasyScienceFitter + The cached Fitter. + """ + if self._fitter_is_dirty or self._fitter is None: + self._fitter = EasyScienceFitter( + fit_object=self, + fit_function=self.as_fit_function(), + ) + self._fitter_is_dirty = False + return self._fitter + + @property + def bayesian(self) -> PosteriorSampler: + """ + Bayesian posterior sampling for this Analysis, created on first use. + + Returns + ------- + PosteriorSampler + The sampler, which holds any chain that has been run. + """ + if self._bayesian is None: + self._bayesian = PosteriorSampler( + analysis=self, + sampling_data=self._sampling_data, + chain_parameters=self._chain_parameters, + parameter_labels=self._parameter_labels, + prepare=self._prepare_for_sampling, + ) + return self._bayesian + + def _invalidate_fitter(self) -> None: + """Mark the Fitter, and the Sampler built from it, as needing a rebuild.""" + self._fitter_is_dirty = True + self._invalidate_bayesian_sampler() + + def _invalidate_bayesian_sampler(self) -> None: + """Mark the Sampler as needing a rebuild, the data having changed.""" + if self._bayesian is not None: + self._bayesian.invalidate() + ############# - # Hooks for BayesianSamplingMixin + # The contract PosteriorSampler relies on ############# - def _build_bayesian_fitter(self) -> EasyScienceFitter: + def _parameter_labels(self) -> ParameterLabels: """ - Build the EasyScience Fitter for this Analysis. + Get labels for the chain's parameters. + + A single Q index holds one copy of each parameter, so nothing needs qualifying. Returns ------- - EasyScienceFitter - A Fitter bound to this Analysis and its fit function. + ParameterLabels + Labels over the current free parameters. """ - return EasyScienceFitter( - fit_object=self, - fit_function=self.as_fit_function(), - ) + return ParameterLabels(self._chain_parameters()) - def _get_sampling_data(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + def _sampling_data(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """ Get the finite data for the chosen Q index, as used by both fitting and sampling. @@ -292,7 +345,7 @@ def _get_sampling_data(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: ) return x, y, weights - def _get_chain_parameters(self) -> list[Parameter]: + def _chain_parameters(self) -> list[Parameter]: """ Get the free parameters of this Analysis. diff --git a/src/easydynamics/analysis/bayesian_sampling.py b/src/easydynamics/analysis/bayesian_sampling.py deleted file mode 100644 index 0ea38d8ef..000000000 --- a/src/easydynamics/analysis/bayesian_sampling.py +++ /dev/null @@ -1,993 +0,0 @@ -# SPDX-FileCopyrightText: 2026 EasyScience contributors -# SPDX-License-Identifier: BSD-3-Clause - -""" -Shared Bayesian MCMC sampling machinery for the Analysis classes. - -Everything that does not depend on how a particular Analysis is wired up lives here: caching the -Fitter and the Sampler, guarding the parameter bounds, restoring parameter values afterwards, and -turning raw draws into a readable summary. A concrete Analysis supplies the three things that do -differ, via :meth:`BayesianSamplingMixin._build_bayesian_fitter`, -:meth:`BayesianSamplingMixin._get_sampling_data`, and -:meth:`BayesianSamplingMixin._get_chain_parameters`. -""" - -from __future__ import annotations - -import json -import warnings -from pathlib import Path -from typing import TYPE_CHECKING -from typing import Any - -import numpy as np -from easyscience.fitting import AvailableMinimizers -from easyscience.fitting import Sampler - -from easydynamics.analysis.posterior import PosteriorSummary -from easydynamics.analysis.posterior import parameters_at_bounds -from easydynamics.analysis.posterior import suggest_bounds_for_parameters -from easydynamics.analysis.posterior import summarize_draws -from easydynamics.analysis.posterior import unbounded_parameters - -if TYPE_CHECKING: - import os - from collections.abc import Callable - - from easyscience.fitting.fitter import Fitter - from easyscience.fitting.sampler import SamplingResults - from easyscience.variable import Parameter - from matplotlib.figure import Figure - - from easydynamics.analysis.posterior import BoundsSuggestions - -# Suffix of the sidecar mapping chain columns to stable parameter names, written next to the BUMPS -# chain files by save_chain(). -_NAME_MAP_SUFFIX = '.parameter-names.json' - - -class BayesianSamplingMixin: - """ - Bayesian MCMC sampling on top of an Analysis, backed by the BUMPS DREAM sampler. - - Sampling explores the full posterior distribution of the free parameters rather than reporting - a single best-fit point, which is worth doing when parameters are correlated or their - uncertainties are strongly non-Gaussian -- both common in QENS. - - Running :meth:`fit` first is not required, but it helps: DREAM seeds its population in a small - ball around the parameters' current values, so starting from fitted values shortens the burn-in - needed to reach the typical set. - - Notes - ----- - All free parameters must have finite bounds before sampling, because in DREAM the bounds are - the prior. :meth:`suggest_bounds` proposes bounds for any parameter still missing one. - - Examples - -------- - ```python - analysis.fit() - analysis.suggest_bounds().apply() - results = analysis.sample_posterior(samples=10000, burn=2000, thin=10) - analysis.posterior_summary() - ``` - """ - - ############# - # Setup - ############# - - def _init_bayesian_state(self) -> None: - """ - Initialize the cached sampling state. - - Must be called by the concrete Analysis before any observer callback can fire, in the same - way as the other cached objects on the class. - """ - self._fitter = None - self._fitter_is_dirty = True - self._bayesian_sampler = None - self._bayesian_sampler_is_dirty = True - self._posterior_result = None - # Maps a chain column's unique_name to the parameter name it had when saved. Only populated - # by load_chain, because unique_names are per-session and do not survive a round trip. - self._chain_name_map = {} - - def _invalidate_fitter(self) -> None: - """ - Mark the cached Fitter and Sampler as needing a rebuild. - - The Sampler binds its data at construction, so anything that invalidates the Fitter - invalidates the Sampler too. - """ - self._fitter_is_dirty = True - self._bayesian_sampler_is_dirty = True - - def _invalidate_bayesian_sampler(self) -> None: - """ - Mark only the cached Sampler as needing a rebuild. - - Used when the data changed but the model did not. - """ - self._bayesian_sampler_is_dirty = True - - ############# - # Hooks for concrete Analysis classes - ############# - - def _build_bayesian_fitter(self) -> Fitter: - """ - Build the EasyScience Fitter (or MultiFitter) for this Analysis. - - Returns - ------- - Fitter - A configured Fitter or MultiFitter. - - Raises - ------ - NotImplementedError - If the concrete Analysis does not implement it. - """ - raise NotImplementedError('Subclasses must implement _build_bayesian_fitter.') - - def _get_sampling_data(self) -> tuple: - """ - Get the ``(x, y, weights)`` to bind to the Sampler. - - Each element is either an array (single dataset) or a list of arrays (MultiFitter). - - Returns - ------- - tuple - The ``(x, y, weights)`` triple. - - Raises - ------ - NotImplementedError - If the concrete Analysis does not implement it. - """ - raise NotImplementedError('Subclasses must implement _get_sampling_data.') - - def _get_chain_parameters(self) -> list[Parameter]: - """ - Get the free parameters that will appear as columns of the chain. - - Returns - ------- - list[Parameter] - The free parameters of the underlying model(s). - - Raises - ------ - NotImplementedError - If the concrete Analysis does not implement it. - """ - raise NotImplementedError('Subclasses must implement _get_chain_parameters.') - - def _prepare_for_sampling(self) -> None: - """ - Bring any cached computation up to date before a sampling run. - - The default does nothing; Analysis classes that cache a convolver override it. - """ - - ############# - # Properties - ############# - - @property - def fitter(self) -> Fitter: - """ - The EasyScience Fitter used for fitting and sampling, built on first use. - - Exposed so the minimizer, tolerance, and maximum evaluation count can be configured - directly, e.g. ``analysis.fitter.switch_minimizer(AvailableMinimizers.Bumps)``. - - Returns - ------- - Fitter - The cached Fitter or MultiFitter. - """ - if self._fitter_is_dirty or self._fitter is None: - self._fitter = self._build_bayesian_fitter() - self._fitter_is_dirty = False - return self._fitter - - @property - def bayesian_sampler(self) -> Sampler | None: - """ - The EasyScience Sampler holding the MCMC chain, or None before the first run. - - Named to avoid confusion with the SampleModel: this samples the posterior, not the sample. - - Returns - ------- - Sampler | None - The cached Sampler, or None if no chain has been started. - """ - return self._bayesian_sampler - - @property - def posterior_result(self) -> SamplingResults | None: - """ - The results of the most recent sampling run, or None if there has not been one. - - Returns - ------- - SamplingResults | None - The most recent sampling results. - """ - return self._posterior_result - - ############# - # Bounds - ############# - - def suggest_bounds( - self, - n_sigma: float = 10.0, - relative_pad: float = 0.2, - absolute_floor: float | None = None, - ) -> BoundsSuggestions: - """ - Propose finite bounds for free parameters that still have an infinite one. - - Nothing is changed until :meth:`BoundsSuggestions.apply` is called, so the proposal can be - reviewed first. Bounds that are already finite are never widened or narrowed, so physical - limits such as a non-negative area are left alone. - - Because the bounds act as a uniform prior in DREAM, a generous width is the safe choice: - too tight a bound truncates the posterior and understates the uncertainty. - - Parameters - ---------- - n_sigma : float, default=10.0 - How many standard deviations of the fitted uncertainty to allow on each side. - relative_pad : float, default=0.2 - Extra half-width as a fraction of the absolute parameter value. This guards against - minimizers that report a zero or absurdly small uncertainty. - absolute_floor : float | None, default=None - A minimum half-width in the parameter's own units, for when neither the uncertainty nor - the value carries the natural scale. - - Returns - ------- - BoundsSuggestions - The proposed bounds, which must be applied explicitly. - """ - return suggest_bounds_for_parameters( - self._get_chain_parameters(), - n_sigma=n_sigma, - relative_pad=relative_pad, - absolute_floor=absolute_floor, - ) - - def check_bounds_for_sampling(self) -> None: - """ - Verify that every free parameter has finite bounds. - - Raises - ------ - ValueError - If any free parameter has an infinite lower or upper bound. - """ - unbounded = unbounded_parameters(self._get_chain_parameters()) - if not unbounded: - return - names = ', '.join(parameter.name for parameter in unbounded) - raise ValueError( - f'Bayesian sampling requires finite bounds on every free parameter, because the ' - f'bounds act as the prior. These parameters are unbounded: {names}. ' - f'Set their min and max, or call suggest_bounds() to propose values.' - ) - - ############# - # Sampling - ############# - - def sample_posterior( - self, - samples: int = 10000, - burn: int = 2000, - thin: int = 10, - population: int | None = None, - parameters: list[Parameter] | list[str] | None = None, - **sampler_options: dict[str, Any], - ) -> SamplingResults: - """ - Draw samples from the posterior distribution of the free parameters. - - This starts a fresh chain, replacing any existing one; use :meth:`extend_sampling` to - continue a chain instead. Parameter values are restored to what they were beforehand, so - sampling never silently moves the model off its fitted values; use - :meth:`set_parameters_to_posterior_median` to adopt the posterior. - - Parameters - ---------- - samples : int, default=10000 - Number of raw samples to draw across all chains, before thinning. This is a guaranteed - minimum rather than an exact count. - burn : int, default=2000 - Burn-in generations to discard before collecting samples. - thin : int, default=10 - Thinning interval, which reduces autocorrelation between retained draws. - population : int | None, default=None - DREAM population scale factor: BUMPS runs ``ceil(population * n_parameters)`` chains. - parameters : list[Parameter] | list[str] | None, default=None - Restrict the chain to these parameters, given as Parameter objects or names. All other - free parameters are held fixed for the duration of the run. Note that holding a - parameter fixed is not the same as marginalizing over it: the resulting credible - intervals are conditional on the fixed values and will be too narrow if the parameters - are correlated. The default samples every free parameter. - **sampler_options : dict[str, Any] - Forwarded to the EasyScience Sampler, e.g. ``sampler_kwargs``, ``progress_callback``, - or ``abort_test``. - - Returns - ------- - SamplingResults - The sampling results, also stored on :attr:`posterior_result`. - """ - return self._run_sampling( - parameters=parameters, - run=lambda sampler: sampler.sample( - samples=samples, - burn=burn, - thin=thin, - population=population, - **sampler_options, - ), - ) - - def extend_sampling( - self, - additional_samples: int = 5000, - thin: int = 10, - parameters: list[Parameter] | list[str] | None = None, - **sampler_options: dict[str, Any], - ) -> SamplingResults: - """ - Continue the existing chain with additional samples. - - Parameters - ---------- - additional_samples : int, default=5000 - Number of additional samples to draw, in the same units as ``samples``. - thin : int, default=10 - Thinning interval for the retained draws. - parameters : list[Parameter] | list[str] | None, default=None - The same restriction as in :meth:`sample_posterior`. Pass the same value that started - the chain, since the chain's columns cannot change on extension. - **sampler_options : dict[str, Any] - Forwarded to the EasyScience Sampler. - - Returns - ------- - SamplingResults - The sampling results for the full extended chain. - - Raises - ------ - RuntimeError - If there is no chain to extend. - """ - if self._bayesian_sampler is None: - raise RuntimeError( - 'No chain to extend. Call sample_posterior() or load_chain() first.' - ) - return self._run_sampling( - parameters=parameters, - run=lambda sampler: sampler.extend( - additional_samples=additional_samples, - thin=thin, - **sampler_options, - ), - reuse_sampler=True, - ) - - def _run_sampling( - self, - parameters: list[Parameter] | list[str] | None, - run: Callable[[Sampler], SamplingResults], - reuse_sampler: bool = False, - ) -> SamplingResults: - """ - Run a sampling operation with all the surrounding guards in place. - - Checks the bounds, switches the minimizer to BUMPS, optionally holds parameters fixed, - runs, and then restores the parameter values, fixed flags, and minimizer. - - Parameters - ---------- - parameters : list[Parameter] | list[str] | None - Parameters to restrict the chain to, or None for all free parameters. - run : Callable[[Sampler], SamplingResults] - The operation to perform on the prepared Sampler. - reuse_sampler : bool, default=False - Whether to reuse the cached Sampler rather than rebuilding it. Required when extending - a chain, since the chain lives on the Sampler. - - Returns - ------- - SamplingResults - The results of the run. - - Raises - ------ - RuntimeError - If the BUMPS sampler fails while removing outlier chains, which points at degenerate - parameters. - """ - held_fixed = self._resolve_parameters_to_hold_fixed(parameters) - self._warn_about_held_parameters(held_fixed) - - with _FixedParameters(held_fixed): - self.check_bounds_for_sampling() - self._prepare_for_sampling() - - chain_parameters = self._get_chain_parameters() - saved_values = [(p, p.value) for p in chain_parameters] - - fitter = self.fitter - original_minimizer = fitter.minimizer.enum - fitter.switch_minimizer(AvailableMinimizers.Bumps) - try: - sampler = self._get_or_build_sampler(reuse_sampler=reuse_sampler) - results = run(sampler) - except IndexError as error: - # BUMPS' own outlier removal indexes past the end of its buffer when chains - # scatter wildly, which in practice means the model is not identifiable. The bare - # IndexError says nothing useful, so point at the likely cause instead. - raise RuntimeError( - 'The BUMPS sampler failed while removing outlier chains. This usually means ' - 'the chains scattered because two or more free parameters are degenerate, so ' - 'the data cannot determine them separately. Check for degenerate parameters ' - "and fix one of them, or retry with sampler_kwargs={'outliers': 'none'}." - ) from error - finally: - fitter.switch_minimizer(original_minimizer) - for parameter, value in saved_values: - parameter.value = value - - # A fresh chain is labelled with this session's unique names, so any mapping left over - # from a loaded chain no longer applies. - self._chain_name_map = { - parameter.unique_name: parameter.name for parameter in chain_parameters - } - - self._posterior_result = results - self._warn_about_bounds_occupancy(results, self._resolve_chain_parameters(results)) - return results - - def _get_or_build_sampler(self, reuse_sampler: bool) -> Sampler: - """ - Get the cached Sampler, rebuilding it if the data or model changed. - - Parameters - ---------- - reuse_sampler : bool - Whether to reuse the cached Sampler even if it is marked dirty. - - Returns - ------- - Sampler - The Sampler to run. - """ - needs_rebuild = self._bayesian_sampler is None or ( - self._bayesian_sampler_is_dirty and not reuse_sampler - ) - if needs_rebuild: - x, y, weights = self._get_sampling_data() - self._bayesian_sampler = Sampler(self.fitter, x, y, weights=weights) - self._bayesian_sampler_is_dirty = False - return self._bayesian_sampler - - def _resolve_parameters_to_hold_fixed( - self, - parameters: list[Parameter] | list[str] | None, - ) -> list[Parameter]: - """ - Work out which free parameters must be held fixed to honour a subset request. - - Parameters - ---------- - parameters : list[Parameter] | list[str] | None - The requested subset, as Parameter objects or names, or None for all free parameters. - - Returns - ------- - list[Parameter] - The free parameters that are not in the requested subset. - - Raises - ------ - TypeError - If parameters is not a list of Parameters or strings, or None. - ValueError - If a requested name does not match any free parameter, or the subset is empty. - """ - if parameters is None: - return [] - if not isinstance(parameters, (list, tuple)): - raise TypeError('parameters must be a list of Parameters, a list of names, or None.') - - free = self._get_chain_parameters() - by_name = {parameter.name: parameter for parameter in free} - requested = [] - for entry in parameters: - if isinstance(entry, str): - if entry not in by_name: - available = ', '.join(sorted(by_name)) - raise ValueError(f'No free parameter named {entry!r}. Available: {available}.') - requested.append(by_name[entry]) - elif hasattr(entry, 'unique_name'): - requested.append(entry) - else: - raise TypeError( - 'parameters must contain Parameter objects or parameter names (strings).' - ) - - requested_unique_names = {parameter.unique_name for parameter in requested} - if not requested_unique_names: - raise ValueError('parameters must name at least one parameter to sample.') - return [ - parameter for parameter in free if parameter.unique_name not in requested_unique_names - ] - - @staticmethod - def _warn_about_held_parameters(held_fixed: list[Parameter]) -> None: - """ - Warn that holding parameters fixed makes the credible intervals conditional. - - Parameters - ---------- - held_fixed : list[Parameter] - The parameters being held fixed for the run. - """ - if not held_fixed: - return - names = ', '.join(parameter.name for parameter in held_fixed) - warnings.warn( - ( - f'Holding these parameters fixed while sampling: {names}. ' - f'Fixing a parameter is not the same as marginalizing over it, so the resulting ' - f'credible intervals are conditional on these values and will be too narrow if ' - f'the parameters are correlated.' - ), - UserWarning, - stacklevel=4, - ) - - @staticmethod - def _warn_about_bounds_occupancy( - results: SamplingResults, - parameters_by_column: list[Parameter | None], - ) -> None: - """ - Warn when the posterior has piled up against a bound. - - Parameters - ---------- - results : SamplingResults - The sampling results to inspect. - parameters_by_column : list[Parameter | None] - The parameter for each column of the chain, or None where none could be matched. - """ - piled_up = parameters_at_bounds(results.draws, parameters_by_column) - if not piled_up: - return - details = ', '.join( - f'{name} ({fraction:.0%} of draws)' for name, fraction in piled_up.items() - ) - warnings.warn( - ( - f'The posterior is piled up against the bounds for: {details}. ' - f'The bounds, rather than the data, are setting these credible intervals. ' - f'Widen the bounds, or check whether these parameters are degenerate with others.' - ), - UserWarning, - stacklevel=4, - ) - - ############# - # Results - ############# - - def posterior_summary(self) -> PosteriorSummary: - """ - Summarize the marginal posterior of each sampled parameter. - - Reports the median and the 68% credible interval under the parameter's own name and unit, - rather than the opaque unique name the sampler uses internally. Requires a completed - sampling run. - - Returns - ------- - PosteriorSummary - One entry per sampled parameter. - """ - results = self._require_posterior_result() - return summarize_draws( - draws=results.draws, - fallback_names=self._chain_display_names(results), - parameters_by_column=self._resolve_chain_parameters(results), - ) - - def set_parameters_to_posterior_median(self) -> list[Parameter]: - """ - Set every sampled parameter to the median of its marginal posterior. - - Note that the vector of marginal medians is not in general the same as the - highest-posterior point, and for strongly correlated parameters it need not even be a good - fit. Requires a completed sampling run. - - Returns - ------- - list[Parameter] - The parameters that were changed. - """ - results = self._require_posterior_result() - changed = [] - for column, parameter in enumerate(self._resolve_chain_parameters(results)): - if parameter is None: - continue - parameter.value = float(np.median(results.draws[:, column])) - changed.append(parameter) - return changed - - def _require_posterior_result(self) -> SamplingResults: - """ - Get the stored sampling results, raising if there are none. - - Returns - ------- - SamplingResults - The most recent sampling results. - - Raises - ------ - RuntimeError - If no sampling has been run yet. - """ - if self._posterior_result is None: - raise RuntimeError( - 'No posterior samples yet. Call sample_posterior() or load_chain() first.' - ) - return self._posterior_result - - ############# - # Persistence - ############# - - def save_chain(self, path: str | os.PathLike) -> None: - """ - Save the MCMC chain to disk. - - Writes the BUMPS chain files alongside a sidecar recording the parameter names and a - fingerprint of the data that was sampled. - - Parameters - ---------- - path : str | os.PathLike - Path prefix for the chain files. - - Raises - ------ - RuntimeError - If there is no chain to save. - """ - if self._bayesian_sampler is None: - raise RuntimeError('No chain to save. Call sample_posterior() first.') - self._bayesian_sampler.save(path) - # The BUMPS sidecar records unique names, which are handed out per session and so mean - # nothing on reload. Record the parameter names alongside them, which are stable. - Path(f'{path}{_NAME_MAP_SUFFIX}').write_text( - json.dumps(self._chain_name_map, indent=2), - encoding='utf-8', - ) - - def load_chain(self, path: str | os.PathLike, skip: int = 0) -> SamplingResults: - """ - Load a previously saved MCMC chain. - - The loaded chain can be inspected, summarized, or continued with :meth:`extend_sampling`. A - chain saved from different data loads with a warning. - - Parameters - ---------- - path : str | os.PathLike - The path prefix the chain was saved under. - skip : int, default=0 - Number of initial samples to skip when reading the chain. - - Returns - ------- - SamplingResults - The loaded sampling results, also stored on :attr:`posterior_result`. - """ - self._prepare_for_sampling() - name_map_path = Path(f'{path}{_NAME_MAP_SUFFIX}') - if name_map_path.is_file(): - self._chain_name_map = json.loads(name_map_path.read_text(encoding='utf-8')) - else: - self._chain_name_map = {} - warnings.warn( - ( - f'No parameter-name sidecar found at {name_map_path}. The chain will be ' - f'reported under the internal names it was saved with, because those cannot ' - f'be matched to this Analysis.' - ), - UserWarning, - stacklevel=2, - ) - - fitter = self.fitter - original_minimizer = fitter.minimizer.enum - fitter.switch_minimizer(AvailableMinimizers.Bumps) - try: - sampler = self._get_or_build_sampler(reuse_sampler=False) - results = sampler.load_state(path, skip=skip) - finally: - fitter.switch_minimizer(original_minimizer) - self._posterior_result = results - return results - - ############# - # Plotting - ############# - - def plot_trace(self, **kwargs: dict[str, Any]) -> Figure: - """ - Plot the chain trace of each sampled parameter. - - A well-mixed chain looks like a "hairy caterpillar" with no drift; visible trends mean the - chain has not converged and needs a longer burn-in. Requires a completed sampling run. - - Parameters - ---------- - **kwargs : dict[str, Any] - Forwarded to :func:`easydynamics.utils.posterior_plotting.plot_trace`. - - Returns - ------- - Figure - The matplotlib Figure. - """ - from easydynamics.utils.posterior_plotting import plot_trace - - results = self._require_posterior_result() - return plot_trace( - draws=results.draws, - logp=results.logp, - names=self._chain_display_names(results), - title=self.display_name, - **kwargs, - ) - - def plot_corner(self, **kwargs: dict[str, Any]) -> Figure: - """ - Plot the marginal and pairwise posterior distributions. - - Diagonal panels show each parameter's marginal distribution; off-diagonal panels show the - joint distribution of a pair, where a strong diagonal ridge means the two are correlated. - Requires a completed sampling run. - - Parameters - ---------- - **kwargs : dict[str, Any] - Forwarded to :func:`easydynamics.utils.posterior_plotting.plot_corner`. - - Returns - ------- - Figure - The matplotlib Figure. - """ - from easydynamics.utils.posterior_plotting import plot_corner - - results = self._require_posterior_result() - return plot_corner( - draws=results.draws, - names=self._chain_display_names(results), - title=self.display_name, - **kwargs, - ) - - def plot_posterior_predictive( - self, - n_draws: int = 200, - credible_interval: float = 68.0, - **kwargs: dict[str, Any], - ) -> Figure: - """ - Plot the data against the credible band implied by the posterior. - - The model is re-evaluated for a random subset of the posterior draws, and the spread of - those curves becomes the band. Data straying outside the band systematically points at a - model that is missing something, rather than at parameters that need tuning. Requires a - completed sampling run. - - Parameters - ---------- - n_draws : int, default=200 - How many posterior draws to evaluate the model for. Each draw costs one full model - evaluation, so this trades smoothness of the band against time. - credible_interval : float, default=68.0 - Width of the credible band, as a percentage. - **kwargs : dict[str, Any] - Forwarded to :func:`easydynamics.utils.posterior_plotting.plot_posterior_predictive`. - - Returns - ------- - Figure - The matplotlib Figure. - - Raises - ------ - NotImplementedError - If this Analysis binds a list of datasets rather than a single one. - ValueError - If n_draws is not a positive integer. - """ - from easydynamics.utils.posterior_plotting import plot_posterior_predictive - - if not isinstance(n_draws, int) or isinstance(n_draws, bool) or n_draws < 1: - raise ValueError(f'n_draws must be a positive integer. Got {n_draws}.') - - results = self._require_posterior_result() - x, y, weights = self._get_sampling_data() - if isinstance(x, (list, tuple)): - raise NotImplementedError( - 'plot_posterior_predictive supports a single dataset only. Plot each dataset ' - 'from its own Analysis1d instead.' - ) - - predictions = self._evaluate_over_draws(results, x, n_draws) - y_err = None if weights is None else 1.0 / np.asarray(weights) - return plot_posterior_predictive( - x=np.asarray(x), - y=np.asarray(y), - predictions=predictions, - y_err=y_err, - title=self.display_name, - credible_interval=credible_interval, - **kwargs, - ) - - def _evaluate_over_draws( - self, - results: SamplingResults, - x: np.ndarray, - n_draws: int, - ) -> np.ndarray: - """ - Evaluate the model once per posterior draw, restoring the parameters afterwards. - - Parameters - ---------- - results : SamplingResults - The sampling results supplying the draws. - x : np.ndarray - The independent variable to evaluate the model on. - n_draws : int - How many draws to evaluate. Draws are taken evenly across the chain. - - Returns - ------- - np.ndarray - Model evaluations, shape ``(n_selected, len(x))``. - """ - self._prepare_for_sampling() - - columns = [ - (parameter, column) - for column, parameter in enumerate(self._resolve_chain_parameters(results)) - if parameter is not None - ] - saved_values = [(parameter, parameter.value) for parameter, _ in columns] - - total = results.draws.shape[0] - indices = np.unique(np.linspace(0, total - 1, min(n_draws, total)).astype(int)) - - fit_function = self.fitter.fit_function - predictions = [] - try: - for index in indices: - for parameter, column in columns: - parameter.value = float(results.draws[index, column]) - predictions.append(np.asarray(fit_function(x))) - finally: - for parameter, value in saved_values: - parameter.value = value - - return np.vstack(predictions) - - def _resolve_chain_parameters(self, results: SamplingResults) -> list[Parameter | None]: - """ - Match each column of the chain to one of this Analysis's parameters. - - Columns are matched on ``unique_name`` first. That fails for a chain loaded from disk, - because unique names are handed out per session, so a saved chain also records the - parameter names and those are used as a fallback. - - Parameters - ---------- - results : SamplingResults - The sampling results whose columns should be matched. - - Returns - ------- - list[Parameter | None] - The parameter for each column, or None where no match could be made. - """ - parameters = self._get_chain_parameters() - by_unique_name = {p.unique_name: p for p in parameters} - by_name = {p.name: p for p in parameters} - resolved = [] - for unique_name in results.param_names: - parameter = by_unique_name.get(unique_name) - if parameter is None: - saved_name = self._chain_name_map.get(unique_name) - parameter = None if saved_name is None else by_name.get(saved_name) - resolved.append(parameter) - return resolved - - def _chain_display_names(self, results: SamplingResults) -> list[str]: - """ - Translate the chain's column names into parameter names. - - Parameters - ---------- - results : SamplingResults - The sampling results whose columns should be named. - - Returns - ------- - list[str] - One name per column of the chain. - """ - resolved = self._resolve_chain_parameters(results) - return [ - self._chain_name_map.get(unique_name, unique_name) - if parameter is None - else parameter.name - for unique_name, parameter in zip(results.param_names, resolved, strict=True) - ] - - -class _FixedParameters: - """ - Context manager that temporarily fixes parameters and restores their flags on exit. - """ - - def __init__(self, parameters: list[Parameter]) -> None: - """ - Initialize the context manager. - - Parameters - ---------- - parameters : list[Parameter] - The parameters to hold fixed for the duration of the block. - """ - self._parameters = list(parameters) - self._saved: list[tuple[Parameter, bool]] = [] - - def __enter__(self) -> None: - """ - Fix the parameters, remembering their previous state. - """ - self._saved = [(parameter, parameter.fixed) for parameter in self._parameters] - for parameter in self._parameters: - parameter.fixed = True - - def __exit__(self, *_exc_info: object) -> None: - """ - Restore the previous fixed state of every parameter. - - Parameters - ---------- - *_exc_info : object - Exception information, ignored. - """ - for parameter, was_fixed in self._saved: - parameter.fixed = was_fixed diff --git a/src/easydynamics/analysis/posterior.py b/src/easydynamics/analysis/posterior.py index e22ae87d9..4853bb1a5 100644 --- a/src/easydynamics/analysis/posterior.py +++ b/src/easydynamics/analysis/posterior.py @@ -40,6 +40,8 @@ class BoundsSuggestion: ---------- parameter : Parameter The parameter the suggestion applies to. + label : str + The name the parameter is reported under, qualified where several share a name. suggested_min : float The proposed lower bound. Equal to the parameter's current lower bound when that is already finite. @@ -51,6 +53,7 @@ class BoundsSuggestion: """ parameter: Parameter + label: str suggested_min: float suggested_max: float reason: str @@ -178,7 +181,8 @@ def __repr__(self) -> str: if not self._suggestions: return 'BoundsSuggestions(no free parameters)' - header = f'{"parameter":<28s} {"current":>26s} {"suggested":>26s}' + width = max(len('parameter'), *(len(s.label) for s in self._suggestions)) + header = f'{"parameter":<{width}s} {"current":>26s} {"suggested":>26s}' lines = ['BoundsSuggestions', header, '-' * len(header)] for s in self._suggestions: current = f'({s.parameter.min:.4g}, {s.parameter.max:.4g})' @@ -186,7 +190,7 @@ def __repr__(self) -> str: suggested = f'-- {s.reason}' else: suggested = f'({s.suggested_min:.4g}, {s.suggested_max:.4g})' - lines.append(f'{s.parameter.name:<28s} {current:>26s} {suggested:>26s}') + lines.append(f'{s.label:<{width}s} {current:>26s} {suggested:>26s}') attention = self.needing_attention if attention: @@ -199,6 +203,7 @@ def __repr__(self) -> str: def suggest_bounds_for_parameters( parameters: list[Parameter], + labels: list[str] | None = None, n_sigma: float = 10.0, relative_pad: float = 0.2, absolute_floor: float | None = None, @@ -223,6 +228,9 @@ def suggest_bounds_for_parameters( ---------- parameters : list[Parameter] The parameters to propose bounds for. + labels : list[str] | None, default=None + The name to report each parameter under, one per parameter. Defaults to the parameters' own + names. n_sigma : float, default=10.0 How many standard deviations of the parameter's fitted uncertainty to allow on each side. relative_pad : float, default=0.2 @@ -242,20 +250,24 @@ def suggest_bounds_for_parameters( if absolute_floor is not None: _verify_nonneg_number(absolute_floor, 'absolute_floor') + if labels is None: + labels = [parameter.name for parameter in parameters] suggestions = [ _suggest_bounds_for_parameter( parameter=parameter, + label=label, n_sigma=n_sigma, relative_pad=relative_pad, absolute_floor=absolute_floor, ) - for parameter in parameters + for parameter, label in zip(parameters, labels, strict=True) ] return BoundsSuggestions(suggestions) def _suggest_bounds_for_parameter( parameter: Parameter, + label: str, n_sigma: float, relative_pad: float, absolute_floor: float | None, @@ -267,6 +279,8 @@ def _suggest_bounds_for_parameter( ---------- parameter : Parameter The parameter to propose bounds for. + label : str + The name to report the parameter under. n_sigma : float How many standard deviations to allow on each side. relative_pad : float @@ -288,6 +302,7 @@ def _suggest_bounds_for_parameter( if min_is_finite and max_is_finite: return BoundsSuggestion( parameter=parameter, + label=label, suggested_min=current_min, suggested_max=current_max, reason='', @@ -298,6 +313,7 @@ def _suggest_bounds_for_parameter( if not np.isfinite(value): return BoundsSuggestion( parameter=parameter, + label=label, suggested_min=current_min, suggested_max=current_max, reason='value is not finite', @@ -312,6 +328,7 @@ def _suggest_bounds_for_parameter( if not np.isfinite(half_width) or half_width <= 0: return BoundsSuggestion( parameter=parameter, + label=label, suggested_min=current_min, suggested_max=current_max, reason='no scale information (zero value and uncertainty)', @@ -319,6 +336,7 @@ def _suggest_bounds_for_parameter( return BoundsSuggestion( parameter=parameter, + label=label, suggested_min=current_min if min_is_finite else value - half_width, suggested_max=current_max if max_is_finite else value + half_width, reason='', @@ -527,13 +545,14 @@ def __repr__(self) -> str: if not self._entries: return 'PosteriorSummary(no parameters)' + width = max(len('parameter'), *(len(e.name) for e in self._entries)) header = ( - f'{"parameter":<28s} {"unit":>10s} {"median":>14s} ' + f'{"parameter":<{width}s} {"unit":>10s} {"median":>14s} ' f'{"-":>12s} {"+":>12s} {"current":>14s}' ) lines = ['PosteriorSummary', header, '-' * len(header)] lines.extend( - f'{e.name:<28s} {e.unit:>10s} {e.median:>14.5g} ' + f'{e.name:<{width}s} {e.unit:>10s} {e.median:>14.5g} ' f'{e.minus:>12.4g} {e.plus:>12.4g} {e.value:>14.5g}' for e in self._entries ) @@ -542,7 +561,7 @@ def __repr__(self) -> str: def summarize_draws( draws: np.ndarray, - fallback_names: list[str], + labels: list[str], parameters_by_column: list[Parameter | None], ) -> PosteriorSummary: """ @@ -556,8 +575,8 @@ def summarize_draws( ---------- draws : np.ndarray Posterior draws, shape ``(n_draws, n_parameters)``. - fallback_names : list[str] - Label to use for any column with no matching parameter, one per column. + labels : list[str] + The label to report each column under, one per column. parameters_by_column : list[Parameter | None] The parameter for each column of ``draws``, or None where none could be matched. @@ -573,7 +592,7 @@ def summarize_draws( ) entries.append( ParameterPosterior( - name=fallback_names[column] if parameter is None else parameter.name, + name=labels[column], unit='' if parameter is None else str(parameter.unit), median=median, lower=lower, diff --git a/src/easydynamics/analysis/posterior_labels.py b/src/easydynamics/analysis/posterior_labels.py new file mode 100644 index 000000000..84ff9212e --- /dev/null +++ b/src/easydynamics/analysis/posterior_labels.py @@ -0,0 +1,182 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +""" +Naming the columns of an MCMC chain. + +The sampler labels its columns with each parameter's ``unique_name`` -- ``Parameter_4`` and the +like -- which is not what a user recognises, and which is handed out per session so it does not +survive a saved chain either. This turns those columns back into readable labels. +""" + +from __future__ import annotations + +from collections import Counter +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Callable + + from easyscience.variable import Parameter + + +class ParameterLabels: + """ + Readable labels and units for the columns of a chain. + + Built once for a fixed set of parameters, so the name counts and lookups are computed a single + time. Doing this per column instead is quadratic in the parameter count, which is seconds of + work for an analysis with many Q values. + + Parameters + ---------- + parameters : list[Parameter] + The parameters that can appear as columns. + qualify : Callable[[Parameter], str | None] | None, default=None + Returns a qualifier for a parameter whose name is shared with another, for example its Q + index. Only consulted when the bare name really is ambiguous, so an analysis with nothing + to disambiguate keeps its short names. Returning None leaves the name unqualified. + """ + + def __init__( + self, + parameters: list[Parameter], + qualify: Callable[[Parameter], str | None] | None = None, + ) -> None: + self._parameters = list(parameters) + self._qualify = qualify + self._counts = Counter(parameter.name for parameter in self._parameters) + self._by_unique_name = {p.unique_name: p for p in self._parameters} + self._by_label = {self.label(p): p for p in self._parameters} + + @property + def parameters(self) -> list[Parameter]: + """ + The parameters these labels describe. + + Returns + ------- + list[Parameter] + The parameters given at construction. + """ + return list(self._parameters) + + def label(self, parameter: Parameter) -> str: + """ + Get the label a parameter is reported under. + + Parameters + ---------- + parameter : Parameter + The parameter to label. + + Returns + ------- + str + The parameter's name, qualified only where that name is shared with another parameter. + """ + if self._counts[parameter.name] <= 1 or self._qualify is None: + return parameter.name + qualifier = self._qualify(parameter) + return parameter.name if qualifier is None else f'{parameter.name} ({qualifier})' + + def name_map(self) -> dict[str, str]: + """ + Map each parameter's ``unique_name`` to its label. + + Saved alongside a chain, because unique names are per-session: without this a reloaded + chain cannot be matched back to any parameter. + + Returns + ------- + dict[str, str] + Mapping of unique name to label. + """ + return {p.unique_name: self.label(p) for p in self._parameters} + + def resolve( + self, + column_names: list[str], + saved_labels: dict[str, str] | None = None, + ) -> list[Parameter | None]: + """ + Match each column of a chain to a parameter. + + Columns are matched on ``unique_name`` first. That fails for a chain loaded from disk, + where the saved labels are used instead. + + Parameters + ---------- + column_names : list[str] + The sampler's name for each column. + saved_labels : dict[str, str] | None, default=None + Mapping of unique name to label, as recorded when a chain was saved. + + Returns + ------- + list[Parameter | None] + The parameter for each column, or None where no match could be made. + """ + saved_labels = saved_labels or {} + resolved = [] + for unique_name in column_names: + parameter = self._by_unique_name.get(unique_name) + if parameter is None: + parameter = self._by_label.get(saved_labels.get(unique_name, '')) + resolved.append(parameter) + return resolved + + def display_names( + self, + column_names: list[str], + saved_labels: dict[str, str] | None = None, + ) -> list[str]: + """ + Get a readable label for each column of a chain. + + Parameters + ---------- + column_names : list[str] + The sampler's name for each column. + saved_labels : dict[str, str] | None, default=None + Mapping of unique name to label, as recorded when a chain was saved. + + Returns + ------- + list[str] + One label per column, falling back to the saved label and then to the raw column name. + """ + saved_labels = saved_labels or {} + return [ + saved_labels.get(unique_name, unique_name) + if parameter is None + else self.label(parameter) + for unique_name, parameter in zip( + column_names, self.resolve(column_names, saved_labels), strict=True + ) + ] + + def units( + self, + column_names: list[str], + saved_labels: dict[str, str] | None = None, + ) -> list[str]: + """ + Get the unit of each column of a chain. + + Parameters + ---------- + column_names : list[str] + The sampler's name for each column. + saved_labels : dict[str, str] | None, default=None + Mapping of unique name to label, as recorded when a chain was saved. + + Returns + ------- + list[str] + One unit per column, empty where no parameter could be matched. + """ + return [ + '' if parameter is None else str(parameter.unit) + for parameter in self.resolve(column_names, saved_labels) + ] diff --git a/src/easydynamics/analysis/posterior_sampling.py b/src/easydynamics/analysis/posterior_sampling.py new file mode 100644 index 000000000..f782b2d00 --- /dev/null +++ b/src/easydynamics/analysis/posterior_sampling.py @@ -0,0 +1,933 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +""" +Bayesian MCMC sampling for the Analysis classes, backed by the BUMPS DREAM sampler. + +The sampler is composed into an Analysis rather than inherited by it: an Analysis exposes one +``bayesian`` property, and everything to do with sampling lives here instead of being mixed into +three classes. Labelling lives in :mod:`easydynamics.analysis.posterior_labels` and the figures in +:mod:`easydynamics.utils.posterior_plotting`; this module only runs chains. +""" + +from __future__ import annotations + +import json +import warnings +from pathlib import Path +from typing import TYPE_CHECKING +from typing import Any + +import numpy as np +from easyscience.fitting import AvailableMinimizers +from easyscience.fitting import Sampler + +from easydynamics.analysis.posterior import parameters_at_bounds +from easydynamics.analysis.posterior import suggest_bounds_for_parameters +from easydynamics.analysis.posterior import summarize_draws +from easydynamics.analysis.posterior import unbounded_parameters + +if TYPE_CHECKING: + import os + from collections.abc import Callable + + from easyscience.fitting.sampler import SamplingResults + from easyscience.variable import Parameter + from matplotlib.figure import Figure + + from easydynamics.analysis.posterior import BoundsSuggestions + from easydynamics.analysis.posterior import PosteriorSummary + from easydynamics.analysis.posterior_labels import ParameterLabels + +# Suffix of the sidecar mapping chain columns to stable labels, written next to the BUMPS chain +# files by save(). +_LABEL_MAP_SUFFIX = '.parameter-names.json' + + +class PosteriorSampler: + """ + Draws samples from the posterior distribution of an Analysis' free parameters. + + Reached as ``analysis.bayesian``. Sampling explores the whole posterior rather than reporting a + single best-fit point, which is worth doing when parameters are correlated or their + uncertainties are strongly non-Gaussian, both common in QENS. + + Running a fit first is not required, but it helps: DREAM seeds its population in a small ball + around the parameters' current values, so starting from fitted values shortens the burn-in. + + The Analysis passes in everything that differs between the Analysis classes, so this class + needs no knowledge of how any of them is built. + + Parameters + ---------- + analysis : object + The Analysis being sampled, used for its ``display_name`` and its ``fitter``. + sampling_data : Callable[[], tuple] + Returns the ``(x, y, weights)`` to bind to the sampler. Each is an array, or a list of + arrays for a multi-dataset fit. + chain_parameters : Callable[[], list[Parameter]] + Returns the free parameters that will form the chain's columns. + parameter_labels : Callable[[], ParameterLabels] + Returns labels for those parameters. + prepare : Callable[[], None] | None, default=None + Brings any cached computation on the Analysis up to date before a run. + + Notes + ----- + Every free parameter must have finite bounds before sampling, because in DREAM the bounds are + the prior. :meth:`suggest_bounds` proposes bounds for any parameter still missing one. + + Examples + -------- + ```python + analysis.fit() + analysis.bayesian.suggest_bounds().apply() + analysis.bayesian.sample(samples=10000, burn=2000, thin=10) + analysis.bayesian.summary() + ``` + """ + + def __init__( + self, + analysis: object, + sampling_data: Callable[[], tuple], + chain_parameters: Callable[[], list[Parameter]], + parameter_labels: Callable[[], ParameterLabels], + prepare: Callable[[], None] | None = None, + ) -> None: + self._analysis = analysis + self._sampling_data = sampling_data + self._chain_parameters = chain_parameters + self._parameter_labels = parameter_labels + self._prepare_hook = prepare + self._sampler: Sampler | None = None + self._sampler_is_dirty = True + self._results: SamplingResults | None = None + # Maps a chain column's unique_name to the label it had when saved. Only populated by + # load(), because unique names are per-session and do not survive a round trip. + self._saved_labels: dict[str, str] = {} + + ############# + # State + ############# + + def invalidate(self) -> None: + """ + Mark the underlying Sampler as needing a rebuild. + + Called by the Analysis when its data changes, since the Sampler binds its data at + construction. + """ + self._sampler_is_dirty = True + + @property + def sampler(self) -> Sampler | None: + """ + The EasyScience Sampler holding the chain, or None before the first run. + + Returns + ------- + Sampler | None + The cached Sampler. + """ + return self._sampler + + @property + def results(self) -> SamplingResults | None: + """ + The results of the most recent run, or None if there has not been one. + + Returns + ------- + SamplingResults | None + The most recent sampling results. + """ + return self._results + + ############# + # Bounds + ############# + + def suggest_bounds( + self, + n_sigma: float = 10.0, + relative_pad: float = 0.2, + absolute_floor: float | None = None, + ) -> BoundsSuggestions: + """ + Propose finite bounds for free parameters that still have an infinite one. + + Nothing changes until :meth:`BoundsSuggestions.apply` is called, so the proposal can be + reviewed first. Bounds that are already finite are never widened or narrowed, so physical + limits such as a non-negative area are left alone. + + Because the bounds act as a uniform prior in DREAM, a generous width is the safe choice: + too tight a bound truncates the posterior and understates the uncertainty. + + Parameters + ---------- + n_sigma : float, default=10.0 + How many standard deviations of the fitted uncertainty to allow on each side. + relative_pad : float, default=0.2 + Extra half-width as a fraction of the absolute parameter value, guarding against + minimizers that report a zero or absurdly small uncertainty. + absolute_floor : float | None, default=None + A minimum half-width in the parameter's own units, for when neither the uncertainty nor + the value carries the natural scale. + + Returns + ------- + BoundsSuggestions + The proposed bounds, which must be applied explicitly. + """ + labels = self._labels() + return suggest_bounds_for_parameters( + labels.parameters, + labels=[labels.label(parameter) for parameter in labels.parameters], + n_sigma=n_sigma, + relative_pad=relative_pad, + absolute_floor=absolute_floor, + ) + + def check_bounds(self) -> None: + """ + Verify that every free parameter has finite bounds. + + Raises + ------ + ValueError + If any free parameter has an infinite lower or upper bound. + """ + labels = self._labels() + unbounded = unbounded_parameters(labels.parameters) + if not unbounded: + return + names = ', '.join(labels.label(parameter) for parameter in unbounded) + raise ValueError( + f'Bayesian sampling requires finite bounds on every free parameter, because the ' + f'bounds act as the prior. These parameters are unbounded: {names}. ' + f'Set their min and max, or call suggest_bounds() to propose values.' + ) + + ############# + # Sampling + ############# + + def sample( + self, + samples: int = 10000, + burn: int = 2000, + thin: int = 10, + population: int | None = None, + parameters: list[Parameter] | list[str] | None = None, + **sampler_options: dict[str, Any], + ) -> SamplingResults: + """ + Draw samples from the posterior distribution of the free parameters. + + Starts a fresh chain, replacing any existing one; use :meth:`extend` to continue one. + Parameter values are restored afterwards, so sampling never silently moves the model off + its fitted values; use :meth:`set_parameters_to_median` to adopt the posterior. + + Parameters + ---------- + samples : int, default=10000 + Number of raw samples to draw across all chains, before thinning. A guaranteed minimum + rather than an exact count. + burn : int, default=2000 + Burn-in generations to discard before collecting samples. + thin : int, default=10 + Thinning interval, which reduces autocorrelation between retained draws. + population : int | None, default=None + DREAM population scale factor: BUMPS runs ``ceil(population * n_parameters)`` chains. + parameters : list[Parameter] | list[str] | None, default=None + Restrict the chain to these parameters, given as Parameter objects or labels. All other + free parameters are held fixed for the run. Holding a parameter fixed is not the same + as marginalizing over it: the resulting intervals are conditional on those values and + will be too narrow if the parameters are correlated. The default samples everything. + **sampler_options : dict[str, Any] + Forwarded to the EasyScience Sampler, e.g. ``sampler_kwargs`` or ``progress_callback``. + + Returns + ------- + SamplingResults + The sampling results, also stored on :attr:`results`. + """ + return self._run( + parameters=parameters, + run=lambda sampler: sampler.sample( + samples=samples, burn=burn, thin=thin, population=population, **sampler_options + ), + ) + + def extend( + self, + additional_samples: int = 5000, + thin: int = 10, + parameters: list[Parameter] | list[str] | None = None, + **sampler_options: dict[str, Any], + ) -> SamplingResults: + """ + Continue the existing chain with additional samples. + + Parameters + ---------- + additional_samples : int, default=5000 + Number of additional samples to draw, in the same units as ``samples``. + thin : int, default=10 + Thinning interval for the retained draws. + parameters : list[Parameter] | list[str] | None, default=None + The same restriction as in :meth:`sample`. It must leave the chain the same width, + since BUMPS resumes from a stored chain whose columns are fixed. + **sampler_options : dict[str, Any] + Forwarded to the EasyScience Sampler. + + Returns + ------- + SamplingResults + The sampling results for the full extended chain. + + Raises + ------ + RuntimeError + If there is no chain to extend. + """ + if self._sampler is None: + raise RuntimeError('No chain to extend. Call sample() or load() first.') + return self._run( + parameters=parameters, + run=lambda sampler: sampler.extend( + additional_samples=additional_samples, thin=thin, **sampler_options + ), + reuse_sampler=True, + ) + + def _run( + self, + parameters: list[Parameter] | list[str] | None, + run: Callable[[Sampler], SamplingResults], + reuse_sampler: bool = False, + ) -> SamplingResults: + """ + Run a sampling operation with the surrounding guards in place. + + Checks the bounds, switches the minimizer to BUMPS, optionally holds parameters fixed, and + restores the parameter values, fixed flags and minimizer afterwards. + + Parameters + ---------- + parameters : list[Parameter] | list[str] | None + Parameters to restrict the chain to, or None for all free parameters. + run : Callable[[Sampler], SamplingResults] + The operation to perform on the prepared Sampler. + reuse_sampler : bool, default=False + Whether to reuse the cached Sampler, as an extension must. + + Returns + ------- + SamplingResults + The results of the run. + + Raises + ------ + IndexError + Re-raised untouched when it did not come from BUMPS, since that is a bug here rather + than a modelling problem. + RuntimeError + If the BUMPS sampler fails while removing outlier chains. + """ + held_fixed = self._resolve_parameters_to_hold_fixed(parameters) + _warn_about_held_parameters(self._labels(), held_fixed) + + with _FixedParameters(held_fixed): + self.check_bounds() + self._prepare() + + chain_parameters = self._chain_parameters() + saved_values = [(p, p.value) for p in chain_parameters] + + if reuse_sampler: + self._verify_chain_shape_unchanged(chain_parameters) + + fitter = self._analysis.fitter + original_minimizer = fitter.minimizer.enum + fitter.switch_minimizer(AvailableMinimizers.Bumps) + try: + results = run(self._get_or_build_sampler(reuse_sampler=reuse_sampler)) + except IndexError as error: + if not _raised_inside_bumps(error): + raise + raise RuntimeError( + 'The BUMPS sampler failed while removing outlier chains. This happens when ' + 'the chains scatter because two or more free parameters are degenerate, and ' + 'also on short chains, where BUMPS has too few generations to work with. ' + 'Check for degenerate parameters, raise samples, or switch the outlier ' + "removal off with sampler_kwargs={'outliers': 'none'}." + ) from error + finally: + fitter.switch_minimizer(original_minimizer) + for parameter, value in saved_values: + parameter.value = value + + # Labelled outside the block above, so a subset run records the labels a full run would. + # Inside it the other parameters are fixed, nothing looks ambiguous, and the sidecar would + # be written with unqualified names that no longer match on reload. + self._saved_labels = self._labels().name_map() + self._results = results + self._warn_about_bounds_occupancy(results) + return results + + def _get_or_build_sampler(self, reuse_sampler: bool) -> Sampler: + """ + Get the cached Sampler, rebuilding it if the data changed. + + Parameters + ---------- + reuse_sampler : bool + Whether to reuse the cached Sampler even if it is marked dirty. + + Returns + ------- + Sampler + The Sampler to run. + """ + if self._sampler is None or (self._sampler_is_dirty and not reuse_sampler): + x, y, weights = self._sampling_data() + self._sampler = Sampler(self._analysis.fitter, x, y, weights=weights) + self._sampler_is_dirty = False + return self._sampler + + def _verify_chain_shape_unchanged(self, chain_parameters: list[Parameter]) -> None: + """ + Check that an extension keeps the chain's columns. + + Parameters + ---------- + chain_parameters : list[Parameter] + The parameters that would form the chain for this run. + + Raises + ------ + ValueError + If the number of parameters differs from the existing chain's. + """ + if self._results is None: + return + existing = self._results.draws.shape[1] + if len(chain_parameters) != existing: + raise ValueError( + f'Cannot extend a chain of {existing} parameters with a run of ' + f'{len(chain_parameters)}. An extension continues the stored chain, whose columns ' + f'are fixed, so it needs the same parameters the chain was started with. Start a ' + f'fresh chain with sample() instead.' + ) + + def _resolve_parameters_to_hold_fixed( + self, + parameters: list[Parameter] | list[str] | None, + ) -> list[Parameter]: + """ + Work out which free parameters must be held fixed to honour a subset request. + + Parameters + ---------- + parameters : list[Parameter] | list[str] | None + The requested subset, as Parameter objects or labels, or None for everything. + + Returns + ------- + list[Parameter] + The free parameters that are not in the requested subset. + + Raises + ------ + TypeError + If parameters is not a list of Parameters or strings, or None. + ValueError + If a requested label matches no free parameter, or the subset is empty. + """ + if parameters is None: + return [] + if not isinstance(parameters, (list, tuple)): + raise TypeError('parameters must be a list of Parameters, a list of labels, or None.') + + labels = self._labels() + by_label = {labels.label(parameter): parameter for parameter in labels.parameters} + requested = [] + for entry in parameters: + if isinstance(entry, str): + if entry not in by_label: + raise ValueError( + f'No free parameter named {entry!r}. ' + f'Available: {", ".join(sorted(by_label))}.' + ) + requested.append(by_label[entry]) + elif hasattr(entry, 'unique_name'): + requested.append(entry) + else: + raise TypeError('parameters must contain Parameter objects or labels (strings).') + + wanted = {parameter.unique_name for parameter in requested} + if not wanted: + raise ValueError('parameters must name at least one parameter to sample.') + return [p for p in labels.parameters if p.unique_name not in wanted] + + def _warn_about_bounds_occupancy(self, results: SamplingResults) -> None: + """ + Warn when the posterior has piled up against a bound. + + Parameters + ---------- + results : SamplingResults + The sampling results to inspect. + """ + piled_up = parameters_at_bounds(results.draws, self._resolve(results)) + if not piled_up: + return + details = ', '.join( + f'{name} ({fraction:.0%} of draws)' for name, fraction in piled_up.items() + ) + warnings.warn( + ( + f'The posterior is piled up against the bounds for: {details}. ' + f'The bounds, rather than the data, are setting these credible intervals. ' + f'Widen the bounds, or check whether these parameters are degenerate with others.' + ), + UserWarning, + stacklevel=4, + ) + + ############# + # Results + ############# + + def summary(self) -> PosteriorSummary: + """ + Summarize the marginal posterior of each sampled parameter. + + Reports the median and the 68% credible interval under the parameter's own label and unit. + + Returns + ------- + PosteriorSummary + One entry per sampled parameter. + """ + results = self._require_results() + labels = self._labels() + return summarize_draws( + draws=results.draws, + labels=labels.display_names(results.param_names, self._saved_labels), + parameters_by_column=self._resolve(results), + ) + + def set_parameters_to_median(self) -> list[Parameter]: + """ + Set every sampled parameter to the median of its marginal posterior. + + The vector of marginal medians is not in general the highest-posterior point, and for + strongly correlated parameters need not even be a good fit. + + Returns + ------- + list[Parameter] + The parameters that were changed. + """ + results = self._require_results() + changed = [] + for column, parameter in enumerate(self._resolve(results)): + if parameter is None: + continue + parameter.value = float(np.median(results.draws[:, column])) + changed.append(parameter) + return changed + + ############# + # Persistence + ############# + + def save(self, path: str | os.PathLike) -> None: + """ + Save the MCMC chain to disk. + + Writes the BUMPS chain files plus a sidecar recording the column labels, because the unique + names BUMPS stores are per-session and cannot be matched up again on their own. + + Parameters + ---------- + path : str | os.PathLike + Path prefix for the chain files. + + Raises + ------ + RuntimeError + If there is no chain to save. + """ + if self._sampler is None: + raise RuntimeError('No chain to save. Call sample() first.') + self._sampler.save(path) + Path(f'{path}{_LABEL_MAP_SUFFIX}').write_text( + json.dumps(self._saved_labels, indent=2), encoding='utf-8' + ) + + def load(self, path: str | os.PathLike, skip: int = 0) -> SamplingResults: + """ + Load a previously saved MCMC chain. + + The loaded chain can be summarized, plotted, or continued with :meth:`extend`. + + Parameters + ---------- + path : str | os.PathLike + The path prefix the chain was saved under. + skip : int, default=0 + Number of initial samples to skip when reading the chain. + + Returns + ------- + SamplingResults + The loaded results, also stored on :attr:`results`. + """ + self._prepare() + sidecar = Path(f'{path}{_LABEL_MAP_SUFFIX}') + if sidecar.is_file(): + self._saved_labels = json.loads(sidecar.read_text(encoding='utf-8')) + else: + self._saved_labels = {} + warnings.warn( + ( + f'No parameter-name sidecar found at {sidecar}. The chain will be reported ' + f'under the internal names it was saved with, because those cannot be matched ' + f'to this Analysis.' + ), + UserWarning, + stacklevel=2, + ) + + fitter = self._analysis.fitter + original_minimizer = fitter.minimizer.enum + fitter.switch_minimizer(AvailableMinimizers.Bumps) + try: + self._results = self._get_or_build_sampler(reuse_sampler=False).load_state( + path, skip=skip + ) + finally: + fitter.switch_minimizer(original_minimizer) + return self._results + + ############# + # Figures, each one a call into posterior_plotting + ############# + + def plot_trace(self, **kwargs: dict[str, Any]) -> Figure: + """ + Plot the chain trace of each sampled parameter. + + Parameters + ---------- + **kwargs : dict[str, Any] + Forwarded to :func:`easydynamics.utils.posterior_plotting.plot_trace`. + + Returns + ------- + Figure + The matplotlib Figure. + """ + from easydynamics.utils.posterior_plotting import plot_trace + + results = self._require_results() + return plot_trace( + draws=results.draws, + names=self._display_names(results), + logp=results.logp, + units=self._units(results), + title=self._analysis.display_name, + **kwargs, + ) + + def plot_corner(self, **kwargs: dict[str, Any]) -> Figure: + """ + Plot the marginal and pairwise posterior distributions. + + Parameters + ---------- + **kwargs : dict[str, Any] + Forwarded to :func:`easydynamics.utils.posterior_plotting.plot_corner`. + + Returns + ------- + Figure + The matplotlib Figure. + """ + from easydynamics.utils.posterior_plotting import plot_corner + + results = self._require_results() + return plot_corner( + draws=results.draws, + names=self._display_names(results), + units=self._units(results), + title=self._analysis.display_name, + **kwargs, + ) + + def plot_posterior_predictive( + self, + n_draws: int = 200, + credible_interval: float = 68.0, + **kwargs: dict[str, Any], + ) -> Figure: + """ + Plot the data against the credible band implied by the posterior. + + Parameters + ---------- + n_draws : int, default=200 + How many posterior draws to evaluate the model for. Each costs a full model evaluation. + credible_interval : float, default=68.0 + Width of the credible band, as a percentage. + **kwargs : dict[str, Any] + Forwarded to :func:`easydynamics.utils.posterior_plotting.plot_posterior_predictive`. + + Returns + ------- + Figure + The matplotlib Figure. + + Raises + ------ + NotImplementedError + If this Analysis binds a list of datasets rather than a single one. + ValueError + If n_draws is not a positive integer. + """ + from easydynamics.utils.posterior_plotting import plot_posterior_predictive + + if not isinstance(n_draws, int) or isinstance(n_draws, bool) or n_draws < 1: + raise ValueError(f'n_draws must be a positive integer. Got {n_draws}.') + + self._require_results() + x, y, weights = self._sampling_data() + if isinstance(x, (list, tuple)): + raise NotImplementedError( + 'plot_posterior_predictive supports a single dataset only. Plot each dataset ' + 'from its own Analysis1d instead.' + ) + + energy = getattr(self._analysis, 'energy', None) + sample_model = getattr(self._analysis, 'sample_model', None) + y_unit = None if sample_model is None else getattr(sample_model, 'y_unit', None) + kwargs.setdefault('xlabel', None if energy is None else f'Energy ({energy.unit})') + kwargs.setdefault('ylabel', 'Intensity' if y_unit is None else f'Intensity ({y_unit})') + + return plot_posterior_predictive( + x=np.asarray(x), + y=np.asarray(y), + predictions=self.predictions(n_draws), + y_err=None if weights is None else 1.0 / np.asarray(weights), + title=self._analysis.display_name, + credible_interval=credible_interval, + **kwargs, + ) + + def predictions(self, n_draws: int = 200) -> np.ndarray: + """ + Evaluate the model once per posterior draw, restoring the parameters afterwards. + + Parameters + ---------- + n_draws : int, default=200 + How many draws to evaluate, taken evenly across the chain. + + Returns + ------- + np.ndarray + Model evaluations, shape ``(n_selected, len(x))``. + """ + results = self._require_results() + self._prepare() + + x, _, _ = self._sampling_data() + columns = [ + (parameter, column) + for column, parameter in enumerate(self._resolve(results)) + if parameter is not None + ] + saved_values = [(parameter, parameter.value) for parameter, _ in columns] + + total = results.draws.shape[0] + indices = np.unique(np.linspace(0, total - 1, min(n_draws, total)).astype(int)) + + fit_function = self._analysis.fitter.fit_function + predictions = [] + try: + for index in indices: + for parameter, column in columns: + parameter.value = float(results.draws[index, column]) + predictions.append(np.asarray(fit_function(x))) + finally: + for parameter, value in saved_values: + parameter.value = value + return np.vstack(predictions) + + ############# + # Talking to the Analysis + ############# + + def _labels(self) -> ParameterLabels: + """ + Get the label helper for the current free parameters. + + Returns + ------- + ParameterLabels + Built fresh, because which parameters are free can change between calls. + """ + return self._parameter_labels() + + def _prepare(self) -> None: + """Bring any cached computation on the Analysis up to date before a run.""" + if self._prepare_hook is not None: + self._prepare_hook() + + def _resolve(self, results: SamplingResults) -> list[Parameter | None]: + """ + Match each column of a chain to a parameter. + + Parameters + ---------- + results : SamplingResults + The results whose columns should be matched. + + Returns + ------- + list[Parameter | None] + The parameter for each column, or None where none could be matched. + """ + return self._labels().resolve(results.param_names, self._saved_labels) + + def _display_names(self, results: SamplingResults) -> list[str]: + """ + Get a readable label for each column of a chain. + + Parameters + ---------- + results : SamplingResults + The results whose columns should be named. + + Returns + ------- + list[str] + One label per column. + """ + return self._labels().display_names(results.param_names, self._saved_labels) + + def _units(self, results: SamplingResults) -> list[str]: + """ + Get the unit of each column of a chain. + + Parameters + ---------- + results : SamplingResults + The results whose columns should be described. + + Returns + ------- + list[str] + One unit per column. + """ + return self._labels().units(results.param_names, self._saved_labels) + + def _require_results(self) -> SamplingResults: + """ + Get the stored results, raising if there are none. + + Returns + ------- + SamplingResults + The most recent sampling results. + + Raises + ------ + RuntimeError + If no sampling has been run yet. + """ + if self._results is None: + raise RuntimeError('No posterior samples yet. Call sample() or load() first.') + return self._results + + +def _warn_about_held_parameters(labels: object, held_fixed: list[Parameter]) -> None: + """ + Warn that holding parameters fixed makes the credible intervals conditional. + + Parameters + ---------- + labels : object + The ParameterLabels used to name them. + held_fixed : list[Parameter] + The parameters being held fixed for the run. + """ + if not held_fixed: + return + names = ', '.join(labels.label(parameter) for parameter in held_fixed) + warnings.warn( + ( + f'Holding these parameters fixed while sampling: {names}. ' + f'Fixing a parameter is not the same as marginalizing over it, so the resulting ' + f'credible intervals are conditional on these values and will be too narrow if the ' + f'parameters are correlated.' + ), + UserWarning, + stacklevel=4, + ) + + +def _raised_inside_bumps(error: BaseException) -> bool: + """ + Check whether an exception came from inside BUMPS. + + Used so only BUMPS' own failures are relabelled, and a bug in this package is not reported as a + modelling problem. + + Parameters + ---------- + error : BaseException + The exception to inspect. + + Returns + ------- + bool + True when any frame of the traceback lies in the bumps package. + """ + traceback = error.__traceback__ + while traceback is not None: + module = traceback.tb_frame.f_globals.get('__name__', '') + if module == 'bumps' or module.startswith('bumps.'): + return True + traceback = traceback.tb_next + return False + + +class _FixedParameters: + """Context manager that temporarily fixes parameters and restores their flags on exit.""" + + def __init__(self, parameters: list[Parameter]) -> None: + self._parameters = list(parameters) + self._saved: list[tuple[Parameter, bool]] = [] + + def __enter__(self) -> None: + """Fix the parameters, remembering their previous state.""" + self._saved = [(parameter, parameter.fixed) for parameter in self._parameters] + for parameter in self._parameters: + parameter.fixed = True + + def __exit__(self, *_exc_info: object) -> None: + """ + Restore the previous fixed state of every parameter. + + Parameters + ---------- + *_exc_info : object + Exception information, ignored. + """ + for parameter, was_fixed in self._saved: + parameter.fixed = was_fixed diff --git a/src/easydynamics/utils/posterior_plotting.py b/src/easydynamics/utils/posterior_plotting.py index bfb1a1320..8e75cf425 100644 --- a/src/easydynamics/utils/posterior_plotting.py +++ b/src/easydynamics/utils/posterior_plotting.py @@ -14,6 +14,7 @@ import matplotlib.pyplot as plt import numpy as np +from matplotlib.ticker import MaxNLocator if TYPE_CHECKING: from matplotlib.figure import Figure @@ -23,6 +24,7 @@ def plot_trace( draws: np.ndarray, names: list[str], logp: np.ndarray | None = None, + units: list[str] | None = None, title: str | None = None, figsize: tuple[float, float] | None = None, ) -> Figure: @@ -44,6 +46,9 @@ def plot_trace( One label per column of ``draws``. logp : np.ndarray | None, default=None Log-posterior values, plotted in an extra panel when given. + units : list[str] | None, default=None + Unit of each column, appended to its label. Entries that are empty or dimensionless are + skipped, since a bare "dimensionless" only adds clutter. title : str | None, default=None Figure title. figsize : tuple[float, float] | None, default=None @@ -66,7 +71,7 @@ def plot_trace( for axis, column, name in zip(axes, range(draws.shape[1]), names, strict=False): axis.plot(draws[:, column], lw=0.5) - axis.set_ylabel(name, fontsize=8) + axis.set_ylabel(_with_unit(name, units, column), fontsize=8) axis.set_xlim(0, len(draws) - 1) if logp is not None: @@ -83,6 +88,7 @@ def plot_trace( def plot_corner( draws: np.ndarray, names: list[str], + units: list[str] | None = None, title: str | None = None, bins: int = 40, figsize: tuple[float, float] | None = None, @@ -103,6 +109,9 @@ def plot_corner( Posterior draws, shape ``(n_draws, n_parameters)``. names : list[str] One label per column of ``draws``. + units : list[str] | None, default=None + Unit of each column, appended to its label. Entries that are empty or dimensionless are + skipped, since a bare "dimensionless" only adds clutter. title : str | None, default=None Figure title. bins : int, default=40 @@ -143,7 +152,25 @@ def plot_corner( axis.set_ylabel(names[row], fontsize=8) else: axis.set_yticklabels([]) + if row == 0 and col == 0: + # The top-left panel is a histogram, so its vertical axis counts draws rather than + # carrying a parameter. Say so, instead of leaving it blank as if by omission. + axis.set_ylabel('counts', fontsize=8) axis.tick_params(labelsize=7) + axis.xaxis.set_major_locator(MaxNLocator(nbins=4)) + if row != col: + axis.yaxis.set_major_locator(MaxNLocator(nbins=4)) + + # Matplotlib parks the shared exponent ("1e-8") at the end of the axis, where it lands on top + # of the axis label. Fold it into the label instead. + fig.canvas.draw() + for row in range(n): + for col in range(row + 1): + axis = axes[row, col] + if row == n - 1: + _absorb_offset(axis.xaxis, axis.set_xlabel, names[col], units, col) + if col == 0 and row != 0: + _absorb_offset(axis.yaxis, axis.set_ylabel, names[row], units, row) if title is not None: fig.suptitle(title) @@ -158,6 +185,8 @@ def plot_posterior_predictive( y_err: np.ndarray | None = None, title: str | None = None, credible_interval: float = 68.0, + xlabel: str | None = None, + ylabel: str | None = None, figsize: tuple[float, float] = (8.0, 5.0), ) -> Figure: """ @@ -181,6 +210,10 @@ def plot_posterior_predictive( Figure title. credible_interval : float, default=68.0 Width of the credible band, as a percentage. + xlabel : str | None, default=None + Label for the independent axis. + ylabel : str | None, default=None + Label for the dependent axis. figsize : tuple[float, float], default=(8.0, 5.0) Figure size in inches. @@ -224,6 +257,10 @@ def plot_posterior_predictive( label=f'{credible_interval:.0f}% credible band', ) axis.plot(x, median, '-', color='C3', label='Posterior median') + if xlabel is not None: + axis.set_xlabel(xlabel) + if ylabel is not None: + axis.set_ylabel(ylabel) axis.legend() if title is not None: axis.set_title(title) @@ -231,6 +268,85 @@ def plot_posterior_predictive( return fig +def _unit_for(units: list[str] | None, column: int) -> str: + """ + Get the unit to show for a column, if it is worth showing. + + Parameters + ---------- + units : list[str] | None + The units, one per column, or None. + column : int + The column to look up. + + Returns + ------- + str + The unit, or an empty string when there is none worth printing. + """ + if units is None or column >= len(units): + return '' + unit = (units[column] or '').strip() + return '' if unit.lower() in ('', 'dimensionless', 'none') else unit + + +def _with_unit(name: str, units: list[str] | None, column: int) -> str: + """ + Append a column's unit to its label. + + Parameters + ---------- + name : str + The label to extend. + units : list[str] | None + The units, one per column, or None. + column : int + The column the label belongs to. + + Returns + ------- + str + The label, with the unit in parentheses when there is one. + """ + unit = _unit_for(units, column) + return f'{name} ({unit})' if unit else name + + +def _absorb_offset( + axis_object: object, + set_label: object, + name: str, + units: list[str] | None = None, + column: int = 0, +) -> None: + """ + Move an axis' shared exponent into its label, so the two stop overlapping. + + The exponent and the unit share one set of parentheses, since two adjacent parentheticals read + badly: ``D (1e-8 m^2/s)`` rather than ``D (1e-8) (m^2/s)``. + + Parameters + ---------- + axis_object : object + The matplotlib ``XAxis`` or ``YAxis`` carrying the offset text. + set_label : object + The corresponding ``set_xlabel`` or ``set_ylabel`` callable. + name : str + The label the axis should carry, before the exponent and unit are appended. + units : list[str] | None, default=None + The units, one per column, or None. + column : int, default=0 + The column the axis belongs to. + """ + offset_text = axis_object.get_offset_text() + offset = offset_text.get_text() + unit = _unit_for(units, column) + suffix = ' '.join(part for part in (offset, unit) if part) + set_label(f'{name} ({suffix})' if suffix else name, fontsize=8) + if offset: + offset_text.set_visible(False) + + def _verify_draws(draws: np.ndarray, names: list[str]) -> None: """ Verify that a draws array is two-dimensional and matches its labels. diff --git a/tests/integration/fitting/test_bayesian_sampling.py b/tests/integration/fitting/test_bayesian_sampling.py index 704a37ea2..b6adb3bd2 100644 --- a/tests/integration/fitting/test_bayesian_sampling.py +++ b/tests/integration/fitting/test_bayesian_sampling.py @@ -76,17 +76,17 @@ def build_analysis(): def sampled_analysis(): analysis = build_analysis() analysis.fit() - analysis.suggest_bounds().apply() + analysis.bayesian.suggest_bounds().apply() with warnings.catch_warnings(): warnings.simplefilter('ignore') - analysis.sample_posterior(**SAMPLE_KWARGS) + analysis.bayesian.sample(**SAMPLE_KWARGS) return analysis class TestRealChain: def test_chain_has_one_column_per_free_parameter(self, sampled_analysis): # EXPECT - results = sampled_analysis.posterior_result + results = sampled_analysis.bayesian.results assert results.draws.shape[1] == len(sampled_analysis.get_free_parameters()) assert results.draws.shape[0] > 0 @@ -96,7 +96,7 @@ def test_chain_has_one_column_per_free_parameter(self, sampled_analysis): ) def test_posterior_recovers_the_true_parameters(self, sampled_analysis, name, truth): # WHEN - entry = sampled_analysis.posterior_summary()[name] + entry = sampled_analysis.bayesian.summary()[name] # EXPECT the truth sits within a few posterior standard deviations of the median. A 68% # interval is deliberately not used: it excludes the truth about a third of the time for @@ -106,7 +106,7 @@ def test_posterior_recovers_the_true_parameters(self, sampled_analysis, name, tr def test_summary_is_reported_under_parameter_names_and_units(self, sampled_analysis): # WHEN - summary = sampled_analysis.posterior_summary() + summary = sampled_analysis.bayesian.summary() # EXPECT assert {entry.name for entry in summary} == { @@ -118,12 +118,12 @@ def test_sampling_leaves_the_fitted_values_untouched(self): # WHEN analysis = build_analysis() analysis.fit() - analysis.suggest_bounds().apply() + analysis.bayesian.suggest_bounds().apply() before = [float(p.value) for p in analysis.get_free_parameters()] with warnings.catch_warnings(): warnings.simplefilter('ignore') - analysis.sample_posterior(**SAMPLE_KWARGS) + analysis.bayesian.sample(**SAMPLE_KWARGS) # EXPECT after = [float(p.value) for p in analysis.get_free_parameters()] @@ -131,11 +131,11 @@ def test_sampling_leaves_the_fitted_values_untouched(self): def test_extend_grows_the_chain(self, sampled_analysis): # WHEN - before = int(sampled_analysis.posterior_result.state.Ngen) + before = int(sampled_analysis.bayesian.results.state.Ngen) with warnings.catch_warnings(): warnings.simplefilter('ignore') - extended = sampled_analysis.extend_sampling( + extended = sampled_analysis.bayesian.extend( additional_samples=500, thin=2, sampler_kwargs={'trim': False} ) @@ -145,17 +145,17 @@ def test_extend_grows_the_chain(self, sampled_analysis): def test_save_and_load_round_trip_keeps_parameter_identity(self, sampled_analysis, tmp_path): # WHEN prefix = str(tmp_path / 'chain') - sampled_analysis.save_chain(prefix) + sampled_analysis.bayesian.save(prefix) fresh = build_analysis() fresh.fit() - fresh.suggest_bounds().apply() + fresh.bayesian.suggest_bounds().apply() with warnings.catch_warnings(): warnings.simplefilter('ignore') - fresh.load_chain(prefix) + fresh.bayesian.load(prefix) # EXPECT the reloaded chain is reported under real names, not internal unique names - summary = fresh.posterior_summary() + summary = fresh.bayesian.summary() assert {entry.name for entry in summary} == {p.name for p in fresh.get_free_parameters()} assert all(np.isfinite(entry.value) for entry in summary) @@ -163,24 +163,24 @@ def test_subset_sampling_produces_a_single_column(self): # WHEN analysis = build_analysis() analysis.fit() - analysis.suggest_bounds().apply() + analysis.bayesian.suggest_bounds().apply() with warnings.catch_warnings(): warnings.simplefilter('ignore') - results = analysis.sample_posterior(parameters=['Gaussian width'], **SAMPLE_KWARGS) + results = analysis.bayesian.sample(parameters=['Gaussian width'], **SAMPLE_KWARGS) # EXPECT assert results.draws.shape[1] == 1 - assert analysis.posterior_summary().entries[0].name == 'Gaussian width' + assert analysis.bayesian.summary().entries[0].name == 'Gaussian width' def test_plots_render(self, sampled_analysis): # WHEN import matplotlib.pyplot as plt n_parameters = len(sampled_analysis.get_free_parameters()) - trace = sampled_analysis.plot_trace() - corner = sampled_analysis.plot_corner() - predictive = sampled_analysis.plot_posterior_predictive(n_draws=20) + trace = sampled_analysis.bayesian.plot_trace() + corner = sampled_analysis.bayesian.plot_corner() + predictive = sampled_analysis.bayesian.plot_posterior_predictive(n_draws=20) # EXPECT assert len(trace.axes) == n_parameters + 1 @@ -192,13 +192,13 @@ def test_posterior_median_is_close_to_the_least_squares_fit(self): # WHEN analysis = build_analysis() analysis.fit() - analysis.suggest_bounds().apply() + analysis.bayesian.suggest_bounds().apply() fitted = {p.name: float(p.value) for p in analysis.get_free_parameters()} with warnings.catch_warnings(): warnings.simplefilter('ignore') - analysis.sample_posterior(**SAMPLE_KWARGS) - summary = analysis.posterior_summary() + analysis.bayesian.sample(**SAMPLE_KWARGS) + summary = analysis.bayesian.summary() # EXPECT the two agree within the posterior's own uncertainty, since with flat priors the # maximum-likelihood point sits inside the bulk of the posterior diff --git a/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py b/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py index b29a2b54a..345bc2724 100644 --- a/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py +++ b/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py @@ -18,7 +18,7 @@ from easydynamics.sample_model import SampleModel from easydynamics.sample_model.components.gaussian import Gaussian -SAMPLER_PATH = 'easydynamics.analysis.bayesian_sampling.Sampler' +SAMPLER_PATH = 'easydynamics.analysis.posterior_sampling.Sampler' def make_analysis(): @@ -113,23 +113,23 @@ class TestBoundsPreflight: def test_sampling_refuses_unbounded_parameters(self, analysis): # EXPECT with pytest.raises(ValueError, match='finite bounds'): - analysis.sample_posterior(samples=10) + analysis.bayesian.sample(samples=10) def test_error_names_the_offending_parameters(self, analysis): # EXPECT with pytest.raises(ValueError, match='Gaussian area'): - analysis.check_bounds_for_sampling() + analysis.bayesian.check_bounds() def test_bounded_parameters_pass(self, analysis): # WHEN bound_all(analysis) # EXPECT: does not raise - analysis.check_bounds_for_sampling() + analysis.bayesian.check_bounds() def test_suggest_bounds_covers_the_free_parameters(self, analysis): # WHEN - suggestions = analysis.suggest_bounds() + suggestions = analysis.bayesian.suggest_bounds() # EXPECT assert len(suggestions) == len(analysis.get_free_parameters()) @@ -150,7 +150,7 @@ def mutate_then_return(**_kwargs): return fake_results(analysis) sampler_class.return_value.sample.side_effect = mutate_then_return - analysis.sample_posterior(samples=10, burn=1, thin=1) + analysis.bayesian.sample(samples=10, burn=1, thin=1) # EXPECT after = [(p.unique_name, p.value) for p in analysis.get_free_parameters()] @@ -167,7 +167,7 @@ def test_switches_to_bumps_for_the_run(self, analysis): seen.append(analysis.fitter.minimizer.enum), fake_results(analysis), )[1] - analysis.sample_posterior(samples=10) + analysis.bayesian.sample(samples=10) # EXPECT assert seen == [AvailableMinimizers.Bumps] @@ -179,7 +179,7 @@ def test_restores_the_minimizer_even_when_sampling_raises(self, analysis): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.side_effect = RuntimeError('boom') with pytest.raises(RuntimeError, match='boom'): - analysis.sample_posterior(samples=10) + analysis.bayesian.sample(samples=10) # EXPECT assert analysis.fitter.minimizer.enum == AvailableMinimizers.LMFit_leastsq @@ -190,7 +190,7 @@ def test_forwards_sampling_arguments(self, analysis): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) - analysis.sample_posterior(samples=123, burn=7, thin=3, population=5) + analysis.bayesian.sample(samples=123, burn=7, thin=3, population=5) # EXPECT kwargs = sampler_class.return_value.sample.call_args.kwargs @@ -206,11 +206,11 @@ def test_stores_the_result(self, analysis): with patch(SAMPLER_PATH) as sampler_class: expected = fake_results(analysis) sampler_class.return_value.sample.return_value = expected - returned = analysis.sample_posterior(samples=10) + returned = analysis.bayesian.sample(samples=10) # EXPECT assert returned is expected - assert analysis.posterior_result is expected + assert analysis.bayesian.results is expected def test_warns_when_the_posterior_piles_up_against_a_bound(self, analysis): # WHEN a parameter's draws span its whole allowed range @@ -224,7 +224,7 @@ def test_warns_when_the_posterior_piles_up_against_a_bound(self, analysis): # EXPECT with pytest.warns(UserWarning, match='piled up'): - analysis.sample_posterior(samples=10) + analysis.bayesian.sample(samples=10) def test_does_not_warn_when_the_posterior_is_well_inside(self, analysis): # WHEN @@ -235,7 +235,7 @@ def test_does_not_warn_when_the_posterior_is_well_inside(self, analysis): # EXPECT with warnings_as_errors(): - analysis.sample_posterior(samples=10) + analysis.bayesian.sample(samples=10) class TestParameterSubset: @@ -253,7 +253,7 @@ def record(**_kwargs): sampler_class.return_value.sample.side_effect = record with pytest.warns(UserWarning, match='Holding these parameters fixed'): - analysis.sample_posterior(samples=10, parameters=[target.name]) + analysis.bayesian.sample(samples=10, parameters=[target.name]) # EXPECT assert seen['free'] == [target.unique_name] @@ -267,7 +267,7 @@ def test_restores_the_fixed_flags_afterwards(self, analysis): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) with pytest.warns(UserWarning): - analysis.sample_posterior(samples=10, parameters=[target]) + analysis.bayesian.sample(samples=10, parameters=[target]) # EXPECT assert [(p.unique_name, p.fixed) for p in analysis.get_all_parameters()] == before @@ -278,17 +278,17 @@ def test_unknown_parameter_name_raises(self, analysis): # EXPECT with pytest.raises(ValueError, match='No free parameter named'): - analysis.sample_posterior(samples=10, parameters=['not a parameter']) + analysis.bayesian.sample(samples=10, parameters=['not a parameter']) def test_non_list_parameters_raises(self, analysis): # EXPECT with pytest.raises(TypeError, match='must be a list'): - analysis.sample_posterior(samples=10, parameters='Gaussian area') + analysis.bayesian.sample(samples=10, parameters='Gaussian area') def test_empty_parameter_list_raises(self, analysis): # EXPECT with pytest.raises(ValueError, match='at least one parameter'): - analysis.sample_posterior(samples=10, parameters=[]) + analysis.bayesian.sample(samples=10, parameters=[]) class TestSamplerCaching: @@ -298,8 +298,8 @@ def test_sampler_is_reused_between_runs(self, analysis): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) - analysis.sample_posterior(samples=10) - analysis.sample_posterior(samples=10) + analysis.bayesian.sample(samples=10) + analysis.bayesian.sample(samples=10) # EXPECT the data is bound once, not per run assert sampler_class.call_count == 1 @@ -310,9 +310,9 @@ def test_changing_the_q_index_rebuilds_the_sampler(self, analysis): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) - analysis.sample_posterior(samples=10) + analysis.bayesian.sample(samples=10) analysis.Q_index = 0 - analysis.sample_posterior(samples=10) + analysis.bayesian.sample(samples=10) # EXPECT the Sampler binds its data at construction, so it must be rebuilt assert sampler_class.call_count == 2 @@ -323,10 +323,10 @@ def test_binds_the_same_data_the_fit_uses(self, analysis): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) - analysis.sample_posterior(samples=10) + analysis.bayesian.sample(samples=10) # EXPECT - expected_x, expected_y, expected_w = analysis._get_sampling_data() + expected_x, expected_y, expected_w = analysis._sampling_data() args, kwargs = sampler_class.call_args assert np.array_equal(args[1], expected_x) assert np.array_equal(args[2], expected_y) @@ -337,7 +337,7 @@ class TestExtendAndPersistence: def test_extend_without_a_chain_raises(self, analysis): # EXPECT with pytest.raises(RuntimeError, match='No chain to extend'): - analysis.extend_sampling() + analysis.bayesian.extend() def test_extend_delegates_to_the_sampler(self, analysis): # WHEN @@ -346,8 +346,8 @@ def test_extend_delegates_to_the_sampler(self, analysis): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) sampler_class.return_value.extend.side_effect = lambda **_k: fake_results(analysis) - analysis.sample_posterior(samples=10) - analysis.extend_sampling(additional_samples=42, thin=2) + analysis.bayesian.sample(samples=10) + analysis.bayesian.extend(additional_samples=42, thin=2) # EXPECT kwargs = sampler_class.return_value.extend.call_args.kwargs @@ -357,7 +357,7 @@ def test_extend_delegates_to_the_sampler(self, analysis): def test_save_without_a_chain_raises(self, analysis): # EXPECT with pytest.raises(RuntimeError, match='No chain to save'): - analysis.save_chain('somewhere') + analysis.bayesian.save('somewhere') def test_save_writes_the_parameter_name_sidecar(self, analysis, tmp_path): # WHEN @@ -366,8 +366,8 @@ def test_save_writes_the_parameter_name_sidecar(self, analysis, tmp_path): bound_all(analysis) with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) - analysis.sample_posterior(samples=10) - analysis.save_chain(str(tmp_path / 'chain')) + analysis.bayesian.sample(samples=10) + analysis.bayesian.save(str(tmp_path / 'chain')) # EXPECT the unique names are recorded against the stable parameter names sidecar = tmp_path / 'chain.parameter-names.json' @@ -384,14 +384,14 @@ def test_load_without_a_sidecar_warns(self, analysis, tmp_path): # EXPECT with pytest.warns(UserWarning, match='No parameter-name sidecar'): - analysis.load_chain(str(tmp_path / 'missing')) + analysis.bayesian.load(str(tmp_path / 'missing')) class TestResults: def test_summary_without_sampling_raises(self, analysis): # EXPECT with pytest.raises(RuntimeError, match='No posterior samples yet'): - analysis.posterior_summary() + analysis.bayesian.summary() def test_summary_uses_parameter_names_and_units(self, analysis): # WHEN @@ -399,10 +399,10 @@ def test_summary_uses_parameter_names_and_units(self, analysis): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) - analysis.sample_posterior(samples=10) + analysis.bayesian.sample(samples=10) # EXPECT - summary = analysis.posterior_summary() + summary = analysis.bayesian.summary() names = {entry.name for entry in summary} assert names == {p.name for p in analysis.get_free_parameters()} assert all(entry.unit == 'meV' for entry in summary) @@ -416,9 +416,9 @@ def test_set_parameters_to_posterior_median(self, analysis): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.return_value = fake_results(analysis, values=draws) expected = [float(p.value) + 2.0 for p in parameters] - analysis.sample_posterior(samples=10) + analysis.bayesian.sample(samples=10) - changed = analysis.set_parameters_to_posterior_median() + changed = analysis.bayesian.set_parameters_to_median() # EXPECT assert len(changed) == len(parameters) @@ -427,7 +427,7 @@ def test_set_parameters_to_posterior_median(self, analysis): def test_median_without_sampling_raises(self, analysis): # EXPECT with pytest.raises(RuntimeError, match='No posterior samples yet'): - analysis.set_parameters_to_posterior_median() + analysis.bayesian.set_parameters_to_median() class TestPlots: @@ -437,11 +437,11 @@ def test_predictive_rejects_a_bad_draw_count(self, analysis): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) - analysis.sample_posterior(samples=10) + analysis.bayesian.sample(samples=10) # EXPECT with pytest.raises(ValueError, match='positive integer'): - analysis.plot_posterior_predictive(n_draws=0) + analysis.bayesian.plot_posterior_predictive(n_draws=0) def test_predictive_restores_parameter_values(self, analysis): # WHEN @@ -451,10 +451,10 @@ def test_predictive_restores_parameter_values(self, analysis): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.return_value = fake_results(analysis, values=draws) - analysis.sample_posterior(samples=10) + analysis.bayesian.sample(samples=10) before = [float(p.value) for p in parameters] - analysis.plot_posterior_predictive(n_draws=5) + analysis.bayesian.plot_posterior_predictive(n_draws=5) # EXPECT assert [float(p.value) for p in parameters] == pytest.approx(before) @@ -462,9 +462,9 @@ def test_predictive_restores_parameter_values(self, analysis): def test_plots_without_sampling_raise(self, analysis): # EXPECT with pytest.raises(RuntimeError): - analysis.plot_trace() + analysis.bayesian.plot_trace() with pytest.raises(RuntimeError): - analysis.plot_corner() + analysis.bayesian.plot_corner() class warnings_as_errors: diff --git a/tests/unit/easydynamics/analysis/test_posterior.py b/tests/unit/easydynamics/analysis/test_posterior.py index 88dc29bf6..659412774 100644 --- a/tests/unit/easydynamics/analysis/test_posterior.py +++ b/tests/unit/easydynamics/analysis/test_posterior.py @@ -251,7 +251,7 @@ def test_reports_parameter_names_units_and_percentiles(self): draws = np.linspace(0.0, 100.0, 101).reshape(-1, 1) # THEN - summary = summarize_draws(draws, ['Parameter_0'], [parameter]) + summary = summarize_draws(draws, ['Gaussian width'], [parameter]) # EXPECT entry = summary['Gaussian width'] @@ -263,6 +263,17 @@ def test_reports_parameter_names_units_and_percentiles(self): assert entry.plus == pytest.approx(34.0) assert entry.value == pytest.approx(1.5) + def test_labels_are_reported_verbatim(self): + # WHEN a caller supplies a qualified label, as a multi-Q analysis does + parameter = make_parameter(name='Gaussian width') + + # THEN + summary = summarize_draws(np.zeros((5, 1)), ['Gaussian width (Q_index=2)'], [parameter]) + + # EXPECT + assert summary.entries[0].name == 'Gaussian width (Q_index=2)' + assert summary.entries[0].unit == 'meV' + def test_unmatched_column_falls_back_to_the_supplied_name(self): # WHEN draws = np.zeros((10, 1)) @@ -289,7 +300,7 @@ def test_repr_contains_the_parameter_name(self): parameter = make_parameter(name='Gaussian area') # THEN - text = repr(summarize_draws(np.zeros((5, 1)), ['x'], [parameter])) + text = repr(summarize_draws(np.zeros((5, 1)), ['Gaussian area'], [parameter])) # EXPECT assert 'Gaussian area' in text diff --git a/tests/unit/easydynamics/analysis/test_posterior_labels.py b/tests/unit/easydynamics/analysis/test_posterior_labels.py new file mode 100644 index 000000000..b4b08f86c --- /dev/null +++ b/tests/unit/easydynamics/analysis/test_posterior_labels.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +import numpy as np +from easyscience.variable import Parameter + +from easydynamics.analysis.posterior_labels import ParameterLabels + + +def make_parameter(name, unit='meV'): + return Parameter(name=name, value=1.0, unit=unit) + + +class TestLabelling: + def test_unique_names_are_left_alone(self): + # WHEN nothing is ambiguous, a qualifier would only cost width + parameters = [make_parameter('area'), make_parameter('width')] + labels = ParameterLabels(parameters, qualify=lambda _p: 'Q_index=0') + + # EXPECT + assert [labels.label(p) for p in parameters] == ['area', 'width'] + + def test_shared_names_are_qualified(self): + # WHEN two parameters share a name + first, second = make_parameter('width'), make_parameter('width') + owners = {first.unique_name: 'Q_index=0', second.unique_name: 'Q_index=1'} + labels = ParameterLabels([first, second], qualify=lambda p: owners[p.unique_name]) + + # EXPECT + assert labels.label(first) == 'width (Q_index=0)' + assert labels.label(second) == 'width (Q_index=1)' + + def test_a_qualifier_that_declines_leaves_the_name_alone(self): + # WHEN the qualifier cannot identify an owner, as for a parameter shared across Q + first, second = make_parameter('width'), make_parameter('width') + labels = ParameterLabels([first, second], qualify=lambda _p: None) + + # EXPECT the plain name rather than an invented qualifier + assert labels.label(first) == 'width' + + def test_without_a_qualifier_names_stay_bare(self): + # WHEN + first, second = make_parameter('width'), make_parameter('width') + labels = ParameterLabels([first, second]) + + # EXPECT + assert labels.label(first) == 'width' + + +class TestChainColumns: + def test_columns_resolve_by_unique_name(self): + # WHEN + parameters = [make_parameter('area'), make_parameter('width')] + labels = ParameterLabels(parameters) + columns = [p.unique_name for p in reversed(parameters)] + + # EXPECT resolution follows the chain's order, not the parameter list's + assert labels.resolve(columns) == list(reversed(parameters)) + assert labels.display_names(columns) == ['width', 'area'] + assert labels.units(columns) == ['meV', 'meV'] + + def test_a_saved_chain_resolves_through_its_labels(self): + # WHEN a chain was saved in another session, so its unique names mean nothing here + original = make_parameter('width') + saved = {original.unique_name: 'width'} + current = make_parameter('width') + labels = ParameterLabels([current]) + + # EXPECT the saved label finds the parameter this session has + assert labels.resolve([original.unique_name], saved) == [current] + assert labels.display_names([original.unique_name], saved) == ['width'] + + def test_an_unknown_column_is_reported_not_guessed(self): + # WHEN + labels = ParameterLabels([make_parameter('area')]) + + # EXPECT None rather than a wrong parameter, and the raw name to show something + assert labels.resolve(['Parameter_999']) == [None] + assert labels.display_names(['Parameter_999']) == ['Parameter_999'] + assert labels.units(['Parameter_999']) == [''] + + def test_name_map_records_labels_against_unique_names(self): + # WHEN + first, second = make_parameter('width'), make_parameter('width') + owners = {first.unique_name: 'Q_index=0', second.unique_name: 'Q_index=1'} + labels = ParameterLabels([first, second], qualify=lambda p: owners[p.unique_name]) + + # EXPECT what save() writes alongside a chain + assert labels.name_map() == { + first.unique_name: 'width (Q_index=0)', + second.unique_name: 'width (Q_index=1)', + } + + +class TestCost: + def test_labelling_does_not_rescan_per_parameter(self): + # WHEN there are many parameters. Computing the name counts per parameter is quadratic, + # which was seconds of work for an analysis with many Q values. + parameters = [make_parameter(f'p{i // 2}') for i in range(400)] + labels = ParameterLabels(parameters, qualify=lambda _p: 'q') + + # EXPECT labelling all of them stays cheap + import time + + start = time.perf_counter() + names = [labels.label(p) for p in parameters] + assert time.perf_counter() - start < 0.5 + assert len(names) == len(parameters) + assert np.all([n.endswith('(q)') for n in names]) From 8297346840b46257e73bbaac7907332eef1641b4 Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Fri, 14 Aug 2026 15:27:04 +0200 Subject: [PATCH 14/29] Export the multi-Q sampler and drop the mixin's name The section headers still pointed at a class that no longer exists, and MultiQPosteriorSampler was reachable only through Analysis.bayesian. Co-Authored-By: Claude Opus 5 (1M context) --- src/easydynamics/analysis/__init__.py | 2 ++ src/easydynamics/analysis/analysis.py | 2 +- src/easydynamics/analysis/parameter_analysis.py | 2 +- 3 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/easydynamics/analysis/__init__.py b/src/easydynamics/analysis/__init__.py index 89126ecdf..c6eb02a92 100644 --- a/src/easydynamics/analysis/__init__.py +++ b/src/easydynamics/analysis/__init__.py @@ -8,12 +8,14 @@ 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 MultiQPosteriorSampler from easydynamics.analysis.posterior_sampling import PosteriorSampler __all__ = [ 'Analysis', 'BoundsSuggestion', 'BoundsSuggestions', + 'MultiQPosteriorSampler', 'ParameterAnalysis', 'ParameterLabels', 'ParameterPosterior', diff --git a/src/easydynamics/analysis/analysis.py b/src/easydynamics/analysis/analysis.py index c68961a81..46645afa1 100644 --- a/src/easydynamics/analysis/analysis.py +++ b/src/easydynamics/analysis/analysis.py @@ -800,7 +800,7 @@ def _create_analysis_list(self) -> None: ############# ############# - # Hooks for BayesianSamplingMixin (simultaneous sampling over all Q) + # The contract PosteriorSampler relies on (simultaneous sampling over all Q) ############# def _build_fitter(self) -> MultiFitter: diff --git a/src/easydynamics/analysis/parameter_analysis.py b/src/easydynamics/analysis/parameter_analysis.py index 0ff4154c6..1ac496842 100644 --- a/src/easydynamics/analysis/parameter_analysis.py +++ b/src/easydynamics/analysis/parameter_analysis.py @@ -328,7 +328,7 @@ def _build_fit_inputs(self) -> tuple[list, list, list, list, list]: return xs, ys, ws, funcs, models ############# - # Hooks for BayesianSamplingMixin + # The contract PosteriorSampler relies on ############# def _build_fitter(self) -> MultiFitter: From 85f8a714df32ce1eafc6f2e2653388df030eb04e Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Thu, 13 Aug 2026 14:26:55 +0200 Subject: [PATCH 15/29] Warm the tutorial data cache before running notebooks in parallel The notebook tests run with '-n auto', and five of the notebooks fetch vanadium_data_example.h5 through pooch. On a cold cache the workers race: one is still writing the file into the cache while another opens it, which fails on Windows with "PermissionError: Permission denied". This failed twice in a row on windows-latest, always on that file, always with the other sixteen notebooks passing. The race is pre-existing, but adding a fifth notebook that wants the same file, and lengthening tutorial 1, made it reliable rather than rare. Fetching every tutorial data file once, before the parallel run starts, leaves the workers with nothing to do but read, which is safe. The prefetch reads the URLs and hashes out of the notebooks themselves, so it cannot drift from what they actually download, and it never fails the run: a file it cannot fetch is left to the notebook that needs it, which reports the problem with far more context. Co-Authored-By: Claude Opus 5 (1M context) (cherry picked from commit 46d745a4a73e8025c37e02e490fab67cfdb5ff22) --- pixi.toml | 8 ++- tools/prefetch_tutorial_data.py | 92 +++++++++++++++++++++++++++++++++ 2 files changed, 99 insertions(+), 1 deletion(-) create mode 100644 tools/prefetch_tutorial_data.py diff --git a/pixi.toml b/pixi.toml index f26b4fb7e..db46523e1 100644 --- a/pixi.toml +++ b/pixi.toml @@ -100,7 +100,13 @@ user = { features = ['py-max', 'user'] } unit-tests = 'python -m pytest tests/unit/ --color=yes -v' functional-tests = 'python -m pytest tests/functional/ --color=yes -v' integration-tests = 'python -m pytest tests/integration/ --color=yes -n auto -v' -notebook-tests = 'python -m pytest --nbmake docs/docs/tutorials/ --nbmake-timeout=1200 --color=yes -n auto -v' +# Warm the pooch cache first. Several notebooks fetch the same file, and running them with +# '-n auto' has the workers race: one writes the file while another opens it, which fails on +# Windows. Fetching up front leaves the parallel run with nothing to do but read. +prefetch-tutorial-data = 'python tools/prefetch_tutorial_data.py' +notebook-tests = { cmd = 'python -m pytest --nbmake docs/docs/tutorials/ --nbmake-timeout=1200 --color=yes -n auto -v', depends-on = [ + 'prefetch-tutorial-data', +] } test = { depends-on = ['unit-tests'] } diff --git a/tools/prefetch_tutorial_data.py b/tools/prefetch_tutorial_data.py new file mode 100644 index 000000000..839b1897a --- /dev/null +++ b/tools/prefetch_tutorial_data.py @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: 2026 EasyScience contributors +# SPDX-License-Identifier: BSD-3-Clause + +""" +Download every data file the tutorial notebooks fetch, once, before they are run. + +The notebooks are executed in parallel with ``pytest -n auto``, and several of them fetch the same +file through ``pooch``. On a cold cache the workers race: one is still writing the file into the +cache while another tries to open it, which fails on Windows with a permission error. Fetching +everything up front leaves the parallel run with nothing to do but read. + +Run as ``python tools/prefetch_tutorial_data.py``; it is wired into the ``notebook-tests`` task. +""" + +from __future__ import annotations + +import json +import re +import sys +from pathlib import Path + +import pooch + +TUTORIALS = Path(__file__).resolve().parent.parent / 'docs' / 'docs' / 'tutorials' + +# Matches the pooch.retrieve(url=..., known_hash=...) calls the notebooks use, in either order. +URL_PATTERN = re.compile(r"url\s*=\s*f?['\"]([^'\"]+)['\"]") +HASH_PATTERN = re.compile(r"known_hash\s*=\s*['\"]([^'\"]+)['\"]") + + +def find_downloads() -> dict[str, str]: + """ + Collect the ``(url, known_hash)`` pairs the notebooks fetch. + + Returns + ------- + dict[str, str] + Mapping of URL to expected hash, deduplicated across notebooks. + """ + downloads: dict[str, str] = {} + for notebook in sorted(TUTORIALS.glob('*.ipynb')): + cells = json.loads(notebook.read_text(encoding='utf-8'))['cells'] + for cell in cells: + if cell['cell_type'] != 'code': + continue + source = ''.join(cell['source']) + if 'pooch.retrieve' not in source: + continue + urls = URL_PATTERN.findall(source) + hashes = HASH_PATTERN.findall(source) + # Only pairs are usable; a templated URL without a literal hash is skipped rather than + # guessed at, and the notebook will simply fetch it itself. + for url, known_hash in zip(urls, hashes, strict=False): + downloads[url] = known_hash + return downloads + + +def main() -> int: + """ + Fetch every tutorial data file into the pooch cache. + + Deliberately never fails: this only warms a cache. A file that cannot be fetched here is left + to the notebook that needs it, which reports the problem with far more context than this script + could, and which is where the failure belongs. + + Returns + ------- + int + Always zero. + """ + downloads = find_downloads() + if not downloads: + sys.stdout.write('No tutorial downloads found.\n') + return 0 + + failures = 0 + for url, known_hash in downloads.items(): + name = url.rsplit('/', 1)[-1] + try: + pooch.retrieve(url=url, known_hash=known_hash) + except Exception as error: # noqa: BLE001 - report and continue, the notebook will retry + failures += 1 + sys.stdout.write(f'could not prefetch {name}, leaving it to the notebook: {error}\n') + else: + sys.stdout.write(f'cached {name}\n') + + sys.stdout.write(f'{len(downloads) - failures}/{len(downloads)} tutorial data files ready.\n') + return 0 + + +if __name__ == '__main__': + sys.exit(main()) From 7e53eb727f67e719d6a167218291c4de7c8da788 Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Sun, 16 Aug 2026 22:16:02 +0200 Subject: [PATCH 16/29] Mark setup, action and expectation apart in the new tests The sampling tests labelled the action WHEN and had no THEN, so a reader could not see where the arrangement stopped and the call under test began. Setup is WHEN, the action is THEN, the assertions are EXPECT, and steps that genuinely collapse onto one statement carry one combined marker instead. Comments only; no test changed what it does. Co-Authored-By: Claude Opus 5 (1M context) --- .../fitting/test_bayesian_sampling.py | 16 ++++- .../analysis/test_analysis1d_bayesian.py | 60 ++++++++++++------- .../easydynamics/analysis/test_posterior.py | 14 +++-- .../analysis/test_posterior_labels.py | 18 +++++- .../utils/test_posterior_plotting.py | 25 ++++---- 5 files changed, 90 insertions(+), 43 deletions(-) diff --git a/tests/integration/fitting/test_bayesian_sampling.py b/tests/integration/fitting/test_bayesian_sampling.py index b6adb3bd2..9b5997388 100644 --- a/tests/integration/fitting/test_bayesian_sampling.py +++ b/tests/integration/fitting/test_bayesian_sampling.py @@ -85,8 +85,10 @@ def sampled_analysis(): class TestRealChain: def test_chain_has_one_column_per_free_parameter(self, sampled_analysis): - # EXPECT + # THEN results = sampled_analysis.bayesian.results + + # EXPECT assert results.draws.shape[1] == len(sampled_analysis.get_free_parameters()) assert results.draws.shape[0] > 0 @@ -95,7 +97,7 @@ def test_chain_has_one_column_per_free_parameter(self, sampled_analysis): [('Gaussian area', TRUE_AREA), ('Gaussian width', TRUE_WIDTH)], ) def test_posterior_recovers_the_true_parameters(self, sampled_analysis, name, truth): - # WHEN + # THEN entry = sampled_analysis.bayesian.summary()[name] # EXPECT the truth sits within a few posterior standard deviations of the median. A 68% @@ -105,7 +107,7 @@ def test_posterior_recovers_the_true_parameters(self, sampled_analysis, name, tr assert abs(entry.median - truth) < 4 * spread def test_summary_is_reported_under_parameter_names_and_units(self, sampled_analysis): - # WHEN + # THEN summary = sampled_analysis.bayesian.summary() # EXPECT @@ -121,6 +123,7 @@ def test_sampling_leaves_the_fitted_values_untouched(self): 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) @@ -133,6 +136,7 @@ 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( @@ -150,6 +154,8 @@ def test_save_and_load_round_trip_keeps_parameter_identity(self, sampled_analysi fresh = build_analysis() fresh.fit() fresh.bayesian.suggest_bounds().apply() + + # THEN with warnings.catch_warnings(): warnings.simplefilter('ignore') fresh.bayesian.load(prefix) @@ -165,6 +171,7 @@ def test_subset_sampling_produces_a_single_column(self): analysis.fit() analysis.bayesian.suggest_bounds().apply() + # THEN with warnings.catch_warnings(): warnings.simplefilter('ignore') results = analysis.bayesian.sample(parameters=['Gaussian width'], **SAMPLE_KWARGS) @@ -178,6 +185,8 @@ def test_plots_render(self, sampled_analysis): 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) @@ -195,6 +204,7 @@ def test_posterior_median_is_close_to_the_least_squares_fit(self): 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) diff --git a/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py b/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py index 345bc2724..6b5d7a5fe 100644 --- a/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py +++ b/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py @@ -78,7 +78,7 @@ def analysis(): class TestFitterExposure: def test_fitter_is_built_lazily_and_cached(self, analysis): - # WHEN + # THEN fitter = analysis.fitter # EXPECT @@ -88,20 +88,22 @@ def test_fitter_is_built_lazily_and_cached(self, analysis): def test_fitter_is_rebuilt_when_the_sample_model_changes(self, analysis): # WHEN original = analysis.fitter + + # THEN analysis.sample_model = SampleModel(components=Gaussian(area=1.0)) # EXPECT assert analysis.fitter is not original def test_minimizer_can_be_switched_through_the_fitter(self, analysis): - # WHEN + # THEN analysis.fitter.switch_minimizer(AvailableMinimizers.Bumps) # EXPECT assert analysis.fitter.minimizer.enum == AvailableMinimizers.Bumps def test_fit_uses_the_persistent_fitter(self, analysis): - # WHEN + # THEN result = analysis.fit() # EXPECT @@ -111,12 +113,12 @@ def test_fit_uses_the_persistent_fitter(self, analysis): class TestBoundsPreflight: def test_sampling_refuses_unbounded_parameters(self, analysis): - # EXPECT + # THEN EXPECT with pytest.raises(ValueError, match='finite bounds'): analysis.bayesian.sample(samples=10) def test_error_names_the_offending_parameters(self, analysis): - # EXPECT + # THEN EXPECT with pytest.raises(ValueError, match='Gaussian area'): analysis.bayesian.check_bounds() @@ -124,11 +126,11 @@ def test_bounded_parameters_pass(self, analysis): # WHEN bound_all(analysis) - # EXPECT: does not raise + # THEN EXPECT: does not raise analysis.bayesian.check_bounds() def test_suggest_bounds_covers_the_free_parameters(self, analysis): - # WHEN + # THEN suggestions = analysis.bayesian.suggest_bounds() # EXPECT @@ -141,6 +143,7 @@ def test_restores_parameter_values_and_minimizer(self, analysis): 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): @@ -162,6 +165,7 @@ def test_switches_to_bumps_for_the_run(self, analysis): 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), @@ -176,6 +180,7 @@ 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'): @@ -188,6 +193,7 @@ 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) @@ -203,6 +209,7 @@ 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 @@ -222,7 +229,7 @@ def test_warns_when_the_posterior_piles_up_against_a_bound(self, analysis): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.return_value = fake_results(analysis, values=draws) - # EXPECT + # THEN EXPECT with pytest.warns(UserWarning, match='piled up'): analysis.bayesian.sample(samples=10) @@ -233,7 +240,7 @@ def test_does_not_warn_when_the_posterior_is_well_inside(self, analysis): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.return_value = fake_results(analysis) - # EXPECT + # THEN EXPECT with warnings_as_errors(): analysis.bayesian.sample(samples=10) @@ -245,6 +252,7 @@ def test_holds_other_parameters_fixed_during_the_run(self, analysis): target = analysis.get_free_parameters()[0] seen = {} + # THEN with patch(SAMPLER_PATH) as sampler_class: def record(**_kwargs): @@ -264,6 +272,7 @@ def test_restores_the_fixed_flags_afterwards(self, 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): @@ -276,17 +285,17 @@ def test_unknown_parameter_name_raises(self, analysis): # WHEN bound_all(analysis) - # EXPECT + # THEN EXPECT with pytest.raises(ValueError, match='No free parameter named'): analysis.bayesian.sample(samples=10, parameters=['not a parameter']) def test_non_list_parameters_raises(self, analysis): - # EXPECT + # 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): - # EXPECT + # THEN EXPECT with pytest.raises(ValueError, match='at least one parameter'): analysis.bayesian.sample(samples=10, parameters=[]) @@ -296,6 +305,7 @@ 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) @@ -308,6 +318,7 @@ 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) @@ -321,6 +332,7 @@ 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) @@ -335,7 +347,7 @@ def test_binds_the_same_data_the_fit_uses(self, analysis): class TestExtendAndPersistence: def test_extend_without_a_chain_raises(self, analysis): - # EXPECT + # THEN EXPECT with pytest.raises(RuntimeError, match='No chain to extend'): analysis.bayesian.extend() @@ -343,6 +355,7 @@ 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) @@ -355,7 +368,7 @@ def test_extend_delegates_to_the_sampler(self, analysis): assert kwargs['thin'] == 2 def test_save_without_a_chain_raises(self, analysis): - # EXPECT + # THEN EXPECT with pytest.raises(RuntimeError, match='No chain to save'): analysis.bayesian.save('somewhere') @@ -364,6 +377,8 @@ def test_save_writes_the_parameter_name_sidecar(self, analysis, tmp_path): 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) @@ -382,14 +397,14 @@ def test_load_without_a_sidecar_warns(self, analysis, tmp_path): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.load_state.return_value = fake_results(analysis) - # EXPECT + # THEN EXPECT with pytest.warns(UserWarning, match='No parameter-name sidecar'): analysis.bayesian.load(str(tmp_path / 'missing')) class TestResults: def test_summary_without_sampling_raises(self, analysis): - # EXPECT + # THEN EXPECT with pytest.raises(RuntimeError, match='No posterior samples yet'): analysis.bayesian.summary() @@ -401,8 +416,10 @@ def test_summary_uses_parameter_names_and_units(self, analysis): sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) analysis.bayesian.sample(samples=10) - # EXPECT + # 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) @@ -418,6 +435,7 @@ def test_set_parameters_to_posterior_median(self, analysis): expected = [float(p.value) + 2.0 for p in parameters] analysis.bayesian.sample(samples=10) + # THEN changed = analysis.bayesian.set_parameters_to_median() # EXPECT @@ -425,7 +443,7 @@ def test_set_parameters_to_posterior_median(self, analysis): assert [float(p.value) for p in parameters] == pytest.approx(expected) def test_median_without_sampling_raises(self, analysis): - # EXPECT + # THEN EXPECT with pytest.raises(RuntimeError, match='No posterior samples yet'): analysis.bayesian.set_parameters_to_median() @@ -439,7 +457,7 @@ def test_predictive_rejects_a_bad_draw_count(self, analysis): sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) analysis.bayesian.sample(samples=10) - # EXPECT + # THEN EXPECT with pytest.raises(ValueError, match='positive integer'): analysis.bayesian.plot_posterior_predictive(n_draws=0) @@ -454,13 +472,15 @@ def test_predictive_restores_parameter_values(self, analysis): 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): - # EXPECT + # THEN EXPECT with pytest.raises(RuntimeError): analysis.bayesian.plot_trace() with pytest.raises(RuntimeError): diff --git a/tests/unit/easydynamics/analysis/test_posterior.py b/tests/unit/easydynamics/analysis/test_posterior.py index 659412774..732dcb8df 100644 --- a/tests/unit/easydynamics/analysis/test_posterior.py +++ b/tests/unit/easydynamics/analysis/test_posterior.py @@ -109,12 +109,12 @@ def test_non_finite_value_is_flagged(self): @pytest.mark.parametrize('kwargs', [{'n_sigma': -1.0}, {'relative_pad': -0.1}]) def test_negative_settings_raise(self, kwargs): - # EXPECT + # THEN EXPECT with pytest.raises(ValueError): suggest_bounds_for_parameters([make_parameter()], **kwargs) def test_non_numeric_setting_raises(self): - # EXPECT + # THEN EXPECT with pytest.raises(TypeError): suggest_bounds_for_parameters([make_parameter()], n_sigma='wide') @@ -125,8 +125,10 @@ def test_apply_sets_bounds_and_reports_changes(self): parameter = make_parameter(value=10.0, error=0.5) suggestions = suggest_bounds_for_parameters([parameter]) - # THEN nothing has changed until apply is called + # WHEN nothing has changed until apply is called assert parameter.max == np.inf + + # THEN changed = suggestions.apply() # EXPECT @@ -160,14 +162,14 @@ def test_repr_lists_parameters_and_flags_attention(self): assert 'need bounds set by hand' in text def test_repr_with_no_parameters(self): - # EXPECT + # 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()]) - # EXPECT + # THEN EXPECT assert len(suggestions) == 2 assert all(isinstance(s, BoundsSuggestion) for s in suggestions) @@ -291,7 +293,7 @@ def test_lookup_of_missing_name_raises(self): # WHEN summary = summarize_draws(np.zeros((5, 1)), ['x'], [None]) - # EXPECT + # THEN EXPECT with pytest.raises(KeyError): summary['not a parameter'] diff --git a/tests/unit/easydynamics/analysis/test_posterior_labels.py b/tests/unit/easydynamics/analysis/test_posterior_labels.py index b4b08f86c..82347db72 100644 --- a/tests/unit/easydynamics/analysis/test_posterior_labels.py +++ b/tests/unit/easydynamics/analysis/test_posterior_labels.py @@ -15,6 +15,8 @@ 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 @@ -24,6 +26,8 @@ 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 @@ -33,6 +37,8 @@ def test_shared_names_are_qualified(self): 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 @@ -41,6 +47,8 @@ def test_a_qualifier_that_declines_leaves_the_name_alone(self): def test_without_a_qualifier_names_stay_bare(self): # WHEN first, second = make_parameter('width'), make_parameter('width') + + # THEN labels = ParameterLabels([first, second]) # EXPECT @@ -54,7 +62,7 @@ def test_columns_resolve_by_unique_name(self): labels = ParameterLabels(parameters) columns = [p.unique_name for p in reversed(parameters)] - # EXPECT resolution follows the chain's order, not the parameter list's + # 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'] @@ -64,6 +72,8 @@ def test_a_saved_chain_resolves_through_its_labels(self): 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 @@ -71,7 +81,7 @@ def test_a_saved_chain_resolves_through_its_labels(self): assert labels.display_names([original.unique_name], saved) == ['width'] def test_an_unknown_column_is_reported_not_guessed(self): - # WHEN + # THEN labels = ParameterLabels([make_parameter('area')]) # EXPECT None rather than a wrong parameter, and the raw name to show something @@ -83,6 +93,8 @@ 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 @@ -99,7 +111,7 @@ def test_labelling_does_not_rescan_per_parameter(self): parameters = [make_parameter(f'p{i // 2}') for i in range(400)] labels = ParameterLabels(parameters, qualify=lambda _p: 'q') - # EXPECT labelling all of them stays cheap + # THEN EXPECT labelling all of them stays cheap import time start = time.perf_counter() diff --git a/tests/unit/easydynamics/utils/test_posterior_plotting.py b/tests/unit/easydynamics/utils/test_posterior_plotting.py index 64d0129e9..1a412cdb2 100644 --- a/tests/unit/easydynamics/utils/test_posterior_plotting.py +++ b/tests/unit/easydynamics/utils/test_posterior_plotting.py @@ -27,14 +27,14 @@ def draws(): class TestPlotTrace: def test_one_panel_per_parameter(self, draws): - # WHEN + # THEN fig = plot_trace(draws=draws, names=['a', 'b', 'c']) # EXPECT assert len(fig.axes) == 3 def test_logp_adds_a_panel(self, draws): - # WHEN + # THEN fig = plot_trace(draws=draws, names=['a', 'b', 'c'], logp=np.zeros(len(draws))) # EXPECT @@ -42,47 +42,47 @@ def test_logp_adds_a_panel(self, draws): assert fig.axes[-1].get_ylabel() == 'log-posterior' def test_names_label_the_panels(self, draws): - # WHEN + # 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): - # WHEN + # THEN fig = plot_trace(draws=np.zeros((10, 1)), names=['only']) # EXPECT assert len(fig.axes) == 1 def test_mismatched_names_raise(self, draws): - # EXPECT + # 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): - # EXPECT + # THEN EXPECT with pytest.raises(ValueError, match='two-dimensional'): plot_trace(draws=np.zeros(10), names=['a']) class TestPlotCorner: def test_grid_is_square_in_the_parameter_count(self, draws): - # WHEN + # THEN fig = plot_corner(draws=draws, names=['a', 'b', 'c']) # EXPECT assert len(fig.axes) == 9 def test_upper_triangle_is_hidden(self, draws): - # WHEN + # 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): - # EXPECT + # THEN EXPECT with pytest.raises(ValueError, match='one entry per column'): plot_corner(draws=draws, names=['a']) @@ -93,6 +93,7 @@ def test_returns_a_figure_with_data_and_band(self): 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 @@ -104,6 +105,7 @@ 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), @@ -115,13 +117,13 @@ def test_error_bars_are_drawn_when_given(self): assert len(fig.axes[0].containers) == 1 def test_wrong_prediction_shape_raises(self): - # EXPECT + # 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): - # EXPECT + # THEN EXPECT with pytest.raises(ValueError, match='credible_interval'): plot_posterior_predictive( x=np.zeros(4), @@ -135,6 +137,7 @@ def test_band_widens_with_the_credible_interval(self): 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 ) From 2ecc29fed9abd8bf675c9a3af8755853cdf5572b Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Sun, 16 Aug 2026 22:24:07 +0200 Subject: [PATCH 17/29] Mark setup, action and expectation apart in the multi-Q tests Same pass as on the single-Q tests: setup is WHEN, the action is THEN, the assertions are EXPECT, and a step that collapses onto one statement carries one combined marker. Comments only; no test changed what it does. Co-Authored-By: Claude Opus 5 (1M context) --- .../fitting/test_bayesian_sampling_multi_q.py | 17 ++++++-- .../analysis/test_analysis1d_bayesian.py | 24 +++++++---- .../analysis/test_analysis_bayesian.py | 42 ++++++++++++------- .../test_parameter_analysis_bayesian.py | 31 ++++++++++---- .../easydynamics/analysis/test_posterior.py | 10 ++--- .../utils/test_posterior_plotting.py | 9 ++-- 6 files changed, 86 insertions(+), 47 deletions(-) diff --git a/tests/integration/fitting/test_bayesian_sampling_multi_q.py b/tests/integration/fitting/test_bayesian_sampling_multi_q.py index 24a6d8531..ef7a284af 100644 --- a/tests/integration/fitting/test_bayesian_sampling_multi_q.py +++ b/tests/integration/fitting/test_bayesian_sampling_multi_q.py @@ -85,13 +85,15 @@ def simultaneously_sampled(): class TestSimultaneousChain: def test_chain_covers_every_q_index(self, simultaneously_sampled): - # EXPECT one column per free parameter across all Q, in one chain + # THEN results = simultaneously_sampled.bayesian.results + + # EXPECT one column per free parameter across all Q, in one chain assert results.draws.shape[1] == len(simultaneously_sampled._chain_parameters()) assert results.draws.shape[1] == 3 * len(Q_VALUES) def test_summary_labels_are_unique_and_q_qualified(self, simultaneously_sampled): - # WHEN + # THEN names = [entry.name for entry in simultaneously_sampled.bayesian.summary()] # EXPECT @@ -100,7 +102,7 @@ def test_summary_labels_are_unique_and_q_qualified(self, simultaneously_sampled) @pytest.mark.parametrize('q_index', range(len(Q_VALUES))) def test_posterior_recovers_the_true_width_at_each_q(self, simultaneously_sampled, q_index): - # WHEN + # THEN entry = simultaneously_sampled.bayesian.summary()[f'Gaussian width (Q_index={q_index})'] # EXPECT the truth within a few posterior standard deviations. A 68% interval is not used @@ -115,6 +117,7 @@ def test_sampling_leaves_the_fitted_values_untouched(self): analysis.bayesian.suggest_bounds().apply() before = [float(p.value) for p in analysis._chain_parameters()] + # THEN with warnings.catch_warnings(): warnings.simplefilter('ignore') analysis.bayesian.sample(fit_method='simultaneous', **SAMPLE_KWARGS) @@ -128,6 +131,8 @@ def test_plots_render(self, simultaneously_sampled): import matplotlib.pyplot as plt n_parameters = len(simultaneously_sampled._chain_parameters()) + + # THEN trace = simultaneously_sampled.bayesian.plot_trace() corner = simultaneously_sampled.bayesian.plot_corner() @@ -145,6 +150,7 @@ def test_one_chain_per_q_index(self): for analysis1d in analysis.analysis_list: analysis1d.bayesian.suggest_bounds().apply() + # THEN with warnings.catch_warnings(): warnings.simplefilter('ignore') results = analysis.bayesian.sample(fit_method='independent', **SAMPLE_KWARGS) @@ -155,12 +161,13 @@ def test_one_chain_per_q_index(self): assert result.draws.shape[1] == len(analysis1d.get_free_parameters()) def test_independent_and_simultaneous_agree_on_the_widths(self, simultaneously_sampled): - # WHEN the same data is sampled per-Q instead of all at once + # WHEN analysis = build_analysis() analysis.fit(fit_method='independent') for analysis1d in analysis.analysis_list: analysis1d.bayesian.suggest_bounds().apply() + # THEN the same data is sampled per-Q instead of all at once with warnings.catch_warnings(): warnings.simplefilter('ignore') analysis.bayesian.sample(fit_method='independent', **SAMPLE_KWARGS) @@ -209,6 +216,7 @@ def test_recovers_a_straight_line_through_the_widths(self): analysis.bayesian.suggest_bounds(absolute_floor=1.0).apply() assert not analysis.bayesian.suggest_bounds().needing_attention + # THEN with warnings.catch_warnings(): warnings.simplefilter('ignore') results = analysis.bayesian.sample(**SAMPLE_KWARGS) @@ -252,6 +260,7 @@ def test_median_applies_each_chain_to_its_own_q(self): warnings.simplefilter('ignore') analysis.bayesian.sample(fit_method='independent', **SAMPLE_KWARGS) + # THEN changed = analysis.bayesian.set_parameters_to_median() # EXPECT every Q's parameters land on that Q's own median diff --git a/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py b/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py index 22d664f72..94f515b60 100644 --- a/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py +++ b/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py @@ -529,7 +529,8 @@ def test_bumps_outlier_crash_is_reported_helpfully(self, analysis): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.side_effect = _bumps_style_index_error() - # EXPECT the bare IndexError is replaced by something actionable, naming both causes + # THEN EXPECT the bare IndexError is replaced by something actionable, naming both + # causes with pytest.raises(RuntimeError, match='degenerate') as raised: analysis.bayesian.sample(samples=10) assert 'short chains' in str(raised.value) @@ -543,12 +544,12 @@ def test_an_index_error_of_our_own_is_not_relabelled(self, analysis): with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.side_effect = IndexError('list index out of range') - # EXPECT it propagates untouched + # THEN EXPECT it propagates untouched with pytest.raises(IndexError, match='list index out of range'): analysis.bayesian.sample(samples=10) def test_parameters_entry_of_the_wrong_type_raises(self, analysis): - # EXPECT + # THEN EXPECT with pytest.raises(TypeError, match='Parameter objects or labels'): analysis.bayesian.sample(samples=10, parameters=[42]) @@ -563,8 +564,10 @@ def test_median_skips_columns_with_no_matching_parameter(self, analysis): sampler_class.return_value.sample.return_value = results analysis.bayesian.sample(samples=10) - # EXPECT the unknown column is skipped rather than crashing + # THEN changed = analysis.bayesian.set_parameters_to_median() + + # EXPECT the unknown column is skipped rather than crashing assert len(changed) == len(analysis.get_free_parameters()) def test_load_chain_uses_the_sidecar_when_present(self, analysis, tmp_path): @@ -578,6 +581,8 @@ def test_load_chain_uses_the_sidecar_when_present(self, analysis, tmp_path): fresh = make_analysis() bound_all(fresh) + + # THEN with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.load_state.return_value = saved fresh.bayesian.load(str(tmp_path / 'chain')) @@ -602,7 +607,7 @@ def test_trace_and_corner_render_from_a_chain(self, analysis): sampler_class.return_value.sample.return_value = fake_results(analysis) analysis.bayesian.sample(samples=10) - # EXPECT + # THEN EXPECT assert len(analysis.bayesian.plot_trace().axes) == n_parameters + 1 assert len(analysis.bayesian.plot_corner().axes) == n_parameters**2 plt.close('all') @@ -619,8 +624,8 @@ def test_extending_with_a_different_subset_is_refused(self, analysis): target = analysis.get_free_parameters()[0] - # EXPECT refused up front, rather than failing obscurely inside BUMPS, which resumes - # from a stored chain whose width is fixed + # THEN EXPECT refused up front, rather than failing obscurely inside BUMPS, which + # resumes from a stored chain whose width is fixed with pytest.warns(UserWarning), pytest.raises(ValueError, match='Cannot extend'): analysis.bayesian.extend(additional_samples=10, parameters=[target.name]) @@ -633,7 +638,7 @@ def test_extending_with_the_same_parameters_is_allowed(self, analysis): sampler_class.return_value.extend.side_effect = lambda **_k: fake_results(analysis) analysis.bayesian.sample(samples=10) - # EXPECT: does not raise + # THEN EXPECT: does not raise analysis.bayesian.extend(additional_samples=10) @@ -645,6 +650,7 @@ def test_a_subset_run_records_the_same_labels_a_full_run_would(self, analysis): bound_all(analysis) 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): @@ -668,6 +674,6 @@ def test_extending_after_a_failed_run_is_allowed(self, analysis): assert analysis.bayesian.sampler is not None assert analysis.bayesian.results is None - # EXPECT the shape guard steps aside rather than comparing against nothing + # THEN EXPECT the shape guard steps aside rather than comparing against nothing sampler_class.return_value.extend.side_effect = lambda **_k: fake_results(analysis) analysis.bayesian.extend(additional_samples=10) diff --git a/tests/unit/easydynamics/analysis/test_analysis_bayesian.py b/tests/unit/easydynamics/analysis/test_analysis_bayesian.py index a263691c8..c552fd7b4 100644 --- a/tests/unit/easydynamics/analysis/test_analysis_bayesian.py +++ b/tests/unit/easydynamics/analysis/test_analysis_bayesian.py @@ -69,7 +69,7 @@ def analysis(): class TestChainParameters: def test_union_covers_every_q_index(self, analysis): - # WHEN + # THEN parameters = analysis._chain_parameters() # EXPECT one copy of each per-Q parameter, with no duplicates @@ -77,7 +77,7 @@ def test_union_covers_every_q_index(self, analysis): assert len({p.unique_name for p in parameters}) == len(parameters) def test_labels_are_qualified_by_q_index(self, analysis): - # WHEN + # THEN labels = [analysis._parameter_labels().label(p) for p in analysis._chain_parameters()] # EXPECT every per-Q copy is distinguishable, which the bare name would not be @@ -86,7 +86,7 @@ def test_labels_are_qualified_by_q_index(self, analysis): assert 'Gaussian width (Q_index=2)' in labels def test_bare_names_would_collide(self, analysis): - # WHEN + # THEN names = [p.name for p in analysis._chain_parameters()] # EXPECT the collision the Q-qualified label exists to solve @@ -95,17 +95,17 @@ def test_bare_names_would_collide(self, analysis): class TestBoundsPreflight: def test_sampling_refuses_unbounded_parameters(self, analysis): - # EXPECT + # THEN EXPECT with pytest.raises(ValueError, match='finite bounds'): analysis.bayesian.sample(fit_method='simultaneous', samples=10) def test_error_names_parameters_by_q_index(self, analysis): - # EXPECT + # THEN EXPECT with pytest.raises(ValueError, match=r'Gaussian width \(Q_index=0\)'): analysis.bayesian.check_bounds() def test_suggest_bounds_labels_every_q(self, analysis): - # WHEN + # THEN suggestions = analysis.bayesian.suggest_bounds() # EXPECT @@ -120,6 +120,7 @@ def test_binds_one_dataset_per_q_index(self, analysis): bound_all(analysis) parameters = analysis._chain_parameters() + # THEN with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.return_value = fake_results(parameters) analysis.bayesian.sample(fit_method='simultaneous', samples=10) @@ -135,6 +136,7 @@ def test_returns_a_single_result(self, analysis): bound_all(analysis) parameters = analysis._chain_parameters() + # THEN with patch(SAMPLER_PATH) as sampler_class: expected = fake_results(parameters) sampler_class.return_value.sample.return_value = expected @@ -149,6 +151,7 @@ def test_summary_is_labelled_by_q_index(self, analysis): bound_all(analysis) parameters = analysis._chain_parameters() + # THEN with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.return_value = fake_results(parameters) analysis.bayesian.sample(fit_method='simultaneous', samples=10) @@ -163,6 +166,7 @@ def test_refreshes_every_convolver_before_sampling(self, analysis): bound_all(analysis) parameters = analysis._chain_parameters() + # THEN with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.return_value = fake_results(parameters) for analysis1d in analysis.analysis_list: @@ -189,6 +193,7 @@ def test_returns_one_result_per_q_index(self, analysis): parameter.min = float(parameter.value) - 5.0 parameter.max = float(parameter.value) + 5.0 + # THEN with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.side_effect = lambda **_k: fake_results( analysis.analysis_list[0].get_free_parameters() @@ -206,6 +211,7 @@ def test_single_q_index_returns_one_result(self, analysis): parameter.min = float(parameter.value) - 5.0 parameter.max = float(parameter.value) + 5.0 + # THEN with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.side_effect = lambda **_k: fake_results( target.get_free_parameters() @@ -217,14 +223,14 @@ def test_single_q_index_returns_one_result(self, analysis): assert result is target.bayesian.results def test_invalid_q_index_raises(self, analysis): - # EXPECT + # THEN EXPECT with pytest.raises((ValueError, IndexError)): analysis.bayesian.sample(fit_method='independent', Q_index=99, samples=10) class TestSamplePosteriorValidation: def test_invalid_fit_method_raises(self, analysis): - # EXPECT + # THEN EXPECT with pytest.raises(ValueError, match='Invalid fit method'): analysis.bayesian.sample(fit_method='nonsense') @@ -232,7 +238,7 @@ def test_missing_q_values_raises(self): # WHEN analysis = edyn.Analysis(display_name='Empty') - # EXPECT + # THEN EXPECT with pytest.raises(ValueError, match='No Q values available'): analysis.bayesian.sample() @@ -247,7 +253,7 @@ def test_predictive_is_not_supported_for_multiple_datasets(self, analysis): sampler_class.return_value.sample.return_value = fake_results(parameters) analysis.bayesian.sample(fit_method='simultaneous', samples=10) - # EXPECT + # THEN EXPECT with pytest.raises(NotImplementedError, match='single dataset only'): analysis.bayesian.plot_posterior_predictive() @@ -309,13 +315,13 @@ def test_operations_needing_one_chain_point_at_the_per_q_chains(self, analysis): ) analysis.bayesian.sample(fit_method='independent', samples=10) - # EXPECT anything that genuinely needs a single chain says where the chains actually are, - # rather than claiming none exist + # THEN EXPECT anything that genuinely needs a single chain says where the chains + # actually are, rather than claiming none exist with pytest.raises(RuntimeError, match='analysis_list'): analysis.bayesian.plot_posterior_predictive() def test_untouched_analysis_still_reports_no_samples(self, analysis): - # EXPECT the plain message when nothing has been sampled anywhere + # THEN EXPECT the plain message when nothing has been sampled anywhere with pytest.raises(RuntimeError, match='No posterior samples yet'): analysis.bayesian.summary() @@ -431,6 +437,7 @@ def test_corner_offers_a_slider_in_a_notebook(self, analysis): # WHEN self._sample_independently(analysis) + # THEN with patch('easydynamics.analysis.posterior_sampling._in_notebook', return_value=True): widget = analysis.bayesian.plot_corner() @@ -452,7 +459,7 @@ def test_corner_without_a_notebook_or_q_index_says_what_to_do(self, analysis): # WHEN self._sample_independently(analysis) - # EXPECT it names the sampled Q indices rather than just refusing + # THEN EXPECT it names the sampled Q indices rather than just refusing with ( patch('easydynamics.analysis.posterior_sampling._in_notebook', return_value=False), pytest.raises(RuntimeError, match=r'sampled Q indices are \[0, 1, 2\]'), @@ -472,6 +479,7 @@ def test_the_slider_only_offers_q_indices_that_were_sampled(self, analysis): ) analysis.bayesian.sample(fit_method='independent', Q_index=2, samples=10) + # THEN with patch('easydynamics.analysis.posterior_sampling._in_notebook', return_value=True): widget = analysis.bayesian.plot_corner() @@ -482,7 +490,7 @@ def test_trace_points_at_the_individual_chains(self, analysis): # WHEN self._sample_independently(analysis) - # EXPECT + # THEN EXPECT with pytest.raises(RuntimeError, match='no single trace'): analysis.bayesian.plot_trace() @@ -513,8 +521,10 @@ def test_only_the_sampled_q_indices_are_gathered(self, analysis): ) analysis.bayesian.sample(fit_method='independent', Q_index=1, samples=10) - # EXPECT the unsampled Q indices are passed over rather than breaking the aggregation + # THEN summary = analysis.bayesian.summary() + + # EXPECT the unsampled Q indices are passed over rather than breaking the aggregation assert len(summary) == len(target.get_free_parameters()) assert all('Q_index=1' in entry.name for entry in summary) assert len(analysis.bayesian.set_parameters_to_median()) == len( diff --git a/tests/unit/easydynamics/analysis/test_parameter_analysis_bayesian.py b/tests/unit/easydynamics/analysis/test_parameter_analysis_bayesian.py index de14ebd64..7966e1f87 100644 --- a/tests/unit/easydynamics/analysis/test_parameter_analysis_bayesian.py +++ b/tests/unit/easydynamics/analysis/test_parameter_analysis_bayesian.py @@ -87,7 +87,7 @@ def test_fitter_is_a_cached_multifitter(self, analysis): assert analysis.fitter is analysis.fitter def test_fit_still_returns_per_target_results(self, analysis): - # WHEN + # THEN results = analysis.fit() # EXPECT one result per fit target, as before @@ -97,6 +97,8 @@ def test_fit_still_returns_per_target_results(self, analysis): def test_changing_bindings_rebuilds_the_fitter(self, analysis): # WHEN original = analysis.fitter + + # THEN analysis.bindings = analysis.bindings[:1] # EXPECT @@ -105,6 +107,8 @@ def test_changing_bindings_rebuilds_the_fitter(self, analysis): def test_changing_parameters_rebuilds_the_fitter(self, analysis): # WHEN original = analysis.fitter + + # THEN analysis.parameters = make_dataset() # EXPECT @@ -113,7 +117,7 @@ def test_changing_parameters_rebuilds_the_fitter(self, analysis): class TestChainParameters: def test_covers_every_binding_model(self, analysis): - # WHEN + # THEN parameters = analysis._chain_parameters() # EXPECT both Polynomials contribute their two coefficients @@ -121,7 +125,7 @@ def test_covers_every_binding_model(self, analysis): assert len({p.unique_name for p in parameters}) == 4 def test_labels_are_unique(self, analysis): - # WHEN + # THEN labels = [analysis._parameter_labels().label(p) for p in analysis._chain_parameters()] # EXPECT @@ -129,6 +133,8 @@ def test_labels_are_unique(self, analysis): def test_model_name_is_not_repeated_in_the_label(self, analysis): # WHEN a model already names its parameters after itself + + # THEN labels = [analysis._parameter_labels().label(p) for p in analysis._chain_parameters()] # EXPECT no 'Width line: Width line_c0' @@ -138,7 +144,7 @@ def test_model_name_is_not_repeated_in_the_label(self, analysis): class TestSampling: def test_refuses_unbounded_parameters(self, analysis): - # EXPECT + # THEN EXPECT with pytest.raises(ValueError, match='finite bounds'): analysis.bayesian.sample(samples=10) @@ -147,6 +153,7 @@ def test_binds_one_dataset_per_target(self, analysis): bound_all(analysis) parameters = analysis._chain_parameters() + # THEN with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.return_value = fake_results(parameters) analysis.bayesian.sample(samples=10) @@ -161,6 +168,7 @@ def test_summary_uses_model_qualified_labels(self, analysis): bound_all(analysis) parameters = analysis._chain_parameters() + # THEN with patch(SAMPLER_PATH) as sampler_class: sampler_class.return_value.sample.return_value = fake_results(parameters) analysis.bayesian.sample(samples=10) @@ -176,6 +184,7 @@ def test_restores_parameter_values(self, analysis): parameters = analysis._chain_parameters() before = [float(p.value) for p in parameters] + # THEN with patch(SAMPLER_PATH) as sampler_class: def mutate(**_kwargs): @@ -193,7 +202,7 @@ def test_missing_parameters_dataset_raises(self): # WHEN analysis = edyn.ParameterAnalysis() - # EXPECT + # THEN EXPECT with pytest.raises(ValueError, match='No parameters Dataset'): analysis.bayesian.sample(samples=10) @@ -201,7 +210,7 @@ def test_missing_bindings_raises(self): # WHEN analysis = edyn.ParameterAnalysis(parameters=make_dataset()) - # EXPECT + # THEN EXPECT with pytest.raises(ValueError, match='No fit bindings'): analysis.bayesian.sample(samples=10) @@ -236,8 +245,10 @@ def test_single_binding_keeps_plain_names(self): # WHEN analysis = make_analysis(two_bindings=False) - # EXPECT no model prefix, since there is nothing to disambiguate + # THEN labels = [analysis._parameter_labels().label(p) for p in analysis._chain_parameters()] + + # EXPECT no model prefix, since there is nothing to disambiguate assert labels == ['Width line_c0', 'Width line_c1'] def test_parameter_from_outside_the_analysis_keeps_its_name(self, analysis): @@ -246,7 +257,7 @@ def test_parameter_from_outside_the_analysis_keeps_its_name(self, analysis): stranger = Parameter(name='Width line_c0', value=1.0) - # EXPECT it is returned unqualified rather than mislabelled + # THEN EXPECT it is returned unqualified rather than mislabelled assert analysis._parameter_labels().label(stranger) == 'Width line_c0' def test_models_without_a_display_name_fall_back_to_the_unique_name(self): @@ -321,7 +332,7 @@ def test_ambiguous_name_owned_by_no_model_keeps_its_name(self): ) stranger = Parameter(name='Line_c0', value=1.0) - # EXPECT it falls back to the plain name rather than claiming an owner + # THEN EXPECT it falls back to the plain name rather than claiming an owner assert analysis._parameter_labels().label(stranger) == 'Line_c0' @@ -342,6 +353,7 @@ def test_changing_the_number_of_targets_rebuilds_the_fitter(self): analysis = edyn.ParameterAnalysis(parameters=make_dataset(), bindings=[binding]) assert len(analysis.fit()) == 1 + # THEN binding.targets = {'width': 'Lorentzian width', 'area': 'Lorentzian area'} # EXPECT the fit follows the binding rather than failing on a stale fitter @@ -361,6 +373,7 @@ def test_shrinking_the_targets_also_rebuilds(self): analysis = edyn.ParameterAnalysis(parameters=make_dataset(), bindings=[binding]) assert len(analysis.fit()) == 2 + # THEN binding.targets = {'width': 'Lorentzian width'} # EXPECT diff --git a/tests/unit/easydynamics/analysis/test_posterior.py b/tests/unit/easydynamics/analysis/test_posterior.py index 9e891e534..04bd60349 100644 --- a/tests/unit/easydynamics/analysis/test_posterior.py +++ b/tests/unit/easydynamics/analysis/test_posterior.py @@ -316,13 +316,13 @@ def test_len_and_iteration(self): parameters = [make_parameter(name='a'), make_parameter(name='b')] summary = summarize_draws(np.zeros((7, 2)), ['a', 'b'], parameters) - # EXPECT + # THEN EXPECT assert len(summary) == 2 assert [entry.name for entry in summary] == ['a', 'b'] assert len(summary.entries) == 2 def test_repr_with_no_entries(self): - # EXPECT + # WHEN THEN EXPECT assert 'no parameters' in repr(summarize_draws(np.zeros((3, 0)), [], [])) @@ -333,7 +333,7 @@ def test_applying_a_wildly_wide_bound_warns(self): parameter = make_parameter(name='Delta area', value=1.0, error=1e9) suggestions = suggest_bounds_for_parameters([parameter]) - # EXPECT it is still applied, since it is what the fit implied, but not silently + # THEN EXPECT it is still applied, since it is what the fit implied, but not silently with pytest.warns(UserWarning, match='far wider than the parameter'): changed = suggestions.apply() assert changed == [parameter] @@ -343,7 +343,7 @@ def test_a_sane_bound_applies_without_warning(self): parameter = make_parameter(name='sane', value=10.0, error=0.5) suggestions = suggest_bounds_for_parameters([parameter]) - # EXPECT + # THEN EXPECT import warnings as warnings_module with warnings_module.catch_warnings(): @@ -355,7 +355,7 @@ def test_a_zero_valued_parameter_is_not_called_absurd(self): parameter = make_parameter(name='zero', value=0.0, error=1.0) suggestions = suggest_bounds_for_parameters([parameter]) - # EXPECT no warning, since the ratio is meaningless rather than alarming + # THEN EXPECT no warning, since the ratio is meaningless rather than alarming import warnings as warnings_module with warnings_module.catch_warnings(): diff --git a/tests/unit/easydynamics/utils/test_posterior_plotting.py b/tests/unit/easydynamics/utils/test_posterior_plotting.py index 296b60911..73ccfd31e 100644 --- a/tests/unit/easydynamics/utils/test_posterior_plotting.py +++ b/tests/unit/easydynamics/utils/test_posterior_plotting.py @@ -66,7 +66,7 @@ def test_one_dimensional_draws_raise(self): plot_trace(draws=np.zeros(10), names=['a']) def test_labels_carry_units(self, draws): - # WHEN + # THEN fig = plot_trace(draws=draws, names=['a', 'b', 'c'], units=['meV', 'm^2/s', '']) # EXPECT the real units are shown, and an empty one is skipped @@ -94,7 +94,7 @@ def test_mismatched_names_raise(self, draws): plot_corner(draws=draws, names=['a']) def test_diagonal_panel_is_labelled_as_counts(self, draws): - # WHEN + # THEN fig = plot_corner(draws=draws, names=['a', 'b', 'c']) # EXPECT the top-left panel says what its vertical axis actually is. It is a histogram, so @@ -102,7 +102,7 @@ def test_diagonal_panel_is_labelled_as_counts(self, draws): assert fig.axes[0].get_ylabel() == 'counts' def test_units_are_appended_to_labels(self, draws): - # WHEN + # THEN fig = plot_corner(draws=draws, names=['a', 'b', 'c'], units=['meV', '', 'dimensionless']) # EXPECT the real unit is shown, and empty or dimensionless ones are skipped @@ -176,7 +176,7 @@ def test_band_widens_with_the_credible_interval(self): assert wide_span > narrow_span def test_axis_labels_are_set_when_given(self): - # WHEN + # THEN fig = plot_posterior_predictive( x=np.zeros(4), y=np.zeros(4), @@ -196,6 +196,7 @@ def test_shared_exponent_is_folded_into_the_label(self): # on top of the axis label draws = np.random.default_rng(0).normal(size=(200, 2)) * 1e-8 + 1.15e-8 + # THEN fig = plot_corner(draws=draws, names=['D', 'scale'], units=['m^2/s', '']) # EXPECT the exponent and the unit share one parenthetical, and the overlapping offset From fbfedd626b31dff4184841e09b93de6291877494 Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Sun, 16 Aug 2026 22:32:18 +0200 Subject: [PATCH 18/29] Give the sampler its own test file Tests were split by feature rather than by the file they exercise, so posterior_sampling.py had no test file of its own and Analysis1d had two. The sampler's tests now live in test_posterior_sampling.py under one TestPosteriorSampler, with the old class names as section banners, and the four tests that are really about Analysis1d's cached fitter move into TestAnalysis1d. No test changed what it does; the same 31 + 4 tests run as before. Co-Authored-By: Claude Opus 5 (1M context) --- .../easydynamics/analysis/test_analysis1d.py | 72 +++++++++++++++++++ ...bayesian.py => test_posterior_sampling.py} | 68 +++++++----------- 2 files changed, 98 insertions(+), 42 deletions(-) rename tests/unit/easydynamics/analysis/{test_analysis1d_bayesian.py => test_posterior_sampling.py} (93%) 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_analysis1d_bayesian.py b/tests/unit/easydynamics/analysis/test_posterior_sampling.py similarity index 93% rename from tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py rename to tests/unit/easydynamics/analysis/test_posterior_sampling.py index 6b5d7a5fe..ee4c1e03c 100644 --- a/tests/unit/easydynamics/analysis/test_analysis1d_bayesian.py +++ b/tests/unit/easydynamics/analysis/test_posterior_sampling.py @@ -1,7 +1,10 @@ # SPDX-FileCopyrightText: 2026 EasyScience contributors # SPDX-License-Identifier: BSD-3-Clause -"""Unit tests for Bayesian sampling on Analysis1d, with the EasyScience Sampler mocked out.""" +""" +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 @@ -76,42 +79,11 @@ def analysis(): return make_analysis() -class TestFitterExposure: - def test_fitter_is_built_lazily_and_cached(self, analysis): - # THEN - fitter = analysis.fitter - - # EXPECT - assert fitter is analysis.fitter - assert fitter.fit_object is analysis - - def test_fitter_is_rebuilt_when_the_sample_model_changes(self, analysis): - # WHEN - original = analysis.fitter - - # THEN - analysis.sample_model = SampleModel(components=Gaussian(area=1.0)) - - # EXPECT - assert analysis.fitter is not original - - def test_minimizer_can_be_switched_through_the_fitter(self, analysis): - # THEN - analysis.fitter.switch_minimizer(AvailableMinimizers.Bumps) - - # EXPECT - assert analysis.fitter.minimizer.enum == AvailableMinimizers.Bumps - - def test_fit_uses_the_persistent_fitter(self, analysis): - # THEN - result = analysis.fit() - - # EXPECT - assert result is analysis._fit_result - assert np.isfinite(result.reduced_chi2) - +class TestPosteriorSampler: + ############# + # Bounds pre-flight + ############# -class TestBoundsPreflight: def test_sampling_refuses_unbounded_parameters(self, analysis): # THEN EXPECT with pytest.raises(ValueError, match='finite bounds'): @@ -136,8 +108,10 @@ def test_suggest_bounds_covers_the_free_parameters(self, analysis): # EXPECT assert len(suggestions) == len(analysis.get_free_parameters()) + ############# + # Sampling + ############# -class TestSamplePosterior: def test_restores_parameter_values_and_minimizer(self, analysis): # WHEN bound_all(analysis) @@ -244,8 +218,10 @@ def test_does_not_warn_when_the_posterior_is_well_inside(self, analysis): with warnings_as_errors(): analysis.bayesian.sample(samples=10) + ############# + # Parameter subsets + ############# -class TestParameterSubset: def test_holds_other_parameters_fixed_during_the_run(self, analysis): # WHEN bound_all(analysis) @@ -299,8 +275,10 @@ def test_empty_parameter_list_raises(self, analysis): with pytest.raises(ValueError, match='at least one parameter'): analysis.bayesian.sample(samples=10, parameters=[]) + ############# + # Sampler caching + ############# -class TestSamplerCaching: def test_sampler_is_reused_between_runs(self, analysis): # WHEN bound_all(analysis) @@ -344,8 +322,10 @@ def test_binds_the_same_data_the_fit_uses(self, analysis): assert np.array_equal(args[2], expected_y) assert np.array_equal(kwargs['weights'], expected_w) + ############# + # Extending and persistence + ############# -class TestExtendAndPersistence: def test_extend_without_a_chain_raises(self, analysis): # THEN EXPECT with pytest.raises(RuntimeError, match='No chain to extend'): @@ -401,8 +381,10 @@ def test_load_without_a_sidecar_warns(self, analysis, tmp_path): with pytest.warns(UserWarning, match='No parameter-name sidecar'): analysis.bayesian.load(str(tmp_path / 'missing')) + ############# + # Results + ############# -class TestResults: def test_summary_without_sampling_raises(self, analysis): # THEN EXPECT with pytest.raises(RuntimeError, match='No posterior samples yet'): @@ -447,8 +429,10 @@ def test_median_without_sampling_raises(self, analysis): with pytest.raises(RuntimeError, match='No posterior samples yet'): analysis.bayesian.set_parameters_to_median() + ############# + # Plots + ############# -class TestPlots: def test_predictive_rejects_a_bad_draw_count(self, analysis): # WHEN bound_all(analysis) From c6d3e1b2968be65ac93e3952cea19783802a946d Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Sun, 16 Aug 2026 22:52:43 +0200 Subject: [PATCH 19/29] Put each test in the file of the class it exercises Analysis and ParameterAnalysis each had a second test file, and the sampler had none of its own. The sampler's tests, whichever analysis drives them, now live in test_posterior_sampling.py under TestPosteriorSampler and TestMultiQPosteriorSampler; the fitter, chain parameter and label tests move into TestAnalysis and TestParameterAnalysis. Old class names became section banners. The multi-Q and ParameterAnalysis helpers keep distinct names in the merged file, since their signatures differ from the single-Q ones. The same 1660 tests run as before. Co-Authored-By: Claude Opus 5 (1M context) --- .../easydynamics/analysis/test_analysis.py | 148 +++++ .../analysis/test_analysis_bayesian.py | 545 --------------- .../analysis/test_parameter_analysis.py | 277 ++++++++ .../test_parameter_analysis_bayesian.py | 380 ----------- .../analysis/test_posterior_sampling.py | 623 +++++++++++++++++- 5 files changed, 1023 insertions(+), 950 deletions(-) delete mode 100644 tests/unit/easydynamics/analysis/test_analysis_bayesian.py delete mode 100644 tests/unit/easydynamics/analysis/test_parameter_analysis_bayesian.py diff --git a/tests/unit/easydynamics/analysis/test_analysis.py b/tests/unit/easydynamics/analysis/test_analysis.py index fdedd845e..b8bda4f69 100644 --- a/tests/unit/easydynamics/analysis/test_analysis.py +++ b/tests/unit/easydynamics/analysis/test_analysis.py @@ -8,6 +8,8 @@ import pytest import scipp as sc +import easydynamics as edyn +import easydynamics.sample_model as sm from easydynamics.analysis.analysis import Analysis from easydynamics.experiment import Experiment from easydynamics.sample_model import InstrumentModel @@ -15,6 +17,8 @@ from easydynamics.sample_model.components.gaussian import Gaussian from easydynamics.settings.convolution_settings import ConvolutionSettings +Q_VALUES = [0.5, 1.0, 1.5] + class TestAnalysis: @pytest.fixture @@ -70,6 +74,33 @@ def analysis_single_Q(self): extra_parameters=None, ) + @pytest.fixture + def multi_q_analysis(self): + # Three Q indices sharing one Gaussian, so the per-Q parameter copies collide by name. + energy_values = np.linspace(-5.0, 5.0, 15) + rows = [2.0 * np.exp(-0.5 * (energy_values / (0.8 + 0.4 * q**2)) ** 2) for q in Q_VALUES] + observed = np.vstack(rows) + experiment = Experiment( + data=sc.DataArray( + data=sc.array( + dims=['Q', 'energy'], + values=observed, + variances=np.full_like(observed, 0.01), + ), + coords={ + 'Q': sc.array(dims=['Q'], values=Q_VALUES, unit='1/Angstrom'), + 'energy': sc.array(dims=['energy'], values=energy_values, unit='meV'), + }, + ) + ) + + return Analysis( + display_name='TestMultiQ', + experiment=experiment, + sample_model=SampleModel(components=Gaussian(area=2.0, width=1.0)), + instrument_model=InstrumentModel(), + ) + def test_init(self, analysis): # WHEN THEN @@ -1141,3 +1172,120 @@ def test_repr(self, analysis): assert 'Analysis' in repr_str assert 'display_name=' in repr_str assert 'n_analyses=' in repr_str + + ############# + # Chain parameters and labels + ############# + + def test_union_covers_every_q_index(self, multi_q_analysis): + # THEN + parameters = multi_q_analysis._chain_parameters() + + # EXPECT one copy of each per-Q parameter, with no duplicates + assert len(parameters) == sum( + len(a.get_free_parameters()) for a in multi_q_analysis.analysis_list + ) + assert len({p.unique_name for p in parameters}) == len(parameters) + + def test_labels_are_qualified_by_q_index(self, multi_q_analysis): + # THEN + labels = [ + multi_q_analysis._parameter_labels().label(p) + for p in multi_q_analysis._chain_parameters() + ] + + # EXPECT every per-Q copy is distinguishable, which the bare name would not be + assert len(set(labels)) == len(labels) + assert 'Gaussian width (Q_index=0)' in labels + assert 'Gaussian width (Q_index=2)' in labels + + def test_bare_names_would_collide(self, multi_q_analysis): + # THEN + names = [p.name for p in multi_q_analysis._chain_parameters()] + + # EXPECT the collision the Q-qualified label exists to solve + assert len(set(names)) < len(names) + + ############# + # Parameter label edge cases + ############# + + def test_single_q_analysis_keeps_plain_names(self): + # WHEN there is only one Q index, nothing needs disambiguating + energy_values = np.linspace(-5.0, 5.0, 15) + intensity = 2.0 * np.exp(-0.5 * (energy_values / 1.2) ** 2) + experiment = edyn.Experiment( + data=sc.DataArray( + data=sc.array( + dims=['Q', 'energy'], + values=intensity[None, :], + variances=np.full_like(intensity, 0.01)[None, :], + ), + coords={ + 'Q': sc.array(dims=['Q'], values=[1.0], unit='1/Angstrom'), + 'energy': sc.array(dims=['energy'], values=energy_values, unit='meV'), + }, + ) + ) + analysis = edyn.Analysis( + display_name='SingleQ', + experiment=experiment, + sample_model=sm.SampleModel(components=sm.Gaussian(area=2.0, width=1.0)), + instrument_model=sm.InstrumentModel(), + ) + + # THEN + labels = [analysis._parameter_labels().label(p) for p in analysis._chain_parameters()] + + # EXPECT the short form, not 'Gaussian width (Q_index=0)' + assert 'Gaussian width' in labels + assert not any('Q_index=' in label for label in labels) + + def test_parameter_from_outside_the_analysis_keeps_its_name(self, multi_q_analysis): + # WHEN a parameter belongs to no Q index of this analysis + from easyscience.variable import Parameter + + stranger = Parameter(name='Gaussian width', value=1.0) + + # EXPECT it is returned unqualified rather than mislabelled + assert multi_q_analysis._parameter_labels().label(stranger) == 'Gaussian width' + + def test_a_parameter_shared_across_q_is_not_tied_to_one_index(self): + # WHEN a diffusion model contributes global parameters, the same objects appear at every Q + energy_values = np.linspace(-5.0, 5.0, 15) + rows = [2.0 * np.exp(-0.5 * (energy_values / 1.2) ** 2) for _ in Q_VALUES] + observed = np.vstack(rows) + experiment = edyn.Experiment( + data=sc.DataArray( + data=sc.array( + dims=['Q', 'energy'], + values=observed, + variances=np.full_like(observed, 0.01), + ), + coords={ + 'Q': sc.array(dims=['Q'], values=Q_VALUES, unit='1/Angstrom'), + 'energy': sc.array(dims=['energy'], values=energy_values, unit='meV'), + }, + ) + ) + analysis = edyn.Analysis( + display_name='Shared', + experiment=experiment, + sample_model=sm.SampleModel( + components=sm.ComponentCollection(components=[sm.DeltaFunction(area=0.2)]), + diffusion_models=sm.BrownianTranslationalDiffusion( + name='Brownian', diffusion_coefficient=2.4e-9, scale=0.5 + ), + ), + instrument_model=sm.InstrumentModel(), + ) + + # THEN + owners = analysis._parameter_owner_index() + shared = [p for p in analysis._chain_parameters() if p.unique_name not in owners] + + # EXPECT the shared parameters are left out of the owner map, since no single Q owns them, + # and so keep their plain names rather than being labelled with an arbitrary Q + assert shared, 'expected the diffusion model to contribute parameters shared across Q' + for parameter in shared: + assert analysis._parameter_labels().label(parameter) == parameter.name diff --git a/tests/unit/easydynamics/analysis/test_analysis_bayesian.py b/tests/unit/easydynamics/analysis/test_analysis_bayesian.py deleted file mode 100644 index c552fd7b4..000000000 --- a/tests/unit/easydynamics/analysis/test_analysis_bayesian.py +++ /dev/null @@ -1,545 +0,0 @@ -# SPDX-FileCopyrightText: 2026 EasyScience contributors -# SPDX-License-Identifier: BSD-3-Clause - -"""Unit tests for Bayesian sampling on the 2D Analysis, with the EasyScience Sampler mocked out.""" - -from types import SimpleNamespace -from unittest.mock import MagicMock -from unittest.mock import patch - -import matplotlib as mpl -import numpy as np -import pytest -import scipp as sc - -mpl.use('Agg') - -import easydynamics as edyn -import easydynamics.sample_model as sm - -SAMPLER_PATH = 'easydynamics.analysis.posterior_sampling.Sampler' -Q_VALUES = [0.5, 1.0, 1.5] - - -def make_analysis(): - energy_values = np.linspace(-5.0, 5.0, 15) - rows = [2.0 * np.exp(-0.5 * (energy_values / (0.8 + 0.4 * q**2)) ** 2) for q in Q_VALUES] - observed = np.vstack(rows) - experiment = edyn.Experiment( - data=sc.DataArray( - data=sc.array( - dims=['Q', 'energy'], - values=observed, - variances=np.full_like(observed, 0.01), - ), - coords={ - 'Q': sc.array(dims=['Q'], values=Q_VALUES, unit='1/Angstrom'), - 'energy': sc.array(dims=['energy'], values=energy_values, unit='meV'), - }, - ) - ) - return edyn.Analysis( - display_name='TestMultiQ', - experiment=experiment, - sample_model=sm.SampleModel(components=sm.Gaussian(area=2.0, width=1.0)), - instrument_model=sm.InstrumentModel(), - ) - - -def bound_all(analysis, half_width=5.0): - for parameter in analysis._chain_parameters(): - parameter.min = float(parameter.value) - half_width - parameter.max = float(parameter.value) + half_width - - -def fake_results(parameters, n_draws=50): - draws = np.tile([float(p.value) for p in parameters], (n_draws, 1)) - return SimpleNamespace( - draws=draws, - param_names=[p.unique_name for p in parameters], - logp=np.zeros(n_draws), - state=MagicMock(Ngen=10, Npop=4), - ) - - -@pytest.fixture -def analysis(): - return make_analysis() - - -class TestChainParameters: - def test_union_covers_every_q_index(self, analysis): - # THEN - parameters = analysis._chain_parameters() - - # EXPECT one copy of each per-Q parameter, with no duplicates - assert len(parameters) == sum(len(a.get_free_parameters()) for a in analysis.analysis_list) - assert len({p.unique_name for p in parameters}) == len(parameters) - - def test_labels_are_qualified_by_q_index(self, analysis): - # THEN - labels = [analysis._parameter_labels().label(p) for p in analysis._chain_parameters()] - - # EXPECT every per-Q copy is distinguishable, which the bare name would not be - assert len(set(labels)) == len(labels) - assert 'Gaussian width (Q_index=0)' in labels - assert 'Gaussian width (Q_index=2)' in labels - - def test_bare_names_would_collide(self, analysis): - # THEN - names = [p.name for p in analysis._chain_parameters()] - - # EXPECT the collision the Q-qualified label exists to solve - assert len(set(names)) < len(names) - - -class TestBoundsPreflight: - def test_sampling_refuses_unbounded_parameters(self, analysis): - # THEN EXPECT - with pytest.raises(ValueError, match='finite bounds'): - analysis.bayesian.sample(fit_method='simultaneous', samples=10) - - def test_error_names_parameters_by_q_index(self, analysis): - # THEN EXPECT - with pytest.raises(ValueError, match=r'Gaussian width \(Q_index=0\)'): - analysis.bayesian.check_bounds() - - def test_suggest_bounds_labels_every_q(self, analysis): - # THEN - suggestions = analysis.bayesian.suggest_bounds() - - # EXPECT - labels = [s.label for s in suggestions] - assert len(set(labels)) == len(labels) - assert 'Gaussian area (Q_index=1)' in labels - - -class TestSimultaneousSampling: - def test_binds_one_dataset_per_q_index(self, analysis): - # WHEN - bound_all(analysis) - parameters = analysis._chain_parameters() - - # THEN - with patch(SAMPLER_PATH) as sampler_class: - sampler_class.return_value.sample.return_value = fake_results(parameters) - analysis.bayesian.sample(fit_method='simultaneous', samples=10) - - # EXPECT - args, kwargs = sampler_class.call_args - assert len(args[1]) == len(Q_VALUES) - assert len(args[2]) == len(Q_VALUES) - assert len(kwargs['weights']) == len(Q_VALUES) - - def test_returns_a_single_result(self, analysis): - # WHEN - bound_all(analysis) - parameters = analysis._chain_parameters() - - # THEN - with patch(SAMPLER_PATH) as sampler_class: - expected = fake_results(parameters) - sampler_class.return_value.sample.return_value = expected - returned = analysis.bayesian.sample(fit_method='simultaneous', samples=10) - - # EXPECT - assert returned is expected - assert analysis.bayesian.results is expected - - def test_summary_is_labelled_by_q_index(self, analysis): - # WHEN - bound_all(analysis) - parameters = analysis._chain_parameters() - - # THEN - with patch(SAMPLER_PATH) as sampler_class: - sampler_class.return_value.sample.return_value = fake_results(parameters) - analysis.bayesian.sample(fit_method='simultaneous', samples=10) - - # EXPECT - names = [entry.name for entry in analysis.bayesian.summary()] - assert len(set(names)) == len(names) - assert all('Q_index=' in name for name in names) - - def test_refreshes_every_convolver_before_sampling(self, analysis): - # WHEN - bound_all(analysis) - parameters = analysis._chain_parameters() - - # THEN - with patch(SAMPLER_PATH) as sampler_class: - sampler_class.return_value.sample.return_value = fake_results(parameters) - for analysis1d in analysis.analysis_list: - analysis1d._convolver_is_dirty = True - analysis.bayesian.sample(fit_method='simultaneous', samples=10) - - # EXPECT the sampler sees the same prepared convolvers a simultaneous fit would - assert all(not a._convolver_is_dirty for a in analysis.analysis_list) - - def test_uses_a_multifitter(self, analysis): - # WHEN - from easyscience.fitting.multi_fitter import MultiFitter - - # EXPECT - assert isinstance(analysis.fitter, MultiFitter) - assert len(analysis.fitter.fit_object) == len(Q_VALUES) - - -class TestIndependentSampling: - def test_returns_one_result_per_q_index(self, analysis): - # WHEN - for analysis1d in analysis.analysis_list: - for parameter in analysis1d.get_free_parameters(): - parameter.min = float(parameter.value) - 5.0 - parameter.max = float(parameter.value) + 5.0 - - # THEN - with patch(SAMPLER_PATH) as sampler_class: - sampler_class.return_value.sample.side_effect = lambda **_k: fake_results( - analysis.analysis_list[0].get_free_parameters() - ) - results = analysis.bayesian.sample(fit_method='independent', samples=10) - - # EXPECT - assert isinstance(results, list) - assert len(results) == len(Q_VALUES) - - def test_single_q_index_returns_one_result(self, analysis): - # WHEN - target = analysis.analysis_list[1] - for parameter in target.get_free_parameters(): - parameter.min = float(parameter.value) - 5.0 - parameter.max = float(parameter.value) + 5.0 - - # THEN - with patch(SAMPLER_PATH) as sampler_class: - sampler_class.return_value.sample.side_effect = lambda **_k: fake_results( - target.get_free_parameters() - ) - result = analysis.bayesian.sample(fit_method='independent', Q_index=1, samples=10) - - # EXPECT - assert not isinstance(result, list) - assert result is target.bayesian.results - - def test_invalid_q_index_raises(self, analysis): - # THEN EXPECT - with pytest.raises((ValueError, IndexError)): - analysis.bayesian.sample(fit_method='independent', Q_index=99, samples=10) - - -class TestSamplePosteriorValidation: - def test_invalid_fit_method_raises(self, analysis): - # THEN EXPECT - with pytest.raises(ValueError, match='Invalid fit method'): - analysis.bayesian.sample(fit_method='nonsense') - - def test_missing_q_values_raises(self): - # WHEN - analysis = edyn.Analysis(display_name='Empty') - - # THEN EXPECT - with pytest.raises(ValueError, match='No Q values available'): - analysis.bayesian.sample() - - -class TestPredictivePlot: - def test_predictive_is_not_supported_for_multiple_datasets(self, analysis): - # WHEN - bound_all(analysis) - parameters = analysis._chain_parameters() - - with patch(SAMPLER_PATH) as sampler_class: - sampler_class.return_value.sample.return_value = fake_results(parameters) - analysis.bayesian.sample(fit_method='simultaneous', samples=10) - - # THEN EXPECT - with pytest.raises(NotImplementedError, match='single dataset only'): - analysis.bayesian.plot_posterior_predictive() - - -class TestParameterLabelEdgeCases: - def test_single_q_analysis_keeps_plain_names(self): - # WHEN there is only one Q index, nothing needs disambiguating - energy_values = np.linspace(-5.0, 5.0, 15) - intensity = 2.0 * np.exp(-0.5 * (energy_values / 1.2) ** 2) - experiment = edyn.Experiment( - data=sc.DataArray( - data=sc.array( - dims=['Q', 'energy'], - values=intensity[None, :], - variances=np.full_like(intensity, 0.01)[None, :], - ), - coords={ - 'Q': sc.array(dims=['Q'], values=[1.0], unit='1/Angstrom'), - 'energy': sc.array(dims=['energy'], values=energy_values, unit='meV'), - }, - ) - ) - analysis = edyn.Analysis( - display_name='SingleQ', - experiment=experiment, - sample_model=sm.SampleModel(components=sm.Gaussian(area=2.0, width=1.0)), - instrument_model=sm.InstrumentModel(), - ) - - # THEN - labels = [analysis._parameter_labels().label(p) for p in analysis._chain_parameters()] - - # EXPECT the short form, not 'Gaussian width (Q_index=0)' - assert 'Gaussian width' in labels - assert not any('Q_index=' in label for label in labels) - - def test_parameter_from_outside_the_analysis_keeps_its_name(self, analysis): - # WHEN a parameter belongs to no Q index of this analysis - from easyscience.variable import Parameter - - stranger = Parameter(name='Gaussian width', value=1.0) - - # EXPECT it is returned unqualified rather than mislabelled - assert analysis._parameter_labels().label(stranger) == 'Gaussian width' - - -class TestIndependentSamplingDiscoverability: - def test_operations_needing_one_chain_point_at_the_per_q_chains(self, analysis): - # WHEN sampling independently, the chains live on the Analysis1d objects, not here - remaining = iter(analysis.analysis_list) - for analysis1d in analysis.analysis_list: - for parameter in analysis1d.get_free_parameters(): - parameter.min = float(parameter.value) - 5.0 - parameter.max = float(parameter.value) + 5.0 - - with patch(SAMPLER_PATH) as sampler_class: - sampler_class.return_value.sample.side_effect = lambda **_k: fake_results( - next(remaining).get_free_parameters() - ) - analysis.bayesian.sample(fit_method='independent', samples=10) - - # THEN EXPECT anything that genuinely needs a single chain says where the chains - # actually are, rather than claiming none exist - with pytest.raises(RuntimeError, match='analysis_list'): - analysis.bayesian.plot_posterior_predictive() - - def test_untouched_analysis_still_reports_no_samples(self, analysis): - # THEN EXPECT the plain message when nothing has been sampled anywhere - with pytest.raises(RuntimeError, match='No posterior samples yet'): - analysis.bayesian.summary() - - -class TestSharedParameterLabels: - def test_a_parameter_shared_across_q_is_not_tied_to_one_index(self): - # WHEN a diffusion model contributes global parameters, the same objects appear at every Q - energy_values = np.linspace(-5.0, 5.0, 15) - rows = [2.0 * np.exp(-0.5 * (energy_values / 1.2) ** 2) for _ in Q_VALUES] - observed = np.vstack(rows) - experiment = edyn.Experiment( - data=sc.DataArray( - data=sc.array( - dims=['Q', 'energy'], - values=observed, - variances=np.full_like(observed, 0.01), - ), - coords={ - 'Q': sc.array(dims=['Q'], values=Q_VALUES, unit='1/Angstrom'), - 'energy': sc.array(dims=['energy'], values=energy_values, unit='meV'), - }, - ) - ) - analysis = edyn.Analysis( - display_name='Shared', - experiment=experiment, - sample_model=sm.SampleModel( - components=sm.ComponentCollection(components=[sm.DeltaFunction(area=0.2)]), - diffusion_models=sm.BrownianTranslationalDiffusion( - name='Brownian', diffusion_coefficient=2.4e-9, scale=0.5 - ), - ), - instrument_model=sm.InstrumentModel(), - ) - - # THEN - owners = analysis._parameter_owner_index() - shared = [p for p in analysis._chain_parameters() if p.unique_name not in owners] - - # EXPECT the shared parameters are left out of the owner map, since no single Q owns them, - # and so keep their plain names rather than being labelled with an arbitrary Q - assert shared, 'expected the diffusion model to contribute parameters shared across Q' - for parameter in shared: - assert analysis._parameter_labels().label(parameter) == parameter.name - - -class TestAggregatingPerQChains: - def _sample_independently(self, analysis): - for analysis1d in analysis.analysis_list: - for parameter in analysis1d.get_free_parameters(): - parameter.min = float(parameter.value) - 5.0 - parameter.max = float(parameter.value) + 5.0 - - # The Q indices sample in order, and each must get a chain over its own parameters. - remaining = iter(analysis.analysis_list) - - with patch(SAMPLER_PATH) as sampler_class: - sampler_class.return_value.sample.side_effect = lambda **_k: fake_results( - next(remaining).get_free_parameters() - ) - analysis.bayesian.sample(fit_method='independent', samples=10) - - def test_posterior_results_holds_one_chain_per_q(self, analysis): - # WHEN - self._sample_independently(analysis) - - # EXPECT - assert len(analysis.bayesian.results_per_q) == len(Q_VALUES) - assert all(result is not None for result in analysis.bayesian.results_per_q) - - def test_posterior_results_is_none_before_sampling(self, analysis): - # EXPECT - assert analysis.bayesian.results_per_q is None - - def test_summary_gathers_every_q(self, analysis): - # WHEN - self._sample_independently(analysis) - - # THEN - summary = analysis.bayesian.summary() - - # EXPECT one entry per free parameter per Q, each labelled by its Q index - expected = sum(len(a.get_free_parameters()) for a in analysis.analysis_list) - names = [entry.name for entry in summary] - assert len(summary) == expected - assert len(set(names)) == len(names) - assert all('Q_index=' in name for name in names) - - def test_median_applies_each_chain_to_its_own_q(self, analysis): - # WHEN - self._sample_independently(analysis) - - # THEN - changed = analysis.bayesian.set_parameters_to_median() - - # EXPECT every Q's parameters are set, from that Q's own chain - expected = sum(len(a.get_free_parameters()) for a in analysis.analysis_list) - assert len(changed) == expected - - def test_corner_plots_one_q_at_a_time(self, analysis): - # WHEN each Q was sampled separately, no draw pairs one Q with another, so a corner plot - # can only show one chain at a time - self._sample_independently(analysis) - - # THEN - figure = analysis.bayesian.plot_corner(Q_index=1) - - # EXPECT that Q's own chain, not a combination across Q - n_parameters = len(analysis.analysis_list[1].get_free_parameters()) - assert len(figure.axes) == n_parameters**2 - - def test_corner_offers_a_slider_in_a_notebook(self, analysis): - # WHEN - self._sample_independently(analysis) - - # THEN - with patch('easydynamics.analysis.posterior_sampling._in_notebook', return_value=True): - widget = analysis.bayesian.plot_corner() - - # EXPECT a slider over the sampled Q indices, and a panel that actually holds a figure. - # The obvious way to build this captures nothing and leaves the panel blank beside the - # slider, so an empty panel is the regression worth guarding. Which mime type arrives - # depends on the environment: a live kernel renders a PNG, plain pytest only the repr. - # The figure comes first and the slider sits under it, where plopp puts its controls. - panel, slider = widget.children - assert list(slider.options) == list(range(len(Q_VALUES))) - assert panel.outputs, 'the initial chain was not drawn' - assert 'Figure' in str(panel.outputs[0]['data']) - - slider.value = 2 - assert panel.outputs, 'changing Q did not redraw' - assert 'Figure' in str(panel.outputs[0]['data']) - - def test_corner_without_a_notebook_or_q_index_says_what_to_do(self, analysis): - # WHEN - self._sample_independently(analysis) - - # THEN EXPECT it names the sampled Q indices rather than just refusing - with ( - patch('easydynamics.analysis.posterior_sampling._in_notebook', return_value=False), - pytest.raises(RuntimeError, match=r'sampled Q indices are \[0, 1, 2\]'), - ): - analysis.bayesian.plot_corner() - - def test_the_slider_only_offers_q_indices_that_were_sampled(self, analysis): - # WHEN only one Q index is sampled - target = analysis.analysis_list[2] - for parameter in target.get_free_parameters(): - parameter.min = float(parameter.value) - 5.0 - parameter.max = float(parameter.value) + 5.0 - - with patch(SAMPLER_PATH) as sampler_class: - sampler_class.return_value.sample.side_effect = lambda **_k: fake_results( - target.get_free_parameters() - ) - analysis.bayesian.sample(fit_method='independent', Q_index=2, samples=10) - - # THEN - with patch('easydynamics.analysis.posterior_sampling._in_notebook', return_value=True): - widget = analysis.bayesian.plot_corner() - - # EXPECT the slider cannot land on a Q with nothing to draw - assert list(widget.children[1].options) == [2] - - def test_trace_points_at_the_individual_chains(self, analysis): - # WHEN - self._sample_independently(analysis) - - # THEN EXPECT - with pytest.raises(RuntimeError, match='no single trace'): - analysis.bayesian.plot_trace() - - def test_a_simultaneous_chain_still_takes_precedence(self, analysis): - # WHEN a simultaneous run follows an independent one - self._sample_independently(analysis) - bound_all(analysis) - parameters = analysis._chain_parameters() - - with patch(SAMPLER_PATH) as sampler_class: - sampler_class.return_value.sample.return_value = fake_results(parameters) - analysis.bayesian.sample(fit_method='simultaneous', samples=10) - - # EXPECT the single chain is summarized, not the stale per-Q ones - assert len(analysis.bayesian.summary()) == len(parameters) - analysis.bayesian.plot_corner() - - def test_only_the_sampled_q_indices_are_gathered(self, analysis): - # WHEN just one Q index is sampled - target = analysis.analysis_list[1] - for parameter in target.get_free_parameters(): - parameter.min = float(parameter.value) - 5.0 - parameter.max = float(parameter.value) + 5.0 - - with patch(SAMPLER_PATH) as sampler_class: - sampler_class.return_value.sample.side_effect = lambda **_k: fake_results( - target.get_free_parameters() - ) - analysis.bayesian.sample(fit_method='independent', Q_index=1, samples=10) - - # THEN - summary = analysis.bayesian.summary() - - # EXPECT the unsampled Q indices are passed over rather than breaking the aggregation - assert len(summary) == len(target.get_free_parameters()) - assert all('Q_index=1' in entry.name for entry in summary) - assert len(analysis.bayesian.set_parameters_to_median()) == len( - target.get_free_parameters() - ) - - def test_a_simultaneous_chain_serves_the_median_and_the_trace(self, analysis): - # WHEN - bound_all(analysis) - parameters = analysis._chain_parameters() - - with patch(SAMPLER_PATH) as sampler_class: - sampler_class.return_value.sample.return_value = fake_results(parameters) - analysis.bayesian.sample(fit_method='simultaneous', samples=10) - - # EXPECT both come from the single chain, with no per-Q gathering involved - assert len(analysis.bayesian.set_parameters_to_median()) == len(parameters) - assert len(analysis.bayesian.plot_trace().axes) == len(parameters) + 1 diff --git a/tests/unit/easydynamics/analysis/test_parameter_analysis.py b/tests/unit/easydynamics/analysis/test_parameter_analysis.py index 031f813cf..1b6b80178 100644 --- a/tests/unit/easydynamics/analysis/test_parameter_analysis.py +++ b/tests/unit/easydynamics/analysis/test_parameter_analysis.py @@ -8,7 +8,10 @@ import numpy as np import pytest import scipp as sc +from easyscience.fitting.multi_fitter import MultiFitter +import easydynamics as edyn +import easydynamics.sample_model as sm from easydynamics.analysis.analysis import Analysis from easydynamics.analysis.fit_binding import FitBinding from easydynamics.analysis.parameter_analysis import ParameterAnalysis @@ -19,6 +22,8 @@ ) from easydynamics.utils.fit_target import FitTarget +Q = np.array([0.5, 0.8, 1.1, 1.4, 1.7, 2.0]) + def make_target(dataset_key, function, label, x_unit=None, y_unit=None, name='value'): """Build a FitTarget for mocking FitBinding.get_targets in tests.""" @@ -32,6 +37,51 @@ def make_target(dataset_key, function, label, x_unit=None, y_unit=None, name='va ) +def make_dataset(): + widths = 0.10 + 0.35 * Q + areas = 2.0 - 0.3 * Q + return sc.Dataset({ + 'Lorentzian width': sc.DataArray( + data=sc.array( + dims=['Q'], values=widths, variances=np.full_like(widths, 1e-4), unit='meV' + ), + coords={'Q': sc.array(dims=['Q'], values=Q, unit='1/angstrom')}, + ), + 'Lorentzian area': sc.DataArray( + data=sc.array( + dims=['Q'], values=areas, variances=np.full_like(areas, 4e-4), unit='meV' + ), + coords={'Q': sc.array(dims=['Q'], values=Q, unit='1/angstrom')}, + ), + }) + + +def make_analysis(two_bindings=True): + bindings = [ + edyn.FitBinding( + model=sm.Polynomial( + coefficients=[0.1, 0.35], x_unit='1/angstrom', y_unit='meV', name='Width line' + ), + targets='Lorentzian width', + ) + ] + if two_bindings: + bindings.append( + edyn.FitBinding( + model=sm.Polynomial( + coefficients=[2.0, -0.3], x_unit='1/angstrom', y_unit='meV', name='Area line' + ), + targets='Lorentzian area', + ) + ) + return edyn.ParameterAnalysis(parameters=make_dataset(), bindings=bindings) + + +@pytest.fixture +def analysis(): + return make_analysis() + + class TestParameterAnalysis: @pytest.fixture def dataset(self): @@ -1150,6 +1200,233 @@ def test_repr(self, parameter_analysis): assert 'parameter_names=' in repr_str assert 'bindings=' in repr_str + ############# + # The cached fitter + ############# + + def test_fitter_is_a_cached_multifitter(self, analysis): + # EXPECT + assert isinstance(analysis.fitter, MultiFitter) + assert analysis.fitter is analysis.fitter + + def test_fit_still_returns_per_target_results(self, analysis): + # THEN + results = analysis.fit() + + # EXPECT one result per fit target, as before + assert isinstance(results, list) + assert len(results) == 2 + + def test_changing_bindings_rebuilds_the_fitter(self, analysis): + # WHEN + original = analysis.fitter + + # THEN + analysis.bindings = analysis.bindings[:1] + + # EXPECT + assert analysis.fitter is not original + + def test_changing_parameters_rebuilds_the_fitter(self, analysis): + # WHEN + original = analysis.fitter + + # THEN + analysis.parameters = make_dataset() + + # EXPECT + assert analysis.fitter is not original + + def test_changing_the_number_of_targets_rebuilds_the_fitter(self): + # WHEN a binding is edited in place so that it resolves to two targets instead of one. + # ParameterAnalysis cannot observe this, and the cached fitter would otherwise still hold + # one fit function against two datasets, which dies inside the minimizer. + binding = edyn.FitBinding( + model=sm.BrownianTranslationalDiffusion( + name='Brownian', + lorentzian_name='Lorentzian', + diffusion_coefficient=2.4e-9, + scale=0.5, + ), + targets={'width': 'Lorentzian width'}, + ) + analysis = edyn.ParameterAnalysis(parameters=make_dataset(), bindings=[binding]) + assert len(analysis.fit()) == 1 + + # THEN + binding.targets = {'width': 'Lorentzian width', 'area': 'Lorentzian area'} + + # EXPECT the fit follows the binding rather than failing on a stale fitter + assert len(analysis.fit()) == 2 + + def test_shrinking_the_targets_also_rebuilds(self): + # WHEN + binding = edyn.FitBinding( + model=sm.BrownianTranslationalDiffusion( + name='Brownian', + lorentzian_name='Lorentzian', + diffusion_coefficient=2.4e-9, + scale=0.5, + ), + targets={'width': 'Lorentzian width', 'area': 'Lorentzian area'}, + ) + analysis = edyn.ParameterAnalysis(parameters=make_dataset(), bindings=[binding]) + assert len(analysis.fit()) == 2 + + # THEN + binding.targets = {'width': 'Lorentzian width'} + + # EXPECT + assert len(analysis.fit()) == 1 + + ############# + # Chain parameters and labels + ############# + + def test_covers_every_binding_model(self, analysis): + # THEN + parameters = analysis._chain_parameters() + + # EXPECT both Polynomials contribute their two coefficients + assert len(parameters) == 4 + assert len({p.unique_name for p in parameters}) == 4 + + def test_labels_are_unique(self, analysis): + # THEN + labels = [analysis._parameter_labels().label(p) for p in analysis._chain_parameters()] + + # EXPECT + assert len(set(labels)) == len(labels) + + def test_model_name_is_not_repeated_in_the_label(self, analysis): + # WHEN a model already names its parameters after itself + + # THEN + labels = [analysis._parameter_labels().label(p) for p in analysis._chain_parameters()] + + # EXPECT no 'Width line: Width line_c0' + assert 'Width line_c0' in labels + assert not any(label.count('Width line') > 1 for label in labels) + + def test_colliding_names_are_qualified_by_model(self): + # WHEN two bindings use models whose parameters share a name + shared_name_model_a = sm.Polynomial( + coefficients=[0.1, 0.35], x_unit='1/angstrom', y_unit='meV', name='Line' + ) + shared_name_model_b = sm.Polynomial( + coefficients=[2.0, -0.3], x_unit='1/angstrom', y_unit='meV', name='Line' + ) + analysis = edyn.ParameterAnalysis( + parameters=make_dataset(), + bindings=[ + edyn.FitBinding(model=shared_name_model_a, targets='Lorentzian width'), + edyn.FitBinding(model=shared_name_model_b, targets='Lorentzian area'), + ], + ) + + # THEN + parameters = analysis._chain_parameters() + names = [p.name for p in parameters] + labels = [analysis._parameter_labels().label(p) for p in parameters] + + # EXPECT the bare names collide, and the labels resolve it + assert len(set(names)) < len(names) + assert len(set(labels)) == len(labels) + + def test_single_binding_keeps_plain_names(self): + # WHEN + analysis = make_analysis(two_bindings=False) + + # THEN + labels = [analysis._parameter_labels().label(p) for p in analysis._chain_parameters()] + + # EXPECT no model prefix, since there is nothing to disambiguate + assert labels == ['Width line_c0', 'Width line_c1'] + + def test_parameter_from_outside_the_analysis_keeps_its_name(self, analysis): + # WHEN a parameter belongs to none of the binding models + from easyscience.variable import Parameter + + stranger = Parameter(name='Width line_c0', value=1.0) + + # THEN EXPECT it is returned unqualified rather than mislabelled + assert analysis._parameter_labels().label(stranger) == 'Width line_c0' + + def test_models_without_a_display_name_fall_back_to_the_unique_name(self): + # WHEN two colliding models have no display name to tell them apart + model_a = sm.Polynomial(coefficients=[0.1, 0.35], x_unit='1/angstrom', y_unit='meV') + model_b = sm.Polynomial(coefficients=[2.0, -0.3], x_unit='1/angstrom', y_unit='meV') + analysis = edyn.ParameterAnalysis( + parameters=make_dataset(), + bindings=[ + edyn.FitBinding(model=model_a, targets='Lorentzian width'), + edyn.FitBinding(model=model_b, targets='Lorentzian area'), + ], + ) + + # THEN + labels = [analysis._parameter_labels().label(p) for p in analysis._chain_parameters()] + + # EXPECT still unambiguous, which is what matters + assert len(set(labels)) == len(labels) + + def test_colliding_names_with_distinct_models_use_the_display_name(self): + # WHEN two diffusion models are bound to different targets. Their parameters are not named + # after the model, so the names collide while the model names do not. + analysis = edyn.ParameterAnalysis( + parameters=make_dataset(), + bindings=[ + edyn.FitBinding( + model=sm.BrownianTranslationalDiffusion( + name='Diffusion A', diffusion_coefficient=2.4e-9, scale=0.5 + ), + targets={'width': 'Lorentzian width'}, + ), + edyn.FitBinding( + model=sm.BrownianTranslationalDiffusion( + name='Diffusion B', diffusion_coefficient=2.4e-9, scale=0.5 + ), + targets={'area': 'Lorentzian area'}, + ), + ], + ) + + # THEN + parameters = analysis._chain_parameters() + labels = [analysis._parameter_labels().label(p) for p in parameters] + + # EXPECT the model's name resolves the collision + assert len({p.name for p in parameters}) < len(parameters) + assert len(set(labels)) == len(labels) + assert any(label.endswith('(Diffusion A)') for label in labels) + assert any(label.endswith('(Diffusion B)') for label in labels) + + def test_ambiguous_name_owned_by_no_model_keeps_its_name(self): + # WHEN a parameter shares an ambiguous name but belongs to none of the models + from easyscience.variable import Parameter + + analysis = edyn.ParameterAnalysis( + parameters=make_dataset(), + bindings=[ + edyn.FitBinding( + model=sm.Polynomial( + coefficients=[0.1, 0.35], x_unit='1/angstrom', y_unit='meV', name='Line' + ), + targets='Lorentzian width', + ), + edyn.FitBinding( + model=sm.Polynomial( + coefficients=[2.0, -0.3], x_unit='1/angstrom', y_unit='meV', name='Line' + ), + targets='Lorentzian area', + ), + ], + ) + stranger = Parameter(name='Line_c0', value=1.0) + + # THEN EXPECT it falls back to the plain name rather than claiming an owner + assert analysis._parameter_labels().label(stranger) == 'Line_c0' + class TestParameterAnalysisWorkflows: """End-to-end fits for the standard workflows on synthetic data.""" diff --git a/tests/unit/easydynamics/analysis/test_parameter_analysis_bayesian.py b/tests/unit/easydynamics/analysis/test_parameter_analysis_bayesian.py deleted file mode 100644 index 7966e1f87..000000000 --- a/tests/unit/easydynamics/analysis/test_parameter_analysis_bayesian.py +++ /dev/null @@ -1,380 +0,0 @@ -# SPDX-FileCopyrightText: 2026 EasyScience contributors -# SPDX-License-Identifier: BSD-3-Clause - -"""Unit tests for Bayesian sampling on ParameterAnalysis, with the 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.multi_fitter import MultiFitter - -import easydynamics as edyn -import easydynamics.sample_model as sm - -SAMPLER_PATH = 'easydynamics.analysis.posterior_sampling.Sampler' -Q = np.array([0.5, 0.8, 1.1, 1.4, 1.7, 2.0]) - - -def make_dataset(): - widths = 0.10 + 0.35 * Q - areas = 2.0 - 0.3 * Q - return sc.Dataset({ - 'Lorentzian width': sc.DataArray( - data=sc.array( - dims=['Q'], values=widths, variances=np.full_like(widths, 1e-4), unit='meV' - ), - coords={'Q': sc.array(dims=['Q'], values=Q, unit='1/angstrom')}, - ), - 'Lorentzian area': sc.DataArray( - data=sc.array( - dims=['Q'], values=areas, variances=np.full_like(areas, 4e-4), unit='meV' - ), - coords={'Q': sc.array(dims=['Q'], values=Q, unit='1/angstrom')}, - ), - }) - - -def make_analysis(two_bindings=True): - bindings = [ - edyn.FitBinding( - model=sm.Polynomial( - coefficients=[0.1, 0.35], x_unit='1/angstrom', y_unit='meV', name='Width line' - ), - targets='Lorentzian width', - ) - ] - if two_bindings: - bindings.append( - edyn.FitBinding( - model=sm.Polynomial( - coefficients=[2.0, -0.3], x_unit='1/angstrom', y_unit='meV', name='Area line' - ), - targets='Lorentzian area', - ) - ) - return edyn.ParameterAnalysis(parameters=make_dataset(), bindings=bindings) - - -def bound_all(analysis, half_width=5.0): - for parameter in analysis._chain_parameters(): - parameter.min = float(parameter.value) - half_width - parameter.max = float(parameter.value) + half_width - - -def fake_results(parameters, n_draws=50): - draws = np.tile([float(p.value) for p in parameters], (n_draws, 1)) - return SimpleNamespace( - draws=draws, - param_names=[p.unique_name for p in parameters], - logp=np.zeros(n_draws), - state=MagicMock(Ngen=10, Npop=4), - ) - - -@pytest.fixture -def analysis(): - return make_analysis() - - -class TestFitterExposure: - def test_fitter_is_a_cached_multifitter(self, analysis): - # EXPECT - assert isinstance(analysis.fitter, MultiFitter) - assert analysis.fitter is analysis.fitter - - def test_fit_still_returns_per_target_results(self, analysis): - # THEN - results = analysis.fit() - - # EXPECT one result per fit target, as before - assert isinstance(results, list) - assert len(results) == 2 - - def test_changing_bindings_rebuilds_the_fitter(self, analysis): - # WHEN - original = analysis.fitter - - # THEN - analysis.bindings = analysis.bindings[:1] - - # EXPECT - assert analysis.fitter is not original - - def test_changing_parameters_rebuilds_the_fitter(self, analysis): - # WHEN - original = analysis.fitter - - # THEN - analysis.parameters = make_dataset() - - # EXPECT - assert analysis.fitter is not original - - -class TestChainParameters: - def test_covers_every_binding_model(self, analysis): - # THEN - parameters = analysis._chain_parameters() - - # EXPECT both Polynomials contribute their two coefficients - assert len(parameters) == 4 - assert len({p.unique_name for p in parameters}) == 4 - - def test_labels_are_unique(self, analysis): - # THEN - labels = [analysis._parameter_labels().label(p) for p in analysis._chain_parameters()] - - # EXPECT - assert len(set(labels)) == len(labels) - - def test_model_name_is_not_repeated_in_the_label(self, analysis): - # WHEN a model already names its parameters after itself - - # THEN - labels = [analysis._parameter_labels().label(p) for p in analysis._chain_parameters()] - - # EXPECT no 'Width line: Width line_c0' - assert 'Width line_c0' in labels - assert not any(label.count('Width line') > 1 for label in labels) - - -class TestSampling: - def test_refuses_unbounded_parameters(self, analysis): - # THEN EXPECT - with pytest.raises(ValueError, match='finite bounds'): - analysis.bayesian.sample(samples=10) - - def test_binds_one_dataset_per_target(self, analysis): - # WHEN - bound_all(analysis) - parameters = analysis._chain_parameters() - - # THEN - with patch(SAMPLER_PATH) as sampler_class: - sampler_class.return_value.sample.return_value = fake_results(parameters) - analysis.bayesian.sample(samples=10) - - # EXPECT - args, kwargs = sampler_class.call_args - assert len(args[1]) == 2 - assert len(kwargs['weights']) == 2 - - def test_summary_uses_model_qualified_labels(self, analysis): - # WHEN - bound_all(analysis) - parameters = analysis._chain_parameters() - - # THEN - with patch(SAMPLER_PATH) as sampler_class: - sampler_class.return_value.sample.return_value = fake_results(parameters) - analysis.bayesian.sample(samples=10) - - # EXPECT - names = [entry.name for entry in analysis.bayesian.summary()] - assert len(set(names)) == len(names) - assert 'Width line_c0' in names - - def test_restores_parameter_values(self, analysis): - # WHEN - bound_all(analysis) - parameters = analysis._chain_parameters() - before = [float(p.value) for p in parameters] - - # THEN - with patch(SAMPLER_PATH) as sampler_class: - - def mutate(**_kwargs): - for parameter in parameters: - parameter.value = float(parameter.value) + 1.0 - return fake_results(parameters) - - sampler_class.return_value.sample.side_effect = mutate - analysis.bayesian.sample(samples=10) - - # EXPECT - assert [float(p.value) for p in parameters] == pytest.approx(before) - - def test_missing_parameters_dataset_raises(self): - # WHEN - analysis = edyn.ParameterAnalysis() - - # THEN EXPECT - with pytest.raises(ValueError, match='No parameters Dataset'): - analysis.bayesian.sample(samples=10) - - def test_missing_bindings_raises(self): - # WHEN - analysis = edyn.ParameterAnalysis(parameters=make_dataset()) - - # THEN EXPECT - with pytest.raises(ValueError, match='No fit bindings'): - analysis.bayesian.sample(samples=10) - - -class TestParameterLabelEdgeCases: - def test_colliding_names_are_qualified_by_model(self): - # WHEN two bindings use models whose parameters share a name - shared_name_model_a = sm.Polynomial( - coefficients=[0.1, 0.35], x_unit='1/angstrom', y_unit='meV', name='Line' - ) - shared_name_model_b = sm.Polynomial( - coefficients=[2.0, -0.3], x_unit='1/angstrom', y_unit='meV', name='Line' - ) - analysis = edyn.ParameterAnalysis( - parameters=make_dataset(), - bindings=[ - edyn.FitBinding(model=shared_name_model_a, targets='Lorentzian width'), - edyn.FitBinding(model=shared_name_model_b, targets='Lorentzian area'), - ], - ) - - # THEN - parameters = analysis._chain_parameters() - names = [p.name for p in parameters] - labels = [analysis._parameter_labels().label(p) for p in parameters] - - # EXPECT the bare names collide, and the labels resolve it - assert len(set(names)) < len(names) - assert len(set(labels)) == len(labels) - - def test_single_binding_keeps_plain_names(self): - # WHEN - analysis = make_analysis(two_bindings=False) - - # THEN - labels = [analysis._parameter_labels().label(p) for p in analysis._chain_parameters()] - - # EXPECT no model prefix, since there is nothing to disambiguate - assert labels == ['Width line_c0', 'Width line_c1'] - - def test_parameter_from_outside_the_analysis_keeps_its_name(self, analysis): - # WHEN a parameter belongs to none of the binding models - from easyscience.variable import Parameter - - stranger = Parameter(name='Width line_c0', value=1.0) - - # THEN EXPECT it is returned unqualified rather than mislabelled - assert analysis._parameter_labels().label(stranger) == 'Width line_c0' - - def test_models_without_a_display_name_fall_back_to_the_unique_name(self): - # WHEN two colliding models have no display name to tell them apart - model_a = sm.Polynomial(coefficients=[0.1, 0.35], x_unit='1/angstrom', y_unit='meV') - model_b = sm.Polynomial(coefficients=[2.0, -0.3], x_unit='1/angstrom', y_unit='meV') - analysis = edyn.ParameterAnalysis( - parameters=make_dataset(), - bindings=[ - edyn.FitBinding(model=model_a, targets='Lorentzian width'), - edyn.FitBinding(model=model_b, targets='Lorentzian area'), - ], - ) - - # THEN - labels = [analysis._parameter_labels().label(p) for p in analysis._chain_parameters()] - - # EXPECT still unambiguous, which is what matters - assert len(set(labels)) == len(labels) - - def test_colliding_names_with_distinct_models_use_the_display_name(self): - # WHEN two diffusion models are bound to different targets. Their parameters are not named - # after the model, so the names collide while the model names do not. - analysis = edyn.ParameterAnalysis( - parameters=make_dataset(), - bindings=[ - edyn.FitBinding( - model=sm.BrownianTranslationalDiffusion( - name='Diffusion A', diffusion_coefficient=2.4e-9, scale=0.5 - ), - targets={'width': 'Lorentzian width'}, - ), - edyn.FitBinding( - model=sm.BrownianTranslationalDiffusion( - name='Diffusion B', diffusion_coefficient=2.4e-9, scale=0.5 - ), - targets={'area': 'Lorentzian area'}, - ), - ], - ) - - # THEN - parameters = analysis._chain_parameters() - labels = [analysis._parameter_labels().label(p) for p in parameters] - - # EXPECT the model's name resolves the collision - assert len({p.name for p in parameters}) < len(parameters) - assert len(set(labels)) == len(labels) - assert any(label.endswith('(Diffusion A)') for label in labels) - assert any(label.endswith('(Diffusion B)') for label in labels) - - def test_ambiguous_name_owned_by_no_model_keeps_its_name(self): - # WHEN a parameter shares an ambiguous name but belongs to none of the models - from easyscience.variable import Parameter - - analysis = edyn.ParameterAnalysis( - parameters=make_dataset(), - bindings=[ - edyn.FitBinding( - model=sm.Polynomial( - coefficients=[0.1, 0.35], x_unit='1/angstrom', y_unit='meV', name='Line' - ), - targets='Lorentzian width', - ), - edyn.FitBinding( - model=sm.Polynomial( - coefficients=[2.0, -0.3], x_unit='1/angstrom', y_unit='meV', name='Line' - ), - targets='Lorentzian area', - ), - ], - ) - stranger = Parameter(name='Line_c0', value=1.0) - - # THEN EXPECT it falls back to the plain name rather than claiming an owner - assert analysis._parameter_labels().label(stranger) == 'Line_c0' - - -class TestInPlaceBindingEdits: - def test_changing_the_number_of_targets_rebuilds_the_fitter(self): - # WHEN a binding is edited in place so that it resolves to two targets instead of one. - # ParameterAnalysis cannot observe this, and the cached fitter would otherwise still hold - # one fit function against two datasets, which dies inside the minimizer. - binding = edyn.FitBinding( - model=sm.BrownianTranslationalDiffusion( - name='Brownian', - lorentzian_name='Lorentzian', - diffusion_coefficient=2.4e-9, - scale=0.5, - ), - targets={'width': 'Lorentzian width'}, - ) - analysis = edyn.ParameterAnalysis(parameters=make_dataset(), bindings=[binding]) - assert len(analysis.fit()) == 1 - - # THEN - binding.targets = {'width': 'Lorentzian width', 'area': 'Lorentzian area'} - - # EXPECT the fit follows the binding rather than failing on a stale fitter - assert len(analysis.fit()) == 2 - - def test_shrinking_the_targets_also_rebuilds(self): - # WHEN - binding = edyn.FitBinding( - model=sm.BrownianTranslationalDiffusion( - name='Brownian', - lorentzian_name='Lorentzian', - diffusion_coefficient=2.4e-9, - scale=0.5, - ), - targets={'width': 'Lorentzian width', 'area': 'Lorentzian area'}, - ) - analysis = edyn.ParameterAnalysis(parameters=make_dataset(), bindings=[binding]) - assert len(analysis.fit()) == 2 - - # THEN - binding.targets = {'width': 'Lorentzian width'} - - # EXPECT - assert len(analysis.fit()) == 1 diff --git a/tests/unit/easydynamics/analysis/test_posterior_sampling.py b/tests/unit/easydynamics/analysis/test_posterior_sampling.py index 0d22ea259..254c76483 100644 --- a/tests/unit/easydynamics/analysis/test_posterior_sampling.py +++ b/tests/unit/easydynamics/analysis/test_posterior_sampling.py @@ -2,8 +2,10 @@ # SPDX-License-Identifier: BSD-3-Clause """ -Unit tests for the posterior sampler, driven through an Analysis1d, with the EasyScience Sampler -mocked out. +Unit tests for the posterior sampler, with the EasyScience Sampler mocked out. + +The sampler is driven through the analyses that hold one: an Analysis1d and a ParameterAnalysis +for PosteriorSampler, and an Analysis for the multi-Q subclass. """ import types @@ -11,11 +13,17 @@ from unittest.mock import MagicMock from unittest.mock import patch +import matplotlib as mpl import numpy as np import pytest import scipp as sc from easyscience.fitting import AvailableMinimizers +from easyscience.fitting.multi_fitter import MultiFitter + +mpl.use('Agg') +import easydynamics as edyn +import easydynamics.sample_model as sm from easydynamics.analysis.analysis1d import Analysis1d from easydynamics.experiment import Experiment from easydynamics.sample_model import InstrumentModel @@ -94,6 +102,103 @@ def analysis(): return make_analysis() +Q_VALUES = [0.5, 1.0, 1.5] + + +def make_multi_q_analysis(): + energy_values = np.linspace(-5.0, 5.0, 15) + rows = [2.0 * np.exp(-0.5 * (energy_values / (0.8 + 0.4 * q**2)) ** 2) for q in Q_VALUES] + observed = np.vstack(rows) + experiment = edyn.Experiment( + data=sc.DataArray( + data=sc.array( + dims=['Q', 'energy'], + values=observed, + variances=np.full_like(observed, 0.01), + ), + coords={ + 'Q': sc.array(dims=['Q'], values=Q_VALUES, unit='1/Angstrom'), + 'energy': sc.array(dims=['energy'], values=energy_values, unit='meV'), + }, + ) + ) + return edyn.Analysis( + display_name='TestMultiQ', + experiment=experiment, + sample_model=sm.SampleModel(components=sm.Gaussian(area=2.0, width=1.0)), + instrument_model=sm.InstrumentModel(), + ) + + +def bound_all_chain(multi_q_analysis, half_width=5.0): + for parameter in multi_q_analysis._chain_parameters(): + parameter.min = float(parameter.value) - half_width + parameter.max = float(parameter.value) + half_width + + +def fake_chain_results(parameters, n_draws=50): + draws = np.tile([float(p.value) for p in parameters], (n_draws, 1)) + return SimpleNamespace( + draws=draws, + param_names=[p.unique_name for p in parameters], + logp=np.zeros(n_draws), + state=MagicMock(Ngen=10, Npop=4), + ) + + +@pytest.fixture +def multi_q_analysis(): + return make_multi_q_analysis() + + +Q = np.array([0.5, 0.8, 1.1, 1.4, 1.7, 2.0]) + + +def make_dataset(): + widths = 0.10 + 0.35 * Q + areas = 2.0 - 0.3 * Q + return sc.Dataset({ + 'Lorentzian width': sc.DataArray( + data=sc.array( + dims=['Q'], values=widths, variances=np.full_like(widths, 1e-4), unit='meV' + ), + coords={'Q': sc.array(dims=['Q'], values=Q, unit='1/angstrom')}, + ), + 'Lorentzian area': sc.DataArray( + data=sc.array( + dims=['Q'], values=areas, variances=np.full_like(areas, 4e-4), unit='meV' + ), + coords={'Q': sc.array(dims=['Q'], values=Q, unit='1/angstrom')}, + ), + }) + + +def make_parameter_analysis(two_bindings=True): + bindings = [ + edyn.FitBinding( + model=sm.Polynomial( + coefficients=[0.1, 0.35], x_unit='1/angstrom', y_unit='meV', name='Width line' + ), + targets='Lorentzian width', + ) + ] + if two_bindings: + bindings.append( + edyn.FitBinding( + model=sm.Polynomial( + coefficients=[2.0, -0.3], x_unit='1/angstrom', y_unit='meV', name='Area line' + ), + targets='Lorentzian area', + ) + ) + return edyn.ParameterAnalysis(parameters=make_dataset(), bindings=bindings) + + +@pytest.fixture +def parameter_analysis(): + return make_parameter_analysis() + + class TestPosteriorSampler: ############# # Bounds pre-flight @@ -485,27 +590,10 @@ def test_plots_without_sampling_raise(self, analysis): with pytest.raises(RuntimeError): analysis.bayesian.plot_corner() + ############# + # Error paths + ############# -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 - - -class TestErrorPaths: def test_bumps_outlier_crash_is_reported_helpfully(self, analysis): # WHEN BUMPS' own outlier removal indexes past the end of its buffer bound_all(analysis) @@ -576,8 +664,10 @@ def test_load_chain_uses_the_sidecar_when_present(self, analysis, tmp_path): 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) + ############# + # Plot rendering + ############# -class TestPlotRendering: def test_trace_and_corner_render_from_a_chain(self, analysis): # WHEN import matplotlib as mpl @@ -596,8 +686,10 @@ def test_trace_and_corner_render_from_a_chain(self, analysis): assert len(analysis.bayesian.plot_corner().axes) == n_parameters**2 plt.close('all') + ############# + # Extend guards + ############# -class TestExtendGuards: def test_extending_with_a_different_subset_is_refused(self, analysis): # WHEN a chain is started over all parameters and then extended over one bound_all(analysis) @@ -625,8 +717,10 @@ def test_extending_with_the_same_parameters_is_allowed(self, analysis): # THEN EXPECT: does not raise analysis.bayesian.extend(additional_samples=10) + ############# + # Sidecar labels + ############# -class TestSidecarLabels: def test_a_subset_run_records_the_same_labels_a_full_run_would(self, analysis): # WHEN only one parameter is sampled. Inside the run the others are fixed, so nothing looks # ambiguous; the recorded labels must still match what a full run would have written, or @@ -661,3 +755,482 @@ def test_extending_after_a_failed_run_is_allowed(self, analysis): # THEN EXPECT the shape guard steps aside rather than comparing against nothing sampler_class.return_value.extend.side_effect = lambda **_k: fake_results(analysis) analysis.bayesian.extend(additional_samples=10) + + ############# + # Driven through a ParameterAnalysis + ############# + + def test_refuses_unbounded_parameters(self, parameter_analysis): + # THEN EXPECT + with pytest.raises(ValueError, match='finite bounds'): + parameter_analysis.bayesian.sample(samples=10) + + def test_binds_one_dataset_per_target(self, parameter_analysis): + # WHEN + bound_all_chain(parameter_analysis) + parameters = parameter_analysis._chain_parameters() + + # THEN + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_chain_results(parameters) + parameter_analysis.bayesian.sample(samples=10) + + # EXPECT + args, kwargs = sampler_class.call_args + assert len(args[1]) == 2 + assert len(kwargs['weights']) == 2 + + def test_summary_uses_model_qualified_labels(self, parameter_analysis): + # WHEN + bound_all_chain(parameter_analysis) + parameters = parameter_analysis._chain_parameters() + + # THEN + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_chain_results(parameters) + parameter_analysis.bayesian.sample(samples=10) + + # EXPECT + names = [entry.name for entry in parameter_analysis.bayesian.summary()] + assert len(set(names)) == len(names) + assert 'Width line_c0' in names + + def test_restores_parameter_values(self, parameter_analysis): + # WHEN + bound_all_chain(parameter_analysis) + parameters = parameter_analysis._chain_parameters() + before = [float(p.value) for p in parameters] + + # THEN + with patch(SAMPLER_PATH) as sampler_class: + + def mutate(**_kwargs): + for parameter in parameters: + parameter.value = float(parameter.value) + 1.0 + return fake_chain_results(parameters) + + sampler_class.return_value.sample.side_effect = mutate + parameter_analysis.bayesian.sample(samples=10) + + # EXPECT + assert [float(p.value) for p in parameters] == pytest.approx(before) + + def test_missing_parameters_dataset_raises(self): + # WHEN + parameter_analysis = edyn.ParameterAnalysis() + + # THEN EXPECT + with pytest.raises(ValueError, match='No parameters Dataset'): + parameter_analysis.bayesian.sample(samples=10) + + def test_missing_bindings_raises(self): + # WHEN + parameter_analysis = edyn.ParameterAnalysis(parameters=make_dataset()) + + # THEN EXPECT + with pytest.raises(ValueError, match='No fit bindings'): + parameter_analysis.bayesian.sample(samples=10) + + +class TestMultiQPosteriorSampler: + ############# + # Bounds pre-flight + ############# + + def test_sampling_refuses_unbounded_parameters(self, multi_q_analysis): + # THEN EXPECT + with pytest.raises(ValueError, match='finite bounds'): + multi_q_analysis.bayesian.sample(fit_method='simultaneous', samples=10) + + def test_error_names_parameters_by_q_index(self, multi_q_analysis): + # THEN EXPECT + with pytest.raises(ValueError, match=r'Gaussian width \(Q_index=0\)'): + multi_q_analysis.bayesian.check_bounds() + + def test_suggest_bounds_labels_every_q(self, multi_q_analysis): + # THEN + suggestions = multi_q_analysis.bayesian.suggest_bounds() + + # EXPECT + labels = [s.label for s in suggestions] + assert len(set(labels)) == len(labels) + assert 'Gaussian area (Q_index=1)' in labels + + ############# + # Simultaneous sampling + ############# + + def test_binds_one_dataset_per_q_index(self, multi_q_analysis): + # WHEN + bound_all_chain(multi_q_analysis) + parameters = multi_q_analysis._chain_parameters() + + # THEN + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_chain_results(parameters) + multi_q_analysis.bayesian.sample(fit_method='simultaneous', samples=10) + + # EXPECT + args, kwargs = sampler_class.call_args + assert len(args[1]) == len(Q_VALUES) + assert len(args[2]) == len(Q_VALUES) + assert len(kwargs['weights']) == len(Q_VALUES) + + def test_returns_a_single_result(self, multi_q_analysis): + # WHEN + bound_all_chain(multi_q_analysis) + parameters = multi_q_analysis._chain_parameters() + + # THEN + with patch(SAMPLER_PATH) as sampler_class: + expected = fake_chain_results(parameters) + sampler_class.return_value.sample.return_value = expected + returned = multi_q_analysis.bayesian.sample(fit_method='simultaneous', samples=10) + + # EXPECT + assert returned is expected + assert multi_q_analysis.bayesian.results is expected + + def test_summary_is_labelled_by_q_index(self, multi_q_analysis): + # WHEN + bound_all_chain(multi_q_analysis) + parameters = multi_q_analysis._chain_parameters() + + # THEN + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_chain_results(parameters) + multi_q_analysis.bayesian.sample(fit_method='simultaneous', samples=10) + + # EXPECT + names = [entry.name for entry in multi_q_analysis.bayesian.summary()] + assert len(set(names)) == len(names) + assert all('Q_index=' in name for name in names) + + def test_refreshes_every_convolver_before_sampling(self, multi_q_analysis): + # WHEN + bound_all_chain(multi_q_analysis) + parameters = multi_q_analysis._chain_parameters() + + # THEN + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_chain_results(parameters) + for analysis1d in multi_q_analysis.analysis_list: + analysis1d._convolver_is_dirty = True + multi_q_analysis.bayesian.sample(fit_method='simultaneous', samples=10) + + # EXPECT the sampler sees the same prepared convolvers a simultaneous fit would + assert all(not a._convolver_is_dirty for a in multi_q_analysis.analysis_list) + + def test_uses_a_multifitter(self, multi_q_analysis): + # WHEN + + # EXPECT + assert isinstance(multi_q_analysis.fitter, MultiFitter) + assert len(multi_q_analysis.fitter.fit_object) == len(Q_VALUES) + + ############# + # Independent sampling + ############# + + def test_returns_one_result_per_q_index(self, multi_q_analysis): + # WHEN + for analysis1d in multi_q_analysis.analysis_list: + for parameter in analysis1d.get_free_parameters(): + parameter.min = float(parameter.value) - 5.0 + parameter.max = float(parameter.value) + 5.0 + + # THEN + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_chain_results( + multi_q_analysis.analysis_list[0].get_free_parameters() + ) + results = multi_q_analysis.bayesian.sample(fit_method='independent', samples=10) + + # EXPECT + assert isinstance(results, list) + assert len(results) == len(Q_VALUES) + + def test_single_q_index_returns_one_result(self, multi_q_analysis): + # WHEN + target = multi_q_analysis.analysis_list[1] + for parameter in target.get_free_parameters(): + parameter.min = float(parameter.value) - 5.0 + parameter.max = float(parameter.value) + 5.0 + + # THEN + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_chain_results( + target.get_free_parameters() + ) + result = multi_q_analysis.bayesian.sample( + fit_method='independent', Q_index=1, samples=10 + ) + + # EXPECT + assert not isinstance(result, list) + assert result is target.bayesian.results + + def test_invalid_q_index_raises(self, multi_q_analysis): + # THEN EXPECT + with pytest.raises((ValueError, IndexError)): + multi_q_analysis.bayesian.sample(fit_method='independent', Q_index=99, samples=10) + + ############# + # Validation + ############# + + def test_invalid_fit_method_raises(self, multi_q_analysis): + # THEN EXPECT + with pytest.raises(ValueError, match='Invalid fit method'): + multi_q_analysis.bayesian.sample(fit_method='nonsense') + + def test_missing_q_values_raises(self): + # WHEN + multi_q_analysis = edyn.Analysis(display_name='Empty') + + # THEN EXPECT + with pytest.raises(ValueError, match='No Q values available'): + multi_q_analysis.bayesian.sample() + + ############# + # Predictive plot + ############# + + def test_predictive_is_not_supported_for_multiple_datasets(self, multi_q_analysis): + # WHEN + bound_all_chain(multi_q_analysis) + parameters = multi_q_analysis._chain_parameters() + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_chain_results(parameters) + multi_q_analysis.bayesian.sample(fit_method='simultaneous', samples=10) + + # THEN EXPECT + with pytest.raises(NotImplementedError, match='single dataset only'): + multi_q_analysis.bayesian.plot_posterior_predictive() + + ############# + # Discoverability + ############# + + def test_operations_needing_one_chain_point_at_the_per_q_chains(self, multi_q_analysis): + # WHEN sampling independently, the chains live on the Analysis1d objects, not here + remaining = iter(multi_q_analysis.analysis_list) + for analysis1d in multi_q_analysis.analysis_list: + for parameter in analysis1d.get_free_parameters(): + parameter.min = float(parameter.value) - 5.0 + parameter.max = float(parameter.value) + 5.0 + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_chain_results( + next(remaining).get_free_parameters() + ) + multi_q_analysis.bayesian.sample(fit_method='independent', samples=10) + + # THEN EXPECT anything that genuinely needs a single chain says where the chains + # actually are, rather than claiming none exist + with pytest.raises(RuntimeError, match='analysis_list'): + multi_q_analysis.bayesian.plot_posterior_predictive() + + def test_untouched_analysis_still_reports_no_samples(self, multi_q_analysis): + # THEN EXPECT the plain message when nothing has been sampled anywhere + with pytest.raises(RuntimeError, match='No posterior samples yet'): + multi_q_analysis.bayesian.summary() + + ############# + # Aggregating the per-Q chains + ############# + + def _sample_independently(self, multi_q_analysis): + for analysis1d in multi_q_analysis.analysis_list: + for parameter in analysis1d.get_free_parameters(): + parameter.min = float(parameter.value) - 5.0 + parameter.max = float(parameter.value) + 5.0 + + # The Q indices sample in order, and each must get a chain over its own parameters. + remaining = iter(multi_q_analysis.analysis_list) + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_chain_results( + next(remaining).get_free_parameters() + ) + multi_q_analysis.bayesian.sample(fit_method='independent', samples=10) + + def test_posterior_results_holds_one_chain_per_q(self, multi_q_analysis): + # WHEN + self._sample_independently(multi_q_analysis) + + # EXPECT + assert len(multi_q_analysis.bayesian.results_per_q) == len(Q_VALUES) + assert all(result is not None for result in multi_q_analysis.bayesian.results_per_q) + + def test_posterior_results_is_none_before_sampling(self, multi_q_analysis): + # EXPECT + assert multi_q_analysis.bayesian.results_per_q is None + + def test_summary_gathers_every_q(self, multi_q_analysis): + # WHEN + self._sample_independently(multi_q_analysis) + + # THEN + summary = multi_q_analysis.bayesian.summary() + + # EXPECT one entry per free parameter per Q, each labelled by its Q index + expected = sum(len(a.get_free_parameters()) for a in multi_q_analysis.analysis_list) + names = [entry.name for entry in summary] + assert len(summary) == expected + assert len(set(names)) == len(names) + assert all('Q_index=' in name for name in names) + + def test_median_applies_each_chain_to_its_own_q(self, multi_q_analysis): + # WHEN + self._sample_independently(multi_q_analysis) + + # THEN + changed = multi_q_analysis.bayesian.set_parameters_to_median() + + # EXPECT every Q's parameters are set, from that Q's own chain + expected = sum(len(a.get_free_parameters()) for a in multi_q_analysis.analysis_list) + assert len(changed) == expected + + def test_corner_plots_one_q_at_a_time(self, multi_q_analysis): + # WHEN each Q was sampled separately, no draw pairs one Q with another, so a corner plot + # can only show one chain at a time + self._sample_independently(multi_q_analysis) + + # THEN + figure = multi_q_analysis.bayesian.plot_corner(Q_index=1) + + # EXPECT that Q's own chain, not a combination across Q + n_parameters = len(multi_q_analysis.analysis_list[1].get_free_parameters()) + assert len(figure.axes) == n_parameters**2 + + def test_corner_offers_a_slider_in_a_notebook(self, multi_q_analysis): + # WHEN + self._sample_independently(multi_q_analysis) + + # THEN + with patch('easydynamics.analysis.posterior_sampling._in_notebook', return_value=True): + widget = multi_q_analysis.bayesian.plot_corner() + + # EXPECT a slider over the sampled Q indices, and a panel that actually holds a figure. + # The obvious way to build this captures nothing and leaves the panel blank beside the + # slider, so an empty panel is the regression worth guarding. Which mime type arrives + # depends on the environment: a live kernel renders a PNG, plain pytest only the repr. + # The figure comes first and the slider sits under it, where plopp puts its controls. + panel, slider = widget.children + assert list(slider.options) == list(range(len(Q_VALUES))) + assert panel.outputs, 'the initial chain was not drawn' + assert 'Figure' in str(panel.outputs[0]['data']) + + slider.value = 2 + assert panel.outputs, 'changing Q did not redraw' + assert 'Figure' in str(panel.outputs[0]['data']) + + def test_corner_without_a_notebook_or_q_index_says_what_to_do(self, multi_q_analysis): + # WHEN + self._sample_independently(multi_q_analysis) + + # THEN EXPECT it names the sampled Q indices rather than just refusing + with ( + patch('easydynamics.analysis.posterior_sampling._in_notebook', return_value=False), + pytest.raises(RuntimeError, match=r'sampled Q indices are \[0, 1, 2\]'), + ): + multi_q_analysis.bayesian.plot_corner() + + def test_the_slider_only_offers_q_indices_that_were_sampled(self, multi_q_analysis): + # WHEN only one Q index is sampled + target = multi_q_analysis.analysis_list[2] + for parameter in target.get_free_parameters(): + parameter.min = float(parameter.value) - 5.0 + parameter.max = float(parameter.value) + 5.0 + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_chain_results( + target.get_free_parameters() + ) + multi_q_analysis.bayesian.sample(fit_method='independent', Q_index=2, samples=10) + + # THEN + with patch('easydynamics.analysis.posterior_sampling._in_notebook', return_value=True): + widget = multi_q_analysis.bayesian.plot_corner() + + # EXPECT the slider cannot land on a Q with nothing to draw + assert list(widget.children[1].options) == [2] + + def test_trace_points_at_the_individual_chains(self, multi_q_analysis): + # WHEN + self._sample_independently(multi_q_analysis) + + # THEN EXPECT + with pytest.raises(RuntimeError, match='no single trace'): + multi_q_analysis.bayesian.plot_trace() + + def test_a_simultaneous_chain_still_takes_precedence(self, multi_q_analysis): + # WHEN a simultaneous run follows an independent one + self._sample_independently(multi_q_analysis) + bound_all_chain(multi_q_analysis) + parameters = multi_q_analysis._chain_parameters() + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_chain_results(parameters) + multi_q_analysis.bayesian.sample(fit_method='simultaneous', samples=10) + + # EXPECT the single chain is summarized, not the stale per-Q ones + assert len(multi_q_analysis.bayesian.summary()) == len(parameters) + multi_q_analysis.bayesian.plot_corner() + + def test_only_the_sampled_q_indices_are_gathered(self, multi_q_analysis): + # WHEN just one Q index is sampled + target = multi_q_analysis.analysis_list[1] + for parameter in target.get_free_parameters(): + parameter.min = float(parameter.value) - 5.0 + parameter.max = float(parameter.value) + 5.0 + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_chain_results( + target.get_free_parameters() + ) + multi_q_analysis.bayesian.sample(fit_method='independent', Q_index=1, samples=10) + + # THEN + summary = multi_q_analysis.bayesian.summary() + + # EXPECT the unsampled Q indices are passed over rather than breaking the aggregation + assert len(summary) == len(target.get_free_parameters()) + assert all('Q_index=1' in entry.name for entry in summary) + assert len(multi_q_analysis.bayesian.set_parameters_to_median()) == len( + target.get_free_parameters() + ) + + def test_a_simultaneous_chain_serves_the_median_and_the_trace(self, multi_q_analysis): + # WHEN + bound_all_chain(multi_q_analysis) + parameters = multi_q_analysis._chain_parameters() + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_chain_results(parameters) + multi_q_analysis.bayesian.sample(fit_method='simultaneous', samples=10) + + # EXPECT both come from the single chain, with no per-Q gathering involved + assert len(multi_q_analysis.bayesian.set_parameters_to_median()) == len(parameters) + assert len(multi_q_analysis.bayesian.plot_trace().axes) == len(parameters) + 1 + + +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 From f6756c71769ad5dc41a705ff3021350a1b58e419 Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Mon, 17 Aug 2026 11:05:06 +0200 Subject: [PATCH 20/29] Refuse silent chain corruption and harden the posterior sampler - extend() now verifies the chain holds the same parameters, not just the same number, and refuses to resume after a failed run or after the model or data changed - Parameter objects passed to sample(parameters=...) are validated against the free set the same way strings are - sampling with no free parameters and degenerate (min >= max) bounds raise clear errors before reaching BUMPS - parameters_at_bounds keys by unique_name so same-named per-Q parameters no longer collide, and guards empty draws - suggest_bounds flags non-finite fitted uncertainties for attention - save() refuses to write an empty label sidecar; loading one warns like a missing sidecar - colliding display labels get positional suffixes in the sidecar so save/load resolves each column to its own parameter - plot_posterior_predictive omits error bars when the data carries no variances (new Experiment.has_variances) - posterior plots validate draws/logp up front, name NaN columns, and share x-limits per corner column - document that sampling runs are not seedable Co-Authored-By: Claude Fable 5 --- src/easydynamics/analysis/posterior.py | 53 ++++- src/easydynamics/analysis/posterior_labels.py | 26 +- .../analysis/posterior_sampling.py | 158 +++++++++++-- src/easydynamics/experiment/experiment.py | 15 ++ src/easydynamics/utils/posterior_plotting.py | 63 ++++- .../easydynamics/analysis/test_posterior.py | 65 ++++- .../analysis/test_posterior_labels.py | 28 +++ .../analysis/test_posterior_sampling.py | 222 +++++++++++++++++- .../experiment/test_experiment.py | 13 + .../utils/test_posterior_plotting.py | 43 ++++ 10 files changed, 638 insertions(+), 48 deletions(-) diff --git a/src/easydynamics/analysis/posterior.py b/src/easydynamics/analysis/posterior.py index 4853bb1a5..d8d65af77 100644 --- a/src/easydynamics/analysis/posterior.py +++ b/src/easydynamics/analysis/posterior.py @@ -319,9 +319,19 @@ def _suggest_bounds_for_parameter( reason='value is not finite', ) - half_width = relative_pad * abs(value) - if np.isfinite(error): - half_width += n_sigma * error + # 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) @@ -364,6 +374,33 @@ def unbounded_parameters(parameters: list[Parameter]) -> list[Parameter]: ] +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], @@ -385,10 +422,14 @@ def parameters_at_bounds( Returns ------- dict[str, float] - Mapping of parameter name to the fraction of draws sitting in the outer + 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``. + 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: @@ -402,7 +443,7 @@ def parameters_at_bounds( at_edge = (values <= low + edge) | (values >= high - edge) fraction = float(np.count_nonzero(at_edge)) / len(values) if fraction > BOUND_OCCUPANCY_THRESHOLD: - piled_up[parameter.name] = fraction + piled_up[parameter.unique_name] = fraction return piled_up diff --git a/src/easydynamics/analysis/posterior_labels.py b/src/easydynamics/analysis/posterior_labels.py index 84ff9212e..e39811718 100644 --- a/src/easydynamics/analysis/posterior_labels.py +++ b/src/easydynamics/analysis/posterior_labels.py @@ -47,7 +47,23 @@ def __init__( self._qualify = qualify self._counts = Counter(parameter.name for parameter in self._parameters) self._by_unique_name = {p.unique_name: p for p in self._parameters} - self._by_label = {self.label(p): p for p in self._parameters} + # 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]: @@ -85,14 +101,16 @@ 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. + 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. + Mapping of unique name to label, collision-free. """ - return {p.unique_name: self.label(p) for p in self._parameters} + return dict(self._storage_labels) def resolve( self, diff --git a/src/easydynamics/analysis/posterior_sampling.py b/src/easydynamics/analysis/posterior_sampling.py index f782b2d00..44ee7dc85 100644 --- a/src/easydynamics/analysis/posterior_sampling.py +++ b/src/easydynamics/analysis/posterior_sampling.py @@ -22,6 +22,7 @@ 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 @@ -196,18 +197,27 @@ def check_bounds(self) -> None: Raises ------ ValueError - If any free parameter has an infinite lower or upper bound. + 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 not unbounded: - return - names = ', '.join(labels.label(parameter) for parameter in unbounded) - raise ValueError( - f'Bayesian sampling requires finite bounds on every free parameter, because the ' - f'bounds act as the prior. These parameters are unbounded: {names}. ' - f'Set their min and max, or call suggest_bounds() to propose values.' - ) + 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 @@ -252,6 +262,13 @@ def sample( ------- 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. """ return self._run( parameters=parameters, @@ -290,7 +307,15 @@ def extend( Raises ------ RuntimeError - If there is no chain to extend. + If there is no chain to extend, or the previous run failed and left no results. + ValueError + If the model or data changed since the chain was started, or this run's parameters + differ from the ones the chain holds. + + Notes + ----- + 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.') @@ -335,6 +360,8 @@ def _run( 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) @@ -344,6 +371,11 @@ def _run( 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: @@ -384,13 +416,25 @@ def _get_or_build_sampler(self, reuse_sampler: bool) -> Sampler: Parameters ---------- reuse_sampler : bool - Whether to reuse the cached Sampler even if it is marked dirty. + 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) @@ -399,7 +443,7 @@ def _get_or_build_sampler(self, reuse_sampler: bool) -> Sampler: def _verify_chain_shape_unchanged(self, chain_parameters: list[Parameter]) -> None: """ - Check that an extension keeps the chain's columns. + Check that an extension keeps the chain's columns, both in count and in identity. Parameters ---------- @@ -408,11 +452,16 @@ def _verify_chain_shape_unchanged(self, chain_parameters: list[Parameter]) -> No Raises ------ + RuntimeError + If there are no stored results to continue from, as after a failed run. ValueError - If the number of parameters differs from the existing chain's. + If the number or the identity of the parameters differs from the existing chain's. """ if self._results is None: - return + 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( @@ -422,6 +471,28 @@ def _verify_chain_shape_unchanged(self, chain_parameters: list[Parameter]) -> No 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, @@ -444,7 +515,8 @@ def _resolve_parameters_to_hold_fixed( TypeError If parameters is not a list of Parameters or strings, or None. ValueError - If a requested label matches no free parameter, or the subset is empty. + If a requested label or Parameter matches no free parameter of this analysis, or the + subset is empty. """ if parameters is None: return [] @@ -453,6 +525,7 @@ def _resolve_parameters_to_hold_fixed( 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): @@ -463,7 +536,17 @@ def _resolve_parameters_to_hold_fixed( ) requested.append(by_label[entry]) elif hasattr(entry, 'unique_name'): - requested.append(entry) + # 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).') @@ -484,8 +567,13 @@ def _warn_about_bounds_occupancy(self, results: SamplingResults) -> None: 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'{name} ({fraction:.0%} of draws)' for name, fraction in piled_up.items() + 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( ( @@ -565,6 +653,20 @@ def save(self, path: str | os.PathLike) -> None: 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' ) @@ -589,15 +691,16 @@ def load(self, path: str | os.PathLike, skip: int = 0) -> SamplingResults: """ self._prepare() sidecar = Path(f'{path}{_LABEL_MAP_SUFFIX}') - if sidecar.is_file(): - self._saved_labels = json.loads(sidecar.read_text(encoding='utf-8')) - else: - self._saved_labels = {} + 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 found at {sidecar}. The chain will be reported ' - f'under the internal names it was saved with, because those cannot be matched ' - f'to this Analysis.' + 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, @@ -718,11 +821,16 @@ def plot_posterior_predictive( 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=None if weights is None else 1.0 / np.asarray(weights), + 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, 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/posterior_plotting.py b/src/easydynamics/utils/posterior_plotting.py index 8e75cf425..bfa26e987 100644 --- a/src/easydynamics/utils/posterior_plotting.py +++ b/src/easydynamics/utils/posterior_plotting.py @@ -35,8 +35,8 @@ def plot_trace( excursions. A visible trend means the chain has not reached the typical set and needs a longer burn-in. - A ``ValueError`` is raised if ``draws`` is not two-dimensional, or if ``names`` does not have - one entry per column. + A ``ValueError`` is raised 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. Parameters ---------- @@ -45,7 +45,7 @@ def plot_trace( names : list[str] One label per column of ``draws``. logp : np.ndarray | None, default=None - Log-posterior values, plotted in an extra panel when given. + 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. @@ -61,6 +61,13 @@ def plot_trace( """ 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: @@ -72,10 +79,13 @@ def plot_trace( for axis, column, name in zip(axes, range(draws.shape[1]), names, strict=False): axis.plot(draws[:, column], lw=0.5) axis.set_ylabel(_with_unit(name, units, column), fontsize=8) - axis.set_xlim(0, len(draws) - 1) + # 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(np.asarray(logp), lw=0.5, color='C4') + axes[-1].plot(logp, lw=0.5, color='C4') axes[-1].set_ylabel('log-posterior', fontsize=8) axes[-1].set_xlabel('sample index') @@ -100,8 +110,8 @@ def plot_corner( distribution of a pair: a compact blob means the two are independent, while a narrow diagonal ridge means they are correlated and cannot be determined separately from this data. - A ``ValueError`` is raised if ``draws`` is not two-dimensional, or if ``names`` does not have - one entry per column. + A ``ValueError`` is raised 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. Parameters ---------- @@ -127,11 +137,22 @@ def plot_corner( 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): @@ -144,6 +165,8 @@ def plot_corner( 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: @@ -268,6 +291,28 @@ def plot_posterior_predictive( 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 _unit_for(units: list[str] | None, column: int) -> str: """ Get the unit to show for a column, if it is worth showing. @@ -365,6 +410,10 @@ def _verify_draws(draws: np.ndarray, names: list[str]) -> None: """ 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. ' diff --git a/tests/unit/easydynamics/analysis/test_posterior.py b/tests/unit/easydynamics/analysis/test_posterior.py index 732dcb8df..28817e340 100644 --- a/tests/unit/easydynamics/analysis/test_posterior.py +++ b/tests/unit/easydynamics/analysis/test_posterior.py @@ -7,6 +7,7 @@ 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 @@ -95,6 +96,18 @@ def test_absolute_floor_rescues_a_scaleless_parameter(self): 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) @@ -187,6 +200,29 @@ def test_finds_parameters_with_an_infinite_side(self): 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 @@ -197,8 +233,8 @@ def test_uniform_posterior_across_the_bounds_is_reported(self): result = parameters_at_bounds(draws, [parameter]) # EXPECT - assert parameter.name in result - assert result[parameter.name] == pytest.approx(0.1, abs=0.01) + 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 @@ -221,7 +257,7 @@ def test_partly_clipped_posterior_is_reported(self): result = parameters_at_bounds(draws, [parameter]) # EXPECT - assert parameter.name in result + assert parameter.unique_name in result def test_posterior_pinned_at_one_bound_is_reported(self): # WHEN @@ -232,7 +268,7 @@ def test_posterior_pinned_at_one_bound_is_reported(self): result = parameters_at_bounds(draws, [parameter]) # EXPECT - assert result[parameter.name] > 0.9 + assert result[parameter.unique_name] > 0.9 def test_unmatched_and_unbounded_columns_are_skipped(self): # WHEN @@ -245,6 +281,27 @@ def test_unmatched_and_unbounded_columns_are_skipped(self): # 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): diff --git a/tests/unit/easydynamics/analysis/test_posterior_labels.py b/tests/unit/easydynamics/analysis/test_posterior_labels.py index 82347db72..618c8284c 100644 --- a/tests/unit/easydynamics/analysis/test_posterior_labels.py +++ b/tests/unit/easydynamics/analysis/test_posterior_labels.py @@ -89,6 +89,34 @@ def test_an_unknown_column_is_reported_not_guessed(self): 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') diff --git a/tests/unit/easydynamics/analysis/test_posterior_sampling.py b/tests/unit/easydynamics/analysis/test_posterior_sampling.py index ee4c1e03c..6f2cad85e 100644 --- a/tests/unit/easydynamics/analysis/test_posterior_sampling.py +++ b/tests/unit/easydynamics/analysis/test_posterior_sampling.py @@ -14,6 +14,7 @@ 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 @@ -24,13 +25,13 @@ SAMPLER_PATH = 'easydynamics.analysis.posterior_sampling.Sampler' -def make_analysis(): +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, :], + variances=np.full_like(intensity, 0.01)[None, :] if with_variances else None, ) experiment = Experiment( data=sc.DataArray( @@ -108,6 +109,17 @@ def test_suggest_bounds_covers_the_free_parameters(self, analysis): # 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 ############# @@ -207,6 +219,15 @@ def test_warns_when_the_posterior_piles_up_against_a_bound(self, analysis): 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) @@ -265,6 +286,26 @@ def test_unknown_parameter_name_raises(self, analysis): 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'): @@ -347,6 +388,51 @@ def test_extend_delegates_to_the_sampler(self, analysis): 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'): @@ -381,6 +467,48 @@ def test_load_without_a_sidecar_warns(self, analysis, tmp_path): 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 ############# @@ -470,6 +598,96 @@ def test_plots_without_sampling_raise(self, analysis): 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 + + ############# + # 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.""" 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 index 1a412cdb2..5f341a24c 100644 --- a/tests/unit/easydynamics/utils/test_posterior_plotting.py +++ b/tests/unit/easydynamics/utils/test_posterior_plotting.py @@ -65,6 +65,29 @@ def test_one_dimensional_draws_raise(self): 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): @@ -86,6 +109,26 @@ def test_mismatched_names_raise(self, draws): 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 TestPlotPosteriorPredictive: def test_returns_a_figure_with_data_and_band(self): From 2839da454a9d31557dc7829c1b21fad0ec6265b8 Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Mon, 17 Aug 2026 11:55:36 +0200 Subject: [PATCH 21/29] Keep the multi-Q sampler pointed at the chain the user actually ran - sampling one Q index independently now clears a stale simultaneous chain, so summary(), set_parameters_to_median() and plot_corner() report the run the user just made instead of the old one - extend() and save() after an independent run explain that the chains live on the per-Q analyses instead of resuming or saving the stale simultaneous chain; a genuinely failed run keeps its own message - Q_index arguments are validated like every Analysis method, so a negative index raises instead of silently wrapping - the gathered summary resolves each per-Q chain through its own saved labels, so chains loaded from disk keep names and units - warnings are attributed to the caller on both the single-Q and multi-Q paths, and the corner-plot slider forwards plot kwargs - the multi-Q integration tests share one independent sampling run, assert the straight line is actually recovered, and the extend test no longer mutates the shared fixture Co-Authored-By: Claude Fable 5 --- .../analysis/posterior_sampling.py | 180 ++++++++++++++++-- .../fitting/test_bayesian_sampling.py | 15 +- .../fitting/test_bayesian_sampling_multi_q.py | 109 +++++------ .../analysis/test_posterior_sampling.py | 125 ++++++++++++ 4 files changed, 351 insertions(+), 78 deletions(-) diff --git a/src/easydynamics/analysis/posterior_sampling.py b/src/easydynamics/analysis/posterior_sampling.py index bb364addc..c0da03f70 100644 --- a/src/easydynamics/analysis/posterior_sampling.py +++ b/src/easydynamics/analysis/posterior_sampling.py @@ -12,6 +12,7 @@ from __future__ import annotations +import inspect import json import warnings from pathlib import Path @@ -29,6 +30,7 @@ from easydynamics.analysis.posterior import summarize_draws from easydynamics.analysis.posterior import unbounded_parameters from easydynamics.utils.utils import _in_notebook +from easydynamics.utils.utils import verify_Q_index if TYPE_CHECKING: import os @@ -588,19 +590,26 @@ def _warn_about_bounds_occupancy(self, results: SamplingResults) -> None: f'Widen the bounds, or check whether these parameters are degenerate with others.' ), UserWarning, - stacklevel=4, + stacklevel=_stacklevel_above_module(), ) ############# # Results ############# - def summary(self) -> PosteriorSummary: + def summary(self, labeller: Callable[[Parameter], str] | None = None) -> 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. + Parameters + ---------- + labeller : Callable[[Parameter], str] | None, default=None + Overrides the label a resolved column is reported under. Used by an Analysis covering + several Q values, whose gathered table qualifies each name with its Q index. Columns + that resolve to no parameter keep their usual fallback name. + Returns ------- PosteriorSummary @@ -608,10 +617,17 @@ def summary(self) -> PosteriorSummary: """ results = self._require_results() labels = self._labels() + names = labels.display_names(results.param_names, self._saved_labels) + parameters = self._resolve(results) + if labeller is not None: + names = [ + name if parameter is None else labeller(parameter) + for parameter, name in zip(parameters, names, strict=True) + ] return summarize_draws( draws=results.draws, - labels=labels.display_names(results.param_names, self._saved_labels), - parameters_by_column=self._resolve(results), + labels=names, + parameters_by_column=parameters, ) def set_parameters_to_median(self) -> list[Parameter]: @@ -670,7 +686,7 @@ def save(self, path: str | os.PathLike) -> None: f'one. A future load() will report the columns under their internal names.' ), UserWarning, - stacklevel=2, + stacklevel=_stacklevel_above_module(), ) return Path(f'{path}{_LABEL_MAP_SUFFIX}').write_text( @@ -1050,6 +1066,10 @@ def sample( Raises ------ + IndexError + If Q_index is negative or out of range. + TypeError + If Q_index is not an int or None. ValueError If fit_method is not "independent" or "simultaneous". """ @@ -1060,12 +1080,18 @@ def sample( raise ValueError( 'No Q values available for sampling. Please check the experiment data.' ) + verify_Q_index(Q_index=Q_index, Q=self._analysis.Q, allow_none=True) if fit_method == 'simultaneous': return super().sample(samples=samples, burn=burn, thin=thin, **sampler_options) if Q_index is not None: - return per_q[Q_index].bayesian.sample( + result = per_q[Q_index].bayesian.sample( samples=samples, burn=burn, thin=thin, **sampler_options ) + # The fresh per-Q chain now outranks any older simultaneous one, exactly as after an + # all-Q independent run; keeping the old chain here would make summary() silently + # report it instead. Cleared only on success, so a failed run changes nothing. + self._results = None + return result # The per-Q chains live on their own samplers; this one then holds nothing of its own. self._results = None return [ @@ -1073,7 +1099,90 @@ def sample( for analysis1d in per_q ] - def summary(self) -> PosteriorSummary: + def extend( + self, + additional_samples: int = 5000, + thin: int = 10, + parameters: list[Parameter] | list[str] | None = None, + **sampler_options: dict[str, Any], + ) -> SamplingResults: + """ + Continue the existing simultaneous chain with additional samples. + + The chains from independent sampling live on the per-Q samplers, so each is extended + there rather than here. + + 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:`PosteriorSampler.extend`. + **sampler_options : dict[str, Any] + Forwarded to the EasyScience Sampler. + + Returns + ------- + SamplingResults + The sampling results for the full extended chain. + + Raises + ------ + RuntimeError + If the latest sampling ran per Q index, so there is no simultaneous chain here to + extend, or if there is no chain at all. + ValueError + If the model or data changed since the chain was started, or this run's parameters + differ from the ones the chain holds. + """ + if self._results is None and self.results_per_q is not None: + # Without this check, a stale simultaneous sampler would either be extended silently + # or misdiagnosed as a failed run. + raise RuntimeError( + 'The latest sampling ran per Q index, so there is no simultaneous chain here to ' + 'extend. Extend a per-Q chain with ' + 'analysis.analysis_list[Q_index].bayesian.extend(), or start a fresh simultaneous ' + "chain with sample(fit_method='simultaneous')." + ) + return super().extend( + additional_samples=additional_samples, + thin=thin, + parameters=parameters, + **sampler_options, + ) + + def save(self, path: str | os.PathLike) -> None: + """ + Save the simultaneous MCMC chain to disk. + + The chains from independent sampling live on the per-Q samplers, so each is saved there + rather than here. + + Parameters + ---------- + path : str | os.PathLike + Path prefix for the chain files. + + Raises + ------ + RuntimeError + If the latest sampling ran per Q index -- there is then no simultaneous chain here to + save -- or if there is no chain at all. + """ + if self._results is None and self.results_per_q is not None: + # Without this check, a stale simultaneous chain would be written to disk as if it + # were the latest sampling. + raise RuntimeError( + 'The latest sampling ran per Q index, and those chains live on the per-Q ' + 'samplers; there is no simultaneous chain here to save. Save each with ' + 'analysis.analysis_list[Q_index].bayesian.save(), or sample with ' + "fit_method='simultaneous' first." + ) + super().save(path) + + def summary(self, labeller: Callable[[Parameter], str] | None = None) -> PosteriorSummary: """ Summarize the posterior, gathering the per-Q chains when sampling was independent. @@ -1081,6 +1190,12 @@ def summary(self) -> PosteriorSummary: within its own chain, so collecting them into one table is sound even though the chains are separate. Labels carry the Q index either way, so the table reads the same. + Parameters + ---------- + labeller : Callable[[Parameter], str] | None, default=None + Overrides the label a resolved column is reported under. The default is this + analysis' own Q-qualified labels. + Returns ------- PosteriorSummary @@ -1088,20 +1203,17 @@ def summary(self) -> PosteriorSummary: """ per_q = self.results_per_q if self._results is not None or per_q is None: - return super().summary() + return super().summary(labeller) - labels = self._labels() + # Each chain is summarized by its own per-Q sampler, whose saved labels can match a chain + # loaded from disk in a fresh session; this sampler's labels then supply the Q-qualified + # display name for every column that resolves to a parameter. + qualify = self._labels().label if labeller is None else labeller entries = [] - for result in per_q: - if result is None: + for analysis1d in self._per_q(): + if analysis1d.bayesian.results is None: continue - entries.extend( - summarize_draws( - draws=result.draws, - labels=labels.display_names(result.param_names), - parameters_by_column=labels.resolve(result.param_names), - ).entries - ) + entries.extend(analysis1d.bayesian.summary(labeller=qualify).entries) return PosteriorSummary(entries) def set_parameters_to_median(self) -> list[Parameter]: @@ -1148,11 +1260,16 @@ def plot_corner(self, Q_index: int | None = None, **kwargs: dict[str, Any]) -> F Raises ------ + IndexError + If Q_index is negative or out of range. RuntimeError If a slider is asked for outside a notebook. + TypeError + If Q_index is not an int or None. """ from easydynamics.utils.posterior_plotting import corner_with_slider + verify_Q_index(Q_index=Q_index, Q=self._analysis.Q, allow_none=True) per_q = self.results_per_q if self._results is not None or per_q is None: return super().plot_corner(**kwargs) @@ -1181,7 +1298,7 @@ def plot_corner(self, Q_index: int | None = None, **kwargs: dict[str, Any]) -> F 'names': [entry.name for entry in entries], 'units': [entry.unit for entry in entries], } - return corner_with_slider(chains, title=self._analysis.display_name) + return corner_with_slider(chains, title=self._analysis.display_name, **kwargs) def plot_trace(self, **kwargs: dict[str, Any]) -> Figure: """ @@ -1238,6 +1355,29 @@ def _require_results(self) -> SamplingResults: return super()._require_results() +def _stacklevel_above_module() -> int: + """ + Compute the stacklevel that points a warning at the first frame outside this module. + + The entry points nest to different depths -- ``MultiQPosteriorSampler.sample`` goes through + ``PosteriorSampler.sample`` and ``_run``, a plain ``sample`` skips the first hop -- so any + fixed stacklevel points warnings at an internal frame on one path or the other. Counting the + in-module frames instead lands the warning on the caller's own line either way. + + Returns + ------- + int + The stacklevel for a ``warnings.warn`` call made directly by this function's caller. + """ + frame = inspect.currentframe() + frame = None if frame is None else frame.f_back + level = 1 + while frame is not None and frame.f_globals.get('__name__') == __name__: + frame = frame.f_back + level += 1 + return level + + def _warn_about_held_parameters(labels: object, held_fixed: list[Parameter]) -> None: """ Warn that holding parameters fixed makes the credible intervals conditional. @@ -1260,7 +1400,7 @@ def _warn_about_held_parameters(labels: object, held_fixed: list[Parameter]) -> f'parameters are correlated.' ), UserWarning, - stacklevel=4, + stacklevel=_stacklevel_above_module(), ) diff --git a/tests/integration/fitting/test_bayesian_sampling.py b/tests/integration/fitting/test_bayesian_sampling.py index 1a0262707..b56501a34 100644 --- a/tests/integration/fitting/test_bayesian_sampling.py +++ b/tests/integration/fitting/test_bayesian_sampling.py @@ -135,14 +135,21 @@ def test_sampling_leaves_the_fitted_values_untouched(self): 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) + def test_extend_grows_the_chain(self): + # WHEN a chain of this test's own: extending mutates the sampler state, so running it on + # the module-scoped fixture would hand every later test the extended chain + analysis = build_analysis() + analysis.fit() + analysis.bayesian.suggest_bounds().apply() + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + analysis.bayesian.sample(**SAMPLE_KWARGS) + before = int(analysis.bayesian.results.state.Ngen) # THEN with warnings.catch_warnings(): warnings.simplefilter('ignore') - extended = sampled_analysis.bayesian.extend( + extended = analysis.bayesian.extend( additional_samples=500, thin=2, sampler_kwargs={'trim': False, 'outliers': 'none'} ) diff --git a/tests/integration/fitting/test_bayesian_sampling_multi_q.py b/tests/integration/fitting/test_bayesian_sampling_multi_q.py index ef7a284af..814725c5d 100644 --- a/tests/integration/fitting/test_bayesian_sampling_multi_q.py +++ b/tests/integration/fitting/test_bayesian_sampling_multi_q.py @@ -83,6 +83,19 @@ def simultaneously_sampled(): return analysis +@pytest.fixture(scope='module') +def independently_sampled(): + """One independent DREAM run shared by every test that only reads the per-Q chains.""" + analysis = build_analysis() + analysis.fit(fit_method='independent') + for analysis1d in analysis.analysis_list: + analysis1d.bayesian.suggest_bounds().apply() + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + results = analysis.bayesian.sample(fit_method='independent', **SAMPLE_KWARGS) + return analysis, results + + class TestSimultaneousChain: def test_chain_covers_every_q_index(self, simultaneously_sampled): # THEN @@ -143,34 +156,20 @@ def test_plots_render(self, simultaneously_sampled): class TestIndependentChains: - def test_one_chain_per_q_index(self): - # WHEN - analysis = build_analysis() - analysis.fit(fit_method='independent') - for analysis1d in analysis.analysis_list: - analysis1d.bayesian.suggest_bounds().apply() - + def test_one_chain_per_q_index(self, independently_sampled): # THEN - with warnings.catch_warnings(): - warnings.simplefilter('ignore') - results = analysis.bayesian.sample(fit_method='independent', **SAMPLE_KWARGS) + analysis, results = independently_sampled # EXPECT assert len(results) == len(Q_VALUES) for analysis1d, result in zip(analysis.analysis_list, results, strict=True): assert result.draws.shape[1] == len(analysis1d.get_free_parameters()) - def test_independent_and_simultaneous_agree_on_the_widths(self, simultaneously_sampled): - # WHEN - analysis = build_analysis() - analysis.fit(fit_method='independent') - for analysis1d in analysis.analysis_list: - analysis1d.bayesian.suggest_bounds().apply() - - # THEN the same data is sampled per-Q instead of all at once - with warnings.catch_warnings(): - warnings.simplefilter('ignore') - analysis.bayesian.sample(fit_method='independent', **SAMPLE_KWARGS) + def test_independent_and_simultaneous_agree_on_the_widths( + self, simultaneously_sampled, independently_sampled + ): + # THEN the same data sampled per-Q is compared with the single simultaneous chain + analysis, _ = independently_sampled # EXPECT both routes land on the same widths, since nothing is shared across Q here for q_index, analysis1d in enumerate(analysis.analysis_list): @@ -206,13 +205,12 @@ def test_recovers_a_straight_line_through_the_widths(self): bindings=edyn.FitBinding(model=model, targets='Gaussian width'), ) analysis.fit() - # The linear coefficient sits at exactly zero with a vanishing uncertainty, so the sigma # rule has no scale to work from and flags it rather than inventing one. absolute_floor - # supplies the scale the data cannot. + # supplies the scale the data cannot; the asserts guard that this setup really leaves + # every coefficient bounded before sampling. flagged = analysis.bayesian.suggest_bounds().needing_attention assert [s.label for s in flagged] == ['Width model_c1'] - analysis.bayesian.suggest_bounds(absolute_floor=1.0).apply() assert not analysis.bayesian.suggest_bounds().needing_attention @@ -221,25 +219,27 @@ def test_recovers_a_straight_line_through_the_widths(self): warnings.simplefilter('ignore') results = analysis.bayesian.sample(**SAMPLE_KWARGS) - # EXPECT a column per polynomial coefficient, and a readable summary + # EXPECT the posterior recovers the generating polynomial within a few posterior + # standard deviations (a 68% interval would exclude the truth too often to be strict), + # with a column per coefficient and a readable, collision-free summary + summary = analysis.bayesian.summary() + for name, truth in ( + ('Width model_c0', 0.8), + ('Width model_c1', 0.0), + ('Width model_c2', 0.4), + ): + entry = summary[name] + spread = max(entry.minus, entry.plus) + assert abs(entry.median - truth) < 4 * spread assert results.draws.shape[1] == len(analysis._chain_parameters()) - names = [entry.name for entry in analysis.bayesian.summary()] + names = [entry.name for entry in summary] assert len(set(names)) == len(names) class TestAggregatedIndependentChains: - def test_summary_gathers_the_real_per_q_chains(self): - # WHEN each Q is sampled on its own - analysis = build_analysis() - analysis.fit(fit_method='independent') - for analysis1d in analysis.analysis_list: - analysis1d.bayesian.suggest_bounds().apply() - - with warnings.catch_warnings(): - warnings.simplefilter('ignore') - analysis.bayesian.sample(fit_method='independent', **SAMPLE_KWARGS) - + def test_summary_gathers_the_real_per_q_chains(self, independently_sampled): # THEN + analysis, _ = independently_sampled summary = analysis.bayesian.summary() # EXPECT one table covering every Q, and the widths still recovered @@ -249,22 +249,23 @@ def test_summary_gathers_the_real_per_q_chains(self): spread = max(entry.minus, entry.plus) assert abs(entry.median - true_width(Q_VALUES[q_index])) < 4 * spread - def test_median_applies_each_chain_to_its_own_q(self): - # WHEN - analysis = build_analysis() - analysis.fit(fit_method='independent') - for analysis1d in analysis.analysis_list: - analysis1d.bayesian.suggest_bounds().apply() - - with warnings.catch_warnings(): - warnings.simplefilter('ignore') - analysis.bayesian.sample(fit_method='independent', **SAMPLE_KWARGS) + def test_median_applies_each_chain_to_its_own_q(self, independently_sampled): + # WHEN the fixture is module-scoped, so the values moved here are restored afterwards + analysis, _ = independently_sampled + parameters = [p for a in analysis.analysis_list for p in a.get_free_parameters()] + saved_values = [(p, float(p.value)) for p in parameters] - # THEN - changed = analysis.bayesian.set_parameters_to_median() + try: + # THEN + changed = analysis.bayesian.set_parameters_to_median() - # EXPECT every Q's parameters land on that Q's own median - assert len(changed) == sum(len(a.get_free_parameters()) for a in analysis.analysis_list) - summary = analysis.bayesian.summary() - for entry in summary: - assert entry.value == pytest.approx(entry.median, rel=1e-6) + # EXPECT every Q's parameters land on that Q's own median + assert len(changed) == sum( + len(a.get_free_parameters()) for a in analysis.analysis_list + ) + summary = analysis.bayesian.summary() + for entry in summary: + assert entry.value == pytest.approx(entry.median, rel=1e-6) + finally: + for parameter, value in saved_values: + parameter.value = value diff --git a/tests/unit/easydynamics/analysis/test_posterior_sampling.py b/tests/unit/easydynamics/analysis/test_posterior_sampling.py index 6777e1a8c..40848f105 100644 --- a/tests/unit/easydynamics/analysis/test_posterior_sampling.py +++ b/tests/unit/easydynamics/analysis/test_posterior_sampling.py @@ -1189,6 +1189,18 @@ def test_invalid_fit_method_raises(self, multi_q_analysis): with pytest.raises(ValueError, match='Invalid fit method'): multi_q_analysis.bayesian.sample(fit_method='nonsense') + def test_negative_q_index_raises(self, multi_q_analysis): + # THEN EXPECT a refusal, rather than silently wrapping around to the last Q + with pytest.raises(IndexError, match='non-negative'): + multi_q_analysis.bayesian.sample(fit_method='independent', Q_index=-1, samples=10) + + def test_corner_q_index_is_validated(self, multi_q_analysis): + # THEN EXPECT both ends of the range are checked before any chain is looked up + with pytest.raises(IndexError, match='non-negative'): + multi_q_analysis.bayesian.plot_corner(Q_index=-1) + with pytest.raises(IndexError, match='out of bounds'): + multi_q_analysis.bayesian.plot_corner(Q_index=99) + def test_missing_q_values_raises(self): # WHEN multi_q_analysis = edyn.Analysis(display_name='Empty') @@ -1385,6 +1397,119 @@ def test_a_simultaneous_chain_still_takes_precedence(self, multi_q_analysis): assert len(multi_q_analysis.bayesian.summary()) == len(parameters) multi_q_analysis.bayesian.plot_corner() + def test_a_fresh_per_q_chain_wins_after_a_simultaneous_run(self, multi_q_analysis): + # WHEN an independent run of one Q follows a simultaneous one + bound_all_chain(multi_q_analysis) + parameters = multi_q_analysis._chain_parameters() + target = multi_q_analysis.analysis_list[2] + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_chain_results(parameters) + multi_q_analysis.bayesian.sample(fit_method='simultaneous', samples=10) + + # THEN + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_chain_results( + target.get_free_parameters() + ) + multi_q_analysis.bayesian.sample(fit_method='independent', Q_index=2, samples=10) + + # EXPECT the fresh per-Q chain is what summary() reports, not the stale simultaneous one + summary = multi_q_analysis.bayesian.summary() + assert len(summary) == len(target.get_free_parameters()) + assert all('Q_index=2' in entry.name for entry in summary) + + def test_gathered_summary_uses_the_per_q_saved_labels(self, multi_q_analysis): + # WHEN the per-Q chains look freshly loaded from disk in a new session: foreign column + # names, matched to parameters only through each per-Q sampler's saved labels + self._sample_independently(multi_q_analysis) + for q_index, analysis1d in enumerate(multi_q_analysis.analysis_list): + sampler = analysis1d.bayesian + name_map = analysis1d._parameter_labels().name_map() + foreign = [f'Loaded_{q_index}_{i}' for i in range(len(sampler.results.param_names))] + sampler._saved_labels = { + foreign_name: name_map[unique_name] + for foreign_name, unique_name in zip( + foreign, sampler.results.param_names, strict=True + ) + } + sampler.results.param_names = foreign + + # THEN + summary = multi_q_analysis.bayesian.summary() + + # EXPECT every column resolves to its parameter: Q-qualified names, real units and finite + # values, rather than raw column names with no unit and NaN + expected = sum(len(a.get_free_parameters()) for a in multi_q_analysis.analysis_list) + assert len(summary) == expected + assert all('Q_index=' in entry.name for entry in summary) + assert all(entry.unit != '' for entry in summary) + assert all(np.isfinite(entry.value) for entry in summary) + + def test_the_slider_path_forwards_plot_kwargs(self, multi_q_analysis): + # WHEN + self._sample_independently(multi_q_analysis) + + # THEN + with ( + patch('easydynamics.analysis.posterior_sampling._in_notebook', return_value=True), + patch('easydynamics.utils.posterior_plotting.corner_with_slider') as slider, + ): + multi_q_analysis.bayesian.plot_corner(bins=13) + + # EXPECT the kwargs the docstring promises to forward reach the slider's corner plots + assert slider.call_args.kwargs['bins'] == 13 + + ############# + # Extending and persistence + ############# + + def test_extend_after_an_independent_run_points_at_the_per_q_chains(self, multi_q_analysis): + # WHEN an independent run follows a simultaneous one, so this sampler still holds the old + # simultaneous chain while the latest chains live per Q + bound_all_chain(multi_q_analysis) + parameters = multi_q_analysis._chain_parameters() + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_chain_results(parameters) + multi_q_analysis.bayesian.sample(fit_method='simultaneous', samples=10) + self._sample_independently(multi_q_analysis) + + # THEN EXPECT the error says where the chains are, rather than extending the stale chain + # or misdiagnosing a failed run + with pytest.raises(RuntimeError, match=r'analysis_list\[Q_index\]\.bayesian\.extend'): + multi_q_analysis.bayesian.extend() + + def test_save_after_an_independent_run_refuses_the_stale_chain( + self, multi_q_analysis, tmp_path + ): + # WHEN an independent run follows a simultaneous one + bound_all_chain(multi_q_analysis) + parameters = multi_q_analysis._chain_parameters() + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_chain_results(parameters) + multi_q_analysis.bayesian.sample(fit_method='simultaneous', samples=10) + stale_sampler = sampler_class.return_value + self._sample_independently(multi_q_analysis) + + # THEN EXPECT save refuses, rather than silently writing the stale simultaneous chain + with pytest.raises(RuntimeError, match='no simultaneous chain here to save'): + multi_q_analysis.bayesian.save(str(tmp_path / 'chain')) + stale_sampler.save.assert_not_called() + + def test_extend_after_a_failed_simultaneous_run_keeps_the_failed_run_message( + self, multi_q_analysis + ): + # WHEN a simultaneous run fails after building the sampler, with no per-Q chains anywhere + bound_all_chain(multi_q_analysis) + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = RuntimeError('boom') + with pytest.raises(RuntimeError, match='boom'): + multi_q_analysis.bayesian.sample(fit_method='simultaneous', samples=10) + + # THEN EXPECT the genuine failed-run diagnosis, not the pointer at per-Q chains + with pytest.raises(RuntimeError, match='left no results'): + multi_q_analysis.bayesian.extend() + def test_only_the_sampled_q_indices_are_gathered(self, multi_q_analysis): # WHEN just one Q index is sampled target = multi_q_analysis.analysis_list[1] From 6ccc0bafd71b61f6d7398b8392734f6c9326ff4f Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Mon, 17 Aug 2026 17:08:24 +0200 Subject: [PATCH 22/29] Add marginal posteriors, correlation heatmaps and sampling progress - plot_marginal(parameter) renders one parameter's posterior histogram with the median and the 16/84 percentile interval summary() reports, resolving labels the same way sample(parameters=...) does - plot_correlations() renders the Pearson correlation matrix of the chain with annotated cells, a diverging colormap and masked cells for constant columns - sample(progress=True) and extend(progress=True) report sampling progress through the Sampler's progress_callback, closing the line with an explicit done marker because BUMPS' own step estimate assumes the wrong chain count - the 95 percent predictive band needed no change: credible_interval already exists on plot_posterior_predictive Co-Authored-By: Claude Fable 5 --- .../analysis/posterior_sampling.py | 265 +++++++++++++++++- src/easydynamics/utils/posterior_plotting.py | 170 +++++++++++ .../fitting/test_bayesian_sampling.py | 17 ++ .../analysis/test_posterior_sampling.py | 197 +++++++++++++ .../utils/test_posterior_plotting.py | 158 +++++++++++ 5 files changed, 794 insertions(+), 13 deletions(-) diff --git a/src/easydynamics/analysis/posterior_sampling.py b/src/easydynamics/analysis/posterior_sampling.py index 44ee7dc85..d1fe04103 100644 --- a/src/easydynamics/analysis/posterior_sampling.py +++ b/src/easydynamics/analysis/posterior_sampling.py @@ -230,6 +230,7 @@ def sample( thin: int = 10, population: int | None = None, parameters: list[Parameter] | list[str] | None = None, + progress: bool = False, **sampler_options: dict[str, Any], ) -> SamplingResults: """ @@ -255,6 +256,10 @@ def sample( 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``. @@ -270,18 +275,28 @@ def sample( 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. """ - return self._run( - parameters=parameters, - run=lambda sampler: sampler.sample( - samples=samples, burn=burn, thin=thin, population=population, **sampler_options - ), - ) + reporter = _install_progress_reporter(progress, sampler_options) + try: + results = self._run( + parameters=parameters, + run=lambda sampler: sampler.sample( + samples=samples, burn=burn, thin=thin, population=population, **sampler_options + ), + ) + except BaseException: + if reporter is not None: + reporter.close(completed=False) + raise + if reporter is not None: + reporter.close(completed=True) + 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: """ @@ -296,6 +311,9 @@ def extend( 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. @@ -319,13 +337,22 @@ def extend( """ if self._sampler is None: raise RuntimeError('No chain to extend. Call sample() or load() first.') - return self._run( - parameters=parameters, - run=lambda sampler: sampler.extend( - additional_samples=additional_samples, thin=thin, **sampler_options - ), - reuse_sampler=True, - ) + reporter = _install_progress_reporter(progress, sampler_options) + try: + results = self._run( + parameters=parameters, + run=lambda sampler: sampler.extend( + additional_samples=additional_samples, thin=thin, **sampler_options + ), + reuse_sampler=True, + ) + except BaseException: + if reporter is not None: + reporter.close(completed=False) + raise + if reporter is not None: + reporter.close(completed=True) + return results def _run( self, @@ -772,6 +799,66 @@ def plot_corner(self, **kwargs: dict[str, Any]) -> Figure: **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, @@ -912,6 +999,54 @@ def _resolve(self, results: SamplingResults) -> list[Parameter | None]: """ 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}. ' + f'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. @@ -989,6 +1124,110 @@ def _warn_about_held_parameters(labels: object, held_fixed: list[Parameter]) -> ) +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') + print(flush=True) + + 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. + """ + print(f'\r{line.ljust(self._line_length)}', end='', flush=True) + 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. diff --git a/src/easydynamics/utils/posterior_plotting.py b/src/easydynamics/utils/posterior_plotting.py index bfa26e987..bc2b4c280 100644 --- a/src/easydynamics/utils/posterior_plotting.py +++ b/src/easydynamics/utils/posterior_plotting.py @@ -10,10 +10,12 @@ 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: @@ -201,6 +203,150 @@ def plot_corner( 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. + + 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. + + Raises + ------ + ValueError + If ``draws`` is not two-dimensional or is empty, or if ``names`` does not have one entry + per column. + """ + 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, @@ -313,6 +459,30 @@ def _column_limits(draws: np.ndarray) -> list[tuple[float, float]]: 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. diff --git a/tests/integration/fitting/test_bayesian_sampling.py b/tests/integration/fitting/test_bayesian_sampling.py index 9b5997388..f341a1c84 100644 --- a/tests/integration/fitting/test_bayesian_sampling.py +++ b/tests/integration/fitting/test_bayesian_sampling.py @@ -197,6 +197,23 @@ def test_plots_render(self, sampled_analysis): 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() diff --git a/tests/unit/easydynamics/analysis/test_posterior_sampling.py b/tests/unit/easydynamics/analysis/test_posterior_sampling.py index 6f2cad85e..f8e19b798 100644 --- a/tests/unit/easydynamics/analysis/test_posterior_sampling.py +++ b/tests/unit/easydynamics/analysis/test_posterior_sampling.py @@ -629,6 +629,203 @@ def test_predictive_omits_error_bars_when_the_data_has_no_variances(self): # 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 ############# diff --git a/tests/unit/easydynamics/utils/test_posterior_plotting.py b/tests/unit/easydynamics/utils/test_posterior_plotting.py index 5f341a24c..c470bdd20 100644 --- a/tests/unit/easydynamics/utils/test_posterior_plotting.py +++ b/tests/unit/easydynamics/utils/test_posterior_plotting.py @@ -10,6 +10,8 @@ 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 @@ -130,6 +132,162 @@ def test_columns_share_limits_between_histogram_and_hexbin_panels(self, draws): 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 From 74fa676de228c321d3ea688ac33898e6fedc956c Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Mon, 17 Aug 2026 17:53:48 +0200 Subject: [PATCH 23/29] Give every posterior plot a Q slider over independent chains After independent per-Q sampling the multi-Q sampler now presents a Q slider instead of refusing: - plot_posterior_predictive builds the per-Q data, median and credible band into a scipp DataGroup and renders it through plopp exactly like plot_data_and_model; plopp cannot shade a band on sliced lines, so the slider view draws labelled band edges while the Q_index path keeps the shaded band - plot_trace, plot_marginal and plot_correlations take Q_index for a single figure, show a slider in a notebook, and otherwise name the sampled Q indices - the matplotlib sliders render every figure once up front and only swap PNG bytes on a move, so dragging tracks smoothly with continuous updates instead of re-rendering per change - per-Q energy grids are NaN-padded onto the common grid through the finite mask, so masked points draw as gaps Co-Authored-By: Claude Fable 5 --- .../analysis/posterior_sampling.py | 402 +++++++++++++++++- src/easydynamics/utils/posterior_plotting.py | 232 ++++++++-- .../fitting/test_bayesian_sampling_multi_q.py | 58 +++ .../analysis/test_posterior_sampling.py | 273 +++++++++++- .../utils/test_posterior_plotting.py | 189 ++++++++ 5 files changed, 1092 insertions(+), 62 deletions(-) diff --git a/src/easydynamics/analysis/posterior_sampling.py b/src/easydynamics/analysis/posterior_sampling.py index caf34b153..4470e3343 100644 --- a/src/easydynamics/analysis/posterior_sampling.py +++ b/src/easydynamics/analysis/posterior_sampling.py @@ -29,6 +29,7 @@ from easydynamics.analysis.posterior import suggest_bounds_for_parameters from easydynamics.analysis.posterior import summarize_draws from easydynamics.analysis.posterior import unbounded_parameters +from easydynamics.analysis.posterior_labels import ParameterLabels from easydynamics.utils.utils import _in_notebook from easydynamics.utils.utils import verify_Q_index @@ -40,9 +41,9 @@ from easyscience.variable import Parameter from ipywidgets import VBox from matplotlib.figure import Figure + from plopp.backends.matplotlib.figure import InteractiveFigure from easydynamics.analysis.posterior import BoundsSuggestions - 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(). @@ -1435,35 +1436,410 @@ def plot_corner(self, Q_index: int | None = None, **kwargs: dict[str, Any]) -> F } return corner_with_slider(chains, title=self._analysis.display_name, **kwargs) - def plot_trace(self, **kwargs: dict[str, Any]) -> Figure: + def plot_trace(self, Q_index: int | None = None, **kwargs: dict[str, Any]) -> Figure | VBox: """ Plot the chain trace of each sampled parameter. - Only available for a simultaneous chain, since the per-Q chains are separate runs of - different lengths rather than one trace. + A simultaneous chain is one trace and is drawn directly. After independent sampling each + Q index has its own chain, so the traces are stepped through one at a time: pick one with + ``Q_index``, or leave it out in a notebook to get a slider. Parameters ---------- + Q_index : int | None, default=None + Which Q index to plot, when the chains are per-Q. If None, a slider is returned. Not + used for a simultaneous chain, which is a single trace already. **kwargs : dict[str, Any] Forwarded to :func:`easydynamics.utils.posterior_plotting.plot_trace`. Returns ------- - Figure - The matplotlib Figure. + Figure | VBox + The matplotlib Figure, or an ipywidgets box with a Q slider. Raises ------ + IndexError + If Q_index is negative or out of range. RuntimeError - If only independent per-Q chains exist. + If a slider is asked for outside a notebook, or nothing has been sampled yet. + TypeError + If Q_index is not an int or None. """ - if self._results is None and self.results_per_q is not None: - raise RuntimeError( - 'Each Q index has its own chain, so there is no single trace to draw. Use ' - 'analysis.analysis_list[Q_index].bayesian.plot_trace() for one of them, or sample ' - "with fit_method='simultaneous'." + verify_Q_index(Q_index=Q_index, Q=self._analysis.Q, allow_none=True) + per_q = self.results_per_q + if self._results is not None or per_q is None: + return super().plot_trace(**kwargs) + if Q_index is not None: + return self._per_q()[Q_index].bayesian.plot_trace(**kwargs) + self._require_notebook_for_slider(per_q) + return self._figures_with_q_slider( + per_q, lambda analysis1d: analysis1d.bayesian.plot_trace(**kwargs) + ) + + def plot_marginal( + self, + parameter: Parameter | str, + Q_index: int | None = None, + **kwargs: dict[str, Any], + ) -> Figure | VBox: + """ + Plot the marginal posterior distribution of a single sampled parameter. + + A simultaneous chain holds every Q's parameters under Q-qualified labels, so the label + picks the Q as well (``'Gaussian width (Q_index=1)'``). After independent sampling the + chains are per-Q and the parameter goes by its plain label in each; pick a chain with + ``Q_index``, or leave it out in a notebook to step through the Q values with a slider. + + Parameters + ---------- + parameter : Parameter | str + The parameter to plot, as a Parameter object or its label. On the slider path a + Parameter object is resolved to its display name first, so the matching parameter of + every Q is shown even though the object itself belongs to one Q. + Q_index : int | None, default=None + Which Q index to plot, when the chains are per-Q. If None, a slider is returned. Not + used for a simultaneous chain, whose labels carry the Q index already. + **kwargs : dict[str, Any] + Forwarded to :func:`easydynamics.utils.posterior_plotting.plot_marginal`. + + Returns + ------- + Figure | VBox + The matplotlib Figure, or an ipywidgets box with a Q slider. + + Raises + ------ + IndexError + If Q_index is negative or out of range. + RuntimeError + If a slider is asked for outside a notebook, or nothing has been sampled yet. + TypeError + If Q_index is not an int or None. + ValueError + If the parameter matches no sampled chain column. + """ + verify_Q_index(Q_index=Q_index, Q=self._analysis.Q, allow_none=True) + per_q = self.results_per_q + if self._results is not None or per_q is None: + return super().plot_marginal(parameter, **kwargs) + if Q_index is not None: + return self._per_q()[Q_index].bayesian.plot_marginal(parameter, **kwargs) + self._require_notebook_for_slider(per_q) + # Resolved to a display name up front, because a Parameter object belongs to one Q only + # and every chain must find its own copy under the shared name. + label = ( + parameter + if isinstance(parameter, str) + else self._shared_display_name(parameter, per_q) + ) + return self._figures_with_q_slider( + per_q, lambda analysis1d: analysis1d.bayesian.plot_marginal(label, **kwargs) + ) + + def plot_correlations( + self, Q_index: int | None = None, **kwargs: dict[str, Any] + ) -> Figure | VBox: + """ + Plot the Pearson correlation matrix of the sampled parameters. + + A simultaneous chain gives one matrix over every Q's parameters at once. After + independent sampling no draw pairs one Q with another, so there is one matrix per chain: + pick one with ``Q_index``, or leave it out in a notebook to get a slider. + + Parameters + ---------- + Q_index : int | None, default=None + Which Q index to plot, when the chains are per-Q. If None, a slider is returned. Not + used for a simultaneous chain, which already covers every Q. + **kwargs : dict[str, Any] + Forwarded to :func:`easydynamics.utils.posterior_plotting.plot_correlations`. + + Returns + ------- + Figure | VBox + The matplotlib Figure, or an ipywidgets box with a Q slider. + + Raises + ------ + IndexError + If Q_index is negative or out of range. + RuntimeError + If a slider is asked for outside a notebook, or nothing has been sampled yet. + TypeError + If Q_index is not an int or None. + """ + verify_Q_index(Q_index=Q_index, Q=self._analysis.Q, allow_none=True) + per_q = self.results_per_q + if self._results is not None or per_q is None: + return super().plot_correlations(**kwargs) + if Q_index is not None: + return self._per_q()[Q_index].bayesian.plot_correlations(**kwargs) + self._require_notebook_for_slider(per_q) + return self._figures_with_q_slider( + per_q, lambda analysis1d: analysis1d.bayesian.plot_correlations(**kwargs) + ) + + def plot_posterior_predictive( + self, + n_draws: int = 200, + credible_interval: float = 68.0, + Q_index: int | None = None, + **kwargs: dict[str, Any], + ) -> Figure | InteractiveFigure: + """ + Plot the data against the credible band implied by the posterior. + + After independent sampling each Q has its own chain, and its own band: pick one with + ``Q_index`` for a single matplotlib figure, or leave it out in a notebook to get a plopp + figure with a Q slider, looking and handling exactly like + ``Analysis.plot_data_and_model``. Plopp draws no filled band, so the slider view shows + the posterior median with a dashed line along each band edge instead of a shaded band. + + Parameters + ---------- + n_draws : int, default=200 + How many posterior draws to evaluate the model for, per Q on the slider path. Each + costs a full model evaluation. + credible_interval : float, default=68.0 + Width of the credible band, as a percentage. + Q_index : int | None, default=None + Which Q index to plot, when the chains are per-Q. If None, a slider is returned. + **kwargs : dict[str, Any] + Forwarded to :func:`easydynamics.utils.posterior_plotting.plot_posterior_predictive` + for a single figure, or to + :func:`easydynamics.utils.posterior_plotting.predictive_with_slider` for the slider. + + Returns + ------- + Figure | InteractiveFigure + The matplotlib Figure for one Q, or the plopp figure with a Q slider. + + Raises + ------ + IndexError + If Q_index is negative or out of range. + NotImplementedError + If the latest chain is simultaneous: it binds every dataset at once, and no per-Q + chain exists for Q_index to pick out. + RuntimeError + If a slider is asked for outside a notebook, or nothing has been sampled yet. + TypeError + If Q_index is not an int or None. + ValueError + If n_draws is not a positive integer, or credible_interval is out of range. + """ + 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}.') + verify_Q_index(Q_index=Q_index, Q=self._analysis.Q, allow_none=True) + per_q = self.results_per_q + if self._results is not None or per_q is None: + return super().plot_posterior_predictive( + n_draws=n_draws, credible_interval=credible_interval, **kwargs ) - return super().plot_trace(**kwargs) + if Q_index is not None: + return self._per_q()[Q_index].bayesian.plot_posterior_predictive( + n_draws=n_draws, credible_interval=credible_interval, **kwargs + ) + self._require_notebook_for_slider(per_q) + return self._predictive_with_q_slider(per_q, n_draws, credible_interval, **kwargs) + + ############# + # Sliders over the independent per-Q chains + ############# + + def _require_notebook_for_slider(self, per_q: list[SamplingResults | None]) -> None: + """ + Refuse the slider path outside a notebook, naming the sampled Q indices. + + Parameters + ---------- + per_q : list[SamplingResults | None] + The per-Q chains, None where a Q index has not been sampled. + + Raises + ------ + RuntimeError + If not running in a Jupyter notebook. + """ + if _in_notebook(): + return + sampled = [index for index, result in enumerate(per_q) if result is not None] + raise RuntimeError( + f'Each Q index has its own chain, and the slider needs a Jupyter notebook. ' + f'Pass Q_index to plot one of them; sampled Q indices are {sampled}.' + ) + + def _figures_with_q_slider( + self, + per_q: list[SamplingResults | None], + plot_one: Callable[[object], Figure], + ) -> VBox: + """ + Render one figure per sampled Q index and put them behind a slider. + + Only the Q indices that actually hold a chain get a figure, so the slider cannot land on + an empty position. Each figure carries its per-Q Analysis' own display name, which names + the Q index. + + Parameters + ---------- + per_q : list[SamplingResults | None] + The per-Q chains, None where a Q index has not been sampled. + plot_one : Callable[[object], Figure] + Renders the figure for one per-Q Analysis. + + Returns + ------- + VBox + An ipywidgets box with the pre-rendered figures behind a Q slider. + """ + from easydynamics.utils.posterior_plotting import figures_with_slider + + figures = {} + for analysis1d, result in zip(self._per_q(), per_q, strict=True): + if result is None: + continue + figures[analysis1d.Q_index] = plot_one(analysis1d) + return figures_with_slider(figures) + + def _shared_display_name( + self, + parameter: Parameter, + per_q: list[SamplingResults | None], + ) -> str: + """ + Find the display name a Parameter goes by within its own Q's chain. + + The same model is repeated per Q, so the name one chain reports a parameter under is the + name every other chain reports its own copy under. Resolving through it lets a slider + show the matching marginal at every Q even though the Parameter object belongs to one. + + Parameters + ---------- + parameter : Parameter + The parameter to resolve. + per_q : list[SamplingResults | None] + The per-Q chains, None where a Q index has not been sampled. + + Returns + ------- + str + The display name of the chain column holding the parameter's draws. + + Raises + ------ + ValueError + If no sampled chain holds draws of the parameter. + """ + for analysis1d, result in zip(self._per_q(), per_q, strict=True): + if result is None: + continue + # The same labels that Q's own sampler reports its chain under: its free parameters, + # unqualified, since a single Q has one copy of each. + labels = ParameterLabels(analysis1d.get_free_parameters()) + if any( + candidate.unique_name == parameter.unique_name for candidate in labels.parameters + ): + return labels.label(parameter) + name = getattr(parameter, 'name', '?') + raise ValueError(f'No sampled parameter named {name!r} in any per-Q chain.') + + def _predictive_with_q_slider( + self, + per_q: list[SamplingResults | None], + n_draws: int, + credible_interval: float, + **kwargs: dict[str, Any], + ) -> InteractiveFigure: + """ + Build the posterior-predictive figure with a Q slider from the per-Q chains. + + Each sampled Q contributes its data, median prediction and band edges, computed from its + own chain with the same machinery the single-Q figure uses. Rows are laid out on the + experiment's common energy grid; a Q's masked-away points stay NaN, leaving a gap rather + than inventing a value there. + + Parameters + ---------- + per_q : list[SamplingResults | None] + The per-Q chains, None where a Q index has not been sampled. + n_draws : int + How many posterior draws to evaluate the model for, per Q. + credible_interval : float + Width of the credible band, as a percentage. + **kwargs : dict[str, Any] + Forwarded to :func:`easydynamics.utils.posterior_plotting.predictive_with_slider`. + + Returns + ------- + InteractiveFigure + The plopp figure with its Q slider. + + Raises + ------ + ValueError + If credible_interval is not between 0 and 100. + """ + from easydynamics.utils.posterior_plotting import predictive_with_slider + + if not 0 < credible_interval < 100: + raise ValueError( + f'credible_interval must be between 0 and 100. Got {credible_interval}.' + ) + + energy = self._analysis.energy + q = self._analysis.Q + energy_values = np.asarray(energy.values, dtype=float) + + # As in the single-Q figure: without variances the weights are all-ones placeholders, and + # inverting them would fabricate error bars the data never had. + experiment = getattr(self._analysis, 'experiment', None) + has_variances = experiment is None or getattr(experiment, 'has_variances', True) + 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('ylabel', 'Intensity' if y_unit is None else f'Intensity ({y_unit})') + + sampled = [ + analysis1d + for analysis1d, result in zip(self._per_q(), per_q, strict=True) + if result is not None + ] + shape = (len(sampled), len(energy_values)) + data = np.full(shape, np.nan) + variances = np.full(shape, np.nan) if has_variances else None + lower = np.full(shape, np.nan) + median = np.full(shape, np.nan) + upper = np.full(shape, np.nan) + tail = (100.0 - credible_interval) / 2.0 + for row, analysis1d in enumerate(sampled): + _, y, weights, mask = analysis1d.experiment.extract_x_y_weights_only_finite( + Q_index=analysis1d.Q_index + ) + predictions = analysis1d.bayesian.predictions(n_draws) + # The mask places every finite point back on the common grid, so the padding stays + # NaN wherever a point was masked away. + data[row, mask] = np.asarray(y) + if variances is not None: + variances[row, mask] = 1.0 / np.asarray(weights) ** 2 + lower[row, mask], median[row, mask], upper[row, mask] = np.percentile( + predictions, [tail, 50.0, 100.0 - tail], axis=0 + ) + + return predictive_with_slider( + energy=energy_values, + q_values=np.asarray([float(q.values[a.Q_index]) for a in sampled]), + y=data, + lower=lower, + median=median, + upper=upper, + y_variances=variances, + energy_unit=str(energy.unit), + q_unit=str(q.unit), + title=self._analysis.display_name, + credible_interval=credible_interval, + **kwargs, + ) def _require_results(self) -> SamplingResults: """ diff --git a/src/easydynamics/utils/posterior_plotting.py b/src/easydynamics/utils/posterior_plotting.py index 28bb1550e..f18266294 100644 --- a/src/easydynamics/utils/posterior_plotting.py +++ b/src/easydynamics/utils/posterior_plotting.py @@ -10,6 +10,7 @@ from __future__ import annotations +import io import warnings from typing import TYPE_CHECKING from typing import Any @@ -22,6 +23,7 @@ if TYPE_CHECKING: from ipywidgets import VBox from matplotlib.figure import Figure + from plopp.backends.matplotlib.figure import InteractiveFigure def plot_trace( @@ -593,6 +595,66 @@ def _verify_draws(draws: np.ndarray, names: list[str]) -> None: ) +def figures_with_slider(figures: dict[int, Figure], description: str = 'Q index') -> VBox: + """ + Show one pre-rendered figure at a time, with a slider choosing which one. + + Every figure is rendered to PNG bytes once, up front, and the slider callback only swaps the + stored bytes into an image widget. Moving the slider therefore costs no matplotlib work at + all, which keeps it as responsive as the plopp slider on the data plots; re-rendering a + figure on every move is what made the previous slider feel sluggish. + + The figures are closed after rendering, so no backend draws them a second time. + + Parameters + ---------- + figures : dict[int, Figure] + Mapping of slider position to the matplotlib Figure shown there. Only these positions are + offered, so the slider cannot land on an index with nothing to show. + description : str, default='Q index' + Label shown next to the slider. + + Returns + ------- + VBox + An ipywidgets box holding the image and, under it, the slider. + + Raises + ------ + ValueError + If no figures are given. + """ + import ipywidgets as widgets + + if not figures: + raise ValueError('No figures to show.') + + indices = sorted(figures) + rendered = {} + for index in indices: + figure = figures[index] + buffer = io.BytesIO() + figure.savefig(buffer, format='png', bbox_inches='tight') + rendered[index] = buffer.getvalue() + # Rendered to bytes already, so the figure is closed rather than left for a backend to + # draw a second time. + plt.close(figure) + + image = widgets.Image(value=rendered[indices[0]], format='png') + image.layout.max_width = '100%' + # Swapping stored bytes is instant, so the image can follow the slider continuously; there is + # no need for the release-to-update behaviour an expensive redraw would force. + slider = widgets.SelectionSlider( + options=indices, + value=indices[0], + description=description, + continuous_update=True, + ) + slider.observe(lambda change: setattr(image, 'value', rendered[change['new']]), names='value') + # Slider under the figure, matching where plopp puts its slicer controls. + return widgets.VBox([image, slider]) + + def corner_with_slider( chains: dict[int, dict], title: str | None = None, @@ -603,7 +665,8 @@ def corner_with_slider( Chains sampled separately share no draws, so there is no joint distribution across them to plot. Stepping through them one at a time shows the correlations that were actually sampled, - which is what a single combined figure could not do honestly. + which is what a single combined figure could not do honestly. The figures are pre-rendered + through :func:`figures_with_slider`, so the slider moves without re-drawing anything. Parameters ---------- @@ -618,53 +681,152 @@ def corner_with_slider( Returns ------- VBox - An ipywidgets box holding the slider and the figure. + An ipywidgets box holding the figure and the slider. Raises ------ ValueError If no chains are given. """ - import ipywidgets as widgets - if not chains: raise ValueError('No chains to plot.') - indices = sorted(chains) - output = widgets.Output() - - def draw(index: int) -> None: - """ - Render the corner plot for one chain. - - Parameters - ---------- - index : int - The chain to draw. - """ - chain = chains[index] - figure = plot_corner( + figures = { + index: plot_corner( draws=chain['draws'], names=chain['names'], units=chain.get('units'), title=title if title is None else f'{title} (Q index {index})', **kwargs, ) - # append_display_data rather than the `with output:` context manager, which captures - # nothing under some kernels and would leave the slider with a blank panel beside it. - output.outputs = () - output.append_display_data(figure) - # Rendered into the widget already, so the figure is closed rather than left for a backend - # to draw a second time. - plt.close(figure) + for index, chain in chains.items() + } + return figures_with_slider(figures) - slider = widgets.SelectionSlider( - options=indices, - value=indices[0], - description='Q index', - continuous_update=False, - ) - slider.observe(lambda change: draw(change['new']), names='value') - draw(indices[0]) - # Slider under the figure, matching where plopp puts its slicer controls. - return widgets.VBox([output, slider]) + +def predictive_with_slider( + energy: np.ndarray, + q_values: np.ndarray, + y: np.ndarray, + lower: np.ndarray, + median: np.ndarray, + upper: np.ndarray, + y_variances: np.ndarray | None = None, + energy_unit: str | None = None, + q_unit: str | None = None, + ylabel: str | None = None, + title: str | None = None, + credible_interval: float = 68.0, + **kwargs: dict[str, Any], +) -> InteractiveFigure: + """ + Plot per-Q posterior-predictive bands behind a plopp Q slider. + + Built on ``plopp.slicer`` over a scipp DataGroup with a Q dimension, so the figure looks and + handles exactly like ``Analysis.plot_data_and_model``: the data with its error bars, the model + curves on top, and a Q slider underneath. Plopp draws no filled band for sliced data -- its + only spread representation is variance-based error bars -- so the credible band is drawn as + the posterior median with a dashed line along each band edge, labelled with the interval. + + Rows are laid out on one common energy grid; where a Q has no point (masked or never + measured), NaN leaves a gap in the lines rather than inventing a value. + + Parameters + ---------- + energy : np.ndarray + The common energy grid, one column per point. + q_values : np.ndarray + The Q value of each row, shown on the slider. + y : np.ndarray + Observed values, shape ``(len(q_values), len(energy))``, NaN where a Q has no point. + lower : np.ndarray + Lower band edge per Q, same shape as ``y``. + median : np.ndarray + Posterior median prediction per Q, same shape as ``y``. + upper : np.ndarray + Upper band edge per Q, same shape as ``y``. + y_variances : np.ndarray | None, default=None + Variances of the observed values, drawn as error bars when given. + energy_unit : str | None, default=None + Unit of the energy grid, shown on the horizontal axis. + q_unit : str | None, default=None + Unit of the Q values, shown beside the slider. + ylabel : str | None, default=None + Label for the dependent axis. + title : str | None, default=None + Figure title. + credible_interval : float, default=68.0 + Width of the credible band the edges enclose, as a percentage, used in their labels. + **kwargs : dict[str, Any] + Forwarded to ``plopp.slicer``, overriding the style defaults. + + Returns + ------- + InteractiveFigure + The plopp figure with its Q slider. + + Raises + ------ + ValueError + If the arrays do not share the shape ``(len(q_values), len(energy))``, or if + ``credible_interval`` is not between 0 and 100. + """ + import plopp as pp + import scipp as sc + + if not 0 < credible_interval < 100: + raise ValueError(f'credible_interval must be between 0 and 100. Got {credible_interval}.') + expected = (len(q_values), len(energy)) + arrays = {'y': y, 'lower': lower, 'median': median, 'upper': upper} + if y_variances is not None: + arrays['y_variances'] = y_variances + for name, array in arrays.items(): + if np.asarray(array).shape != expected: + raise ValueError(f'{name} must have shape {expected}. Got {np.asarray(array).shape}.') + + coords = { + 'Q': sc.array(dims=['Q'], values=np.asarray(q_values, dtype=float), unit=q_unit), + 'energy': sc.array( + dims=['energy'], values=np.asarray(energy, dtype=float), unit=energy_unit + ), + } + + def data_array(values: np.ndarray, variances: np.ndarray | None = None) -> sc.DataArray: + return sc.DataArray( + data=sc.array( + dims=['Q', 'energy'], + values=np.asarray(values, dtype=float), + variances=None if variances is None else np.asarray(variances, dtype=float), + ), + coords=coords, + ) + + lower_key = f'{credible_interval:.0f}% band (lower)' + upper_key = f'{credible_interval:.0f}% band (upper)' + data_group = sc.DataGroup({ + 'Data': data_array(y, y_variances), + 'Posterior median': data_array(median), + lower_key: data_array(lower), + upper_key: data_array(upper), + }) + + # The same styling plot_data_and_model gives its DataGroup: data as open black circles, the + # model curves as lines, with the band edges dashed to read as edges rather than curves. + style = { + 'keep': 'energy', + 'linestyle': {'Data': 'none', 'Posterior median': '-', lower_key: '--', upper_key: '--'}, + 'marker': {'Data': 'o', 'Posterior median': None, lower_key: None, upper_key: None}, + 'color': {'Data': 'black', 'Posterior median': 'C3', lower_key: 'C3', upper_key: 'C3'}, + 'markerfacecolor': {'Data': 'none'}, + } + if title is not None: + style['title'] = title + style.update(kwargs) + + fig = pp.slicer(data_group, **style) + for widget in fig.bottom_bar[0].controls.values(): + widget.slider_toggler.value = '-o-' + if ylabel is not None: + fig.ax.set_ylabel(ylabel) + fig.autoscale() + return fig diff --git a/tests/integration/fitting/test_bayesian_sampling_multi_q.py b/tests/integration/fitting/test_bayesian_sampling_multi_q.py index 814725c5d..74061b57e 100644 --- a/tests/integration/fitting/test_bayesian_sampling_multi_q.py +++ b/tests/integration/fitting/test_bayesian_sampling_multi_q.py @@ -10,6 +10,7 @@ """ import warnings +from unittest.mock import patch import matplotlib as mpl import numpy as np @@ -181,6 +182,63 @@ def test_independent_and_simultaneous_agree_on_the_widths( assert abs(independent.median - simultaneous.median) < 4 * spread +class TestIndependentChainWidgets: + def test_corner_slider_renders_from_the_real_chains(self, independently_sampled): + # WHEN + analysis, _ = independently_sampled + + # THEN + with patch('easydynamics.analysis.posterior_sampling._in_notebook', return_value=True): + widget = analysis.bayesian.plot_corner() + + # EXPECT every real chain pre-rendered behind the slider, and moving the slider swapping + # the stored renderings rather than drawing anything new + image, slider = widget.children + assert list(slider.options) == list(range(len(Q_VALUES))) + assert bytes(image.value).startswith(b'\x89PNG') + first_bytes = image.value + slider.value = 1 + assert image.value != first_bytes + slider.value = 0 + assert image.value == first_bytes + + def test_predictive_slider_renders_from_the_real_chains(self, independently_sampled): + # WHEN the plopp slicer needs an interactive matplotlib backend, switched in for the test + import matplotlib.pyplot as plt + + analysis, _ = independently_sampled + plt.switch_backend('module://ipympl.backend_nbagg') + try: + # THEN + with patch('easydynamics.analysis.posterior_sampling._in_notebook', return_value=True): + fig = analysis.bayesian.plot_posterior_predictive(n_draws=10) + + # EXPECT a plopp figure whose one slider spans the sampled Q values, labelled like + # the single-Q predictive plot + controls = list(fig.bottom_bar[0].controls.values()) + assert len(controls) == 1 + assert controls[0].slider.min == 0 + assert controls[0].slider.max == len(Q_VALUES) - 1 + assert fig.ax.get_ylabel().startswith('Intensity') + finally: + plt.switch_backend('Agg') + + def test_predictive_q_index_plots_one_q_from_its_own_chain(self, independently_sampled): + # WHEN + import matplotlib.pyplot as plt + + analysis, _ = independently_sampled + + # THEN + figure = analysis.bayesian.plot_posterior_predictive(Q_index=1, n_draws=10) + + # EXPECT the single-Q matplotlib figure, with its data and credible band + labels = [text.get_text() for text in figure.axes[0].get_legend().get_texts()] + assert 'Data' in labels + assert any('credible band' in label for label in labels) + plt.close('all') + + class TestParameterAnalysisChain: def test_recovers_a_straight_line_through_the_widths(self): # WHEN the fitted widths are themselves fitted against a model of their Q dependence diff --git a/tests/unit/easydynamics/analysis/test_posterior_sampling.py b/tests/unit/easydynamics/analysis/test_posterior_sampling.py index a142005c4..7069935fe 100644 --- a/tests/unit/easydynamics/analysis/test_posterior_sampling.py +++ b/tests/unit/easydynamics/analysis/test_posterior_sampling.py @@ -1423,6 +1423,81 @@ def test_predictive_is_not_supported_for_multiple_datasets(self, multi_q_analysi with pytest.raises(NotImplementedError, match='single dataset only'): multi_q_analysis.bayesian.plot_posterior_predictive() + def test_predictive_q_index_plots_that_q_alone(self, multi_q_analysis): + # WHEN + self._sample_independently(multi_q_analysis) + + # THEN + figure = multi_q_analysis.bayesian.plot_posterior_predictive(Q_index=1, n_draws=3) + + # EXPECT a single matplotlib figure from that Q's own chain + assert len(figure.axes) == 1 + labels = [text.get_text() for text in figure.axes[0].get_legend().get_texts()] + assert 'Data' in labels + assert any('credible band' in label for label in labels) + + def test_predictive_offers_a_plopp_slider_in_a_notebook(self, multi_q_analysis): + # WHEN + self._sample_independently(multi_q_analysis) + + # THEN the per-Q predictive data is assembled and handed to the plopp-backed slider + with ( + patch('easydynamics.analysis.posterior_sampling._in_notebook', return_value=True), + patch('easydynamics.utils.posterior_plotting.predictive_with_slider') as slicer, + ): + multi_q_analysis.bayesian.plot_posterior_predictive(n_draws=3) + + # EXPECT one row per sampled Q on the common energy grid, each Q's own data in its row, + # a band that encloses its median, and the labelling of plot_data_and_model + kwargs = slicer.call_args.kwargs + n_energy = len(multi_q_analysis.energy.values) + assert kwargs['y'].shape == (len(Q_VALUES), n_energy) + assert list(kwargs['q_values']) == pytest.approx(Q_VALUES) + for row, analysis1d in enumerate(multi_q_analysis.analysis_list): + _, y, _ = analysis1d._sampling_data() + assert kwargs['y'][row] == pytest.approx(np.asarray(y)) + assert np.all(kwargs['lower'] <= kwargs['median']) + assert np.all(kwargs['median'] <= kwargs['upper']) + assert kwargs['y_variances'].shape == (len(Q_VALUES), n_energy) + assert kwargs['energy_unit'] == 'meV' + assert kwargs['q_unit'] == '1/Å' + assert kwargs['ylabel'].startswith('Intensity') + assert kwargs['title'] == multi_q_analysis.display_name + + def test_predictive_pads_a_masked_point_with_nan(self, multi_q_analysis): + # WHEN one Q's data has a NaN point, so its masked grid is shorter than the common grid + multi_q_analysis.experiment.binned_data.values[1, 4] = np.nan + self._sample_independently(multi_q_analysis) + + # THEN + with ( + patch('easydynamics.analysis.posterior_sampling._in_notebook', return_value=True), + patch('easydynamics.utils.posterior_plotting.predictive_with_slider') as slicer, + ): + multi_q_analysis.bayesian.plot_posterior_predictive(n_draws=3) + + # EXPECT the gap stays NaN in every per-Q array, and only there + kwargs = slicer.call_args.kwargs + for key in ('y', 'lower', 'median', 'upper'): + assert np.isnan(kwargs[key][1, 4]) + assert np.isfinite(np.delete(kwargs[key], 4, axis=1)).all() + + def test_predictive_without_a_notebook_or_q_index_says_what_to_do(self, multi_q_analysis): + # WHEN + self._sample_independently(multi_q_analysis) + + # THEN EXPECT it names the sampled Q indices rather than just refusing + with ( + patch('easydynamics.analysis.posterior_sampling._in_notebook', return_value=False), + pytest.raises(RuntimeError, match=r'sampled Q indices are \[0, 1, 2\]'), + ): + multi_q_analysis.bayesian.plot_posterior_predictive() + + def test_predictive_rejects_a_bad_draw_count(self, multi_q_analysis): + # THEN EXPECT the count is checked before any chain is looked up + with pytest.raises(ValueError, match='positive integer'): + multi_q_analysis.bayesian.plot_posterior_predictive(n_draws=0) + ############# # Discoverability ############# @@ -1444,7 +1519,7 @@ def test_operations_needing_one_chain_point_at_the_per_q_chains(self, multi_q_an # THEN EXPECT anything that genuinely needs a single chain says where the chains # actually are, rather than claiming none exist with pytest.raises(RuntimeError, match='analysis_list'): - multi_q_analysis.bayesian.plot_posterior_predictive() + multi_q_analysis.bayesian.predictions() def test_untouched_analysis_still_reports_no_samples(self, multi_q_analysis): # THEN EXPECT the plain message when nothing has been sampled anywhere @@ -1527,19 +1602,36 @@ def test_corner_offers_a_slider_in_a_notebook(self, multi_q_analysis): with patch('easydynamics.analysis.posterior_sampling._in_notebook', return_value=True): widget = multi_q_analysis.bayesian.plot_corner() - # EXPECT a slider over the sampled Q indices, and a panel that actually holds a figure. - # The obvious way to build this captures nothing and leaves the panel blank beside the - # slider, so an empty panel is the regression worth guarding. Which mime type arrives - # depends on the environment: a live kernel renders a PNG, plain pytest only the repr. - # The figure comes first and the slider sits under it, where plopp puts its controls. - panel, slider = widget.children + # EXPECT a slider over the sampled Q indices, and an image that actually holds a + # pre-rendered figure: every chain is rendered to PNG bytes once, up front, so an empty + # image is the regression worth guarding. The figure comes first and the slider sits + # under it, where plopp puts its controls. + image, slider = widget.children assert list(slider.options) == list(range(len(Q_VALUES))) - assert panel.outputs, 'the initial chain was not drawn' - assert 'Figure' in str(panel.outputs[0]['data']) + assert bytes(image.value).startswith(b'\x89PNG'), 'the initial chain was not rendered' slider.value = 2 - assert panel.outputs, 'changing Q did not redraw' - assert 'Figure' in str(panel.outputs[0]['data']) + assert bytes(image.value).startswith(b'\x89PNG'), 'changing Q did not swap in a rendering' + + def test_the_corner_slider_swaps_bytes_without_redrawing(self, multi_q_analysis): + # WHEN every chain's figure was rendered once, at construction + self._sample_independently(multi_q_analysis) + with patch('easydynamics.analysis.posterior_sampling._in_notebook', return_value=True): + widget = multi_q_analysis.bayesian.plot_corner() + image, slider = widget.children + first_bytes = image.value + + # THEN the slider moves with matplotlib rendering forbidden + with patch('easydynamics.utils.posterior_plotting.plot_corner') as render: + slider.value = 1 + changed_bytes = image.value + slider.value = 0 + + # EXPECT the callback only swapped stored bytes: nothing was drawn on a move, the image + # followed the slider, and coming back restored the identical rendering + render.assert_not_called() + assert changed_bytes != first_bytes + assert image.value == first_bytes def test_corner_without_a_notebook_or_q_index_says_what_to_do(self, multi_q_analysis): # WHEN @@ -1572,14 +1664,167 @@ def test_the_slider_only_offers_q_indices_that_were_sampled(self, multi_q_analys # EXPECT the slider cannot land on a Q with nothing to draw assert list(widget.children[1].options) == [2] - def test_trace_points_at_the_individual_chains(self, multi_q_analysis): + ############# + # Per-Q sliders for trace, marginal and correlations + ############# + + def test_trace_q_index_plots_that_qs_chain(self, multi_q_analysis): # WHEN self._sample_independently(multi_q_analysis) - # THEN EXPECT - with pytest.raises(RuntimeError, match='no single trace'): + # THEN + figure = multi_q_analysis.bayesian.plot_trace(Q_index=1) + + # EXPECT that Q's own trace: one panel per parameter plus the log-posterior + n_parameters = len(multi_q_analysis.analysis_list[1].get_free_parameters()) + assert len(figure.axes) == n_parameters + 1 + + def test_trace_offers_a_slider_in_a_notebook(self, multi_q_analysis): + # WHEN + self._sample_independently(multi_q_analysis) + + # THEN + with patch('easydynamics.analysis.posterior_sampling._in_notebook', return_value=True): + widget = multi_q_analysis.bayesian.plot_trace() + + # EXPECT the pre-rendered image-and-slider box, offering every sampled Q index + image, slider = widget.children + assert list(slider.options) == list(range(len(Q_VALUES))) + assert bytes(image.value).startswith(b'\x89PNG') + + def test_trace_without_a_notebook_or_q_index_says_what_to_do(self, multi_q_analysis): + # WHEN + self._sample_independently(multi_q_analysis) + + # THEN EXPECT it names the sampled Q indices rather than just refusing + with ( + patch('easydynamics.analysis.posterior_sampling._in_notebook', return_value=False), + pytest.raises(RuntimeError, match=r'sampled Q indices are \[0, 1, 2\]'), + ): multi_q_analysis.bayesian.plot_trace() + def test_marginal_q_index_plots_that_qs_chain(self, multi_q_analysis): + # WHEN + self._sample_independently(multi_q_analysis) + + # THEN + figure = multi_q_analysis.bayesian.plot_marginal('Gaussian width', Q_index=2) + + # EXPECT a single-axis marginal under the parameter's plain per-Q label + assert len(figure.axes) == 1 + assert figure.axes[0].get_xlabel() == 'Gaussian width (meV)' + + def test_marginal_offers_a_slider_in_a_notebook(self, multi_q_analysis): + # WHEN + self._sample_independently(multi_q_analysis) + + # THEN + with patch('easydynamics.analysis.posterior_sampling._in_notebook', return_value=True): + widget = multi_q_analysis.bayesian.plot_marginal('Gaussian width') + + # EXPECT + image, slider = widget.children + assert list(slider.options) == list(range(len(Q_VALUES))) + assert bytes(image.value).startswith(b'\x89PNG') + + def test_marginal_slider_resolves_a_parameter_object_across_q(self, multi_q_analysis): + # WHEN the Parameter object belongs to one Q's model only + self._sample_independently(multi_q_analysis) + parameters = multi_q_analysis.analysis_list[1].get_free_parameters() + target = next(p for p in parameters if p.name == 'Gaussian width') + + # THEN + with patch('easydynamics.analysis.posterior_sampling._in_notebook', return_value=True): + widget = multi_q_analysis.bayesian.plot_marginal(target) + + # EXPECT the slider still covers every Q, through the shared display name + image, slider = widget.children + assert list(slider.options) == list(range(len(Q_VALUES))) + assert bytes(image.value).startswith(b'\x89PNG') + + def test_marginal_without_a_notebook_or_q_index_says_what_to_do(self, multi_q_analysis): + # WHEN + self._sample_independently(multi_q_analysis) + + # THEN EXPECT + with ( + patch('easydynamics.analysis.posterior_sampling._in_notebook', return_value=False), + pytest.raises(RuntimeError, match=r'sampled Q indices are \[0, 1, 2\]'), + ): + multi_q_analysis.bayesian.plot_marginal('Gaussian width') + + def test_correlations_q_index_plots_that_qs_chain(self, multi_q_analysis): + # WHEN + self._sample_independently(multi_q_analysis) + + # THEN + figure = multi_q_analysis.bayesian.plot_correlations(Q_index=0) + + # EXPECT that Q's own matrix and its colorbar, under the plain per-Q labels + assert len(figure.axes) == 2 + labels = [text.get_text() for text in figure.axes[0].get_xticklabels()] + assert 'Gaussian width' in labels + assert all('Q_index=' not in label for label in labels) + + def test_correlations_offers_a_slider_in_a_notebook(self, multi_q_analysis): + # WHEN + self._sample_independently(multi_q_analysis) + + # THEN + with patch('easydynamics.analysis.posterior_sampling._in_notebook', return_value=True): + widget = multi_q_analysis.bayesian.plot_correlations() + + # EXPECT + image, slider = widget.children + assert list(slider.options) == list(range(len(Q_VALUES))) + assert bytes(image.value).startswith(b'\x89PNG') + + def test_correlations_without_a_notebook_or_q_index_says_what_to_do(self, multi_q_analysis): + # WHEN + self._sample_independently(multi_q_analysis) + + # THEN EXPECT + with ( + patch('easydynamics.analysis.posterior_sampling._in_notebook', return_value=False), + pytest.raises(RuntimeError, match=r'sampled Q indices are \[0, 1, 2\]'), + ): + multi_q_analysis.bayesian.plot_correlations() + + def test_chain_figure_q_indices_are_validated(self, multi_q_analysis): + # THEN EXPECT both ends of the range are checked before any chain is looked up + for plot in ( + multi_q_analysis.bayesian.plot_trace, + multi_q_analysis.bayesian.plot_correlations, + ): + with pytest.raises(IndexError, match='non-negative'): + plot(Q_index=-1) + with pytest.raises(IndexError, match='out of bounds'): + plot(Q_index=99) + with pytest.raises(IndexError, match='non-negative'): + multi_q_analysis.bayesian.plot_marginal('Gaussian width', Q_index=-1) + with pytest.raises(IndexError, match='out of bounds'): + multi_q_analysis.bayesian.plot_posterior_predictive(Q_index=99) + + def test_a_simultaneous_chain_serves_marginal_and_correlations(self, multi_q_analysis): + # WHEN + bound_all_chain(multi_q_analysis) + parameters = multi_q_analysis._chain_parameters() + + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.return_value = fake_chain_results(parameters) + multi_q_analysis.bayesian.sample(fit_method='simultaneous', samples=10) + + # THEN + marginal = multi_q_analysis.bayesian.plot_marginal('Gaussian width (Q_index=0)') + correlations = multi_q_analysis.bayesian.plot_correlations() + + # EXPECT single figures over the joint chain, under its Q-qualified labels + assert len(marginal.axes) == 1 + assert marginal.axes[0].get_xlabel().startswith('Gaussian width (Q_index=0)') + labels = [text.get_text() for text in correlations.axes[0].get_xticklabels()] + assert len(labels) == len(parameters) + assert all('Q_index=' in label for label in labels) + def test_a_simultaneous_chain_still_takes_precedence(self, multi_q_analysis): # WHEN a simultaneous run follows an independent one self._sample_independently(multi_q_analysis) diff --git a/tests/unit/easydynamics/utils/test_posterior_plotting.py b/tests/unit/easydynamics/utils/test_posterior_plotting.py index 3130ac4a5..1bb3bc4c3 100644 --- a/tests/unit/easydynamics/utils/test_posterior_plotting.py +++ b/tests/unit/easydynamics/utils/test_posterior_plotting.py @@ -1,6 +1,9 @@ # SPDX-FileCopyrightText: 2026 EasyScience contributors # SPDX-License-Identifier: BSD-3-Clause +from unittest.mock import MagicMock +from unittest.mock import patch + import matplotlib as mpl import numpy as np import pytest @@ -9,11 +12,14 @@ import matplotlib.pyplot as plt +from easydynamics.utils.posterior_plotting import corner_with_slider +from easydynamics.utils.posterior_plotting import figures_with_slider 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 +from easydynamics.utils.posterior_plotting import predictive_with_slider @pytest.fixture(autouse=True) @@ -391,6 +397,189 @@ def test_axis_labels_are_set_when_given(self): assert fig.axes[0].get_ylabel() == 'Intensity' +class TestFiguresWithSlider: + @staticmethod + def _figure(value): + fig, axis = plt.subplots(figsize=(2.0, 1.5)) + axis.plot([0.0, 1.0], [0.0, value]) + return fig + + def test_returns_an_image_above_a_slider_over_the_given_indices(self): + # WHEN figures exist for a sparse set of indices + figures = {0: self._figure(0.0), 2: self._figure(2.0)} + + # THEN + widget = figures_with_slider(figures) + + # EXPECT the pre-rendered PNG of the first index, and only positions that hold a figure + image, slider = widget.children + assert bytes(image.value).startswith(b'\x89PNG') + assert list(slider.options) == [0, 2] + assert slider.value == 0 + + def test_moving_the_slider_swaps_stored_bytes_without_rendering(self): + # WHEN every figure was rendered once, at construction + widget = figures_with_slider({0: self._figure(0.0), 1: self._figure(1.0)}) + image, slider = widget.children + first_bytes = image.value + + # THEN the slider moves with no figures left to draw from + open_before = plt.get_fignums() + slider.value = 1 + changed_bytes = image.value + slider.value = 0 + + # EXPECT the image followed the slider by swapping stored bytes: no new matplotlib work, + # and coming back restores the identical rendering + assert plt.get_fignums() == open_before + assert changed_bytes != first_bytes + assert image.value == first_bytes + + def test_figures_are_closed_after_rendering(self): + # WHEN + figures = {0: self._figure(0.0), 1: self._figure(1.0)} + + # THEN + figures_with_slider(figures) + + # EXPECT no figure is left for a backend to draw a second time + assert plt.get_fignums() == [] + + def test_no_figures_raise(self): + # THEN EXPECT + with pytest.raises(ValueError, match='No figures'): + figures_with_slider({}) + + +class TestCornerWithSlider: + @pytest.fixture + def chains(self, draws): + return { + index: {'draws': draws + index, 'names': ['a', 'b', 'c'], 'units': ['meV', '', '']} + for index in (0, 2) + } + + def test_renders_one_corner_per_chain_behind_the_slider(self, chains): + # THEN + with patch( + 'easydynamics.utils.posterior_plotting.plot_corner', wraps=plot_corner + ) as render: + widget = corner_with_slider(chains, title='Fit', bins=13) + + # EXPECT every chain rendered once, up front, with the kwargs and per-index titles + # forwarded, and only the given indices on the slider + assert render.call_count == len(chains) + titles = {call.kwargs['title'] for call in render.call_args_list} + assert titles == {'Fit (Q index 0)', 'Fit (Q index 2)'} + assert all(call.kwargs['bins'] == 13 for call in render.call_args_list) + image, slider = widget.children + assert bytes(image.value).startswith(b'\x89PNG') + assert list(slider.options) == [0, 2] + + def test_no_chains_raise(self): + # THEN EXPECT + with pytest.raises(ValueError, match='No chains'): + corner_with_slider({}) + + +class TestPredictiveWithSlider: + @pytest.fixture + def arrays(self): + energy = np.linspace(-5.0, 5.0, 10) + q_values = np.array([0.5, 1.0]) + median = np.tile(np.exp(-0.5 * energy**2), (2, 1)) + return { + 'energy': energy, + 'q_values': q_values, + 'y': median + 0.01, + 'lower': median - 0.1, + 'median': median, + 'upper': median + 0.1, + } + + @staticmethod + def _fake_slicer_figure(): + control = MagicMock() + fig = MagicMock() + fig.bottom_bar = [MagicMock()] + fig.bottom_bar[0].controls = {'Q': control} + return fig, control + + def test_builds_the_datagroup_and_style_plopp_slices(self, arrays): + # WHEN pp.slicer is mocked out, since the real one needs an interactive backend + fake_fig, control = self._fake_slicer_figure() + + # THEN + with patch('plopp.slicer', return_value=fake_fig) as slicer: + fig = predictive_with_slider( + **arrays, + y_variances=np.full((2, 10), 0.01), + energy_unit='meV', + q_unit='1/angstrom', + ylabel='Intensity', + title='Fit', + credible_interval=68.0, + ) + + # EXPECT a Q/energy DataGroup sliced along energy, styled like plot_data_and_model: + # data as open black circles with error bars, the median a solid line, the band edges + # dashed and labelled with the interval + assert fig is fake_fig + args, kwargs = slicer.call_args + data_group = args[0] + assert set(data_group.keys()) == { + 'Data', + 'Posterior median', + '68% band (lower)', + '68% band (upper)', + } + assert data_group['Data'].dims == ('Q', 'energy') + assert data_group['Data'].variances is not None + assert str(data_group['Data'].coords['energy'].unit) == 'meV' + assert kwargs['keep'] == 'energy' + assert kwargs['title'] == 'Fit' + assert kwargs['linestyle']['Data'] == 'none' + assert kwargs['marker']['Data'] == 'o' + assert kwargs['color']['Data'] == 'black' + assert kwargs['linestyle']['Posterior median'] == '-' + assert kwargs['linestyle']['68% band (lower)'] == '--' + assert kwargs['linestyle']['68% band (upper)'] == '--' + # The plopp slider is switched to its single-value mode, as plot_data_and_model does, + # and the y label lands on the axis + assert control.slider_toggler.value == '-o-' + fake_fig.ax.set_ylabel.assert_called_once_with('Intensity') + fake_fig.autoscale.assert_called_once() + + def test_nan_padding_survives_into_the_datagroup(self, arrays): + # WHEN one Q is missing a point on the common grid + arrays['y'][1, 3] = np.nan + arrays['median'][1, 3] = np.nan + fake_fig, _ = self._fake_slicer_figure() + + # THEN + with patch('plopp.slicer', return_value=fake_fig) as slicer: + predictive_with_slider(**arrays) + + # EXPECT the gap reaches plopp as NaN, drawn as a break rather than an invented value + data_group = slicer.call_args.args[0] + assert np.isnan(data_group['Data'].values[1, 3]) + assert np.isnan(data_group['Posterior median'].values[1, 3]) + + def test_mismatched_shapes_raise(self, arrays): + # WHEN + arrays['median'] = arrays['median'][:, :-1] + + # THEN EXPECT + with pytest.raises(ValueError, match='median must have shape'): + predictive_with_slider(**arrays) + + @pytest.mark.parametrize('interval', [0.0, 100.0, -5.0]) + def test_invalid_credible_interval_raises(self, arrays, interval): + # THEN EXPECT + with pytest.raises(ValueError, match='credible_interval'): + predictive_with_slider(**arrays, credible_interval=interval) + + class TestScientificNotation: def test_shared_exponent_is_folded_into_the_label(self): # WHEN the values are small enough that matplotlib factors out an exponent, which it parks From 26c168b5fd610a438f1b2e68f5155e1ee42964d2 Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Mon, 17 Aug 2026 18:02:26 +0200 Subject: [PATCH 24/29] Write the progress line through sys.stdout Co-Authored-By: Claude Fable 5 --- src/easydynamics/analysis/posterior_sampling.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/easydynamics/analysis/posterior_sampling.py b/src/easydynamics/analysis/posterior_sampling.py index d1fe04103..2d33535f3 100644 --- a/src/easydynamics/analysis/posterior_sampling.py +++ b/src/easydynamics/analysis/posterior_sampling.py @@ -13,6 +13,7 @@ from __future__ import annotations import json +import sys import warnings from pathlib import Path from typing import TYPE_CHECKING @@ -1212,7 +1213,8 @@ def close(self, completed: bool) -> None: return if completed: self._write('Sampling: done') - print(flush=True) + sys.stdout.write('\n') + sys.stdout.flush() def _write(self, line: str) -> None: """ @@ -1223,7 +1225,8 @@ def _write(self, line: str) -> None: line : str The text to show, padded so it fully overwrites a longer previous line. """ - print(f'\r{line.ljust(self._line_length)}', end='', flush=True) + 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 From fa94bb2f7bb1674409a3fb6a0478f6d275ec67e5 Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Mon, 17 Aug 2026 18:08:33 +0200 Subject: [PATCH 25/29] Show the new posterior plots in the Bayesian tutorial The tutorial now demonstrates plot_marginal and plot_correlations from the sampled chain, progress=True on the sampling call, the 95 percent predictive band option, the Q slider that every posterior plot offers over independent chains, and notes that runs are not seedable. Co-Authored-By: Claude Fable 5 --- docs/docs/tutorials/bayesian.ipynb | 72 ++++++++++++++++++++++++++++-- 1 file changed, 68 insertions(+), 4 deletions(-) diff --git a/docs/docs/tutorials/bayesian.ipynb b/docs/docs/tutorials/bayesian.ipynb index b4ffa3fd6..e0d233ae0 100644 --- a/docs/docs/tutorials/bayesian.ipynb +++ b/docs/docs/tutorials/bayesian.ipynb @@ -152,7 +152,9 @@ "- `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." + "Sampling never moves your parameters: their values are restored afterwards, so the model is left exactly as the fit left it.\n", + "\n", + "For a long run, `progress=True` shows a single self-updating line with the percentage of generations completed, closed with `Sampling: done`. The percentage is based on the backend's own estimate of the run length, which can be too high, so a finished run may close the line before reaching 100%." ] }, { @@ -162,7 +164,7 @@ "metadata": {}, "outputs": [], "source": [ - "results = analysis.bayesian.sample(samples=4000, burn=300, thin=2)\n", + "results = analysis.bayesian.sample(samples=4000, burn=300, thin=2, progress=True)\n", "\n", "print(f'Collected {results.draws.shape[0]} draws for {results.draws.shape[1]} parameters.')" ] @@ -227,6 +229,46 @@ "analysis.bayesian.plot_corner()" ] }, + { + "cell_type": "markdown", + "id": "e52ab7b4", + "metadata": {}, + "source": [ + "### One parameter at a time\n", + "\n", + "`plot_marginal()` pulls a single parameter's posterior out of the chain: a histogram of its draws, with the median and the 16/84 percentiles — the same numbers `summary()` reports — marked on it." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2cdf2451", + "metadata": {}, + "outputs": [], + "source": [ + "analysis.bayesian.plot_marginal('Res. Gauss width')" + ] + }, + { + "cell_type": "markdown", + "id": "683943ef", + "metadata": {}, + "source": [ + "### The correlation matrix at a glance\n", + "\n", + "Where the corner plot shows every pairwise distribution, `plot_correlations()` reduces each panel to a single number — the Pearson correlation between the two parameters — and colour-codes the grid. It is the quickest way to spot which parameters the data cannot tell apart." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d577ef22", + "metadata": {}, + "outputs": [], + "source": [ + "analysis.bayesian.plot_correlations()" + ] + }, { "cell_type": "markdown", "id": "8f819c75", @@ -234,7 +276,9 @@ "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." + "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.\n", + "\n", + "The band defaults to the 68% credible interval; `credible_interval=95.0` widens it to 95%." ] }, { @@ -358,6 +402,24 @@ "full_analysis.bayesian.plot_corner()" ] }, + { + "cell_type": "markdown", + "id": "2c16b5e4", + "metadata": {}, + "source": [ + "The other plots work the same way over independent chains: `plot_posterior_predictive()`, `plot_trace()`, `plot_marginal()` and `plot_correlations()` all show a Q slider in a notebook — the predictive plot through the same slider machinery as `plot_data_and_model()` — take `Q_index=` to go straight to one Q, and outside a notebook name the sampled Q indices instead." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "86b57835", + "metadata": {}, + "outputs": [], + "source": [ + "full_analysis.bayesian.plot_posterior_predictive(n_draws=100)" + ] + }, { "cell_type": "markdown", "id": "a3448aee", @@ -373,7 +435,9 @@ "\n", "**Several Q values at once.** An `Analysis` can sample either way. `fit_method='independent'` gives each Q its own chain, which is cheaper; `fit_method='simultaneous'` runs one chain over every Q, which is what you need when parameters are shared across Q. `bayesian.summary()` gathers the per-Q chains into one table either way.\n", "\n", - "Corner plots are the exception. Independent chains share no draws, so nothing pairs a parameter at one Q with a parameter at another, and combining them would show correlations that came from how the sampling was run rather than from the data. `analysis.bayesian.plot_corner()` therefore shows one Q at a time: pass `Q_index`, or leave it out in a notebook to get a slider across the sampled Q values." + "Corner plots are the exception. Independent chains share no draws, so nothing pairs a parameter at one Q with a parameter at another, and combining them would show correlations that came from how the sampling was run rather than from the data. `analysis.bayesian.plot_corner()` therefore shows one Q at a time: pass `Q_index`, or leave it out in a notebook to get a slider across the sampled Q values.\n", + "\n", + "**Reproducibility.** Two identical `sample()` calls will not give identical chains: the DREAM backend draws from global random state and exposes no seed. Judge results by whether the summary is stable when the chain is extended, not by exact repetition." ] } ], From 7a7dfe3f576d8fbcb77d03939bddc874c71bb5e1 Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Mon, 17 Aug 2026 18:11:19 +0200 Subject: [PATCH 26/29] Apply the formatting fixes Co-Authored-By: Claude Fable 5 --- src/easydynamics/analysis/posterior_labels.py | 6 +-- .../analysis/posterior_sampling.py | 42 +++++++++---------- src/easydynamics/utils/posterior_plotting.py | 21 +++++----- 3 files changed, 33 insertions(+), 36 deletions(-) diff --git a/src/easydynamics/analysis/posterior_labels.py b/src/easydynamics/analysis/posterior_labels.py index e39811718..fe406f706 100644 --- a/src/easydynamics/analysis/posterior_labels.py +++ b/src/easydynamics/analysis/posterior_labels.py @@ -101,9 +101,9 @@ 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. + 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 ------- diff --git a/src/easydynamics/analysis/posterior_sampling.py b/src/easydynamics/analysis/posterior_sampling.py index 2d33535f3..c7a3f0731 100644 --- a/src/easydynamics/analysis/posterior_sampling.py +++ b/src/easydynamics/analysis/posterior_sampling.py @@ -258,8 +258,8 @@ def sample( 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 + 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``. @@ -271,10 +271,10 @@ def sample( 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. + 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) try: @@ -313,8 +313,7 @@ def extend( 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`. + 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. @@ -333,8 +332,8 @@ def extend( Notes ----- - Like :meth:`sample`, extensions are not reproducible: the sampler draws from NumPy's - global random state and exposes no seed control. + 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.') @@ -804,9 +803,8 @@ def plot_marginal(self, parameter: Parameter | str, **kwargs: dict[str, Any]) -> """ 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. + 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 ---------- @@ -1145,8 +1143,8 @@ def _install_progress_reporter( Returns ------- _SamplingProgress | None - The installed reporter, which the caller must close after the run, or None when nothing - was installed. + 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 @@ -1159,13 +1157,13 @@ 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. + 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. + 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: @@ -1206,8 +1204,8 @@ def close(self, completed: bool) -> None: 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. + 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 diff --git a/src/easydynamics/utils/posterior_plotting.py b/src/easydynamics/utils/posterior_plotting.py index bc2b4c280..ebaff33ac 100644 --- a/src/easydynamics/utils/posterior_plotting.py +++ b/src/easydynamics/utils/posterior_plotting.py @@ -214,8 +214,8 @@ def plot_marginal( """ 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. + 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 ---------- @@ -224,8 +224,8 @@ def plot_marginal( 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. + 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 @@ -282,9 +282,8 @@ def plot_correlations( 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. + 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. Parameters ---------- @@ -295,8 +294,8 @@ def plot_correlations( 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. + Figure size in inches. Defaults to a square that scales with the parameter count, plus room + for the colorbar. Returns ------- @@ -473,8 +472,8 @@ def _correlation_matrix(draws: np.ndarray) -> np.ndarray: 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. + 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) From 132fc5901e86095fb832c968bf0da7f70dd96c5a Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Mon, 17 Aug 2026 18:27:05 +0200 Subject: [PATCH 27/29] Satisfy the docstring and formatting checks The progress reporter closes through try/finally instead of a bare re-raise, and the plotting validation errors are documented in the form the docstring linter expects. Co-Authored-By: Claude Fable 5 --- .../analysis/posterior_sampling.py | 27 +++++++++---------- src/easydynamics/utils/posterior_plotting.py | 26 +++++++++--------- .../analysis/test_posterior_sampling.py | 8 +++--- 3 files changed, 31 insertions(+), 30 deletions(-) diff --git a/src/easydynamics/analysis/posterior_sampling.py b/src/easydynamics/analysis/posterior_sampling.py index c7a3f0731..2c23a9235 100644 --- a/src/easydynamics/analysis/posterior_sampling.py +++ b/src/easydynamics/analysis/posterior_sampling.py @@ -277,6 +277,7 @@ def sample( 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, @@ -284,12 +285,10 @@ def sample( samples=samples, burn=burn, thin=thin, population=population, **sampler_options ), ) - except BaseException: + completed = True + finally: if reporter is not None: - reporter.close(completed=False) - raise - if reporter is not None: - reporter.close(completed=True) + reporter.close(completed=completed) return results def extend( @@ -326,18 +325,19 @@ def extend( ------ RuntimeError If there is no chain to extend, or the previous run failed and left no results. - ValueError - If the model or data changed since the chain was started, or this run's parameters - differ from the ones the chain holds. 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, @@ -346,12 +346,10 @@ def extend( ), reuse_sampler=True, ) - except BaseException: + completed = True + finally: if reporter is not None: - reporter.close(completed=False) - raise - if reporter is not None: - reporter.close(completed=True) + reporter.close(completed=completed) return results def _run( @@ -1041,8 +1039,7 @@ def _resolve_column(self, results: SamplingResults, parameter: Parameter | str) parameter if isinstance(parameter, str) else getattr(parameter, 'name', '?') ) raise ValueError( - f'No sampled parameter named {requested!r}. ' - f'Available: {", ".join(sorted(names))}.' + f'No sampled parameter named {requested!r}. Available: {", ".join(sorted(names))}.' ) return matches[0] diff --git a/src/easydynamics/utils/posterior_plotting.py b/src/easydynamics/utils/posterior_plotting.py index ebaff33ac..c3576d164 100644 --- a/src/easydynamics/utils/posterior_plotting.py +++ b/src/easydynamics/utils/posterior_plotting.py @@ -37,9 +37,6 @@ def plot_trace( excursions. A visible trend means the chain has not reached the typical set and needs a longer burn-in. - A ``ValueError`` is raised if ``draws`` is not two-dimensional or is empty, if ``names`` does - not have one entry per column, or if ``logp`` does not have one entry per draw. - Parameters ---------- draws : np.ndarray @@ -60,6 +57,12 @@ def plot_trace( ------- 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) @@ -112,9 +115,6 @@ def plot_corner( distribution of a pair: a compact blob means the two are independent, while a narrow diagonal ridge means they are correlated and cannot be determined separately from this data. - A ``ValueError`` is raised if ``draws`` is not two-dimensional or is empty, if ``names`` does - not have one entry per column, or if any column contains non-finite values. - Parameters ---------- draws : np.ndarray @@ -135,6 +135,12 @@ def plot_corner( ------- 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) @@ -284,6 +290,8 @@ def plot_correlations( 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 ---------- @@ -301,12 +309,6 @@ def plot_correlations( ------- Figure The matplotlib Figure. - - Raises - ------ - ValueError - If ``draws`` is not two-dimensional or is empty, or if ``names`` does not have one entry - per column. """ draws = np.asarray(draws) _verify_draws(draws, names) diff --git a/tests/unit/easydynamics/analysis/test_posterior_sampling.py b/tests/unit/easydynamics/analysis/test_posterior_sampling.py index f8e19b798..5fb5cd24e 100644 --- a/tests/unit/easydynamics/analysis/test_posterior_sampling.py +++ b/tests/unit/easydynamics/analysis/test_posterior_sampling.py @@ -762,9 +762,11 @@ def test_progress_reports_and_finishes_the_line(self, analysis, capsys): def run_reporting_progress(**kwargs): for iteration in (1, 5, 10): - kwargs['progress_callback']( - {'iteration': iteration, 'total_steps': 10, 'sampling': True} - ) + kwargs['progress_callback']({ + 'iteration': iteration, + 'total_steps': 10, + 'sampling': True, + }) return fake_results(analysis) sampler_class.return_value.sample.side_effect = run_reporting_progress From fbe8dd45097e734506969b368b7ab3bec68c51ac Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Mon, 17 Aug 2026 18:34:14 +0200 Subject: [PATCH 28/29] Document propagated exceptions the way the docstring linter expects Co-Authored-By: Claude Fable 5 --- .../analysis/posterior_sampling.py | 121 +++++++++--------- src/easydynamics/utils/posterior_plotting.py | 18 +-- 2 files changed, 66 insertions(+), 73 deletions(-) diff --git a/src/easydynamics/analysis/posterior_sampling.py b/src/easydynamics/analysis/posterior_sampling.py index 533007e3f..04ca3104a 100644 --- a/src/easydynamics/analysis/posterior_sampling.py +++ b/src/easydynamics/analysis/posterior_sampling.py @@ -1198,12 +1198,13 @@ def sample( Raises ------ - IndexError - If Q_index is negative or out of range. - TypeError - If Q_index is not an int or None. ValueError - If fit_method is not "independent" or "simultaneous". + If fit_method is not "independent" or "simultaneous", or there are no Q values. + + Notes + ----- + An ``IndexError`` or ``TypeError`` propagates from the Q_index validation if Q_index is + out of range or not an int. """ if fit_method not in ('independent', 'simultaneous'): raise ValueError("Invalid fit method. Choose 'independent' or 'simultaneous'.") @@ -1241,8 +1242,8 @@ def extend( """ Continue the existing simultaneous chain with additional samples. - The chains from independent sampling live on the per-Q samplers, so each is extended - there rather than here. + The chains from independent sampling live on the per-Q samplers, so each is extended there + rather than here. Parameters ---------- @@ -1265,9 +1266,11 @@ def extend( RuntimeError If the latest sampling ran per Q index, so there is no simultaneous chain here to extend, or if there is no chain at all. - ValueError - If the model or data changed since the chain was started, or this run's parameters - differ from the ones the chain holds. + + 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. """ if self._results is None and self.results_per_q is not None: # Without this check, a stale simultaneous sampler would either be extended silently @@ -1325,8 +1328,8 @@ def summary(self, labeller: Callable[[Parameter], str] | None = None) -> Posteri Parameters ---------- labeller : Callable[[Parameter], str] | None, default=None - Overrides the label a resolved column is reported under. The default is this - analysis' own Q-qualified labels. + Overrides the label a resolved column is reported under. The default is this analysis' + own Q-qualified labels. Returns ------- @@ -1392,12 +1395,13 @@ def plot_corner(self, Q_index: int | None = None, **kwargs: dict[str, Any]) -> F Raises ------ - IndexError - If Q_index is negative or out of range. RuntimeError If a slider is asked for outside a notebook. - TypeError - If Q_index is not an int or None. + + Notes + ----- + An ``IndexError`` or ``TypeError`` propagates from the Q_index validation if Q_index is + out of range or not an int. """ from easydynamics.utils.posterior_plotting import corner_with_slider @@ -1436,8 +1440,8 @@ def plot_trace(self, Q_index: int | None = None, **kwargs: dict[str, Any]) -> Fi """ Plot the chain trace of each sampled parameter. - A simultaneous chain is one trace and is drawn directly. After independent sampling each - Q index has its own chain, so the traces are stepped through one at a time: pick one with + A simultaneous chain is one trace and is drawn directly. After independent sampling each Q + index has its own chain, so the traces are stepped through one at a time: pick one with ``Q_index``, or leave it out in a notebook to get a slider. Parameters @@ -1453,14 +1457,11 @@ def plot_trace(self, Q_index: int | None = None, **kwargs: dict[str, Any]) -> Fi Figure | VBox The matplotlib Figure, or an ipywidgets box with a Q slider. - Raises - ------ - IndexError - If Q_index is negative or out of range. - RuntimeError - If a slider is asked for outside a notebook, or nothing has been sampled yet. - TypeError - If Q_index is not an int or None. + Notes + ----- + A ``RuntimeError`` propagates if a slider is asked for outside a notebook or nothing has + been sampled yet, and an ``IndexError`` or ``TypeError`` from the Q_index validation if + Q_index is out of range or not an int. """ verify_Q_index(Q_index=Q_index, Q=self._analysis.Q, allow_none=True) per_q = self.results_per_q @@ -1504,16 +1505,12 @@ def plot_marginal( Figure | VBox The matplotlib Figure, or an ipywidgets box with a Q slider. - Raises - ------ - IndexError - If Q_index is negative or out of range. - RuntimeError - If a slider is asked for outside a notebook, or nothing has been sampled yet. - TypeError - If Q_index is not an int or None. - ValueError - If the parameter matches no sampled chain column. + Notes + ----- + A ``ValueError`` propagates if the parameter matches no sampled chain column, a + ``RuntimeError`` if a slider is asked for outside a notebook or nothing has been sampled + yet, and an ``IndexError`` or ``TypeError`` from the Q_index validation if Q_index is out + of range or not an int. """ verify_Q_index(Q_index=Q_index, Q=self._analysis.Q, allow_none=True) per_q = self.results_per_q @@ -1539,9 +1536,9 @@ def plot_correlations( """ Plot the Pearson correlation matrix of the sampled parameters. - A simultaneous chain gives one matrix over every Q's parameters at once. After - independent sampling no draw pairs one Q with another, so there is one matrix per chain: - pick one with ``Q_index``, or leave it out in a notebook to get a slider. + A simultaneous chain gives one matrix over every Q's parameters at once. After independent + sampling no draw pairs one Q with another, so there is one matrix per chain: pick one with + ``Q_index``, or leave it out in a notebook to get a slider. Parameters ---------- @@ -1556,14 +1553,11 @@ def plot_correlations( Figure | VBox The matplotlib Figure, or an ipywidgets box with a Q slider. - Raises - ------ - IndexError - If Q_index is negative or out of range. - RuntimeError - If a slider is asked for outside a notebook, or nothing has been sampled yet. - TypeError - If Q_index is not an int or None. + Notes + ----- + A ``RuntimeError`` propagates if a slider is asked for outside a notebook or nothing has + been sampled yet, and an ``IndexError`` or ``TypeError`` from the Q_index validation if + Q_index is out of range or not an int. """ verify_Q_index(Q_index=Q_index, Q=self._analysis.Q, allow_none=True) per_q = self.results_per_q @@ -1588,9 +1582,9 @@ def plot_posterior_predictive( After independent sampling each Q has its own chain, and its own band: pick one with ``Q_index`` for a single matplotlib figure, or leave it out in a notebook to get a plopp - figure with a Q slider, looking and handling exactly like - ``Analysis.plot_data_and_model``. Plopp draws no filled band, so the slider view shows - the posterior median with a dashed line along each band edge instead of a shaded band. + figure with a Q slider, looking and handling exactly like ``Analysis.plot_data_and_model``. + Plopp draws no filled band, so the slider view shows the posterior median with a dashed + line along each band edge instead of a shaded band. Parameters ---------- @@ -1613,17 +1607,16 @@ def plot_posterior_predictive( Raises ------ - IndexError - If Q_index is negative or out of range. - NotImplementedError - If the latest chain is simultaneous: it binds every dataset at once, and no per-Q - chain exists for Q_index to pick out. - RuntimeError - If a slider is asked for outside a notebook, or nothing has been sampled yet. - TypeError - If Q_index is not an int or None. ValueError If n_draws is not a positive integer, or credible_interval is out of range. + + Notes + ----- + A ``NotImplementedError`` propagates when the latest chain is simultaneous: it binds every + dataset at once, and no per-Q chain exists for Q_index to pick out. A ``RuntimeError`` + propagates if a slider is asked for outside a notebook or nothing has been sampled yet, + and an ``IndexError`` or ``TypeError`` from the Q_index validation if Q_index is out of + range or not an int. """ 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}.') @@ -1674,9 +1667,9 @@ def _figures_with_q_slider( """ Render one figure per sampled Q index and put them behind a slider. - Only the Q indices that actually hold a chain get a figure, so the slider cannot land on - an empty position. Each figure carries its per-Q Analysis' own display name, which names - the Q index. + Only the Q indices that actually hold a chain get a figure, so the slider cannot land on an + empty position. Each figure carries its per-Q Analysis' own display name, which names the Q + index. Parameters ---------- @@ -1708,8 +1701,8 @@ def _shared_display_name( Find the display name a Parameter goes by within its own Q's chain. The same model is repeated per Q, so the name one chain reports a parameter under is the - name every other chain reports its own copy under. Resolving through it lets a slider - show the matching marginal at every Q even though the Parameter object belongs to one. + name every other chain reports its own copy under. Resolving through it lets a slider show + the matching marginal at every Q even though the Parameter object belongs to one. Parameters ---------- diff --git a/src/easydynamics/utils/posterior_plotting.py b/src/easydynamics/utils/posterior_plotting.py index a03c7658d..d1e45b146 100644 --- a/src/easydynamics/utils/posterior_plotting.py +++ b/src/easydynamics/utils/posterior_plotting.py @@ -294,8 +294,8 @@ def plot_correlations( 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. + 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 ---------- @@ -601,9 +601,9 @@ def figures_with_slider(figures: dict[int, Figure], description: str = 'Q index' Show one pre-rendered figure at a time, with a slider choosing which one. Every figure is rendered to PNG bytes once, up front, and the slider callback only swaps the - stored bytes into an image widget. Moving the slider therefore costs no matplotlib work at - all, which keeps it as responsive as the plopp slider on the data plots; re-rendering a - figure on every move is what made the previous slider feel sluggish. + stored bytes into an image widget. Moving the slider therefore costs no matplotlib work at all, + which keeps it as responsive as the plopp slider on the data plots; re-rendering a figure on + every move is what made the previous slider feel sluggish. The figures are closed after rendering, so no backend draws them a second time. @@ -726,11 +726,11 @@ def predictive_with_slider( Built on ``plopp.slicer`` over a scipp DataGroup with a Q dimension, so the figure looks and handles exactly like ``Analysis.plot_data_and_model``: the data with its error bars, the model curves on top, and a Q slider underneath. Plopp draws no filled band for sliced data -- its - only spread representation is variance-based error bars -- so the credible band is drawn as - the posterior median with a dashed line along each band edge, labelled with the interval. + only spread representation is variance-based error bars -- so the credible band is drawn as the + posterior median with a dashed line along each band edge, labelled with the interval. - Rows are laid out on one common energy grid; where a Q has no point (masked or never - measured), NaN leaves a gap in the lines rather than inventing a value. + Rows are laid out on one common energy grid; where a Q has no point (masked or never measured), + NaN leaves a gap in the lines rather than inventing a value. Parameters ---------- From 6310e7253d3ad56ae6b27f3285355b12f90bac64 Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Mon, 17 Aug 2026 21:41:32 +0200 Subject: [PATCH 29/29] Give the Bayesian tutorial the widget backend its sliders need The Q-slider cells go through the plopp slicer, which refuses the inline backend; every plopp-using tutorial already runs %matplotlib widget. Co-Authored-By: Claude Fable 5 --- docs/docs/tutorials/bayesian.ipynb | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/docs/tutorials/bayesian.ipynb b/docs/docs/tutorials/bayesian.ipynb index e0d233ae0..bc92e21fc 100644 --- a/docs/docs/tutorials/bayesian.ipynb +++ b/docs/docs/tutorials/bayesian.ipynb @@ -27,7 +27,8 @@ "import easydynamics.sample_model as sm\n", "from easydynamics.analysis.analysis1d import Analysis1d\n", "\n", - "%matplotlib inline" + "# Make the plots interactive; the Q sliders need the widget backend\n", + "%matplotlib widget" ] }, {