diff --git a/bayes_opt/bayesian_optimization.py b/bayes_opt/bayesian_optimization.py index 24397b50..4cd4eae4 100644 --- a/bayes_opt/bayesian_optimization.py +++ b/bayes_opt/bayesian_optimization.py @@ -406,13 +406,13 @@ def set_gp_params(self, **params: Any) -> None: 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. + def _state_to_dict(self) -> dict[str, Any]: + """Convert optimizer state to a dictionary. - Parameters - ---------- - path : str or PathLike - Path to save the optimization state + Returns + ------- + dict + Dictionary containing the complete optimizer state. """ random_state = None if self._random_state is not None: @@ -428,7 +428,7 @@ def save_state(self, path: str | PathLike[str]) -> None: # 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 = { + return { "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), @@ -448,20 +448,14 @@ def save_state(self, path: str | PathLike[str]) -> None: "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: - """Load optimizer state from a JSON file. + def _load_state_dict(self, state: dict[str, Any]) -> None: + """Load optimizer state from a dictionary. Parameters ---------- - path : str or PathLike - Path to the JSON file containing the optimizer state. + state : dict + Dictionary containing the optimizer state. """ - 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 = ( @@ -504,3 +498,27 @@ def load_state(self, path: str | PathLike[str]) -> None: state["random_state"]["cached_gaussian"], ) self._random_state.set_state(random_state_tuple) + + 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 + """ + state = self._state_to_dict() + with Path(path).open("w") as f: + json.dump(state, f, indent=2) + + def load_state(self, path: str | PathLike[str]) -> None: + """Load optimizer state from a JSON file. + + Parameters + ---------- + path : str or PathLike + Path to the JSON file containing the optimizer state. + """ + with Path(path).open("r") as file: + state = json.load(file) + self._load_state_dict(state)