def test_single_builder_using_passes_on_correct_dataset_and_model(
single_builder: Union[SingleModelAcquisitionBuilder, SingleModelBatchAcquisitionBuilder]
) -> None:
builder = single_builder.using("foo")
data = {"foo": mk_dataset([[0.0]], [[0.0]]), "bar": mk_dataset([[1.0]], [[1.0]])}
models = {"foo": QuadraticMeanAndRBFKernel(0.0), "bar": QuadraticMeanAndRBFKernel(1.0)}
builder.prepare_acquisition_function(data, models)
def test_bayesian_optimizer_optimize_tracked_state() -> None:
class _CountingRule(AcquisitionRule[State[Optional[int], TensorType],
Box]):
def acquire(
self,
search_space: Box,
datasets: Mapping[str, Dataset],
models: Mapping[str, ProbabilisticModel],
) -> State[int | None, TensorType]:
def go(state: int | None) -> tuple[int | None, TensorType]:
new_state = 0 if state is None else state + 1
return new_state, tf.constant([[10.0]], tf.float64) + new_state
return go
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
initial_data = mk_dataset([[0.0]], [[0.0]])
model = _DecreasingVarianceModel(initial_data)
_, history = (BayesianOptimizer(_quadratic_observer,
Box([0],
[1])).optimize(3, {
"": initial_data
}, {
"": model
}, _CountingRule()).astuple())
assert [record.acquisition_state for record in history] == [None, 0, 1]
assert_datasets_allclose(history[0].datasets[""], initial_data)
assert_datasets_allclose(history[1].datasets[""],
mk_dataset([[0.0], [10.0]], [[0.0], [100.0]]))
assert_datasets_allclose(
history[2].datasets[""],
mk_dataset([[0.0], [10.0], [11.0]], [[0.0], [100.0], [121.0]]))
for step in range(3):
assert history[step].model == history[step].models[""]
assert history[step].dataset == history[step].datasets[""]
_, variance_from_saved_model = (history[step].models[""].predict(
tf.constant([[0.0]], tf.float64)))
npt.assert_allclose(variance_from_saved_model, 1.0 / (step + 1))
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([[0.0]])
result, history = (BayesianOptimizer(_quadratic_observer, Box(
[0], [1])).optimize(10, {
"": mk_dataset([[0.0]], [[0.0]])
}, {
"": _UncopyableModel()
}, rule).astuple())
with pytest.raises(_Whoops):
result.unwrap()
assert len(history) == 4
def test_bayesian_optimizer_optimizes_initial_model(
fit_initial_model: bool) -> None:
class _CountingOptimizerModel(_PseudoTrainableQuadratic):
_optimize_count = 0
def optimize(self, dataset: Dataset) -> None:
self._optimize_count += 1
rule = FixedAcquisitionRule([[0.0]])
model = _CountingOptimizerModel()
final_opt_state, _ = (BayesianOptimizer(_quadratic_observer, Box(
[0], [1])).optimize(
1,
{
"": mk_dataset([[0.0]], [[0.0]])
},
{
"": model
},
rule,
fit_initial_model=fit_initial_model,
).astuple())
final_model = final_opt_state.unwrap().model
if fit_initial_model: # optimized at start and end of first BO step
assert final_model._optimize_count == 2 # type: ignore
else: # optimized just at end of first BO step
assert final_model._optimize_count == 1 # type: ignore
def test_batch_monte_carlo_expected_improvement_raises_for_model_with_wrong_event_shape(
) -> None:
builder = BatchMonteCarloExpectedImprovement(100)
data = mk_dataset([[0.0, 0.0]], [[0.0, 0.0]])
model = _dim_two_gp()
with pytest.raises(TF_DEBUGGING_ERROR_TYPES):
builder.prepare_acquisition_function(data, model)
def test_batch_monte_carlo_expected_improvement_raises_for_model_with_wrong_event_shape() -> None:
builder = BatchMonteCarloExpectedImprovement(100)
data = mk_dataset([(0.0, 0.0)], [(0.0, 0.0)])
matern52 = tfp.math.psd_kernels.MaternFiveHalves(
amplitude=tf.cast(2.3, tf.float64), length_scale=tf.cast(0.5, tf.float64)
)
model = GaussianProcess([lambda x: branin(x), lambda x: quadratic(x)], [matern52, rbf()])
with pytest.raises(TF_DEBUGGING_ERROR_TYPES):
builder.prepare_acquisition_function(data, model)
def test_bayesian_optimizer_optimize_for_failed_step(
observer: Observer, model: TrainableProbabilisticModel,
rule: AcquisitionRule) -> None:
optimizer = BayesianOptimizer(observer, Box([0], [1]))
data, models = {"": mk_dataset([[0.0]], [[0.0]])}, {"": model}
result, history = optimizer.optimize(3, data, models, rule).astuple()
with pytest.raises(_Whoops):
result.unwrap()
assert len(history) == 1
def test_ask_tell_optimizer_updates_state_with_new_data(
search_space: Box,
init_dataset: Dataset,
model: TrainableProbabilisticModel,
acquisition_rule: AcquisitionRule[TensorType, Box],
) -> None:
new_data = mk_dataset([[1.0]], [[1.0]])
ask_tell = AskTellOptimizer(search_space, init_dataset, model, acquisition_rule)
ask_tell.tell(new_data)
state_record: Record[None] = ask_tell.to_record()
assert_datasets_allclose(state_record.dataset, init_dataset + new_data)
def test_ask_tell_optimizer_tell_validates_keys(
search_space: Box,
init_dataset: Dataset,
model: TrainableProbabilisticModel,
acquisition_rule: AcquisitionRule[TensorType, Box],
) -> None:
dataset_with_key_1 = {"1": init_dataset}
model_with_key_1 = {"1": model}
new_data_with_key_2 = {"2": mk_dataset([[1.0]], [[1.0]])}
ask_tell = AskTellOptimizer(search_space, dataset_with_key_1,
model_with_key_1, acquisition_rule)
with pytest.raises(ValueError):
ask_tell.tell(new_data_with_key_2)
def test_ask_tell_optimizer_trains_model(
search_space: Box,
init_dataset: Dataset,
model: TrainableProbabilisticModel,
acquisition_rule: AcquisitionRule[TensorType, Box],
) -> None:
new_data = mk_dataset([[1.0]], [[1.0]])
ask_tell = AskTellOptimizer(
search_space, init_dataset, model, acquisition_rule, fit_model=False
)
ask_tell.tell(new_data)
state_record: Record[None] = ask_tell.to_record()
assert state_record.model.optimize_count == 1 # type: ignore
def test_bayesian_optimizer_calls_observer_once_per_iteration(steps: int) -> None:
class _CountingObserver:
call_count = 0
def __call__(self, x: tf.Tensor) -> Dataset:
self.call_count += 1
return Dataset(x, tf.reduce_sum(x ** 2, axis=-1, keepdims=True))
observer = _CountingObserver()
optimizer = BayesianOptimizer(observer, Box([-1], [1]))
data = mk_dataset([[0.5]], [[0.25]])
optimizer.optimize(steps, data, _PseudoTrainableQuadratic()).final_result.unwrap()
assert observer.call_count == steps
def test_batch_monte_carlo_expected_improvement() -> None:
xs = tf.random.uniform([3, 5, 7, 2], dtype=tf.float64)
model = QuadraticMeanAndRBFKernel()
mean, cov = model.predict_joint(xs)
mvn = tfp.distributions.MultivariateNormalFullCovariance(tf.linalg.matrix_transpose(mean), cov)
mvn_samples = mvn.sample(10_000)
min_predictive_mean_at_known_points = 0.09
# fmt: off
expected = tf.reduce_mean(tf.reduce_max(tf.maximum(
min_predictive_mean_at_known_points - mvn_samples, 0.0
), axis=-1), axis=0)
# fmt: on
builder = BatchMonteCarloExpectedImprovement(10_000)
acq = builder.prepare_acquisition_function(mk_dataset([[0.3], [0.5]], [[0.09], [0.25]]), model)
npt.assert_allclose(acq(xs), expected, rtol=0.05)
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 = None,
) -> NoReturn:
assert False
def _unusable_observer(x: tf.Tensor) -> NoReturn:
assert False
data = {"": mk_dataset([[0.0]], [[0.0]])}
result, history = (BayesianOptimizer(_unusable_observer,
Box([-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: int | None,
expected_states_received: list[int | None],
final_acquisition_state: int | None,
) -> None:
class Rule(AcquisitionRule[State[Optional[int], TensorType], Box]):
def __init__(self) -> None:
self.states_received: list[int | None] = []
def acquire(
self,
search_space: Box,
datasets: Mapping[str, Dataset],
models: Mapping[str, ProbabilisticModel],
) -> State[int | None, TensorType]:
def go(state: int | None) -> tuple[int | None, TensorType]:
self.states_received.append(state)
if state is None:
state = 0
return state + 1, tf.constant([[0.0]], tf.float64)
return go
rule = Rule()
data, models = {
"": mk_dataset([[0.0]], [[0.0]])
}, {
"": _PseudoTrainableQuadratic()
}
final_state, history = (BayesianOptimizer(lambda x: {
"": Dataset(x, x**2)
}, Box([-1], [1])).optimize(3, data, models, 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 init_dataset() -> Dataset:
return mk_dataset([[0.0]], [[0.0]])
