Refactor botorch/sampling/pathwise and add support for product kernels

botorch/acquisition/knowledge_gradient.py

@@ -226,7 +226,7 @@ def evaluate(self, X: Tensor, bounds: Tensor, **kwargs: Any) -> Tensor:
             kwargs: Additional keyword arguments. This includes the options for
                 optimization of the inner problem, i.e. `num_restarts`, `raw_samples`,
                 an `options` dictionary to be passed on to the optimization helpers, and
-                a `scipy_options` dictionary to be passed to `scipy.optimize.minimize`.
+                a `scipy_options` dictionary to be passed to `scipy.minimize`.
 
         Returns:
             A Tensor of shape `b`. For t-batch b, the q-KG value of the design

botorch/generation/gen.py

@@ -72,10 +72,6 @@ def gen_candidates_scipy(
 
     Optimizes an acquisition function starting from a set of initial candidates
     using `scipy.optimize.minimize` via a numpy converter.
-    We use SLSQP, if constraints are present, and LBFGS-B otherwise.
-    As `scipy.optimize.minimize` does not support optimizating a batch of problems, we
-    treat optimizing a set of candidates as a single optimization problem by
-    summing together their acquisition values.
 
     Args:
         initial_conditions: Starting points for optimization, with shape
@@ -102,7 +98,7 @@ def gen_candidates_scipy(
             `optimize_acqf()`. The constraints will later be passed to the scipy
             solver.
         options: Options used to control the optimization including "method"
-            and "maxiter". Select method for `scipy.optimize.minimize` using the
+            and "maxiter". Select method for `scipy.minimize` using the
             "method" key. By default uses L-BFGS-B for box-constrained problems
             and SLSQP if inequality or equality constraints are present. If
             `with_grad=False`, then we use a two-point finite difference estimate
@@ -664,13 +660,13 @@ def _process_scipy_result(res: OptimizeResult, options: dict[str, Any]) -> None:
             or "Iteration limit reached" in res.message
         ):
             logger.info(
-                "`scipy.optimize.minimize` exited by reaching the iteration limit of "
+                "`scipy.minimize` exited by reaching the iteration limit of "
                 f"`maxiter: {options.get('maxiter')}`."
             )
         elif "EVALUATIONS EXCEEDS LIMIT" in res.message:
             logger.info(
-                "`scipy.optimize.minimize` exited by reaching the function evaluation "
-                f"limit of `maxfun: {options.get('maxfun')}`."
+                "`scipy.minimize` exited by reaching the function evaluation limit of "
+                f"`maxfun: {options.get('maxfun')}`."
             )
         elif "Optimization timed out after" in res.message:
             logger.info(res.message)

botorch/models/gpytorch.py

@@ -816,6 +816,7 @@ def _apply_noise(
         self,
         X: Tensor,
         mvn: MultivariateNormal,
+        num_outputs: int,
         observation_noise: bool | Tensor,
     ) -> MultivariateNormal:
         """Adds the observation noise to the posterior.
@@ -947,6 +948,7 @@ def posterior(
             mvn = self._apply_noise(
                 X=X_full,
                 mvn=mvn,
+                num_outputs=num_outputs,
                 observation_noise=observation_noise,
             )
         # If single-output, return the posterior of a single-output model

botorch/optim/core.py

@@ -79,8 +79,8 @@ def scipy_minimize(
         bounds: A dictionary mapping parameter names to lower and upper bounds.
         callback: A callable taking `parameters` and an OptimizationResult as arguments.
         x0: An optional initialization vector passed to scipy.optimize.minimize.
-        method: Solver type, passed along to scipy.optimize.minimize.
-        options: Dictionary of solver options, passed along to scipy.optimize.minimize.
+        method: Solver type, passed along to scipy.minimize.
+        options: Dictionary of solver options, passed along to scipy.minimize.
         timeout_sec: Timeout in seconds to wait before aborting the optimization loop
             if not converged (will return the best found solution thus far).
 

botorch/optim/fit.py

@@ -69,8 +69,8 @@ def fit_gpytorch_mll_scipy(
             Responsible for setting the `grad` attributes of `parameters`. If no closure
             is provided, one will be obtained by calling `get_loss_closure_with_grads`.
         closure_kwargs: Keyword arguments passed to `closure`.
-        method: Solver type, passed along to scipy.optimize.minimize.
-        options: Dictionary of solver options, passed along to scipy.optimize.minimize.
+        method: Solver type, passed along to scipy.minimize.
+        options: Dictionary of solver options, passed along to scipy.minimize.
         callback: Optional callback taking `parameters` and an OptimizationResult as its
             sole arguments.
         timeout_sec: Timeout in seconds after which to terminate the fitting loop

botorch/optim/optimize.py

@@ -539,29 +539,7 @@ def optimize_acqf(
     retry_on_optimization_warning: bool = True,
     **ic_gen_kwargs: Any,
 ) -> tuple[Tensor, Tensor]:
-    r"""Optimize the acquisition function for a single or multiple joint candidates.
-
-    A high-level description (missing exceptions for special setups):
-
-    This function optimizes the acquisition function `acq_function` in two steps:
-
-    i) It will sample `raw_samples` random points using Sobol sampling in the bounds
-    `bounds` and pass on the "best" `num_restarts` many.
-    The default way to find these "best" is via `gen_batch_initial_conditions`
-    (deviating for some acq functions, see `get_ic_generator`),
-    which by default performs Boltzmann sampling on the acquisition function value
-    (The behavior of step (i) can be further controlled by specifying `ic_generator`
-    or `batch_initial_conditions`.)
-
-    ii) A batch of the `num_restarts` points (or joint sets of points)
-    with the highest acquisition values in the previous step are then further
-    optimized. This is by default done by LBFGS-B optimization, if no constraints are
-    present, and SLSQP, if constraints are present (can be changed to
-    other optmizers via `gen_candidates`).
-
-    While the optimization procedure runs on CPU by default for this function,
-    the acq_function can be implemented on GPU and simply move the inputs
-    to GPU internally.
+    r"""Generate a set of candidates via multi-start optimization.
 
     Args:
         acq_function: An AcquisitionFunction.
@@ -570,13 +548,10 @@ def optimize_acqf(
             +inf, respectively).
         q: The number of candidates.
         num_restarts: The number of starting points for multistart acquisition
-            function optimization. Even though the name suggests this happens
-            sequentually, it is done in parallel (using batched evaluations)
-            for up to `options.batch_limit` candidates (by default completely parallel).
+            function optimization.
         raw_samples: The number of samples for initialization. This is required
             if `batch_initial_conditions` is not specified.
-        options: Options for both optimization, passed to `gen_candidates`,
-            and initialization, passed to the `ic_generator` via the `options` kwarg.
+        options: Options for candidate generation.
         inequality_constraints: A list of tuples (indices, coefficients, rhs),
             with each tuple encoding an inequality constraint of the form
             `\sum_i (X[indices[i]] * coefficients[i]) >= rhs`. `indices` and
@@ -622,11 +597,10 @@ def optimize_acqf(
             acquisition values) given a tensor of initial conditions and an
             acquisition function. Other common inputs include lower and upper bounds
             and a dictionary of options, but refer to the documentation of specific
-            generation functions (e.g., botorch.optim.optimize.gen_candidates_scipy
-            and botorch.generation.gen.gen_candidates_torch) for method-specific
-            inputs. Default: `gen_candidates_scipy`
+            generation functions (e.g gen_candidates_scipy and gen_candidates_torch)
+            for method-specific inputs. Default: `gen_candidates_scipy`
         sequential: If False, uses joint optimization, otherwise uses sequential
-            optimization for optimizing multiple joint candidates (q > 1).
+            optimization.
         ic_generator: Function for generating initial conditions. Not needed when
             `batch_initial_conditions` are provided. Defaults to
             `gen_one_shot_kg_initial_conditions` for `qKnowledgeGradient` acquisition

botorch/optim/parameter_constraints.py

@@ -91,7 +91,7 @@ def make_scipy_linear_constraints(
     Returns:
         A list of dictionaries containing callables for constraint function
         values and Jacobians and a string indicating the associated constraint
-        type ("eq", "ineq"), as expected by `scipy.optimize.minimize`.
+        type ("eq", "ineq"), as expected by `scipy.minimize`.
 
     This function assumes that constraints are the same for each input batch,
     and broadcasts the constraints accordingly to the input batch shape. This
@@ -222,7 +222,7 @@ def _make_linear_constraints(
     shapeX: torch.Size,
     eq: bool = False,
 ) -> list[ScipyConstraintDict]:
-    r"""Create linear constraints to be used by `scipy.optimize.minimize`.
+    r"""Create linear constraints to be used by `scipy.minimize`.
 
     Encodes constraints of the form
     `\sum_i (coefficients[i] * X[..., indices[i]]) ? rhs`
@@ -317,7 +317,7 @@ def _make_linear_constraints(
 def _make_nonlinear_constraints(
     f_np_wrapper: Callable, nlc: Callable, is_intrapoint: bool, shapeX: torch.Size
 ) -> list[ScipyConstraintDict]:
-    """Create nonlinear constraints to be used by `scipy.optimize.minimize`.
+    """Create nonlinear constraints to be used by `scipy.minimize`.
 
     Args:
         f_np_wrapper: A wrapper function that given a constraint evaluates
@@ -578,7 +578,7 @@ def make_scipy_nonlinear_inequality_constraints(
     Returns:
         A list of dictionaries containing callables for constraint function
         values and Jacobians and a string indicating the associated constraint
-        type ("eq", "ineq"), as expected by `scipy.optimize.minimize`.
+        type ("eq", "ineq"), as expected by `scipy.minimize`.
     """
 
     scipy_nonlinear_inequality_constraints = []

