Multi trust regions

docs/notebooks/multi_trust_region.pct.py

@@ -0,0 +1,202 @@
+# %%
+# from aim.ext.tensorboard_tracker import Run
+from datetime import datetime
+
+import numpy as np
+import tensorflow as tf
+from matplotlib import pyplot as plt
+from matplotlib.patches import Rectangle
+from matplotlib.pyplot import cm
+
+import trieste
+from trieste.acquisition import ParallelContinuousThompsonSampling
+from trieste.acquisition.optimizer import automatic_optimizer_selector
+from trieste.acquisition.rule import MultiTrustRegionBox, TrustRegionBox
+from trieste.acquisition.utils import split_acquisition_function_calls
+from trieste.ask_tell_optimization import AskTellOptimizer
+from trieste.experimental.plotting import plot_regret
+from trieste.experimental.plotting.plotting import create_grid
+from trieste.logging import pyplot
+from trieste.models.gpflow import GaussianProcessRegression, build_gpr
+from trieste.objectives import Hartmann6, ScaledBranin
+from trieste.types import TensorType
+
+# %%
+np.random.seed(179)
+tf.random.set_seed(179)
+
+# CONFIG
+tensorboard_dir_1 = f"./results/{datetime.now()}/tensorboard"
+
+summary_writer = tf.summary.create_file_writer(tensorboard_dir_1)
+trieste.logging.set_tensorboard_writer(summary_writer)
+
+# %%
+obj = ScaledBranin.objective
+search_space = ScaledBranin.search_space
+
+
+def obj_fun(
+    x: TensorType,
+) -> TensorType:  # contaminate observations with Gaussian noise
+    return obj(x)  # + tf.random.normal([len(x), 1], 0, .1, tf.float64)
+
+
+num_initial_data_points = 6
+num_query_points = 3
+num_steps = 10
+
+# %%
+initial_query_points = search_space.sample(num_initial_data_points)
+observer = trieste.objectives.utils.mk_observer(obj_fun)
+initial_data = observer(initial_query_points)
+
+
+# gpflow_model = build_svgp(
+#     initial_data, search_space, likelihood_variance=0.001, num_inducing_points=50
+# )
+#
+# inducing_point_selector = ConditionalImprovementReduction()
+#
+# model = SparseVariational(
+#     gpflow_model,
+#     num_rff_features=1000,
+#     inducing_point_selector=inducing_point_selector,
+#     optimizer=BatchOptimizer(
+#         tf.optimizers.Adam(0.05), max_iter=100, batch_size=50, compile=True
+#     ),
+# )
+gpflow_model = build_gpr(
+    initial_data,
+    search_space,
+    likelihood_variance=1e-4,
+    trainable_likelihood=False,
+)
+model = GaussianProcessRegression(gpflow_model)
+
+base_rule = trieste.acquisition.rule.EfficientGlobalOptimization(
+    builder=ParallelContinuousThompsonSampling(),
+    num_query_points=num_query_points,
+    optimizer=split_acquisition_function_calls(
+        automatic_optimizer_selector, split_size=100_000
+    ),
+)
+
+acq_rule = MultiTrustRegionBox(
+    TrustRegionBox, base_rule, number_of_tr=num_query_points
+)
+
+ask_tell = AskTellOptimizer(
+    search_space, initial_data, model, fit_model=True, acquisition_rule=acq_rule
+)
+
+# %%
+color = cm.rainbow(np.linspace(0, 1, num_query_points))
+
+Xplot, xx, yy = create_grid(
+    mins=search_space.lower, maxs=search_space.upper, grid_density=90
+)
+ff = obj_fun(Xplot).numpy()
+
+for step in range(num_steps):
+    print(f"step number {step}")
+    trieste.logging.set_step_number(step)
+
+    new_points = ask_tell.ask()
+    new_data = observer(new_points)
+    # monitor models after each tell
+    if summary_writer:
+        models = ask_tell._models  # pylint: disable=protected-access
+        trieste.logging.set_step_number(step)
+
+        with summary_writer.as_default(step=step):
+            for tag, model in models.items():
+                with tf.name_scope(f"{tag}.model"):
+                    model.log()
+
+            fig, ax = plt.subplots(1, 2, squeeze=False, figsize=(15, 5))
+            fig.suptitle(f"step number {step}")
+            ax[0, 0].scatter(
+                ask_tell.dataset.query_points[:, 0].numpy(),
+                ask_tell.dataset.query_points[:, 1].numpy(),
+                color="blue",
+            )
+            ax[0, 0].scatter(
+                new_points[:, 0].numpy(), new_points[:, 1].numpy(), color="red"
+            )
+
+            state = ask_tell.acquisition_state
+            assert state is not None
+            assert isinstance(state, MultiTrustRegionBox.State)
+
+            xmin = {
+                tag: state.acquisition_space.get_subspace(  # type: ignore[attr-defined]
+                    tag
+                ).get_local_min(
+                    ask_tell.dataset
+                )[
+                    0
+                ]
+                for tag in state.acquisition_space.subspace_tags
+            }
+            i = 0
+
+            ax[0, 1].contour(xx, yy, ff.reshape(*xx.shape), 80, alpha=0.5)
+
+            for tag in state.acquisition_space.subspace_tags:
+                ax[0, 1].scatter(
+                    xmin[tag].numpy()[0],
+                    xmin[tag].numpy()[1],
+                    color=color[i],
+                    marker="x",
+                    alpha=0.5,
+                )
+                lb = state.acquisition_space.get_subspace(tag).lower
+                ub = state.acquisition_space.get_subspace(tag).upper
+                ax[0, 1].add_patch(
+                    Rectangle(
+                        (lb[0], lb[1]),
+                        ub[0] - lb[0],
+                        ub[1] - lb[1],
+                        facecolor=color[i],
+                        edgecolor=color[i],
+                        alpha=0.3,
+                    )
+                )
+                ax[0, 1].scatter(
+                    new_points[i, 0].numpy(),
+                    new_points[i, 1].numpy(),
+                    color=color[i],
+                    alpha=0.5,
+                )
+                ax[0, 1].scatter(
+                    ask_tell.dataset.query_points[:, 0].numpy(),
+                    ask_tell.dataset.query_points[:, 1].numpy(),
+                    color="black",
+                    alpha=0.2,
+                )
+                i = i + 1
+
+            pyplot("Query points", fig)
+            plt.close(fig)
+
+    ask_tell.tell(new_data)
+
+# %%
+dataset = ask_tell.dataset
+
+ground_truth_regret = obj(dataset.query_points) - Hartmann6.minimum
+best_found_truth_idx = tf.squeeze(tf.argmin(ground_truth_regret, axis=0))
+
+fig, ax = plt.subplots()
+plot_regret(
+    ground_truth_regret.numpy(), ax, num_init=10, idx_best=best_found_truth_idx
+)
+
+ax.set_yscale("log")
+ax.set_ylabel("Regret")
+ax.set_xlabel("# evaluations")
+
+# %%
+fig, ax = plt.subplots()
+ax.scatter(dataset.query_points[:, 0], dataset.query_points[:, 1])

