Updatable discrete trust region

docs/notebooks/mixed_search_spaces.pct.py

@@ -251,6 +251,10 @@
 # `SingleObjectiveTrustRegionDiscrete`, which follows a very similar algorithm to the continuous
 # one, but with a region defined by a set of neighboring points. Both the continuous and
 # discrete sub-regions are updated at each step of the optimization.
+#
+# Note that `SingleObjectiveTrustRegionDiscrete` is designed for discrete numerical
+# variables only, which we have in this example. It is not suitable for qualitative (categorical,
+# ordinal and binary) variables.
 
 # %%
 from trieste.acquisition import ParallelContinuousThompsonSampling

tests/integration/test_mixed_space_bayesian_optimization.py

@@ -31,6 +31,7 @@
     AcquisitionRule,
     BatchTrustRegionProduct,
     EfficientGlobalOptimization,
+    FixedPointTrustRegionDiscrete,
     SingleObjectiveTrustRegionBox,
     SingleObjectiveTrustRegionDiscrete,
     UpdatableTrustRegionProduct,
@@ -96,6 +97,35 @@ def _get_mixed_search_space() -> TaggedProductSearchSpace:
             ),
             id="LocalPenalization",
         ),
+        pytest.param(
+            8,
+            BatchTrustRegionProduct(
+                [
+                    UpdatableTrustRegionProduct(
+                        [
+                            FixedPointTrustRegionDiscrete(
+                                cast(
+                                    DiscreteSearchSpace, mixed_search_space.get_subspace("discrete")
+                                )
+                            ),
+                            SingleObjectiveTrustRegionBox(
+                                mixed_search_space.get_subspace("continuous")
+                            ),
+                        ],
+                        tags=mixed_search_space.subspace_tags,
+                    )
+                    for _ in range(10)
+                ],
+                EfficientGlobalOptimization(
+                    ParallelContinuousThompsonSampling(),
+                    # Use a large batch to ensure discrete init finds a good point.
+                    # We are using a fixed point trust region for the discrete space, so
+                    # the init point is randomly chosen and then never updated.
+                    num_query_points=10,
+                ),
+            ),
+            id="TrustRegionSingleObjectiveFixed",
+        ),
         pytest.param(
             8,
             BatchTrustRegionProduct(

tests/unit/acquisition/test_optimizer.py

@@ -36,6 +36,7 @@
     generate_initial_points,
     generate_random_search_optimizer,
     get_bounds_of_box_relaxation_around_point,
+    get_bounds_of_optimization,
     optimize_discrete,
     sample_from_space,
 )
@@ -567,6 +568,121 @@ def test_get_bounds_of_box_relaxation_around_point(
     npt.assert_array_equal(bounds.ub, upper)
 
 
+def test_get_bounds_of_optimization_raises_on_incorrect_num_subspaces() -> None:
+    search_space = TaggedMultiSearchSpace([Box([-1], [2]), Box([3], [4])])
+    points = tf.ones([4, 3, 1], dtype=tf.float64)
+    with pytest.raises(TF_DEBUGGING_ERROR_TYPES, match="The vectorization of the target function"):
+        get_bounds_of_optimization(search_space, points)
+
+
+@pytest.mark.parametrize(
+    "search_space, exp_bounds",
+    [
+        (Box([-1], [2]), [(-1, 2)] * 12),
+        (TaggedProductSearchSpace([Box([-1], [2]), Box([3], [4])]), [([-1, 3], [2, 4])] * 12),
+        (TaggedMultiSearchSpace([Box([-1], [2]), Box([3], [4])]), [(-1, 2), (3, 4)] * 6),
+        (
+            TaggedMultiSearchSpace([TaggedProductSearchSpace([Box([-1], [2]), Box([3], [4])])]),
+            [([-1, 3], [2, 4])] * 12,
+        ),
+        (
+            TaggedMultiSearchSpace(
+                [
+                    TaggedProductSearchSpace([Box([-1], [2]), Box([3], [4])]),
+                    TaggedProductSearchSpace([Box([-3], [0]), Box([7], [8])]),
+                ]
+            ),
+            [([-1, 3], [2, 4]), ([-3, 7], [0, 8])] * 6,
+        ),
+        (
+            TaggedProductSearchSpace([Box([-1], [2]), DiscreteSearchSpace([[-0.5], [0.2]])]),
+            [  # Expect use of get_bounds_of_box_relaxation_around_point.
+                ([-1, 11], [2, 11]),
+                ([-1, 13], [2, 13]),
+                ([-1, 15], [2, 15]),
+                ([-1, 17], [2, 17]),
+                ([-1, 21], [2, 21]),
+                ([-1, 23], [2, 23]),
+                ([-1, 25], [2, 25]),
+                ([-1, 27], [2, 27]),
+                ([-1, 31], [2, 31]),
+                ([-1, 33], [2, 33]),
+                ([-1, 35], [2, 35]),
+                ([-1, 37], [2, 37]),
+            ],
+        ),
+        (
+            TaggedMultiSearchSpace(
+                [DiscreteSearchSpace([[-0.5], [0.2]]), DiscreteSearchSpace([[1.0], [-0.2], [3.7]])]
+            ),
+            [(-0.5, 0.2), (-0.2, 3.7)] * 6,
+        ),
+        (
+            TaggedMultiSearchSpace(
+                [TaggedProductSearchSpace([Box([-1], [2]), DiscreteSearchSpace([[-0.5], [0.2]])])]
+            ),
+            [  # Expect use of get_bounds_of_box_relaxation_around_point.
+                ([-1, 11], [2, 11]),
+                ([-1, 13], [2, 13]),
+                ([-1, 15], [2, 15]),
+                ([-1, 17], [2, 17]),
+                ([-1, 21], [2, 21]),
+                ([-1, 23], [2, 23]),
+                ([-1, 25], [2, 25]),
+                ([-1, 27], [2, 27]),
+                ([-1, 31], [2, 31]),
+                ([-1, 33], [2, 33]),
+                ([-1, 35], [2, 35]),
+                ([-1, 37], [2, 37]),
+            ],
+        ),
+        (
+            TaggedMultiSearchSpace(
+                [
+                    TaggedProductSearchSpace(
+                        [Box([-1], [2]), DiscreteSearchSpace([[-0.5], [0.2]])]
+                    ),
+                    TaggedProductSearchSpace(
+                        [Box([-3], [0]), DiscreteSearchSpace([[1.0], [-0.2], [3.7]])]
+                    ),
+                ]
+            ),
+            [  # Expect use of get_bounds_of_box_relaxation_around_point.
+                ([-1, 11], [2, 11]),
+                ([-3, 13], [0, 13]),
+                ([-1, 15], [2, 15]),
+                ([-3, 17], [0, 17]),
+                ([-1, 21], [2, 21]),
+                ([-3, 23], [0, 23]),
+                ([-1, 25], [2, 25]),
+                ([-3, 27], [0, 27]),
+                ([-1, 31], [2, 31]),
+                ([-3, 33], [0, 33]),
+                ([-1, 35], [2, 35]),
+                ([-3, 37], [0, 37]),
+            ],
+        ),
+    ],
+)
+def test_get_bounds_of_optimization(
+    search_space: SearchSpace, exp_bounds: list[Tuple[TensorType, TensorType]]
+) -> None:
+    points = tf.constant(
+        [
+            [[10, 11], [12, 13], [14, 15], [16, 17]],
+            [[20, 21], [22, 23], [24, 25], [26, 27]],
+            [[30, 31], [32, 33], [34, 35], [36, 37]],
+        ],
+        dtype=tf.float64,
+    )
+    bounds = get_bounds_of_optimization(search_space, points)
+
+    assert len(bounds) == 12
+    for exp, b in zip(exp_bounds, bounds):
+        npt.assert_array_equal(exp[0], b.lb)
+        npt.assert_array_equal(exp[1], b.ub)
+
+
 def test_batchify_joint_raises_with_invalid_batch_size() -> None:
     batch_size_one_optimizer = generate_continuous_optimizer()
     with pytest.raises(ValueError):

trieste/acquisition/optimizer.py

@@ -596,45 +596,8 @@ def _objective_value(vectorized_x: TensorType) -> TensorType:  # [N, D] -> [N, 1
     def _objective_value_and_gradient(x: TensorType) -> Tuple[TensorType, TensorType]:
         return tfp.math.value_and_gradient(_objective_value, x)  # [len(x), 1], [len(x), D]
 
-    if isinstance(
-        space, TaggedProductSearchSpace
-    ):  # build continuous relaxation of discrete subspaces
-        bounds = [
-            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):
-        subspaces = [space.get_subspace(tag) for tag in space.subspace_tags]
-        # Create a sequence of bounds for each optimization run, per subspace.
-        # For product subspaces, we build a continuous relaxation of the discrete subspaces.
-        # Otherwise, we use the original bounds.
-        bounds = [
-            [
-                get_bounds_of_box_relaxation_around_point(ss, starting_points[i : i + 1, j])
-                if isinstance(ss, TaggedProductSearchSpace)
-                else spo.Bounds(ss.lower, ss.upper)
-                for j, ss in enumerate(subspaces)
-            ]
-            for i in tf.range(num_optimization_runs_per_function)
-        ]
-        bounds = list(chain.from_iterable(bounds))
-        # The bounds is a sequence of tensors, stack them into a single tensor. In this case
-        # the vectorization of the target function must be a multple of the length of the sequence.
-        remainder = num_optimization_runs % len(bounds)
-        tf.debugging.assert_equal(
-            remainder,
-            tf.cast(0, dtype=remainder.dtype),
-            message=(
-                f"""
-                The number of optimization runs {num_optimization_runs} must be a multiple of the
-                length of the bounds sequence {len(bounds)}.
-                """
-            ),
-        )
-        multiple = num_optimization_runs // len(bounds)
-        bounds = bounds * multiple
-    else:
-        bounds = [spo.Bounds(space.lower, space.upper)] * num_optimization_runs
+    # Get the bounds for the optimization
+    bounds = get_bounds_of_optimization(space, starting_points)
 
     # Initialize the numpy arrays to be passed to the greenlets
     np_batch_x = np.zeros((num_optimization_runs, tf.shape(starting_points)[-1]), dtype=np.float64)
@@ -777,6 +740,70 @@ def get_bounds_of_box_relaxation_around_point(
     return spo.Bounds(space_with_fixed_discrete.lower, space_with_fixed_discrete.upper)
 
 
+def get_bounds_of_optimization(space: SearchSpace, starting_points: TensorType) -> List[spo.Bounds]:
+    """
+    A function to return the bounds of the optimization for each of the
+    individual optimization runs. This function is used to provide the
+    bounds for the Scipy optimizer.
+
+    :param space: The original search space.
+    :param starting_points: The points at which to begin our optimizations of shape
+        [num_optimization_runs, V, D]. The leading dimension of
+        `starting_points` controls the number of individual optimization runs
+        for each of the V target functions.
+    :return: A list of bounds for the Scipy optimizer. The length of the list
+        is equal to the number of individual optimization runs, i.e. `num_optimization_runs` x `V`.
+    """
+
+    num_optimization_runs_per_function, V, D = starting_points.shape
+    num_total_optimization_runs = num_optimization_runs_per_function * V
+
+    if isinstance(
+        space, TaggedProductSearchSpace
+    ):  # build continuous relaxation of discrete subspaces
+        vectorized_starting_points = tf.reshape(
+            starting_points, [-1, D]
+        )  # [num_total_optimization_runs, D]
+        bounds = [
+            get_bounds_of_box_relaxation_around_point(space, vectorized_starting_points[i : i + 1])
+            for i in tf.range(num_total_optimization_runs)
+        ]
+    elif isinstance(space, TaggedMultiSearchSpace):
+        subspaces = [space.get_subspace(tag) for tag in space.subspace_tags]
+        # The vectorization `V` of the target function must be a multple of the number of subspaces.
+        remainder = V % len(subspaces)
+        tf.debugging.assert_equal(
+            remainder,
+            0,
+            message=(
+                f"""
+                The vectorization of the target function {V} must be a multiple of the number of
+                subspaces {len(subspaces)}.
+                """
+            ),
+        )
+        multiple = V // len(subspaces)
+        subspaces = subspaces * multiple
+
+        # Create a sequence of bounds for each optimization run, per subspace.
+        # For product subspaces, we build a continuous relaxation of the discrete subspaces.
+        # Otherwise, we use the original bounds.
+        bounds = [
+            [
+                get_bounds_of_box_relaxation_around_point(ss, starting_points[i : i + 1, j])
+                if isinstance(ss, TaggedProductSearchSpace)
+                else spo.Bounds(ss.lower, ss.upper)
+                for j, ss in enumerate(subspaces)
+            ]
+            for i in tf.range(num_optimization_runs_per_function)
+        ]
+        bounds = list(chain.from_iterable(bounds))
+    else:
+        bounds = [spo.Bounds(space.lower, space.upper)] * num_total_optimization_runs
+
+    return bounds
+
+
 def batchify_joint(
     batch_size_one_optimizer: AcquisitionOptimizer[SearchSpaceType],
     batch_size: int,

trieste/acquisition/rule.py

@@ -2076,7 +2076,18 @@ def update(
 class SingleObjectiveTrustRegionDiscrete(UpdatableTrustRegionDiscrete):
     """
     An updatable discrete trust region that maintains a set of neighboring points around a
-    single location point. The region is updated based on the best point found in the region.
+    single location point, allowing for local exploration of the search space. The region is
+    updated based on the best point found in the region.
+
+    This trust region is designed for discrete numerical variables. As it uses axis-aligned
+    Euclidean distance to determine the neighbors within the region, it is not suitable for
+    qualitative (categorical, ordinal and binary) variables.
+
+    When using this trust region, it is important to consider the scaling of the number of value
+    combinations. Since the region computes pairwise distances between points, the computational
+    and memory complexity increases quadratically with the number of points. For example,
+    1000 3D points will result in the distances matrix containing 1000x1000x3 entries. Therefore,
+    this trust region is not suitable for problems with a large number of points.
     """
 
     def __init__(
@@ -2110,36 +2121,41 @@ def __init__(
         self._min_eps = min_eps
         self._initialized = False
         self._step_is_success = False
+        self._init_location()
+        self._compute_global_distances()
+        self._init_eps()
+        self._update_neighbors()
+        self._y_min = np.inf
+
+    @property
+    def location(self) -> TensorType:
+        """Center point of the region."""
+        return tf.reshape(tf.gather(self.global_search_space.points, self._location_ix), (-1,))
 
+    def _init_location(self) -> None:
         # Random initial location index from the global search space.
         self._location_ix = tf.random.categorical(
             tf.ones(
                 (
                     1,
-                    global_search_space.points.shape[0],
+                    self.global_search_space.points.shape[0],
                 )
             ),
             1,
         )[0]
 
-        # Pairwise distances for each axis in the global search space.
-        points = global_search_space.points
-        self._global_distances = tf.abs(tf.expand_dims(points, -2) - tf.expand_dims(points, -3))
-
-        self._init_eps()
-        self._update_neighbors()
-        self._y_min = np.inf
-
-    @property
-    def location(self) -> TensorType:
-        """Center point of the region."""
-        return tf.reshape(tf.gather(self.global_search_space.points, self._location_ix), (-1,))
-
     def _init_eps(self) -> None:
         global_lower = self.global_search_space.lower
         global_upper = self.global_search_space.upper
         self.eps = 0.5 * (global_upper - global_lower) / (5.0 ** (1.0 / global_lower.shape[-1]))
 
+    def _compute_global_distances(self) -> None:
+        # Pairwise distances along each axis in the global search space.
+        points = self.global_search_space.points
+        self._global_distances = tf.abs(
+            tf.expand_dims(points, -2) - tf.expand_dims(points, -3)
+        )  # [num_points, num_points, D]
+
     def _update_neighbors(self) -> None:
         # Indices of the neighbors within the trust region.
         neighbors_mask = tf.reduce_all(
@@ -2165,15 +2181,7 @@ def initialize(
         """
 
         datasets = self.select_in_region(datasets)
-        self._location_ix = tf.random.categorical(
-            tf.ones(
-                (
-                    1,
-                    self.global_search_space.points.shape[0],
-                )
-            ),
-            1,
-        )[0]
+        self._init_location()
         self._step_is_success = False
         self._init_eps()
         self._update_neighbors()