Repository: secondmind-labs/trieste

Issue:
remove redundancy in `optimize_random`
`optimize_random` is identical to `optimize_discrete`, barring a single line - we could reduce the code by adding an argument random in `optimize_discrete` and use it in both places


Discussion:

Resolving patch:
diff --git a/trieste/acquisition/optimizer.py b/trieste/acquisition/optimizer.py
--- a/trieste/acquisition/optimizer.py
+++ b/trieste/acquisition/optimizer.py
@@ -108,24 +108,11 @@ def automatic_optimizer_selector(
         )
 
 
-def optimize_discrete(
-    space: DiscreteSearchSpace,
-    target_func: Union[AcquisitionFunction, Tuple[AcquisitionFunction, int]],
+def _get_max_discrete_points(
+    points: TensorType, target_func: Union[AcquisitionFunction, Tuple[AcquisitionFunction, int]]
 ) -> TensorType:
-    """
-    An :const:`AcquisitionOptimizer` for :class:'DiscreteSearchSpace' spaces.
-
-    When this functions receives an acquisition-integer tuple as its `target_func`,it evaluates
-    all the points in the search space for each of the individual V functions making
-    up `target_func`.
-
-    :param space: The space of points over which to search, for points with shape [D].
-    :param target_func: The function to maximise, with input shape [..., V, D] and output shape
-            [..., V].
-    :return: The V points in ``space`` that maximises ``target_func``, with shape [V, D].
-    """
-
-    if isinstance(target_func, tuple):  # check if we need a vectorized optimizer
+    # check if we need a vectorized optimizer
+    if isinstance(target_func, tuple):
         target_func, V = target_func
     else:
         V = 1
@@ -133,7 +120,6 @@ def optimize_discrete(
     if V < 0:
         raise ValueError(f"vectorization must be positive, got {V}")
 
-    points = space.points[:, None, :]
     tiled_points = tf.tile(points, [1, V, 1])
     target_func_values = target_func(tiled_points)
     tf.debugging.assert_shapes(
@@ -151,6 +137,26 @@ def optimize_discrete(
     return tf.gather(tf.transpose(tiled_points, [1, 0, 2]), best_indices, batch_dims=1)  # [V, D]
 
 
+def optimize_discrete(
+    space: DiscreteSearchSpace,
+    target_func: Union[AcquisitionFunction, Tuple[AcquisitionFunction, int]],
+) -> TensorType:
+    """
+    An :const:`AcquisitionOptimizer` for :class:'DiscreteSearchSpace' spaces.
+
+    When this functions receives an acquisition-integer tuple as its `target_func`,it evaluates
+    all the points in the search space for each of the individual V functions making
+    up `target_func`.
+
+    :param space: The space of points over which to search, for points with shape [D].
+    :param target_func: The function to maximise, with input shape [..., V, D] and output shape
+            [..., V].
+    :return: The V points in ``space`` that maximises ``target_func``, with shape [V, D].
+    """
+    points = space.points[:, None, :]
+    return _get_max_discrete_points(points, target_func)
+
+
 def generate_continuous_optimizer(
     num_initial_samples: int = NUM_SAMPLES_MIN,
     num_optimization_runs: int = 10,
@@ -624,31 +630,7 @@ def optimize_random(
                 [..., V].
         :return: The V points in ``space`` that maximises ``target_func``, with shape [V, D].
         """
-        if isinstance(target_func, tuple):  # check if we need a vectorized optimizer
-            target_func, V = target_func
-        else:
-            V = 1
-
-        if V < 0:
-            raise ValueError(f"vectorization must be positive, got {V}")
-
         points = space.sample(num_samples)[:, None, :]
-        tiled_points = tf.tile(points, [1, V, 1])
-        target_func_values = target_func(tiled_points)
-        tf.debugging.assert_shapes(
-            [(target_func_values, ("_", V))],
-            message=(
-                f"""
-                The result of function target_func has shape
-                {tf.shape(target_func_values)}, however, expected a trailing
-                dimension of size {V}.
-                """
-            ),
-        )
-
-        best_indices = tf.math.argmax(target_func_values, axis=0)  # [V]
-        return tf.gather(
-            tf.transpose(tiled_points, [1, 0, 2]), best_indices, batch_dims=1
-        )  # [V, D]
+        return _get_max_discrete_points(points, target_func)
 
     return optimize_random


Test patch: