Add logging with state to Bayesian Optimizer (#547)

* Add functionality to save and load state of the BayesianOptimization

* Update basic-tour with new save and load functionality

* move load stateful path to optional argument in class instantiation

* add test for string params, update tests with new load functionality

* updated basic tour with updated paths

* add the random state to the set of things to list of saved items

* move state loading to separate function, add functionality for saving acquisition function state

* use new loading schema

* update tests, add integration tests for saving and loading acquisition functions

* undo abstractmethod implementation for get and set state saving functionality

* reorganize state saving and loading for consistency

* move integration tests into acquisition

* remove unndecessary test, add tests for domain reduction and custom parameters

* make test more comprehensive

* add test logs

* sync execution counts from basic tour

* linting, whitespace removal, import structuring

* ruff fix for string literal in error message

* fix ruff complaints

* make all side param comparisons almost equal to account for slight numpy differences

* reformat array comparison check

* upgrade poetry2.0 & apply pep621 (#545)

* chore: upgrade poetry2.0 & apply pep621

* fix: replace poetry action(not support 2.0)

* fix: exclude one matrix

* chore: split numpy deps

* fix: numpy constraints

* fix: install root

* chore: use install-poetry

* Fix coverage report (#552)

* remove unnecessary files, have acquisition baseclass functions raise errors

* remove duplicate acquisition functions random state

* ruff format

* add  type hints for base acquisition get/set functions

* remove noreturn

* remove former saving functionality from notebooks

* increase legibility of custom acquisition example

* explicitly stating the optionality of the saving and loading in custom acq functions

---------

Co-authored-by: phi-friday <[email protected]>
Co-authored-by: till-m <[email protected]>

Author: 3 people

Diff for: .gitignore

@@ -36,4 +36,8 @@ docsrc/*.ipynb
 docsrc/static/*
 docsrc/README.md
 
-poetry.lock
+poetry.lock
+
+# Add log files and optimizer state files to gitignore
+examples/logs.log
+examples/optimizer_state.json

Diff for: bayes_opt/acquisition.py

@@ -69,6 +69,33 @@ 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."""
@@ -82,6 +109,34 @@ def _fit_gp(self, gp: GaussianProcessRegressor, target_space: TargetSpace) -> No
             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 parameters of the acquisition function.
+
+        Returns
+        -------
+        dict
+            The parameters of the acquisition function.
+        """
+        error_msg = (
+            "Custom AcquisitionFunction subclasses must implement their own get_acquisition_params method."
+        )
+        raise NotImplementedError(error_msg)
+
+    def set_acquisition_params(self, **params) -> None:
+        """
+        Set the parameters of the acquisition function.
+
+        Parameters
+        ----------
+        **params : dict
+            The parameters of the acquisition function.
+        """
+        error_msg = (
+            "Custom AcquisitionFunction subclasses must implement their own set_acquisition_params method."
+        )
+        raise NotImplementedError(error_msg)
+
     def suggest(
         self,
         gp: GaussianProcessRegressor,
@@ -454,6 +509,34 @@ def decay_exploration(self) -> None:
         ):
             self.kappa = self.kappa * self.exploration_decay
 
+    def get_acquisition_params(self) -> dict:
+        """Get the current acquisition function parameters.
+
+        Returns
+        -------
+        dict
+            Dictionary containing the current acquisition function parameters.
+        """
+        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:
+        """Set the acquisition function parameters.
+
+        Parameters
+        ----------
+        params : dict
+            Dictionary containing the acquisition function parameters.
+        """
+        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):
     r"""Probability of Improvement acqusition function.
@@ -587,6 +670,34 @@ def decay_exploration(self) -> None:
         ):
             self.xi = self.xi * self.exploration_decay
 
+    def get_acquisition_params(self) -> dict:
+        """Get the current acquisition function parameters.
+
+        Returns
+        -------
+        dict
+            Dictionary containing the current 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:
+        """Set the acquisition function parameters.
+
+        Parameters
+        ----------
+        params : dict
+            Dictionary containing the acquisition function parameters.
+        """
+        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):
     r"""Expected Improvement acqusition function.
@@ -728,6 +839,34 @@ def decay_exploration(self) -> None:
         ):
             self.xi = self.xi * self.exploration_decay
 
+    def get_acquisition_params(self) -> dict:
+        """Get the current acquisition function parameters.
+
+        Returns
+        -------
+        dict
+            Dictionary containing the current 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:
+        """Set the acquisition function parameters.
+
+        Parameters
+        ----------
+        params : dict
+            Dictionary containing the acquisition function parameters.
+        """
+        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):
     """Constant Liar acquisition function.
@@ -918,6 +1057,38 @@ def suggest(
 
         return x_max
 
+    def get_acquisition_params(self) -> dict:
+        """Get the current acquisition function parameters.
+
+        Returns
+        -------
+        dict
+            Dictionary containing the current acquisition function parameters.
+        """
+        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:
+        """Set the acquisition function parameters.
+
+        Parameters
+        ----------
+        params : dict
+            Dictionary containing the acquisition function parameters.
+        """
+        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):
     """GPHedge acquisition function.
@@ -1035,3 +1206,38 @@ 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:
+        """Get the current acquisition function parameters.
+
+        Returns
+        -------
+        dict
+            Dictionary containing the current acquisition function parameters.
+        """
+        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,
+            "gphedge_random_state": self._serialize_random_state(),
+        }
+
+    def set_acquisition_params(self, params: dict) -> None:
+        """Set the acquisition function parameters.
+
+        Parameters
+        ----------
+        params : dict
+            Dictionary containing the acquisition function parameters.
+        """
+        for acq, acq_params in zip(self.base_acquisitions, params["base_acquisitions_params"]):
+            acq.set_acquisition_params(acq_params)
+
+        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["gphedge_random_state"])