Add logging with state to Bayesian Optimizer

bayes_opt/acquisition.py

@@ -69,18 +69,57 @@ def __init__(self, random_state: int | RandomState | None = None) -> None:
             self.random_state = RandomState()
         self.i = 0
 
+    def _serialize_random_state(self) -> dict | None:
+        """Convert random state to JSON serializable format."""
+        if self.random_state is not None:
+            state = self.random_state.get_state()
+            return {
+                'bit_generator': state[0],
+                'state': state[1].tolist(),  # Convert numpy array to list
+                'pos': state[2],
+                'has_gauss': state[3],
+                'cached_gaussian': state[4]
+            }
+        return None
+
+    def _deserialize_random_state(self, state_dict: dict | None) -> None:
+        """Restore random state from JSON serializable format."""
+        if state_dict is not None:
+            if self.random_state is None:
+                self.random_state = RandomState()
+            state = (
+                state_dict['bit_generator'],
+                np.array(state_dict['state'], dtype=np.uint32),
+                state_dict['pos'],
+                state_dict['has_gauss'],
+                state_dict['cached_gaussian']
+            )
+            self.random_state.set_state(state)
+
     @abc.abstractmethod
     def base_acq(self, *args: Any, **kwargs: Any) -> NDArray[Float]:
         """Provide access to the base acquisition function."""
 
-    def _fit_gp(self, gp: GaussianProcessRegressor, target_space: TargetSpace) -> None:
-        # Sklearn's GP throws a large number of warnings at times, but
-        # we don't really need to see them here.
-        with warnings.catch_warnings():
-            warnings.simplefilter("ignore")
-            gp.fit(target_space.params, target_space.target)
-            if target_space.constraint is not None:
-                target_space.constraint.fit(target_space.params, target_space._constraint_values)
+    def get_acquisition_params(self) -> dict[str, Any]:
+        """Get the acquisition function parameters.
+
+        Returns
+        -------
+        dict
+            Dictionary containing the acquisition function parameters.
+            All values must be JSON serializable.
+        """
+        return {}
+
+    def set_acquisition_params(self, params: dict[str, Any]) -> None:
+        """Set the acquisition function parameters.
+
+        Parameters
+        ----------
+        params : dict
+            Dictionary containing the acquisition function parameters.
+        """
+        pass
 
     def suggest(
         self,
@@ -128,6 +167,15 @@ def suggest(
 
         acq = self._get_acq(gp=gp, constraint=target_space.constraint)
         return self._acq_min(acq, target_space, n_random=n_random, n_l_bfgs_b=n_l_bfgs_b)
+
+    def _fit_gp(self, gp: GaussianProcessRegressor, target_space: TargetSpace) -> None:
+        # Sklearn's GP throws a large number of warnings at times, but
+        # we don't really need to see them here.
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            gp.fit(target_space.params, target_space.target)
+            if target_space.constraint is not None:
+                target_space.constraint.fit(target_space.params, target_space._constraint_values)
 
     def _get_acq(
         self, gp: GaussianProcessRegressor, constraint: ConstraintModel | None = None
@@ -453,6 +501,20 @@ def decay_exploration(self) -> None:
             self.exploration_decay_delay is None or self.exploration_decay_delay <= self.i
         ):
             self.kappa = self.kappa * self.exploration_decay
+
+    def get_acquisition_params(self) -> dict:
+        return {
+            "kappa": self.kappa,
+            "exploration_decay": self.exploration_decay,
+            "exploration_decay_delay": self.exploration_decay_delay,
+            "random_state": self._serialize_random_state()
+        }
+
+    def set_acquisition_params(self, params: dict) -> None:
+        self.kappa = params["kappa"]
+        self.exploration_decay = params["exploration_decay"]
+        self.exploration_decay_delay = params["exploration_decay_delay"]
+        self._deserialize_random_state(params["random_state"])
 
 
 class ProbabilityOfImprovement(AcquisitionFunction):
@@ -586,6 +648,21 @@ def decay_exploration(self) -> None:
             self.exploration_decay_delay is None or self.exploration_decay_delay <= self.i
         ):
             self.xi = self.xi * self.exploration_decay
+
+    def get_acquisition_params(self) -> dict:
+        """Get the acquisition function parameters."""
+        return {
+            "xi": self.xi,
+            "exploration_decay": self.exploration_decay,
+            "exploration_decay_delay": self.exploration_decay_delay,
+            "random_state": self._serialize_random_state()
+        }
+
+    def set_acquisition_params(self, params: dict) -> None:
+        self.xi = params["xi"]
+        self.exploration_decay = params["exploration_decay"]
+        self.exploration_decay_delay = params["exploration_decay_delay"]
+        self._deserialize_random_state(params["random_state"])
 
 
 class ExpectedImprovement(AcquisitionFunction):
@@ -727,6 +804,20 @@ def decay_exploration(self) -> None:
             self.exploration_decay_delay is None or self.exploration_decay_delay <= self.i
         ):
             self.xi = self.xi * self.exploration_decay
