Multi trust regions

docs/notebooks/multi_trust_region.pct.py

@@ -0,0 +1,200 @@
+# %%
+# 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
+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(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])

trieste/acquisition/optimizer.py

@@ -30,10 +30,12 @@
 from .. import logging
 from ..space import (
     Box,
+    CollectionSearchSpace,
     Constraint,
     DiscreteSearchSpace,
     SearchSpace,
     SearchSpaceType,
+    TaggedMultiSearchSpace,
     TaggedProductSearchSpace,
 )
 from ..types import TensorType
@@ -100,9 +102,10 @@ def automatic_optimizer_selector(
     if isinstance(space, DiscreteSearchSpace):
         return optimize_discrete(space, target_func)
 
-    elif isinstance(space, (Box, TaggedProductSearchSpace)):
-        num_samples = tf.maximum(NUM_SAMPLES_MIN, NUM_SAMPLES_DIM * tf.shape(space.lower)[-1])
-        num_runs = NUM_RUNS_DIM * tf.shape(space.lower)[-1]
+    elif isinstance(space, (Box, CollectionSearchSpace)):
+        space_dim = space.dimension
+        num_samples = tf.maximum(NUM_SAMPLES_MIN, NUM_SAMPLES_DIM * space_dim)
+        num_runs = NUM_RUNS_DIM * space_dim
         return generate_continuous_optimizer(
             num_initial_samples=num_samples,
             num_optimization_runs=num_runs,
@@ -170,10 +173,10 @@ def generate_continuous_optimizer(
     num_optimization_runs: int = 10,
     num_recovery_runs: int = 10,
     optimizer_args: Optional[dict[str, Any]] = None,
-) -> AcquisitionOptimizer[Box | TaggedProductSearchSpace]:
+) -> AcquisitionOptimizer[Box | CollectionSearchSpace]:
     """
-    Generate a gradient-based optimizer for :class:'Box' and :class:'TaggedProductSearchSpace'
-    spaces and batches of size 1. In the case of a :class:'TaggedProductSearchSpace', We perform
+    Generate a gradient-based optimizer for :class:'Box' and :class:'CollectionSearchSpace'
+    spaces and batches of size 1. In the case of a :class:'CollectionSearchSpace', We perform
     gradient-based optimization across all :class:'Box' subspaces, starting from the best location
     found across a sample of `num_initial_samples` random points.
 
@@ -216,14 +219,14 @@ def generate_continuous_optimizer(
         raise ValueError(f"num_recovery_runs must be zero or greater, got {num_recovery_runs}")
 
     def optimize_continuous(
-        space: Box | TaggedProductSearchSpace,
+        space: Box | CollectionSearchSpace,
         target_func: Union[AcquisitionFunction, Tuple[AcquisitionFunction, int]],
     ) -> TensorType:
         """
         A gradient-based :const:`AcquisitionOptimizer` for :class:'Box'
-        and :class:`TaggedProductSearchSpace' spaces.
+        and :class:`CollectionSearchSpace' spaces.
 
-        For :class:'TaggedProductSearchSpace' we only apply gradient updates to
+        For :class:'CollectionSearchSpace' we only apply gradient updates to
         its class:'Box' subspaces.
 
         When this functions receives an acquisition-integer tuple as its `target_func`,it
@@ -246,8 +249,26 @@ def optimize_continuous(
         if V < 0:
             raise ValueError(f"vectorization must be positive, got {V}")
 
-        candidates = space.sample(num_initial_samples)[:, None, :]  # [num_initial_samples, 1, D]
-        tiled_candidates = tf.tile(candidates, [1, V, 1])  # [num_initial_samples, V, D]
+        candidates = space.sample(num_initial_samples)
+        if tf.rank(candidates) == 3:
+            # If samples is a tensor of rank 3, then it is a batch of samples. In this case
+            # the length of the second dimension must be equal to the vectorization of the target
+            # function.
+            tf.debugging.assert_equal(
+                tf.shape(candidates)[1],
+                V,
+                message=(
+                    f"""
+                    The batch shape of initial samples {tf.shape(candidates)[1]} must be equal to
+                    the vectorization of the target function {V}.
+                    """
+                ),
+            )
+            tiled_candidates = candidates  # [num_initial_samples, V, D]
+        else:
+            tiled_candidates = tf.tile(
+                candidates[:, None, :], [1, V, 1]
+            )  # [num_initial_samples, V, D]
 
         target_func_values = target_func(tiled_candidates)  # [num_samples, V]
         tf.debugging.assert_shapes(
@@ -289,11 +310,28 @@ def optimize_continuous(
         total_nfev = tf.reduce_max(nfev)  # acquisition function is evaluated in parallel
 
         recovery_run = False
-        if (
-            num_recovery_runs and not successful_optimization
-        ):  # if all optimizations failed for a function then try again from random starts
-            random_points = space.sample(num_recovery_runs)[:, None, :]  # [num_recovery_runs, 1, D]
-            tiled_random_points = tf.tile(random_points, [1, V, 1])  # [num_recovery_runs, V, D]
+        if num_recovery_runs and not successful_optimization:
+            # if all optimizations failed for a function then try again from random starts
+            random_points = space.sample(num_recovery_runs)
+            if tf.rank(random_points) == 3:
+                # If samples is a tensor of rank 3, then it is a batch of samples. In this case
+                # the length of the second dimension must be equal to the vectorization of the
+                # target function.
+                tf.debugging.assert_equal(
+                    tf.shape(random_points)[1],
+                    V,
+                    message=(
+                        f"""
+                        The batch shape of random samples {tf.shape(random_points)[1]} must be
+                        equal to the vectorization of the target function {V}.
+                        """
+                    ),
+                )
+                tiled_random_points = random_points  # [num_recovery_runs, V, D]
+            else:
+                tiled_random_points = tf.tile(
+                    random_points[:, None, :], [1, V, 1]
+                )  # [num_recovery_runs, V, D]
 
             (
                 recovery_successes,
@@ -383,7 +421,7 @@ def _perform_parallel_continuous_optimization(
     of our acquisition function (and its gradients) is calculated in parallel
     (for each optimization step) using Tensorflow.
 
-    For :class:'TaggedProductSearchSpace' we only apply gradient updates to
+    For :class:'CollectionSearchSpace' we only apply gradient updates to
     its :class:'Box' subspaces, fixing the discrete elements to the best values
     found across the initial random search. To fix these discrete elements, we
     optimize over a continuous :class:'Box' relaxation of the discrete subspaces
@@ -434,6 +472,24 @@ def _objective_value_and_gradient(x: TensorType) -> Tuple[TensorType, TensorType
             get_bounds_of_box_relaxation_around_point(space, vectorized_starting_points[i : i + 1])
             for i in tf.range(num_optimization_runs)
         ]
+    elif isinstance(space, TaggedMultiSearchSpace):
+        bounds = [
+            spo.Bounds(lower, upper)
+            for lower, upper in zip(space.subspace_lower, space.subspace_upper)
+        ]
+        # If bounds is a sequence of tensors, stack them into a single tensor. In this case
+        # the length of the sequence must be equal to the vectorization of the target function.
+        tf.debugging.assert_equal(
+            len(bounds),
+            V,
+            message=(
+                f"""
+                The length of bounds sequence {len(bounds)} must be equal to the
+                vectorization of the target function {V}.
+                """
+            ),
+        )
+        bounds = bounds * num_optimization_runs_per_function
     else:
         bounds = [spo.Bounds(space.lower, space.upper)] * num_optimization_runs