Refactor recommender

baybe/recommenders/pure/bayesian/base.py

@@ -20,6 +20,7 @@
 )
 from baybe.objectives.base import Objective
 from baybe.recommenders.pure.base import PureRecommender
+from baybe.recommenders.pure.bayesian.botorch.optimizers.base import OptimizerProtocol
 from baybe.searchspace import SearchSpace
 from baybe.settings import Settings
 from baybe.surrogates import GaussianProcessSurrogate
@@ -56,6 +57,12 @@ class BayesianRecommender(PureRecommender, ABC):
     )
     """The acquisition function. When omitted, a default is used."""
 
+    optimizer: OptimizerProtocol | None = field(
+        alias="optimizer",
+        default=None,
+    )
+    """The acquisition function optimizer."""
+
     # TODO: The objective is currently only required for validating the recommendation
     #   context. Once multi-target support is complete, we might want to refactor
     #   the validation mechanism, e.g. by

baybe/recommenders/pure/bayesian/botorch/continuous.py

@@ -16,7 +16,7 @@
 )
 from baybe.parameters.numerical import _FixedNumericalContinuousParameter
 from baybe.searchspace import SubspaceContinuous
-from baybe.utils.basic import flatten
+from baybe.searchspace.core import SearchSpace
 
 if TYPE_CHECKING:
     from torch import Tensor
