Interim updates

src/easyscience/fitting/fitter.py

@@ -209,6 +209,7 @@ def inner_fit_callable(
             y: np.ndarray,
             weights: Optional[np.ndarray] = None,
             vectorized: bool = False,
+            progress_callback: Optional[Callable[[dict], Optional[bool]]] = None,
             **kwargs,
         ) -> FitResults:
             """This is a wrapped callable which performs the actual
@@ -237,6 +238,7 @@ def inner_fit_callable(
                 weights=weights,
                 tolerance=self._tolerance,
                 max_evaluations=self._max_evaluations,
+                progress_callback=progress_callback,
                 **kwargs,
             )
 

src/easyscience/fitting/minimizers/minimizer_base.py

@@ -9,7 +9,6 @@
 from typing import Callable
 from typing import Dict
 from typing import List
-from typing import Optional
 from typing import Tuple
 from typing import Union
 
@@ -58,17 +57,23 @@ def enum(self) -> AvailableMinimizers:
     def name(self) -> str:
         return self._minimizer_enum.name
 
+    def _restore_parameter_values(self) -> None:
+        for key in self._cached_pars.keys():
+            self._cached_pars[key].value = self._cached_pars_vals[key][0]
+            self._cached_pars[key].error = self._cached_pars_vals[key][1]
+
     @abstractmethod
     def fit(
         self,
         x: np.ndarray,
         y: np.ndarray,
         weights: np.ndarray,
-        model: Optional[Callable] = None,
-        parameters: Optional[Parameter] = None,
-        method: Optional[str] = None,
-        tolerance: Optional[float] = None,
-        max_evaluations: Optional[int] = None,
+        model: Callable | None = None,
+        parameters: List[Parameter] | None = None,
+        method: str | None = None,
+        tolerance: float | None = None,
+        max_evaluations: int | None = None,
+        progress_callback: Callable[[dict], bool | None] | None = None,
         **kwargs,
     ) -> FitResults:
         """Perform a fit using the  engine.
