From 9d28b4902c1e8af1288876b4cfa538c62a6176ab Mon Sep 17 00:00:00 2001 From: Di Jin Date: Fri, 12 Jul 2024 09:22:24 +0200 Subject: [PATCH] Refactor botorch+cardinality constraint --- baybe/recommenders/pure/bayesian/botorch.py | 306 ++++++++++++-------- baybe/searchspace/continuous.py | 4 +- 2 files changed, 182 insertions(+), 128 deletions(-) diff --git a/baybe/recommenders/pure/bayesian/botorch.py b/baybe/recommenders/pure/bayesian/botorch.py index 31f287e659..91a28ec1a6 100644 --- a/baybe/recommenders/pure/bayesian/botorch.py +++ b/baybe/recommenders/pure/bayesian/botorch.py @@ -1,11 +1,15 @@ """Botorch recommender.""" + +from __future__ import annotations + import math -from typing import Any, ClassVar +from typing import TYPE_CHECKING, Any, ClassVar import pandas as pd from attr.converters import optional from attrs import define, field +from baybe.constraints import ContinuousCardinalityConstraint from baybe.exceptions import NoMCAcquisitionFunctionError from baybe.recommenders.pure.bayesian.base import BayesianRecommender from baybe.searchspace import ( @@ -20,7 +24,14 @@ sample_numerical_df, ) -N_ITER_THRESHOLD = 10 +if TYPE_CHECKING: + from torch import Tensor + +N_THRESHOLD_INACTIVE_PARAMETERS_GENERATOR: int = 10 +"""This threshold controls which inactive parameters generator is chosen. There are +two mechanisms: +* Iterating the combinatorial list of all possible inactive parameters, +* Iterate a fixed number of randomly generated inactive parameter configurations.""" @define(kw_only=True) @@ -141,7 +152,6 @@ def _recommend_continuous( Raises: NoMCAcquisitionFunctionError: If a non-Monte Carlo acquisition function is used with a batch size > 1. - RuntimeError: If the combinatorial list of inactive parameters is None. Returns: A dataframe containing the recommendations as individual rows. @@ -153,138 +163,184 @@ def _recommend_continuous( f"acquisition functions for batch sizes > 1." ) + if len(subspace_continuous.constraints_cardinality): + points, _ = self._recommend_continuous_with_cardinality_constraints( + subspace_continuous, + batch_size, + ) + else: + points, _ = self._recommend_continuous_without_cardinality_constraints( + subspace_continuous, + batch_size, + ) + + # Return optimized points as dataframe + rec = pd.DataFrame(points, columns=subspace_continuous.param_names) + return rec + + def _recommend_continuous_with_cardinality_constraints( + self, + subspace_continuous: SubspaceContinuous, + batch_size: int, + ) -> tuple[Tensor, Tensor]: + """Recommend from a continuous search space with cardinality constraints. + + Args: + subspace_continuous: The continuous subspace from which to generate + recommendations. + batch_size: The size of the recommendation batch. + + Returns: + The recommendations. + The acquisition values. + + Raises: + RuntimeError: If the continuous search space has no cardinality constraint. + """ import torch - from botorch.optim import optimize_acqf - from torch import Tensor - - def _recommend_continuous_with_inactive_parameters( - _subspace_continuous: SubspaceContinuous, - inactive_parameters: tuple[str, ...] | None = None, - ) -> tuple[Tensor, Tensor]: - """Define a helper function that can deal with inactive parameters.""" - if _subspace_continuous.constraints_cardinality: - # When there are cardinality constraints present. - if inactive_parameters is None: - # When no parameters are constrained to zeros - inactive_parameters = () - fixed_parameters = None - else: - # When certain parameters are constrained to zeros. - - # Cast the inactive parameters to the format of fixed features used - # in optimize_acqf()) - indices_inactive_params = [ - _subspace_continuous.param_names.index(key) - for key in _subspace_continuous.param_names - if key in inactive_parameters - ] - fixed_parameters = {ind: 0.0 for ind in indices_inactive_params} - - # Create a new subspace by ensuring all active parameters are non-zeros - _subspace_continuous = _subspace_continuous._ensure_nonzero_parameters( - inactive_parameters - ) - else: - # When there is no cardinality constraint - fixed_parameters = None - - _points, _acqf_values = optimize_acqf( - acq_function=self._botorch_acqf, - bounds=torch.from_numpy(_subspace_continuous.param_bounds_comp), - q=batch_size, - num_restarts=5, # TODO make choice for num_restarts - raw_samples=10, # TODO make choice for raw_samples - fixed_features=fixed_parameters, - equality_constraints=[ - c.to_botorch(_subspace_continuous.parameters) - for c in _subspace_continuous.constraints_lin_eq - ] - or None, # TODO: https://github.com/pytorch/botorch/issues/2042 - inequality_constraints=[ - c.to_botorch(_subspace_continuous.parameters) - for c in _subspace_continuous.constraints_lin_ineq - ] - or None, # TODO: https://github.com/pytorch/botorch/issues/2042 - sequential=self.sequential_continuous, + + if not subspace_continuous.constraints_cardinality: + raise RuntimeError( + f"This method expects a subspace object with constraints of type " + f"{ContinuousCardinalityConstraint.__name__}. For a subspace object " + f"without constraints of type" + f" {ContinuousCardinalityConstraint.__name__}, " + f"{self._recommend_continuous_without_cardinality_constraints.__name__}." # noqa ) - return _points, _acqf_values - if len(subspace_continuous.constraints_cardinality): - acqf_values_all: list[Tensor] = [] - points_all: list[Tensor] = [] - - # The key steps of handling cardinality constraint are - # * Determine several configurations of inactive parameters based on the - # cardinality constraints. - # * Optimize the acquisition function for different configurations and - # pick the best one. - # There are two mechanisms for the inactive parameter configurations. The - # full list of different inactive parameter configurations is used, - # when its size is not too large; otherwise we randomly pick a - # fixed number of inactive parameter configurations. - - if ( - subspace_continuous.n_combinatorial_inactive_parameters - > N_ITER_THRESHOLD - ): - # When the size of full list is too large, randomly set some - # parameters inactive. - for _ in range(N_ITER_THRESHOLD): - inactive_params_sample = ( - subspace_continuous._sample_inactive_parameters(1)[0] - ) - - ( - points_i, - acqf_values_i, - ) = _recommend_continuous_with_inactive_parameters( - subspace_continuous, - tuple(inactive_params_sample), - ) - - points_all.append(points_i.unsqueeze(0)) - acqf_values_all.append(acqf_values_i.unsqueeze(0)) - - elif subspace_continuous.combinatorial_inactive_parameters is not None: - # When the size of full list is not too large, iterate the combinations - # of all possible inactive parameters. - for ( - inactive_params_generator - ) in subspace_continuous.combinatorial_inactive_parameters: - # flatten inactive parameters - inactive_params_sample = { + acqf_values_all: list[Tensor] = [] + points_all: list[Tensor] = [] + + def append_recommendation_for_inactive_parameters_setting( + inactive_parameters: tuple[str, ...], + ): + """Append the recommendation for each inactive parameter configuration. + + Args: + inactive_parameters: A list of inactive parameters. + """ + # Create a new subspace by ensuring all active parameters being + # non-zeros. + subspace_continuous_with_active_params = ( + subspace_continuous._ensure_nonzero_parameters(inactive_parameters) + ) + # Optimize the acquisition function + ( + points_i, + acqf_values_i, + ) = self._recommend_continuous_without_cardinality_constraints( + subspace_continuous_with_active_params, + batch_size, + inactive_parameters, + ) + # Append recommendation list and acquisition function values + points_all.append(points_i.unsqueeze(0)) + acqf_values_all.append(acqf_values_i.unsqueeze(0)) + + # Below we start recommendation + if ( + subspace_continuous.n_combinatorial_inactive_parameters + > N_THRESHOLD_INACTIVE_PARAMETERS_GENERATOR + ): + # When the combinatorial list is too large, randomly set some parameters + # inactive. + for _ in range(N_THRESHOLD_INACTIVE_PARAMETERS_GENERATOR): + inactive_params_sample = tuple( + subspace_continuous._sample_inactive_parameters(1)[0] + ) + append_recommendation_for_inactive_parameters_setting( + inactive_params_sample + ) + else: + # When the combinatorial list is not too large, iterate the combinatorial + # list of all possible inactive parameters. + for ( + inactive_params_generator + ) in subspace_continuous.combinatorial_inactive_parameters: + # Flatten inactive parameter generator + inactive_params_sample = tuple( + { param for sublist in inactive_params_generator for param in sublist } - - ( - points_i, - acqf_values_i, - ) = _recommend_continuous_with_inactive_parameters( - subspace_continuous, - tuple(inactive_params_sample), - ) - - points_all.append(points_i.unsqueeze(0)) - acqf_values_all.append(acqf_values_i.unsqueeze(0)) - else: - raise RuntimeError( - f"The attribute" - f"{SubspaceContinuous.combinatorial_inactive_parameters.__name__}" - f"should not be None." ) - # Find the best option - points = torch.cat(points_all)[torch.argmax(torch.cat(acqf_values_all)), :] - else: - # When there is no cardinality constraint - points, _ = _recommend_continuous_with_inactive_parameters( - subspace_continuous + append_recommendation_for_inactive_parameters_setting( + inactive_params_sample + ) + + # Find the best option + points = torch.cat(points_all)[torch.argmax(torch.cat(acqf_values_all)), :] + acqf_values = torch.max(torch.cat(acqf_values_all)) + return points, acqf_values + + def _recommend_continuous_without_cardinality_constraints( + self, + subspace_continuous: SubspaceContinuous, + batch_size: int, + inactive_parameters: tuple[str, ...] | None = None, + ) -> tuple[Tensor, Tensor]: + """Recommend from a continuous search space without cardinality constraints. + + Args: + subspace_continuous: The continuous subspace from which to generate + recommendations. + batch_size: The size of the recommendation batch. + inactive_parameters: A list of inactive parameters. + + Returns: + The recommendations. + The acquisition values. + + Raises: + RuntimeError: If the continuous search space has any cardinality + constraints. + """ + import torch + from botorch.optim import optimize_acqf + + if subspace_continuous.constraints_cardinality: + raise RuntimeError( + f"This method expects only subspace object without constraints of type " + f"{ContinuousCardinalityConstraint.__name__}. For a subspace object " + f"with constraints of type {ContinuousCardinalityConstraint.__name__}, " + f"try method {self._recommend_continuous.__name__}." ) - # Return optimized points as dataframe - rec = pd.DataFrame(points, columns=subspace_continuous.param_names) - return rec + if not inactive_parameters: + fixed_parameters = None + else: + # Cast the inactive parameters to the format of fixed features used + # in optimize_acqf()) + indices_inactive_params = [ + subspace_continuous.param_names.index(key) + for key in subspace_continuous.param_names + if key in inactive_parameters + ] + fixed_parameters = {ind: 0.0 for ind in indices_inactive_params} + + points, acqf_values = optimize_acqf( + acq_function=self._botorch_acqf, + bounds=torch.from_numpy(subspace_continuous.param_bounds_comp), + q=batch_size, + num_restarts=5, # TODO make choice for num_restarts + raw_samples=10, # TODO make choice for raw_samples + fixed_features=fixed_parameters, + equality_constraints=[ + c.to_botorch(subspace_continuous.parameters) + for c in subspace_continuous.constraints_lin_eq + ] + or None, + # TODO: https://github.com/pytorch/botorch/issues/2042 + inequality_constraints=[ + c.to_botorch(subspace_continuous.parameters) + for c in subspace_continuous.constraints_lin_ineq + ] + or None, + # TODO: https://github.com/pytorch/botorch/issues/2042 + sequential=self.sequential_continuous, + ) + return points, acqf_values def _recommend_hybrid( self, diff --git a/baybe/searchspace/continuous.py b/baybe/searchspace/continuous.py index 8ed977abaa..a2a197ef8d 100644 --- a/baybe/searchspace/continuous.py +++ b/baybe/searchspace/continuous.py @@ -131,12 +131,10 @@ def n_combinatorial_inactive_parameters(self) -> int: @property def combinatorial_inactive_parameters( self - ) -> Iterable[tuple[tuple[str, ...], ...]] | None: + ) -> Iterable[tuple[tuple[str, ...], ...]]: """Combinatorial list of inactive parameters on subspace.""" # The comments on the difference in `n_combinatorial_inactive_parameters` # applies here as well. - if not self.constraints_cardinality: - return None return product( *[