@@ -147,9 +147,6 @@ def recommend_continuous_without_cardinality_constraints(
     Raises:
         ValueError: If the continuous search space has cardinality constraints.
     """
-    import torch
-    from botorch.optim import optimize_acqf
-
     if subspace_continuous.n_subsets > 0:
         raise ValueError(
             f"'{recommend_continuous_without_cardinality_constraints.__name__}' "
@@ -181,31 +178,10 @@ def recommend_continuous_without_cardinality_constraints(
     #   because it is unclear if the corresponding presence checks for these
     #   arguments is correctly implemented in all invoked BoTorch subroutines.
     #   For details: https://github.com/pytorch/botorch/issues/2042
-    points, acqf_values = optimize_acqf(
-        acq_function=recommender._botorch_acqf,
-        bounds=torch.from_numpy(
-            subspace_continuous.comp_rep_bounds.to_numpy(copy=True)
-        ),
-        q=batch_size,
-        num_restarts=recommender.n_restarts,
-        raw_samples=recommender.n_raw_samples,
-        fixed_features=fixed_parameters or None,
-        equality_constraints=flatten(
-            c.to_botorch(
-                subspace_continuous.parameters,
-                batch_size=batch_size if c.is_interpoint else None,
-            )
-            for c in subspace_continuous.constraints_lin_eq
-        )
-        or None,
-        inequality_constraints=flatten(
-            c.to_botorch(
-                subspace_continuous.parameters,
-                batch_size=batch_size if c.is_interpoint else None,
-            )
-            for c in subspace_continuous.constraints_lin_ineq
-        )
-        or None,
-        sequential=recommender.sequential_continuous,
+    points, acqf_values = recommender.optimizer(
+        batch_size=batch_size,
+        acquisition_function=recommender._botorch_acqf,
+        searchspace=SearchSpace(continuous=subspace_continuous),
+        fixed_parameters=fixed_parameters,
     )
     return points, acqf_values

baybe/recommenders/pure/bayesian/botorch/core.py

@@ -30,6 +30,7 @@
     recommend_hybrid_with_subsets,
     recommend_hybrid_without_subsets,
 )
+from baybe.recommenders.pure.bayesian.botorch.optimizers.basic import GradientOptimizer
 from baybe.searchspace import (
     SearchSpace,
     SearchSpaceType,
@@ -213,6 +214,13 @@ def _recommend_continuous(
                 f"acquisition functions for batch sizes > 1."
             )
 
+        if self.optimizer is None:
+            self.optimizer = GradientOptimizer(
+                sequential_continuous=self.sequential_continuous,
+                n_restarts=self.n_restarts,
+                n_raw_samples=self.n_raw_samples,
+            )
+
         points, _ = recommend_continuous_torch(self, subspace_continuous, batch_size)
 
         return pd.DataFrame(points, columns=subspace_continuous.parameter_names)

baybe/recommenders/pure/bayesian/botorch/optimizers/__init__.py

@@ -0,0 +1,7 @@
+"""Acquisition function optimizers."""
+
+from baybe.recommenders.pure.bayesian.botorch.optimizers.basic import GradientOptimizer
+
+__all__ = [
+    "GradientOptimizer",
+]

baybe/recommenders/pure/bayesian/botorch/optimizers/base.py

@@ -0,0 +1,40 @@
+"""Base protocol for all optimizers."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
+
+from baybe.acquisition.base import AcquisitionFunction
+from baybe.searchspace import SearchSpace
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+
+@runtime_checkable
+class OptimizerProtocol(Protocol):
+    """Type protocol specifying the interface optimizers need to implement."""
+
+    # Use slots so that derived classes also remain slotted
+    # See also: https://www.attrs.org/en/stable/glossary.html#term-slotted-classes
+    __slots__ = ()
+
+    def __call__(
+        self,
+        batch_size: int,
+        acquisition_function: AcquisitionFunction,
+        searchspace: SearchSpace,
+        fixed_parameters: dict[int, float] | None = None,
+    ) -> tuple[Tensor, Tensor]:
+        """Recommend a batch of points from the given search space.
+
+        Args:
+            batch_size: The size of the recommendation batch.
+            acquisition_function: The acquisition function to be optimized.
+            searchspace: The search space from which to generate recommendations.
+            fixed_parameters: A dictionary mapping parameter indices to fixed values.
+
+        Returns:
+            The recommendations and corresponding acquisition values.
+        """
+        ...

baybe/recommenders/pure/bayesian/botorch/optimizers/basic.py

@@ -0,0 +1,109 @@
+"""Low-level optimizers of acquisition functions."""
+
+from __future__ import annotations
+
+import gc
+from typing import TYPE_CHECKING
+
+from attrs import define, field
+from attrs.validators import gt, instance_of
+
+from baybe.recommenders.pure.bayesian.botorch.optimizers.base import OptimizerProtocol
+from baybe.searchspace import SearchSpace
+from baybe.utils.basic import flatten
+
+if TYPE_CHECKING:
+    from botorch.acquisition import AcquisitionFunction as BoAcquisitionFunction
+    from torch import Tensor
+
+
+@define(kw_only=True)
+class GradientOptimizer(OptimizerProtocol):
+    """Acquisition function optimizer using gradient-based optimization."""
+
+    n_restarts: int = field(validator=[instance_of(int), gt(0)], default=10)
+    """Number of times gradient-based optimization is restarted from different initial
+    points. **Does not affect purely discrete optimization**.
+    """
+
+    n_raw_samples: int = field(validator=[instance_of(int), gt(0)], default=64)
+    """Number of raw samples drawn for the initialization heuristic in gradient-based
+    optimization. **Does not affect purely discrete optimization**.
+    """
+
+    sequential_continuous: bool = field(default=True)
+    """Flag defining whether to apply sequential greedy or batch optimization in
+    **continuous** search spaces. In discrete/hybrid spaces, sequential greedy
+    optimization is applied automatically.
+    """
+
+    def __call__(
+        self,
+        batch_size: int,
+        acquisition_function: BoAcquisitionFunction,
+        searchspace: SearchSpace,
+        fixed_parameters: dict[int, float] | None = None,
+    ) -> tuple[Tensor, Tensor]:
+        """Recommend from a search space using gradient-based optimization.
+
+        Args:
+            batch_size: The size of the recommendation batch.
+            acquisition_function: The acquisition function to be optimized.
+            searchspace: The search space from which to generate recommendations.
+            fixed_parameters: A dictionary mapping parameter indices to fixed values.
+
+        Returns:
+            The recommendations and corresponding acquisition values.
+
+        Raises:
+            NotImplementedError: If the search space has a discrete component.
+            ValueError: If the search space has cardinality constraints.
+        """
+        import torch
+        from botorch.optim import optimize_acqf
+
+        if not searchspace.discrete.is_empty:
+            raise NotImplementedError(
+                "Gradient-based optimization is not implemented "
+                "for non-empty discrete search spaces."
+            )
+
+        if searchspace.continuous.n_subsets > 0:
+            raise ValueError(
+                f"'{self.__class__.__name__}' "
+                f"expects a continuous subspace without cardinality constraints."
+            )
+
+        points, acqf_values = optimize_acqf(
+            acq_function=acquisition_function,
+            bounds=torch.from_numpy(
+                searchspace.continuous.comp_rep_bounds.to_numpy(copy=True)
+            ),
+            q=batch_size,
+            num_restarts=self.n_restarts,
+            raw_samples=self.n_raw_samples,
+            fixed_features=fixed_parameters or None,
+            equality_constraints=flatten(
+                c.to_botorch(
+                    searchspace.continuous.parameters,
+                    batch_size=batch_size if c.is_interpoint else None,
+                )
+                for c in searchspace.continuous.constraints_lin_eq
+            )
+            or None,
+            inequality_constraints=flatten(
+                c.to_botorch(
+                    searchspace.continuous.parameters,
+                    batch_size=batch_size if c.is_interpoint else None,
+                )
+                for c in searchspace.continuous.constraints_lin_ineq
+            )
+            or None,
+            sequential=self.sequential_continuous,
+        )
+
+        return points, acqf_values
+
+
+# Collect leftover original slotted classes processed by `attrs.define`
+gc.collect()