@@ -88,7 +93,7 @@ def fit(
         """
 
     def evaluate(
-        self, x: np.ndarray, minimizer_parameters: Optional[dict[str, float]] = None, **kwargs
+        self, x: np.ndarray, minimizer_parameters: dict[str, float] | None = None, **kwargs
     ) -> np.ndarray:
         """Evaluate the fit function for values of x. Parameters used
         are either the latest or user supplied. If the parameters are
@@ -117,7 +122,7 @@ def evaluate(
 
         return self._fit_function(x, **minimizer_parameters, **kwargs)
 
-    def _get_method_kwargs(self, passed_method: Optional[str] = None) -> dict[str, str]:
+    def _get_method_kwargs(self, passed_method: str | None = None) -> dict[str, str]:
         if passed_method is not None:
             if passed_method not in self.supported_methods():
                 raise FitError(f'Method {passed_method} not available in {self.__class__}')
@@ -129,7 +134,7 @@ def _get_method_kwargs(self, passed_method: Optional[str] = None) -> dict[str, s
         return {}
 
     @abstractmethod
-    def convert_to_pars_obj(self, par_list: Optional[Union[list]] = None):
+    def convert_to_pars_obj(self, par_list: List[Parameter] | None = None):
         """Create an engine compatible container with the `Parameters`
         converted from the base object.
 

src/easyscience/fitting/minimizers/minimizer_bumps.py

@@ -7,11 +7,12 @@
 import warnings
 from typing import Callable
 from typing import List
-from typing import Optional
 
 import numpy as np
 from bumps.fitters import FIT_AVAILABLE_IDS
-from bumps.fitters import fit as bumps_fit
+from bumps.fitters import FITTERS
+from bumps.fitters import FitDriver
+from bumps.monitor import Monitor
 from bumps.names import Curve
 from bumps.names import FitProblem
 from bumps.parameter import Parameter as BumpsParameter
@@ -31,6 +32,21 @@
 FIT_AVAILABLE_IDS_FILTERED.remove('pt')
 
 
+class _StepCounterMonitor(Monitor):
+    """Lightweight monitor that ensures step count is recorded in
+    history.
+    """
+
+    def __init__(self):
+        self.last_step = 0
+
+    def config_history(self, history):
+        history.requires(step=1)
+
+    def __call__(self, history):
+        self.last_step = int(history.step[0])
+
+
 class _EvalCounter:
     def __init__(self, fn: Callable):
         self._fn = fn
@@ -44,6 +60,25 @@
         return self._fn(*args, **kwargs)
 
 
+class _BumpsProgressMonitor(Monitor):
+    def __init__(self, problem, callback, payload_builder):
+        self._problem = problem
+        self._callback = callback
+        self._payload_builder = payload_builder
+
+    def config_history(self, history):
+        history.requires(step=1, point=1, value=1)
+
+    def __call__(self, history):
+        payload = self._payload_builder(
+            problem=self._problem,
+            iteration=int(history.step[0]),
+            point=np.asarray(history.point[0]),
+            nllf=float(history.value[0]),
+        )
+        self._callback(payload)
+
+
 class Bumps(MinimizerBase):
     """
     This is a wrapper to Bumps: https://bumps.readthedocs.io/
@@ -56,7 +91,7 @@
         self,
         obj,  #: ObjBase,
         fit_function: Callable,
-        minimizer_enum: Optional[AvailableMinimizers] = None,
+        minimizer_enum: AvailableMinimizers | None = None,
     ):  # todo after constraint changes, add type hint: obj: ObjBase  # noqa: E501
         """Initialize the fitting engine with a `ObjBase` and an
         arbitrary fitting function.
@@ -70,7 +105,7 @@
         """
         super().__init__(obj=obj, fit_function=fit_function, minimizer_enum=minimizer_enum)
         self._p_0 = {}
-        self._eval_counter: Optional[_EvalCounter] = None
+        self._eval_counter: _EvalCounter | None = None
 
     @staticmethod
     def all_methods() -> List[str]:
@@ -87,16 +122,17 @@
         x: np.ndarray,
         y: np.ndarray,
         weights: np.ndarray,
-        model: Optional[Callable] = None,
-        parameters: Optional[Parameter] = None,
-        method: Optional[str] = None,
-        tolerance: Optional[float] = None,
-        max_evaluations: Optional[int] = None,
-        minimizer_kwargs: Optional[dict] = None,
-        engine_kwargs: Optional[dict] = None,
+        model: Callable | None = None,
+        parameters: List[Parameter] | None = None,
+        method: str | None = None,
+        tolerance: float | None = None,
+        max_evaluations: int | None = None,
+        progress_callback: Callable[[dict], bool | None] | None = None,
+        minimizer_kwargs: dict | None = None,
+        engine_kwargs: dict | None = None,
         **kwargs,
     ) -> FitResults:
-        """Perform a fit using the lmfit engine.
+        """Perform a fit using the BUMPS engine.
 
         :param x: points to be calculated at
         :type x: np.ndarray
@@ -105,17 +141,14 @@
         :param weights: Weights for supplied measured points
         :type weights: np.ndarray
         :param model: Optional Model which is being fitted to
-        :type model: lmModel
         :param parameters: Optional parameters for the fit
         :type parameters: List[BumpsParameter]
-        :param kwargs: Additional arguments for the fitting function.
         :param method: Method for minimization
         :type method: str
+        :param progress_callback: Optional callback for progress updates
+        :type progress_callback: Callable
         :return: Fit results
-        :rtype: ModelResult For standard least squares, the weights
-            should be 1/sigma, where sigma is the standard deviation of
-            the measurement. For unweighted least squares, these should
-            be 1.
+        :rtype: FitResults
         """
         method_dict = self._get_method_kwargs(method)
 
@@ -156,27 +189,100 @@
         self._p_0 = {f'p{key}': self._cached_pars[key].value for key in self._cached_pars.keys()}
 
         problem = FitProblem(model)
+
+        method_str = method_dict.get('method', self._method)
+        fitclass = self._resolve_fitclass(method_str)
+
+        step_counter = _StepCounterMonitor()
+        monitors = [step_counter]
+        if progress_callback is not None:
+            if not callable(progress_callback):
+                raise ValueError('progress_callback must be callable')
+            monitors.append(
+                _BumpsProgressMonitor(problem, progress_callback, self._build_progress_payload)
+            )
+
+        driver = FitDriver(
+            fitclass=fitclass,
+            problem=problem,
+            monitors=monitors,
+            **minimizer_kwargs,
+            **kwargs,
+        )
+        driver.clip()
+
         # Why do we do this? Because a fitting template has to have global_object instantiated outside pre-runtime
         from easyscience import global_object
 
         stack_status = global_object.stack.enabled
         global_object.stack.enabled = False
 
         try:
-            model_results = bumps_fit(problem, **method_dict, **minimizer_kwargs, **kwargs)
+            # Drive the fit through the local FitDriver instance so the supplied
+            # `monitors` (including the optional progress callback monitor) are
+            # invoked. `bumps.fitters.fit` constructs its own driver.
+            x, fx = driver.fit()
+            from scipy.optimize import OptimizeResult
+
+            model_results = OptimizeResult(
+                x=x,
+                dx=driver.stderr(),
+                fun=fx,
+                success=True,
+                status=0,
+                message='successful termination',
+                nit=driver.monitor_runner.history.step[0],
+            )
+            model_results.state = driver.fitter.state
             self._set_parameter_fit_result(model_results, stack_status, problem._parameters)
             results = self._gen_fit_results(
                 model_results,
                 max_evaluations=max_evaluations,
                 tolerance=tolerance,
             )
         except Exception as e:
-            for key in self._cached_pars.keys():
-                self._cached_pars[key].value = self._cached_pars_vals[key][0]
+            self._restore_parameter_values()
             raise FitError(e)
+        finally:
+            global_object.stack.enabled = stack_status
         return results
 
-    def convert_to_pars_obj(self, par_list: Optional[List] = None) -> List[BumpsParameter]:
+    @staticmethod
+    def _resolve_fitclass(method: str):
+        for fitclass in FITTERS:
+            if fitclass.id == method:
+                return fitclass
+        raise FitError(f'Unknown BUMPS fitting method: {method}')
+
+    def _build_progress_payload(
+        self, problem, iteration: int, point: np.ndarray, nllf: float
+    ) -> dict:
+        # Use the nllf already computed by the fitter to avoid a costly
+        # model re-evaluation, and let BUMPS apply its own chisq scaling.
+        chi2 = float(problem.chisq(nllf=nllf, norm=False))
+        reduced_chi2 = float(problem.chisq(nllf=nllf, norm=True))
+
+        parameter_values = self._current_parameter_snapshot(problem, point)
+
+        return {
+            'iteration': iteration,
+            'chi2': chi2,
+            'reduced_chi2': reduced_chi2,
+            'parameter_values': parameter_values,
+            'refresh_plots': False,
+            'finished': False,
+        }
+
+    def _current_parameter_snapshot(self, problem, point: np.ndarray) -> dict:
+        labels = problem.labels()
+        values = problem.getp() if point is None else point
+        snapshot = {}
+        for label, value in zip(labels, values):
+            dict_name = label[len(MINIMIZER_PARAMETER_PREFIX) :]
+            snapshot[dict_name] = float(value)
+        return snapshot
+
+    def convert_to_pars_obj(self, par_list: List[Parameter] | None = None) -> List[BumpsParameter]:
         """Create a container with the `Parameters` converted from the
         base object.
 
@@ -211,7 +317,7 @@
             fixed=obj.fixed,
         )
 
