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
# 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 = []
# a helper function to launch a worker for a numpy array representing a single point
def launch_worker(x):
worker = ray_objective.remote(np.atleast_2d(x), enable_sleep_delays)
workers.append(worker)
# get first couple of batches of points and init all workers
for _ in range(int(num_workers / batch_size)):
points = async_bo.ask().numpy()
np.apply_along_axis(launch_worker, axis=1, arr=points)
finished_workers = []
while points_observed < num_observations:
ready_workers, remaining_workers = ray.wait(workers, timeout=0)
finished_workers += ready_workers
workers = remaining_workers
if len(finished_workers) < batch_size:
continue
# we saw enough results to ask for a new batch
new_observations = [
observation for worker in finished_workers
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]
#
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 plot_ask_tell_regret(ask_tell_result):
observations = ask_tell_result.try_get_final_dataset().observations.numpy()
arg_min_idx = tf.squeeze(tf.argmin(observations, axis=0))
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)
