diff --git a/baybe/kernels/__init__.py b/baybe/kernels/__init__.py index 9323a2b631..4e50c5728c 100644 --- a/baybe/kernels/__init__.py +++ b/baybe/kernels/__init__.py @@ -5,6 +5,7 @@ """ from baybe.kernels.basic import ( + IndexKernel, LinearKernel, MaternKernel, PeriodicKernel, @@ -18,6 +19,7 @@ __all__ = [ "AdditiveKernel", + "IndexKernel", "LinearKernel", "MaternKernel", "PeriodicKernel", diff --git a/baybe/parameters/fidelity.py b/baybe/parameters/fidelity.py index 15235c9c78..14ecf1a51a 100644 --- a/baybe/parameters/fidelity.py +++ b/baybe/parameters/fidelity.py @@ -21,8 +21,8 @@ validate_is_finite, validate_unique_values, ) +from baybe.settings import active_settings from baybe.utils.conversion import nonstring_to_tuple -from baybe.utils.numerical import DTypeFloatNumpy def _convert_zeta( @@ -107,7 +107,9 @@ def values(self) -> tuple[str | bool, ...]: @cached_property def comp_df(self) -> pd.DataFrame: return pd.DataFrame( - range(len(self.values)), dtype=DTypeFloatNumpy, columns=[self.name] + range(len(self.values)), + dtype=active_settings.DTypeFloatNumpy, + columns=[self.name], ) @@ -159,5 +161,7 @@ def values(self) -> tuple[float, ...]: @cached_property def comp_df(self) -> pd.DataFrame: return pd.DataFrame( - {self.name: self.values}, index=self.values, dtype=DTypeFloatNumpy + {self.name: self.values}, + index=self.values, + dtype=active_settings.DTypeFloatNumpy, ) diff --git a/baybe/searchspace/core.py b/baybe/searchspace/core.py index 5510af704f..b0a979810b 100644 --- a/baybe/searchspace/core.py +++ b/baybe/searchspace/core.py @@ -15,6 +15,10 @@ from baybe.constraints.base import Constraint from baybe.parameters import TaskParameter from baybe.parameters.base import Parameter +from baybe.parameters.fidelity import ( + CategoricalFidelityParameter, + NumericalDiscreteFidelityParameter, +) from baybe.searchspace.continuous import SubspaceContinuous from baybe.searchspace.discrete import ( MemorySize, @@ -48,6 +52,29 @@ class SearchSpaceType(Enum): """Flag for hybrid search spaces resp. compatibility with hybrid search spaces.""" +class SearchSpaceTaskType(Enum): + """Enum class for different types of task and/or fidelity subspaces.""" + + SINGLETASK = "SINGLETASK" + """Flag for search spaces with a single task, meaning no task parameter.""" + + CATEGORICALMULTITASK = "CATEGORICALMULTITASK" + """Flag for search spaces with a categorical task parameter.""" + + +class SearchSpaceFidelityType(Enum): + """Enum class for different types of task and/or fidelity subspaces.""" + + SINGLEFIDELITY = "SINGLEFIDELITY" + """Flag for search spaces with a single fidelity, meaning no fidelity parameter.""" + + NUMERICALDISCRETEMULTIFIDELITY = "NUMERICALDISCRETEMULTIFIDELITY" + """Flag for search spaces with a discrete numerical (ordered) fidelity parameter.""" + + CATEGORICALMULTIFIDELITY = "CATEGORICALMULTIFIDELITY" + """Flag for search spaces with a categorical (unordered) fidelity parameter.""" + + @define class SearchSpace(SerialMixin): """Class for managing the overall search space. @@ -258,15 +285,32 @@ def _task_parameter(self) -> TaskParameter | None: if not params: return None - assert len(params) == 1 # currently ensured by parameter validation step + return params[0] + + @property + def _fidelity_parameter( + self, + ) -> NumericalDiscreteFidelityParameter | CategoricalFidelityParameter | None: + """The (single) fidelity parameter of the space, if it exists.""" + # Currently private, see comment above + fidelity_parameters = ( + NumericalDiscreteFidelityParameter, + CategoricalFidelityParameter, + ) + + params = [p for p in self.parameters if isinstance(p, fidelity_parameters)] + + if not params: + return None + return params[0] @property def task_idx(self) -> int | None: - """The column index of the task parameter in computational representation.""" + """Column index of the task parameter in computational representation.""" if (task_param := self._task_parameter) is None: return None - # TODO[11611]: The current approach has three limitations: + # TODO [11611]: The current approach has three limitations: # 1. It matches by column name and thus assumes that the parameter name # is used as the column name. # 2. It relies on the current implementation detail that discrete parameters @@ -275,6 +319,14 @@ def task_idx(self) -> int | None: # --> Fix this when refactoring the data return cast(int, self.discrete.comp_rep.columns.get_loc(task_param.name)) + @property + def fidelity_idx(self) -> int | None: + """Column index of the fidelity parameter in computational representation.""" + if (fidelity_param := self._fidelity_parameter) is None: + return None + # See TODO [11611] above + return cast(int, self.discrete.comp_rep.columns.get_loc(fidelity_param.name)) + @property def n_tasks(self) -> int: """The number of tasks encoded in the search space.""" @@ -287,6 +339,54 @@ def n_tasks(self) -> int: return 1 return len(task_param.values) + @property + def n_fidelities(self) -> int: + """The number of fidelities encoded in the search space.""" + # See TODO [16932] above + if (fidelity_param := self._fidelity_parameter) is None: + # When there are no task parameters, we effectively have a single task + return 1 + return len(fidelity_param.values) + + @property + def task_type(self) -> SearchSpaceTaskType: + """Return the task type of the search space.""" + task_parameters = (p for p in self.parameters if isinstance(p, TaskParameter)) + + if len(task_parameters) == 0: + return SearchSpaceTaskType.SINGLETASK + elif len(task_parameters) == 1: + return SearchSpaceTaskType.CATEGORICALMULTITASK + else: + raise NotImplementedError( + "BayBE does not currently support search" + "spaces with multiple task parameters." + ) + + def fidelity_type(self) -> SearchSpaceFidelityType: + """Return the fidelity type of the search space.""" + fidelity_parameters = ( + CategoricalFidelityParameter, + NumericalDiscreteFidelityParameter, + ) + + fidelity_parameters = ( + p for p in self.parameters if isinstance(p, fidelity_parameters) + ) + + if len(fidelity_parameters) == 0: + return SearchSpaceFidelityType.SINGLEFIDELITY + elif len(fidelity_parameters) == 1: + if isinstance(fidelity_parameters[0], CategoricalFidelityParameter): + return SearchSpaceFidelityType.CATEGORICALMULTIFIDELITY + if isinstance(fidelity_parameters[0], NumericalDiscreteFidelityParameter): + return SearchSpaceFidelityType.NUMERICALDISCRETEMULTIFIDELITY + else: + raise NotImplementedError( + "BayBE does not currently support search" + "spaces with multiple fidelity parameters." + ) + def get_comp_rep_parameter_indices(self, name: str, /) -> tuple[int, ...]: """Find a parameter's column indices in the computational representation. @@ -386,7 +486,7 @@ def transform( @property def constraints_augmentable(self) -> tuple[Constraint, ...]: - """The searchspace constraints that can be considered during augmentation.""" + """The search space constraints that can be considered during augmentation.""" return tuple(c for c in self.constraints if c.eval_during_augmentation) def get_parameters_by_name(self, names: Sequence[str]) -> tuple[Parameter, ...]: diff --git a/baybe/surrogates/bandit.py b/baybe/surrogates/bandit.py index ad6563cc43..f003f44c64 100644 --- a/baybe/surrogates/bandit.py +++ b/baybe/surrogates/bandit.py @@ -3,7 +3,7 @@ from __future__ import annotations import gc -from typing import TYPE_CHECKING, Any, ClassVar +from typing import TYPE_CHECKING, Any from attrs import define, field from typing_extensions import override @@ -29,9 +29,6 @@ class BetaBernoulliMultiArmedBanditSurrogate(Surrogate): """A multi-armed bandit model with Bernoulli likelihood and beta prior.""" - supports_transfer_learning: ClassVar[bool] = False - # See base class. - prior: BetaPrior = field(factory=lambda: BetaPrior(1, 1)) """The beta prior for the win rates of the bandit arms. Uniform by default.""" diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py index 205e32f703..fc9d4755a6 100644 --- a/baybe/surrogates/base.py +++ b/baybe/surrogates/base.py @@ -82,10 +82,14 @@ def to_botorch(self) -> Model: class Surrogate(ABC, SurrogateProtocol, SerialMixin): """Abstract base class for all surrogate models.""" - supports_transfer_learning: ClassVar[bool] + supports_transfer_learning: ClassVar[bool] = False """Class variable encoding whether or not the surrogate supports transfer learning.""" + supports_multi_fidelity: ClassVar[bool] = False + """Class variable encoding whether or not the surrogate supports multi fidelity + Bayesian optimization.""" + supports_multi_output: ClassVar[bool] = False """Class variable encoding whether or not the surrogate is multi-output compatible.""" @@ -428,6 +432,14 @@ def fit( f"support transfer learning." ) + # Check if multi fidelity capabilities are needed + if (searchspace.n_fidelities > 1) and (not self.supports_multi_fidelity): + raise ValueError( + f"The search space contains fidelity parameters but the selected " + f"surrogate model type ({self.__class__.__name__}) does not " + f"support multi fidelity Bayesian optimisation." + ) + # Block partial measurements handle_missing_values(measurements, [t.name for t in objective.targets]) @@ -472,6 +484,11 @@ def __str__(self) -> str: self.supports_transfer_learning, single_line=True, ), + to_string( + "Supports Multi Fidelity", + self.supports_multi_fidelity, + single_line=True, + ), ] return to_string(self.__class__.__name__, *fields) diff --git a/baybe/surrogates/custom.py b/baybe/surrogates/custom.py index 79c4c6ea1c..00e346319d 100644 --- a/baybe/surrogates/custom.py +++ b/baybe/surrogates/custom.py @@ -11,7 +11,7 @@ from __future__ import annotations import gc -from typing import TYPE_CHECKING, Any, ClassVar, NoReturn +from typing import TYPE_CHECKING, Any, NoReturn import cattrs from attrs import define, field, validators @@ -67,9 +67,6 @@ class CustomONNXSurrogate(IndependentGaussianSurrogate): Note that these surrogates cannot be retrained. """ - supports_transfer_learning: ClassVar[bool] = False - # See base class. - onnx_input_name: str = field(validator=validators.instance_of(str)) """The input name used for constructing the ONNX str.""" diff --git a/baybe/surrogates/gaussian_process/components/kernel.py b/baybe/surrogates/gaussian_process/components/kernel.py index 8b6ddff280..f69130dab8 100644 --- a/baybe/surrogates/gaussian_process/components/kernel.py +++ b/baybe/surrogates/gaussian_process/components/kernel.py @@ -16,6 +16,7 @@ from baybe.kernels.base import Kernel from baybe.parameters.categorical import TaskParameter from baybe.parameters.enum import _ParameterKind +from baybe.parameters.fidelity import CategoricalFidelityParameter from baybe.parameters.selectors import ( ParameterSelectorProtocol, TypeSelector, @@ -157,7 +158,15 @@ def _default_base_kernel_factory(self) -> KernelFactoryProtocol: BayBENumericalKernelFactory, ) - return BayBENumericalKernelFactory(TypeSelector((TaskParameter,), exclude=True)) + return BayBENumericalKernelFactory( + TypeSelector( + ( + TaskParameter, + CategoricalFidelityParameter, + ), + exclude=True, + ) + ) @task_kernel_factory.default def _default_task_kernel_factory(self) -> KernelFactoryProtocol: @@ -165,7 +174,14 @@ def _default_task_kernel_factory(self) -> KernelFactoryProtocol: BayBETaskKernelFactory, ) - return BayBETaskKernelFactory(TypeSelector((TaskParameter,))) + return BayBETaskKernelFactory( + TypeSelector( + ( + TaskParameter, + CategoricalFidelityParameter, + ) + ) + ) @override def __call__( diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py index 2b1a3f361e..a5e43b145c 100644 --- a/baybe/surrogates/gaussian_process/core.py +++ b/baybe/surrogates/gaussian_process/core.py @@ -10,14 +10,13 @@ from attrs import Converter, define, field from attrs.converters import pipe -from attrs.validators import instance_of, is_callable +from attrs.validators import is_callable from typing_extensions import Self, override from baybe.exceptions import DeprecationError from baybe.kernels.base import Kernel from baybe.parameters.base import Parameter from baybe.parameters.categorical import TaskParameter -from baybe.searchspace.core import SearchSpace from baybe.surrogates.base import Surrogate from baybe.surrogates.gaussian_process.components.generic import ( GPComponentType, @@ -39,6 +38,7 @@ BayBELikelihoodFactory, BayBEMeanFactory, ) +from baybe.surrogates.gaussian_process.utils import _ModelContext from baybe.utils.boolean import strtobool from baybe.utils.conversion import to_string @@ -53,51 +53,6 @@ from torch import Tensor -@define -class _ModelContext: - """Model context for :class:`GaussianProcessSurrogate`.""" - - searchspace: SearchSpace = field(validator=instance_of(SearchSpace)) - """The search space the model is trained on.""" - - @property - def task_idx(self) -> int | None: - """The computational column index of the task parameter, if available.""" - return self.searchspace.task_idx - - @property - def is_multitask(self) -> bool: - """Indicates if model is to be operated in a multi-task context.""" - return self.n_task_dimensions > 0 - - @property - def n_task_dimensions(self) -> int: - """The number of task dimensions.""" - # TODO: Generalize to multiple task parameters - return 1 if self.task_idx is not None else 0 - - @property - def n_tasks(self) -> int: - """The number of tasks.""" - return self.searchspace.n_tasks - - @property - def parameter_bounds(self) -> Tensor: - """Get the search space parameter bounds in BoTorch Format.""" - import torch - - return torch.from_numpy(self.searchspace.scaling_bounds.values) - - @property - def numerical_indices(self) -> list[int]: - """The indices of the regular numerical model inputs.""" - return [ - i - for i in range(len(self.searchspace.comp_rep_columns)) - if i != self.task_idx - ] - - def _mark_custom_kernel( value: Kernel | KernelFactoryProtocol, self: GaussianProcessSurrogate ) -> Kernel | KernelFactoryProtocol: @@ -130,6 +85,9 @@ class GaussianProcessSurrogate(Surrogate): supports_transfer_learning: ClassVar[bool] = True # See base class. + supports_multi_fidelity = True + # See base class. + _custom_kernel: bool = field(init=False, default=False, repr=False, eq=False) # For deprecation only! diff --git a/baybe/surrogates/gaussian_process/multi_fidelity.py b/baybe/surrogates/gaussian_process/multi_fidelity.py new file mode 100644 index 0000000000..6e5fa3cdef --- /dev/null +++ b/baybe/surrogates/gaussian_process/multi_fidelity.py @@ -0,0 +1,101 @@ +"""Multi-fidelity Gaussian process surrogates.""" + +from __future__ import annotations + +import gc +from typing import TYPE_CHECKING, ClassVar + +from attrs import define, field +from typing_extensions import override + +from baybe.parameters.base import Parameter +from baybe.surrogates.base import Surrogate +from baybe.surrogates.gaussian_process.utils import _ModelContext + +if TYPE_CHECKING: + from botorch.models.gpytorch import GPyTorchModel + from botorch.models.transforms.input import InputTransform + from botorch.models.transforms.outcome import OutcomeTransform + from botorch.posteriors import Posterior + from torch import Tensor + + +@define +class GaussianProcessSurrogateSTMF(Surrogate): + """Botorch's single task multi fidelity Gaussian process.""" + + supports_multi_fidelity: ClassVar[bool] = True + # See base class. + + _model = field(init=False, default=None, eq=False) + """The actual model.""" + + @override + def to_botorch(self) -> GPyTorchModel: + return self._model + + @override + @staticmethod + def _make_parameter_scaler_factory( + parameter: Parameter, + ) -> type[InputTransform] | None: + return None + + @override + @staticmethod + def _make_target_scaler_factory() -> type[OutcomeTransform] | None: + # For GPs, we let botorch handle the scaling. See [Scaling Workaround] above. + return None + + @override + def _posterior(self, candidates_comp_scaled: Tensor, /) -> Posterior: + return self._model.posterior(candidates_comp_scaled) + + @override + def _fit(self, train_x: Tensor, train_y: Tensor) -> None: + import botorch + import gpytorch + + assert self._searchspace is not None + + context = _ModelContext(self._searchspace) + + assert context.n_fidelity_dimensions > 0, ( + f"{self.__class__.__name__} can only be fit on multi fidelity searchspaces." + ) + + # For GPs, we let botorch handle the scaling. See [Scaling Workaround] above. + input_transform = botorch.models.transforms.Normalize( # type: ignore[attr-defined] + train_x.shape[-1], + bounds=context.parameter_bounds, + indices=context.numerical_indices, + ) + outcome_transform = botorch.models.transforms.Standardize(train_y.shape[-1]) # type: ignore[attr-defined] + + # construct and fit the Gaussian process + self._model = botorch.models.SingleTaskMultiFidelityGP( + train_x, + train_y, + input_transform=input_transform, + outcome_transform=outcome_transform, + data_fidelities=None + if context.fidelity_idx is None + else (context.fidelity_idx,), + ) + + mll = gpytorch.ExactMarginalLogLikelihood(self._model.likelihood, self._model) + + botorch.fit.fit_gpytorch_mll(mll) + + @override + def __str__(self) -> str: + return ( + "Wrapper for a" + ":class:`~botorch.models.gp_regression_fidelity.SingleTaskMultiFidelityGP`," + "used as the default GP for discrete numerical fidelity parameters in," + "e.g., multi fidelity knowledge gradient." + ) + + +# Collect leftover original slotted classes processed by `attrs.define` +gc.collect() diff --git a/baybe/surrogates/gaussian_process/presets/__init__.py b/baybe/surrogates/gaussian_process/presets/__init__.py index deb7de9e64..017e673550 100644 --- a/baybe/surrogates/gaussian_process/presets/__init__.py +++ b/baybe/surrogates/gaussian_process/presets/__init__.py @@ -7,9 +7,6 @@ BayBEMeanFactory, ) -# Core -from baybe.surrogates.gaussian_process.presets.core import GaussianProcessPreset - # EDBO preset from baybe.surrogates.gaussian_process.presets.edbo import ( EDBOKernelFactory, diff --git a/baybe/surrogates/gaussian_process/presets/core.py b/baybe/surrogates/gaussian_process/presets/core.py deleted file mode 100644 index ad77df0b4d..0000000000 --- a/baybe/surrogates/gaussian_process/presets/core.py +++ /dev/null @@ -1,18 +0,0 @@ -"""Preset configurations for Gaussian process surrogates.""" - -from __future__ import annotations - -from enum import Enum - - -class GaussianProcessPreset(Enum): - """Available Gaussian process surrogate presets.""" - - BAYBE = "BAYBE" - """The default BayBE settings of the Gaussian process surrogate class.""" - - EDBO = "EDBO" - """The EDBO settings.""" - - EDBO_SMOOTHED = "EDBO_SMOOTHED" - """A smoothed version of the EDBO settings.""" diff --git a/baybe/surrogates/gaussian_process/utils.py b/baybe/surrogates/gaussian_process/utils.py new file mode 100644 index 0000000000..a8b35bf6fe --- /dev/null +++ b/baybe/surrogates/gaussian_process/utils.py @@ -0,0 +1,72 @@ +"""Gaussian process utilities.""" + +from typing import TYPE_CHECKING + +from attrs import define, field +from attrs.validators import instance_of + +from baybe.searchspace.core import SearchSpace + +if TYPE_CHECKING: + from torch import Tensor + + +@define +class _ModelContext: + """Model context for Gaussian process surrogates.""" + + searchspace: SearchSpace = field(validator=instance_of(SearchSpace)) + """The search space the model is trained on.""" + + @property + def task_idx(self) -> int | None: + """The computational column index of the task parameter, if available.""" + return self.searchspace.task_idx + + @property + def is_multitask(self) -> bool: + """Indicates if model is to be operated in a multi-task context.""" + return self.n_task_dimensions > 0 + + @property + def n_task_dimensions(self) -> int: + """The number of task dimensions.""" + # TODO: Generalize to multiple task parameters + return 1 if self.task_idx is not None else 0 + + @property + def n_tasks(self) -> int: + """The number of tasks.""" + return self.searchspace.n_tasks + + @property + def n_fidelity_dimensions(self) -> int: + """The number of fidelity dimensions.""" + # TODO: Generalize to multiple fidelity parameters + return 1 if self.searchspace.fidelity_idx is not None else 0 + + @property + def fidelity_idx(self) -> int | None: + """The computational column index of the fidelity parameter, if available.""" + return self.searchspace.fidelity_idx + + @property + def n_fidelities(self) -> int: + """The number of fidelities.""" + return self.searchspace.n_fidelities + + @property + def parameter_bounds(self) -> Tensor: + """Get the search space parameter bounds in BoTorch Format.""" + import torch + + return torch.from_numpy(self.searchspace.scaling_bounds.values) + + @property + def numerical_indices(self) -> list[int]: + """The indices of the regular numerical model inputs.""" + return [ + i + for i in range(len(self.searchspace.comp_rep_columns)) + if i not in (self.task_idx, self.fidelity_idx) + ] diff --git a/baybe/surrogates/linear.py b/baybe/surrogates/linear.py index ba1fed5b2c..3c4ded8f45 100644 --- a/baybe/surrogates/linear.py +++ b/baybe/surrogates/linear.py @@ -3,7 +3,7 @@ from __future__ import annotations import gc -from typing import TYPE_CHECKING, ClassVar, TypedDict +from typing import TYPE_CHECKING, TypedDict from attrs import define, field from typing_extensions import override @@ -41,9 +41,6 @@ class _ARDRegressionParams(TypedDict, total=False): class BayesianLinearSurrogate(IndependentGaussianSurrogate): """A Bayesian linear regression surrogate model.""" - supports_transfer_learning: ClassVar[bool] = False - # See base class. - model_params: _ARDRegressionParams = field( factory=dict, converter=dict, diff --git a/baybe/surrogates/naive.py b/baybe/surrogates/naive.py index 3912c6b128..ef7d624c5b 100644 --- a/baybe/surrogates/naive.py +++ b/baybe/surrogates/naive.py @@ -3,7 +3,7 @@ from __future__ import annotations import gc -from typing import TYPE_CHECKING, ClassVar +from typing import TYPE_CHECKING from attrs import define, field from typing_extensions import override @@ -23,9 +23,6 @@ class MeanPredictionSurrogate(IndependentGaussianSurrogate): as posterior mean and a (data-independent) constant posterior variance. """ - supports_transfer_learning: ClassVar[bool] = False - # See base class. - _model: float | None = field(init=False, default=None, eq=False) """The estimated posterior mean value of the training targets.""" diff --git a/baybe/surrogates/ngboost.py b/baybe/surrogates/ngboost.py index 6cb1cee135..d4a4363e3f 100644 --- a/baybe/surrogates/ngboost.py +++ b/baybe/surrogates/ngboost.py @@ -46,9 +46,6 @@ class _NGBRegressorParams(TypedDict, total=False): class NGBoostSurrogate(IndependentGaussianSurrogate): """A natural-gradient-boosting surrogate model.""" - supports_transfer_learning: ClassVar[bool] = False - # See base class. - _default_model_params: ClassVar[dict] = {"n_estimators": 25, "verbose": False} """Class variable encoding the default model parameters.""" diff --git a/baybe/surrogates/random_forest.py b/baybe/surrogates/random_forest.py index 91ad599bb1..f7726828be 100644 --- a/baybe/surrogates/random_forest.py +++ b/baybe/surrogates/random_forest.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import Collection -from typing import TYPE_CHECKING, ClassVar, Literal, Protocol, TypedDict +from typing import TYPE_CHECKING, Literal, Protocol, TypedDict import numpy as np import numpy.typing as npt @@ -61,9 +61,6 @@ def predict(self, x: np.ndarray, /) -> np.ndarray: ... class RandomForestSurrogate(Surrogate): """A random forest surrogate model.""" - supports_transfer_learning: ClassVar[bool] = False - # See base class. - model_params: _RandomForestRegressorParams = field( factory=dict, converter=dict,