-    def _make_model(self, parameters: Optional[List[BumpsParameter]] = None) -> Callable:
+    def _make_model(self, parameters: List[BumpsParameter] | None = None) -> Callable:
         """Generate a bumps model from the supplied `fit_function` and
         parameters in the base object. Note that this makes a callable
         as it needs to be initialized with *x*, *y*, *weights*
@@ -244,43 +350,51 @@
         return _outer(self)
 
     def _set_parameter_fit_result(
-        self, fit_result, stack_status: bool, par_list: List[BumpsParameter]
+        self,
+        fit_result,
+        stack_status: bool,
+        par_list: List[BumpsParameter],
     ):
         """Update parameters to their final values and assign a std
         error to them.
 
-        :param fit_result: Fit object which contains info on the fit
-        :return: None
-        :rtype: noneType
+        :param fit_result: BUMPS OptimizeResult containing best-fit
+            values and errors
+        :param stack_status: Whether the undo stack was enabled
+        :param par_list: List of BUMPS parameter objects
         """
         from easyscience import global_object
 
         pars = self._cached_pars
+        x_result = np.asarray(fit_result.x)
+        stderr = np.asarray(fit_result.dx)
 
         if stack_status:
-            for name in pars.keys():
-                pars[name].value = self._cached_pars_vals[name][0]
-                pars[name].error = self._cached_pars_vals[name][1]
+            self._restore_parameter_values()
             global_object.stack.enabled = True
             global_object.stack.beginMacro('Fitting routine')
 
         for index, name in enumerate([par.name for par in par_list]):
             dict_name = name[len(MINIMIZER_PARAMETER_PREFIX) :]
-            pars[dict_name].value = fit_result.x[index]
-            pars[dict_name].error = fit_result.dx[index]
+            pars[dict_name].value = x_result[index]
+            pars[dict_name].error = stderr[index]
         if stack_status:
             global_object.stack.endMacro()
 
     def _gen_fit_results(
         self,
         fit_results,
-        max_evaluations: Optional[int] = None,
-        tolerance: Optional[float] = None,
+        max_evaluations: int | None = None,
+        tolerance: float | None = None,
         **kwargs,
     ) -> FitResults:
         """Convert fit results into the unified `FitResults` format.
 
-        :param fit_result: Fit object which contains info on the fit
+        :param x_result: Optimized parameter values from FitDriver
+        :param fx: Final objective function value
+        :param driver: The FitDriver instance
+        :param n_evaluations: Number of iterations completed
+        :param max_evaluations: Maximum evaluations budget (if set)
         :return: fit results container
         :rtype: FitResults
         """
@@ -314,7 +428,6 @@
         item = {}
         for index, name in enumerate(self._cached_model.pars.keys()):
             dict_name = name[len(MINIMIZER_PARAMETER_PREFIX) :]
-
             item[name] = pars[dict_name].value
 
         results.p0 = self._p_0
@@ -351,5 +464,4 @@
         results.minimizer_engine = self.__class__
         results.fit_args = None
         results.engine_result = fit_results
-        # results.check_sanity()
         return results