+
+    def get_acquisition_params(self) -> dict:
+        return {
+            "xi": self.xi,
+            "exploration_decay": self.exploration_decay,
+            "exploration_decay_delay": self.exploration_decay_delay,
+            "random_state": self._serialize_random_state()
+        }
+
+    def set_acquisition_params(self, params: dict) -> None:
+        self.xi = params["xi"]
+        self.exploration_decay = params["exploration_decay"]
+        self.exploration_decay_delay = params["exploration_decay_delay"]
+        self._deserialize_random_state(params["random_state"])
 
 
 class ConstantLiar(AcquisitionFunction):
@@ -917,6 +1008,24 @@ def suggest(
         self.dummies.append(x_max)
 
         return x_max
+
+    def get_acquisition_params(self) -> dict:
+        return {
+            "dummies": [dummy.tolist() for dummy in self.dummies],
+            "base_acquisition_params": self.base_acquisition.get_acquisition_params(),
+            "strategy": self.strategy,
+            "atol": self.atol,
+            "rtol": self.rtol,
+            "random_state": self._serialize_random_state()
+        }
+
+    def set_acquisition_params(self, params: dict) -> None:
+        self.dummies = [np.array(dummy) for dummy in params["dummies"]]
+        self.base_acquisition.set_acquisition_params(params["base_acquisition_params"])
+        self.strategy = params["strategy"]
+        self.atol = params["atol"]
+        self.rtol = params["rtol"]
+        self._deserialize_random_state(params["random_state"])
 
 
 class GPHedge(AcquisitionFunction):
@@ -1035,3 +1144,28 @@ def suggest(
         self.previous_candidates = np.array(x_max)
         idx = self._sample_idx_from_softmax_gains()
         return x_max[idx]
+
+    def get_acquisition_params(self) -> dict:
+        return {
+            "base_acquisitions_params": [acq.get_acquisition_params() for acq in self.base_acquisitions],
+            "gains": self.gains.tolist(),
+            "previous_candidates": self.previous_candidates.tolist() if self.previous_candidates is not None else None,
+            "random_states": [acq._serialize_random_state() for acq in self.base_acquisitions] + [self._serialize_random_state()]
+        }
+
+    def set_acquisition_params(self, params: dict) -> None:
+        for acq, acq_params, random_state in zip(
+            self.base_acquisitions, 
+            params["base_acquisitions_params"],
+            params["random_states"][:-1]
+        ):
+            acq.set_acquisition_params(acq_params)
+            acq._deserialize_random_state(random_state)
+
+        self.gains = np.array(params["gains"])
+        self.previous_candidates = (np.array(params["previous_candidates"]) 
+                                  if params["previous_candidates"] is not None 
+                                  else None)
+
+        self._deserialize_random_state(params["random_states"][-1])
+

bayes_opt/bayesian_optimization.py

@@ -10,10 +10,16 @@
 from typing import TYPE_CHECKING, Any
 from warnings import warn
 
+import json
+from pathlib import Path
+from os import PathLike
+
 import numpy as np
 from sklearn.gaussian_process import GaussianProcessRegressor
 from sklearn.gaussian_process.kernels import Matern
 
+from scipy.optimize import NonlinearConstraint
+
 from bayes_opt import acquisition
 from bayes_opt.constraint import ConstraintModel
 from bayes_opt.domain_reduction import DomainTransformer
@@ -28,7 +34,6 @@
 
     from numpy.random import RandomState
     from numpy.typing import NDArray
-    from scipy.optimize import NonlinearConstraint
 
     from bayes_opt.acquisition import AcquisitionFunction
     from bayes_opt.constraint import ConstraintModel
@@ -356,3 +361,122 @@ def set_gp_params(self, **params: Any) -> None:
         if "kernel" in params:
             params["kernel"] = wrap_kernel(kernel=params["kernel"], transform=self._space.kernel_transform)
         self._gp.set_params(**params)
+
+    def save_state(self, path: str | PathLike[str]) -> None:
+        """Save complete state for reconstruction of the optimizer.
+
+        Parameters
+        ----------
+        path : str or PathLike
+            Path to save the optimization state
+
+        Raises
+        ------
+        ValueError
+            If attempting to save state before collecting any samples.
+        """
+        if len(self._space) == 0:
+            raise ValueError(
+                "Cannot save optimizer state before collecting any samples. "
+                "Please probe or register at least one point before saving."
+            )
+
+        random_state = None
+        if self._random_state is not None:
+            state_tuple = self._random_state.get_state()
+            random_state = {
+                'bit_generator': state_tuple[0],
+                'state': state_tuple[1].tolist(),
+                'pos': state_tuple[2],
+                'has_gauss': state_tuple[3],
+                'cached_gaussian': state_tuple[4],
+            }
+
+        # Get constraint values if they exist
+        constraint_values = (self._space._constraint_values.tolist() 
+                           if self.is_constrained 
+                           else None)
+        acquisition_params = self._acquisition_function.get_acquisition_params()
+        state = {
+            "pbounds": {
+                key: self._space._bounds[i].tolist() 
+                for i, key in enumerate(self._space.keys)
+            },
+            # Add current transformed bounds if using bounds transformer
+            "transformed_bounds": (
+                self._space.bounds.tolist() 
+                if self._bounds_transformer 
+                else None
+            ),
+            "keys": self._space.keys,
+            "params": np.array(self._space.params).tolist(),
+            "target": self._space.target.tolist(),
+            "constraint_values": constraint_values,
+            "gp_params": {
+                "kernel": self._gp.kernel.get_params(),
+                "alpha": self._gp.alpha,
+                "normalize_y": self._gp.normalize_y,
+                "n_restarts_optimizer": self._gp.n_restarts_optimizer,
+            },
+            "allow_duplicate_points": self._allow_duplicate_points,
+            "verbose": self._verbose,
+            "random_state": random_state,
+            "acquisition_params": acquisition_params,
+        }
+
+        with Path(path).open('w') as f:
+            json.dump(state, f, indent=2)
+
+    def load_state(self, path: str | PathLike[str]) -> None:
+        with Path(path).open('r') as file:
+            state = json.load(file)
+
+        params_array = np.asarray(state["params"], dtype=np.float64)
+        target_array = np.asarray(state["target"], dtype=np.float64)
+        constraint_array = (np.array(state["constraint_values"]) 
+                          if state["constraint_values"] is not None 
+                          else None)
+
+        for i in range(len(params_array)):
+            params = self._space.array_to_params(params_array[i])
+            target = target_array[i]
+            constraint = constraint_array[i] if constraint_array is not None else None
+            self.register(
+                params=params,
+                target=target,
+                constraint_value=constraint
+            )
+
+        self._acquisition_function.set_acquisition_params(state["acquisition_params"])
+
+        if state.get("transformed_bounds") and self._bounds_transformer:
+            new_bounds = {
+                key: bounds for key, bounds in zip(
+                    self._space.keys, 
+                    np.array(state["transformed_bounds"])
+                )
+            }
+            self._space.set_bounds(new_bounds)
+            self._bounds_transformer.initialize(self._space)
+
+        self._gp.set_params(**state["gp_params"])
+        if isinstance(self._gp.kernel, dict):
+            kernel_params = self._gp.kernel
+            self._gp.kernel = Matern(
+                length_scale=kernel_params['length_scale'],
+                length_scale_bounds=tuple(kernel_params['length_scale_bounds']),
+                nu=kernel_params['nu']
+            )
+
+        self._gp.fit(self._space.params, self._space.target)
+
+        if state["random_state"] is not None:
+            random_state_tuple = (
+                state["random_state"]["bit_generator"],
+                np.array(state["random_state"]["state"], dtype=np.uint32),
+                state["random_state"]["pos"],
+                state["random_state"]["has_gauss"],
+                state["random_state"]["cached_gaussian"],
+            )
+            self._random_state.set_state(random_state_tuple)
+