Adaptive hyper-prior: CHENKernelFactory

CHANGELOG.md

@@ -16,9 +16,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - 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`
+- `CHEN`, `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

CONTRIBUTORS.md

@@ -39,4 +39,10 @@
 - Kathrin Skubch (Merck KGaA, Darmstadt, Germany):\
   Transfer learning regression benchmarks infrastructure
 - Myra Zmarsly (Merck Life Science KGaA, Darmstadt, Germany):\
-  Identification of non-dominated parameter configurations
+  Identification of non-dominated parameter configurations
+- Thijs Stuyver (PSL University, Paris, France):\  
+  Adaptive hyper-prior tailored for reaction yield optimization tasks
+- Maximilian Fleck (PSL University, Paris, France):\
+  Adaptive hyper-prior tailored for reaction yield optimization tasks
+- Guanming Chen (PSL University, Paris, France):\
+  Adaptive hyper-prior tailored for reaction yield optimization tasks

baybe/parameters/__init__.py

@@ -5,6 +5,7 @@
 from baybe.parameters.enum import (
     CategoricalEncoding,
     CustomEncoding,
+    ParameterKind,
     SubstanceEncoding,
 )
 from baybe.parameters.numerical import (
@@ -22,6 +23,7 @@
     "MeasurableMetadata",
     "NumericalContinuousParameter",
     "NumericalDiscreteParameter",
+    "ParameterKind",
     "SubstanceEncoding",
     "SubstanceParameter",
     "TaskParameter",

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, ...]:

baybe/parameters/enum.py

@@ -1,6 +1,38 @@
 """Parameter-related enumerations."""
 
-from enum import Enum
+from __future__ import annotations
+
+from enum import Enum, Flag, auto
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from baybe.parameters.base import Parameter
+
+
+class ParameterKind(Flag):
+    """Flag enum encoding the kind of a parameter.
+
+    Can be used to express compatibility (e.g. Gaussian process kernel factories)
+    with different parameter types via bitwise combination of flags.
+    """
+
+    REGULAR = auto()
+    """Regular parameter undergoing no special treatment."""
+
+    TASK = auto()
+    """Task parameter for transfer learning."""
+
+    FIDELITY = auto()
+    """Fidelity parameter for multi-fidelity modelling."""
+
+    @staticmethod
+    def from_parameter(parameter: Parameter) -> ParameterKind:
+        """Determine the kind of a parameter from its type."""
+        from baybe.parameters.categorical import TaskParameter
+
+        if isinstance(parameter, TaskParameter):
+            return ParameterKind.TASK
+        return ParameterKind.REGULAR
 
 
 class ParameterEncoding(Enum):

baybe/parameters/selectors.py

@@ -3,15 +3,13 @@
 import re
 from abc import ABC, abstractmethod
 from collections.abc import Collection
-from typing import ClassVar, Protocol
+from typing import Protocol
 
 from attrs import Converter, define, field
-from attrs.converters import optional
 from attrs.validators import deep_iterable, instance_of, min_len
 from typing_extensions import override
 
 from baybe.parameters.base import Parameter
-from baybe.searchspace.core import SearchSpace
 from baybe.utils.basic import to_tuple
 from baybe.utils.conversion import nonstring_to_tuple
 