botorch/sampling/pathwise/__init__.py

@@ -6,9 +6,16 @@
 
 
 from botorch.sampling.pathwise.features import (
-    gen_kernel_features,
+    DirectSumFeatureMap,
+    FeatureMap,
+    FourierFeatureMap,
+    gen_kernel_feature_map,
+    IndexKernelFeatureMap,
     KernelEvaluationMap,
     KernelFeatureMap,
+    LinearKernelFeatureMap,
+    MultitaskKernelFeatureMap,
+    OuterProductFeatureMap,
 )
 from botorch.sampling.pathwise.paths import (
     GeneralizedLinearPath,
@@ -26,15 +33,22 @@
 
 
 __all__ = [
+    "DirectSumFeatureMap",
     "draw_matheron_paths",
     "draw_kernel_feature_paths",
-    "gen_kernel_features",
+    "FeatureMap",
+    "FourierFeatureMap",
+    "gen_kernel_feature_map",
     "get_matheron_path_model",
     "gaussian_update",
     "GeneralizedLinearPath",
+    "IndexKernelFeatureMap",
     "KernelEvaluationMap",
     "KernelFeatureMap",
+    "LinearKernelFeatureMap",
     "MatheronPath",
+    "MultitaskKernelFeatureMap",
+    "OuterProductFeatureMap",
     "SamplePath",
     "PathDict",
     "PathList",

botorch/sampling/pathwise/__init__.py,cover

@@ -0,0 +1,55 @@
+  #!/usr/bin/env python3
+  # Copyright (c) Meta Platforms, Inc. and affiliates.
+  #
+  # This source code is licensed under the MIT license found in the
+  # LICENSE file in the root directory of this source tree.
+
+
+> from botorch.sampling.pathwise.features import (
+>     DirectSumFeatureMap,
+>     FeatureMap,
+>     FourierFeatureMap,
+>     gen_kernel_feature_map,
+>     IndexKernelFeatureMap,
+>     KernelEvaluationMap,
+>     KernelFeatureMap,
+>     LinearKernelFeatureMap,
+>     MultitaskKernelFeatureMap,
+>     OuterProductFeatureMap,
+> )
+> from botorch.sampling.pathwise.paths import (
+>     GeneralizedLinearPath,
+>     PathDict,
+>     PathList,
+>     SamplePath,
+> )
+> from botorch.sampling.pathwise.posterior_samplers import (
+>     draw_matheron_paths,
+>     get_matheron_path_model,
+>     MatheronPath,
+> )
+> from botorch.sampling.pathwise.prior_samplers import draw_kernel_feature_paths
+> from botorch.sampling.pathwise.update_strategies import gaussian_update
+
+
+> __all__ = [
+>     "DirectSumFeatureMap",
+>     "draw_matheron_paths",
+>     "draw_kernel_feature_paths",
+>     "FeatureMap",
+>     "FourierFeatureMap",
+>     "gen_kernel_feature_map",
+>     "get_matheron_path_model",
+>     "gaussian_update",
+>     "GeneralizedLinearPath",
+>     "IndexKernelFeatureMap",
+>     "KernelEvaluationMap",
+>     "KernelFeatureMap",
+>     "LinearKernelFeatureMap",
+>     "MatheronPath",
+>     "MultitaskKernelFeatureMap",
+>     "OuterProductFeatureMap",
+>     "SamplePath",
+>     "PathDict",
+>     "PathList",
+> ]

botorch/sampling/pathwise/features/__init__.py

@@ -5,16 +5,28 @@
 # LICENSE file in the root directory of this source tree.
 
 
-from botorch.sampling.pathwise.features.generators import gen_kernel_features
+from botorch.sampling.pathwise.features.generators import gen_kernel_feature_map
 from botorch.sampling.pathwise.features.maps import (
+    DirectSumFeatureMap,
     FeatureMap,
+    FourierFeatureMap,
+    IndexKernelFeatureMap,
     KernelEvaluationMap,
     KernelFeatureMap,
+    LinearKernelFeatureMap,
+    MultitaskKernelFeatureMap,
+    OuterProductFeatureMap,
 )
 
 __all__ = [
+    "DirectSumFeatureMap",
     "FeatureMap",
-    "gen_kernel_features",
+    "FourierFeatureMap",
+    "gen_kernel_feature_map",
+    "IndexKernelFeatureMap",
     "KernelEvaluationMap",
     "KernelFeatureMap",
+    "LinearKernelFeatureMap",
+    "MultitaskKernelFeatureMap",
+    "OuterProductFeatureMap",
 ]

botorch/sampling/pathwise/features/__init__.py,cover

@@ -0,0 +1,32 @@
+  #!/usr/bin/env python3
+  # Copyright (c) Meta Platforms, Inc. and affiliates.
+  #
+  # This source code is licensed under the MIT license found in the
+  # LICENSE file in the root directory of this source tree.
+
+
+> from botorch.sampling.pathwise.features.generators import gen_kernel_feature_map
+> from botorch.sampling.pathwise.features.maps import (
+>     DirectSumFeatureMap,
+>     FeatureMap,
+>     FourierFeatureMap,
+>     IndexKernelFeatureMap,
+>     KernelEvaluationMap,
+>     KernelFeatureMap,
+>     LinearKernelFeatureMap,
+>     MultitaskKernelFeatureMap,
+>     OuterProductFeatureMap,
+> )
+
+> __all__ = [
+>     "DirectSumFeatureMap",
+>     "FeatureMap",
+>     "FourierFeatureMap",
+>     "gen_kernel_feature_map",
+>     "IndexKernelFeatureMap",
+>     "KernelEvaluationMap",
+>     "KernelFeatureMap",
+>     "LinearKernelFeatureMap",
+>     "MultitaskKernelFeatureMap",
+>     "OuterProductFeatureMap",
+> ]