From 59e969792d5118ae13fc67f8514e09f3bff083b5 Mon Sep 17 00:00:00 2001 From: AdrianSosic Date: Wed, 15 Apr 2026 11:20:16 +0200 Subject: [PATCH 01/13] Move _ParameterSelectorMixin to components/kernel.py --- baybe/parameters/selectors.py | 38 +----------------- .../gaussian_process/components/kernel.py | 39 ++++++++++++++++++- .../gaussian_process/presets/baybe.py | 6 ++- .../gaussian_process/presets/edbo.py | 6 ++- .../gaussian_process/presets/edbo_smoothed.py | 6 ++- 5 files changed, 51 insertions(+), 44 deletions(-) diff --git a/baybe/parameters/selectors.py b/baybe/parameters/selectors.py index 6b1a4ae16d..e72b6fe4d9 100644 --- a/baybe/parameters/selectors.py +++ b/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." - ) diff --git a/baybe/surrogates/gaussian_process/components/kernel.py b/baybe/surrogates/gaussian_process/components/kernel.py index ca4da066e7..bc9b215465 100644 --- a/baybe/surrogates/gaussian_process/components/kernel.py +++ b/baybe/surrogates/gaussian_process/components/kernel.py @@ -3,9 +3,10 @@ from __future__ import annotations 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 @@ -13,7 +14,9 @@ from baybe.kernels.composite import ProductKernel from baybe.parameters.categorical import TaskParameter from baybe.parameters.selectors import ( + ParameterSelectorProtocol, TypeSelector, + to_parameter_selector, ) from baybe.searchspace.core import SearchSpace from baybe.surrogates.gaussian_process.components.generic import ( @@ -35,6 +38,40 @@ PlainKernelFactory = PlainGPComponentFactory[Kernel] +@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." + ) + + @define class ICMKernelFactory(KernelFactoryProtocol): """A kernel factory that constructs an ICM kernel for transfer learning. diff --git a/baybe/surrogates/gaussian_process/presets/baybe.py b/baybe/surrogates/gaussian_process/presets/baybe.py index bc820ae52b..cd19242998 100644 --- a/baybe/surrogates/gaussian_process/presets/baybe.py +++ b/baybe/surrogates/gaussian_process/presets/baybe.py @@ -13,11 +13,13 @@ 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 ( + KernelFactoryProtocol, + _ParameterSelectorMixin, +) from baybe.surrogates.gaussian_process.components.mean import LazyConstantMeanFactory from baybe.surrogates.gaussian_process.presets.edbo_smoothed import ( SmoothedEDBOKernelFactory, diff --git a/baybe/surrogates/gaussian_process/presets/edbo.py b/baybe/surrogates/gaussian_process/presets/edbo.py index 36611c258e..8204913a41 100644 --- a/baybe/surrogates/gaussian_process/presets/edbo.py +++ b/baybe/surrogates/gaussian_process/presets/edbo.py @@ -16,13 +16,15 @@ from baybe.parameters.selectors import ( ParameterSelectorProtocol, TypeSelector, - _ParameterSelectorMixin, to_parameter_selector, ) from baybe.parameters.substance import SubstanceParameter from baybe.priors.basic import GammaPrior from baybe.searchspace.discrete import SubspaceDiscrete -from baybe.surrogates.gaussian_process.components.kernel import KernelFactoryProtocol +from baybe.surrogates.gaussian_process.components.kernel import ( + KernelFactoryProtocol, + _ParameterSelectorMixin, +) from baybe.surrogates.gaussian_process.components.likelihood import ( LikelihoodFactoryProtocol, ) diff --git a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py index eb8d57491a..77fc74354b 100644 --- a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py +++ b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py @@ -15,11 +15,13 @@ from baybe.parameters.selectors import ( ParameterSelectorProtocol, TypeSelector, - _ParameterSelectorMixin, to_parameter_selector, ) from baybe.priors.basic import GammaPrior -from baybe.surrogates.gaussian_process.components.kernel import KernelFactoryProtocol +from baybe.surrogates.gaussian_process.components.kernel import ( + KernelFactoryProtocol, + _ParameterSelectorMixin, +) from baybe.surrogates.gaussian_process.components.likelihood import ( LikelihoodFactoryProtocol, ) From bcc51bbb04e8f80220db64f5f63e28e85b2b157c Mon Sep 17 00:00:00 2001 From: AdrianSosic Date: Wed, 15 Apr 2026 13:23:08 +0200 Subject: [PATCH 02/13] Turn mixin class back into ABC --- .../gaussian_process/components/kernel.py | 20 +++++++++++++------ .../gaussian_process/presets/baybe.py | 4 ++-- .../gaussian_process/presets/edbo.py | 5 ++--- .../gaussian_process/presets/edbo_smoothed.py | 5 ++--- docs/conf.py | 2 +- 5 files changed, 21 insertions(+), 15 deletions(-) diff --git a/baybe/surrogates/gaussian_process/components/kernel.py b/baybe/surrogates/gaussian_process/components/kernel.py index bc9b215465..5d9aecfc35 100644 --- a/baybe/surrogates/gaussian_process/components/kernel.py +++ b/baybe/surrogates/gaussian_process/components/kernel.py @@ -2,6 +2,7 @@ from __future__ import annotations +from abc import ABC, abstractmethod from functools import partial from typing import TYPE_CHECKING, ClassVar @@ -39,21 +40,21 @@ @define -class _ParameterSelectorMixin: - """A mixin class to enable parameter selection.""" +class _KernelFactory(KernelFactoryProtocol, ABC): + """Base class for kernel factories.""" # For internal use only: sanity check mechanism to remind developers of new - # subclasses to actually use the parameter selector when it is provided + # 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 parameter_selector: ParameterSelectorProtocol | None = field( - default=None, converter=optional(to_parameter_selector), kw_only=True + default=None, converter=optional(to_parameter_selector) ) - """An optional selector to specify which parameters are to be considered.""" + """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.""" + """Get the names of the parameters to be considered by the kernel.""" if self.parameter_selector is None: return None @@ -71,6 +72,13 @@ def __attrs_post_init__(self): f"they actually use the selected parameter names." ) + @override + @abstractmethod + def __call__( + self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor + ) -> Kernel | GPyTorchKernel: + pass + @define class ICMKernelFactory(KernelFactoryProtocol): diff --git a/baybe/surrogates/gaussian_process/presets/baybe.py b/baybe/surrogates/gaussian_process/presets/baybe.py index cd19242998..be20f5cd19 100644 --- a/baybe/surrogates/gaussian_process/presets/baybe.py +++ b/baybe/surrogates/gaussian_process/presets/baybe.py @@ -18,7 +18,7 @@ from baybe.searchspace.core import SearchSpace from baybe.surrogates.gaussian_process.components.kernel import ( KernelFactoryProtocol, - _ParameterSelectorMixin, + _KernelFactory, ) from baybe.surrogates.gaussian_process.components.mean import LazyConstantMeanFactory from baybe.surrogates.gaussian_process.presets.edbo_smoothed import ( @@ -50,7 +50,7 @@ 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 diff --git a/baybe/surrogates/gaussian_process/presets/edbo.py b/baybe/surrogates/gaussian_process/presets/edbo.py index 8204913a41..6e0dd06eeb 100644 --- a/baybe/surrogates/gaussian_process/presets/edbo.py +++ b/baybe/surrogates/gaussian_process/presets/edbo.py @@ -22,8 +22,7 @@ from baybe.priors.basic import GammaPrior from baybe.searchspace.discrete import SubspaceDiscrete from baybe.surrogates.gaussian_process.components.kernel import ( - KernelFactoryProtocol, - _ParameterSelectorMixin, + _KernelFactory, ) from baybe.surrogates.gaussian_process.components.likelihood import ( LikelihoodFactoryProtocol, @@ -58,7 +57,7 @@ def _contains_encoding( @define -class EDBOKernelFactory(KernelFactoryProtocol, _ParameterSelectorMixin): +class EDBOKernelFactory(_KernelFactory): """A factory providing EDBO kernels. References: diff --git a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py index 77fc74354b..eb7eca36da 100644 --- a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py +++ b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py @@ -19,8 +19,7 @@ ) from baybe.priors.basic import GammaPrior from baybe.surrogates.gaussian_process.components.kernel import ( - KernelFactoryProtocol, - _ParameterSelectorMixin, + _KernelFactory, ) from baybe.surrogates.gaussian_process.components.likelihood import ( LikelihoodFactoryProtocol, @@ -40,7 +39,7 @@ @define -class SmoothedEDBOKernelFactory(KernelFactoryProtocol, _ParameterSelectorMixin): +class SmoothedEDBOKernelFactory(_KernelFactory): """A factory providing smoothed versions of EDBO kernels. Takes the low and high dimensional limits of diff --git a/docs/conf.py b/docs/conf.py index 5251bb8654..6096369714 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -150,7 +150,7 @@ ("py:class", "baybe.acquisition.acqfs._ExpectedHypervolumeImprovement"), ("py:class", "baybe.settings._SlottedContextDecorator"), ("py:class", "baybe.surrogates.gaussian_process.components.PlainKernelFactory"), - ("py:class", "baybe.parameters.selectors._ParameterSelectorMixin"), + ("py:class", "baybe.surrogates.gaussian_process.components.kernel._KernelFactory"), # Deprecation ("py:.*", "baybe.targets._deprecated.*"), ] From 9c6051c771e5935b2f94320bc5bdbdfad8ef7671 Mon Sep 17 00:00:00 2001 From: AdrianSosic Date: Thu, 2 Apr 2026 09:45:50 +0200 Subject: [PATCH 03/13] Add ParameterKind Flag enum and Parameter.kind property --- baybe/parameters/__init__.py | 2 ++ baybe/parameters/base.py | 8 ++++++++ baybe/parameters/enum.py | 34 +++++++++++++++++++++++++++++++++- 3 files changed, 43 insertions(+), 1 deletion(-) diff --git a/baybe/parameters/__init__.py b/baybe/parameters/__init__.py index 93e62b6ee9..6fe2c0d8f4 100644 --- a/baybe/parameters/__init__.py +++ b/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", diff --git a/baybe/parameters/base.py b/baybe/parameters/base.py index 2d4df2bc77..7986a21124 100644 --- a/baybe/parameters/base.py +++ b/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, ...]: diff --git a/baybe/parameters/enum.py b/baybe/parameters/enum.py index 3161f67dc7..07e9f9fb34 100644 --- a/baybe/parameters/enum.py +++ b/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): From 447cf6afd10ce3f4aee221239a88946b11913d31 Mon Sep 17 00:00:00 2001 From: AdrianSosic Date: Wed, 15 Apr 2026 13:26:51 +0200 Subject: [PATCH 04/13] Wire parameter kind validation into kernel factories via template method --- .../gaussian_process/components/kernel.py | 52 ++++++++++++++++--- .../gaussian_process/presets/baybe.py | 20 ++++--- .../gaussian_process/presets/edbo.py | 2 +- .../gaussian_process/presets/edbo_smoothed.py | 2 +- 4 files changed, 60 insertions(+), 16 deletions(-) diff --git a/baybe/surrogates/gaussian_process/components/kernel.py b/baybe/surrogates/gaussian_process/components/kernel.py index 5d9aecfc35..350a8405bf 100644 --- a/baybe/surrogates/gaussian_process/components/kernel.py +++ b/baybe/surrogates/gaussian_process/components/kernel.py @@ -3,6 +3,7 @@ from __future__ import annotations from abc import ABC, abstractmethod +from collections.abc import Iterable from functools import partial from typing import TYPE_CHECKING, ClassVar @@ -11,9 +12,11 @@ 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, @@ -31,6 +34,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: @@ -48,6 +53,9 @@ class _KernelFactory(KernelFactoryProtocol, ABC): # 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) ) @@ -62,6 +70,43 @@ def get_parameter_names(self, searchspace: SearchSpace) -> tuple[str, ...] | Non 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( @@ -72,13 +117,6 @@ def __attrs_post_init__(self): f"they actually use the selected parameter names." ) - @override - @abstractmethod - def __call__( - self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor - ) -> Kernel | GPyTorchKernel: - pass - @define class ICMKernelFactory(KernelFactoryProtocol): diff --git a/baybe/surrogates/gaussian_process/presets/baybe.py b/baybe/surrogates/gaussian_process/presets/baybe.py index be20f5cd19..e96a08e6cc 100644 --- a/baybe/surrogates/gaussian_process/presets/baybe.py +++ b/baybe/surrogates/gaussian_process/presets/baybe.py @@ -10,16 +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, to_parameter_selector, ) from baybe.searchspace.core import SearchSpace -from baybe.surrogates.gaussian_process.components.kernel import ( - KernelFactoryProtocol, - _KernelFactory, -) +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, @@ -31,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 @@ -56,6 +59,9 @@ class BayBETaskKernelFactory(_KernelFactory): _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, @@ -63,7 +69,7 @@ class BayBETaskKernelFactory(_KernelFactory): # 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( diff --git a/baybe/surrogates/gaussian_process/presets/edbo.py b/baybe/surrogates/gaussian_process/presets/edbo.py index 6e0dd06eeb..758c643770 100644 --- a/baybe/surrogates/gaussian_process/presets/edbo.py +++ b/baybe/surrogates/gaussian_process/presets/edbo.py @@ -75,7 +75,7 @@ class EDBOKernelFactory(_KernelFactory): # 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: effective_dims = train_x.shape[-1] - len( diff --git a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py index eb7eca36da..ad5bcf7f53 100644 --- a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py +++ b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py @@ -57,7 +57,7 @@ class SmoothedEDBOKernelFactory(_KernelFactory): # 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: effective_dims = train_x.shape[-1] - len( From 5845f0fbf74c1963da00e77859accf5ca9ca664b Mon Sep 17 00:00:00 2001 From: AdrianSosic Date: Thu, 2 Apr 2026 13:37:43 +0200 Subject: [PATCH 05/13] Remove dead TaskParameter filtering in kernel factories The new parameter kind validation step guards us --- baybe/surrogates/gaussian_process/presets/edbo.py | 4 +--- baybe/surrogates/gaussian_process/presets/edbo_smoothed.py | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/baybe/surrogates/gaussian_process/presets/edbo.py b/baybe/surrogates/gaussian_process/presets/edbo.py index 758c643770..ecf24de106 100644 --- a/baybe/surrogates/gaussian_process/presets/edbo.py +++ b/baybe/surrogates/gaussian_process/presets/edbo.py @@ -78,9 +78,7 @@ class EDBOKernelFactory(_KernelFactory): def _make( self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor ) -> Kernel: - effective_dims = train_x.shape[-1] - len( - [p for p in searchspace.parameters if isinstance(p, TaskParameter)] - ) + effective_dims = train_x.shape[-1] switching_condition = _contains_encoding( searchspace.discrete, _EDBO_ENCODINGS diff --git a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py index ad5bcf7f53..89cf9a9ef2 100644 --- a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py +++ b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py @@ -60,9 +60,7 @@ class SmoothedEDBOKernelFactory(_KernelFactory): def _make( self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor ) -> Kernel: - effective_dims = train_x.shape[-1] - len( - [p for p in searchspace.parameters if isinstance(p, TaskParameter)] - ) + effective_dims = train_x.shape[-1] # Interpolate prior moments linearly between low D and high D regime. # The high D regime itself is the average of the EDBO OHE and Mordred regime. From 741997b9ffda88ecbbdaeecbe8f1acc5157e2511 Mon Sep 17 00:00:00 2001 From: AdrianSosic Date: Thu, 2 Apr 2026 13:46:28 +0200 Subject: [PATCH 06/13] Update CHANGELOG.md --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 680f11da28..c7715778a3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,6 +19,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `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 From 762b8b90450d60a6a4b973a6402b7d02874ebcb0 Mon Sep 17 00:00:00 2001 From: AdrianSosic Date: Thu, 16 Apr 2026 09:08:33 +0200 Subject: [PATCH 07/13] Add parameter support test --- tests/test_kernel_factories.py | 75 ++++++++++++++++++++++++++++++++++ 1 file changed, 75 insertions(+) create mode 100644 tests/test_kernel_factories.py diff --git a/tests/test_kernel_factories.py b/tests/test_kernel_factories.py new file mode 100644 index 0000000000..6a8acda6a6 --- /dev/null +++ b/tests/test_kernel_factories.py @@ -0,0 +1,75 @@ +"""Tests for kernel factories.""" + +from contextlib import nullcontext + +import pytest +import torch +from pytest import param + +from baybe.exceptions import IncompatibleSearchSpaceError +from baybe.parameters.categorical import CategoricalParameter, TaskParameter +from baybe.parameters.numerical import ( + NumericalContinuousParameter, + NumericalDiscreteParameter, +) +from baybe.searchspace.core import SearchSpace +from baybe.surrogates.gaussian_process.presets.baybe import ( + BayBEKernelFactory, + BayBENumericalKernelFactory, + BayBETaskKernelFactory, +) + +# A selector that accepts all parameters +_SELECT_ALL = lambda parameter: True # noqa: E731 + + +@pytest.mark.parametrize( + ("factory", "parameters", "error"), + [ + param( + BayBENumericalKernelFactory(parameter_selector=_SELECT_ALL), + [TaskParameter("task", ["t1", "t2"])], + IncompatibleSearchSpaceError, + id="regular_rejects_task", + ), + param( + BayBETaskKernelFactory(parameter_selector=_SELECT_ALL), + [CategoricalParameter("cat", ["a", "b"])], + IncompatibleSearchSpaceError, + id="task_rejects_categorical", + ), + param( + BayBETaskKernelFactory(parameter_selector=_SELECT_ALL), + [NumericalDiscreteParameter("num", [1, 2, 3])], + IncompatibleSearchSpaceError, + id="task_rejects_numerical_discrete", + ), + param( + BayBETaskKernelFactory(parameter_selector=_SELECT_ALL), + [NumericalContinuousParameter("cont", (0, 1))], + IncompatibleSearchSpaceError, + id="task_rejects_numerical_continuous", + ), + param( + BayBEKernelFactory(), + [ + NumericalContinuousParameter("cont", (0, 1)), + TaskParameter("task", ["t1", "t2"]), + ], + None, + id="combined_accepts_both", + ), + ], +) +def test_factory_parameter_kind_validation(factory, parameters, error): + """Factories reject unsupported parameter kinds and accept supported ones.""" + ss = SearchSpace.from_product(parameters) + train_x = torch.zeros(2, len(ss.comp_rep_columns)) + train_y = torch.zeros(2, 1) + + with ( + nullcontext() + if error is None + else pytest.raises(error, match="does not support") + ): + factory(ss, train_x, train_y) From d6f3128998031059f0532d30faa2f4cb8af2ab38 Mon Sep 17 00:00:00 2001 From: AdrianSosic Date: Thu, 16 Apr 2026 10:11:26 +0200 Subject: [PATCH 08/13] Make ICMKernelFactory gpytorch-compatible --- baybe/surrogates/gaussian_process/components/kernel.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/baybe/surrogates/gaussian_process/components/kernel.py b/baybe/surrogates/gaussian_process/components/kernel.py index 350a8405bf..65b6e5e454 100644 --- a/baybe/surrogates/gaussian_process/components/kernel.py +++ b/baybe/surrogates/gaussian_process/components/kernel.py @@ -14,7 +14,6 @@ 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 ( @@ -161,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 From 6ee7b701d30a55eecffe5e041ae5d624177db89c Mon Sep 17 00:00:00 2001 From: AdrianSosic Date: Wed, 22 Apr 2026 15:18:41 +0200 Subject: [PATCH 09/13] Move __attrs_post_init__ to the right place --- .../gaussian_process/components/kernel.py | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/baybe/surrogates/gaussian_process/components/kernel.py b/baybe/surrogates/gaussian_process/components/kernel.py index 65b6e5e454..10318d993b 100644 --- a/baybe/surrogates/gaussian_process/components/kernel.py +++ b/baybe/surrogates/gaussian_process/components/kernel.py @@ -60,6 +60,16 @@ class _KernelFactory(KernelFactoryProtocol, ABC): ) """An optional selector to specify which parameters are considered by 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." + ) + 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: @@ -106,16 +116,6 @@ def _make( ) -> 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): From fad7525645b7904c7e65921de73c3f65b52b91e1 Mon Sep 17 00:00:00 2001 From: AdrianSosic Date: Fri, 24 Apr 2026 09:53:46 +0200 Subject: [PATCH 10/13] Rename _KernelFactory to _PureKernelFactory --- baybe/surrogates/gaussian_process/components/kernel.py | 4 ++-- baybe/surrogates/gaussian_process/presets/baybe.py | 6 +++--- baybe/surrogates/gaussian_process/presets/edbo.py | 4 ++-- baybe/surrogates/gaussian_process/presets/edbo_smoothed.py | 4 ++-- 4 files changed, 9 insertions(+), 9 deletions(-) diff --git a/baybe/surrogates/gaussian_process/components/kernel.py b/baybe/surrogates/gaussian_process/components/kernel.py index 10318d993b..319b729c9d 100644 --- a/baybe/surrogates/gaussian_process/components/kernel.py +++ b/baybe/surrogates/gaussian_process/components/kernel.py @@ -44,8 +44,8 @@ @define -class _KernelFactory(KernelFactoryProtocol, ABC): - """Base class for kernel factories.""" +class _PureKernelFactory(KernelFactoryProtocol, ABC): + """Base class for pure 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 diff --git a/baybe/surrogates/gaussian_process/presets/baybe.py b/baybe/surrogates/gaussian_process/presets/baybe.py index e96a08e6cc..00852d8439 100644 --- a/baybe/surrogates/gaussian_process/presets/baybe.py +++ b/baybe/surrogates/gaussian_process/presets/baybe.py @@ -17,7 +17,7 @@ to_parameter_selector, ) from baybe.searchspace.core import SearchSpace -from baybe.surrogates.gaussian_process.components.kernel import _KernelFactory +from baybe.surrogates.gaussian_process.components.kernel import _PureKernelFactory from baybe.surrogates.gaussian_process.components.mean import LazyConstantMeanFactory from baybe.surrogates.gaussian_process.presets.edbo_smoothed import ( SmoothedEDBOKernelFactory, @@ -29,7 +29,7 @@ @define -class BayBEKernelFactory(_KernelFactory): +class BayBEKernelFactory(_PureKernelFactory): """The default kernel factory for Gaussian process surrogates.""" supported_parameter_kinds: ClassVar[ParameterKind] = ( @@ -53,7 +53,7 @@ def _make( @define -class BayBETaskKernelFactory(_KernelFactory): +class BayBETaskKernelFactory(_PureKernelFactory): """The factory providing the default task kernel for Gaussian process surrogates.""" _uses_parameter_names: ClassVar[bool] = True diff --git a/baybe/surrogates/gaussian_process/presets/edbo.py b/baybe/surrogates/gaussian_process/presets/edbo.py index ecf24de106..6db0af8a30 100644 --- a/baybe/surrogates/gaussian_process/presets/edbo.py +++ b/baybe/surrogates/gaussian_process/presets/edbo.py @@ -22,7 +22,7 @@ from baybe.priors.basic import GammaPrior from baybe.searchspace.discrete import SubspaceDiscrete from baybe.surrogates.gaussian_process.components.kernel import ( - _KernelFactory, + _PureKernelFactory, ) from baybe.surrogates.gaussian_process.components.likelihood import ( LikelihoodFactoryProtocol, @@ -57,7 +57,7 @@ def _contains_encoding( @define -class EDBOKernelFactory(_KernelFactory): +class EDBOKernelFactory(_PureKernelFactory): """A factory providing EDBO kernels. References: diff --git a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py index 89cf9a9ef2..1c2718adb6 100644 --- a/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py +++ b/baybe/surrogates/gaussian_process/presets/edbo_smoothed.py @@ -19,7 +19,7 @@ ) from baybe.priors.basic import GammaPrior from baybe.surrogates.gaussian_process.components.kernel import ( - _KernelFactory, + _PureKernelFactory, ) from baybe.surrogates.gaussian_process.components.likelihood import ( LikelihoodFactoryProtocol, @@ -39,7 +39,7 @@ @define -class SmoothedEDBOKernelFactory(_KernelFactory): +class SmoothedEDBOKernelFactory(_PureKernelFactory): """A factory providing smoothed versions of EDBO kernels. Takes the low and high dimensional limits of From 58130853ae1a22ba87a8fe4693e4673402cf9267 Mon Sep 17 00:00:00 2001 From: AdrianSosic Date: Fri, 24 Apr 2026 09:56:12 +0200 Subject: [PATCH 11/13] Add intermediate _MetaKernelFactory base class --- .../gaussian_process/components/kernel.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/baybe/surrogates/gaussian_process/components/kernel.py b/baybe/surrogates/gaussian_process/components/kernel.py index 319b729c9d..a14c0f677a 100644 --- a/baybe/surrogates/gaussian_process/components/kernel.py +++ b/baybe/surrogates/gaussian_process/components/kernel.py @@ -118,7 +118,18 @@ def _make( @define -class ICMKernelFactory(KernelFactoryProtocol): +class _MetaKernelFactory(KernelFactoryProtocol, ABC): + """Base class for meta kernel factories that orchestrate other kernel factories.""" + + @override + @abstractmethod + def __call__( + self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor + ) -> Kernel: ... + + +@define +class ICMKernelFactory(_MetaKernelFactory): """A kernel factory that constructs an ICM kernel for transfer learning. ICM: Intrinsic Coregionalization Model :cite:p:`NIPS2007_66368270` From 616109744c465967dcfec8b2c5b7cd7857497551 Mon Sep 17 00:00:00 2001 From: AdrianSosic Date: Fri, 24 Apr 2026 09:59:50 +0200 Subject: [PATCH 12/13] Use regex for private classes in conf.py nitpick ignore --- docs/conf.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/docs/conf.py b/docs/conf.py index 6096369714..11ff636a77 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -146,11 +146,9 @@ (r"py:obj", "baybe.utils.boolean.UncertainBool.*"), ("py:obj", "baybe.targets.botorch.*"), ("py:obj", "baybe.objectives.botorch.*"), - ("py:class", "baybe.parameters.base._DiscreteLabelLikeParameter"), - ("py:class", "baybe.acquisition.acqfs._ExpectedHypervolumeImprovement"), - ("py:class", "baybe.settings._SlottedContextDecorator"), ("py:class", "baybe.surrogates.gaussian_process.components.PlainKernelFactory"), - ("py:class", "baybe.surrogates.gaussian_process.components.kernel._KernelFactory"), + # Private classes + (r"py:class", r"baybe\..*\._.*"), # Deprecation ("py:.*", "baybe.targets._deprecated.*"), ] From c92d3c7a519527a19802548b907e51f84b2f3614 Mon Sep 17 00:00:00 2001 From: AdrianSosic Date: Fri, 24 Apr 2026 10:05:00 +0200 Subject: [PATCH 13/13] Make ParameterKind machinery private Until we are certain that we keep it --- baybe/parameters/__init__.py | 2 -- baybe/parameters/base.py | 8 ++++---- baybe/parameters/enum.py | 8 ++++---- baybe/surrogates/gaussian_process/components/kernel.py | 10 ++++++---- baybe/surrogates/gaussian_process/presets/baybe.py | 8 ++++---- 5 files changed, 18 insertions(+), 18 deletions(-) diff --git a/baybe/parameters/__init__.py b/baybe/parameters/__init__.py index 6fe2c0d8f4..93e62b6ee9 100644 --- a/baybe/parameters/__init__.py +++ b/baybe/parameters/__init__.py @@ -5,7 +5,6 @@ from baybe.parameters.enum import ( CategoricalEncoding, CustomEncoding, - ParameterKind, SubstanceEncoding, ) from baybe.parameters.numerical import ( @@ -23,7 +22,6 @@ "MeasurableMetadata", "NumericalContinuousParameter", "NumericalDiscreteParameter", - "ParameterKind", "SubstanceEncoding", "SubstanceParameter", "TaskParameter", diff --git a/baybe/parameters/base.py b/baybe/parameters/base.py index 7986a21124..8fbd489a41 100644 --- a/baybe/parameters/base.py +++ b/baybe/parameters/base.py @@ -21,7 +21,7 @@ from baybe.utils.metadata import MeasurableMetadata, to_metadata if TYPE_CHECKING: - from baybe.parameters.enum import ParameterKind + from baybe.parameters.enum import _ParameterKind from baybe.searchspace.continuous import SubspaceContinuous from baybe.searchspace.core import SearchSpace from baybe.searchspace.discrete import SubspaceDiscrete @@ -79,11 +79,11 @@ def is_discrete(self) -> bool: return isinstance(self, DiscreteParameter) @property - def kind(self) -> ParameterKind: + def _kind(self) -> _ParameterKind: """The kind of the parameter.""" - from baybe.parameters.enum import ParameterKind + from baybe.parameters.enum import _ParameterKind - return ParameterKind.from_parameter(self) + return _ParameterKind.from_parameter(self) @property @abstractmethod diff --git a/baybe/parameters/enum.py b/baybe/parameters/enum.py index 07e9f9fb34..622fa4af58 100644 --- a/baybe/parameters/enum.py +++ b/baybe/parameters/enum.py @@ -9,7 +9,7 @@ from baybe.parameters.base import Parameter -class ParameterKind(Flag): +class _ParameterKind(Flag): """Flag enum encoding the kind of a parameter. Can be used to express compatibility (e.g. Gaussian process kernel factories) @@ -26,13 +26,13 @@ class ParameterKind(Flag): """Fidelity parameter for multi-fidelity modelling.""" @staticmethod - def from_parameter(parameter: Parameter) -> ParameterKind: + 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 + return _ParameterKind.TASK + return _ParameterKind.REGULAR class ParameterEncoding(Enum): diff --git a/baybe/surrogates/gaussian_process/components/kernel.py b/baybe/surrogates/gaussian_process/components/kernel.py index a14c0f677a..8b6ddff280 100644 --- a/baybe/surrogates/gaussian_process/components/kernel.py +++ b/baybe/surrogates/gaussian_process/components/kernel.py @@ -15,7 +15,7 @@ from baybe.exceptions import IncompatibleSearchSpaceError from baybe.kernels.base import Kernel from baybe.parameters.categorical import TaskParameter -from baybe.parameters.enum import ParameterKind +from baybe.parameters.enum import _ParameterKind from baybe.parameters.selectors import ( ParameterSelectorProtocol, TypeSelector, @@ -52,7 +52,7 @@ class _PureKernelFactory(KernelFactoryProtocol, ABC): # 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 + _supported_parameter_kinds: ClassVar[_ParameterKind] = _ParameterKind.REGULAR """The parameter kinds supported by the kernel factory.""" parameter_selector: ParameterSelectorProtocol | None = field( @@ -89,12 +89,14 @@ def _validate_parameter_kinds(self, parameters: Iterable[Parameter]) -> None: IncompatibleSearchSpaceError: If unsupported parameter kinds are found. """ if unsupported := [ - p.name for p in parameters if not (p.kind & self.supported_parameter_kinds) + 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}." + f"{self._supported_parameter_kinds}." ) @override diff --git a/baybe/surrogates/gaussian_process/presets/baybe.py b/baybe/surrogates/gaussian_process/presets/baybe.py index 00852d8439..690fc318cd 100644 --- a/baybe/surrogates/gaussian_process/presets/baybe.py +++ b/baybe/surrogates/gaussian_process/presets/baybe.py @@ -10,7 +10,7 @@ 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.enum import _ParameterKind from baybe.parameters.selectors import ( ParameterSelectorProtocol, TypeSelector, @@ -32,8 +32,8 @@ class BayBEKernelFactory(_PureKernelFactory): """The default kernel factory for Gaussian process surrogates.""" - supported_parameter_kinds: ClassVar[ParameterKind] = ( - ParameterKind.REGULAR | ParameterKind.TASK + _supported_parameter_kinds: ClassVar[_ParameterKind] = ( + _ParameterKind.REGULAR | _ParameterKind.TASK ) # See base class. @@ -59,7 +59,7 @@ class BayBETaskKernelFactory(_PureKernelFactory): _uses_parameter_names: ClassVar[bool] = True # See base class. - supported_parameter_kinds: ClassVar[ParameterKind] = ParameterKind.TASK + _supported_parameter_kinds: ClassVar[_ParameterKind] = _ParameterKind.TASK # See base class. parameter_selector: ParameterSelectorProtocol | None = field(