Beef up GaussianProcessSurrogate

CHANGELOG.md

@@ -12,17 +12,43 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Settings option for random seed control
 - `identify_non_dominated_configurations` method to `Campaign` and `Objective`
   for determining the Pareto front
+- Gaussian process component factories
+- Support for GPyTorch objects (kernels, means, likelihood) as Gaussian process
+  components, enabling full low-level customization
+- Factories for all Gaussian process components
+- `EDBO` and `EDBO_SMOOTHED` presets for `GaussianProcessSurrogate`
+- `TypeSelector` and `NameSelector` classes for parameter selection in kernel factories
+- `parameter_names` attribute to basic kernels for controlling the considered parameters
+- `ParameterKind` flag enum for classifying parameters by their role and automatic
+  parameter kind validation in kernel factories
+- `IndexKernel` and `PositiveIndexKernel` classes
 - Interpoint constraints for continuous search spaces
+- `IndexKernel` and `PositiveIndexKernel` classes
+- Addition and multiplication operators for kernel objects, enabling kernel
+  composition via `+` (sum) and `*` (product), as well as `constant * kernel`
+  for creating a `ScaleKernel` with a fixed output scale
 
 ### Breaking Changes
 - `ContinuousLinearConstraint.to_botorch` now returns a collection of constraint tuples
   instead of a single tuple (needed for interpoint constraints)
+- `Kernel.to_gpytorch` now takes a `SearchSpace` instead of explicit `ard_num_dims`,
+  `batch_shape` and `active_dims` arguments, as kernels now automatically adjust this
+  configuration to the given search space
+- `GaussianProcessSurrogate` no longer automatically adds a task kernel in multi-task
+  scenarios. Custom kernel architectures must now explicitly include the task kernel,
+  e.g. via `ICMKernelFactory`
 
 ### Removed
 - `parallel_runs` argument from `simulate_scenarios`, since parallelization
   can now be conveniently controlled via the new `Settings` mechanism
+- `make_gp_from_preset` utility function, since the same functionality is offered by
+  `GaussianProcessSurrogate.from_preset` 
 
 ### Deprecations
+- Using a custom kernel with `GaussianProcessSurrogate` in a multi-task context now
+  raises a `DeprecationError` to alert users about the changed kernel logic. This can
+  be suppressed by setting the `BAYBE_DISABLE_CUSTOM_KERNEL_WARNING` environment
+  variable to a truthy value
 - `set_random_seed` and `temporary_seed` utility functions
 - The environment variables
   `BAYBE_NUMPY_USE_SINGLE_PRECISION`/`BAYBE_TORCH_USE_SINGLE_PRECISION` have been

baybe/kernels/base.py

@@ -3,54 +3,123 @@
 from __future__ import annotations
 
 import gc
-from abc import ABC
+from abc import ABC, abstractmethod
+from itertools import chain
 from typing import TYPE_CHECKING, Any
 
-from attrs import define
+from attrs import define, field
+from attrs.converters import optional as optional_c
+from attrs.validators import deep_iterable, instance_of
+from attrs.validators import optional as optional_v
+from typing_extensions import override
 
 from baybe.exceptions import UnmatchedAttributeError
 from baybe.priors.base import Prior
+from baybe.searchspace.core import SearchSpace
 from baybe.serialization.mixin import SerialMixin
 from baybe.settings import active_settings
-from baybe.utils.basic import get_baseclasses, match_attributes
+from baybe.utils.basic import classproperty, get_baseclasses, match_attributes, to_tuple
 
 if TYPE_CHECKING:
-    import torch
-
-    from baybe.surrogates.gaussian_process.kernel_factory import PlainKernelFactory
+    from baybe.surrogates.gaussian_process.components.kernel import PlainKernelFactory
 
 
 @define(frozen=True)
 class Kernel(ABC, SerialMixin):
     """Abstract base class for all kernels."""
 
