Repository: secondmind-labs/trieste

Issue:
Allow specification of scipy optimizer kwargs when building continuous optimizers
Allow specification of scipy optimizer kwargs when building continuous optimizers


Discussion:
It might be nice to be able to specify some of the scipy optimizer behaviour from within Trieste. For example, we might want to set max number of iterations or desired optimization tolerance. At the moment we just use the defaults (https://docs.scipy.org/doc/scipy/reference/optimize.minimize-lbfgsb.html).

Resolving patch:
diff --git a/trieste/acquisition/optimizer.py b/trieste/acquisition/optimizer.py
--- a/trieste/acquisition/optimizer.py
+++ b/trieste/acquisition/optimizer.py
@@ -1,4 +1,4 @@
-# Copyright 2020 The Trieste Contributors
+# Copyright 2021 The Trieste Contributors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -11,13 +11,16 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 r"""
 This module contains functionality for optimizing
 :data:`~trieste.acquisition.AcquisitionFunction`\ s over :class:`~trieste.space.SearchSpace`\ s.
 """
+
 from __future__ import annotations
 
-from typing import Callable, TypeVar
+import copy
+from typing import Any, Callable, TypeVar
 
 import gpflow
 import scipy.optimize as spo
@@ -32,6 +35,22 @@
 """ Type variable bound to :class:`~trieste.space.SearchSpace`. """
 
 
+NUM_SAMPLES_MIN: int = 5000
+"""
+The default minimum number of initial samples for :func:`generate_continuous_optimizer` function in
+:func:`automatic_optimizer_selector`, used for determining the number of initial samples in the
+multi-start acquisition function optimization.
+"""
+
+
+NUM_SAMPLES_DIM: int = 1000
+"""
+The default minimum number of initial samples per dimension of the search space for
+:func:`generate_continuous_optimizer` function in :func:`automatic_optimizer_selector`, used for
+determining the number of initial samples in the multi-start acquisition function optimization.
+"""
+
+
 class FailedOptimizationError(Exception):
     """Raised when an acquisition optimizer fails to optimize"""
 
@@ -63,7 +82,7 @@ def automatic_optimizer_selector(
         return optimize_discrete(space, target_func)
 
     elif isinstance(space, (Box, TaggedProductSearchSpace)):
-        num_samples = tf.minimum(5000, 1000 * tf.shape(space.lower)[-1])
+        num_samples = tf.minimum(NUM_SAMPLES_MIN, NUM_SAMPLES_DIM * tf.shape(space.lower)[-1])
         return generate_continuous_optimizer(num_samples)(space, target_func)
 
     else:
@@ -100,6 +119,7 @@ def generate_continuous_optimizer(
     num_initial_samples: int = 1000,
     num_optimization_runs: int = 1,
     num_recovery_runs: int = 5,
+    optimizer_args: dict[str, Any] = dict(),
 ) -> AcquisitionOptimizer[Box | TaggedProductSearchSpace]:
     """
     Generate a gradient-based optimizer for :class:'Box' and :class:'TaggedProductSearchSpace'
@@ -107,11 +127,11 @@ def generate_continuous_optimizer(
     gradient-based optimization across all :class:'Box' subspaces, starting from the best location
     found across a sample of `num_initial_samples` random points.
 
-    This optimizer supports Scipy's L-BFGS-B optimizer, which optimizes directly within and up to
-    the bounds of the search space.
+    This optimizer supports Scipy's L-BFGS-B optimizer (used via GPflow's Scipy optimizer wrapper),
+    which optimizes directly within and up to the bounds of the search space.
 
     For challenging acquisition function optimizations, we run `num_optimization_runs` separate
-    optimizations, each starting from one of the top  `num_optimization_runs` initial query points.
+    optimizations, each starting from one of the top `num_optimization_runs` initial query points.
 
     If all `num_optimization_runs` optimizations fail to converge then we run up to
     `num_recovery_runs` starting from random locations.
@@ -124,6 +144,9 @@ def generate_continuous_optimizer(
         the optimization.
     :param num_optimization_runs: The number of separate optimizations to run.
     :param num_recovery_runs: The maximum number of recovery optimization runs in case of failure.
+    :param optimizer_args: The keyword arguments to pass to the GPflow's Scipy optimizer wrapper.
+        Check `minimize` method  of :class:`~gpflow.optimizers.Scipy` for details what arguments
+        can be passed.
     :return: The acquisition optimizer.
     """
     if num_initial_samples <= 0:
@@ -171,13 +194,17 @@ def optimize_continuous(
         initial_points = tf.gather(trial_search_space, top_k_indicies)  # [num_optimization_runs, D]
 
         variable = tf.Variable(initial_points[0:1])  # [1, D]
+        optimizer_args_local = copy.deepcopy(optimizer_args)
 
         def _objective() -> TensorType:
             return -target_func(variable[:, None, :])  # [1]
 
-        def _perform_optimization(starting_point: TensorType, bounds: spo.Bounds) -> OptimizeResult:
+        def _perform_optimization(
+            starting_point: TensorType,
+            optimizer_args: dict[str, Any],
+        ) -> OptimizeResult:
             variable.assign(starting_point)  # [1, D]
-            return gpflow.optimizers.Scipy().minimize(_objective, (variable,), bounds=bounds)
+            return gpflow.optimizers.Scipy().minimize(_objective, (variable,), **optimizer_args)
 
         successful_optimization = False
         chosen_point = variable  # [1, D]
@@ -190,8 +217,8 @@ def _perform_optimization(starting_point: TensorType, bounds: spo.Bounds) -> Opt
                 bounds = get_bounds_of_box_relaxation_around_point(space, initial_points[i : i + 1])
             else:
                 bounds = spo.Bounds(space.lower, space.upper)
-
-            opt_result = _perform_optimization(initial_points[i : i + 1], bounds)
+            optimizer_args_local["bounds"] = bounds
+            opt_result = _perform_optimization(initial_points[i : i + 1], optimizer_args_local)
             if opt_result.success:
                 successful_optimization = True
 
@@ -209,7 +236,8 @@ def _perform_optimization(starting_point: TensorType, bounds: spo.Bounds) -> Opt
                     bounds = get_bounds_of_box_relaxation_around_point(space, random_start)
                 else:
                     bounds = spo.Bounds(space.lower, space.upper)
-                opt_result = _perform_optimization(random_start, bounds)
+                optimizer_args_local["bounds"] = bounds
+                opt_result = _perform_optimization(random_start, optimizer_args_local)
                 if opt_result.success:
                     chosen_point = variable  # [1, D]
                     successful_optimization = True


Test patch: