Bayesian post-analysis using BUMPS DREAM

pyproject.toml

@@ -96,7 +96,7 @@ requires = ['hatchling', 'versioningit']
 
 [tool.hatch.build.targets.wheel]
 packages = ['src/easyscience']
-exclude = ['src/easyscience/legacy']
+# exclude = ['src/easyscience/legacy']
 
 [tool.hatch.metadata]
 allow-direct-references = true
@@ -211,8 +211,7 @@ select = [
 # Ignore specific rules globally
 ignore = [
   'COM812', # https://docs.astral.sh/ruff/rules/missing-trailing-comma/
-  # The following is replaced by 'D'/[tool.ruff.lint.pydocstyle] and [tool.pydoclint]
-  'DOC', # https://docs.astral.sh/ruff/rules/#pydoclint-doc
+  # The following is replaced by 'D'/[tool.ruff.lint.pydocstyle] and [tool.pydoclint]  'DOC', # https://docs.astral.sh/ruff/rules/#pydoclint-doc
   # Disable, as [tool.format_docstring] split one-line docstrings into the canonical multi-line layout
   'D200', # https://docs.astral.sh/ruff/rules/unnecessary-multiline-docstring/
 ]

src/easyscience/fitting/minimizers/minimizer_bumps.py

@@ -84,6 +84,7 @@
         tolerance: float | None = None,
         max_evaluations: int | None = None,
         progress_callback: Callable[[dict], bool | None] | None = None,
+        abort_test: Callable[[], bool] | None = None,
         minimizer_kwargs: dict | None = None,
         engine_kwargs: dict | None = None,
         **kwargs: Any,
@@ -116,6 +117,10 @@
             Optional callback for progress updates. The payload field
             ``iteration`` carries the BUMPS optimizer step index. By
             default, None.
+        abort_test : Callable[[], bool] | None, default=None
+            Optional callback that returns ``True`` to signal that the fit
+            should be aborted.  Called periodically during the
+            BUMPS optimizer iteration loop.
         minimizer_kwargs : dict | None, default=None
             Additional keyword arguments passed to the BUMPS minimizer.
             By default, None.
@@ -203,6 +208,7 @@
             fitclass=fitclass,
             problem=problem,
             monitors=monitors,
+            abort_test=abort_test or (lambda: False),
             **minimizer_kwargs,
             **kwargs,
         )
@@ -374,6 +380,200 @@
 
         return _outer(self)
 