+    def __add__(self, other: Any) -> Kernel:
+        """Create a sum kernel from two kernels.
+
+        Flattens nested sums so that ``(a + b) + c`` yields
+        ``AdditiveKernel([a, b, c])`` instead of
+        ``AdditiveKernel([AdditiveKernel([a, b]), c])``.
+        """
+        if isinstance(other, Kernel):
+            from baybe.kernels.composite import AdditiveKernel
+
+            left = self.base_kernels if isinstance(self, AdditiveKernel) else (self,)
+            right = (
+                other.base_kernels if isinstance(other, AdditiveKernel) else (other,)
+            )
+            return AdditiveKernel([*left, *right])
+        return NotImplemented
+
+    def __radd__(self, other: Any) -> Kernel:
+        """Support right-hand addition for kernel objects."""
+        # Enable use with built-in sum(), which starts with 0 + first_element.
+        if other == 0:
+            return self
+        if isinstance(other, Kernel):
+            return self.__add__(other)
+        return NotImplemented
+
+    def __mul__(self, other: Any) -> Kernel:
+        """Create a product kernel or scale kernel.
+
+        When multiplied with another kernel, a product kernel is created. Nested
+        products are flattened so that ``(a * b) * c`` yields
+        ``ProductKernel([a, b, c])``. When multiplied with a numeric constant, a scale
+        kernel with a fixed (non-trainable) output scale is created.
+        """
+        if isinstance(other, Kernel):
+            from baybe.kernels.composite import ProductKernel
+
+            left = self.base_kernels if isinstance(self, ProductKernel) else (self,)
+            right = other.base_kernels if isinstance(other, ProductKernel) else (other,)
+            return ProductKernel([*left, *right])
+        if isinstance(other, (int, float)):
+            if other == 1:
+                return self
+
+            from baybe.kernels.composite import ScaleKernel
+
+            return ScaleKernel(
+                base_kernel=self,
+                outputscale_initial_value=float(other),
+                outputscale_trainable=False,
+            )
+        return NotImplemented
+
+    def __rmul__(self, other: Any) -> Kernel:
+        """Support right-hand multiplication, enabling ``constant * kernel``."""
+        # Enable use with math.prod(), which starts with 1 * first_element.
+        return self.__mul__(other)
+
+    @classproperty
+    def _whitelisted_attributes(cls) -> frozenset[str]:
+        """Attribute names to exclude from gpytorch matching."""
+        return frozenset()
+
     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
+        """Wrap the kernel in a :class:`baybe.surrogates.gaussian_process.components.PlainKernelFactory`."""  # noqa: E501
+        from baybe.surrogates.gaussian_process.components.kernel import (
+            PlainKernelFactory,
+        )
 
         return PlainKernelFactory(self)
 
-    def to_gpytorch(
-        self,
-        *,
-        ard_num_dims: int | None = None,
-        batch_shape: torch.Size | None = None,
-        active_dims: tuple[int, ...] | None = None,
-    ):
+    @abstractmethod
+    def _get_dimensions(
+        self, searchspace: SearchSpace
+    ) -> tuple[tuple[int, ...] | None, int | None]:
+        """Get the active dimensions and the number of ARD dimensions."""
+
+    def to_gpytorch(self, searchspace: SearchSpace):
         """Create the gpytorch representation of the kernel."""
         import gpytorch.kernels
 
-        # Extract keywords with non-default values. This is required since gpytorch
-        # makes use of kwargs, i.e. differentiates if certain keywords are explicitly
-        # passed or not. For instance, `ard_num_dims = kwargs.get("ard_num_dims", 1)`
-        # fails if we explicitly pass `ard_num_dims=None`.
-        kw: dict[str, Any] = dict(
-            ard_num_dims=ard_num_dims, batch_shape=batch_shape, active_dims=active_dims
-        )
-        kw = {k: v for k, v in kw.items() if v is not None}
+        active_dims, ard_num_dims = self._get_dimensions(searchspace)
 
         # Get corresponding gpytorch kernel class and its base classes
-        kernel_cls = getattr(gpytorch.kernels, self.__class__.__name__)
+        try:
+            kernel_cls = getattr(gpytorch.kernels, self.__class__.__name__)
+        except AttributeError:
+            import botorch.models.kernels.positive_index
+
+            kernel_cls = getattr(
+                botorch.models.kernels.positive_index, self.__class__.__name__
+            )
         base_classes = get_baseclasses(kernel_cls, abstract=True)
 
         # Fetch the necessary gpytorch constructor parameters of the kernel.
