diff --git a/bayes_opt/acquisition.py b/bayes_opt/acquisition.py index 4cd6828a..7765330f 100644 --- a/bayes_opt/acquisition.py +++ b/bayes_opt/acquisition.py @@ -565,6 +565,7 @@ def get_acquisition_params(self) -> dict[str, Any]: "kappa": self.kappa, "exploration_decay": self.exploration_decay, "exploration_decay_delay": self.exploration_decay_delay, + "i": self.i, } def set_acquisition_params(self, params: dict[str, Any]) -> None: @@ -578,6 +579,7 @@ def set_acquisition_params(self, params: dict[str, Any]) -> None: self.kappa = params["kappa"] self.exploration_decay = params["exploration_decay"] self.exploration_decay_delay = params["exploration_decay_delay"] + self.i = params.get("i", self.i) class ProbabilityOfImprovement(AcquisitionFunction): @@ -745,6 +747,7 @@ def get_acquisition_params(self) -> dict[str, Any]: "xi": self.xi, "exploration_decay": self.exploration_decay, "exploration_decay_delay": self.exploration_decay_delay, + "i": self.i, } def set_acquisition_params(self, params: dict[str, Any]) -> None: @@ -758,6 +761,7 @@ def set_acquisition_params(self, params: dict[str, Any]) -> None: self.xi = params["xi"] self.exploration_decay = params["exploration_decay"] self.exploration_decay_delay = params["exploration_decay_delay"] + self.i = params.get("i", self.i) class ExpectedImprovement(AcquisitionFunction): @@ -934,6 +938,7 @@ def get_acquisition_params(self) -> dict[str, Any]: "xi": self.xi, "exploration_decay": self.exploration_decay, "exploration_decay_delay": self.exploration_decay_delay, + "i": self.i, } def set_acquisition_params(self, params: dict[str, Any]) -> None: @@ -947,6 +952,7 @@ def set_acquisition_params(self, params: dict[str, Any]) -> None: self.xi = params["xi"] self.exploration_decay = params["exploration_decay"] self.exploration_decay_delay = params["exploration_decay_delay"] + self.i = params.get("i", self.i) class ConstantLiar(AcquisitionFunction): diff --git a/tests/test_acquisition_state.py b/tests/test_acquisition_state.py new file mode 100644 index 00000000..e9ad8d8d --- /dev/null +++ b/tests/test_acquisition_state.py @@ -0,0 +1,87 @@ +from __future__ import annotations + +import json + +import pytest + +from bayes_opt import BayesianOptimization +from bayes_opt.acquisition import ExpectedImprovement, ProbabilityOfImprovement, UpperConfidenceBound + + +@pytest.mark.parametrize( + ("acquisition_type", "parameter"), + [(UpperConfidenceBound, "kappa"), (ProbabilityOfImprovement, "xi"), (ExpectedImprovement, "xi")], +) +@pytest.mark.parametrize( + ("delay", "suggestions_before_save"), + [(3, 2), (3, 4), (None, 2), (0, 2), (3, 0)], + ids=["before-delay", "after-delay", "no-delay", "zero-delay", "unused-acquisition"], +) +def test_save_load_preserves_exploration_schedule( + tmp_path, acquisition_type, parameter, delay, suggestions_before_save +): + def make_optimizer(): + acquisition = acquisition_type( + **{parameter: 1.0}, exploration_decay=0.5, exploration_decay_delay=delay + ) + optimizer = BayesianOptimization( + f=lambda x: -((x - 0.3) ** 2), + pbounds={"x": (-1.0, 1.0)}, + acquisition_function=acquisition, + random_state=7, + verbose=0, + ) + optimizer.set_gp_params(n_restarts_optimizer=0) + return optimizer + + original = make_optimizer() + for x in (-1.0, 0.0, 1.0): + original.probe({"x": x}, lazy=False) + for _ in range(suggestions_before_save): + original.suggest() + + state_path = tmp_path / "optimizer.json" + original.save_state(state_path) + restored = make_optimizer() + restored.load_state(state_path) + + assert restored.res == original.res + assert getattr(restored.acquisition_function, parameter) == getattr( + original.acquisition_function, parameter + ) + + for _ in range(2): + assert restored.suggest() == pytest.approx(original.suggest()) + assert getattr(restored.acquisition_function, parameter) == getattr( + original.acquisition_function, parameter + ) + + +@pytest.mark.parametrize( + ("acquisition_type", "parameter"), + [(UpperConfidenceBound, "kappa"), (ProbabilityOfImprovement, "xi"), (ExpectedImprovement, "xi")], +) +def test_load_legacy_acquisition_state(tmp_path, acquisition_type, parameter): + acquisition = acquisition_type(**{parameter: 1.0}, exploration_decay=0.5, exploration_decay_delay=3) + original = BayesianOptimization( + f=lambda x: x, pbounds={"x": (-1.0, 1.0)}, acquisition_function=acquisition, random_state=7, verbose=0 + ) + state_path = tmp_path / "legacy.json" + original.save_state(state_path) + state = json.loads(state_path.read_text()) + state["acquisition_params"].pop("i", None) + state_path.write_text(json.dumps(state)) + + restored_acquisition = acquisition_type(**{parameter: 2.0}) + restored = BayesianOptimization( + f=lambda x: x, + pbounds={"x": (-1.0, 1.0)}, + acquisition_function=restored_acquisition, + random_state=7, + verbose=0, + ) + restored.load_state(state_path) + assert getattr(restored_acquisition, parameter) == 1.0 + assert restored_acquisition.exploration_decay == 0.5 + assert restored_acquisition.exploration_decay_delay == 3 + assert restored.suggest() == original.suggest()