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_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_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
#
# Now let's suppose you are optimizing a process that takes hours or even days to complete, e.g. a lab experiment or a hyperparameter optimization of a big machine learning model. This time you cannot even express the objective function in Python code. Instead you would like to ask Trieste what configuration to run next, go to the lab, perform the experiment, collect data, feed it back to Trieste and ask for the next configuration, and so on. It would be very convenient to be able to store intermediate optimization state to disk or database or other storage, so that your machine can be switched off while you are waiting for observation results.
#
# In this section we'll show how you could do it with Ask-Tell in Trieste. Of course we cannot perform a real physical experiment within this notebook, so we will just mimick it by using pickle to write optimization state and read it back.
# %%
import pickle
model = build_model(initial_data)
ask_tell = AskTellOptimizer(search_space, initial_data, model)
for step in range(n_steps):
print(f"Ask Trieste for configuration #{step}")
new_config = ask_tell.ask()
print("Saving Trieste state to re-use later")
state: Record[None] = ask_tell.to_record()
saved_state = pickle.dumps(state)
print(f"In the lab running the experiment #{step}.")
new_datapoint = scaled_branin(new_config)
print("Back from the lab")
print("Restore optimizer from the saved state")
loaded_state = pickle.loads(saved_state)
ask_tell = AskTellOptimizer.from_record(loaded_state, search_space)
ask_tell.tell(Dataset(new_config, new_datapoint))
plot_ask_tell_regret(ask_tell.to_result())
# %% [markdown]
# A word of warning. This serialization technique is not guaranteed to work smoothly with every Tensorflow-based model, so apply to your own problems with caution.
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)