+    def sample(
+        self,
+        x: np.ndarray,
+        y: np.ndarray,
+        weights: np.ndarray,
+        samples: int = 10000,
+        burn: int = 2000,
+        thin: int = 10,
+        chains: int | None = None,
+        population: int | None = None,
+        seed: int | None = None,
+        sampler_kwargs: dict | None = None,
+        progress_callback: Callable[[dict], bool | None] | None = None,
+        abort_test: Callable[[], bool] | None = None,
+    ) -> dict:
+        """Run Bayesian MCMC sampling using the BUMPS DREAM sampler.
+
+        Builds a BUMPS `FitProblem` from the current model and runs the DREAM
+        sampler.  This is the public minimizer-level entry point for Bayesian
+        sampling; the higher-level `MultiFitter.sample` delegates to this
+        method after flattening multi-dataset arrays.
+
+        Parameters
+        ----------
+        x : np.ndarray
+            Flattened independent variable array.
+        y : np.ndarray
+            Flattened dependent variable array.
+        weights : np.ndarray
+            Flattened weight array.
+        samples : int, default=10000
+            Number of retained DREAM samples requested from BUMPS.
+        burn : int, default=2000
+            Burn-in steps.
+        thin : int, default=10
+            Thinning interval.
+        chains : int | None, default=None
+            User-friendly alias for BUMPS DREAM population count.
+        population : int | None, default=None
+            BUMPS DREAM population count for advanced users.
+        seed : int | None, default=None
+            Best-effort random seed.
+        sampler_kwargs : dict | None, default=None
+            Additional keyword arguments forwarded to `bumps.fitters.fit`.
+        progress_callback : Callable[[dict], bool | None] | None, default=None
+            Optional callback for progress updates during sampling.  The
+            payload dict includes ``iteration`` (DREAM generation number) and
+            ``sampling: True``.
+        abort_test : Callable[[], bool] | None, default=None
+            Optional callback that returns ``True`` to signal that sampling
+            should be aborted. Called periodically during the DREAM sampling
+            loop.
+
+        Returns
+        -------
+        dict
+            Dictionary with keys ``'draws'``, ``'param_names'``, ``'state'``,
+            and ``'logp'``.
+
+        Raises
+        ------
+        ValueError
+            If the input shapes or weights are invalid, if both ``chains``
+            and ``population`` are provided with different values, or if
+            ``progress_callback`` is not callable.
+        FitError
+            If DREAM sampling was aborted by the user (via ``abort_test``).
+        Exception
+            Re-raised from DREAM fitting if any unexpected error occurs
+            (parameter values are restored beforehand).
+        """
+        from bumps.fitters import DreamFit
+        from bumps.names import FitProblem
+
+        x, y, weights = np.asarray(x), np.asarray(y), np.asarray(weights)
+
+        if y.shape != x.shape:
+            raise ValueError('x and y must have the same shape.')
+
+        if weights.shape != x.shape:
+            raise ValueError('Weights must have the same shape as x and y.')
+
+        if not np.isfinite(weights).all():
+            raise ValueError('Weights cannot be NaN or infinite.')
+
+        if (weights <= 0).any():
+            raise ValueError('Weights must be strictly positive and non-zero.')
+
+        # Build the BUMPS Curve model using the minimizer's existing machinery
+        model_func = self._make_model()
+        curve = model_func(x, y, weights)
+        problem = FitProblem(curve)
+
+        # Best-effort seed: sets numpy's global RNG state just before DREAM starts.
+        if seed is not None:
+            np.random.seed(seed)
+
+        # Resolve population parameter
+        if chains is not None and population is not None:
+            if chains != population:
+                raise ValueError(
+                    f'Conflicting population arguments: chains={chains}, population={population}. '
+                    'Only provide one.'
+                )
+            pop = chains
+        elif chains is not None:
+            pop = chains
+        elif population is not None:
+            pop = population
+        else:
+            pop = None
+
+        # Build DREAM kwargs
+        dream_kwargs: dict = {'samples': samples, 'burn': burn, 'thin': thin}
+        if pop is not None:
+            dream_kwargs['pop'] = pop
+        if sampler_kwargs:
+            dream_kwargs.update(sampler_kwargs)
+
+        # Build monitors (same pattern as classical Bumps.fit())
+        monitors = []
+        if progress_callback is not None:
+            if not callable(progress_callback):
+                raise ValueError('progress_callback must be callable')
+            # Compute total DREAM steps for progress display (burn + sampling generations)
+            pop_val = pop if pop else 10
+            _total_steps = burn + (samples + pop_val - 1) // pop_val
+            monitors.append(
+                BumpsProgressMonitor(
+                    problem,
+                    progress_callback,
+                    lambda problem, iteration, point, nllf: {
+                        **self._build_sample_progress_payload(problem, iteration, point, nllf),
+                        'total_steps': _total_steps,
+                    },
+                )
+            )
+
+        driver = FitDriver(
+            fitclass=DreamFit,
+            problem=problem,
+            monitors=monitors,
+            abort_test=abort_test or (lambda: False),
+            **dream_kwargs,
+        )
+        driver.clip()
+
+        from easyscience import global_object
+
+        stack_status = global_object.stack.enabled
+        global_object.stack.enabled = False
+
+        try:
+            x_opt, fx = driver.fit()
+            result_state = getattr(driver.fitter, 'state', None)
+            if result_state is None:
+                raise FitError('Sampling aborted by user')
+        except Exception:
+            self._restore_parameter_values()
+            raise
+        finally:
+            global_object.stack.enabled = stack_status
+
+        draws = result_state.draw().points
+        param_names = [p.name[len(MINIMIZER_PARAMETER_PREFIX) :] for p in problem._parameters]
+        logp = getattr(result_state, 'logp', None)
+
+        return {
+            'draws': draws,
+            'param_names': param_names,
+            'state': result_state,
+            'logp': logp,
+        }
+
+    def _build_sample_progress_payload(
+        self, problem, iteration: int, point: np.ndarray, nllf: float
+    ) -> dict:
+        """Build a progress payload for Bayesian DREAM sampling steps.
+
+        Called by :class:`BumpsProgressMonitor` at each DREAM generation.
+        The payload includes ``sampling: True`` so downstream consumers can
+        distinguish sampling progress from classical fitting progress.
+        """
+        parameter_values = self._current_parameter_snapshot(problem, point)
+        return {
+            'iteration': iteration,
+            'chi2': float(problem.chisq(nllf=nllf, norm=False)),
+            'reduced_chi2': float(problem.chisq(nllf=nllf, norm=True)),
+            'parameter_values': parameter_values,
+            'refresh_plots': False,
+            'finished': False,
+            'sampling': True,
+        }
+
     def _set_parameter_fit_result(
         self,
         fit_result: Any,

src/easyscience/fitting/multi_fitter.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 from typing import Callable
+from typing import Dict
 from typing import List
 from typing import Optional
 
@@ -188,3 +189,133 @@ def _post_compute_reshaping(
             fit_results_list.append(current_results)
             sp = ep
         return fit_results_list
+
+    def sample(
+        self,
+        x: List[np.ndarray],
+        y: List[np.ndarray],
+        weights: List[np.ndarray],
+        samples: int = 10000,
+        burn: int = 1000,
+        thin: int = 10,
+        chains: int | None = None,
+        population: int | None = None,
+        seed: int | None = None,
+        vectorized: bool = False,
+        sampler_kwargs: dict | None = None,
+        progress_callback: Callable[[dict], bool | None] | None = None,
+        abort_test: Callable[[], bool] | None = None,
+    ) -> Dict:
+        """Run Bayesian MCMC sampling using the BUMPS DREAM sampler.
+
+        Requires that the current minimizer is a BUMPS instance (i.e. the
+        minimizer was switched to ``AvailableMinimizers.Bumps`` or equivalent).
+
+        Parameters
+        ----------
+        x : List[np.ndarray]
+            List of independent variable arrays (one per dataset).
+        y : List[np.ndarray]
+            List of dependent variable arrays (one per dataset).
+        weights : List[np.ndarray]
+            List of weight arrays (one per dataset).
+        samples : int, default=10000
+            Number of retained DREAM samples requested from BUMPS.
+        burn : int, default=1000
+            Burn-in steps.
+        thin : int, default=10
+            Thinning interval.
+        chains : int | None, default=None
+            User-friendly alias for BUMPS DREAM population count.
+        population : int | None, default=None
+            BUMPS DREAM population count (``pop``) for advanced users.
+        seed : int | None, default=None
+            Best-effort random seed. BUMPS DREAM may use additional internal
+            RNG state that is not controlled by this seed, so exact
+            reproducibility is not guaranteed.
+        vectorized : bool, default=False
+            Whether the fit function expects vectorized (multidimensional)
+            input.
+        sampler_kwargs : dict | None, default=None
+            Additional keyword arguments forwarded to the BUMPS DREAM sampler
+            via `bumps.fitters.fit`.
+        progress_callback : Callable[[dict], bool | None] | None, default=None
+            Optional callback for progress updates during sampling.  The
+            payload dict includes ``iteration`` (DREAM generation number) and
+            ``sampling: True``.
+        abort_test : Callable[[], bool] | None, default=None
+            Optional callback that returns ``True`` to signal that sampling
+            should be aborted. Called periodically during the DREAM sampling
+            loop.
+
+        Returns
+        -------
+        Dict
+            Dictionary with keys ``'draws'``, ``'param_names'``, ``'state'``,
+            and ``'logp'``.
+
+        Raises
+        ------
+        RuntimeError
+            If the current minimizer is not a BUMPS instance.
+        ValueError
+            If both ``chains`` and ``population`` are provided with different
+            values.
+        """
+        # --- Alias resolution ---
+        if chains is not None and population is not None:
+            if chains != population:
+                raise ValueError(
+                    f'Conflicting population arguments: chains={chains}, population={population}. '
+                    'Only provide one.'
+                )
+            pop = chains
+        elif chains is not None:
+            pop = chains
+        elif population is not None:
+            pop = population
+        else:
+            pop = None
+
+        # Flatten multi-dataset arrays
+        x_fit, x_new, y_new, w_new, _dims = self._precompute_reshaping(
+            x, y, weights, vectorized=vectorized
+        )
+        self._dependent_dims = _dims
+
+        # Wrap fit functions for multi-dataset flattening, mirroring the
+        # ``Fitter.fit`` lifecycle: use the property setter so the minimizer
+        # is re-created with the wrapped fit function.
+        original_fit_func = self.fit_function
+        fit_fun_wrap = self._fit_function_wrapper(x_new, flatten=True)
+        self.fit_function = fit_fun_wrap
+
+        try:
+            minimizer = self.minimizer
+
+            # Verify it's a BUMPS minimizer (sampling only works with BUMPS/DREAM)
+            if not (hasattr(minimizer, 'package') and minimizer.package == 'bumps'):
+                raise RuntimeError(
+                    'Bayesian sampling requires a BUMPS minimizer. '
+                    'Use ``fitter.switch_minimizer(AvailableMinimizers.Bumps)`` first.'
+                )
+
+            # Delegate to the BUMPS minimizer's public sample method
+            result = minimizer.sample(
+                x=x_fit,
+                y=y_new,
+                weights=w_new,
+                samples=samples,
+                burn=burn,
+                thin=thin,
+                chains=None,  # alias already resolved into `pop`
+                population=pop,
+                seed=seed,
+                sampler_kwargs=sampler_kwargs,
+                progress_callback=progress_callback,
+                abort_test=abort_test,
+            )
+        finally:
+            self.fit_function = original_fit_func
+
+        return result