diff --git a/docs/docs/tutorials/bayesian.ipynb b/docs/docs/tutorials/bayesian.ipynb index 51a68acd..bc92e21f 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" ] }, { @@ -152,7 +153,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 +165,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 +230,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 +277,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%." ] }, { @@ -276,6 +321,106 @@ "Chains are expensive, so they can be saved and reloaded with `analysis.bayesian.save(path)` and `analysis.bayesian.load(path)`. A reloaded chain can be summarized, plotted, or extended further, exactly like a fresh one." ] }, + { + "cell_type": "markdown", + "id": "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].bayesian.suggest_bounds().apply()\n", + " full_analysis.bayesian.sample(\n", + " fit_method='independent', Q_index=Q_index, samples=3000, burn=200, thin=2\n", + " )" + ] + }, + { + "cell_type": "markdown", + "id": "740fa625", + "metadata": {}, + "source": [ + "`bayesian.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.bayesian.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.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", @@ -287,7 +432,13 @@ "\n", "**Sampling only some parameters.** `bayesian.sample(parameters=[...])` restricts the chain to a subset, which is faster because the number of chains scales with the number of parameters. Be careful with the result: the other parameters are held *fixed*, which is not the same as averaging over them. The intervals you get are conditional on those fixed values, and will be too narrow whenever the parameters are correlated.\n", "\n", - "**Degenerate parameters.** As seen above, they are a modelling problem rather than a sampling one. `bayesian.suggest_bounds()` returning absurd values, or a warning that the posterior has piled up against its bounds, are both signs to go back and look at the model." + "**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.\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. `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.\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." ] } ], diff --git a/docs/docs/tutorials/tutorial1_brownian.ipynb b/docs/docs/tutorials/tutorial1_brownian.ipynb index bb640325..3c902ea1 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. `bayesian.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.bayesian.suggest_bounds()\n", + "print(suggestions)\n", + "suggestions.apply()" + ] + }, + { + "cell_type": "code", + "id": "604cd4e9", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "parameter_analysis.bayesian.sample(samples=4000, burn=200, thin=2)\n", + "parameter_analysis.bayesian.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.bayesian.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/__init__.py b/src/easydynamics/analysis/__init__.py index 89126ecd..c6eb02a9 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 8fb3d703..46645afa 100644 --- a/src/easydynamics/analysis/analysis.py +++ b/src/easydynamics/analysis/analysis.py @@ -14,6 +14,8 @@ from easydynamics.analysis.analysis1d import Analysis1d from easydynamics.analysis.analysis_base import AnalysisBase +from easydynamics.analysis.posterior_labels import ParameterLabels +from easydynamics.analysis.posterior_sampling import MultiQPosteriorSampler from easydynamics.experiment import Experiment from easydynamics.sample_model import SampleModel from easydynamics.sample_model.instrument_model import InstrumentModel @@ -30,6 +32,10 @@ class Analysis(AnalysisBase): 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 through :attr:`bayesian`; see + :class:`~easydynamics.analysis.posterior_sampling.MultiQPosteriorSampler`. + Examples -------- **Fitting vanadium data for instrument calibration** @@ -117,6 +123,11 @@ 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._fitter = None + self._fitter_is_dirty = True + self._bayesian = None super().__init__( display_name=display_name, unique_name=unique_name, @@ -170,6 +181,70 @@ def analysis_list(self, _value: list[Analysis1d]) -> None: 'or instrument model.' ) + @property + def fitter(self) -> MultiFitter: + """ + The EasyScience MultiFitter covering every Q index, built on first use. + + Returns + ------- + MultiFitter + The cached MultiFitter. + """ + if self._fitter_is_dirty or self._fitter is None: + self._fitter = self._build_fitter() + self._fitter_is_dirty = False + return self._fitter + + @property + def bayesian(self) -> MultiQPosteriorSampler: + """ + Bayesian posterior sampling for this Analysis, created on first use. + + Returns + ------- + MultiQPosteriorSampler + The sampler, which can run per Q index or over all of them at once. + """ + if self._bayesian is None: + self._bayesian = MultiQPosteriorSampler( + analysis=self, + sampling_data=self._sampling_data, + chain_parameters=self._chain_parameters, + parameter_labels=self._parameter_labels, + prepare=self._prepare_for_sampling, + per_q=lambda: self.analysis_list, + ) + return self._bayesian + + def _invalidate_fitter(self) -> None: + """Mark the MultiFitter, and the Sampler built from it, as needing a rebuild.""" + self._fitter_is_dirty = True + if self._bayesian is not None: + self._bayesian.invalidate() + + def _parameter_labels(self) -> ParameterLabels: + """ + Get labels for the chain's parameters, qualified by Q index where needed. + + Every Q index carries its own copy of each model parameter, all sharing a name, so a bare + name would produce several identical rows in a summary and could not pick a parameter out. + + Returns + ------- + ParameterLabels + Labels over the current free parameters. + """ + owners = self._parameter_owner_index() + return ParameterLabels( + self._chain_parameters(), + qualify=lambda parameter: ( + None + if owners.get(parameter.unique_name) is None + else f'Q_index={owners[parameter.unique_name]}' + ), + ) + ############# # Other methods ############# @@ -283,8 +358,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: @@ -661,6 +737,8 @@ 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: """ @@ -668,6 +746,8 @@ 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: """ @@ -675,6 +755,8 @@ 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: """ @@ -682,6 +764,8 @@ 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: """Rebuild the analysis list if any dependency has changed since it was last built.""" @@ -695,6 +779,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. @@ -714,6 +799,102 @@ def _create_analysis_list(self) -> None: # Private methods ############# + ############# + # The contract PosteriorSampler relies on (simultaneous sampling over all Q) + ############# + + def _build_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 _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 _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_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 _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/parameter_analysis.py b/src/easydynamics/analysis/parameter_analysis.py index 7e24108e..1ac49684 100644 --- a/src/easydynamics/analysis/parameter_analysis.py +++ b/src/easydynamics/analysis/parameter_analysis.py @@ -9,11 +9,14 @@ 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.fit_binding import FitBinding +from easydynamics.analysis.posterior_labels import ParameterLabels +from easydynamics.analysis.posterior_sampling import PosteriorSampler from easydynamics.base_classes.easydynamics_modelbase import EasyDynamicsModelBase from easydynamics.utils.fit_target import FitTarget from easydynamics.utils.utils import _in_notebook @@ -98,6 +101,13 @@ def __init__( default, None. """ + self._fitter = None + self._fitter_is_dirty = True + self._bayesian = None + # 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) self._parameters = self._verify_parameters(parameters) @@ -130,6 +140,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 +165,94 @@ 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() + + @property + def fitter(self) -> MultiFitter: + """ + The EasyScience MultiFitter over the binding models, built on first use. + + Returns + ------- + MultiFitter + The cached MultiFitter. + """ + if self._fitter_is_dirty or self._fitter is None: + self._fitter = self._build_fitter() + 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, + ) + return self._bayesian + + def _invalidate_fitter(self) -> None: + """Mark the MultiFitter, and the Sampler built from it, as needing a rebuild.""" + self._fitter_is_dirty = True + if self._bayesian is not None: + self._bayesian.invalidate() + + def _parameter_labels(self) -> ParameterLabels: + """ + Get labels for the chain's parameters, qualified by binding model where needed. + + Two bindings can use models of the same kind, whose parameters would then share a name. The + prefix is the model's name, matching the choice to report parameters under their name + rather than their display name, since for several models the display name is just the class + name. If two models share a name as well, the unique name is used: a label that does not + disambiguate is worse than a long one. + + Returns + ------- + ParameterLabels + Labels over the free parameters of the binding models. + """ + models = {binding.model.unique_name: binding.model for binding in self.bindings} + owners = {} + for model in models.values(): + for parameter in model.get_free_parameters(): + owners.setdefault(parameter.unique_name, model) + model_names = [getattr(m, 'name', None) or m.display_name for m in models.values()] + + def qualify(parameter: Parameter) -> str | None: + """ + Get the model name a parameter belongs to. + + Parameters + ---------- + parameter : Parameter + The parameter to qualify. + + Returns + ------- + str | None + The owning model's name, its unique name if that name is shared, or None if the + parameter belongs to no binding model. + """ + owner = owners.get(parameter.unique_name) + if owner is None: + return None + name = getattr(owner, 'name', None) or owner.display_name + if name is None or model_names.count(name) > 1: + return owner.unique_name + return name + + return ParameterLabels(self._chain_parameters(), qualify=qualify) ############# # Other methods @@ -163,18 +262,37 @@ 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, _, 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]: + """ + 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 +325,94 @@ 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, - ) + ############# + # The contract PosteriorSampler relies on + ############# + + def _build_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() + 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 _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, _, models = self._build_fit_inputs() + self._invalidate_fitter_if_targets_changed(models) + return xs, ys, ws + + def _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 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 d8d65af7..e2c197a8 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 @@ -41,7 +47,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. + 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. @@ -131,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 ------- @@ -139,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: @@ -201,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. + """ + 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 + + def suggest_bounds_for_parameters( parameters: list[Parameter], labels: list[str] | None = None, @@ -230,7 +277,8 @@ def suggest_bounds_for_parameters( 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. + 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 @@ -606,11 +654,12 @@ def summarize_draws( 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 ---------- diff --git a/src/easydynamics/analysis/posterior_sampling.py b/src/easydynamics/analysis/posterior_sampling.py index 2c23a923..04ca3104 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 sys import warnings @@ -23,11 +24,15 @@ from easyscience.fitting import AvailableMinimizers from easyscience.fitting import Sampler +from easydynamics.analysis.posterior import PosteriorSummary from easydynamics.analysis.posterior import degenerate_parameters from easydynamics.analysis.posterior import parameters_at_bounds from easydynamics.analysis.posterior import suggest_bounds_for_parameters from easydynamics.analysis.posterior import summarize_draws from easydynamics.analysis.posterior import unbounded_parameters +from easydynamics.analysis.posterior_labels import ParameterLabels +from easydynamics.utils.utils import _in_notebook +from easydynamics.utils.utils import verify_Q_index if TYPE_CHECKING: import os @@ -35,11 +40,11 @@ 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 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(). @@ -397,6 +402,10 @@ def _run( chain_parameters = self._chain_parameters() if not chain_parameters: + # Let the analysis raise its own, more specific complaint first — e.g. a + # ParameterAnalysis without a parameters Dataset or bindings has no free + # parameters either, but "every parameter is fixed" would mislead there. + self._sampling_data() raise ValueError( 'There are no free parameters to sample: every parameter is fixed. ' 'Free at least one parameter before sampling.' @@ -607,19 +616,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 @@ -627,10 +643,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]: @@ -689,7 +712,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( @@ -1094,6 +1117,767 @@ def _require_results(self) -> SamplingResults: return self._results +class MultiQPosteriorSampler(PosteriorSampler): + """ + Posterior sampling for an Analysis covering several Q values. + + Reached as ``analysis.bayesian``. Sampling can run either way round: + + - ``fit_method='independent'`` gives each Q index its own chain, which is cheaper and keeps the + Q values from influencing one another. + - ``fit_method='simultaneous'`` runs a single chain over every Q at once, which is what is + needed when parameters are shared across Q, and costs considerably more: DREAM runs a number + of chains proportional to the parameter count, and a simultaneous run has every Q's + parameters in play together. + + Results from independent runs stay on the per-Q samplers. This class gathers them where + gathering is sound, and declines where it is not; see :meth:`summary` and :meth:`plot_corner`. + + Parameters + ---------- + per_q : Callable[[], list] + Returns the per-Q Analysis objects, each exposing ``Q_index`` and its own ``bayesian``. + **kwargs : dict[str, Any] + Forwarded to :class:`PosteriorSampler`. + """ + + def __init__(self, per_q: Callable[[], list], **kwargs: dict[str, Any]) -> None: + super().__init__(**kwargs) + self._per_q = per_q + + @property + def results_per_q(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:`results`. + + Returns + ------- + list[SamplingResults | None] | None + One entry per Q index, None where that Q has not been sampled, or None overall if no Q + index has been sampled. + """ + results = [analysis1d.bayesian.results for analysis1d in self._per_q()] + return results if any(result is not None for result in results) else None + + def sample( + 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, per Q index or over all of them at once. + + 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 underlying sampler. + + Returns + ------- + SamplingResults | list[SamplingResults] + A single result when a specific Q index was sampled or when sampling simultaneously, + and otherwise one result per Q index. + + Raises + ------ + ValueError + 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'.") + per_q = self._per_q() + if not per_q: + 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: + 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 [ + analysis1d.bayesian.sample(samples=samples, burn=burn, thin=thin, **sampler_options) + for analysis1d in per_q + ] + + 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. + + 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 + # 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. + + 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 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 + One entry per sampled parameter, across every Q index that has been sampled. + """ + per_q = self.results_per_q + if self._results is not None or per_q is None: + return super().summary(labeller) + + # 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 analysis1d in self._per_q(): + if analysis1d.bayesian.results is None: + continue + entries.extend(analysis1d.bayesian.summary(labeller=qualify).entries) + return PosteriorSummary(entries) + + def set_parameters_to_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. + + Returns + ------- + list[Parameter] + The parameters that were changed. + """ + if self._results is not None or self.results_per_q is None: + return super().set_parameters_to_median() + changed = [] + for analysis1d in self._per_q(): + if analysis1d.bayesian.results is not None: + changed.extend(analysis1d.bayesian.set_parameters_to_median()) + return changed + + def plot_corner(self, Q_index: int | None = None, **kwargs: dict[str, Any]) -> Figure | VBox: + """ + Plot the marginal and pairwise posterior distributions. + + 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 came from 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 | VBox + The matplotlib Figure, or an ipywidgets box with a Q slider. + + Raises + ------ + RuntimeError + If a slider is asked for outside a notebook. + + 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 + + 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) + + analyses = self._per_q() + if Q_index is not None: + return analyses[Q_index].bayesian.plot_corner(**kwargs) + + if not _in_notebook(): + 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}.' + ) + + chains = {} + for analysis1d, result in zip(analyses, per_q, strict=True): + if result is None: + continue + # Named by the per-Q sampler, so the labels match that Q's own summary and stay short: + # the Q index is on the slider, and repeating it in every axis label would only cost + # width. The summary entries follow the draw columns, so the order lines up. + entries = list(analysis1d.bayesian.summary()) + chains[analysis1d.Q_index] = { + 'draws': result.draws, + 'names': [entry.name for entry in entries], + 'units': [entry.unit for entry in entries], + } + return corner_with_slider(chains, title=self._analysis.display_name, **kwargs) + + def plot_trace(self, Q_index: int | None = None, **kwargs: dict[str, Any]) -> Figure | VBox: + """ + 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 + ``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 | VBox + The matplotlib Figure, or an ipywidgets box with a Q slider. + + 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 + 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. + + 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 + 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. + + 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 + 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 + ------ + 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}.') + 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 + ) + 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: + """ + Get the stored results, pointing at the per-Q chains when those are what exist. + + Returns + ------- + SamplingResults + The results of the most recent simultaneous run. + + Raises + ------ + RuntimeError + If no simultaneous sampling has been run. + """ + if self._results is None and self.results_per_q 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. summary() and " + 'set_parameters_to_median() gather those up; for anything needing a single chain, ' + 'use analysis.analysis_list[Q_index].bayesian, or sample with ' + "fit_method='simultaneous'." + ) + 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. @@ -1116,7 +1900,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/src/easydynamics/utils/posterior_plotting.py b/src/easydynamics/utils/posterior_plotting.py index c3576d16..d1e45b14 100644 --- a/src/easydynamics/utils/posterior_plotting.py +++ b/src/easydynamics/utils/posterior_plotting.py @@ -10,8 +10,10 @@ from __future__ import annotations +import io import warnings from typing import TYPE_CHECKING +from typing import Any import matplotlib.pyplot as plt import numpy as np @@ -19,7 +21,9 @@ from matplotlib.ticker import MaxNLocator if TYPE_CHECKING: + from ipywidgets import VBox from matplotlib.figure import Figure + from plopp.backends.matplotlib.figure import InteractiveFigure def plot_trace( @@ -290,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 ---------- @@ -590,3 +594,240 @@ 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 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, + **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. The figures are pre-rendered + through :func:`figures_with_slider`, so the slider moves without re-drawing anything. + + 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 figure and the slider. + + Raises + ------ + ValueError + If no chains are given. + """ + if not chains: + raise ValueError('No chains to plot.') + + 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, + ) + for index, chain in chains.items() + } + return figures_with_slider(figures) + + +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.py b/tests/integration/fitting/test_bayesian_sampling.py index f341a1c8..58fc48ab 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'}, } @@ -132,15 +135,22 @@ 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( - additional_samples=500, thin=2, sampler_kwargs={'trim': False} + extended = analysis.bayesian.extend( + 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 new file mode 100644 index 00000000..74061b57 --- /dev/null +++ b/tests/integration/fitting/test_bayesian_sampling_multi_q.py @@ -0,0 +1,329 @@ +# 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 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 +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 + +Q_VALUES = [0.5, 1.0, 1.5] +NOISE = 0.02 +TRUE_AREA = 2.0 + +SAMPLE_KWARGS = { + 'samples': 2000, + 'burn': 100, + 'thin': 2, + # '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'}, +} + + +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.bayesian.suggest_bounds().apply() + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + analysis.bayesian.sample(fit_method='simultaneous', **SAMPLE_KWARGS) + 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 + 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): + # THEN + names = [entry.name for entry in simultaneously_sampled.bayesian.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): + # 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 + # 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.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) + + # EXPECT + after = [float(p.value) for p in analysis._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._chain_parameters()) + + # THEN + trace = simultaneously_sampled.bayesian.plot_trace() + corner = simultaneously_sampled.bayesian.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, independently_sampled): + # THEN + 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, 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): + independent = analysis1d.bayesian.summary()['Gaussian width'] + simultaneous = simultaneously_sampled.bayesian.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 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 + 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; 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 + + # THEN + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + results = analysis.bayesian.sample(**SAMPLE_KWARGS) + + # 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 summary] + assert len(set(names)) == len(names) + + +class TestAggregatedIndependentChains: + 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 + 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, 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] + + 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) + finally: + for parameter, value in saved_values: + parameter.value = value diff --git a/tests/unit/easydynamics/analysis/test_analysis.py b/tests/unit/easydynamics/analysis/test_analysis.py index fdedd845..b8bda4f6 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_parameter_analysis.py b/tests/unit/easydynamics/analysis/test_parameter_analysis.py index 031f813c..1b6b8017 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_posterior.py b/tests/unit/easydynamics/analysis/test_posterior.py index 28817e34..436283a4 100644 --- a/tests/unit/easydynamics/analysis/test_posterior.py +++ b/tests/unit/easydynamics/analysis/test_posterior.py @@ -322,8 +322,9 @@ 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 + 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 @@ -364,3 +365,56 @@ 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) + + # 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): + # WHEN THEN 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]) + + # 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] + + 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]) + + # THEN 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]) + + # THEN 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() diff --git a/tests/unit/easydynamics/analysis/test_posterior_sampling.py b/tests/unit/easydynamics/analysis/test_posterior_sampling.py index 5fb5cd24..ab3759f9 100644 --- a/tests/unit/easydynamics/analysis/test_posterior_sampling.py +++ b/tests/unit/easydynamics/analysis/test_posterior_sampling.py @@ -2,20 +2,29 @@ # 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 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 from easyscience.fitting import AvailableMinimizers +from easyscience.fitting.multi_fitter import MultiFitter from easyscience.variable import Parameter +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 @@ -75,11 +84,122 @@ 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() +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 @@ -598,6 +718,106 @@ def test_plots_without_sampling_raise(self, analysis): with pytest.raises(RuntimeError): analysis.bayesian.plot_corner() + ############# + # Error paths + ############# + + 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() + + # 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) + 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('list index out of range') + + # 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): + # THEN EXPECT + with pytest.raises(TypeError, match='Parameter objects or labels'): + analysis.bayesian.sample(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.bayesian.sample(samples=10) + + # 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): + # 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.bayesian.sample(samples=10) + analysis.bayesian.save(str(tmp_path / 'chain')) + + 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')) + + # EXPECT the sidecar maps the old unique names onto the new analysis's parameters + 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) + + ############# + # Plot rendering + ############# + + 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.bayesian.sample(samples=10) + + # 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') + + ############# + # Predictive error bars + ############# + 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) @@ -887,6 +1107,889 @@ def test_predictions_take_draws_evenly_across_the_chain(self, analysis): expected = draws[[0, 24, 49, 74, 99], column] assert amplitudes / amplitudes[0] == pytest.approx(expected / expected[0]) + ############# + # Extend guards + ############# + + 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.bayesian.sample(samples=10) + + target = analysis.get_free_parameters()[0] + + # 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]) + + 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.bayesian.sample(samples=10) + + # THEN EXPECT: does not raise + analysis.bayesian.extend(additional_samples=10) + + ############# + # Sidecar labels + ############# + + 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] + + # THEN + with patch(SAMPLER_PATH) as sampler_class: + sampler_class.return_value.sample.side_effect = lambda **_k: fake_results(analysis) + with pytest.warns(UserWarning): + analysis.bayesian.sample(samples=10, parameters=[target.name]) + + # EXPECT + assert analysis.bayesian._saved_labels[ + target.unique_name + ] == analysis._parameter_labels().label(target) + + ############# + # 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_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') + + # 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() + + 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 + ############# + + 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.predictions() + + 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 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 bytes(image.value).startswith(b'\x89PNG'), 'the initial chain was not rendered' + + slider.value = 2 + 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 + 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] + + ############# + # 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 + 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) + 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_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] + 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.""" diff --git a/tests/unit/easydynamics/utils/test_posterior_plotting.py b/tests/unit/easydynamics/utils/test_posterior_plotting.py index c470bdd2..1bb3bc4c 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) @@ -67,6 +73,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): + # 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 + assert [axis.get_ylabel() for axis in fig.axes] == ['a (meV)', 'b (m^2/s)', 'c'] + def test_zero_row_draws_raise(self): # THEN EXPECT with pytest.raises(ValueError, match='no samples'): @@ -111,6 +124,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): + # 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 + # 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): + # 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 + 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' + def test_non_finite_draws_raise_naming_the_column(self, draws): # WHEN one column contains a NaN draws[5, 1] = np.nan @@ -350,3 +381,217 @@ 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): + # THEN + 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' + + +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 + # 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 + # 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()