tests/unit/acquisition/test_rule.py

@@ -18,6 +18,7 @@
 from typing import Callable, Optional
 
 import gpflow
+import numpy as np
 import numpy.testing as npt
 import pytest
 import tensorflow as tf
@@ -48,6 +49,7 @@
     EfficientGlobalOptimization,
     RandomSampling,
     TrustRegion,
+    TrustRegionBox,
 )
 from trieste.acquisition.sampler import (
     ExactThompsonSampler,
@@ -1112,6 +1114,168 @@ def test_turbo_state_deepcopy() -> None:
     npt.assert_allclose(tr_state_copy.y_min, tr_state.y_min)
 
 
+# get_local_min raises if dataset is None.
+def test_trust_region_box_get_local_min_raises_if_dataset_is_none() -> None:
+    search_space = Box([0.0, 0.0], [1.0, 1.0])
+    trb = TrustRegionBox(search_space)
+    with pytest.raises(ValueError, match="dataset must be provided"):
+        trb.get_local_min(None)
+
+
+# get_local_min picks the minimum x and y values from the dataset.
+def test_trust_region_box_get_local_min() -> None:
+    search_space = Box([0.0, 0.0], [1.0, 1.0])
+    dataset = Dataset(
+        tf.constant([[0.1, 0.1], [0.5, 0.5], [0.3, 0.4], [0.8, 0.8], [0.4, 0.4]], dtype=tf.float64),
+        tf.constant([[0.0], [0.5], [0.2], [0.1], [1.0]], dtype=tf.float64),
+    )
+    trb = TrustRegionBox(search_space)
+    trb._lower = tf.constant([0.2, 0.2], dtype=tf.float64)
+    trb._upper = tf.constant([0.7, 0.7], dtype=tf.float64)
+    x_min, y_min = trb.get_local_min(dataset)
+    npt.assert_array_equal(x_min, tf.constant([0.3, 0.4], dtype=tf.float64))
+    npt.assert_array_equal(y_min, tf.constant([0.2], dtype=tf.float64))
+
+
+# get_local_min returns first x value and inf y value when points in dataset are outside the
+# search space.
+def test_trust_region_box_get_local_min_outside_search_space() -> None:
+    search_space = Box([0.0, 0.0], [1.0, 1.0])
+    dataset = Dataset(
+        tf.constant([[1.2, 1.3], [-0.4, -0.5]], dtype=tf.float64),
+        tf.constant([[0.7], [0.9]], dtype=tf.float64),
+    )
+    trb = TrustRegionBox(search_space)
+    x_min, y_min = trb.get_local_min(dataset)
+    npt.assert_array_equal(x_min, tf.constant([1.2, 1.3], dtype=tf.float64))
+    npt.assert_array_equal(y_min, tf.constant([np.inf], dtype=tf.float64))
+
+
+# get_single_model_and_dataset returns model and dataset with the OBJECTIVE tag.
+def test_trust_region_box_get_single_model_dataset() -> None:
+    search_space = Box([0.0, 0.0], [1.0, 1.0])
+    dataset = Dataset(tf.zeros([1, 2], dtype=tf.float64), tf.zeros([1, 1], dtype=tf.float64))
+    models = {
+        "foo": QuadraticMeanAndRBFKernel(),
+        OBJECTIVE: QuadraticMeanAndRBFKernelWithSamplers(dataset),
+    }
+    datasets = {"foo": empty_dataset([2], [1]), OBJECTIVE: dataset}
+    trb = TrustRegionBox(search_space)
+    _model, _dataset = trb.get_single_model_and_dataset(models, datasets)
+    assert isinstance(_model, QuadraticMeanAndRBFKernelWithSamplers)
+    assert _dataset is dataset
+
+
+# Reinitialize sets the box to a random location, and sets the eps and y_min values.
+def test_trust_region_box_reinitialize() -> None:
+    search_space = Box([0.0, 0.0], [1.0, 1.0])
+    datasets = {
+        OBJECTIVE: Dataset(  # Points outside the search space should be ignored.
+            tf.constant([[1.2, 1.3], [-0.4, -0.5]], dtype=tf.float64),
+            tf.constant([[0.7], [0.9]], dtype=tf.float64),
+        )
+    }
+    trb = TrustRegionBox(search_space)
+    trb.reinitialize(datasets=datasets)
+
+    exp_eps = 0.5 * (search_space.upper - search_space.lower) / 5.0 ** (1.0 / 2.0)
+    npt.assert_array_equal(trb._eps, exp_eps)
+    npt.assert_array_compare(np.less_equal, search_space.lower, trb.location)
+    npt.assert_array_compare(np.less_equal, trb.location, search_space.upper)
+    npt.assert_array_compare(np.less_equal, search_space.lower, trb.lower)
+    npt.assert_array_compare(np.less_equal, trb.upper, search_space.upper)
+    npt.assert_array_compare(np.less_equal, trb.upper - trb.lower, 2 * exp_eps)
+    npt.assert_array_equal(trb.y_min, tf.constant([np.inf], dtype=tf.float64))
+
+
+# Update call reintializes the box if eps is smaller than min_eps.
+def test_trust_region_box_update_reinitialize() -> None:
+    search_space = Box([0.0, 0.0], [1.0, 1.0])
+    datasets = {
+        OBJECTIVE: Dataset(  # Points outside the search space should be ignored.
+            tf.constant([[1.2, 1.3], [-0.4, -0.5]], dtype=tf.float64),
+            tf.constant([[0.7], [0.9]], dtype=tf.float64),
+        )
+    }
+    trb = TrustRegionBox(search_space, min_eps=0.5)
+    trb.reinitialize(datasets=datasets)
+    location = trb.location
+
+    trb.update(datasets=datasets)
+    npt.assert_array_compare(np.not_equal, location, trb.location)
+    location = trb.location
+
+    trb.update(datasets=datasets)
+    npt.assert_array_compare(np.not_equal, location, trb.location)
+
+
+# Update call does not reintialize the box if eps is larger than min_eps.
+def test_trust_region_box_update_no_reinitialize() -> None:
+    search_space = Box([0.0, 0.0], [1.0, 1.0])
+    datasets = {
+        OBJECTIVE: Dataset(
+            tf.constant([[0.5, 0.5], [0.0, 0.0], [1.0, 1.0]], dtype=tf.float64),
+            tf.constant([[0.5], [0.0], [1.0]], dtype=tf.float64),
+        )
+    }
+    trb = TrustRegionBox(search_space, min_eps=0.1)
+    trb.reinitialize(datasets=datasets)
+    trb.location = tf.constant([0.5, 0.5], dtype=tf.float64)
+    location = trb.location
+
+    trb.update(datasets=datasets)
+    npt.assert_array_equal(location, trb.location)
+
+
+# Update shrinks/expands box on successful/unsuccessful step.
+@pytest.mark.parametrize("success", [True, False])
+def test_trust_region_box_update_size(success: bool) -> None:
+    search_space = Box([0.0, 0.0], [1.0, 1.0])
+    datasets = {
+        OBJECTIVE: Dataset(
+            tf.constant([[0.5, 0.5], [0.0, 0.0], [1.0, 1.0]], dtype=tf.float64),
+            tf.constant([[0.5], [0.0], [1.0]], dtype=tf.float64),
+        )
+    }
+    trb = TrustRegionBox(search_space, min_eps=0.1)
+    trb.reinitialize(datasets=datasets)
+    eps = trb._eps
+
+    if success:
+        # Sample a point from the box.
+        point = trb.sample(1)
+    else:
+        # Pick point outside the box.
+        point = tf.constant([[1.2, 1.3]], dtype=tf.float64)
+
+    # Add a new min point to the dataset.
+    datasets[OBJECTIVE] = Dataset(
+        np.concatenate([datasets[OBJECTIVE].query_points, point], axis=0),
+        np.concatenate([datasets[OBJECTIVE].observations, [[-0.1]]], axis=0),
+    )
+    # Update the box.
+    trb.update(datasets=datasets)
+
+    if success:
+        # Check that the location is the new min point.
+        point = np.squeeze(point)
+        npt.assert_allclose(point, trb.location)
+        npt.assert_allclose(tf.constant([-0.1], dtype=tf.float64), trb.y_min)
+        # Check that the box is smaller by beta.
+        npt.assert_allclose(eps / trb._beta, trb._eps)
+    else:
+        # Check that the location is the old min point.
+        point, y_min = trb.get_local_min(datasets[OBJECTIVE])
+        npt.assert_allclose(point, trb.location)
+        npt.assert_allclose(y_min, trb.y_min)
+        # Check that the box is larger by beta.
+        npt.assert_allclose(eps * trb._beta, trb._eps)
+
+    # Check the new box bounds.
+    npt.assert_allclose(trb.lower, np.maximum(point - trb._eps, search_space.lower))
+    npt.assert_allclose(trb.upper, np.minimum(point + trb._eps, search_space.upper))
+
+
 def test_asynchronous_rule_state_pending_points() -> None:
     pending_points = tf.constant([[1], [2], [3]])