def test_bayesian_optimizer_optimize_tracked_state() -> None:
class _CountingRule(AcquisitionRule[int, Box]):
def acquire(
self,
search_space: Box,
datasets: Mapping[str, Dataset],
models: Mapping[str, ProbabilisticModel],
state: Optional[int],
) -> Tuple[TensorType, int]:
new_state = 0 if state is None else state + 1
return tf.constant([[10.0]]) + new_state, new_state
class _DecreasingVarianceModel(QuadraticMeanAndRBFKernel,
TrainableProbabilisticModel):
def __init__(self, data: Dataset):
super().__init__()
self._data = data
def predict(self,
query_points: TensorType) -> Tuple[TensorType, TensorType]:
mean, var = super().predict(query_points)
return mean, var / len(self._data)
def update(self, dataset: Dataset) -> None:
self._data = dataset
def optimize(self, dataset: Dataset) -> None:
pass
_, history = (BayesianOptimizer(
_quadratic_observer, one_dimensional_range(0, 1)).optimize(
3, {
"": zero_dataset()
}, {
"": _DecreasingVarianceModel(zero_dataset())
}, _CountingRule()).astuple())
assert [record.acquisition_state for record in history] == [None, 0, 1]
assert_datasets_allclose(history[0].datasets[""], zero_dataset())
assert_datasets_allclose(
history[1].datasets[""],
Dataset(tf.constant([[0.0], [10.0]]), tf.constant([[0.0], [100.0]])),
)
assert_datasets_allclose(
history[2].datasets[""],
Dataset(tf.constant([[0.0], [10.0], [11.0]]),
tf.constant([[0.0], [100.0], [121.0]])),
)
for step in range(3):
_, variance_from_saved_model = history[step].models[""].predict(
tf.constant([[0.0]]))
npt.assert_allclose(variance_from_saved_model, 1.0 / (step + 1))
def test_single_builder_using_passes_on_correct_dataset_and_model() -> None:
class _Mock(SingleModelAcquisitionBuilder):
def prepare_acquisition_function(
self, dataset: Dataset, model: ProbabilisticModel
) -> AcquisitionFunction:
assert dataset is data["foo"]
assert model is models["foo"]
return lambda at: at
builder = _Mock().using("foo")
data = {"foo": zero_dataset(), "bar": zero_dataset()}
models = {"foo": QuadraticWithUnitVariance(), "bar": QuadraticWithUnitVariance()}
builder.prepare_acquisition_function(data, models)
def test_bayesian_optimizer_optimize_raises_for_invalid_rule_keys_and_default_acquisition(
) -> None:
optimizer = BayesianOptimizer(lambda x: x[:1],
one_dimensional_range(-1, 1))
with pytest.raises(ValueError):
optimizer.optimize(3, {"foo": zero_dataset()},
{"foo": _PseudoTrainableQuadratic()})
def test_single_builder_raises_immediately_for_wrong_key() -> None:
builder = _IdentitySingleBuilder().using("foo")
with pytest.raises(KeyError):
builder.prepare_acquisition_function(
{"bar": zero_dataset()}, {"bar": QuadraticWithUnitVariance()}
)
def test_bayesian_optimizer_optimize_for_uncopyable_model() -> None:
class _UncopyableModel(_PseudoTrainableQuadratic):
_optimize_count = 0
def optimize(self, dataset: Dataset) -> None:
self._optimize_count += 1
def __deepcopy__(self, memo: Dict[int, object]) -> _UncopyableModel:
if self._optimize_count >= 3:
raise _Whoops
return self
rule = FixedAcquisitionRule(tf.constant([[0.0]]))
result, history = (BayesianOptimizer(_quadratic_observer,
one_dimensional_range(0, 1)).optimize(
10, {
"": zero_dataset()
}, {
"": _UncopyableModel()
}, rule).astuple())
with pytest.raises(_Whoops):
result.unwrap()
assert len(history) == 4
def test_probability_of_feasibility_builder_builds_pof(threshold: float,
at: tf.Tensor) -> None:
builder = ProbabilityOfFeasibility(threshold)
acq = builder.prepare_acquisition_function(zero_dataset(),
QuadraticWithUnitVariance())
expected = probability_of_feasibility(QuadraticWithUnitVariance(),
threshold, at)
npt.assert_allclose(acq(at), expected)
def test_bayesian_optimizer_optimize_raises_for_negative_steps(
num_steps: int) -> None:
optimizer = BayesianOptimizer(_quadratic_observer,
one_dimensional_range(-1, 1))
with pytest.raises(ValueError, match="num_steps"):
optimizer.optimize(num_steps, {"": zero_dataset()},
{"": _PseudoTrainableQuadratic()})
def test_single_builder_raises_immediately_for_wrong_key(
single_builder: Union[SingleModelAcquisitionBuilder, SingleModelBatchAcquisitionBuilder]
) -> None:
builder = single_builder.using("foo")
with pytest.raises(KeyError):
builder.prepare_acquisition_function(
{"bar": zero_dataset()}, {"bar": QuadraticMeanAndRBFKernel()}
)
def test_bayesian_optimizer_optimize_returns_default_acquisition_state_of_correct_type(
) -> None:
optimizer = BayesianOptimizer(lambda x: {OBJECTIVE: Dataset(x, x[:1])},
one_dimensional_range(-1, 1))
res: OptimizationResult[None] = optimizer.optimize(
3, {OBJECTIVE: zero_dataset()},
{OBJECTIVE: _PseudoTrainableQuadratic()})
if res.error is not None:
raise res.error
assert all(logging_state.acquisition_state is None
for logging_state in res.history)
def test_bayesian_optimizer_optimize_for_failed_step(
observer: Observer, model: TrainableProbabilisticModel,
rule: AcquisitionRule) -> None:
optimizer = BayesianOptimizer(observer, one_dimensional_range(0, 1))
result, history = optimizer.optimize(3, {
"": zero_dataset()
}, {
"": model
}, rule).astuple()
with pytest.raises(_Whoops):
result.unwrap()
assert len(history) == 1
def test_bayesian_optimizer_optimize_doesnt_track_state_if_told_not_to(
) -> None:
class _UncopyableModel(_PseudoTrainableQuadratic):
def __deepcopy__(self, memo: Dict[int, object]) -> NoReturn:
assert False
history = (BayesianOptimizer(_quadratic_observer,
one_dimensional_range(-1, 1)).optimize(
5, {
OBJECTIVE: zero_dataset()
}, {
OBJECTIVE: _UncopyableModel()
},
track_state=False).history)
assert len(history) == 0
def test_bayesian_optimizer_optimize_is_noop_for_zero_steps() -> None:
class _UnusableModel(TrainableProbabilisticModel):
def predict(self, query_points: TensorType) -> NoReturn:
assert False
def predict_joint(self, query_points: TensorType) -> NoReturn:
assert False
def sample(self, query_points: TensorType,
num_samples: int) -> NoReturn:
assert False
def update(self, dataset: Dataset) -> NoReturn:
assert False
def optimize(self, dataset: Dataset) -> NoReturn:
assert False
class _UnusableRule(AcquisitionRule[None, Box]):
def acquire(
self,
search_space: Box,
datasets: Mapping[str, Dataset],
models: Mapping[str, ProbabilisticModel],
state: None,
) -> NoReturn:
assert False
def _unusable_observer(x: tf.Tensor) -> NoReturn:
assert False
data = {"": zero_dataset()}
result, history = (BayesianOptimizer(
_unusable_observer,
one_dimensional_range(-1, 1)).optimize(0, data, {
"": _UnusableModel()
}, _UnusableRule()).astuple())
assert history == []
final_data = result.unwrap().datasets
assert len(final_data) == 1
assert_datasets_allclose(final_data[""], data[""])
def test_bayesian_optimizer_uses_specified_acquisition_state(
starting_state: Optional[int],
expected_states_received: List[Optional[int]],
final_acquisition_state: Optional[int],
) -> None:
class Rule(AcquisitionRule[int, Box]):
def __init__(self):
self.states_received = []
def acquire(
self,
search_space: Box,
datasets: Mapping[str, Dataset],
models: Mapping[str, ProbabilisticModel],
state: Optional[int],
) -> Tuple[TensorType, int]:
self.states_received.append(state)
if state is None:
state = 0
return tf.constant([[0.0]]), state + 1
rule = Rule()
final_state, history = (BayesianOptimizer(lambda x: {
"": Dataset(x, x**2)
}, one_dimensional_range(-1, 1)).optimize(3, {
"": zero_dataset()
}, {
"": _PseudoTrainableQuadratic()
}, rule, starting_state).astuple())
assert rule.states_received == expected_states_received
assert final_state.unwrap().acquisition_state == final_acquisition_state
assert [record.acquisition_state
for record in history] == expected_states_received
def test_bayesian_optimizer_uses_specified_acquisition_state(
starting_state: Optional[int],
expected_states: List[Optional[int]]) -> None:
class Rule(AcquisitionRule[int, Box]):
def __init__(self):
self.states_received = []
def acquire(
self,
search_space: Box,
datasets: Mapping[str, Dataset],
models: Mapping[str, ProbabilisticModel],
state: Optional[int],
) -> Tuple[QueryPoints, int]:
self.states_received.append(state)
if state is None:
state = 0
return tf.constant([[0.0]]), state + 1
rule = Rule()
res = BayesianOptimizer(lambda x: {
"": Dataset(x, x**2)
}, one_dimensional_range(-1,
1)).optimize(3, {"": zero_dataset()},
{"": _PseudoTrainableQuadratic()},
rule, starting_state)
if res.error is not None:
raise res.error
assert rule.states_received == expected_states
assert [state.acquisition_state
for state in res.history] == expected_states
@pytest.mark.parametrize(
"models",
[
{},
{
"foo": QuadraticMeanAndRBFKernel()
},
{
"foo": QuadraticMeanAndRBFKernel(),
OBJECTIVE: QuadraticMeanAndRBFKernel()
},
],
)
@pytest.mark.parametrize("datasets", [{}, {OBJECTIVE: zero_dataset()}])
def test_thompson_sampling_raises_for_invalid_models_keys(
datasets: dict[str, Dataset],
models: dict[str, ProbabilisticModel]) -> None:
search_space = one_dimensional_range(-1, 1)
rule = ThompsonSampling(100, 10)
with pytest.raises(ValueError):
rule.acquire(search_space, datasets, models)
@pytest.mark.parametrize("datasets", [{}, {"foo": zero_dataset()}])
@pytest.mark.parametrize("models", [{}, {
"foo": QuadraticMeanAndRBFKernel()
}, {
OBJECTIVE: QuadraticMeanAndRBFKernel()
}])
def test_optimization_result_try_get_final_models_for_successful_optimization(
) -> None:
models = {"foo": _PseudoTrainableQuadratic()}
result: OptimizationResult[None] = OptimizationResult(
Ok(Record({"foo": zero_dataset()}, models, None)), [])
assert result.try_get_final_models() is models
optimizer.optimize(steps, {
OBJECTIVE: data
}, {
OBJECTIVE: _PseudoTrainableQuadratic()
}).final_result.unwrap()
assert observer.call_count == steps
@pytest.mark.parametrize(
"datasets, models",
[
({}, {}),
({
"foo": zero_dataset()
}, {}),
({
"foo": zero_dataset()
}, {
"bar": _PseudoTrainableQuadratic()
}),
(
{
"foo": zero_dataset()
},
{
"foo": _PseudoTrainableQuadratic(),
"bar": _PseudoTrainableQuadratic()
},
),
