def test_gumbel_sampler_returns_correctly_shaped_samples(
sample_size: int) -> None:
search_space = Box([0, 0], [1, 1])
gumbel_sampler = GumbelSampler(sample_size, QuadraticMeanAndRBFKernel())
query_points = search_space.sample(5)
gumbel_samples = gumbel_sampler.sample(query_points)
tf.debugging.assert_shapes([(gumbel_samples, [sample_size, 1])])
def test_rff_sampler_returns_correctly_shaped_samples(
sample_min_value: bool, sample_size: int
) -> None:
search_space = Box([0.0, 0.0], [1.0, 1.0])
model = QuadraticMeanAndRBFKernel(noise_variance=tf.constant(1.0, dtype=tf.float64))
model.kernel = (
gpflow.kernels.RBF()
) # need a gpflow kernel object for random feature decompositions
x_range = tf.linspace(0.0, 1.0, 5)
x_range = tf.cast(x_range, dtype=tf.float64)
xs = tf.reshape(tf.stack(tf.meshgrid(x_range, x_range, indexing="ij"), axis=-1), (-1, 2))
ys = quadratic(xs)
dataset = Dataset(xs, ys)
sampler = RandomFourierFeatureThompsonSampler(
sample_size, model, dataset, num_features=100, sample_min_value=sample_min_value
)
query_points = search_space.sample(100)
thompson_samples = sampler.sample(query_points)
if sample_min_value:
tf.debugging.assert_shapes([(thompson_samples, [sample_size, 1])])
else:
tf.debugging.assert_shapes([(thompson_samples, [sample_size, 2])])
def test_optimizer_finds_minima_of_the_branin_function(
num_steps: int, acquisition_rule: AcquisitionRule
) -> None:
search_space = Box([0, 0], [1, 1])
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))
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(branin, OBJECTIVE)
initial_data = observer(initial_query_points)
model = build_model(initial_data[OBJECTIVE])
dataset = (
BayesianOptimizer(observer, search_space)
.optimize(num_steps, initial_data, {OBJECTIVE: model}, acquisition_rule)
.try_get_final_datasets()[OBJECTIVE]
)
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.03, axis=-1), axis=0)
npt.assert_allclose(best_y, BRANIN_MINIMUM, rtol=0.03)
def test_discrete_thompson_sampler_returns_correctly_shaped_samples(
sample_size: int) -> None:
search_space = Box([0, 0], [1, 1])
thompson_sampler = DiscreteThompsonSampler(sample_size,
QuadraticMeanAndRBFKernel())
query_points = search_space.sample(100)
thompson_samples = thompson_sampler.sample(query_points)
tf.debugging.assert_shapes([(thompson_samples, ["N", 2])])
def test_gumbel_samples_are_minima() -> None:
dataset = Dataset(tf.zeros([3, 2], dtype=tf.float64),
tf.ones([3, 2], dtype=tf.float64))
search_space = Box([0, 0], [1, 1])
model = QuadraticMeanAndRBFKernel()
gumbel_sampler = GumbelSampler(5, model)
query_points = search_space.sample(100)
query_points = tf.concat([dataset.query_points, query_points], 0)
gumbel_samples = gumbel_sampler.sample(query_points)
fmean, _ = model.predict(dataset.query_points)
assert max(gumbel_samples) < min(fmean)
def test_exact_thompson_sampler_returns_correctly_shaped_samples(
sample_min_value: bool, sample_size: int
) -> None:
search_space = Box([0, 0], [1, 1])
thompson_sampler = ExactThompsonSampler(
sample_size, QuadraticMeanAndRBFKernel(), sample_min_value=sample_min_value
)
query_points = search_space.sample(500)
thompson_samples = thompson_sampler.sample(query_points)
if sample_min_value:
tf.debugging.assert_shapes([(thompson_samples, [sample_size, 1])])
else:
tf.debugging.assert_shapes([(thompson_samples, [sample_size, 2])])
def branin_dataset(num_query_points: int) -> Dataset:
"""
Generate example dataset based on Hartmann 6 objective function.
:param num_query_points: A number of samples from the objective function.
:return: A dataset.
"""
search_space = Box([0, 0], [1, 1])
query_points = search_space.sample(num_query_points)
observer = mk_observer(branin, OBJECTIVE)
data = observer(query_points)
return data[OBJECTIVE]
def test_two_layer_dgp_optimizer_finds_minima_of_michalewicz_function(
num_steps: int, acquisition_rule: AcquisitionRule[TensorType, SearchSpace], keras_float: None
) -> None:
# this unit test fails sometimes for
# normal search space used with MICHALEWICZ function
# so for stability we reduce its size here
search_space = Box(MICHALEWICZ_2_MINIMIZER[0] - 0.5, MICHALEWICZ_2_MINIMIZER[0] + 0.5)
def build_model(data: Dataset) -> DeepGaussianProcess:
epochs = int(2e3)
batch_size = 100
dgp = two_layer_dgp_model(data.query_points)
def scheduler(epoch: int, lr: float) -> float:
if epoch == epochs // 2:
return lr * 0.1
else:
return lr
optimizer = tf.optimizers.Adam(0.01)
fit_args = {
"batch_size": batch_size,
"epochs": epochs,
"verbose": 0,
"callbacks": tf.keras.callbacks.LearningRateScheduler(scheduler),
}
return DeepGaussianProcess(model=dgp, optimizer=optimizer, fit_args=fit_args)
initial_query_points = search_space.sample(50)
observer = mk_observer(michalewicz, OBJECTIVE)
initial_data = observer(initial_query_points)
model = build_model(initial_data[OBJECTIVE])
dataset = (
BayesianOptimizer(observer, search_space)
.optimize(num_steps, initial_data, {OBJECTIVE: model}, acquisition_rule, track_state=False)
.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 - MICHALEWICZ_2_MINIMIZER) / MICHALEWICZ_2_MINIMIZER)
assert tf.reduce_all(relative_minimizer_err < 0.03, axis=-1)
npt.assert_allclose(best_y, MICHALEWICZ_2_MINIMUM, rtol=0.03)
def test_gumbel_samples_are_minima() -> None:
search_space = Box([0, 0], [1, 1])
x_range = tf.linspace(0.0, 1.0, 5)
x_range = tf.cast(x_range, dtype=tf.float64)
xs = tf.reshape(tf.stack(tf.meshgrid(x_range, x_range, indexing="ij"), axis=-1), (-1, 2))
ys = quadratic(xs)
dataset = Dataset(xs, ys)
model = QuadraticMeanAndRBFKernel()
gumbel_sampler = GumbelSampler(5, model)
query_points = search_space.sample(100)
query_points = tf.concat([dataset.query_points, query_points], 0)
gumbel_samples = gumbel_sampler.sample(query_points)
fmean, _ = model.predict(dataset.query_points)
assert max(gumbel_samples) < min(fmean)
def test_multi_objective_optimizer_finds_pareto_front_of_the_VLMOP2_function(
num_steps: int, acquisition_rule: AcquisitionRule[TensorType, Box],
convergence_threshold: float) -> None:
search_space = Box([-2, -2], [2, 2])
def build_stacked_independent_objectives_model(
data: Dataset) -> ModelStack:
gprs = []
for idx in range(2):
single_obj_data = Dataset(
data.query_points, tf.gather(data.observations, [idx], axis=1))
variance = tf.math.reduce_variance(single_obj_data.observations)
kernel = gpflow.kernels.Matern52(
variance, tf.constant([0.2, 0.2], tf.float64))
gpr = gpflow.models.GPR(single_obj_data.astuple(),
kernel,
noise_variance=1e-5)
gpflow.utilities.set_trainable(gpr.likelihood, False)
gprs.append((GaussianProcessRegression(gpr), 1))
return ModelStack(*gprs)
observer = mk_observer(VLMOP2().objective(), OBJECTIVE)
initial_query_points = search_space.sample(10)
initial_data = observer(initial_query_points)
model = build_stacked_independent_objectives_model(initial_data[OBJECTIVE])
dataset = (BayesianOptimizer(observer, search_space).optimize(
num_steps, initial_data, {
OBJECTIVE: model
}, acquisition_rule).try_get_final_datasets()[OBJECTIVE])
# A small log hypervolume difference corresponds to a succesful optimization.
ref_point = get_reference_point(dataset.observations)
obs_hv = Pareto(dataset.observations).hypervolume_indicator(ref_point)
ideal_pf = tf.cast(VLMOP2().gen_pareto_optimal_points(100),
dtype=tf.float64)
ideal_hv = Pareto(ideal_pf).hypervolume_indicator(ref_point)
assert tf.math.log(ideal_hv - obs_hv) < convergence_threshold
def acquire(
self,
search_space: Box,
datasets: Mapping[str, Dataset],
models: Mapping[str, ProbabilisticModel],
previous_state: int | None = None,
) -> tuple[TensorType, int]:
if previous_state is None:
previous_state = 1
candidate_query_points = search_space.sample(previous_state)
linear_predictions, _ = models[LINEAR].predict(
candidate_query_points)
exponential_predictions, _ = models[EXPONENTIAL].predict(
candidate_query_points)
target = linear_predictions + exponential_predictions
optimum_idx = tf.argmin(target, axis=0)[0]
next_query_points = tf.expand_dims(
candidate_query_points[optimum_idx, ...], axis=0)
return next_query_points, previous_state * 2
def test_rff_thompson_samples_are_minima() -> None:
search_space = Box([0.0, 0.0], [1.0, 1.0])
model = QuadraticMeanAndRBFKernel(noise_variance=tf.constant(1e-5, dtype=tf.float64))
model.kernel = (
gpflow.kernels.RBF()
) # need a gpflow kernel object for random feature decompositions
x_range = tf.linspace(0.0, 1.0, 5)
x_range = tf.cast(x_range, dtype=tf.float64)
xs = tf.reshape(tf.stack(tf.meshgrid(x_range, x_range, indexing="ij"), axis=-1), (-1, 2))
ys = quadratic(xs)
dataset = Dataset(xs, ys)
sampler = RandomFourierFeatureThompsonSampler(
1, model, dataset, num_features=100, sample_min_value=True
)
query_points = search_space.sample(100)
query_points = tf.concat([dataset.query_points, query_points], 0)
thompson_samples = sampler.sample(query_points)
fmean, _ = model.predict(dataset.query_points)
assert max(thompson_samples) < min(fmean)
def test_discrete_thompson_sampler_returns_unique_samples() -> None:
search_space = Box([0, 0], [1, 1])
thompson_sampler = DiscreteThompsonSampler(20, QuadraticMeanAndRBFKernel())
query_points = search_space.sample(10)
thompson_samples = thompson_sampler.sample(query_points)
assert len(thompson_samples) < 10
# %% [markdown]
# ## Sample the observer over the search space
#
# Sometimes we don't have direct access to the objective function. We only have an observer that indirectly observes it. In _Trieste_, the observer outputs a number of datasets, each of which must be labelled so the optimization process knows which is which. In our case, we only have one dataset, the objective. We'll use _Trieste_'s default label for single-model setups, `OBJECTIVE`. We can convert a function with `branin`'s signature to a single-output observer using `mk_observer`.
#
# The optimization procedure will benefit from having some starting data from the objective function to base its search on. We sample five points from the search space and evaluate them on the observer.
# %%
import trieste
from trieste.acquisition.rule import OBJECTIVE
observer = trieste.utils.objectives.mk_observer(branin, OBJECTIVE)
num_initial_points = 5
initial_query_points = search_space.sample(num_initial_points)
initial_data = observer(initial_query_points)
# %% [markdown]
# ## Model the objective function
#
# The Bayesian optimization procedure estimates the next best points to query by using a probabilistic model of the objective. We'll use Gaussian process regression for this, provided by GPflow. The model will need to be trained on each step as more points are evaluated, so we'll package it with GPflow's Scipy optimizer.
#
# Just like the data output by the observer, the optimization process assumes multiple models, so we'll need to label the model in the same way.
# %%
import gpflow
def build_model(data):
variance = tf.math.reduce_variance(data.observations)
kernel = gpflow.kernels.Matern52(variance=variance, lengthscales=[0.2, 0.2])
def test_optimizer_finds_minima_of_Gardners_Simulation_1(
num_steps: int, acquisition_function_builder
) -> None:
"""
Test that tests the covergence of constrained BO algorithms on the
synthetic "simulation 1" experiment of :cite:`gardner14`.
"""
search_space = Box([0, 0], [6, 6])
def objective(input_data):
x, y = input_data[..., -2], input_data[..., -1]
z = tf.cos(2.0 * x) * tf.cos(y) + tf.sin(x)
return z[:, None]
def constraint(input_data):
x, y = input_data[:, -2], input_data[:, -1]
z = tf.cos(x) * tf.cos(y) - tf.sin(x) * tf.sin(y)
return z[:, None]
MINIMUM = -2.0
MINIMIZER = [math.pi * 1.5, 0.0]
OBJECTIVE = "OBJECTIVE"
CONSTRAINT = "CONSTRAINT"
def observer(query_points): # observe both objective and constraint data
return {
OBJECTIVE: Dataset(query_points, objective(query_points)),
CONSTRAINT: Dataset(query_points, constraint(query_points)),
}
num_initial_points = 5
initial_data = observer(search_space.sample(num_initial_points))
def build_model(data):
variance = tf.math.reduce_variance(data.observations)
kernel = gpflow.kernels.Matern52(variance, tf.constant([0.2, 0.2], tf.float64))
gpr = gpflow.models.GPR((data.query_points, data.observations), kernel, noise_variance=1e-5)
gpflow.utilities.set_trainable(gpr.likelihood, False)
return GaussianProcessRegression(gpr)
models = map_values(build_model, initial_data)
pof = ProbabilityOfFeasibility(threshold=0.5)
acq = acquisition_function_builder(OBJECTIVE, pof.using(CONSTRAINT))
rule: EfficientGlobalOptimization[Box] = EfficientGlobalOptimization(acq)
dataset = (
BayesianOptimizer(observer, search_space)
.optimize(num_steps, initial_data, models, rule)
.try_get_final_datasets()[OBJECTIVE]
)
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 - MINIMIZER)
# these accuracies are the current best for the given number of optimization steps, which makes
# this is a regression test
assert tf.reduce_all(relative_minimizer_err < 0.03, axis=-1)
npt.assert_allclose(best_y, MINIMUM, rtol=0.03)
def observer(x):
y = masked_branin(x)
mask = np.isfinite(y).reshape(-1)
return {
OBJECTIVE: trieste.data.Dataset(x[mask], y[mask]),
FAILURE: trieste.data.Dataset(x, tf.cast(np.isfinite(y), tf.float64))
}
# %% [markdown]
# We can evaluate the observer at points sampled from the search space.
# %%
num_init_points = 15
initial_data = observer(search_space.sample(num_init_points))
# %% [markdown]
# ## Build GPflow models
#
# We'll model the data on the objective with a regression model, and the data on which points failed with a classification model. The regression model will be a `GaussianProcessRegression` wrapping a GPflow `GPR`, and the classification model a `VariationalGaussianProcess` wrapping a GPflow `VGP` with Bernoulli likelihood.
# %%
import gpflow
def create_regression_model(data):
variance = tf.math.reduce_variance(data.observations)
kernel = gpflow.kernels.Matern52(variance=variance,
lengthscales=[0.2, 0.2])
gpr = gpflow.models.GPR(data.astuple(), kernel, noise_variance=1e-5)
def test_no_function_values_are_less_than_global_minimum(
objective: Callable[[TensorType], TensorType], space: Box,
minimum: TensorType) -> None:
samples = space.sample(1000 * len(space.lower))
npt.assert_array_less(tf.squeeze(minimum) - 1e-6, objective(samples))
def test_box_sampling(num_samples: int) -> None:
box = Box(tf.zeros((3, )), tf.ones((3, )))
samples = box.sample(num_samples)
_assert_correct_number_of_unique_constrained_samples(
num_samples, box, samples)
def test_box_sampling_raises_for_invalid_sample_size(num_samples: int) -> None:
with pytest.raises(TF_DEBUGGING_ERROR_TYPES):
box = Box(tf.zeros((3,)), tf.ones((3,)))
box.sample(num_samples)