@@ -131,37 +129,3 @@ def to_parameter_selector(
         return TypeSelector(items)
 
     raise TypeError(f"Cannot convert {x!r} to a parameter selector.")
-
-
-@define
-class _ParameterSelectorMixin:
-    """A mixin class to enable parameter selection."""
-
-    # For internal use only: sanity check mechanism to remind developers of new
-    # subclasses to actually use the parameter selector when it is provided
-    # TODO: Perhaps we can find a more elegant way to enforce this by design
-    _uses_parameter_names: ClassVar[bool] = False
-
-    parameter_selector: ParameterSelectorProtocol | None = field(
-        default=None, converter=optional(to_parameter_selector), kw_only=True
-    )
-    """An optional selector to specify which parameters are to be considered."""
-
-    def get_parameter_names(self, searchspace: SearchSpace) -> tuple[str, ...] | None:
-        """Get the names of the parameters to be considered."""
-        if self.parameter_selector is None:
-            return None
-
-        return tuple(
-            p.name for p in searchspace.parameters if self.parameter_selector(p)
-        )
-
-    def __attrs_post_init__(self):
-        if self.parameter_selector is not None and not self._uses_parameter_names:
-            raise AssertionError(
-                f"A `parameter_selector` was provided to "
-                f"`{type(self).__name__}`, but the class does not set "
-                f"`_uses_parameter_names = True`. Subclasses that accept a "
-                f"parameter selector must explicitly set this flag to confirm "
-                f"they actually use the selected parameter names."
-            )

baybe/surrogates/gaussian_process/components/kernel.py

@@ -2,18 +2,24 @@
 
 from __future__ import annotations
 
+from abc import ABC, abstractmethod
+from collections.abc import Iterable
 from functools import partial
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, ClassVar
 
 from attrs import define, field
+from attrs.converters import optional
 from attrs.validators import is_callable
 from typing_extensions import override
 
+from baybe.exceptions import IncompatibleSearchSpaceError
 from baybe.kernels.base import Kernel
-from baybe.kernels.composite import ProductKernel
 from baybe.parameters.categorical import TaskParameter
+from baybe.parameters.enum import ParameterKind
 from baybe.parameters.selectors import (
+    ParameterSelectorProtocol,
     TypeSelector,
+    to_parameter_selector,
 )
 from baybe.searchspace.core import SearchSpace
 from baybe.surrogates.gaussian_process.components.generic import (
@@ -27,6 +33,8 @@
     from gpytorch.kernels import Kernel as GPyTorchKernel
     from torch import Tensor
 
+    from baybe.parameters.base import Parameter
+
     KernelFactoryProtocol = GPComponentFactoryProtocol[Kernel | GPyTorchKernel]
     PlainKernelFactory = PlainGPComponentFactory[Kernel | GPyTorchKernel]
 else:
@@ -35,6 +43,80 @@
     PlainKernelFactory = PlainGPComponentFactory[Kernel]
 
 
+@define
+class _KernelFactory(KernelFactoryProtocol, ABC):
+    """Base class for kernel factories."""
+
+    # For internal use only: sanity check mechanism to remind developers of new
+    # factories to actually use the parameter selector when it is provided
+    # TODO: Perhaps we can find a more elegant way to enforce this by design
+    _uses_parameter_names: ClassVar[bool] = False
+
+    supported_parameter_kinds: ClassVar[ParameterKind] = ParameterKind.REGULAR
+    """The parameter kinds supported by the kernel factory."""
+
+    parameter_selector: ParameterSelectorProtocol | None = field(
+        default=None, converter=optional(to_parameter_selector)
+    )
+    """An optional selector to specify which parameters are considered by the kernel."""
+
+    def get_parameter_names(self, searchspace: SearchSpace) -> tuple[str, ...] | None:
+        """Get the names of the parameters to be considered by the kernel."""
+        if self.parameter_selector is None:
+            return None
+
+        return tuple(
+            p.name for p in searchspace.parameters if self.parameter_selector(p)
+        )
+
+    def _validate_parameter_kinds(self, parameters: Iterable[Parameter]) -> None:
+        """Validate that the given parameters are supported by the factory.
+
+        Args:
+            parameters: The parameters to validate.
+
+        Raises:
+            IncompatibleSearchSpaceError: If unsupported parameter kinds are found.
+        """
+        if unsupported := [
+            p.name for p in parameters if not (p.kind & self.supported_parameter_kinds)
+        ]:
+            raise IncompatibleSearchSpaceError(
+                f"'{type(self).__name__}' does not support parameter kind(s) for "
+                f"parameter(s) {unsupported}. Supported kinds: "
+                f"{self.supported_parameter_kinds}."
+            )
+
+    @override
+    def __call__(
+        self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
+    ) -> Kernel:
+        """Construct the kernel, validating parameter kinds before construction."""
+        if self.parameter_selector is not None:
+            params = [p for p in searchspace.parameters if self.parameter_selector(p)]
+        else:
+            params = list(searchspace.parameters)
+        self._validate_parameter_kinds(params)
+
+        return self._make(searchspace, train_x, train_y)
+
+    @abstractmethod
+    def _make(
+        self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
+    ) -> Kernel:
+        """Construct the kernel."""
+
+    def __attrs_post_init__(self):
+        if self.parameter_selector is not None and not self._uses_parameter_names:
+            raise AssertionError(
+                f"A `parameter_selector` was provided to "
+                f"`{type(self).__name__}`, but the class does not set "
+                f"`_uses_parameter_names = True`. Subclasses that accept a "
+                f"parameter selector must explicitly set this flag to confirm "
+                f"they actually use the selected parameter names."
+            )
+
+
 @define
 class ICMKernelFactory(KernelFactoryProtocol):
     """A kernel factory that constructs an ICM kernel for transfer learning.
@@ -78,4 +160,8 @@ def __call__(
     ) -> Kernel:
         base_kernel = self.base_kernel_factory(searchspace, train_x, train_y)
         task_kernel = self.task_kernel_factory(searchspace, train_x, train_y)
-        return ProductKernel([base_kernel, task_kernel])
+        if isinstance(base_kernel, Kernel):
+            base_kernel = base_kernel.to_gpytorch(searchspace)
+        if isinstance(task_kernel, Kernel):
+            task_kernel = task_kernel.to_gpytorch(searchspace)
+        return base_kernel * task_kernel

baybe/surrogates/gaussian_process/presets/__init__.py

@@ -7,6 +7,9 @@
     BayBEMeanFactory,
 )
 
+# Chen preset
+from baybe.surrogates.gaussian_process.presets.chen import CHENKernelFactory
+
 # Core
 from baybe.surrogates.gaussian_process.presets.core import GaussianProcessPreset
 
@@ -31,6 +34,8 @@
     "BayBEKernelFactory",
     "BayBELikelihoodFactory",
     "BayBEMeanFactory",
+    # Chen preset
+    "CHENKernelFactory",
     # EDBO preset
     "EDBOKernelFactory",
     "EDBOLikelihoodFactory",

baybe/surrogates/gaussian_process/presets/baybe.py

@@ -10,14 +10,14 @@
 from baybe.kernels.base import Kernel
 from baybe.kernels.basic import IndexKernel
 from baybe.parameters.categorical import TaskParameter
+from baybe.parameters.enum import ParameterKind
 from baybe.parameters.selectors import (
     ParameterSelectorProtocol,
     TypeSelector,
-    _ParameterSelectorMixin,
     to_parameter_selector,
 )
 from baybe.searchspace.core import SearchSpace
-from baybe.surrogates.gaussian_process.components.kernel import KernelFactoryProtocol
+from baybe.surrogates.gaussian_process.components.kernel import _KernelFactory
 from baybe.surrogates.gaussian_process.components.mean import LazyConstantMeanFactory
 from baybe.surrogates.gaussian_process.presets.edbo_smoothed import (
     SmoothedEDBOKernelFactory,
@@ -29,11 +29,16 @@
 
 
 @define
-class BayBEKernelFactory(KernelFactoryProtocol):
+class BayBEKernelFactory(_KernelFactory):
     """The default kernel factory for Gaussian process surrogates."""
 
+    supported_parameter_kinds: ClassVar[ParameterKind] = (
+        ParameterKind.REGULAR | ParameterKind.TASK
+    )
+    # See base class.
+
     @override
-    def __call__(
+    def _make(
         self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
     ) -> Kernel:
         from baybe.surrogates.gaussian_process.components.kernel import ICMKernelFactory
@@ -48,20 +53,23 @@ def __call__(
 
 
 @define
-class BayBETaskKernelFactory(KernelFactoryProtocol, _ParameterSelectorMixin):
+class BayBETaskKernelFactory(_KernelFactory):
     """The factory providing the default task kernel for Gaussian process surrogates."""
 
     _uses_parameter_names: ClassVar[bool] = True
     # See base class.
 
+    supported_parameter_kinds: ClassVar[ParameterKind] = ParameterKind.TASK
+    # See base class.
+
     parameter_selector: ParameterSelectorProtocol | None = field(
         factory=lambda: TypeSelector([TaskParameter]),
         converter=to_parameter_selector,
     )
     # TODO: Reuse base attribute (https://github.com/python-attrs/attrs/pull/1429)
 
     @override
-    def __call__(
+    def _make(
         self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
     ) -> Kernel:
         return IndexKernel(