pass gen_candidates callable in optimize_acqf

Summary: see title. This will support using stochastic optimization

Differential Revision: D41629164

fbshipit-source-id: 8c6e6f46f7f605ba5bb18291880915fe23dc2f1a

Author: sdaulton

botorch/generation/gen.py

@@ -36,6 +36,8 @@
 
 logger = _get_logger()
 
+TGenCandidates = Callable[[Tensor, AcquisitionFunction, Any], Tuple[Tensor, Tensor]]
+
 
 def gen_candidates_scipy(
     initial_conditions: Tensor,
@@ -47,6 +49,7 @@ def gen_candidates_scipy(
     nonlinear_inequality_constraints: Optional[List[Callable]] = None,
     options: Optional[Dict[str, Any]] = None,
     fixed_features: Optional[Dict[int, Optional[float]]] = None,
+    **kwargs,
 ) -> Tuple[Tensor, Tensor]:
     r"""Generate a set of candidates using `scipy.optimize.minimize`.
 
@@ -253,6 +256,7 @@ def gen_candidates_torch(
     options: Optional[Dict[str, Union[float, str]]] = None,
     callback: Optional[Callable[[int, Tensor, Tensor], NoReturn]] = None,
     fixed_features: Optional[Dict[int, Optional[float]]] = None,
+    **kwargs,
 ) -> Tuple[Tensor, Tensor]:
     r"""Generate a set of candidates using a `torch.optim` optimizer.
 

botorch/optim/optimize.py

@@ -22,7 +22,7 @@
 from botorch.acquisition.knowledge_gradient import qKnowledgeGradient
 from botorch.exceptions import InputDataError, UnsupportedError
 from botorch.exceptions.warnings import OptimizationWarning
-from botorch.generation.gen import gen_candidates_scipy
+from botorch.generation.gen import gen_candidates_scipy, TGenCandidates
 from botorch.logging import logger
 from botorch.optim.initializers import (
     gen_batch_initial_conditions,
@@ -64,6 +64,7 @@ def optimize_acqf(
     post_processing_func: Optional[Callable[[Tensor], Tensor]] = None,
     batch_initial_conditions: Optional[Tensor] = None,
     return_best_only: bool = True,
+    gen_candidates: TGenCandidates = gen_candidates_scipy,
     sequential: bool = False,
     **kwargs: Any,
 ) -> Tuple[Tensor, Tensor]:
@@ -103,6 +104,8 @@ def optimize_acqf(
             this if you do not want to use default initialization strategy.
         return_best_only: If False, outputs the solutions corresponding to all
             random restart initializations of the optimization.
+        gen_candidates: A callable for generating candidates given initial
+            conditions. Default: `gen_candidates_scipy`
         sequential: If False, uses joint optimization, otherwise uses sequential
             optimization.
         kwargs: Additonal keyword arguments.
@@ -258,23 +261,23 @@ def _optimize_batch_candidates() -> Tuple[Tensor, Tensor, List[Warning]]:
         batched_ics = batch_initial_conditions.split(batch_limit)
         opt_warnings = []
 
-        scipy_kws = dict(
-            acquisition_function=acq_function,
-            lower_bounds=None if bounds[0].isinf().all() else bounds[0],
-            upper_bounds=None if bounds[1].isinf().all() else bounds[1],
-            options={k: v for k, v in options.items() if k not in INIT_OPTION_KEYS},
-            inequality_constraints=inequality_constraints,
-            equality_constraints=equality_constraints,
-            nonlinear_inequality_constraints=nonlinear_inequality_constraints,
-            fixed_features=fixed_features,
-        )
+        gen_kwargs = {
+            "acquisition_function": acq_function,
+            "lower_bounds": None if bounds[0].isinf().all() else bounds[0],
+            "upper_bounds": None if bounds[1].isinf().all() else bounds[1],
+            "options": {k: v for k, v in options.items() if k not in INIT_OPTION_KEYS},
+            "fixed_features": fixed_features,
+            "inequality_constraints": inequality_constraints,
+            "equality_constraints": equality_constraints,
+            "nonlinear_inequality_constraints": nonlinear_inequality_constraints,
+        }
 
         for i, batched_ics_ in enumerate(batched_ics):
             # optimize using random restart optimization
             with warnings.catch_warnings(record=True) as ws:
                 warnings.simplefilter("always", category=OptimizationWarning)
-                batch_candidates_curr, batch_acq_values_curr = gen_candidates_scipy(
-                    initial_conditions=batched_ics_, **scipy_kws
+                batch_candidates_curr, batch_acq_values_curr = gen_candidates(
+                    initial_conditions=batched_ics_, **gen_kwargs
                 )
             opt_warnings += ws
             batch_candidates_list.append(batch_candidates_curr)