@@ -66,11 +135,17 @@ def to_gpytorch(
             unmatched_attrs.update(unmatched)
 
         # Sanity check: all attributes of the BayBE kernel need a corresponding match
-        # in the gpytorch kernel (otherwise, the BayBE kernel class is misconfigured).
-        # Exception: initial values are not used during construction but are set
-        # on the created object (see code at the end of the method).
-        missing = set(unmatched) - set(kernel_attrs)
-        if leftover := {m for m in missing if not m.endswith("_initial_value")}:
+        # with the gpytorch kernel signature (otherwise, the BayBE kernel class is
+        # misconfigured). Exceptions: initial values and trainability flags are not used
+        # during construction but are set on the created object after construction.
+        missing = (
+            set(unmatched_attrs) - set(kernel_attrs) - self._whitelisted_attributes
+        )
+        if leftover := {
+            m
+            for m in missing
+            if not m.endswith("_initial_value") and not m.endswith("_trainable")
+        }:
             raise UnmatchedAttributeError(leftover)
 
         # Convert specified priors to gpytorch, if provided
@@ -82,15 +157,17 @@ def to_gpytorch(
 
         # Convert specified inner kernels to gpytorch, if provided
         kernel_dict = {
-            key: value.to_gpytorch(**kw)
+            key: value.to_gpytorch(searchspace)
             for key, value in kernel_attrs.items()
             if isinstance(value, Kernel)
         }
 
         # Create the kernel with all its inner gpytorch objects
         kernel_attrs.update(kernel_dict)
         kernel_attrs.update(prior_dict)
-        gpytorch_kernel = kernel_cls(**kernel_attrs, **kw)
+        gpytorch_kernel = kernel_cls(
+            **kernel_attrs, ard_num_dims=ard_num_dims, active_dims=active_dims
+        )
 
         # If the kernel has a lengthscale, set its initial value
         if kernel_cls.has_lengthscale:
@@ -113,11 +190,58 @@ def to_gpytorch(
 class BasicKernel(Kernel, ABC):
     """Abstract base class for all basic kernels."""
 
+    parameter_names: tuple[str, ...] | None = field(
+        default=None,
+        converter=optional_c(to_tuple),
+        validator=optional_v(
+            deep_iterable(
+                iterable_validator=instance_of(tuple),
+                member_validator=instance_of(str),
+            )
+        ),
+        kw_only=True,
+    )
+    """An optional set of names specifiying the parameters the kernel should act on."""
+
+    @override
+    @classproperty
+    def _whitelisted_attributes(cls) -> frozenset[str]:
+        return frozenset({"parameter_names"})
+
+    @override
+    def _get_dimensions(
+        self, searchspace: SearchSpace
+    ) -> tuple[tuple[int, ...] | None, int | None]:
+        if self.parameter_names is None:
+            # `None` is gpytorch's default indicating that all dimensions are active
+            active_dims = None
+        else:
+            active_dims = tuple(
+                chain.from_iterable(
+                    searchspace.get_comp_rep_parameter_indices(name)
+                    for name in self.parameter_names
+                )
+            )
+
+        # We use automatic relevance determination for all kernels
+        ard_num_dims = (
+            len(active_dims)
+            if active_dims is not None
+            else len(searchspace.comp_rep_columns)
+        )
+        return active_dims, ard_num_dims
+
 
 @define(frozen=True)
 class CompositeKernel(Kernel, ABC):
     """Abstract base class for all composite kernels."""
 
+    @override
+    def _get_dimensions(
+        self, searchspace: SearchSpace
+    ) -> tuple[tuple[int, ...] | None, int | None]:
+        return None, None
+
 
 # Collect leftover original slotted classes processed by `attrs.define`
 gc.collect()

baybe/kernels/basic.py

@@ -215,5 +215,29 @@ class RQKernel(BasicKernel):
     """An optional initial value for the kernel lengthscale."""
 
 
+@define(frozen=True)
+class IndexKernel(BasicKernel):
+    """An index kernel for transfer learning across tasks."""
+
+    num_tasks: int = field(validator=[instance_of(int), ge(2)])
+    """The number of tasks."""
+
+    rank: int = field(validator=[instance_of(int), ge(1)])
+    """The rank of the task covariance matrix."""
+
+    @rank.validator
+    def _validate_rank(self, _, rank: int):
+        if rank > self.num_tasks:
+            raise ValueError(
+                f"The rank of the task covariance matrix must be smaller than "
+                f"the number of tasks. Got rank {rank} > {self.num_tasks} tasks."
+            )
+
+
+@define(frozen=True)
+class PositiveIndexKernel(IndexKernel):
+    """A positive index kernel for transfer learning across tasks."""
+
+
 # Collect leftover original slotted classes processed by `attrs.define`
 gc.collect()

baybe/kernels/composite.py

@@ -36,6 +36,12 @@ class ScaleKernel(CompositeKernel):
     )
     """An optional initial value for the output scale."""
 
+    outputscale_trainable: bool = field(default=True, validator=instance_of(bool))
+    """Boolean flag indicating whether the output scale is trainable.
+
+    If ``False``, the output scale is frozen at its initial value and excluded from
+    optimization."""
+
     @override
     def to_gpytorch(self, *args, **kwargs):
         import torch
@@ -45,6 +51,8 @@ def to_gpytorch(self, *args, **kwargs):
             gpytorch_kernel.outputscale = torch.tensor(
                 initial_value, dtype=active_settings.DTypeFloatTorch
             )
+        if not self.outputscale_trainable:
+            gpytorch_kernel.raw_outputscale.requires_grad_(False)
         return gpytorch_kernel
 
 

baybe/parameters/base.py

@@ -21,6 +21,7 @@
 from baybe.utils.metadata import MeasurableMetadata, to_metadata
 
 if TYPE_CHECKING:
+    from baybe.parameters.enum import _ParameterKind
     from baybe.searchspace.continuous import SubspaceContinuous
     from baybe.searchspace.core import SearchSpace
     from baybe.searchspace.discrete import SubspaceDiscrete
@@ -77,6 +78,13 @@ def is_discrete(self) -> bool:
         """Boolean indicating if this is a discrete parameter."""
         return isinstance(self, DiscreteParameter)
 
+    @property
+    def _kind(self) -> _ParameterKind:
+        """The kind of the parameter."""
+        from baybe.parameters.enum import _ParameterKind
+
+        return _ParameterKind.from_parameter(self)
+
     @property
     @abstractmethod
     def comp_rep_columns(self) -> tuple[str, ...]: