diff --git a/baybe/recommenders/pure/bayesian/base.py b/baybe/recommenders/pure/bayesian/base.py index 4ac5c1eed2..55e652dc85 100644 --- a/baybe/recommenders/pure/bayesian/base.py +++ b/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 diff --git a/baybe/recommenders/pure/bayesian/botorch/continuous.py b/baybe/recommenders/pure/bayesian/botorch/continuous.py index a63aa77651..0fc7a9f7f3 100644 --- a/baybe/recommenders/pure/bayesian/botorch/continuous.py +++ b/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 diff --git a/baybe/recommenders/pure/bayesian/botorch/core.py b/baybe/recommenders/pure/bayesian/botorch/core.py index 7953d5ca74..d948f661c4 100644 --- a/baybe/recommenders/pure/bayesian/botorch/core.py +++ b/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) diff --git a/baybe/recommenders/pure/bayesian/botorch/optimizers/__init__.py b/baybe/recommenders/pure/bayesian/botorch/optimizers/__init__.py new file mode 100644 index 0000000000..eaf5773900 --- /dev/null +++ b/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", +] diff --git a/baybe/recommenders/pure/bayesian/botorch/optimizers/base.py b/baybe/recommenders/pure/bayesian/botorch/optimizers/base.py new file mode 100644 index 0000000000..2e801119b8 --- /dev/null +++ b/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. + """ + ... diff --git a/baybe/recommenders/pure/bayesian/botorch/optimizers/basic.py b/baybe/recommenders/pure/bayesian/botorch/optimizers/basic.py new file mode 100644 index 0000000000..9ff8a592a2 --- /dev/null +++ b/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()