diff --git a/CHANGELOG.md b/CHANGELOG.md index efd679ddc0..26e24b008c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Class hierarchy for objectives - Deserialization is now also possible from optional class name abbreviations - `Kernel`, `MaternKernel`, and `ScaleKernel` classes for specifying kernels +- `KernelFactory` protocol enabling context-dependent construction of kernels +- Preset mechanism for `GaussianProcessSurrogate` - `hypothesis` strategies and roundtrip test for kernels, constraints, objectives, priors and acquisition functions - New acquisition functions: `qSR`, `qNEI`, `LogEI`, `qLogEI`, `qLogNEI` @@ -21,6 +23,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed - Reorganized acquisition.py into `acquisition` subpackage - Reorganized simulation.py into `simulation` subpackage +- Reorganized gaussian_process.py into `gaussian_process` subpackage - Acquisition functions are now their own objects - `acquisition_function_cls` constructor parameter renamed to `acquisition_function` - User guide now explains the new objective classes diff --git a/baybe/kernels/base.py b/baybe/kernels/base.py index 4ccde20078..bbcbeda1eb 100644 --- a/baybe/kernels/base.py +++ b/baybe/kernels/base.py @@ -19,11 +19,19 @@ if TYPE_CHECKING: import torch + from baybe.surrogates.gaussian_process.kernel_factory import PlainKernelFactory + @define(frozen=True) class Kernel(ABC, SerialMixin): """Abstract base class for all kernels.""" + def to_factory(self) -> PlainKernelFactory: + """Wrap the kernel in a :class:`baybe.surrogates.gaussian_process.kernel_factory.PlainKernelFactory`.""" # noqa: E501 + from baybe.surrogates.gaussian_process.kernel_factory import PlainKernelFactory + + return PlainKernelFactory(self) + def to_gpytorch( self, *, diff --git a/baybe/surrogates/__init__.py b/baybe/surrogates/__init__.py index 4f032db707..1128442b8c 100644 --- a/baybe/surrogates/__init__.py +++ b/baybe/surrogates/__init__.py @@ -1,7 +1,7 @@ """BayBE surrogates.""" from baybe.surrogates.custom import _ONNX_INSTALLED, register_custom_architecture -from baybe.surrogates.gaussian_process import GaussianProcessSurrogate +from baybe.surrogates.gaussian_process.core import GaussianProcessSurrogate from baybe.surrogates.linear import BayesianLinearSurrogate from baybe.surrogates.naive import MeanPredictionSurrogate from baybe.surrogates.ngboost import NGBoostSurrogate diff --git a/baybe/surrogates/gaussian_process/__init__.py b/baybe/surrogates/gaussian_process/__init__.py new file mode 100644 index 0000000000..a47b756b91 --- /dev/null +++ b/baybe/surrogates/gaussian_process/__init__.py @@ -0,0 +1,7 @@ +"""Gaussian process surrogates.""" + +from baybe.surrogates.gaussian_process.core import GaussianProcessSurrogate + +__all__ = [ + "GaussianProcessSurrogate", +] diff --git a/baybe/surrogates/gaussian_process.py b/baybe/surrogates/gaussian_process/core.py similarity index 59% rename from baybe/surrogates/gaussian_process.py rename to baybe/surrogates/gaussian_process/core.py index 49679cc367..bcb7eb361a 100644 --- a/baybe/surrogates/gaussian_process.py +++ b/baybe/surrogates/gaussian_process/core.py @@ -2,15 +2,24 @@ from __future__ import annotations -from typing import TYPE_CHECKING, ClassVar, Optional +from typing import TYPE_CHECKING, ClassVar -from attr import define, field +from attrs import define, field -from baybe.kernels import MaternKernel, ScaleKernel -from baybe.kernels.base import Kernel -from baybe.priors import GammaPrior from baybe.searchspace import SearchSpace from baybe.surrogates.base import Surrogate +from baybe.surrogates.gaussian_process.kernel_factory import ( + KernelFactory, + to_kernel_factory, +) +from baybe.surrogates.gaussian_process.presets import ( + GaussianProcessPreset, + make_gp_from_preset, +) +from baybe.surrogates.gaussian_process.presets.default import ( + DefaultKernelFactory, + _default_noise_factory, +) if TYPE_CHECKING: from torch import Tensor @@ -28,14 +37,28 @@ class GaussianProcessSurrogate(Surrogate): # See base class. # Object variables - kernel: Optional[Kernel] = field(default=None) - """The kernel used by the Gaussian Process.""" + kernel_factory: KernelFactory = field( + alias="kernel_or_factory", + factory=DefaultKernelFactory, + converter=to_kernel_factory, + ) + """The factory used to create the kernel of the Gaussian process. + + Accepts either a :class:`baybe.kernels.base.Kernel` or a + :class:`.kernel_factory.KernelFactory`. + When passing a :class:`baybe.kernels.base.Kernel`, it gets automatically wrapped + into a :class:`.kernel_factory.PlainKernelFactory`.""" # TODO: type should be Optional[botorch.models.SingleTaskGP] but is currently # omitted due to: https://github.com/python-attrs/cattrs/issues/531 _model = field(init=False, default=None, eq=False) """The actual model.""" + @classmethod + def from_preset(preset: GaussianProcessPreset) -> GaussianProcessSurrogate: + """Create a Gaussian process surrogate from one of the defined presets.""" + return make_gp_from_preset(preset) + def _posterior(self, candidates: Tensor) -> tuple[Tensor, Tensor]: # See base class. posterior = self._model.posterior(candidates) @@ -65,60 +88,16 @@ def _fit(self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor) -> No ) outcome_transform = botorch.models.transforms.Standardize(train_y.shape[1]) - # ---------- GP prior selection ---------- # - # TODO: temporary prior choices adapted from edbo, replace later on - - mordred = (searchspace.contains_mordred or searchspace.contains_rdkit) and ( - train_x.shape[-1] >= 50 - ) - - # TODO Until now, only the kernels use our custom priors, hence the explicit - # to_gpytorch() calls for all others - # low D priors - if train_x.shape[-1] < 10: - lengthscale_prior = [GammaPrior(1.2, 1.1), 0.2] - outputscale_prior = [GammaPrior(5.0, 0.5), 8.0] - noise_prior = [GammaPrior(1.05, 0.5), 0.1] - - # DFT optimized priors - elif mordred and train_x.shape[-1] < 100: - lengthscale_prior = [GammaPrior(2.0, 0.2), 5.0] - outputscale_prior = [GammaPrior(5.0, 0.5), 8.0] - noise_prior = [GammaPrior(1.5, 0.1), 5.0] - - # Mordred optimized priors - elif mordred: - lengthscale_prior = [GammaPrior(2.0, 0.1), 10.0] - outputscale_prior = [GammaPrior(2.0, 0.1), 10.0] - noise_prior = [GammaPrior(1.5, 0.1), 5.0] - - # OHE optimized priors - else: - lengthscale_prior = [GammaPrior(3.0, 1.0), 2.0] - outputscale_prior = [GammaPrior(5.0, 0.2), 20.0] - noise_prior = [GammaPrior(1.5, 0.1), 5.0] - - # ---------- End: GP prior selection ---------- # - # extract the batch shape of the training data batch_shape = train_x.shape[:-2] # create GP mean mean_module = gpytorch.means.ConstantMean(batch_shape=batch_shape) - # If no kernel is provided, we construct one from our priors - if self.kernel is None: - self.kernel = ScaleKernel( - base_kernel=MaternKernel( - lengthscale_prior=lengthscale_prior[0], - lengthscale_initial_value=lengthscale_prior[1], - ), - outputscale_prior=outputscale_prior[0], - outputscale_initial_value=outputscale_prior[1], - ) - # define the covariance module for the numeric dimensions - base_covar_module = self.kernel.to_gpytorch( + base_covar_module = self.kernel_factory( + searchspace, train_x, train_y + ).to_gpytorch( ard_num_dims=train_x.shape[-1] - n_task_params, active_dims=numeric_idxs, batch_shape=batch_shape, @@ -136,11 +115,11 @@ def _fit(self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor) -> No covar_module = base_covar_module * task_covar_module # create GP likelihood + noise_prior = _default_noise_factory(searchspace, train_x, train_y) likelihood = gpytorch.likelihoods.GaussianLikelihood( noise_prior=noise_prior[0].to_gpytorch(), batch_shape=batch_shape ) - if noise_prior[1] is not None: - likelihood.noise = torch.tensor([noise_prior[1]]) + likelihood.noise = torch.tensor([noise_prior[1]]) # construct and fit the Gaussian process self._model = botorch.models.SingleTaskGP( diff --git a/baybe/surrogates/gaussian_process/kernel_factory.py b/baybe/surrogates/gaussian_process/kernel_factory.py new file mode 100644 index 0000000000..002ea271d2 --- /dev/null +++ b/baybe/surrogates/gaussian_process/kernel_factory.py @@ -0,0 +1,55 @@ +"""Kernel factories for the Gaussian process surrogate.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Protocol, Union + +from attrs import define, field +from attrs.validators import instance_of + +from baybe.kernels.base import Kernel +from baybe.searchspace import SearchSpace +from baybe.serialization.core import ( + converter, + get_base_structure_hook, + unstructure_base, +) +from baybe.serialization.mixin import SerialMixin + +if TYPE_CHECKING: + from torch import Tensor + + +class KernelFactory(Protocol): + """A protocol defining the interface expected for kernel factories.""" + + def __call__( + self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor + ) -> Kernel: + """Create a :class:`baybe.kernels.base.Kernel` for the given DOE context.""" + ... + + +# Register de-/serialization hooks +converter.register_structure_hook(KernelFactory, get_base_structure_hook(KernelFactory)) +converter.register_unstructure_hook(KernelFactory, unstructure_base) + + +@define(frozen=True) +class PlainKernelFactory(KernelFactory, SerialMixin): + """A trivial factory that returns a fixed pre-defined kernel upon request.""" + + kernel: Kernel = field(validator=instance_of(Kernel)) + """The fixed kernel to be returned by the factory.""" + + def __call__( # noqa: D102 + self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor + ) -> Kernel: + # See base class. + + return self.kernel + + +def to_kernel_factory(x: Union[Kernel, KernelFactory], /) -> KernelFactory: + """Wrap a kernel into a plain kernel factory (with factory passthrough).""" + return x.to_factory() if isinstance(x, Kernel) else x diff --git a/baybe/surrogates/gaussian_process/presets/__init__.py b/baybe/surrogates/gaussian_process/presets/__init__.py new file mode 100644 index 0000000000..c301cd87fc --- /dev/null +++ b/baybe/surrogates/gaussian_process/presets/__init__.py @@ -0,0 +1,11 @@ +"""Gaussian process surrogate presets.""" + +from baybe.surrogates.gaussian_process.presets.core import ( + GaussianProcessPreset, + make_gp_from_preset, +) + +__all__ = [ + "make_gp_from_preset", + "GaussianProcessPreset", +] diff --git a/baybe/surrogates/gaussian_process/presets/core.py b/baybe/surrogates/gaussian_process/presets/core.py new file mode 100644 index 0000000000..df3276e63d --- /dev/null +++ b/baybe/surrogates/gaussian_process/presets/core.py @@ -0,0 +1,26 @@ +"""Preset configurations for Gaussian process surrogates.""" + +from __future__ import annotations + +from enum import Enum +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from baybe.surrogates.gaussian_process.core import GaussianProcessSurrogate + + +class GaussianProcessPreset(Enum): + """Available Gaussian process surrogate presets.""" + + BAYBE = "BAYBE" + """Recreates the default settings of the Gaussian process surrogate class.""" + + +def make_gp_from_preset(preset: GaussianProcessPreset) -> GaussianProcessSurrogate: + """Create a :class:`GaussianProcessSurrogate` from a :class:`GaussianProcessPreset.""" # noqa: E501 + if preset is GaussianProcessPreset.BAYBE: + return GaussianProcessSurrogate() + + raise ValueError( + f"Unknown '{GaussianProcessPreset.__name__}' with name '{preset.name}'." + ) diff --git a/baybe/surrogates/gaussian_process/presets/default.py b/baybe/surrogates/gaussian_process/presets/default.py new file mode 100644 index 0000000000..0e81c0a10e --- /dev/null +++ b/baybe/surrogates/gaussian_process/presets/default.py @@ -0,0 +1,110 @@ +"""Default preset for Gaussian process surrogates.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from attrs import define + +from baybe.kernels.basic import MaternKernel, ScaleKernel +from baybe.priors.basic import GammaPrior +from baybe.surrogates.gaussian_process.kernel_factory import KernelFactory + +if TYPE_CHECKING: + from torch import Tensor + + from baybe.kernels.base import Kernel + from baybe.searchspace.core import SearchSpace + + +@define +class DefaultKernelFactory(KernelFactory): + """A factory providing the default kernel for Gaussian process surrogates. + + The logic is adapted from EDBO (Experimental Design via Bayesian Optimization). + + References: + * https://github.com/b-shields/edbo + * https://doi.org/10.1038/s41586-021-03213-y + """ + + def __call__( # noqa: D102 + self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor + ) -> Kernel: + # See base class. + + mordred = (searchspace.contains_mordred or searchspace.contains_rdkit) and ( + train_x.shape[-1] >= 50 + ) + + # low D priors + if train_x.shape[-1] < 10: # <-- different condition compared to EDBO + lengthscale_prior = GammaPrior(1.2, 1.1) + lengthscale_initial_value = 0.2 + outputscale_prior = GammaPrior(5.0, 0.5) + outputscale_initial_value = 8.0 + + # DFT optimized priors + elif mordred and train_x.shape[-1] < 100: + lengthscale_prior = GammaPrior(2.0, 0.2) + lengthscale_initial_value = 5.0 + outputscale_prior = GammaPrior(5.0, 0.5) + outputscale_initial_value = 8.0 + + # Mordred optimized priors + elif mordred: + lengthscale_prior = GammaPrior(2.0, 0.1) + lengthscale_initial_value = 10.0 + outputscale_prior = GammaPrior(2.0, 0.1) + outputscale_initial_value = 10.0 + + # OHE optimized priors + else: + lengthscale_prior = GammaPrior(3.0, 1.0) + lengthscale_initial_value = 2.0 + outputscale_prior = GammaPrior(5.0, 0.2) + outputscale_initial_value = 20.0 + + return ScaleKernel( + MaternKernel( + nu=2.5, + lengthscale_prior=lengthscale_prior, + lengthscale_initial_value=lengthscale_initial_value, + ), + outputscale_prior=outputscale_prior, + outputscale_initial_value=outputscale_initial_value, + ) + + +def _default_noise_factory( + searchspace: SearchSpace, train_x: Tensor, train_y: Tensor +) -> tuple[GammaPrior, float]: + """Create the default noise settings for the Gaussian process surrogate. + + The logic is adapted from EDBO (Experimental Design via Bayesian Optimization). + + References: + * https://github.com/b-shields/edbo + * https://doi.org/10.1038/s41586-021-03213-y + """ + # TODO: Replace this function with a proper likelihood factory + + uses_descriptors = ( + searchspace.contains_mordred or searchspace.contains_rdkit + ) and (train_x.shape[-1] >= 50) + + # low D priors + if train_x.shape[-1] < 10: # <-- different condition compared to EDBO + return [GammaPrior(1.05, 0.5), 0.1] + + # DFT optimized priors + elif uses_descriptors and train_x.shape[-1] < 100: + return [GammaPrior(1.5, 0.1), 5.0] + + # Mordred optimized priors + elif uses_descriptors: + return [GammaPrior(1.5, 0.1), 5.0] + + # OHE optimized priors + else: + return [GammaPrior(1.5, 0.1), 5.0] diff --git a/baybe/utils/boolean.py b/baybe/utils/boolean.py index 7bec444dbb..e1d1318fea 100644 --- a/baybe/utils/boolean.py +++ b/baybe/utils/boolean.py @@ -3,7 +3,8 @@ from abc import ABC from typing import Any -from attr import cmp_using +from attrs import cmp_using +from typing_extensions import is_protocol # Used for comparing pandas dataframes in attrs classes eq_dataframe = cmp_using(lambda x, y: x.equals(y)) @@ -16,7 +17,9 @@ def is_abstract(cls: Any) -> bool: if a class has abstract methods. The latter can be problematic when the class has no abstract methods but is nevertheless not directly usable, for example, because it has uninitialized members, which are only covered in its non-"abstract" subclasses. - By contrast, this method simply checks if the class derives from ``abc.ABC``. + + By contrast, this method simply checks if the class derives from ``abc.ABC`` or + is a protocol class. Args: cls: The class to be inspected. @@ -24,7 +27,7 @@ def is_abstract(cls: Any) -> bool: Returns: ``True`` if the class is "abstract" (see definition above), ``False`` else. """ - return ABC in cls.__bases__ + return ABC in cls.__bases__ or is_protocol(cls) def strtobool(val: str) -> bool: diff --git a/docs/userguide/surrogates.md b/docs/userguide/surrogates.md index 7b34ec1589..44a7051843 100644 --- a/docs/userguide/surrogates.md +++ b/docs/userguide/surrogates.md @@ -6,7 +6,7 @@ Surrogate models are used to model and estimate the unknown objective function o BayBE provides a comprehensive selection of surrogate models, empowering you to choose the most suitable option for your specific needs. The following surrogate models are available within BayBE: -* [`GaussianProcessSurrogate`](baybe.surrogates.gaussian_process.GaussianProcessSurrogate) +* [`GaussianProcessSurrogate`](baybe.surrogates.gaussian_process.core.GaussianProcessSurrogate) * [`BayesianLinearSurrogate`](baybe.surrogates.linear.BayesianLinearSurrogate) * [`MeanPredictionSurrogate`](baybe.surrogates.naive.MeanPredictionSurrogate) * [`NGBoostSurrogate`](baybe.surrogates.ngboost.NGBoostSurrogate) diff --git a/pyproject.toml b/pyproject.toml index 6458e38144..d0f8fdbd91 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,6 +44,7 @@ dependencies = [ "scipy>=1.10.1", "setuptools-scm>=7.1.0", "torch>=1.13.1", + "typing_extensions>=4.7.0", # Telemetry: "opentelemetry-sdk>=1.16.0", diff --git a/tests/conftest.py b/tests/conftest.py index d94c2bbcc5..5bee25f32d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -619,7 +619,7 @@ def fixture_default_surrogate_model(request, onnx_surrogate, kernel): """The default surrogate model to be used if not specified differently.""" if hasattr(request, "param") and request.param == "onnx": return onnx_surrogate - return GaussianProcessSurrogate(kernel=kernel) + return GaussianProcessSurrogate(kernel_or_factory=kernel) @pytest.fixture(name="initial_recommender")