Skip to content

Commit

Permalink
Refactor botorch+cardinality constraint
Browse files Browse the repository at this point in the history
  • Loading branch information
@Waschenbacher
Waschenbacher committed Jul 12, 2024
1 parent 153e25b commit 9d28b49
Show file tree
Hide file tree
Showing 2 changed files with 182 additions and 128 deletions.
306 changes: 181 additions & 125 deletions baybe/recommenders/pure/bayesian/botorch.py
View file
Original file line number Diff line number Diff line change
@@ -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 (
Expand All @@ -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)
Expand Down Expand Up @@ -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.
Expand All @@ -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,
Expand Down
4 changes: 1 addition & 3 deletions baybe/searchspace/continuous.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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(
*[
Expand Down

0 comments on commit 9d28b49

Please sign in to comment.