Refactor GP Surrogates

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

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,
         *,

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

baybe/surrogates/gaussian_process/__init__.py

@@ -0,0 +1,7 @@
+"""Gaussian process surrogates."""
+
+from baybe.surrogates.gaussian_process.core import GaussianProcessSurrogate
+
+__all__ = [
+    "GaussianProcessSurrogate",
+]

baybe/surrogates/gaussian_process.py → 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(

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

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",
+]

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}'."
+    )

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]