def test_ask_tell_optimizer_uses_specified_acquisition_state(
search_space: Box,
init_dataset: Dataset,
model: TrainableProbabilisticModel,
starting_state: int | None,
expected_state: int,
) -> 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,
models: Mapping[str, ProbabilisticModel],
datasets: Optional[Mapping[str, Dataset]] = None,
) -> 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()
ask_tell = AskTellOptimizer(search_space, init_dataset, model, rule, starting_state)
_ = ask_tell.ask()
state_record: Record[State[int, TensorType]] = ask_tell.to_record()
# mypy cannot see that this is in fact int
assert state_record.acquisition_state == expected_state # type: ignore
def test_ask_tell_optimizer_with_default_acquisition_suggests_new_point(
search_space: Box,
init_dataset: Dataset,
model: TrainableProbabilisticModel,
) -> None:
ask_tell = AskTellOptimizer(search_space, init_dataset, model)
new_point = ask_tell.ask()
assert len(new_point) == 1
def test_ask_tell_optimizer_suggests_new_point(
search_space: Box,
init_dataset: Dataset,
model: TrainableProbabilisticModel,
acquisition_rule: AcquisitionRule[TensorType, Box],
) -> None:
ask_tell = AskTellOptimizer(search_space, init_dataset, model, acquisition_rule)
new_point = ask_tell.ask()
assert len(new_point) == 1
def test_ask_tell_optimizer_does_not_accept_empty_datasets_or_models(
search_space: Box,
init_dataset: Dataset,
model: TrainableProbabilisticModel,
acquisition_rule: AcquisitionRule[TensorType, Box],
) -> None:
with pytest.raises(ValueError):
AskTellOptimizer(search_space, {}, model, acquisition_rule) # type: ignore
with pytest.raises(ValueError):
AskTellOptimizer(search_space, init_dataset, {}, acquisition_rule) # type: ignore
def test_ask_tell_optimizer_returns_optimization_result(
search_space: Box,
init_dataset: Dataset,
model: TrainableProbabilisticModel,
acquisition_rule: AcquisitionRule[TensorType, Box],
) -> None:
ask_tell = AskTellOptimizer(search_space, init_dataset, model, acquisition_rule)
result: OptimizationResult[None] = ask_tell.to_result()
assert_datasets_allclose(result.try_get_final_dataset(), init_dataset)
assert isinstance(result.try_get_final_model(), type(model))
def test_ask_tell_optimizer_returns_complete_state(
search_space: Box,
init_dataset: Dataset,
model: TrainableProbabilisticModel,
acquisition_rule: AcquisitionRule[TensorType, Box],
) -> None:
ask_tell = AskTellOptimizer(search_space, init_dataset, model, acquisition_rule)
state_record: Record[None] = ask_tell.to_record()
assert_datasets_allclose(state_record.dataset, init_dataset)
assert isinstance(state_record.model, type(model))
assert state_record.acquisition_state is None
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_ask_tell_optimizer_optimizes_initial_model(
search_space: Box,
init_dataset: Dataset,
model: TrainableProbabilisticModel,
acquisition_rule: AcquisitionRule[TensorType, Box],
fit_initial_model: bool,
) -> None:
ask_tell = AskTellOptimizer(
search_space, init_dataset, model, acquisition_rule, fit_model=fit_initial_model
)
state_record: Record[None] = ask_tell.to_record()
if fit_initial_model:
assert state_record.model.optimize_count == 1 # type: ignore
else:
assert state_record.model.optimize_count == 0 # type: ignore
def test_ask_tell_optimizer_default_acquisition_requires_objective_tag(
search_space: Box,
init_dataset: Dataset,
model: TrainableProbabilisticModel,
) -> None:
wrong_tag = OBJECTIVE + "_WRONG"
wrong_datasets = {wrong_tag: init_dataset}
wrong_models = {wrong_tag: model}
with pytest.raises(ValueError):
AskTellOptimizer(search_space, wrong_datasets, wrong_models)
def test_ask_tell_optimizer_from_state_does_not_train_model(
search_space: Box,
init_dataset: Dataset,
model: TrainableProbabilisticModel,
acquisition_rule: AcquisitionRule[TensorType, Box],
) -> None:
old_state: Record[None] = Record({OBJECTIVE: init_dataset}, {OBJECTIVE: model}, None)
ask_tell = AskTellOptimizer.from_record(old_state, search_space, acquisition_rule)
state_record: Record[None] = ask_tell.to_record()
assert state_record.model.optimize_count == 0 # type: ignore
def test_ask_tell_optimizer_loads_from_state(
search_space: Box,
init_dataset: Dataset,
model: TrainableProbabilisticModel,
acquisition_rule: AcquisitionRule[TensorType, Box],
) -> None:
old_state: Record[None] = Record({OBJECTIVE: init_dataset}, {OBJECTIVE: model}, None)
ask_tell = AskTellOptimizer.from_record(old_state, search_space, acquisition_rule)
new_state: Record[None] = ask_tell.to_record()
assert_datasets_allclose(old_state.dataset, new_state.dataset)
assert isinstance(new_state.model, type(old_state.model))
# Set this flag to False to disable sleep delays in case you want the notebook to execute quickly
enable_sleep_delays = True
# %% [markdown]
# Now we are ready to define the optimizer. Notice how we set the acquisition function to be `BatchMonteCarloExpectedImprovement`. It is also the default function used by the `AsynchronousOptimization` rule, but here we specify it explicitly for clarity. We also set the batch size.
# %%
from trieste.acquisition.rule import AsynchronousOptimization
from trieste.acquisition.function import BatchMonteCarloExpectedImprovement
from trieste.ask_tell_optimization import AskTellOptimizer
model = build_model(initial_data)
acquisition_function = BatchMonteCarloExpectedImprovement(sample_size=10000)
async_rule = AsynchronousOptimization(
acquisition_function, num_query_points=batch_size) # type: ignore
async_bo = AskTellOptimizer(search_space, initial_data, model, async_rule)
# %% [markdown]
# Initialize Ray. This line will output the dashboard URL, which you can open in a separate tab to watch workers doing observations.
# %%
ray.init(ignore_reinit_error=True)
# %% [markdown]
# Here is the main optimization loop. First we ask for several batches of points to make sure all allocated workers are busy. Then we keep waiting for the workers to complete their tasks. Whenever `batch_size` of tasks came back, we tell Trieste new observations and ask for another batch of points.
# %%
points_observed = 0
workers = []
x_sta, x_mean, x_std = normalise(initial_data.query_points)
y_sta, y_mean, y_std = normalise(initial_data.observations)
normalised_data = Dataset(query_points=x_sta, observations=y_sta)
dataset = initial_data
for step in range(num_acquisitions):
if step == 0:
model = build_gp_model(normalised_data)
model.optimize(normalised_data)
else:
model.update(normalised_data)
model.optimize(normalised_data)
# Asking for a new point to observe
ask_tell = AskTellOptimizer(search_space, normalised_data, model)
query_point = ask_tell.ask()
# Transforming the query point back to the non-normalised space
query_point = x_std * query_point + x_mean
# Evaluating the function at the new query point
new_data_point = observer(query_point)
dataset = dataset + new_data_point
# Normalize the dataset with the new query point and observation
x_sta, _, _ = normalise(dataset.query_points, x_mean, x_std)
y_sta, _, _ = normalise(dataset.observations, y_mean, y_std)
normalised_data = Dataset(query_points=x_sta, observations=y_sta)
# %% [markdown]
num_initial_points = 5
initial_query_points = search_space.sample(num_initial_points)
observer = mk_observer(scaled_branin)
initial_data = observer(initial_query_points)
# %% [markdown]
# ## Timing acquisition function: simple use case for Ask-Tell
#
# Let's say we are very concerned with the performance of the acqusition function, and want a simple way of measuring its performance over the course of the optimization. At the time of writing these lines, regular Trieste's optimizer does not provide such customization functionality, and this is where Ask-Tell comes in handy.
# %%
import timeit
model = build_model(initial_data)
ask_tell = AskTellOptimizer(search_space, initial_data, model)
for step in range(n_steps):
start = timeit.default_timer()
new_point = ask_tell.ask()
stop = timeit.default_timer()
print(f"Time at step {step + 1}: {stop - start}")
new_data = observer(new_point)
ask_tell.tell(new_data)
# %% [markdown]
# Once ask-tell optimization is over, you can extract an optimization result object and perform whatever analysis you need, just like with regular Trieste optimization interface. For instance, here we will plot regret for each optimization step.
def test_ask_tell_optimization_finds_minima_of_the_scaled_branin_function(
num_steps: int,
reload_state: bool,
acquisition_rule_fn: Callable[[], AcquisitionRule[TensorType, SearchSpace]]
| Callable[[], AcquisitionRule[State[TensorType, AsynchronousGreedy.State
| TrustRegion.State], Box], ],
) -> None:
# For the case when optimization state is saved and reload on each iteration
# we need to use new acquisition function object to imitate real life usage
# hence acquisition rule factory method is passed in, instead of a rule object itself
# it is then called to create a new rule whenever needed in the test
search_space = BRANIN_SEARCH_SPACE
def build_model(data: Dataset) -> GaussianProcessRegression:
variance = tf.math.reduce_variance(data.observations)
kernel = gpflow.kernels.Matern52(variance,
tf.constant([0.2, 0.2], tf.float64))
scale = tf.constant(1.0, dtype=tf.float64)
kernel.variance.prior = tfp.distributions.LogNormal(
tf.constant(-2.0, dtype=tf.float64), scale)
kernel.lengthscales.prior = tfp.distributions.LogNormal(
tf.math.log(kernel.lengthscales), scale)
gpr = gpflow.models.GPR((data.query_points, data.observations),
kernel,
noise_variance=1e-5)
gpflow.utilities.set_trainable(gpr.likelihood, False)
return GaussianProcessRegression(gpr)
initial_query_points = search_space.sample(5)
observer = mk_observer(scaled_branin)
initial_data = observer(initial_query_points)
model = build_model(initial_data)
ask_tell = AskTellOptimizer(search_space, initial_data, model,
acquisition_rule_fn())
for _ in range(num_steps):
# two scenarios are tested here, depending on `reload_state` parameter
# in first the same optimizer object is always used
# in second new optimizer is created at each step from saved state
new_point = ask_tell.ask()
if reload_state:
state: Record[None
| State[TensorType, AsynchronousGreedy.State
| TrustRegion.State]] = ask_tell.to_record()
written_state = pickle.dumps(state)
new_data_point = observer(new_point)
if reload_state:
state = pickle.loads(written_state)
ask_tell = AskTellOptimizer.from_record(state, search_space,
acquisition_rule_fn())
ask_tell.tell(new_data_point)
result: OptimizationResult[None | State[
TensorType,
AsynchronousGreedy.State | TrustRegion.State]] = ask_tell.to_result()
dataset = result.try_get_final_dataset()
arg_min_idx = tf.squeeze(tf.argmin(dataset.observations, axis=0))
best_y = dataset.observations[arg_min_idx]
best_x = dataset.query_points[arg_min_idx]
relative_minimizer_err = tf.abs(
(best_x - BRANIN_MINIMIZERS) / BRANIN_MINIMIZERS)
# these accuracies are the current best for the given number of optimization steps, which makes
# this is a regression test
assert tf.reduce_any(tf.reduce_all(relative_minimizer_err < 0.05, axis=-1),
axis=0)
npt.assert_allclose(best_y, SCALED_BRANIN_MINIMUM, rtol=0.005)
