From 7a6d68f31521bd48695afd19396334db1a137e31 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Fri, 21 Aug 2026 17:50:07 +0200 Subject: [PATCH 01/73] Support scoping kernels to parameters --- baybe/kernels/base.py | 22 +++++++++++++++++++++- baybe/kernels/composite.py | 25 ++++++++++++++++++++++++- 2 files changed, 45 insertions(+), 2 deletions(-) diff --git a/baybe/kernels/base.py b/baybe/kernels/base.py index 7a07f97096..52730bcdef 100644 --- a/baybe/kernels/base.py +++ b/baybe/kernels/base.py @@ -15,12 +15,12 @@ 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 classproperty, get_baseclasses, match_attributes, to_tuple if TYPE_CHECKING: + from baybe.searchspace.core import SearchSpace from baybe.surrogates.gaussian_process.components.kernel import PlainKernelFactory @@ -121,6 +121,22 @@ def _without_parameter( f"Cannot remove a parameter from kernel '{self.__class__.__name__}'. " ) + def _with_parameter(self, name: str, /) -> Kernel: + """Return a copy of the kernel scoped to a single parameter. + + Args: + name: The name of the parameter to scope the kernel to. + + Raises: + TypeError: If the kernel structure cannot be scoped unambiguously. + + Returns: + The scoped kernel. + """ + raise TypeError( + f"Cannot scope kernel '{self.__class__.__name__}' to a parameter." + ) + @abstractmethod def _get_dimensions( self, searchspace: SearchSpace @@ -273,6 +289,10 @@ def _without_parameter( return self return evolve(self, parameter_names=remaining) if remaining else None + @override + def _with_parameter(self, name: str, /) -> Kernel: + return evolve(self, parameter_names=(name,)) + @define(frozen=True) class CompositeKernel(Kernel, ABC): diff --git a/baybe/kernels/composite.py b/baybe/kernels/composite.py index 431f0fd953..a50c9aed62 100644 --- a/baybe/kernels/composite.py +++ b/baybe/kernels/composite.py @@ -1,8 +1,11 @@ """Composite kernels (that is, kernels composed of other kernels).""" +from __future__ import annotations + import gc from functools import reduce from operator import add, mul +from typing import TYPE_CHECKING from attrs import define, evolve, field from attrs.converters import optional as optional_c @@ -12,11 +15,13 @@ from baybe.kernels.base import CompositeKernel, Kernel from baybe.priors.base import Prior -from baybe.searchspace.core import SearchSpace from baybe.settings import active_settings from baybe.utils.basic import to_tuple from baybe.utils.validation import finite_float +if TYPE_CHECKING: + from baybe.searchspace.core import SearchSpace + @define(frozen=True) class ScaleKernel(CompositeKernel): @@ -50,6 +55,10 @@ def _without_parameter( stripped = self.base_kernel._without_parameter(name, searchspace) return None if stripped is None else evolve(self, base_kernel=stripped) + @override + def _with_parameter(self, name: str, /) -> Kernel: + return evolve(self, base_kernel=self.base_kernel._with_parameter(name)) + @override def to_gpytorch(self, *args, **kwargs): import torch @@ -76,6 +85,13 @@ class AdditiveKernel(CompositeKernel): ) """The individual kernels to be summed.""" + @override + def _with_parameter(self, name: str, /) -> Kernel: + return evolve( + self, + base_kernels=tuple(k._with_parameter(name) for k in self.base_kernels), + ) + @override def to_gpytorch(self, *args, **kwargs): return reduce(add, (k.to_gpytorch(*args, **kwargs) for k in self.base_kernels)) @@ -93,6 +109,13 @@ class ProductKernel(CompositeKernel): ) """The individual kernels to be multiplied.""" + @override + def _with_parameter(self, name: str, /) -> Kernel: + return evolve( + self, + base_kernels=tuple(k._with_parameter(name) for k in self.base_kernels), + ) + @override def to_gpytorch(self, *args, **kwargs): return reduce(mul, (k.to_gpytorch(*args, **kwargs) for k in self.base_kernels)) From e1086615aa385f05f994ba75797ac4f75dd5a9ef Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Fri, 21 Aug 2026 17:50:08 +0200 Subject: [PATCH 02/73] Add kernel overrides to parameters --- baybe/parameters/base.py | 100 +++++++++++++++++- baybe/parameters/categorical.py | 3 + tests/validation/test_parameter_validation.py | 57 ++++++++++ 3 files changed, 157 insertions(+), 3 deletions(-) diff --git a/baybe/parameters/base.py b/baybe/parameters/base.py index 4fd8143627..d4e97a86ce 100644 --- a/baybe/parameters/base.py +++ b/baybe/parameters/base.py @@ -3,17 +3,19 @@ from __future__ import annotations import gc +import sys from abc import ABC, abstractmethod from functools import cached_property -from typing import TYPE_CHECKING, Any, ClassVar +from typing import TYPE_CHECKING, Any, ClassVar, TypeAlias import attrs import pandas as pd -from attrs import define, field +from attrs import Converter, define, field from attrs.converters import optional as optional_c from attrs.validators import instance_of, min_len from typing_extensions import override +from baybe.kernels.base import Kernel from baybe.parameters.enum import ParameterEncoding from baybe.serialization import ( SerialMixin, @@ -22,15 +24,93 @@ from baybe.utils.metadata import MeasurableMetadata, to_metadata if TYPE_CHECKING: + from gpytorch.kernels import Kernel as GPyTorchKernel + from baybe.parameters.enum import _ParameterKind from baybe.searchspace.continuous import SubspaceContinuous from baybe.searchspace.core import SearchSpace from baybe.searchspace.discrete import SubspaceDiscrete + KernelOverride: TypeAlias = Kernel | GPyTorchKernel +else: + KernelOverride: TypeAlias = Kernel + # TODO: Reactive slots in all classes once cached_property is supported: # https://github.com/python-attrs/attrs/issues/164 +def _iter_basic_kernels(kernel: Kernel): + """Iterate over the basic kernel leaves of a BayBE kernel.""" + from baybe.kernels.base import BasicKernel + from baybe.kernels.composite import AdditiveKernel, ProductKernel, ScaleKernel + + if isinstance(kernel, BasicKernel): + yield kernel + elif isinstance(kernel, ScaleKernel): + yield from _iter_basic_kernels(kernel.base_kernel) + elif isinstance(kernel, (AdditiveKernel, ProductKernel)): + for sub in kernel.base_kernels: + yield from _iter_basic_kernels(sub) + + +def _to_kernel_override( + value: KernelOverride | None, instance: Parameter +) -> KernelOverride | None: + """Validate a kernel override and scope BayBE kernels to their parameter. + + Args: + value: The provided kernel override. + instance: The parameter the override belongs to. + + Raises: + ValueError: If a BayBE kernel targets a different parameter or a GPyTorch + kernel specifies explicit active dimensions. + TypeError: If the object is neither a BayBE nor a GPyTorch kernel. + + Returns: + The validated override, with BayBE kernels scoped to the parameter. + """ + if value is None: + return None + + # BayBE kernels: every basic leaf must be unscoped or scoped to the owner. The + # kernel is then rebound to the owning parameter (dropping unspecified names). + if isinstance(value, Kernel): + if any( + leaf.parameter_names not in (None, (instance.name,)) + for leaf in _iter_basic_kernels(value) + ): + raise ValueError( + f"The kernel provided for the kernel override of " + f"'{instance.__class__.__name__}' may only act on the parameter " + f"itself. Its basic kernels must specify 'parameter_names' as " + f"``None`` or ({instance.name!r},)." + ) + return value._with_parameter(instance.name) + + # GPyTorch kernels: no explicit active dimensions allowed anywhere in the tree. + if sys.modules.get("gpytorch") is not None: + from gpytorch.kernels import Kernel as GPyTorchKernel + + if isinstance(value, GPyTorchKernel): + if any( + k.active_dims is not None + for k in value.modules() + if isinstance(k, GPyTorchKernel) + ): + raise ValueError( + "The GPyTorch kernel provided for the kernel override must not " + "specify 'active_dims'." + ) + return value + + raise TypeError( + f"The object provided for the kernel override of " + f"'{instance.__class__.__name__}' must be a BayBE or GPyTorch kernel. " + f"Got: {type(value)}" + ) + + @define(frozen=True, slots=False) class Parameter(ABC, SerialMixin): """Abstract base class for all parameters. @@ -47,6 +127,13 @@ class Parameter(ABC, SerialMixin): name: str = field(validator=(instance_of(str), min_len(1))) """The name of the parameter""" + kernel_override: KernelOverride | None = field( + default=None, + converter=Converter(_to_kernel_override, takes_self=True), # type: ignore[misc, call-overload] + kw_only=True, + ) + """An optional kernel replacing the overall kernel for this parameter.""" + metadata: MeasurableMetadata = field( factory=MeasurableMetadata, converter=lambda x: to_metadata(x, MeasurableMetadata), @@ -111,7 +198,14 @@ def is_equivalent(self, other: Parameter) -> bool: """ if type(self) is not type(other): return False - return attrs.evolve(self, name=other.name) == other + # The override is owner-scoped, so rebind it to the other parameter's name. + kernel_override = self.kernel_override + if isinstance(kernel_override, Kernel): + kernel_override = kernel_override._with_parameter(other.name) + return ( + attrs.evolve(self, name=other.name, kernel_override=kernel_override) + == other + ) @abstractmethod def summary(self) -> dict: diff --git a/baybe/parameters/categorical.py b/baybe/parameters/categorical.py index a98f4dc1a9..8284c6ac28 100644 --- a/baybe/parameters/categorical.py +++ b/baybe/parameters/categorical.py @@ -88,6 +88,9 @@ class TaskParameter(CategoricalParameter): encoding: CategoricalEncoding = field(default=CategoricalEncoding.INT, init=False) # See base class. + kernel_override: None = field(init=False, default=None) + """Task parameters do not support parameter-specific kernel overrides.""" + override_transfer_learning_mode: TransferLearningMode | None = field( default=None, converter=optional_c(TransferLearningMode), diff --git a/tests/validation/test_parameter_validation.py b/tests/validation/test_parameter_validation.py index 3e1d9f3893..35064e97f2 100644 --- a/tests/validation/test_parameter_validation.py +++ b/tests/validation/test_parameter_validation.py @@ -1,5 +1,6 @@ """Validation tests for parameters.""" +from functools import partial from unittest.mock import Mock import numpy as np @@ -7,9 +8,12 @@ import pytest from cattrs.errors import IterableValidationError from exceptiongroup import ExceptionGroup +from gpytorch import kernels as gk from pytest import param from baybe._optional.info import CHEM_INSTALLED +from baybe.kernels import RBFKernel +from baybe.kernels.composite import ScaleKernel from baybe.parameters.categorical import ( CategoricalParameter, TaskParameter, @@ -37,6 +41,59 @@ def test_invalid_parameter_name(name, error): NumericalDiscreteParameter(name=name, values=[1, 2, 3]) +@pytest.mark.parametrize( + ("constructor", "override", "error", "match"), + [ + param( + partial(TaskParameter, "task", ["a", "b"]), + RBFKernel(), + TypeError, + "kernel_override", + id="task_parameter", + ), + param( + partial(NumericalContinuousParameter, "x", (0, 1)), + RBFKernel(parameter_names=("other",)), + ValueError, + "may only act on the parameter itself", + id="baybe_different_name", + ), + param( + partial(NumericalContinuousParameter, "x", (0, 1)), + RBFKernel(parameter_names=("x", "other")), + ValueError, + "may only act on the parameter itself", + id="baybe_multiple_names", + ), + param( + partial(NumericalContinuousParameter, "x", (0, 1)), + ScaleKernel(RBFKernel(parameter_names=("other",))), + ValueError, + "may only act on the parameter itself", + id="baybe_composite_leaf", + ), + param( + partial(NumericalContinuousParameter, "x", (0, 1)), + gk.RBFKernel(active_dims=[0]), + ValueError, + "must not specify 'active_dims'", + id="gpytorch_active_dims", + ), + param( + partial(NumericalContinuousParameter, "x", (0, 1)), + gk.ScaleKernel(gk.RBFKernel(active_dims=[0])), + ValueError, + "must not specify 'active_dims'", + id="gpytorch_nested_active_dims", + ), + ], +) +def test_invalid_kernel_override(constructor, override, error, match): + """Invalid parameter kernel overrides raise the expected exceptions.""" + with pytest.raises(error, match=match): + constructor(kernel_override=override) + + @pytest.mark.parametrize( ("decorrelation", "error"), [ From 42118b9971b3152bb2fc3f75fbd46763c811e35b Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Fri, 21 Aug 2026 17:50:17 +0200 Subject: [PATCH 03/73] Generate valid parameter kernel overrides --- tests/hypothesis_strategies/parameters.py | 53 ++++++++++++++++++++++- 1 file changed, 51 insertions(+), 2 deletions(-) diff --git a/tests/hypothesis_strategies/parameters.py b/tests/hypothesis_strategies/parameters.py index 0b34632f52..a2638d8fa6 100644 --- a/tests/hypothesis_strategies/parameters.py +++ b/tests/hypothesis_strategies/parameters.py @@ -4,6 +4,7 @@ import hypothesis.strategies as st import numpy as np +from attrs import evolve from hypothesis.extra.pandas import columns, data_frames from baybe.parameters.categorical import ( @@ -38,6 +39,39 @@ """A strategy that generates parameter categories.""" +def _remove_kernel_parameter_names(kernel): + """Remove explicit parameter names from all basic kernel leaves.""" + from baybe.kernels.base import BasicKernel + from baybe.kernels.composite import AdditiveKernel, ProductKernel, ScaleKernel + + if isinstance(kernel, BasicKernel): + return evolve(kernel, parameter_names=None) + if isinstance(kernel, ScaleKernel): + return evolve( + kernel, + base_kernel=_remove_kernel_parameter_names(kernel.base_kernel), + ) + if isinstance(kernel, (AdditiveKernel, ProductKernel)): + return evolve( + kernel, + base_kernels=tuple( + _remove_kernel_parameter_names(k) for k in kernel.base_kernels + ), + ) + raise TypeError(f"Unsupported kernel type: {type(kernel)}") + + +@st.composite +def kernel_overrides(draw: st.DrawFn, parameter_name: str): + """Generate valid kernel overrides for a parameter.""" + from tests.hypothesis_strategies.kernels import kernels + + kernel = draw(st.one_of(st.none(), kernels(parameter_names=(parameter_name,)))) + if kernel is not None and draw(st.booleans()): + kernel = _remove_kernel_parameter_names(kernel) + return kernel + + @st.composite def smiles(draw: st.DrawFn): """Generate short SMILES strings.""" @@ -121,8 +155,13 @@ def numerical_discrete_parameters( ) ) param_metadata = draw(measurable_metadata()) + kernel_override = draw(kernel_overrides(name)) return NumericalDiscreteParameter( - name=name, values=values, tolerance=tolerance, metadata=param_metadata + name=name, + values=values, + tolerance=tolerance, + metadata=param_metadata, + kernel_override=kernel_override, ) @@ -132,8 +171,12 @@ def numerical_continuous_parameters(draw: st.DrawFn): name = draw(parameter_names) bounds = draw(intervals(exclude_half_bounded=True, exclude_fully_unbounded=True)) param_metadata = draw(measurable_metadata()) + kernel_override = draw(kernel_overrides(name)) return NumericalContinuousParameter( - name=name, bounds=bounds, metadata=param_metadata + name=name, + bounds=bounds, + metadata=param_metadata, + kernel_override=kernel_override, ) @@ -145,12 +188,14 @@ def categorical_parameters(draw: st.DrawFn): encoding = draw(st.sampled_from(CategoricalEncoding)) active_values = draw(_active_values(values)) param_metadata = draw(measurable_metadata()) + kernel_override = draw(kernel_overrides(name)) return CategoricalParameter( name=name, values=values, encoding=encoding, active_values=active_values, metadata=param_metadata, + kernel_override=kernel_override, ) @@ -188,6 +233,7 @@ def substance_parameters(draw: st.DrawFn): encoding = draw(st.sampled_from(encodings)) param_metadata = draw(measurable_metadata()) + kernel_override = draw(kernel_overrides(name)) return SubstanceParameter( name=name, @@ -196,6 +242,7 @@ def substance_parameters(draw: st.DrawFn): encoding=encoding, active_values=active_values, metadata=param_metadata, + kernel_override=kernel_override, ) @@ -207,12 +254,14 @@ def custom_parameters(draw: st.DrawFn): decorrelate = draw(decorrelations) active_values = draw(_active_values(list(data.index))) param_metadata = draw(measurable_metadata()) + kernel_override = draw(kernel_overrides(name)) return CustomDiscreteParameter( name=name, data=data, decorrelate=decorrelate, active_values=active_values, metadata=param_metadata, + kernel_override=kernel_override, ) From 7314c9a1a2fa91982673a4b1d5e72b4fa81b1e72 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Fri, 21 Aug 2026 17:50:18 +0200 Subject: [PATCH 04/73] Warn when surrogates ignore kernel overrides --- baybe/surrogates/base.py | 19 ++++++++++++++++++- tests/test_surrogate.py | 16 +++++++++++++++- 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py index d76aea9f51..df4ce48662 100644 --- a/baybe/surrogates/base.py +++ b/baybe/surrogates/base.py @@ -3,6 +3,7 @@ from __future__ import annotations import gc +import warnings from abc import ABC, abstractmethod from collections.abc import Sequence from enum import Enum, auto @@ -13,7 +14,11 @@ from joblib.hashing import hash from typing_extensions import override -from baybe.exceptions import IncompatibleSurrogateError, ModelNotTrainedError +from baybe.exceptions import ( + IncompatibleSurrogateError, + ModelNotTrainedError, + UnusedObjectWarning, +) from baybe.objectives.base import Objective from baybe.parameters.base import Parameter from baybe.searchspace import SearchSpace @@ -82,6 +87,9 @@ def to_botorch(self) -> Model: class Surrogate(ABC, SurrogateProtocol, SerialMixin): """Abstract base class for all surrogate models.""" + supports_kernel_overrides: ClassVar[bool] = False + """Whether parameter-specific kernel overrides are supported.""" + supports_transfer_learning: ClassVar[bool] """Class variable encoding whether or not the surrogate supports transfer learning.""" @@ -431,6 +439,15 @@ def fit( f"support transfer learning." ) + if not self.supports_kernel_overrides and ( + names := [p.name for p in searchspace.parameters if p.kernel_override] + ): + warnings.warn( + f"The selected surrogate model ({self.__class__.__name__}) does not " + f"use the kernel overrides specified for parameters {names}.", + UnusedObjectWarning, + ) + # Block partial measurements handle_missing_values(measurements, [t.name for t in objective.targets]) diff --git a/tests/test_surrogate.py b/tests/test_surrogate.py index 321c090b34..5302dde0f7 100644 --- a/tests/test_surrogate.py +++ b/tests/test_surrogate.py @@ -8,7 +8,8 @@ import pytest from pytest import param -from baybe.exceptions import IncompatibleSurrogateError +from baybe.exceptions import IncompatibleSurrogateError, UnusedObjectWarning +from baybe.kernels import RBFKernel from baybe.objectives.base import Objective from baybe.objectives.desirability import DesirabilityObjective from baybe.objectives.pareto import ParetoObjective @@ -48,6 +49,19 @@ def test_caching(patched, searchspace, objective, fake_measurements): patched.assert_not_called() +def test_ignored_kernel_override_warning(): + """Non-GP surrogates warn when ignoring a parameter kernel override.""" + parameter = NumericalDiscreteParameter("x", [0, 1], kernel_override=RBFKernel()) + measurements = pd.DataFrame({"x": [0, 1], "y": [0.0, 1.0]}) + + with pytest.warns(UnusedObjectWarning, match=r"kernel overrides.*\['x'\]"): + MeanPredictionSurrogate().fit( + parameter.to_searchspace(), + NumericalTarget("y").to_objective(), + measurements, + ) + + @pytest.mark.parametrize( "objective", [ From 3743956ed2a8ef38f2e12b2f1d1090d00f1d8f2b Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Mon, 24 Aug 2026 14:50:08 +0200 Subject: [PATCH 05/73] Partition GP kernels around parameter overrides Construct default numerical and task factors separately: overrides can leave a numerical-only, task-only, or empty residual, which ICM does not directly assemble without additional special handling. A restricted BayBEKernelFactory selector still attaches the default task kernel, duplicating an explicit TL override. Numerical factories also cannot always accept an empty selection because dimension-scaled priors require nonzero width. Construct on the full search space to preserve computational column indices and preset dispatch, combining the original selector with override exclusions. --- baybe/surrogates/gaussian_process/core.py | 343 +++++++++++++++------- 1 file changed, 233 insertions(+), 110 deletions(-) diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py index 3d4e93447d..f7fc7fa416 100644 --- a/baybe/surrogates/gaussian_process/core.py +++ b/baybe/surrogates/gaussian_process/core.py @@ -4,13 +4,15 @@ import gc import importlib +import operator import os import warnings -from functools import partial -from typing import TYPE_CHECKING, ClassVar +from copy import deepcopy +from functools import partial, reduce +from typing import TYPE_CHECKING, ClassVar, NoReturn import pandas as pd -from attrs import Converter, define, field, fields +from attrs import Converter, define, field from attrs.converters import optional as optional_c from attrs.converters import pipe from attrs.validators import instance_of, is_callable, optional @@ -72,6 +74,8 @@ from gpytorch.means import Mean as GPyTorchMean from torch import Tensor + from baybe.parameters.selectors import ParameterSelectorProtocol + @define class _ModelContext: @@ -142,6 +146,45 @@ def _mark_custom_kernel( return value +def _bind_gpytorch_override( + override: GPyTorchKernel, indices: tuple[int, ...], name: str +) -> GPyTorchKernel: + """Copy a raw GPyTorch override and bind it to the given dimensions. + + Args: + override: The provided GPyTorch kernel (must not specify active dimensions). + indices: The computational column indices of the owning parameter. + name: The owning parameter name (for error messages). + + Raises: + IncompatibleOverrideError: If the kernel specifies active dimensions or an + incompatible number of ARD dimensions. + + Returns: + A copy of the kernel bound to the given dimensions. + """ + import torch + from gpytorch.kernels import Kernel as GPyTorchKernel + + for kernel in override.modules(): + if not isinstance(kernel, GPyTorchKernel): + continue + if kernel.active_dims is not None: + raise IncompatibleOverrideError( + f"The GPyTorch kernel override for parameter '{name}' must not " + f"specify 'active_dims'." + ) + if kernel.ard_num_dims not in (None, len(indices)): + raise IncompatibleOverrideError( + f"The GPyTorch kernel override for parameter '{name}' specifies " + f"{kernel.ard_num_dims} ARD dimensions, but the parameter has " + f"{len(indices)} computational dimensions." + ) + result = deepcopy(override) + result.active_dims = torch.tensor(indices) + return result + + @define class GaussianProcessSurrogate(Surrogate): """A Gaussian process surrogate model.""" @@ -161,6 +204,9 @@ class GaussianProcessSurrogate(Surrogate): # to `optimize_acqf_*`, which is configured to be called on the original scale. # Moving the scaling operation into the botorch GP object avoids this conflict. + supports_kernel_overrides: ClassVar[bool] = True + # See base class. + supports_transfer_learning: ClassVar[bool] = True # See base class. @@ -396,129 +442,206 @@ def _posterior(self, candidates_comp_scaled: Tensor, /) -> Posterior: return self._model.posterior(candidates_comp_scaled) def _resolve_kernel(self, context: _ModelContext) -> GPyTorchKernel: - """Resolve the GP kernel, dispatching on task parameter overrides. + """Resolve the GP kernel, applying parameter and transfer overrides. - Args: - context: The model context providing searchspace information. + The effective kernel is the surrogate kernel restricted to the + non-overridden dimensions, multiplied by one factor per override. + """ + overrides = self._extract_parameter_overrides(context) + overrides.extend(self._extract_transfer_learning_overrides(context)) + + # No overrides: let the surrogate kernel (factory) handle everything. + if not overrides: + factory = self.kernel_factory or BayBEKernelFactory() + kernel = factory( + context.searchspace, context.objective, context.measurements + ) + return self._as_gpytorch(kernel, context.searchspace) - Raises: - IncompatibleOverrideError: If a transfer learning override is combined - with a kernel or kernel factory that cannot be reduced to a task-free - base kernel operating on parameter names. + excluded_names = {name for name, _ in overrides} + residual = self._resolve_residual_kernel(context, excluded_names) + factors = ([] if residual is None else [residual]) + [k for _, k in overrides] + return reduce(operator.mul, factors) - Returns: - The constructed gpytorch kernel. - """ - searchspace = context.searchspace - task_param = searchspace._task_parameter - tl_override = context.tl_override + def _extract_parameter_overrides( + self, context: _ModelContext + ) -> list[tuple[str, GPyTorchKernel]]: + """Extract regular parameter-specific kernel overrides.""" + return [ + (p.name, self._make_parameter_override_kernel(p, context.searchspace)) + for p in context.searchspace.parameters + if p.kernel_override is not None + ] - if tl_override is None: - # No override: let the factory handle everything (default path) - kernel_factory = self.kernel_factory or BayBEKernelFactory() - kernel = kernel_factory( - searchspace, context.objective, context.measurements - ) - if isinstance(kernel, Kernel): - kernel = kernel.to_gpytorch(searchspace=searchspace) - return kernel + def _extract_transfer_learning_overrides( + self, context: _ModelContext + ) -> list[tuple[str, GPyTorchKernel]]: + """Extract the transfer-learning kernel override, if any.""" + task_param = context.searchspace._task_parameter + if task_param is None or context.tl_override is None: + return [] + kernel = self._make_transfer_learning_override_kernel(context) + return [(task_param.name, kernel)] - assert task_param is not None # a set override implies a task parameter + @staticmethod + def _make_parameter_override_kernel( + parameter: Parameter, searchspace: SearchSpace + ) -> GPyTorchKernel: + """Create the kernel factor for a parameter's override.""" + override = parameter.kernel_override + assert override is not None + indices = searchspace.get_comp_rep_parameter_indices(parameter.name) + + # BayBE kernels resolve their own dimensions; raw kernels are bound manually. + if isinstance(override, Kernel): + return override.to_gpytorch(searchspace) + return _bind_gpytorch_override(override, indices, parameter.name) - # Override is set: assemble the prescribed task kernel - n_tasks = searchspace.n_tasks - task_kernel_cls: type[IndexKernel] - match tl_override: + @staticmethod + def _make_transfer_learning_override_kernel( + context: _ModelContext, + ) -> GPyTorchKernel: + """Create the task kernel factor requested by a transfer-learning override.""" + task_param = context.searchspace._task_parameter + override = context.tl_override + assert task_param is not None and override is not None + n_tasks, names = context.n_tasks, (task_param.name,) + match override: case TransferLearningMode.POSITIVE_INDEX_KERNEL: - task_kernel_cls = PositiveIndexKernel + spec: Kernel = PositiveIndexKernel( + num_tasks=n_tasks, rank=n_tasks, parameter_names=names + ) case TransferLearningMode.INDEX_KERNEL: - task_kernel_cls = IndexKernel + spec = IndexKernel( + num_tasks=n_tasks, rank=n_tasks, parameter_names=names + ) case _: - assert_never(tl_override) - task_kernel_spec = task_kernel_cls( - num_tasks=n_tasks, rank=n_tasks, parameter_names=(task_param.name,) + assert_never(override) + return spec.to_gpytorch(context.searchspace) + + def _resolve_residual_kernel( + self, context: _ModelContext, excluded_names: set[str] + ) -> GPyTorchKernel | None: + """Resolve the surrogate kernel restricted to the non-excluded dimensions. + + This is intentionally task-agnostic: it only removes ``excluded_names``. + A task parameter is excluded here only if a transfer-learning override + already replaced it; otherwise it is preserved by the surrogate kernel. + """ + searchspace = context.searchspace + factory = self.kernel_factory or BayBEKernelFactory() + remaining = [p for p in searchspace.parameters if p.name not in excluded_names] + if not remaining: + return None + + # Default kernel: reuse the default machinery, but restrict it to the + # remaining (non-overridden) parameters via the factory's selector. + if isinstance(factory, BayBEKernelFactory): + return self._resolve_default_base(context, excluded_names) + + # A fixed kernel is reduced directly by removing the excluded parameters. + if isinstance(factory, PlainGPComponentFactory): + spec = self._reduce_kernel_spec( + factory.component, excluded_names, searchspace, factory + ) + return None if spec is None else spec.to_gpytorch(searchspace) + + # Any other callable factory: call it on the reduced space and reduce the + # returned BayBE kernel. Factories needing full-space information or + # returning a raw kernel are unsupported. + reduced_space = searchspace._drop_parameters(excluded_names) + try: + returned = factory(reduced_space, context.objective, context.measurements) + except ( + IncompatibleSearchSpaceError, + _UnsupportedSearchSpaceAttributeError, + ) as ex: + self._raise_incompatible_override(excluded_names, factory, ex) + spec = self._reduce_kernel_spec(returned, excluded_names, searchspace, factory) + return None if spec is None else spec.to_gpytorch(searchspace) + + def _resolve_default_base( + self, context: _ModelContext, excluded_names: set[str] + ) -> GPyTorchKernel | None: + """Resolve the default kernel base, excluding the given parameters. + + The numerical base always excludes the task; the default task kernel is + re-added unless the task is itself overridden (i.e. already excluded). + """ + from baybe.surrogates.gaussian_process.components.kernel import ( + _PureKernelFactory, + ) + from baybe.surrogates.gaussian_process.presets.baybe import ( + _BayBENumericalKernelFactory, + _BayBETaskKernelFactory, ) - # Default factory (None or a `BayBEKernelFactory` without a custom parameter - # selector): reuse the ICM machinery on the full searchspace, which builds the - # task-excluded base kernel and combines it with the prescribed task kernel. - # This avoids the reduced searchspace, on which the default factory's numerical - # kernel cannot resolve its active dimensions. A `BayBEKernelFactory` carrying a - # custom selector falls through to the general factory path below, so the - # selector is honored (or loudly rejected when it cannot be). - if self.kernel_factory is None or ( - isinstance(self.kernel_factory, BayBEKernelFactory) - and self.kernel_factory.parameter_selector is None - ): - icm = ICMKernelFactory(task_kernel_or_factory=task_kernel_spec) - kernel = icm(searchspace, context.objective, context.measurements) - if isinstance(kernel, Kernel): - kernel = kernel.to_gpytorch(searchspace=searchspace) - return kernel - - # Otherwise, build a task-free base kernel and attach the prescribed task - # kernel manually. - effective_factory = self.kernel_factory - incompatible_message = ( - f"The '{TaskParameter.__name__}' '{task_param.name}' specifies " - f"'{fields(TaskParameter).override_transfer_learning_mode.name}=" - f"{tl_override.name}', which requires a " - f"kernel (factory) that yields a task-free BayBE kernel operating on " - f"parameter names. The provided kernel factory " - f"'{type(effective_factory).__name__}' does not satisfy this (e.g., it " - f"returns a raw gpytorch kernel or already operates on the task " - f"parameter)." + searchspace = context.searchspace + factory = self.kernel_factory or BayBEKernelFactory() + assert isinstance(factory, _PureKernelFactory) + selector = factory.parameter_selector + task_param = searchspace._task_parameter + + keep: ParameterSelectorProtocol = lambda parameter: ( # noqa: E731 + parameter.name not in excluded_names + and parameter is not task_param + and (selector is None or selector(parameter)) ) + factors: list[GPyTorchKernel] = [] + if any(keep(p) for p in searchspace.parameters): + base = _BayBENumericalKernelFactory(parameter_selector=keep)( + searchspace, context.objective, context.measurements + ) + factors.append(self._as_gpytorch(base, searchspace)) + if task_param is not None and task_param.name not in excluded_names: + task = _BayBETaskKernelFactory()( + searchspace, context.objective, context.measurements + ) + factors.append(self._as_gpytorch(task, searchspace)) - if isinstance(self.kernel_factory, PlainGPComponentFactory): - # A fixed kernel was provided: strip the task parameter directly. - component = self.kernel_factory.component - if not isinstance(component, Kernel): - raise IncompatibleOverrideError(incompatible_message) + return reduce(operator.mul, factors) if factors else None + + @staticmethod + def _reduce_kernel_spec( + component: object, + excluded_names: set[str], + searchspace: SearchSpace, + factory: object, + ) -> Kernel | None: + """Remove the excluded parameters from a fixed BayBE kernel.""" + if not isinstance(component, Kernel): + GaussianProcessSurrogate._raise_incompatible_override( + excluded_names, factory + ) + spec: Kernel | None = component + for name in excluded_names: + if spec is None: + break try: - base_spec = component._without_parameter(task_param.name, searchspace) + spec = spec._without_parameter(name, searchspace) except TypeError as ex: - raise IncompatibleOverrideError(incompatible_message) from ex - else: - # Call the factory on a reduced (task-free) searchspace so that it - # produces only the base kernel. Factories that need computational - # information unavailable on the reduced space, or that return a raw - # gpytorch kernel, are not supported. - reduced_searchspace = searchspace._drop_parameters({task_param.name}) - try: - factory_kernel = effective_factory( - reduced_searchspace, context.objective, context.measurements + GaussianProcessSurrogate._raise_incompatible_override( + excluded_names, factory, ex ) - except ( - IncompatibleSearchSpaceError, - _UnsupportedSearchSpaceAttributeError, - ) as ex: - raise IncompatibleOverrideError(incompatible_message) from ex - if not isinstance(factory_kernel, Kernel): - raise IncompatibleOverrideError(incompatible_message) - # Normalize to an explicitly task-free spec. - try: - base_spec = factory_kernel._without_parameter( - task_param.name, searchspace - ) - except TypeError as ex: - raise IncompatibleOverrideError(incompatible_message) from ex - - # Stripping left no base kernel: return only the prescribed task kernel. - if base_spec is None: - return task_kernel_spec.to_gpytorch(searchspace=searchspace) - - # Combine base and task kernel via the ICM machinery (as in the default-factory - # branch above), which converts both on the full searchspace and validates the - # dimension partitioning between base and task kernel. - icm = ICMKernelFactory( - base_kernel_or_factory=base_spec, - task_kernel_or_factory=task_kernel_spec, - ) - kernel = icm(searchspace, context.objective, context.measurements) - if isinstance(kernel, Kernel): - kernel = kernel.to_gpytorch(searchspace=searchspace) - return kernel + return spec + + @staticmethod + def _as_gpytorch( + kernel: Kernel | GPyTorchKernel, searchspace: SearchSpace + ) -> GPyTorchKernel: + """Convert a BayBE kernel to GPyTorch, passing through raw kernels.""" + return kernel.to_gpytorch(searchspace) if isinstance(kernel, Kernel) else kernel + + @staticmethod + def _raise_incompatible_override( + parameter_names: set[str], factory: object, cause: Exception | None = None + ) -> NoReturn: + """Raise an error for a surrogate kernel that cannot be reduced.""" + raise IncompatibleOverrideError( + f"Kernel overrides for {sorted(parameter_names)} require a surrogate " + f"kernel (factory) that can exclude these parameters. " + f"'{type(factory).__name__}' does not satisfy this requirement." + ) from cause def _resolve_components( self, context: _ModelContext From 09fc1139403e72e1a612178ddd3244b224557718 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Mon, 24 Aug 2026 14:50:08 +0200 Subject: [PATCH 06/73] Test parameter-specific kernel overrides --- tests/test_parameter_kernel_overrides.py | 296 +++++++++++++++++++++++ 1 file changed, 296 insertions(+) create mode 100644 tests/test_parameter_kernel_overrides.py diff --git a/tests/test_parameter_kernel_overrides.py b/tests/test_parameter_kernel_overrides.py new file mode 100644 index 0000000000..2395b151f3 --- /dev/null +++ b/tests/test_parameter_kernel_overrides.py @@ -0,0 +1,296 @@ +"""Functional tests for parameter-specific kernel overrides. + +These tests resolve or fit kernels to check binding and composition behavior. +Construction-time input validation lives in ``tests/validation``. +""" + +from copy import deepcopy + +import pandas as pd +import pytest +import torch +from botorch.models.kernels.positive_index import PositiveIndexKernel +from gpytorch import kernels as gk +from pytest import param + +from baybe.exceptions import IncompatibleOverrideError +from baybe.kernels import MaternKernel, RBFKernel +from baybe.kernels.composite import AdditiveKernel, ProductKernel, ScaleKernel +from baybe.parameters import ( + CategoricalParameter, + NumericalContinuousParameter, + NumericalDiscreteParameter, + TaskParameter, +) +from baybe.parameters.enum import TransferLearningMode +from baybe.parameters.selectors import NameSelector +from baybe.searchspace import SearchSpace +from baybe.surrogates import GaussianProcessSurrogate +from baybe.surrogates.gaussian_process.core import _ModelContext +from baybe.surrogates.gaussian_process.presets import BayBEKernelFactory +from baybe.targets import NumericalTarget + + +def _resolve(parameters, kernel_or_factory=None): + """Resolve the default GP kernel for the given parameters.""" + searchspace = SearchSpace.from_product(parameters) + context = _ModelContext( + searchspace, NumericalTarget("y").to_objective(), pd.DataFrame() + ) + surrogate = GaussianProcessSurrogate(kernel_or_factory=kernel_or_factory) + return surrogate._resolve_kernel(context), searchspace + + +def _leaf_kernels(kernel): + """Return the leaf kernels of a GPyTorch kernel tree.""" + children = tuple(kernel.sub_kernels()) + return ( + (kernel,) + if not children + else tuple(k for c in children for k in _leaf_kernels(c)) + ) + + +def _fit(parameters): + """Fit a GP and return its installed covariance module and searchspace.""" + searchspace = SearchSpace.from_product(parameters) + measurements = pd.DataFrame( + [ + {**{p.name: p.values[0] for p in parameters}, "y": 0.0}, + {**{p.name: p.values[-1] for p in parameters}, "y": 1.0}, + ] + ) + surrogate = GaussianProcessSurrogate() + surrogate.fit(searchspace, NumericalTarget("y").to_objective(), measurements) + assert surrogate._model is not None + return surrogate._model.covar_module, searchspace + + +def test_default_factory_selector_is_preserved(): + """Partitioning preserves an explicit default-factory parameter selector.""" + kernel, _ = _resolve( + [ + NumericalContinuousParameter("x1", (0, 1)), + NumericalContinuousParameter("x2", (0, 1), kernel_override=RBFKernel()), + NumericalContinuousParameter("x3", (0, 1)), + ], + BayBEKernelFactory(parameter_selector=NameSelector(("x1",), regex=False)), + ) + + assert {tuple(k.active_dims.tolist()) for k in _leaf_kernels(kernel)} == { + (0,), + (1,), + } + + +@pytest.mark.parametrize( + ("parameters", "expected_types", "expected_names"), + [ + param( + [ + CategoricalParameter("base", ["a", "b"]), + NumericalDiscreteParameter( + "override", [0, 1, 2], kernel_override=RBFKernel() + ), + ], + [gk.MaternKernel, gk.RBFKernel], + ["base", "override"], + id="single-dim", + ), + param( + [ + CategoricalParameter("base", ["a", "b"]), + CategoricalParameter( + "override", ["a", "b", "c"], kernel_override=RBFKernel() + ), + ], + [gk.MaternKernel, gk.RBFKernel], + ["base", "override"], + id="multi-dim", + ), + param( + [ + CategoricalParameter("first", ["a", "b"], kernel_override=RBFKernel()), + CategoricalParameter( + "second", ["a", "b"], kernel_override=MaternKernel() + ), + ], + [gk.RBFKernel, gk.MaternKernel], + ["first", "second"], + id="all-overridden", + ), + param( + [ + CategoricalParameter("base", ["a", "b"]), + CategoricalParameter( + "override", ["a", "b"], kernel_override=RBFKernel() + ), + TaskParameter("task", ["a", "b"]), + ], + [ + gk.MaternKernel, + PositiveIndexKernel, + gk.RBFKernel, + ], + ["base", "task", "override"], + id="task-without-tl-override", + ), + param( + [ + CategoricalParameter("base", ["a", "b"]), + CategoricalParameter( + "override", ["a", "b", "c"], kernel_override=RBFKernel() + ), + TaskParameter( + "task", + ["a", "b"], + override_transfer_learning_mode=TransferLearningMode.INDEX_KERNEL, + ), + ], + [ + gk.MaternKernel, + gk.RBFKernel, + gk.IndexKernel, + ], + ["base", "override", "task"], + id="override-and-tl-index", + ), + param( + [ + CategoricalParameter("base", ["a", "b"]), + CategoricalParameter( + "override", ["a", "b", "c"], kernel_override=RBFKernel() + ), + TaskParameter( + "task", + ["a", "b"], + override_transfer_learning_mode=( + TransferLearningMode.POSITIVE_INDEX_KERNEL + ), + ), + ], + [ + gk.MaternKernel, + gk.RBFKernel, + PositiveIndexKernel, + ], + ["base", "override", "task"], + id="override-and-tl-pos-index", + ), + ], +) +def test_fitted_model_uses_parameter_kernel_overrides( + parameters, expected_types, expected_names +): + """The fitted model uses each configured kernel on the intended dimensions.""" + kernel, searchspace = _fit(parameters) + leaves = _leaf_kernels(kernel) + expected_dimensions = [ + searchspace.get_comp_rep_parameter_indices(name) for name in expected_names + ] + + assert isinstance(kernel, gk.ProductKernel) + assert [type(k) for k in leaves] == expected_types + assert [tuple(k.active_dims.tolist()) for k in leaves] == expected_dimensions + assert [k.ard_num_dims for k in leaves] == [len(d) for d in expected_dimensions] + + +@pytest.mark.parametrize( + "override", + [ + param(RBFKernel(), id="unscoped"), + param(RBFKernel(parameter_names=("override",)), id="owner-scoped"), + param(ScaleKernel(RBFKernel()), id="scale"), + param(AdditiveKernel([RBFKernel(), MaternKernel()]), id="additive"), + param(ProductKernel([RBFKernel(), MaternKernel()]), id="product"), + ], +) +def test_valid_baybe_kernel_override(override): + """Valid BayBE overrides bind all their leaves to the owning parameter.""" + parameters = [ + CategoricalParameter("base", ["a", "b"]), + CategoricalParameter("override", ["a", "b", "c"], kernel_override=override), + ] + kernel, searchspace = _resolve(parameters) + expected = set(searchspace.get_comp_rep_parameter_indices("override")) + + override_leaves = [ + k + for k in _leaf_kernels(kernel) + if set(k.active_dims.tolist()) + != set(searchspace.get_comp_rep_parameter_indices("base")) + ] + assert override_leaves + assert all(set(k.active_dims.tolist()) == expected for k in override_leaves) + + +@pytest.mark.parametrize( + ("categories", "override"), + [ + param(["a", "b"], gk.RBFKernel(), id="no-ard"), + param(["a", "b"], gk.RBFKernel(ard_num_dims=2), id="ard"), + param(["a", "b", "c"], gk.RBFKernel(ard_num_dims=3), id="multi-column-ard"), + param(["a", "b"], gk.ScaleKernel(gk.RBFKernel()), id="nested"), + ], +) +def test_valid_gpytorch_kernel_override(categories, override): + """Valid GPyTorch overrides are bound to the owner without mutating the input.""" + snapshot = deepcopy(override.state_dict()) + base = NumericalContinuousParameter("base", (0, 1)) + kernel, searchspace = _resolve( + [base, CategoricalParameter("override", categories, kernel_override=override)] + ) + expected = set(searchspace.get_comp_rep_parameter_indices("override")) + + # The override factor is the product factor acting on the override dimensions. + factor = next(k for k in kernel.kernels if set(k.active_dims.tolist()) == expected) + assert type(factor) is type(override) + + # The provided kernel is copied, not mutated. + assert override.active_dims is None + assert all(torch.equal(snapshot[k], override.state_dict()[k]) for k in snapshot) + + +def test_gpytorch_ard_mismatch_rejected(): + """A GPyTorch override with mismatched ARD size is rejected during resolution.""" + override = gk.RBFKernel(ard_num_dims=2) + parameter = CategoricalParameter("p", ["a", "b", "c"], kernel_override=override) + with pytest.raises(IncompatibleOverrideError, match="has 3 computational"): + _resolve([parameter]) + + +def test_parameter_equivalence_ignores_override_parameter_name(): + """Equivalent parameters may scope the same override to their own names.""" + left = NumericalContinuousParameter( + "left", (0, 1), kernel_override=RBFKernel(parameter_names=("left",)) + ) + right = NumericalContinuousParameter( + "right", (0, 1), kernel_override=RBFKernel(parameter_names=("right",)) + ) + + assert left.is_equivalent(right) + + +@pytest.mark.parametrize( + "kernel_or_factory", + [ + param( + AdditiveKernel([MaternKernel(), MaternKernel()]), + id="non-reducible-baybe-kernel", + ), + param(gk.MaternKernel(), id="raw-gpytorch-kernel"), + param( + lambda s, o, m: gk.MaternKernel(), + id="factory-returning-raw-kernel", + ), + ], +) +def test_incompatible_surrogate_kernel_is_rejected(kernel_or_factory): + """Surrogate kernels that cannot exclude overridden dimensions are rejected.""" + parameters = [ + NumericalContinuousParameter("x1", (0, 1)), + NumericalContinuousParameter("x2", (0, 1), kernel_override=RBFKernel()), + ] + + with pytest.raises(IncompatibleOverrideError): + _resolve(parameters, kernel_or_factory) From 1f65a5e6190d653548a1529d47eba634e1dcdc64 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Fri, 21 Aug 2026 17:50:28 +0200 Subject: [PATCH 07/73] Test parameter kernel overrides end to end --- tests/test_iterations.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tests/test_iterations.py b/tests/test_iterations.py index c3491f46bc..e5ceb8ceca 100644 --- a/tests/test_iterations.py +++ b/tests/test_iterations.py @@ -27,6 +27,7 @@ ) from baybe.kernels.composite import ScaleKernel from baybe.objectives.pareto import ParetoObjective +from baybe.parameters import NumericalContinuousParameter from baybe.priors import ( GammaPrior, HalfCauchyPrior, @@ -398,6 +399,25 @@ def test_kernels(ongoing_campaign, n_iterations, batch_size): run_iterations(ongoing_campaign, n_iterations, batch_size) +@pytest.mark.slow +@pytest.mark.parametrize( + "parameters", + [ + [ + NumericalContinuousParameter("x1", (0, 1)), + NumericalContinuousParameter("x2", (0, 1), kernel_override=RBFKernel()), + ] + ], + ids=["parameter_kernel_override"], +) +@pytest.mark.parametrize("n_iterations", [3], ids=["i3"]) +def test_parameter_kernel_override_iteration( + ongoing_campaign, n_iterations, batch_size +): + """A complete optimization iteration supports a parameter kernel override.""" + run_iterations(ongoing_campaign, n_iterations, batch_size) + + @pytest.mark.parametrize("kernel", valid_kernel_factories) @pytest.mark.parametrize("n_iterations", [3], ids=["i3"]) def test_kernel_factories(ongoing_campaign, n_iterations, batch_size): From a97dbb60b9fe0520400b33763c64593767ed6e4c Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Fri, 21 Aug 2026 17:50:42 +0200 Subject: [PATCH 08/73] Document parameter-specific kernel overrides --- docs/components/parameters.md | 60 ++++++++++++++++++++++++++++++++++- docs/components/surrogates.md | 5 ++- 2 files changed, 63 insertions(+), 2 deletions(-) diff --git a/docs/components/parameters.md b/docs/components/parameters.md index 11288c04b6..36545de29e 100644 --- a/docs/components/parameters.md +++ b/docs/components/parameters.md @@ -283,4 +283,62 @@ enlarged data corpus. ```{seealso} For details, refer to [transfer learning](../concepts/transfer_learning.md). -``` \ No newline at end of file +``` + +(parameter_kernel_overrides)= +## Parameter-Specific Kernel Overrides + +Regular parameters can replace the surrogate kernel on their computational dimensions +using ``kernel_override``: + +```python +from baybe.kernels import RBFKernel +from baybe.parameters import NumericalContinuousParameter + +x3 = NumericalContinuousParameter( + name="x3", + bounds=(0.0, 1.0), + kernel_override=RBFKernel(), +) +``` + +The surrogate kernel is the kernel supplied to the surrogate, or BayBE's default kernel +if none is supplied. Overrides remove their parameter dimensions from that kernel and +replace them with multiplicative parameter-specific factors: + +```{math} +k_{\mathrm{effective}} += k_{\mathrm{surrogate}}(X_{\mathrm{remaining}}) +\prod_{p \in P_{\mathrm{override}}} k_p(X_p). +``` + +For parameters ``x1``, ``x2``, and ``x3`` in the example above, this produces +``surrogate kernel(x1, x2) * RBF(x3)``. The override does not replace a dimension inside +the surrogate kernel itself. It partitions the dimensions and creates a separate kernel +factor. Multiple overrides produce multiple factors. For a parameter with several +computational dimensions, the override acts jointly on the complete computational +representation. + +An override can be a configured BayBE or GPyTorch kernel instance. For a BayBE kernel, +``parameter_names`` must be ``None`` or contain exactly the owning parameter. Here, +``None`` is rebound to the owning parameter instead of acting on all dimensions. + +```{admonition} Limitations +:class: warning +Parameter-local kernel factories are not supported. The surrogate kernel or factory of +the {class}`~baybe.surrogates.gaussian_process.core.GaussianProcessSurrogate` must allow +the overridden dimensions to be excluded. Otherwise, BayBE raises +{class}`~baybe.exceptions.IncompatibleOverrideError`. + +Raw GPyTorch overrides must not specify ``active_dims``, and their ``ard_num_dims`` must +be ``None`` or match the number of computational dimensions of the parameter. ``None`` +is preserved and does not automatically enable ARD. Like raw GPyTorch kernels used +elsewhere, they cannot be serialized. +``` + +The feature applies only to Gaussian process surrogates. Other surrogates ignore kernel +overrides and emit an {class}`~baybe.exceptions.UnusedObjectWarning`. Task parameters do +not expose {attr}`~baybe.parameters.base.Parameter.kernel_override`; use +{attr}`~baybe.parameters.categorical.TaskParameter.override_transfer_learning_mode` +instead. Both override mechanisms can be used in the same searchspace, in which case +their kernel factors are multiplied. diff --git a/docs/components/surrogates.md b/docs/components/surrogates.md index 02d938a95a..e153a3b366 100644 --- a/docs/components/surrogates.md +++ b/docs/components/surrogates.md @@ -161,6 +161,9 @@ kernel = ScaleKernel( kernel = MaternKernel(parameter_names=["Param_A", "Param_B"]) ``` +For a convenience interface that assigns a kernel directly to one parameter, see +{ref}`parameter_kernel_overrides`. + #### Presets Providing customized GP configurations on the component level is powerful for expert @@ -355,4 +358,4 @@ posterior_mean = lambda x: campaign.get_surrogate().posterior(x).mean Note that BayBE has a dedicated [Insights module](./insights.md) that offers certain analytical tools and visualizations out-of-the-box, without requiring manual extraction of the surrogate model to begin with. -``` \ No newline at end of file +``` From 21ca20248fe085bd69188511e386b6063543e986 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Mon, 24 Aug 2026 19:24:11 +0200 Subject: [PATCH 09/73] Group kernel override helpers into a subpackage --- .../gaussian_process/_override/__init__.py | 19 ++ .../gaussian_process/_override/core.py | 61 ++++++ .../gaussian_process/_override/parameter.py | 96 +++++++++ .../gaussian_process/_override/tl.py | 62 ++++++ baybe/surrogates/gaussian_process/core.py | 184 +++--------------- 5 files changed, 261 insertions(+), 161 deletions(-) create mode 100644 baybe/surrogates/gaussian_process/_override/__init__.py create mode 100644 baybe/surrogates/gaussian_process/_override/core.py create mode 100644 baybe/surrogates/gaussian_process/_override/parameter.py create mode 100644 baybe/surrogates/gaussian_process/_override/tl.py diff --git a/baybe/surrogates/gaussian_process/_override/__init__.py b/baybe/surrogates/gaussian_process/_override/__init__.py new file mode 100644 index 0000000000..081d26149d --- /dev/null +++ b/baybe/surrogates/gaussian_process/_override/__init__.py @@ -0,0 +1,19 @@ +"""Kernel override resolution for Gaussian process surrogates.""" + +from baybe.surrogates.gaussian_process._override.core import ( + raise_incompatible_override, + reduce_kernel_spec, +) +from baybe.surrogates.gaussian_process._override.parameter import ( + extract_parameter_overrides, +) +from baybe.surrogates.gaussian_process._override.tl import ( + extract_transfer_learning_overrides, +) + +__all__ = [ + "extract_parameter_overrides", + "extract_transfer_learning_overrides", + "raise_incompatible_override", + "reduce_kernel_spec", +] diff --git a/baybe/surrogates/gaussian_process/_override/core.py b/baybe/surrogates/gaussian_process/_override/core.py new file mode 100644 index 0000000000..2be7a1c312 --- /dev/null +++ b/baybe/surrogates/gaussian_process/_override/core.py @@ -0,0 +1,61 @@ +"""Shared helpers for kernel override resolution.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, NoReturn + +from baybe.exceptions import IncompatibleOverrideError +from baybe.kernels.base import Kernel + +if TYPE_CHECKING: + from baybe.searchspace import SearchSpace + + +def reduce_kernel_spec( + component: object, + excluded_names: set[str], + searchspace: SearchSpace, + factory: object, +) -> Kernel | None: + """Remove the excluded parameters from a fixed BayBE kernel. + + Args: + component: The kernel specification to reduce. + excluded_names: The names of the parameters to remove. + searchspace: The search space the kernel operates on. + factory: The originating factory (for error messages). + + Returns: + The reduced kernel, or ``None`` if nothing remains. + """ + if not isinstance(component, Kernel): + raise_incompatible_override(excluded_names, factory) + spec: Kernel | None = component + for name in excluded_names: + if spec is None: + break + try: + spec = spec._without_parameter(name, searchspace) + except TypeError as ex: + raise_incompatible_override(excluded_names, factory, ex) + return spec + + +def raise_incompatible_override( + parameter_names: set[str], factory: object, cause: Exception | None = None +) -> NoReturn: + """Raise an error for a surrogate kernel that cannot be reduced. + + Args: + parameter_names: The overridden parameter names. + factory: The offending kernel factory. + cause: The underlying exception, if any. + + Raises: + IncompatibleOverrideError: Always. + """ + raise IncompatibleOverrideError( + f"Kernel overrides for {sorted(parameter_names)} require a surrogate " + f"kernel (factory) that can exclude these parameters. " + f"'{type(factory).__name__}' does not satisfy this requirement." + ) from cause diff --git a/baybe/surrogates/gaussian_process/_override/parameter.py b/baybe/surrogates/gaussian_process/_override/parameter.py new file mode 100644 index 0000000000..a15c4bedcf --- /dev/null +++ b/baybe/surrogates/gaussian_process/_override/parameter.py @@ -0,0 +1,96 @@ +"""Resolution of parameter-specific kernel overrides.""" + +from __future__ import annotations + +from copy import deepcopy +from typing import TYPE_CHECKING + +from baybe.exceptions import IncompatibleOverrideError +from baybe.kernels.base import Kernel + +if TYPE_CHECKING: + from gpytorch.kernels import Kernel as GPyTorchKernel + + from baybe.parameters.base import Parameter + from baybe.searchspace import SearchSpace + from baybe.surrogates.gaussian_process.core import _ModelContext + + +def extract_parameter_overrides( + context: _ModelContext, +) -> list[tuple[str, GPyTorchKernel]]: + """Extract the regular parameter-specific kernel overrides. + + Args: + context: The model context providing the search space. + + Returns: + A ``(parameter_name, kernel)`` pair for each parameter with an override. + """ + return [ + (p.name, make_parameter_override_kernel(p, context.searchspace)) + for p in context.searchspace.parameters + if p.kernel_override is not None + ] + + +def make_parameter_override_kernel( + parameter: Parameter, searchspace: SearchSpace +) -> GPyTorchKernel: + """Create the kernel factor for a parameter's override. + + Args: + parameter: The parameter carrying the override. + searchspace: The search space the kernel operates on. + + Returns: + The GPyTorch kernel bound to the parameter's dimensions. + """ + override = parameter.kernel_override + assert override is not None + + # BayBE kernels resolve their own dimensions; raw kernels are bound manually. + if isinstance(override, Kernel): + return override.to_gpytorch(searchspace) + indices = searchspace.get_comp_rep_parameter_indices(parameter.name) + return bind_gpytorch_override(override, indices, parameter.name) + + +def bind_gpytorch_override( + override: GPyTorchKernel, indices: tuple[int, ...], name: str +) -> GPyTorchKernel: + """Copy a raw GPyTorch override and bind it to the given dimensions. + + Args: + override: The provided GPyTorch kernel (must not specify active dimensions). + indices: The computational column indices of the owning parameter. + name: The owning parameter name (for error messages). + + Raises: + IncompatibleOverrideError: If the kernel specifies active dimensions or an + incompatible number of ARD dimensions. + + Returns: + A copy of the kernel bound to the given dimensions. + """ + import torch + from gpytorch.kernels import Kernel as GPyTorchKernel + + for kernel in override.modules(): + if not isinstance(kernel, GPyTorchKernel): + continue + if kernel.active_dims is not None: + raise IncompatibleOverrideError( + f"The GPyTorch kernel override for parameter '{name}' must not " + f"specify 'active_dims'." + ) + if kernel.ard_num_dims not in (None, len(indices)): + raise IncompatibleOverrideError( + f"The GPyTorch kernel override for parameter '{name}' specifies " + f"{kernel.ard_num_dims} ARD dimensions, but the parameter has " + f"{len(indices)} computational dimensions." + ) + + result = deepcopy(override) + result.active_dims = torch.tensor(indices) + return result diff --git a/baybe/surrogates/gaussian_process/_override/tl.py b/baybe/surrogates/gaussian_process/_override/tl.py new file mode 100644 index 0000000000..d796fb553e --- /dev/null +++ b/baybe/surrogates/gaussian_process/_override/tl.py @@ -0,0 +1,62 @@ +"""Resolution of the transfer-learning kernel override.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from typing_extensions import assert_never + +from baybe.kernels.base import Kernel +from baybe.parameters.enum import TransferLearningMode + +if TYPE_CHECKING: + from gpytorch.kernels import Kernel as GPyTorchKernel + + from baybe.surrogates.gaussian_process.core import _ModelContext + + +def extract_transfer_learning_overrides( + context: _ModelContext, +) -> list[tuple[str, GPyTorchKernel]]: + """Extract the transfer-learning kernel override, if any. + + Args: + context: The model context providing the task parameter and override mode. + + Returns: + A single ``(task_name, kernel)`` pair if an override is set, else an empty + list. + """ + task_param = context.searchspace._task_parameter + if task_param is None or context.tl_override is None: + return [] + return [(task_param.name, make_transfer_learning_override_kernel(context))] + + +def make_transfer_learning_override_kernel( + context: _ModelContext, +) -> GPyTorchKernel: + """Create the task kernel factor requested by a transfer-learning override. + + Args: + context: The model context providing the task parameter and override mode. + + Returns: + The GPyTorch task kernel for the requested transfer-learning mode. + """ + from baybe.kernels.basic import IndexKernel, PositiveIndexKernel + + task_param = context.searchspace._task_parameter + override = context.tl_override + assert task_param is not None and override is not None + n_tasks, names = context.n_tasks, (task_param.name,) + match override: + case TransferLearningMode.POSITIVE_INDEX_KERNEL: + spec: Kernel = PositiveIndexKernel( + num_tasks=n_tasks, rank=n_tasks, parameter_names=names + ) + case TransferLearningMode.INDEX_KERNEL: + spec = IndexKernel(num_tasks=n_tasks, rank=n_tasks, parameter_names=names) + case _: + assert_never(override) + return spec.to_gpytorch(context.searchspace) diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py index f7fc7fa416..f61445871b 100644 --- a/baybe/surrogates/gaussian_process/core.py +++ b/baybe/surrogates/gaussian_process/core.py @@ -9,30 +9,29 @@ import warnings from copy import deepcopy from functools import partial, reduce -from typing import TYPE_CHECKING, ClassVar, NoReturn +from typing import TYPE_CHECKING, ClassVar import pandas as pd from attrs import Converter, define, field from attrs.converters import optional as optional_c from attrs.converters import pipe from attrs.validators import instance_of, is_callable, optional -from typing_extensions import Self, assert_never, override +from typing_extensions import Self, override from baybe.exceptions import ( DeprecationError, - IncompatibleOverrideError, IncompatibleSearchSpaceError, ModelNotTrainedError, _UnsupportedSearchSpaceAttributeError, ) from baybe.kernels.base import Kernel -from baybe.kernels.basic import IndexKernel, PositiveIndexKernel from baybe.objectives.base import Objective from baybe.parameters.base import Parameter from baybe.parameters.categorical import TaskParameter from baybe.parameters.enum import TransferLearningMode from baybe.searchspace.core import SearchSpace from baybe.surrogates.base import Surrogate +from baybe.surrogates.gaussian_process import _override from baybe.surrogates.gaussian_process.components.fit_criterion import ( FitCriterion, FitCriterionFactoryProtocol, @@ -146,45 +145,6 @@ def _mark_custom_kernel( return value -def _bind_gpytorch_override( - override: GPyTorchKernel, indices: tuple[int, ...], name: str -) -> GPyTorchKernel: - """Copy a raw GPyTorch override and bind it to the given dimensions. - - Args: - override: The provided GPyTorch kernel (must not specify active dimensions). - indices: The computational column indices of the owning parameter. - name: The owning parameter name (for error messages). - - Raises: - IncompatibleOverrideError: If the kernel specifies active dimensions or an - incompatible number of ARD dimensions. - - Returns: - A copy of the kernel bound to the given dimensions. - """ - import torch - from gpytorch.kernels import Kernel as GPyTorchKernel - - for kernel in override.modules(): - if not isinstance(kernel, GPyTorchKernel): - continue - if kernel.active_dims is not None: - raise IncompatibleOverrideError( - f"The GPyTorch kernel override for parameter '{name}' must not " - f"specify 'active_dims'." - ) - if kernel.ard_num_dims not in (None, len(indices)): - raise IncompatibleOverrideError( - f"The GPyTorch kernel override for parameter '{name}' specifies " - f"{kernel.ard_num_dims} ARD dimensions, but the parameter has " - f"{len(indices)} computational dimensions." - ) - result = deepcopy(override) - result.active_dims = torch.tensor(indices) - return result - - @define class GaussianProcessSurrogate(Surrogate): """A Gaussian process surrogate model.""" @@ -447,8 +407,8 @@ def _resolve_kernel(self, context: _ModelContext) -> GPyTorchKernel: The effective kernel is the surrogate kernel restricted to the non-overridden dimensions, multiplied by one factor per override. """ - overrides = self._extract_parameter_overrides(context) - overrides.extend(self._extract_transfer_learning_overrides(context)) + overrides = _override.extract_parameter_overrides(context) + overrides.extend(_override.extract_transfer_learning_overrides(context)) # No overrides: let the surrogate kernel (factory) handle everything. if not overrides: @@ -456,69 +416,15 @@ def _resolve_kernel(self, context: _ModelContext) -> GPyTorchKernel: kernel = factory( context.searchspace, context.objective, context.measurements ) - return self._as_gpytorch(kernel, context.searchspace) + if isinstance(kernel, Kernel): + return kernel.to_gpytorch(context.searchspace) + return kernel excluded_names = {name for name, _ in overrides} residual = self._resolve_residual_kernel(context, excluded_names) factors = ([] if residual is None else [residual]) + [k for _, k in overrides] return reduce(operator.mul, factors) - def _extract_parameter_overrides( - self, context: _ModelContext - ) -> list[tuple[str, GPyTorchKernel]]: - """Extract regular parameter-specific kernel overrides.""" - return [ - (p.name, self._make_parameter_override_kernel(p, context.searchspace)) - for p in context.searchspace.parameters - if p.kernel_override is not None - ] - - def _extract_transfer_learning_overrides( - self, context: _ModelContext - ) -> list[tuple[str, GPyTorchKernel]]: - """Extract the transfer-learning kernel override, if any.""" - task_param = context.searchspace._task_parameter - if task_param is None or context.tl_override is None: - return [] - kernel = self._make_transfer_learning_override_kernel(context) - return [(task_param.name, kernel)] - - @staticmethod - def _make_parameter_override_kernel( - parameter: Parameter, searchspace: SearchSpace - ) -> GPyTorchKernel: - """Create the kernel factor for a parameter's override.""" - override = parameter.kernel_override - assert override is not None - indices = searchspace.get_comp_rep_parameter_indices(parameter.name) - - # BayBE kernels resolve their own dimensions; raw kernels are bound manually. - if isinstance(override, Kernel): - return override.to_gpytorch(searchspace) - return _bind_gpytorch_override(override, indices, parameter.name) - - @staticmethod - def _make_transfer_learning_override_kernel( - context: _ModelContext, - ) -> GPyTorchKernel: - """Create the task kernel factor requested by a transfer-learning override.""" - task_param = context.searchspace._task_parameter - override = context.tl_override - assert task_param is not None and override is not None - n_tasks, names = context.n_tasks, (task_param.name,) - match override: - case TransferLearningMode.POSITIVE_INDEX_KERNEL: - spec: Kernel = PositiveIndexKernel( - num_tasks=n_tasks, rank=n_tasks, parameter_names=names - ) - case TransferLearningMode.INDEX_KERNEL: - spec = IndexKernel( - num_tasks=n_tasks, rank=n_tasks, parameter_names=names - ) - case _: - assert_never(override) - return spec.to_gpytorch(context.searchspace) - def _resolve_residual_kernel( self, context: _ModelContext, excluded_names: set[str] ) -> GPyTorchKernel | None: @@ -530,18 +436,17 @@ def _resolve_residual_kernel( """ searchspace = context.searchspace factory = self.kernel_factory or BayBEKernelFactory() - remaining = [p for p in searchspace.parameters if p.name not in excluded_names] - if not remaining: + if all(p.name in excluded_names for p in searchspace.parameters): return None - # Default kernel: reuse the default machinery, but restrict it to the - # remaining (non-overridden) parameters via the factory's selector. + # Default kernel: build the (task-free) numerical base restricted to the + # remaining parameters and re-add the default task kernel unless overridden. if isinstance(factory, BayBEKernelFactory): return self._resolve_default_base(context, excluded_names) # A fixed kernel is reduced directly by removing the excluded parameters. if isinstance(factory, PlainGPComponentFactory): - spec = self._reduce_kernel_spec( + spec = _override.reduce_kernel_spec( factory.component, excluded_names, searchspace, factory ) return None if spec is None else spec.to_gpytorch(searchspace) @@ -556,8 +461,10 @@ def _resolve_residual_kernel( IncompatibleSearchSpaceError, _UnsupportedSearchSpaceAttributeError, ) as ex: - self._raise_incompatible_override(excluded_names, factory, ex) - spec = self._reduce_kernel_spec(returned, excluded_names, searchspace, factory) + _override.raise_incompatible_override(excluded_names, factory, ex) + spec = _override.reduce_kernel_spec( + returned, excluded_names, searchspace, factory + ) return None if spec is None else spec.to_gpytorch(searchspace) def _resolve_default_base( @@ -568,18 +475,13 @@ def _resolve_default_base( The numerical base always excludes the task; the default task kernel is re-added unless the task is itself overridden (i.e. already excluded). """ - from baybe.surrogates.gaussian_process.components.kernel import ( - _PureKernelFactory, - ) from baybe.surrogates.gaussian_process.presets.baybe import ( _BayBENumericalKernelFactory, _BayBETaskKernelFactory, ) searchspace = context.searchspace - factory = self.kernel_factory or BayBEKernelFactory() - assert isinstance(factory, _PureKernelFactory) - selector = factory.parameter_selector + selector = getattr(self.kernel_factory, "parameter_selector", None) task_param = searchspace._task_parameter keep: ParameterSelectorProtocol = lambda parameter: ( # noqa: E731 @@ -592,56 +494,18 @@ def _resolve_default_base( base = _BayBENumericalKernelFactory(parameter_selector=keep)( searchspace, context.objective, context.measurements ) - factors.append(self._as_gpytorch(base, searchspace)) + factors.append( + base.to_gpytorch(searchspace) if isinstance(base, Kernel) else base + ) if task_param is not None and task_param.name not in excluded_names: task = _BayBETaskKernelFactory()( searchspace, context.objective, context.measurements ) - factors.append(self._as_gpytorch(task, searchspace)) - - return reduce(operator.mul, factors) if factors else None - - @staticmethod - def _reduce_kernel_spec( - component: object, - excluded_names: set[str], - searchspace: SearchSpace, - factory: object, - ) -> Kernel | None: - """Remove the excluded parameters from a fixed BayBE kernel.""" - if not isinstance(component, Kernel): - GaussianProcessSurrogate._raise_incompatible_override( - excluded_names, factory + factors.append( + task.to_gpytorch(searchspace) if isinstance(task, Kernel) else task ) - spec: Kernel | None = component - for name in excluded_names: - if spec is None: - break - try: - spec = spec._without_parameter(name, searchspace) - except TypeError as ex: - GaussianProcessSurrogate._raise_incompatible_override( - excluded_names, factory, ex - ) - return spec - @staticmethod - def _as_gpytorch( - kernel: Kernel | GPyTorchKernel, searchspace: SearchSpace - ) -> GPyTorchKernel: - """Convert a BayBE kernel to GPyTorch, passing through raw kernels.""" - return kernel.to_gpytorch(searchspace) if isinstance(kernel, Kernel) else kernel - - @staticmethod - def _raise_incompatible_override( - parameter_names: set[str], factory: object, cause: Exception | None = None - ) -> NoReturn: - """Raise an error for a surrogate kernel that cannot be reduced.""" - raise IncompatibleOverrideError( - f"Kernel overrides for {sorted(parameter_names)} require a surrogate " - f"kernel (factory) that can exclude these parameters. " - f"'{type(factory).__name__}' does not satisfy this requirement." - ) from cause + return reduce(operator.mul, factors) if factors else None def _resolve_components( self, context: _ModelContext @@ -798,8 +662,6 @@ def _make_posterior_mean_module( Returns: A mean module ready for use in a new GP. """ - from copy import deepcopy - import gpytorch frozen_model = deepcopy(model) From c0165b17298c225aff0b48c905c999343ca867ea Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Wed, 9 Sep 2026 12:45:41 +0200 Subject: [PATCH 10/73] Test selector precedence with kernel overrides The unified resolver constructs the default numerical kernel on the full search space with a composed selector. This removes PR #868's reduced-space limitation for selector-bearing default factories. Replace the old rejection expectation and verify that selectors restrict the residual while explicit regular and TL overrides still apply. --- tests/test_kernel_factories.py | 97 ++++++++++++++++++++++++++++++---- 1 file changed, 86 insertions(+), 11 deletions(-) diff --git a/tests/test_kernel_factories.py b/tests/test_kernel_factories.py index a4904b4358..cda0d80fd6 100644 --- a/tests/test_kernel_factories.py +++ b/tests/test_kernel_factories.py @@ -12,7 +12,7 @@ from pytest import param from baybe.exceptions import IncompatibleOverrideError, IncompatibleSearchSpaceError -from baybe.kernels.basic import IndexKernel, MaternKernel +from baybe.kernels.basic import IndexKernel, MaternKernel, RBFKernel from baybe.kernels.composite import ScaleKernel from baybe.parameters.categorical import ( CategoricalParameter, @@ -23,7 +23,7 @@ NumericalContinuousParameter, NumericalDiscreteParameter, ) -from baybe.parameters.selectors import TypeSelector +from baybe.parameters.selectors import NameSelector, TypeSelector from baybe.searchspace.core import SearchSpace from baybe.surrogates import GaussianProcessSurrogate from baybe.surrogates.gaussian_process.components.kernel import ICMKernelFactory @@ -248,6 +248,24 @@ def _make_dispatch_context(override_mode): True, id="index_override+default_factory", ), + param( + TransferLearningMode.POSITIVE_INDEX_KERNEL, + BayBEKernelFactory( + parameter_selector=TypeSelector((TaskParameter,), exclude=True) + ), + GPyTorchPositiveIndexKernel, + True, + id="positive_index_override+default_factory_with_selector", + ), + param( + TransferLearningMode.INDEX_KERNEL, + BayBEKernelFactory( + parameter_selector=TypeSelector((TaskParameter,), exclude=True) + ), + GPyTorchIndexKernel, + True, + id="index_override+default_factory_with_selector", + ), ], ) def test_resolve_kernel_dispatch_success( @@ -307,15 +325,6 @@ def test_resolve_kernel_dispatch_success( _gpytorch_returning_factory, id="override+factory_returning_gpytorch", ), - param( - TransferLearningMode.INDEX_KERNEL, - BayBEKernelFactory( - parameter_selector=TypeSelector((TaskParameter,), exclude=True) - ), - # On a searchspace without a substance parameter, the default factory - # yields a gpytorch kernel and cannot be combined with an override. - id="override+default_factory_with_selector_unsupported", - ), ], ) def test_resolve_kernel_dispatch_raises(monkeypatch, override_mode, kernel_or_factory): @@ -333,3 +342,69 @@ def test_resolve_kernel_dispatch_raises(monkeypatch, override_mode, kernel_or_fa with pytest.raises(IncompatibleOverrideError): surrogate._resolve_kernel(context) + + +@pytest.mark.parametrize( + ("override_mode", "expected_task_kernel_cls"), + [ + param( + TransferLearningMode.POSITIVE_INDEX_KERNEL, + GPyTorchPositiveIndexKernel, + id="positive-index", + ), + param(TransferLearningMode.INDEX_KERNEL, GPyTorchIndexKernel, id="index"), + ], +) +@pytest.mark.parametrize( + "has_regular_override", [False, True], ids=["tl-only", "tl-and-regular"] +) +def test_default_factory_selector_override_precedence( + override_mode: TransferLearningMode, + expected_task_kernel_cls: type[gpytorch.kernels.Kernel], + has_regular_override: bool, +) -> None: + """Selectors restrict the base kernel but cannot exclude parameter overrides.""" + from baybe.surrogates.gaussian_process.core import _ModelContext + + searchspace = SearchSpace.from_product( + [ + NumericalContinuousParameter("x1", (0, 1)), + NumericalContinuousParameter( + "x2", + (0, 1), + kernel_override=RBFKernel() if has_regular_override else None, + ), + NumericalContinuousParameter("omitted", (0, 1)), + TaskParameter( + "Task", ["A", "B"], override_transfer_learning_mode=override_mode + ), + ] + ) + selector = NameSelector(("x1",), regex=False) + assert [p.name for p in searchspace.parameters if selector(p)] == ["x1"] + surrogate = GaussianProcessSurrogate( + kernel_or_factory=BayBEKernelFactory(parameter_selector=selector) + ) + context = _ModelContext( + searchspace, NumericalTarget("y").to_objective(), pd.DataFrame() + ) + + kernel = surrogate._resolve_kernel(context) + + expected_names = ["x1"] + (["x2"] if has_regular_override else []) + ["Task"] + expected_types = ( + [gpytorch.kernels.MaternKernel] + + ([gpytorch.kernels.RBFKernel] if has_regular_override else []) + + [expected_task_kernel_cls] + ) + expected_dimensions = [ + searchspace.get_comp_rep_parameter_indices(name) for name in expected_names + ] + assert isinstance(kernel, gpytorch.kernels.ProductKernel) + assert [type(k) for k in kernel.kernels] == expected_types + assert [ + tuple(k.active_dims.tolist()) for k in kernel.kernels + ] == expected_dimensions + assert [k.ard_num_dims for k in kernel.kernels] == [ + len(d) for d in expected_dimensions + ] From a306cfdf99c59d98b67abccf83b06ac7f17c0f3b Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Wed, 9 Sep 2026 12:58:16 +0200 Subject: [PATCH 11/73] Preserve kernel partition validation in unified overrides Final PR #868 validates resolved base and task dimensions through ICM. Restore these checks in the generalized resolver and extend them to regular overrides, while accepting standard composites and selector omissions. Compare rejection cases against ICM and cover combined override iterations and raw-kernel boundaries. --- .../gaussian_process/_override/__init__.py | 2 + .../gaussian_process/_override/core.py | 38 +++++- baybe/surrogates/gaussian_process/core.py | 54 ++++++-- tests/test_iterations.py | 18 ++- tests/test_parameter_kernel_overrides.py | 122 ++++++++++++++++++ 5 files changed, 217 insertions(+), 17 deletions(-) diff --git a/baybe/surrogates/gaussian_process/_override/__init__.py b/baybe/surrogates/gaussian_process/_override/__init__.py index 081d26149d..579dc9db3c 100644 --- a/baybe/surrogates/gaussian_process/_override/__init__.py +++ b/baybe/surrogates/gaussian_process/_override/__init__.py @@ -1,6 +1,7 @@ """Kernel override resolution for Gaussian process surrogates.""" from baybe.surrogates.gaussian_process._override.core import ( + get_active_dimensions, raise_incompatible_override, reduce_kernel_spec, ) @@ -12,6 +13,7 @@ ) __all__ = [ + "get_active_dimensions", "extract_parameter_overrides", "extract_transfer_learning_overrides", "raise_incompatible_override", diff --git a/baybe/surrogates/gaussian_process/_override/core.py b/baybe/surrogates/gaussian_process/_override/core.py index 2be7a1c312..33cb651520 100644 --- a/baybe/surrogates/gaussian_process/_override/core.py +++ b/baybe/surrogates/gaussian_process/_override/core.py @@ -8,14 +8,43 @@ from baybe.kernels.base import Kernel if TYPE_CHECKING: + from gpytorch.kernels import Kernel as GPyTorchKernel + from baybe.searchspace import SearchSpace + from baybe.surrogates.gaussian_process.components.kernel import ( + KernelFactoryProtocol, + ) + + +def get_active_dimensions(kernel: GPyTorchKernel, searchspace: SearchSpace) -> set[int]: + """Get the input columns used by a kernel, including standard composites. + + Args: + kernel: The resolved kernel to inspect. + searchspace: The search space defining the full model input. + + Returns: + The active input column indices. + """ + from gpytorch.kernels import AdditiveKernel, ProductKernel, ScaleKernel + + if kernel.active_dims is not None: + # TODO[typing]: https://github.com/facebook/pyrefly/issues/3988 + return set(kernel.active_dims.tolist()) # pyrefly: ignore[not-callable] + if isinstance(kernel, (AdditiveKernel, ProductKernel)): + return set().union( + *(get_active_dimensions(k, searchspace) for k in kernel.kernels) + ) + if isinstance(kernel, ScaleKernel): + return get_active_dimensions(kernel.base_kernel, searchspace) + return set(range(len(searchspace.comp_rep_columns))) def reduce_kernel_spec( - component: object, + component: Kernel | GPyTorchKernel, excluded_names: set[str], searchspace: SearchSpace, - factory: object, + factory: KernelFactoryProtocol, ) -> Kernel | None: """Remove the excluded parameters from a fixed BayBE kernel. @@ -38,11 +67,14 @@ def reduce_kernel_spec( spec = spec._without_parameter(name, searchspace) except TypeError as ex: raise_incompatible_override(excluded_names, factory, ex) + return spec def raise_incompatible_override( - parameter_names: set[str], factory: object, cause: Exception | None = None + parameter_names: set[str], + factory: KernelFactoryProtocol, + cause: Exception | None = None, ) -> NoReturn: """Raise an error for a surrogate kernel that cannot be reduced. diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py index f61445871b..ab09de793e 100644 --- a/baybe/surrogates/gaussian_process/core.py +++ b/baybe/surrogates/gaussian_process/core.py @@ -7,7 +7,6 @@ import operator import os import warnings -from copy import deepcopy from functools import partial, reduce from typing import TYPE_CHECKING, ClassVar @@ -406,22 +405,43 @@ def _resolve_kernel(self, context: _ModelContext) -> GPyTorchKernel: The effective kernel is the surrogate kernel restricted to the non-overridden dimensions, multiplied by one factor per override. + + Args: + context: The model context providing the inputs and override settings. + + Raises: + ValueError: If the resolved kernels violate the override partition. + + Returns: + The resolved and partition-validated kernel. """ overrides = _override.extract_parameter_overrides(context) - overrides.extend(_override.extract_transfer_learning_overrides(context)) - - # No overrides: let the surrogate kernel (factory) handle everything. - if not overrides: - factory = self.kernel_factory or BayBEKernelFactory() - kernel = factory( - context.searchspace, context.objective, context.measurements - ) - if isinstance(kernel, Kernel): - return kernel.to_gpytorch(context.searchspace) - return kernel + overrides += _override.extract_transfer_learning_overrides(context) excluded_names = {name for name, _ in overrides} residual = self._resolve_residual_kernel(context, excluded_names) + searchspace = context.searchspace + excluded_dimensions: set[int] = set() + for name, kernel in overrides: + expected = set(searchspace.get_comp_rep_parameter_indices(name)) + actual = _override.get_active_dimensions(kernel, searchspace) + if actual != expected: + raise ValueError( + f"The kernel override for '{name}' has 'active_dims' {actual}, " + f"but must use exactly the parameter indices {expected}." + ) + excluded_dimensions.update(expected) + + if residual is not None and excluded_names: + allowed = ( + set(range(len(searchspace.comp_rep_columns))) - excluded_dimensions + ) + actual = _override.get_active_dimensions(residual, searchspace) + if not actual <= allowed: + raise ValueError( + f"The residual kernel's 'active_dims' {actual} must be a subset " + f"of the non-overridden indices {allowed}." + ) factors = ([] if residual is None else [residual]) + [k for _, k in overrides] return reduce(operator.mul, factors) @@ -436,6 +456,14 @@ def _resolve_residual_kernel( """ searchspace = context.searchspace factory = self.kernel_factory or BayBEKernelFactory() + if not excluded_names: + kernel = factory(searchspace, context.objective, context.measurements) + return ( + kernel.to_gpytorch(searchspace) + if isinstance(kernel, Kernel) + else kernel + ) + if all(p.name in excluded_names for p in searchspace.parameters): return None @@ -662,6 +690,8 @@ def _make_posterior_mean_module( Returns: A mean module ready for use in a new GP. """ + from copy import deepcopy + import gpytorch frozen_model = deepcopy(model) diff --git a/tests/test_iterations.py b/tests/test_iterations.py index e5ceb8ceca..164f4adbbe 100644 --- a/tests/test_iterations.py +++ b/tests/test_iterations.py @@ -27,7 +27,11 @@ ) from baybe.kernels.composite import ScaleKernel from baybe.objectives.pareto import ParetoObjective -from baybe.parameters import NumericalContinuousParameter +from baybe.parameters import ( + NumericalContinuousParameter, + TaskParameter, + TransferLearningMode, +) from baybe.priors import ( GammaPrior, HalfCauchyPrior, @@ -406,9 +410,19 @@ def test_kernels(ongoing_campaign, n_iterations, batch_size): [ NumericalContinuousParameter("x1", (0, 1)), NumericalContinuousParameter("x2", (0, 1), kernel_override=RBFKernel()), + *( + [] + if mode is None + else [ + TaskParameter( + "task", ["a", "b"], override_transfer_learning_mode=mode + ) + ] + ), ] + for mode in [None, *TransferLearningMode] ], - ids=["parameter_kernel_override"], + ids=["parameter_kernel_override", *(mode.name for mode in TransferLearningMode)], ) @pytest.mark.parametrize("n_iterations", [3], ids=["i3"]) def test_parameter_kernel_override_iteration( diff --git a/tests/test_parameter_kernel_overrides.py b/tests/test_parameter_kernel_overrides.py index 2395b151f3..5979186e56 100644 --- a/tests/test_parameter_kernel_overrides.py +++ b/tests/test_parameter_kernel_overrides.py @@ -15,6 +15,7 @@ from baybe.exceptions import IncompatibleOverrideError from baybe.kernels import MaternKernel, RBFKernel +from baybe.kernels.basic import IndexKernel from baybe.kernels.composite import AdditiveKernel, ProductKernel, ScaleKernel from baybe.parameters import ( CategoricalParameter, @@ -26,6 +27,7 @@ from baybe.parameters.selectors import NameSelector from baybe.searchspace import SearchSpace from baybe.surrogates import GaussianProcessSurrogate +from baybe.surrogates.gaussian_process.components.kernel import ICMKernelFactory from baybe.surrogates.gaussian_process.core import _ModelContext from baybe.surrogates.gaussian_process.presets import BayBEKernelFactory from baybe.targets import NumericalTarget @@ -294,3 +296,123 @@ def test_incompatible_surrogate_kernel_is_rejected(kernel_or_factory): with pytest.raises(IncompatibleOverrideError): _resolve(parameters, kernel_or_factory) + + +@pytest.mark.parametrize("resolver", ["icm", "unified"], ids=["icm", "unified"]) +@pytest.mark.parametrize( + ("factor", "active_dims"), + [ + param("residual", None, id="unrestricted-residual"), + param("residual", (0, 1), id="overlapping-residual"), + param("task", None, id="unrestricted-task"), + param("task", (0, 1), id="overlapping-task"), + param("task", (), id="empty-task"), + ], +) +def test_tl_partition_validation(monkeypatch, resolver, factor, active_dims): + """The unified resolver retains ICM's rejection of misbound kernel outputs.""" + parameters = [ + NumericalContinuousParameter("x", (0, 1)), + TaskParameter( + "task", + ["a", "b"], + override_transfer_learning_mode=TransferLearningMode.INDEX_KERNEL, + ), + ] + searchspace = SearchSpace.from_product(parameters) + objective = NumericalTarget("y").to_objective() + measurements = pd.DataFrame() + kernel_cls = MaternKernel if factor == "residual" else IndexKernel + + def misbound_kernel(self, searchspace): + return gk.RBFKernel(active_dims=active_dims) + + monkeypatch.setattr(kernel_cls, "to_gpytorch", misbound_kernel) + with pytest.raises(ValueError, match="active_dims"): + if resolver == "icm": + ICMKernelFactory( + base_kernel_or_factory=MaternKernel(parameter_names=("x",)), + task_kernel_or_factory=IndexKernel( + num_tasks=2, rank=2, parameter_names=("task",) + ), + )(searchspace, objective, measurements) + else: + GaussianProcessSurrogate(kernel_or_factory=MaternKernel())._resolve_kernel( + _ModelContext(searchspace, objective, measurements) + ) + + +@pytest.mark.parametrize( + "factor", ["residual", "override"], ids=["residual", "override"] +) +def test_regular_partition_validation(monkeypatch, factor): + """Partition validation also rejects overlap for regular parameter overrides.""" + + def misbound_kernel(self, searchspace): + return gk.RBFKernel() + + monkeypatch.setattr( + MaternKernel if factor == "residual" else RBFKernel, + "to_gpytorch", + misbound_kernel, + ) + with pytest.raises(ValueError, match="active_dims"): + _resolve( + [ + NumericalContinuousParameter("x", (0, 1)), + NumericalContinuousParameter("y", (0, 1), kernel_override=RBFKernel()), + ], + MaternKernel(), + ) + + +@pytest.mark.parametrize( + "override", + [ + param(ScaleKernel(RBFKernel()), id="scale"), + param(AdditiveKernel([RBFKernel(), MaternKernel()]), id="additive"), + param(ProductKernel([RBFKernel(), MaternKernel()]), id="product"), + param(gk.ScaleKernel(gk.RBFKernel()), id="raw-scale"), + ], +) +def test_composite_partition_omits_unselected_columns(override): + """Composite factors may share their owner while unused columns stay irrelevant.""" + kernel, searchspace = _resolve( + [ + NumericalContinuousParameter("x", (0, 1)), + NumericalContinuousParameter("y", (0, 1), kernel_override=override), + NumericalContinuousParameter("omitted", (0, 1)), + TaskParameter("task", ["a", "b"]), + ], + BayBEKernelFactory(parameter_selector=NameSelector(("x",), regex=False)), + ) + inputs = torch.zeros(2, len(searchspace.comp_rep_columns)) + baseline = kernel(inputs).to_dense() + omitted = searchspace.get_comp_rep_parameter_indices("omitted") + inputs[1, list(omitted)] = 1.0 + torch.testing.assert_close(kernel(inputs).to_dense(), baseline) + overridden = searchspace.get_comp_rep_parameter_indices("y") + inputs[1, list(overridden)] = 1.0 + assert not torch.allclose(kernel(inputs).to_dense(), baseline) + + +@pytest.mark.parametrize("as_factory", [False, True], ids=["fixed", "callable"]) +def test_raw_surrogate_kernel_without_overrides(as_factory): + """Without overrides, raw surrogate kernels are passed through unchanged.""" + raw = gk.RBFKernel(active_dims=(0,)) + specification = (lambda s, o, m: raw) if as_factory else raw + kernel, _ = _resolve([NumericalContinuousParameter("x", (0, 1))], specification) + assert kernel is raw + + +def test_all_overridden_bypasses_surrogate_factory(): + """A fully overridden space does not construct an unused residual kernel.""" + + def unused_factory(searchspace, objective, measurements): + pytest.fail("The residual factory must not be called.") + + kernel, _ = _resolve( + [NumericalContinuousParameter("x", (0, 1), kernel_override=RBFKernel())], + unused_factory, + ) + assert isinstance(kernel, gk.RBFKernel) From 361d431565aebf702ac3ab5b5fff8d8a9e37edfe Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Wed, 9 Sep 2026 14:25:18 +0200 Subject: [PATCH 12/73] Consolidate kernel override regression tests --- tests/test_kernel_factories.py | 70 +---- tests/test_parameter_kernel_overrides.py | 339 ++++++++--------------- 2 files changed, 112 insertions(+), 297 deletions(-) diff --git a/tests/test_kernel_factories.py b/tests/test_kernel_factories.py index cda0d80fd6..a8c12b7df5 100644 --- a/tests/test_kernel_factories.py +++ b/tests/test_kernel_factories.py @@ -12,7 +12,7 @@ from pytest import param from baybe.exceptions import IncompatibleOverrideError, IncompatibleSearchSpaceError -from baybe.kernels.basic import IndexKernel, MaternKernel, RBFKernel +from baybe.kernels.basic import IndexKernel, MaternKernel from baybe.kernels.composite import ScaleKernel from baybe.parameters.categorical import ( CategoricalParameter, @@ -23,7 +23,7 @@ NumericalContinuousParameter, NumericalDiscreteParameter, ) -from baybe.parameters.selectors import NameSelector, TypeSelector +from baybe.parameters.selectors import TypeSelector from baybe.searchspace.core import SearchSpace from baybe.surrogates import GaussianProcessSurrogate from baybe.surrogates.gaussian_process.components.kernel import ICMKernelFactory @@ -342,69 +342,3 @@ def test_resolve_kernel_dispatch_raises(monkeypatch, override_mode, kernel_or_fa with pytest.raises(IncompatibleOverrideError): surrogate._resolve_kernel(context) - - -@pytest.mark.parametrize( - ("override_mode", "expected_task_kernel_cls"), - [ - param( - TransferLearningMode.POSITIVE_INDEX_KERNEL, - GPyTorchPositiveIndexKernel, - id="positive-index", - ), - param(TransferLearningMode.INDEX_KERNEL, GPyTorchIndexKernel, id="index"), - ], -) -@pytest.mark.parametrize( - "has_regular_override", [False, True], ids=["tl-only", "tl-and-regular"] -) -def test_default_factory_selector_override_precedence( - override_mode: TransferLearningMode, - expected_task_kernel_cls: type[gpytorch.kernels.Kernel], - has_regular_override: bool, -) -> None: - """Selectors restrict the base kernel but cannot exclude parameter overrides.""" - from baybe.surrogates.gaussian_process.core import _ModelContext - - searchspace = SearchSpace.from_product( - [ - NumericalContinuousParameter("x1", (0, 1)), - NumericalContinuousParameter( - "x2", - (0, 1), - kernel_override=RBFKernel() if has_regular_override else None, - ), - NumericalContinuousParameter("omitted", (0, 1)), - TaskParameter( - "Task", ["A", "B"], override_transfer_learning_mode=override_mode - ), - ] - ) - selector = NameSelector(("x1",), regex=False) - assert [p.name for p in searchspace.parameters if selector(p)] == ["x1"] - surrogate = GaussianProcessSurrogate( - kernel_or_factory=BayBEKernelFactory(parameter_selector=selector) - ) - context = _ModelContext( - searchspace, NumericalTarget("y").to_objective(), pd.DataFrame() - ) - - kernel = surrogate._resolve_kernel(context) - - expected_names = ["x1"] + (["x2"] if has_regular_override else []) + ["Task"] - expected_types = ( - [gpytorch.kernels.MaternKernel] - + ([gpytorch.kernels.RBFKernel] if has_regular_override else []) - + [expected_task_kernel_cls] - ) - expected_dimensions = [ - searchspace.get_comp_rep_parameter_indices(name) for name in expected_names - ] - assert isinstance(kernel, gpytorch.kernels.ProductKernel) - assert [type(k) for k in kernel.kernels] == expected_types - assert [ - tuple(k.active_dims.tolist()) for k in kernel.kernels - ] == expected_dimensions - assert [k.ard_num_dims for k in kernel.kernels] == [ - len(d) for d in expected_dimensions - ] diff --git a/tests/test_parameter_kernel_overrides.py b/tests/test_parameter_kernel_overrides.py index 5979186e56..f3ec7be1ec 100644 --- a/tests/test_parameter_kernel_overrides.py +++ b/tests/test_parameter_kernel_overrides.py @@ -9,6 +9,7 @@ import pandas as pd import pytest import torch +from attrs import evolve from botorch.models.kernels.positive_index import PositiveIndexKernel from gpytorch import kernels as gk from pytest import param @@ -53,148 +54,98 @@ def _leaf_kernels(kernel): ) -def _fit(parameters): - """Fit a GP and return its installed covariance module and searchspace.""" - searchspace = SearchSpace.from_product(parameters) - measurements = pd.DataFrame( - [ - {**{p.name: p.values[0] for p in parameters}, "y": 0.0}, - {**{p.name: p.values[-1] for p in parameters}, "y": 1.0}, - ] - ) - surrogate = GaussianProcessSurrogate() - surrogate.fit(searchspace, NumericalTarget("y").to_objective(), measurements) - assert surrogate._model is not None - return surrogate._model.covar_module, searchspace - - -def test_default_factory_selector_is_preserved(): - """Partitioning preserves an explicit default-factory parameter selector.""" - kernel, _ = _resolve( - [ - NumericalContinuousParameter("x1", (0, 1)), - NumericalContinuousParameter("x2", (0, 1), kernel_override=RBFKernel()), - NumericalContinuousParameter("x3", (0, 1)), - ], +@pytest.mark.parametrize( + "mode", + [None, *TransferLearningMode], + ids=lambda mode: mode.name if mode else "no-task", +) +def test_default_factory_selector_is_preserved(mode): + """Selectors omit unselected dimensions but cannot exclude parameter overrides.""" + parameters = [ + NumericalContinuousParameter("x1", (0, 1)), + NumericalContinuousParameter("x2", (0, 1), kernel_override=RBFKernel()), + NumericalContinuousParameter("omitted", (0, 1)), + ] + expected_names = ["x1", "x2"] + if mode is not None: + parameters.append( + TaskParameter("task", ["a", "b"], override_transfer_learning_mode=mode) + ) + expected_names.append("task") + kernel, searchspace = _resolve( + parameters, BayBEKernelFactory(parameter_selector=NameSelector(("x1",), regex=False)), ) - assert {tuple(k.active_dims.tolist()) for k in _leaf_kernels(kernel)} == { - (0,), - (1,), + searchspace.get_comp_rep_parameter_indices(name) for name in expected_names } @pytest.mark.parametrize( - ("parameters", "expected_types", "expected_names"), + ("override_parameter", "base_override", "task_parameter"), [ - param( - [ - CategoricalParameter("base", ["a", "b"]), - NumericalDiscreteParameter( - "override", [0, 1, 2], kernel_override=RBFKernel() - ), - ], - [gk.MaternKernel, gk.RBFKernel], - ["base", "override"], - id="single-dim", - ), - param( - [ - CategoricalParameter("base", ["a", "b"]), - CategoricalParameter( - "override", ["a", "b", "c"], kernel_override=RBFKernel() - ), - ], - [gk.MaternKernel, gk.RBFKernel], - ["base", "override"], - id="multi-dim", - ), - param( - [ - CategoricalParameter("first", ["a", "b"], kernel_override=RBFKernel()), - CategoricalParameter( - "second", ["a", "b"], kernel_override=MaternKernel() - ), - ], - [gk.RBFKernel, gk.MaternKernel], - ["first", "second"], - id="all-overridden", - ), - param( - [ - CategoricalParameter("base", ["a", "b"]), - CategoricalParameter( - "override", ["a", "b"], kernel_override=RBFKernel() - ), - TaskParameter("task", ["a", "b"]), - ], - [ - gk.MaternKernel, - PositiveIndexKernel, - gk.RBFKernel, - ], - ["base", "task", "override"], - id="task-without-tl-override", - ), - param( - [ - CategoricalParameter("base", ["a", "b"]), - CategoricalParameter( - "override", ["a", "b", "c"], kernel_override=RBFKernel() - ), - TaskParameter( - "task", - ["a", "b"], - override_transfer_learning_mode=TransferLearningMode.INDEX_KERNEL, - ), - ], - [ - gk.MaternKernel, - gk.RBFKernel, - gk.IndexKernel, - ], - ["base", "override", "task"], - id="override-and-tl-index", - ), - param( - [ - CategoricalParameter("base", ["a", "b"]), - CategoricalParameter( - "override", ["a", "b", "c"], kernel_override=RBFKernel() - ), - TaskParameter( - "task", - ["a", "b"], - override_transfer_learning_mode=( - TransferLearningMode.POSITIVE_INDEX_KERNEL - ), - ), - ], - [ - gk.MaternKernel, - gk.RBFKernel, - PositiveIndexKernel, - ], - ["base", "override", "task"], - id="override-and-tl-pos-index", + (NumericalDiscreteParameter("override", [0, 1, 2]), None, None), + (CategoricalParameter("override", ["a", "b", "c"]), None, None), + (CategoricalParameter("override", ["a", "b"]), MaternKernel(), None), + ( + CategoricalParameter("override", ["a", "b"]), + None, + TaskParameter("task", ["a", "b"]), ), + *[ + ( + CategoricalParameter("override", ["a", "b", "c"]), + None, + TaskParameter("task", ["a", "b"], override_transfer_learning_mode=mode), + ) + for mode in TransferLearningMode + ], + ], + ids=[ + "single-dim", + "multi-dim", + "all-overridden", + "task-without-tl-override", + *(mode.name for mode in TransferLearningMode), ], ) def test_fitted_model_uses_parameter_kernel_overrides( - parameters, expected_types, expected_names + override_parameter, base_override, task_parameter ): """The fitted model uses each configured kernel on the intended dimensions.""" - kernel, searchspace = _fit(parameters) - leaves = _leaf_kernels(kernel) - expected_dimensions = [ - searchspace.get_comp_rep_parameter_indices(name) for name in expected_names + parameters = [ + CategoricalParameter("base", ["a", "b"], kernel_override=base_override), + evolve(override_parameter, kernel_override=RBFKernel()), ] + expected = {"base": gk.MaternKernel, "override": gk.RBFKernel} + if task_parameter is not None: + parameters.append(task_parameter) + expected["task"] = ( + gk.IndexKernel + if task_parameter.override_transfer_learning_mode + == TransferLearningMode.INDEX_KERNEL + else PositiveIndexKernel + ) + searchspace = SearchSpace.from_product(parameters) + measurements = pd.DataFrame( + [ + {**{p.name: p.values[i] for p in parameters}, "y": y} + for i, y in [(0, 0.0), (-1, 1.0)] + ] + ) + surrogate = GaussianProcessSurrogate() + surrogate.fit(searchspace, NumericalTarget("y").to_objective(), measurements) + assert surrogate._model is not None + kernel = surrogate._model.covar_module + leaves = _leaf_kernels(kernel) assert isinstance(kernel, gk.ProductKernel) - assert [type(k) for k in leaves] == expected_types - assert [tuple(k.active_dims.tolist()) for k in leaves] == expected_dimensions - assert [k.ard_num_dims for k in leaves] == [len(d) for d in expected_dimensions] + assert len(leaves) == len(expected) + assert {tuple(k.active_dims.tolist()): type(k) for k in leaves} == { + searchspace.get_comp_rep_parameter_indices(name): cls + for name, cls in expected.items() + } + assert all(k.ard_num_dims == len(k.active_dims) for k in leaves) @pytest.mark.parametrize( @@ -215,12 +166,10 @@ def test_valid_baybe_kernel_override(override): ] kernel, searchspace = _resolve(parameters) expected = set(searchspace.get_comp_rep_parameter_indices("override")) + base_dims = set(searchspace.get_comp_rep_parameter_indices("base")) override_leaves = [ - k - for k in _leaf_kernels(kernel) - if set(k.active_dims.tolist()) - != set(searchspace.get_comp_rep_parameter_indices("base")) + k for k in _leaf_kernels(kernel) if set(k.active_dims.tolist()) != base_dims ] assert override_leaves assert all(set(k.active_dims.tolist()) == expected for k in override_leaves) @@ -261,18 +210,6 @@ def test_gpytorch_ard_mismatch_rejected(): _resolve([parameter]) -def test_parameter_equivalence_ignores_override_parameter_name(): - """Equivalent parameters may scope the same override to their own names.""" - left = NumericalContinuousParameter( - "left", (0, 1), kernel_override=RBFKernel(parameter_names=("left",)) - ) - right = NumericalContinuousParameter( - "right", (0, 1), kernel_override=RBFKernel(parameter_names=("right",)) - ) - - assert left.is_equivalent(right) - - @pytest.mark.parametrize( "kernel_or_factory", [ @@ -281,10 +218,7 @@ def test_parameter_equivalence_ignores_override_parameter_name(): id="non-reducible-baybe-kernel", ), param(gk.MaternKernel(), id="raw-gpytorch-kernel"), - param( - lambda s, o, m: gk.MaternKernel(), - id="factory-returning-raw-kernel", - ), + param(lambda s, o, m: gk.MaternKernel(), id="factory-returning-raw-kernel"), ], ) def test_incompatible_surrogate_kernel_is_rejected(kernel_or_factory): @@ -294,125 +228,72 @@ def test_incompatible_surrogate_kernel_is_rejected(kernel_or_factory): NumericalContinuousParameter("x2", (0, 1), kernel_override=RBFKernel()), ] - with pytest.raises(IncompatibleOverrideError): + with pytest.raises(IncompatibleOverrideError, match="can exclude these parameters"): _resolve(parameters, kernel_or_factory) -@pytest.mark.parametrize("resolver", ["icm", "unified"], ids=["icm", "unified"]) @pytest.mark.parametrize( - ("factor", "active_dims"), + ("override_kind", "kernel_cls", "active_dims"), [ - param("residual", None, id="unrestricted-residual"), - param("residual", (0, 1), id="overlapping-residual"), - param("task", None, id="unrestricted-task"), - param("task", (0, 1), id="overlapping-task"), - param("task", (), id="empty-task"), + param("tl", MaternKernel, None, id="tl-unrestricted-residual"), + param("tl", MaternKernel, (0, 1), id="tl-overlapping-residual"), + param("tl", IndexKernel, None, id="unrestricted-task"), + param("tl", IndexKernel, (0, 1), id="overlapping-task"), + param("tl", IndexKernel, (), id="empty-task"), + param("regular", MaternKernel, None, id="regular-unrestricted-residual"), + param("regular", RBFKernel, None, id="unrestricted-override"), ], ) -def test_tl_partition_validation(monkeypatch, resolver, factor, active_dims): - """The unified resolver retains ICM's rejection of misbound kernel outputs.""" +def test_partition_validation(monkeypatch, override_kind, kernel_cls, active_dims): + """Reject misbound regular and TL factors, matching ICM for TL partitions.""" parameters = [ NumericalContinuousParameter("x", (0, 1)), TaskParameter( "task", ["a", "b"], override_transfer_learning_mode=TransferLearningMode.INDEX_KERNEL, + ) + if override_kind == "tl" + else NumericalContinuousParameter( + "override", (0, 1), kernel_override=RBFKernel() ), ] searchspace = SearchSpace.from_product(parameters) objective = NumericalTarget("y").to_objective() measurements = pd.DataFrame() - kernel_cls = MaternKernel if factor == "residual" else IndexKernel def misbound_kernel(self, searchspace): return gk.RBFKernel(active_dims=active_dims) monkeypatch.setattr(kernel_cls, "to_gpytorch", misbound_kernel) with pytest.raises(ValueError, match="active_dims"): - if resolver == "icm": + GaussianProcessSurrogate(kernel_or_factory=MaternKernel())._resolve_kernel( + _ModelContext(searchspace, objective, measurements) + ) + if override_kind == "tl": + with pytest.raises(ValueError, match="active_dims"): ICMKernelFactory( base_kernel_or_factory=MaternKernel(parameter_names=("x",)), task_kernel_or_factory=IndexKernel( num_tasks=2, rank=2, parameter_names=("task",) ), )(searchspace, objective, measurements) - else: - GaussianProcessSurrogate(kernel_or_factory=MaternKernel())._resolve_kernel( - _ModelContext(searchspace, objective, measurements) - ) - - -@pytest.mark.parametrize( - "factor", ["residual", "override"], ids=["residual", "override"] -) -def test_regular_partition_validation(monkeypatch, factor): - """Partition validation also rejects overlap for regular parameter overrides.""" - - def misbound_kernel(self, searchspace): - return gk.RBFKernel() - - monkeypatch.setattr( - MaternKernel if factor == "residual" else RBFKernel, - "to_gpytorch", - misbound_kernel, - ) - with pytest.raises(ValueError, match="active_dims"): - _resolve( - [ - NumericalContinuousParameter("x", (0, 1)), - NumericalContinuousParameter("y", (0, 1), kernel_override=RBFKernel()), - ], - MaternKernel(), - ) - - -@pytest.mark.parametrize( - "override", - [ - param(ScaleKernel(RBFKernel()), id="scale"), - param(AdditiveKernel([RBFKernel(), MaternKernel()]), id="additive"), - param(ProductKernel([RBFKernel(), MaternKernel()]), id="product"), - param(gk.ScaleKernel(gk.RBFKernel()), id="raw-scale"), - ], -) -def test_composite_partition_omits_unselected_columns(override): - """Composite factors may share their owner while unused columns stay irrelevant.""" - kernel, searchspace = _resolve( - [ - NumericalContinuousParameter("x", (0, 1)), - NumericalContinuousParameter("y", (0, 1), kernel_override=override), - NumericalContinuousParameter("omitted", (0, 1)), - TaskParameter("task", ["a", "b"]), - ], - BayBEKernelFactory(parameter_selector=NameSelector(("x",), regex=False)), - ) - inputs = torch.zeros(2, len(searchspace.comp_rep_columns)) - baseline = kernel(inputs).to_dense() - omitted = searchspace.get_comp_rep_parameter_indices("omitted") - inputs[1, list(omitted)] = 1.0 - torch.testing.assert_close(kernel(inputs).to_dense(), baseline) - overridden = searchspace.get_comp_rep_parameter_indices("y") - inputs[1, list(overridden)] = 1.0 - assert not torch.allclose(kernel(inputs).to_dense(), baseline) @pytest.mark.parametrize("as_factory", [False, True], ids=["fixed", "callable"]) def test_raw_surrogate_kernel_without_overrides(as_factory): - """Without overrides, raw surrogate kernels are passed through unchanged.""" + """Preserve raw kernels and the factory's original searchspace without overrides.""" raw = gk.RBFKernel(active_dims=(0,)) - specification = (lambda s, o, m: raw) if as_factory else raw - kernel, _ = _resolve([NumericalContinuousParameter("x", (0, 1))], specification) - assert kernel is raw - - -def test_all_overridden_bypasses_surrogate_factory(): - """A fully overridden space does not construct an unused residual kernel.""" + searchspace = SearchSpace.from_product([NumericalContinuousParameter("x", (0, 1))]) - def unused_factory(searchspace, objective, measurements): - pytest.fail("The residual factory must not be called.") + def factory(received_searchspace, objective, measurements): + assert received_searchspace is searchspace + return raw - kernel, _ = _resolve( - [NumericalContinuousParameter("x", (0, 1), kernel_override=RBFKernel())], - unused_factory, + surrogate = GaussianProcessSurrogate( + kernel_or_factory=factory if as_factory else raw + ) + context = _ModelContext( + searchspace, NumericalTarget("y").to_objective(), pd.DataFrame() ) - assert isinstance(kernel, gk.RBFKernel) + assert surrogate._resolve_kernel(context) is raw From 3d5b2560709409962add6f70b14e92a4a7939bc2 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Wed, 9 Sep 2026 14:39:19 +0200 Subject: [PATCH 13/73] Align kernel scoping docstring with parameter removal --- baybe/kernels/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/baybe/kernels/base.py b/baybe/kernels/base.py index 52730bcdef..7721f8a819 100644 --- a/baybe/kernels/base.py +++ b/baybe/kernels/base.py @@ -122,7 +122,7 @@ def _without_parameter( ) def _with_parameter(self, name: str, /) -> Kernel: - """Return a copy of the kernel scoped to a single parameter. + """Return a copy of the kernel that acts only on the given parameter. Args: name: The name of the parameter to scope the kernel to. From 429f4d6ec375b4c545289322892b36730a7dcb64 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Wed, 9 Sep 2026 14:39:33 +0200 Subject: [PATCH 14/73] Explain the lazy GPyTorch import guard --- baybe/parameters/base.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/baybe/parameters/base.py b/baybe/parameters/base.py index d4e97a86ce..0cb3b9d725 100644 --- a/baybe/parameters/base.py +++ b/baybe/parameters/base.py @@ -89,6 +89,8 @@ def _to_kernel_override( return value._with_parameter(instance.name) # GPyTorch kernels: no explicit active dimensions allowed anywhere in the tree. + # An existing GPyTorch instance implies the module is already imported. Avoid + # importing it (and Torch) solely to validate other parameter inputs. if sys.modules.get("gpytorch") is not None: from gpytorch.kernels import Kernel as GPyTorchKernel From 27470649822daec21518ef9b9cf3cff7ffbc1398 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Wed, 9 Sep 2026 14:40:00 +0200 Subject: [PATCH 15/73] Clarify empty kernel reduction semantics --- baybe/surrogates/gaussian_process/_override/core.py | 1 + 1 file changed, 1 insertion(+) diff --git a/baybe/surrogates/gaussian_process/_override/core.py b/baybe/surrogates/gaussian_process/_override/core.py index 33cb651520..31aff9204a 100644 --- a/baybe/surrogates/gaussian_process/_override/core.py +++ b/baybe/surrogates/gaussian_process/_override/core.py @@ -59,6 +59,7 @@ def reduce_kernel_spec( """ if not isinstance(component, Kernel): raise_incompatible_override(excluded_names, factory) + # Reduction can exhaust the scope and return None despite a non-None input. spec: Kernel | None = component for name in excluded_names: if spec is None: From 389c7fe6c008fbd9dbd4c56fc10e2753594393e8 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Wed, 9 Sep 2026 14:41:23 +0200 Subject: [PATCH 16/73] Annotate the basic kernel iterator return type --- baybe/parameters/base.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/baybe/parameters/base.py b/baybe/parameters/base.py index 0cb3b9d725..9322b2a00d 100644 --- a/baybe/parameters/base.py +++ b/baybe/parameters/base.py @@ -24,8 +24,11 @@ from baybe.utils.metadata import MeasurableMetadata, to_metadata if TYPE_CHECKING: + from collections.abc import Iterator + from gpytorch.kernels import Kernel as GPyTorchKernel + from baybe.kernels.base import BasicKernel from baybe.parameters.enum import _ParameterKind from baybe.searchspace.continuous import SubspaceContinuous from baybe.searchspace.core import SearchSpace @@ -39,8 +42,15 @@ # https://github.com/python-attrs/attrs/issues/164 -def _iter_basic_kernels(kernel: Kernel): - """Iterate over the basic kernel leaves of a BayBE kernel.""" +def _iter_basic_kernels(kernel: Kernel) -> Iterator[BasicKernel]: + """Iterate over the basic kernel leaves of a BayBE kernel. + + Args: + kernel: The kernel to traverse. + + Yields: + The basic kernel leaves. + """ from baybe.kernels.base import BasicKernel from baybe.kernels.composite import AdditiveKernel, ProductKernel, ScaleKernel From e8e0cf05b0cce6a605bc0d6d446955938e989004 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Wed, 9 Sep 2026 14:41:43 +0200 Subject: [PATCH 17/73] Reject unsupported kernel structures during leaf traversal --- baybe/parameters/base.py | 6 ++++++ tests/validation/test_parameter_validation.py | 11 +++++++++++ 2 files changed, 17 insertions(+) diff --git a/baybe/parameters/base.py b/baybe/parameters/base.py index 9322b2a00d..ef42ef0260 100644 --- a/baybe/parameters/base.py +++ b/baybe/parameters/base.py @@ -50,6 +50,9 @@ def _iter_basic_kernels(kernel: Kernel) -> Iterator[BasicKernel]: Yields: The basic kernel leaves. + + Raises: + TypeError: If the kernel structure is unsupported. """ from baybe.kernels.base import BasicKernel from baybe.kernels.composite import AdditiveKernel, ProductKernel, ScaleKernel @@ -62,6 +65,9 @@ def _iter_basic_kernels(kernel: Kernel) -> Iterator[BasicKernel]: for sub in kernel.base_kernels: yield from _iter_basic_kernels(sub) + else: + raise TypeError(f"Cannot traverse kernel '{type(kernel).__name__}'.") + def _to_kernel_override( value: KernelOverride | None, instance: Parameter diff --git a/tests/validation/test_parameter_validation.py b/tests/validation/test_parameter_validation.py index 35064e97f2..60d2286f7e 100644 --- a/tests/validation/test_parameter_validation.py +++ b/tests/validation/test_parameter_validation.py @@ -13,6 +13,7 @@ from baybe._optional.info import CHEM_INSTALLED from baybe.kernels import RBFKernel +from baybe.kernels.base import Kernel from baybe.kernels.composite import ScaleKernel from baybe.parameters.categorical import ( CategoricalParameter, @@ -94,6 +95,16 @@ def test_invalid_kernel_override(constructor, override, error, match): constructor(kernel_override=override) +@pytest.mark.parametrize("scaled", [False, True], ids=["direct", "nested"]) +def test_unsupported_kernel_override_structure(scaled): + """Unknown kernel structures cannot silently bypass owner-scope validation.""" + kernel = Mock(spec=Kernel) + with pytest.raises(TypeError, match="Cannot traverse kernel"): + NumericalContinuousParameter( + "x", (0, 1), kernel_override=ScaleKernel(kernel) if scaled else kernel + ) + + @pytest.mark.parametrize( ("decorrelation", "error"), [ From 8562c1272ba80a6d33a102dab1b9a5054a3b8f8f Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Wed, 9 Sep 2026 14:42:01 +0200 Subject: [PATCH 18/73] Document the parameter override helper precondition --- baybe/surrogates/gaussian_process/_override/parameter.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/baybe/surrogates/gaussian_process/_override/parameter.py b/baybe/surrogates/gaussian_process/_override/parameter.py index a15c4bedcf..23286cc0ce 100644 --- a/baybe/surrogates/gaussian_process/_override/parameter.py +++ b/baybe/surrogates/gaussian_process/_override/parameter.py @@ -40,7 +40,8 @@ def make_parameter_override_kernel( """Create the kernel factor for a parameter's override. Args: - parameter: The parameter carrying the override. + parameter: The parameter carrying a non-``None`` kernel override, as ensured + by :func:`extract_parameter_overrides`. searchspace: The search space the kernel operates on. Returns: From 8d1c6d00c48be1b9f5f0aca91e705ee9ac61d3c9 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Wed, 9 Sep 2026 14:42:51 +0200 Subject: [PATCH 19/73] Support parameter removal from product kernels Reduce each factor recursively, drop exhausted factors, and collapse empty or single-factor products. Additive reduction remains unsupported, including additive factors nested inside products. --- baybe/kernels/composite.py | 15 ++++++++++++ tests/test_kernel_factories.py | 16 +++++++------ tests/test_kernels.py | 42 ++++++++++++++++++++++++++++++++-- 3 files changed, 64 insertions(+), 9 deletions(-) diff --git a/baybe/kernels/composite.py b/baybe/kernels/composite.py index a50c9aed62..ce16226260 100644 --- a/baybe/kernels/composite.py +++ b/baybe/kernels/composite.py @@ -109,6 +109,21 @@ class ProductKernel(CompositeKernel): ) """The individual kernels to be multiplied.""" + @override + def _without_parameter( + self, name: str, searchspace: SearchSpace, / + ) -> Kernel | None: + remaining = tuple( + reduced + for kernel in self.base_kernels + if (reduced := kernel._without_parameter(name, searchspace)) is not None + ) + if not remaining: + return None + if len(remaining) == 1: + return remaining[0] + return evolve(self, base_kernels=remaining) + @override def _with_parameter(self, name: str, /) -> Kernel: return evolve( diff --git a/tests/test_kernel_factories.py b/tests/test_kernel_factories.py index a8c12b7df5..603f0c9557 100644 --- a/tests/test_kernel_factories.py +++ b/tests/test_kernel_factories.py @@ -146,6 +146,14 @@ def _make_dispatch_context(override_mode): @pytest.mark.parametrize( ("override_mode", "kernel_or_factory", "expected_task_kernel_cls", "has_base"), [ + param( + TransferLearningMode.POSITIVE_INDEX_KERNEL, + MaternKernel(parameter_names=("x",)) + * IndexKernel(num_tasks=3, rank=3, parameter_names=("Task",)), + GPyTorchPositiveIndexKernel, + True, + id="override+product_kernel", + ), param( None, None, @@ -305,12 +313,6 @@ def test_resolve_kernel_dispatch_success( gpytorch.kernels.MaternKernel(nu=2.5), id="index_override+bare_gpytorch_matern", ), - param( - TransferLearningMode.POSITIVE_INDEX_KERNEL, - MaternKernel(parameter_names=("x",)) - * IndexKernel(num_tasks=3, rank=3, parameter_names=("Task",)), - id="override+product_kernel", - ), param( TransferLearningMode.POSITIVE_INDEX_KERNEL, ICMKernelFactory( @@ -330,7 +332,7 @@ def test_resolve_kernel_dispatch_success( def test_resolve_kernel_dispatch_raises(monkeypatch, override_mode, kernel_or_factory): """`_resolve_kernel` raises for inputs incompatible with an override. - This covers raw gpytorch kernels, composite kernels, and task-aware factories. + This covers raw gpytorch kernels and task-aware factories. """ monkeypatch.setenv("BAYBE_DISABLE_CUSTOM_KERNEL_WARNING", "True") diff --git a/tests/test_kernels.py b/tests/test_kernels.py index b8d54b7b17..1de2b306b7 100644 --- a/tests/test_kernels.py +++ b/tests/test_kernels.py @@ -273,18 +273,56 @@ def fixture_task_searchspace() -> SearchSpace: None, id="scale_kernel_over_sole_parameter_collapses_to_none", ), + param( + MaternKernel() * RBFKernel(), + MaternKernel(parameter_names=("x",)) * RBFKernel(parameter_names=("x",)), + id="product_reduces_each_factor", + ), + param( + MaternKernel() * RBFKernel(parameter_names=("Task",)), + MaternKernel(parameter_names=("x",)), + id="product_collapses_to_single_factor", + ), + param( + MaternKernel(parameter_names=("Task",)) + * RBFKernel(parameter_names=("Task",)), + None, + id="product_collapses_to_none", + ), + param( + ScaleKernel( + ProductKernel( + [ + MaternKernel(), + ProductKernel( + [ + RBFKernel(parameter_names=("Task",)), + RBFKernel(parameter_names=("x",)), + ] + ), + ] + ) + ), + ScaleKernel( + MaternKernel(parameter_names=("x",)) * RBFKernel(parameter_names=("x",)) + ), + id="nested_product_under_scale", + ), ], ) def test_without_parameter(kernel, expected, task_searchspace): - """Removing a parameter reduces basic and scaled kernels as expected.""" + """Removing a parameter reduces basic, scaled, and product kernels as expected.""" assert kernel._without_parameter("Task", task_searchspace) == expected @pytest.mark.parametrize( "kernel", [ - param(MaternKernel() * RBFKernel(), id="product_kernel"), param(MaternKernel() + RBFKernel(), id="additive_kernel"), + param( + MaternKernel() * (MaternKernel() + RBFKernel()), + id="product_with_unsupported_factor", + ), ], ) def test_without_parameter_unsupported(kernel, task_searchspace): From bc2c5cb50f3ce957b33e45eacf336ee43d0d4f13 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Wed, 9 Sep 2026 14:43:07 +0200 Subject: [PATCH 20/73] Clarify the custom kernel deprecation flag --- baybe/surrogates/gaussian_process/core.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py index ab09de793e..def326d86f 100644 --- a/baybe/surrogates/gaussian_process/core.py +++ b/baybe/surrogates/gaussian_process/core.py @@ -170,7 +170,12 @@ class GaussianProcessSurrogate(Surrogate): # See base class. _custom_kernel: bool = field(init=False, default=False, repr=False, eq=False) - # For deprecation only! + """Legacy flag for custom surrogate kernels affected by task-attachment changes. + + Parameter overrides do not set this flag: default residual construction still + attaches the default task kernel unless an explicit TL override replaces it. + This flag does not classify the final composed covariance. + """ kernel_factory: KernelFactoryProtocol | None = field( alias="kernel_or_factory", From 1dceb3fcda98c4d7a6a1812937110b1fb17fa2fd Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Wed, 9 Sep 2026 14:43:30 +0200 Subject: [PATCH 21/73] Test nondefault residual indices around a parameter override --- tests/test_parameter_kernel_overrides.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tests/test_parameter_kernel_overrides.py b/tests/test_parameter_kernel_overrides.py index f3ec7be1ec..35de8e91de 100644 --- a/tests/test_parameter_kernel_overrides.py +++ b/tests/test_parameter_kernel_overrides.py @@ -81,6 +81,26 @@ def test_default_factory_selector_is_preserved(mode): } +@pytest.mark.parametrize( + ("kernel_or_factory", "expected_residual_dims"), + [param(MaternKernel(), (0, 2), id="unnamed")], +) +def test_nondefault_residual_indices(kernel_or_factory, expected_residual_dims): + """Removing a middle parameter preserves the original computational indices.""" + kernel, _ = _resolve( + [ + NumericalContinuousParameter("x1", (0, 1)), + NumericalContinuousParameter("x2", (0, 1), kernel_override=RBFKernel()), + NumericalContinuousParameter("x3", (0, 1)), + ], + kernel_or_factory, + ) + residual, override = kernel.kernels + assert tuple(residual.active_dims.tolist()) == expected_residual_dims + assert tuple(override.active_dims.tolist()) == (1,) + assert isinstance(override, gk.RBFKernel) + + @pytest.mark.parametrize( ("override_parameter", "base_override", "task_parameter"), [ From b15768473eefaba2f7c1ea1b254fa1b7ffc889b4 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Wed, 9 Sep 2026 14:43:53 +0200 Subject: [PATCH 22/73] Extend reducible kernel override success cases --- tests/test_parameter_kernel_overrides.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/test_parameter_kernel_overrides.py b/tests/test_parameter_kernel_overrides.py index 35de8e91de..e8e2b253b4 100644 --- a/tests/test_parameter_kernel_overrides.py +++ b/tests/test_parameter_kernel_overrides.py @@ -83,7 +83,16 @@ def test_default_factory_selector_is_preserved(mode): @pytest.mark.parametrize( ("kernel_or_factory", "expected_residual_dims"), - [param(MaternKernel(), (0, 2), id="unnamed")], + [ + param(MaternKernel(), (0, 2), id="unnamed"), + param(MaternKernel(parameter_names=("x1", "x2")), (0,), id="named"), + param(ScaleKernel(MaternKernel()), (0, 2), id="scaled"), + param( + lambda s, o, m: MaternKernel(parameter_names=s.parameter_names), + (0, 2), + id="callable", + ), + ], ) def test_nondefault_residual_indices(kernel_or_factory, expected_residual_dims): """Removing a middle parameter preserves the original computational indices.""" From 6e424f95c432b74f84bd8bbe53639a89da6b4e31 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Wed, 9 Sep 2026 14:44:24 +0200 Subject: [PATCH 23/73] Document surrogate kernel override behavior --- baybe/surrogates/gaussian_process/core.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py index def326d86f..49a63ff4e4 100644 --- a/baybe/surrogates/gaussian_process/core.py +++ b/baybe/surrogates/gaussian_process/core.py @@ -195,9 +195,16 @@ class GaussianProcessSurrogate(Surrogate): * :obj:`.components.kernel.KernelFactoryProtocol` * :class:`gpytorch.kernels.Kernel` - If a :class:`.TaskParameter` sets ``override_transfer_learning_mode``, this must - reduce to a task-free BayBE kernel or an :class:`.IncompatibleOverrideError` is - raised. + A :attr:`~baybe.parameters.base.Parameter.kernel_override` removes its parameter + from this kernel and contributes a separate multiplicative factor. A + :attr:`~baybe.parameters.categorical.TaskParameter.override_transfer_learning_mode` + replaces the task factor in the same way. When a residual kernel is needed, + the configured kernel or factory must support excluding the overridden parameters; + otherwise, :class:`~baybe.exceptions.IncompatibleOverrideError` is raised. + If all parameters are overridden, this kernel or factory is not used. Without + overrides, it is used unchanged. + + See :ref:`parameter_kernel_overrides` for details and limitations. """ mean_factory: MeanFactoryProtocol | None = field( From 891b3b296fae59d02648258587260802b945c38f Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Wed, 9 Sep 2026 16:13:49 +0200 Subject: [PATCH 24/73] Fix documentation reference for KernelOverride alias Exclude the exact alias name from class-reference checks, matching existing type-alias exceptions. Attrs-generated parameter constructors expose KernelOverride without a linkable class target. --- docs/conf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/conf.py b/docs/conf.py index a381683946..cda98df971 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -128,7 +128,7 @@ (r"py:class", r"^(dtype=np\.int64|if metric == 'precomputed')$"), ##### Type aliases in TYPE_CHECKING blocks ##### # These exist only at type-checking time and cannot be resolved by Sphinx. - (r"py:class", r"^(GPComponent|TensorCallable|ConvertibleToFloat)$"), + (r"py:class", r"^(GPComponent|KernelOverride|TensorCallable|ConvertibleToFloat)$"), (r"py:class", r"^(GPyTorchKernel|GPyTorchLikelihood|GPyTorchMean|GPyTorchModel)$"), (r"py:class", r"^(pd\.DataFrame|pl\.Expr)$"), (r"py:class", r"^(TypeAliasForwardRef|P)$"), From 26ac43625fa52285434a10292e5cc55d962e3e6b Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Tue, 22 Sep 2026 16:53:39 +0200 Subject: [PATCH 25/73] Remove comment Co-authored-by: AdrianSosic --- tests/test_parameter_kernel_overrides.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_parameter_kernel_overrides.py b/tests/test_parameter_kernel_overrides.py index e8e2b253b4..4241f577f7 100644 --- a/tests/test_parameter_kernel_overrides.py +++ b/tests/test_parameter_kernel_overrides.py @@ -1,7 +1,6 @@ """Functional tests for parameter-specific kernel overrides. These tests resolve or fit kernels to check binding and composition behavior. -Construction-time input validation lives in ``tests/validation``. """ from copy import deepcopy From 8f749b45c4b6699d395f3a40ab954f19fea787c6 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Tue, 22 Sep 2026 16:54:17 +0200 Subject: [PATCH 26/73] Improve variable names Co-authored-by: AdrianSosic --- tests/test_parameter_kernel_overrides.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_parameter_kernel_overrides.py b/tests/test_parameter_kernel_overrides.py index 4241f577f7..22f2330395 100644 --- a/tests/test_parameter_kernel_overrides.py +++ b/tests/test_parameter_kernel_overrides.py @@ -87,7 +87,7 @@ def test_default_factory_selector_is_preserved(mode): param(MaternKernel(parameter_names=("x1", "x2")), (0,), id="named"), param(ScaleKernel(MaternKernel()), (0, 2), id="scaled"), param( - lambda s, o, m: MaternKernel(parameter_names=s.parameter_names), + lambda space, _, __: MaternKernel(parameter_names=space.parameter_names), (0, 2), id="callable", ), From 33ee99e75e857b5ee147bc141d92f12d09793af3 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Tue, 22 Sep 2026 17:57:46 +0200 Subject: [PATCH 27/73] Clarify override condition in kernel factory docstring --- baybe/surrogates/gaussian_process/core.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py index 49a63ff4e4..7fd436b8ae 100644 --- a/baybe/surrogates/gaussian_process/core.py +++ b/baybe/surrogates/gaussian_process/core.py @@ -198,11 +198,12 @@ class GaussianProcessSurrogate(Surrogate): A :attr:`~baybe.parameters.base.Parameter.kernel_override` removes its parameter from this kernel and contributes a separate multiplicative factor. A :attr:`~baybe.parameters.categorical.TaskParameter.override_transfer_learning_mode` - replaces the task factor in the same way. When a residual kernel is needed, - the configured kernel or factory must support excluding the overridden parameters; - otherwise, :class:`~baybe.exceptions.IncompatibleOverrideError` is raised. - If all parameters are overridden, this kernel or factory is not used. Without - overrides, it is used unchanged. + replaces the task factor in the same way. When not all parameters are removed by + an override, the configured kernel or factory must support excluding the + overridden parameters; otherwise, + :class:`~baybe.exceptions.IncompatibleOverrideError` is raised. If all parameters + are overridden, this kernel or factory is not used. Without overrides, it is used + unchanged. See :ref:`parameter_kernel_overrides` for details and limitations. """ From 6789cfe95a1f60dce4da704af5e84fbddf7202ba Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Tue, 22 Sep 2026 17:58:40 +0200 Subject: [PATCH 28/73] Rename kernel scoping method to _scope_to_parameter --- baybe/kernels/base.py | 4 ++-- baybe/kernels/composite.py | 12 ++++++------ baybe/parameters/base.py | 4 ++-- 3 files changed, 10 insertions(+), 10 deletions(-) diff --git a/baybe/kernels/base.py b/baybe/kernels/base.py index 7721f8a819..7ee7e44282 100644 --- a/baybe/kernels/base.py +++ b/baybe/kernels/base.py @@ -121,7 +121,7 @@ def _without_parameter( f"Cannot remove a parameter from kernel '{self.__class__.__name__}'. " ) - def _with_parameter(self, name: str, /) -> Kernel: + def _scope_to_parameter(self, name: str, /) -> Kernel: """Return a copy of the kernel that acts only on the given parameter. Args: @@ -290,7 +290,7 @@ def _without_parameter( return evolve(self, parameter_names=remaining) if remaining else None @override - def _with_parameter(self, name: str, /) -> Kernel: + def _scope_to_parameter(self, name: str, /) -> Kernel: return evolve(self, parameter_names=(name,)) diff --git a/baybe/kernels/composite.py b/baybe/kernels/composite.py index ce16226260..f61fbcf71e 100644 --- a/baybe/kernels/composite.py +++ b/baybe/kernels/composite.py @@ -56,8 +56,8 @@ def _without_parameter( return None if stripped is None else evolve(self, base_kernel=stripped) @override - def _with_parameter(self, name: str, /) -> Kernel: - return evolve(self, base_kernel=self.base_kernel._with_parameter(name)) + def _scope_to_parameter(self, name: str, /) -> Kernel: + return evolve(self, base_kernel=self.base_kernel._scope_to_parameter(name)) @override def to_gpytorch(self, *args, **kwargs): @@ -86,10 +86,10 @@ class AdditiveKernel(CompositeKernel): """The individual kernels to be summed.""" @override - def _with_parameter(self, name: str, /) -> Kernel: + def _scope_to_parameter(self, name: str, /) -> Kernel: return evolve( self, - base_kernels=tuple(k._with_parameter(name) for k in self.base_kernels), + base_kernels=tuple(k._scope_to_parameter(name) for k in self.base_kernels), ) @override @@ -125,10 +125,10 @@ def _without_parameter( return evolve(self, base_kernels=remaining) @override - def _with_parameter(self, name: str, /) -> Kernel: + def _scope_to_parameter(self, name: str, /) -> Kernel: return evolve( self, - base_kernels=tuple(k._with_parameter(name) for k in self.base_kernels), + base_kernels=tuple(k._scope_to_parameter(name) for k in self.base_kernels), ) @override diff --git a/baybe/parameters/base.py b/baybe/parameters/base.py index ef42ef0260..b62ff77f4b 100644 --- a/baybe/parameters/base.py +++ b/baybe/parameters/base.py @@ -102,7 +102,7 @@ def _to_kernel_override( f"itself. Its basic kernels must specify 'parameter_names' as " f"``None`` or ({instance.name!r},)." ) - return value._with_parameter(instance.name) + return value._scope_to_parameter(instance.name) # GPyTorch kernels: no explicit active dimensions allowed anywhere in the tree. # An existing GPyTorch instance implies the module is already imported. Avoid @@ -219,7 +219,7 @@ def is_equivalent(self, other: Parameter) -> bool: # The override is owner-scoped, so rebind it to the other parameter's name. kernel_override = self.kernel_override if isinstance(kernel_override, Kernel): - kernel_override = kernel_override._with_parameter(other.name) + kernel_override = kernel_override._scope_to_parameter(other.name) return ( attrs.evolve(self, name=other.name, kernel_override=kernel_override) == other From a67f896d024a3a927af0190c36a0b080836f95c7 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Tue, 22 Sep 2026 17:59:33 +0200 Subject: [PATCH 29/73] Absorb None handling via converters.optional --- baybe/parameters/base.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/baybe/parameters/base.py b/baybe/parameters/base.py index b62ff77f4b..709c354471 100644 --- a/baybe/parameters/base.py +++ b/baybe/parameters/base.py @@ -69,9 +69,7 @@ def _iter_basic_kernels(kernel: Kernel) -> Iterator[BasicKernel]: raise TypeError(f"Cannot traverse kernel '{type(kernel).__name__}'.") -def _to_kernel_override( - value: KernelOverride | None, instance: Parameter -) -> KernelOverride | None: +def _to_kernel_override(value: KernelOverride, instance: Parameter) -> KernelOverride: """Validate a kernel override and scope BayBE kernels to their parameter. Args: @@ -86,9 +84,6 @@ def _to_kernel_override( Returns: The validated override, with BayBE kernels scoped to the parameter. """ - if value is None: - return None - # BayBE kernels: every basic leaf must be unscoped or scoped to the owner. The # kernel is then rebound to the owning parameter (dropping unspecified names). if isinstance(value, Kernel): @@ -147,7 +142,7 @@ class Parameter(ABC, SerialMixin): kernel_override: KernelOverride | None = field( default=None, - converter=Converter(_to_kernel_override, takes_self=True), # type: ignore[misc, call-overload] + converter=optional_c(Converter(_to_kernel_override, takes_self=True)), # type: ignore[misc, call-overload] kw_only=True, ) """An optional kernel replacing the overall kernel for this parameter.""" From 4ca445c6bbe305add2f0c33e075dbb0ab1914a1d Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Tue, 22 Sep 2026 17:59:59 +0200 Subject: [PATCH 30/73] Use field alias in kernel override error message --- baybe/parameters/base.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/baybe/parameters/base.py b/baybe/parameters/base.py index 709c354471..d77a0af06e 100644 --- a/baybe/parameters/base.py +++ b/baybe/parameters/base.py @@ -87,6 +87,9 @@ def _to_kernel_override(value: KernelOverride, instance: Parameter) -> KernelOve # BayBE kernels: every basic leaf must be unscoped or scoped to the owner. The # kernel is then rebound to the owning parameter (dropping unspecified names). if isinstance(value, Kernel): + from baybe.kernels.base import BasicKernel + + names_alias = attrs.fields(BasicKernel).parameter_names.alias if any( leaf.parameter_names not in (None, (instance.name,)) for leaf in _iter_basic_kernels(value) @@ -94,7 +97,7 @@ def _to_kernel_override(value: KernelOverride, instance: Parameter) -> KernelOve raise ValueError( f"The kernel provided for the kernel override of " f"'{instance.__class__.__name__}' may only act on the parameter " - f"itself. Its basic kernels must specify 'parameter_names' as " + f"itself. Its basic kernels must specify '{names_alias}' as " f"``None`` or ({instance.name!r},)." ) return value._scope_to_parameter(instance.name) From 898a4432b45cb51f0dbba935f52bad14ca5e6f77 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Tue, 22 Sep 2026 18:00:10 +0200 Subject: [PATCH 31/73] Align kernel override flag docstring with sibling flags --- baybe/surrogates/base.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py index df4ce48662..bb1d1cd1d0 100644 --- a/baybe/surrogates/base.py +++ b/baybe/surrogates/base.py @@ -88,7 +88,8 @@ class Surrogate(ABC, SurrogateProtocol, SerialMixin): """Abstract base class for all surrogate models.""" supports_kernel_overrides: ClassVar[bool] = False - """Whether parameter-specific kernel overrides are supported.""" + """Class variable encoding whether or not the surrogate supports + parameter-specific kernel overrides.""" supports_transfer_learning: ClassVar[bool] """Class variable encoding whether or not the surrogate supports transfer From 7f7e2e616b59f8841d953ca5455cf316d9f845a2 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Tue, 22 Sep 2026 18:00:22 +0200 Subject: [PATCH 32/73] Remove self-referential surrogate kernel definition --- docs/components/parameters.md | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/docs/components/parameters.md b/docs/components/parameters.md index 36545de29e..7f1bd5e31a 100644 --- a/docs/components/parameters.md +++ b/docs/components/parameters.md @@ -302,9 +302,11 @@ x3 = NumericalContinuousParameter( ) ``` -The surrogate kernel is the kernel supplied to the surrogate, or BayBE's default kernel -if none is supplied. Overrides remove their parameter dimensions from that kernel and -replace them with multiplicative parameter-specific factors: +Here, "surrogate kernel" refers to the ``kernel_or_factory`` of the +{class}`~baybe.surrogates.gaussian_process.core.GaussianProcessSurrogate`, falling back +to BayBE's default kernel when none is given. Overrides remove their parameter +dimensions from that kernel and replace them with multiplicative parameter-specific +factors: ```{math} k_{\mathrm{effective}} From 72c2a5513bddd51bc36b5515fa7a37bc5da4628b Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Tue, 22 Sep 2026 18:00:33 +0200 Subject: [PATCH 33/73] Link kernel_override attribute on first mention --- docs/components/parameters.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/components/parameters.md b/docs/components/parameters.md index 7f1bd5e31a..b2eca64289 100644 --- a/docs/components/parameters.md +++ b/docs/components/parameters.md @@ -289,7 +289,7 @@ For details, refer to [transfer learning](../concepts/transfer_learning.md). ## Parameter-Specific Kernel Overrides Regular parameters can replace the surrogate kernel on their computational dimensions -using ``kernel_override``: +using {attr}`~baybe.parameters.base.Parameter.kernel_override`: ```python from baybe.kernels import RBFKernel From 42b3934fc5cead0a727e3d744f43c20773a5f025 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Tue, 22 Sep 2026 18:02:09 +0200 Subject: [PATCH 34/73] Lead kernel override section with the mechanism --- docs/components/parameters.md | 29 ++++++++++++++++------------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/docs/components/parameters.md b/docs/components/parameters.md index b2eca64289..5dc9089e5f 100644 --- a/docs/components/parameters.md +++ b/docs/components/parameters.md @@ -288,6 +288,21 @@ For details, refer to [transfer learning](../concepts/transfer_learning.md). (parameter_kernel_overrides)= ## Parameter-Specific Kernel Overrides +By default, the Gaussian process models all parameters with a single surrogate kernel. +A parameter can instead be assigned its own kernel, which then replaces the surrogate +kernel on that parameter's computational dimensions. The remaining dimensions stay with +the surrogate kernel, and the two are combined multiplicatively: + +```{math} +k_{\mathrm{effective}} += k_{\mathrm{surrogate}}(X_{\mathrm{remaining}}) +\prod_{p \in P_{\mathrm{override}}} k_p(X_p). +``` + +Here, "surrogate kernel" refers to the ``kernel_or_factory`` of the +{class}`~baybe.surrogates.gaussian_process.core.GaussianProcessSurrogate`, falling back +to BayBE's default kernel when none is given. + Regular parameters can replace the surrogate kernel on their computational dimensions using {attr}`~baybe.parameters.base.Parameter.kernel_override`: @@ -302,19 +317,7 @@ x3 = NumericalContinuousParameter( ) ``` -Here, "surrogate kernel" refers to the ``kernel_or_factory`` of the -{class}`~baybe.surrogates.gaussian_process.core.GaussianProcessSurrogate`, falling back -to BayBE's default kernel when none is given. Overrides remove their parameter -dimensions from that kernel and replace them with multiplicative parameter-specific -factors: - -```{math} -k_{\mathrm{effective}} -= k_{\mathrm{surrogate}}(X_{\mathrm{remaining}}) -\prod_{p \in P_{\mathrm{override}}} k_p(X_p). -``` - -For parameters ``x1``, ``x2``, and ``x3`` in the example above, this produces +In a search space with parameters ``x1``, ``x2``, and ``x3``, this produces ``surrogate kernel(x1, x2) * RBF(x3)``. The override does not replace a dimension inside the surrogate kernel itself. It partitions the dimensions and creates a separate kernel factor. Multiple overrides produce multiple factors. For a parameter with several From 1c87ebee4fbdade2ca7572b5aced9fc47e438897 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Tue, 22 Sep 2026 18:02:22 +0200 Subject: [PATCH 35/73] Clarify that overrides accept kernel instances --- docs/components/parameters.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/docs/components/parameters.md b/docs/components/parameters.md index 5dc9089e5f..41de4e6b9d 100644 --- a/docs/components/parameters.md +++ b/docs/components/parameters.md @@ -330,8 +330,9 @@ An override can be a configured BayBE or GPyTorch kernel instance. For a BayBE k ```{admonition} Limitations :class: warning -Parameter-local kernel factories are not supported. The surrogate kernel or factory of -the {class}`~baybe.surrogates.gaussian_process.core.GaussianProcessSurrogate` must allow +{attr}`~baybe.parameters.base.Parameter.kernel_override` accepts kernel instances only, +not factories. The surrogate kernel or factory of the +{class}`~baybe.surrogates.gaussian_process.core.GaussianProcessSurrogate` must allow the overridden dimensions to be excluded. Otherwise, BayBE raises {class}`~baybe.exceptions.IncompatibleOverrideError`. From c3bcbed4c5dd6e4dd3593eaa289680c84e77b4dd Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Tue, 22 Sep 2026 18:05:15 +0200 Subject: [PATCH 36/73] Support parameter removal from additive kernels A summand scoped exclusively to the removed parameter is dropped, which treats it as the additive identity. The alternative, retaining the degenerate constant, is not expressible since its value is a learned quantity and no constant kernel exists. Refusing the reduction does not preserve correctness either, so all composites now reduce uniformly. --- baybe/kernels/composite.py | 43 ++++++++++++++++++------ tests/test_kernels.py | 42 +++++++++++++---------- tests/test_parameter_kernel_overrides.py | 4 --- 3 files changed, 58 insertions(+), 31 deletions(-) diff --git a/baybe/kernels/composite.py b/baybe/kernels/composite.py index f61fbcf71e..5017abec21 100644 --- a/baybe/kernels/composite.py +++ b/baybe/kernels/composite.py @@ -23,6 +23,32 @@ from baybe.searchspace.core import SearchSpace +def _reduce_base_kernels( + kernel: AdditiveKernel | ProductKernel, name: str, searchspace: SearchSpace, / +) -> Kernel | None: + """Remove a parameter from all base kernels of a composite kernel. + + Args: + kernel: The composite kernel whose base kernels are to be reduced. + name: The name of the parameter to remove. + searchspace: The search space the kernel operates on. + + Returns: + The reduced composite kernel, the sole remaining base kernel, or ``None`` + if no base kernel remains. + """ + remaining = tuple( + reduced + for k in kernel.base_kernels + if (reduced := k._without_parameter(name, searchspace)) is not None + ) + if not remaining: + return None + if len(remaining) == 1: + return remaining[0] + return evolve(kernel, base_kernels=remaining) + + @define(frozen=True) class ScaleKernel(CompositeKernel): """A kernel for decorating existing kernels with an outputscale.""" @@ -85,6 +111,12 @@ class AdditiveKernel(CompositeKernel): ) """The individual kernels to be summed.""" + @override + def _without_parameter( + self, name: str, searchspace: SearchSpace, / + ) -> Kernel | None: + return _reduce_base_kernels(self, name, searchspace) + @override def _scope_to_parameter(self, name: str, /) -> Kernel: return evolve( @@ -113,16 +145,7 @@ class ProductKernel(CompositeKernel): def _without_parameter( self, name: str, searchspace: SearchSpace, / ) -> Kernel | None: - remaining = tuple( - reduced - for kernel in self.base_kernels - if (reduced := kernel._without_parameter(name, searchspace)) is not None - ) - if not remaining: - return None - if len(remaining) == 1: - return remaining[0] - return evolve(self, base_kernels=remaining) + return _reduce_base_kernels(self, name, searchspace) @override def _scope_to_parameter(self, name: str, /) -> Kernel: diff --git a/tests/test_kernels.py b/tests/test_kernels.py index 1de2b306b7..38d71cbce8 100644 --- a/tests/test_kernels.py +++ b/tests/test_kernels.py @@ -273,6 +273,30 @@ def fixture_task_searchspace() -> SearchSpace: None, id="scale_kernel_over_sole_parameter_collapses_to_none", ), + param( + MaternKernel() + RBFKernel(), + MaternKernel(parameter_names=("x",)) + RBFKernel(parameter_names=("x",)), + id="sum_reduces_each_summand", + ), + param( + MaternKernel() + RBFKernel(parameter_names=("Task",)), + MaternKernel(parameter_names=("x",)), + id="sum_collapses_to_single_summand", + ), + param( + MaternKernel(parameter_names=("Task",)) + + RBFKernel(parameter_names=("Task",)), + None, + id="sum_collapses_to_none", + ), + param( + MaternKernel() * (MaternKernel() + RBFKernel()), + MaternKernel(parameter_names=("x",)) + * ( + MaternKernel(parameter_names=("x",)) + RBFKernel(parameter_names=("x",)) + ), + id="nested_sum_under_product", + ), param( MaternKernel() * RBFKernel(), MaternKernel(parameter_names=("x",)) * RBFKernel(parameter_names=("x",)), @@ -311,21 +335,5 @@ def fixture_task_searchspace() -> SearchSpace: ], ) def test_without_parameter(kernel, expected, task_searchspace): - """Removing a parameter reduces basic, scaled, and product kernels as expected.""" + """Removing a parameter reduces basic, scaled, and composite kernels as expected.""" assert kernel._without_parameter("Task", task_searchspace) == expected - - -@pytest.mark.parametrize( - "kernel", - [ - param(MaternKernel() + RBFKernel(), id="additive_kernel"), - param( - MaternKernel() * (MaternKernel() + RBFKernel()), - id="product_with_unsupported_factor", - ), - ], -) -def test_without_parameter_unsupported(kernel, task_searchspace): - """Kernels that cannot be reduced unambiguously raise ``TypeError``.""" - with pytest.raises(TypeError, match="Cannot remove a parameter"): - kernel._without_parameter("Task", task_searchspace) diff --git a/tests/test_parameter_kernel_overrides.py b/tests/test_parameter_kernel_overrides.py index 22f2330395..f5c1bb5a64 100644 --- a/tests/test_parameter_kernel_overrides.py +++ b/tests/test_parameter_kernel_overrides.py @@ -241,10 +241,6 @@ def test_gpytorch_ard_mismatch_rejected(): @pytest.mark.parametrize( "kernel_or_factory", [ - param( - AdditiveKernel([MaternKernel(), MaternKernel()]), - id="non-reducible-baybe-kernel", - ), param(gk.MaternKernel(), id="raw-gpytorch-kernel"), param(lambda s, o, m: gk.MaternKernel(), id="factory-returning-raw-kernel"), ], From 29d2e1f8fa2d0e006ae2e09c93a1ff54984771e2 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Tue, 22 Sep 2026 18:06:29 +0200 Subject: [PATCH 37/73] Allow unscoping kernels via _scope_to_parameter Passing None removes all leaf scoping, which lets the Hypothesis strategy reuse the production traversal instead of duplicating it. --- baybe/kernels/base.py | 9 ++++---- baybe/kernels/composite.py | 6 +++--- tests/hypothesis_strategies/parameters.py | 25 +---------------------- 3 files changed, 9 insertions(+), 31 deletions(-) diff --git a/baybe/kernels/base.py b/baybe/kernels/base.py index 7ee7e44282..d88dad44df 100644 --- a/baybe/kernels/base.py +++ b/baybe/kernels/base.py @@ -121,11 +121,12 @@ def _without_parameter( f"Cannot remove a parameter from kernel '{self.__class__.__name__}'. " ) - def _scope_to_parameter(self, name: str, /) -> Kernel: + def _scope_to_parameter(self, name: str | None, /) -> Kernel: """Return a copy of the kernel that acts only on the given parameter. Args: - name: The name of the parameter to scope the kernel to. + name: The name of the parameter to scope the kernel to, or ``None`` to + remove all scoping so that the kernel acts on all dimensions. Raises: TypeError: If the kernel structure cannot be scoped unambiguously. @@ -290,8 +291,8 @@ def _without_parameter( return evolve(self, parameter_names=remaining) if remaining else None @override - def _scope_to_parameter(self, name: str, /) -> Kernel: - return evolve(self, parameter_names=(name,)) + def _scope_to_parameter(self, name: str | None, /) -> Kernel: + return evolve(self, parameter_names=None if name is None else (name,)) @define(frozen=True) diff --git a/baybe/kernels/composite.py b/baybe/kernels/composite.py index 5017abec21..7aeb7e93a1 100644 --- a/baybe/kernels/composite.py +++ b/baybe/kernels/composite.py @@ -82,7 +82,7 @@ def _without_parameter( return None if stripped is None else evolve(self, base_kernel=stripped) @override - def _scope_to_parameter(self, name: str, /) -> Kernel: + def _scope_to_parameter(self, name: str | None, /) -> Kernel: return evolve(self, base_kernel=self.base_kernel._scope_to_parameter(name)) @override @@ -118,7 +118,7 @@ def _without_parameter( return _reduce_base_kernels(self, name, searchspace) @override - def _scope_to_parameter(self, name: str, /) -> Kernel: + def _scope_to_parameter(self, name: str | None, /) -> Kernel: return evolve( self, base_kernels=tuple(k._scope_to_parameter(name) for k in self.base_kernels), @@ -148,7 +148,7 @@ def _without_parameter( return _reduce_base_kernels(self, name, searchspace) @override - def _scope_to_parameter(self, name: str, /) -> Kernel: + def _scope_to_parameter(self, name: str | None, /) -> Kernel: return evolve( self, base_kernels=tuple(k._scope_to_parameter(name) for k in self.base_kernels), diff --git a/tests/hypothesis_strategies/parameters.py b/tests/hypothesis_strategies/parameters.py index a2638d8fa6..13f1cecc90 100644 --- a/tests/hypothesis_strategies/parameters.py +++ b/tests/hypothesis_strategies/parameters.py @@ -4,7 +4,6 @@ import hypothesis.strategies as st import numpy as np -from attrs import evolve from hypothesis.extra.pandas import columns, data_frames from baybe.parameters.categorical import ( @@ -39,28 +38,6 @@ """A strategy that generates parameter categories.""" -def _remove_kernel_parameter_names(kernel): - """Remove explicit parameter names from all basic kernel leaves.""" - from baybe.kernels.base import BasicKernel - from baybe.kernels.composite import AdditiveKernel, ProductKernel, ScaleKernel - - if isinstance(kernel, BasicKernel): - return evolve(kernel, parameter_names=None) - if isinstance(kernel, ScaleKernel): - return evolve( - kernel, - base_kernel=_remove_kernel_parameter_names(kernel.base_kernel), - ) - if isinstance(kernel, (AdditiveKernel, ProductKernel)): - return evolve( - kernel, - base_kernels=tuple( - _remove_kernel_parameter_names(k) for k in kernel.base_kernels - ), - ) - raise TypeError(f"Unsupported kernel type: {type(kernel)}") - - @st.composite def kernel_overrides(draw: st.DrawFn, parameter_name: str): """Generate valid kernel overrides for a parameter.""" @@ -68,7 +45,7 @@ def kernel_overrides(draw: st.DrawFn, parameter_name: str): kernel = draw(st.one_of(st.none(), kernels(parameter_names=(parameter_name,)))) if kernel is not None and draw(st.booleans()): - kernel = _remove_kernel_parameter_names(kernel) + kernel = kernel._scope_to_parameter(None) return kernel From 755d5095d9ee0b40aaba06c3b80bf5666d125fe2 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Tue, 22 Sep 2026 18:06:54 +0200 Subject: [PATCH 38/73] Clarify selector override test --- tests/test_parameter_kernel_overrides.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/tests/test_parameter_kernel_overrides.py b/tests/test_parameter_kernel_overrides.py index f5c1bb5a64..bc458a314f 100644 --- a/tests/test_parameter_kernel_overrides.py +++ b/tests/test_parameter_kernel_overrides.py @@ -58,8 +58,12 @@ def _leaf_kernels(kernel): [None, *TransferLearningMode], ids=lambda mode: mode.name if mode else "no-task", ) -def test_default_factory_selector_is_preserved(mode): - """Selectors omit unselected dimensions but cannot exclude parameter overrides.""" +def test_selector_does_not_exclude_overridden_parameters(mode): + """A selector narrows the surrogate kernel, but overrides apply regardless. + + The factory selects only ``x1``, so ``omitted`` contributes no kernel factor. + ``x2`` is likewise unselected, yet its override still yields its own factor. + """ parameters = [ NumericalContinuousParameter("x1", (0, 1)), NumericalContinuousParameter("x2", (0, 1), kernel_override=RBFKernel()), @@ -71,13 +75,17 @@ def test_default_factory_selector_is_preserved(mode): TaskParameter("task", ["a", "b"], override_transfer_learning_mode=mode) ) expected_names.append("task") + kernel, searchspace = _resolve( parameters, BayBEKernelFactory(parameter_selector=NameSelector(("x1",), regex=False)), ) - assert {tuple(k.active_dims.tolist()) for k in _leaf_kernels(kernel)} == { + + actual_dimensions = {tuple(k.active_dims.tolist()) for k in _leaf_kernels(kernel)} + expected_dimensions = { searchspace.get_comp_rep_parameter_indices(name) for name in expected_names } + assert actual_dimensions == expected_dimensions @pytest.mark.parametrize( From 705012ee86b6ec82c52db8bd22a631645cce45fa Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Tue, 22 Sep 2026 18:23:00 +0200 Subject: [PATCH 39/73] Retarget pyrefly suppressions for stricter active_dims inference --- baybe/surrogates/gaussian_process/_override/core.py | 8 +++++--- baybe/surrogates/gaussian_process/_override/parameter.py | 4 +++- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/baybe/surrogates/gaussian_process/_override/core.py b/baybe/surrogates/gaussian_process/_override/core.py index 31aff9204a..bf7d883728 100644 --- a/baybe/surrogates/gaussian_process/_override/core.py +++ b/baybe/surrogates/gaussian_process/_override/core.py @@ -29,11 +29,13 @@ def get_active_dimensions(kernel: GPyTorchKernel, searchspace: SearchSpace) -> s from gpytorch.kernels import AdditiveKernel, ProductKernel, ScaleKernel if kernel.active_dims is not None: - # TODO[typing]: https://github.com/facebook/pyrefly/issues/3988 - return set(kernel.active_dims.tolist()) # pyrefly: ignore[not-callable] + # TODO[typing]: GPyTorch annotates `active_dims` with the constructor's tuple + # type, but `register_buffer` stores a tensor at runtime. + return set(kernel.active_dims.tolist()) # pyrefly: ignore[missing-attribute] if isinstance(kernel, (AdditiveKernel, ProductKernel)): return set().union( - *(get_active_dimensions(k, searchspace) for k in kernel.kernels) + # TODO[typing]: Iterating a `ModuleList` yields the `Module` base type. + *(get_active_dimensions(k, searchspace) for k in kernel.kernels) # pyrefly: ignore[bad-argument-type] ) if isinstance(kernel, ScaleKernel): return get_active_dimensions(kernel.base_kernel, searchspace) diff --git a/baybe/surrogates/gaussian_process/_override/parameter.py b/baybe/surrogates/gaussian_process/_override/parameter.py index 23286cc0ce..342ecbf0b2 100644 --- a/baybe/surrogates/gaussian_process/_override/parameter.py +++ b/baybe/surrogates/gaussian_process/_override/parameter.py @@ -93,5 +93,7 @@ def bind_gpytorch_override( ) result = deepcopy(override) - result.active_dims = torch.tensor(indices) + # TODO[typing]: GPyTorch annotates `active_dims` with the constructor's tuple + # type, but `register_buffer` stores a tensor at runtime. + result.active_dims = torch.tensor(indices) # pyrefly: ignore[bad-assignment] return result From 26fe66e6bf9d0b7d24405338719810843cff6d78 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Fri, 21 Aug 2026 17:50:42 +0200 Subject: [PATCH 40/73] Update CHANGELOG --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index a5b91f6979..b143763a6c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -48,6 +48,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `DiscreteLinearConstraint` for (optionally weighted) sum constraints on discrete parameters, supporting `coefficients` and mirroring `ContinuousLinearConstraint`'s `operator`/`rhs`/`coefficients` interface (replaces `DiscreteSumConstraint`) +- Parameter-specific kernel overrides for conveniently composing Gaussian process + kernels on individual parameter dimensions ### Changed - `BOTORCH` GP preset now includes `BetaPrior(2.5, 1.5)` for the task covariance From 4d3690cde15a0cf4f37d585a9b01894da919f645 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Thu, 24 Sep 2026 13:59:33 +0200 Subject: [PATCH 41/73] Compare GPyTorch kernel overrides structurally --- baybe/parameters/base.py | 62 +++++++++++++++++++++++- tests/test_parameter_kernel_overrides.py | 53 ++++++++++++++++++++ 2 files changed, 114 insertions(+), 1 deletion(-) diff --git a/baybe/parameters/base.py b/baybe/parameters/base.py index d77a0af06e..7687f13ccd 100644 --- a/baybe/parameters/base.py +++ b/baybe/parameters/base.py @@ -6,6 +6,7 @@ import sys from abc import ABC, abstractmethod from functools import cached_property +from itertools import chain from typing import TYPE_CHECKING, Any, ClassVar, TypeAlias import attrs @@ -69,6 +70,52 @@ def _iter_basic_kernels(kernel: Kernel) -> Iterator[BasicKernel]: raise TypeError(f"Cannot traverse kernel '{type(kernel).__name__}'.") +def _is_gpytorch_kernel_equivalent( + kernel: GPyTorchKernel, other: GPyTorchKernel, / +) -> bool: + """Check if two GPyTorch kernels are equivalent. + + Two kernels are considered equivalent if they have the same module structure, + the same public module attributes, and agree on the values of all tensors that + are not changed by fitting. Tensors changed by fitting are only compared by + shape, since their values merely serve as initialization. + + Args: + kernel: The first kernel. + other: The kernel to compare against. + + Returns: + ``True`` if the kernels are equivalent, ``False`` otherwise. + """ + modules = list(kernel.named_modules()) + other_modules = list(other.named_modules()) + if [(n, type(m)) for n, m in modules] != [(n, type(m)) for n, m in other_modules]: + return False + for (_, module), (_, other_module) in zip(modules, other_modules): + public = {k: v for k, v in vars(module).items() if not k.startswith("_")} + other_public = { + k: v for k, v in vars(other_module).items() if not k.startswith("_") + } + if public != other_public: + return False + + # Tensors requiring gradients are changed by fitting, so only shapes must agree + tensors = dict(chain(kernel.named_parameters(), kernel.named_buffers())) + other_tensors = dict(chain(other.named_parameters(), other.named_buffers())) + if tensors.keys() != other_tensors.keys(): + return False + for name, tensor in tensors.items(): + other_tensor = other_tensors[name] + if (tensor.shape, tensor.requires_grad) != ( + other_tensor.shape, + other_tensor.requires_grad, + ): + return False + if not tensor.requires_grad and not tensor.equal(other_tensor): + return False + return True + + def _to_kernel_override(value: KernelOverride, instance: Parameter) -> KernelOverride: """Validate a kernel override and scope BayBE kernels to their parameter. @@ -206,6 +253,10 @@ def is_equivalent(self, other: Parameter) -> bool: Two parameters are considered equivalent if they have the same type and all attributes are equal except for the name. + GPyTorch kernel overrides are compared structurally, since GPyTorch kernels + define no equality. Values of tensors changed by fitting (i.e., initial + hyperparameter values) do not affect equivalence. + Args: other: The parameter to compare against. @@ -214,10 +265,19 @@ def is_equivalent(self, other: Parameter) -> bool: """ if type(self) is not type(other): return False - # The override is owner-scoped, so rebind it to the other parameter's name. kernel_override = self.kernel_override + other_override = other.kernel_override if isinstance(kernel_override, Kernel): + # The override is owner-scoped, so rebind it to the other parameter's name kernel_override = kernel_override._scope_to_parameter(other.name) + elif ( + kernel_override is not None + and other_override is not None + and not isinstance(other_override, Kernel) + and _is_gpytorch_kernel_equivalent(kernel_override, other_override) + ): + # GPyTorch kernels define no equality, so substitute equivalent overrides + kernel_override = other_override return ( attrs.evolve(self, name=other.name, kernel_override=kernel_override) == other diff --git a/tests/test_parameter_kernel_overrides.py b/tests/test_parameter_kernel_overrides.py index bc458a314f..8fbd81bb8f 100644 --- a/tests/test_parameter_kernel_overrides.py +++ b/tests/test_parameter_kernel_overrides.py @@ -11,6 +11,7 @@ from attrs import evolve from botorch.models.kernels.positive_index import PositiveIndexKernel from gpytorch import kernels as gk +from gpytorch.priors import GammaPrior from pytest import param from baybe.exceptions import IncompatibleOverrideError @@ -53,6 +54,14 @@ def _leaf_kernels(kernel): ) +def _rbf_kernel(lengthscale, *, frozen=False): + """Create a GPyTorch RBF kernel with a given, optionally frozen, lengthscale.""" + kernel = gk.RBFKernel() + kernel.lengthscale = lengthscale + kernel.raw_lengthscale.requires_grad_(not frozen) + return kernel + + @pytest.mark.parametrize( "mode", [None, *TransferLearningMode], @@ -246,6 +255,50 @@ def test_gpytorch_ard_mismatch_rejected(): _resolve([parameter]) +@pytest.mark.parametrize( + ("left", "right", "expected"), + [ + param(None, None, True, id="none"), + param(RBFKernel(), RBFKernel(), True, id="baybe"), + param(gk.MaternKernel(), gk.MaternKernel(), True, id="gpytorch"), + param( + gk.ScaleKernel(gk.RBFKernel(lengthscale_prior=GammaPrior(3.0, 6.0))), + gk.ScaleKernel(gk.RBFKernel(lengthscale_prior=GammaPrior(3.0, 6.0))), + True, + id="gpytorch-nested", + ), + param( + gk.IndexKernel(num_tasks=2), + gk.IndexKernel(num_tasks=2), + True, + id="gpytorch-random-init", + ), + param(_rbf_kernel(2.0), _rbf_kernel(3.0), True, id="gpytorch-init-value"), + param(RBFKernel(), MaternKernel(), False, id="baybe-class"), + param(gk.MaternKernel(nu=0.5), gk.MaternKernel(), False, id="gpytorch-attr"), + param( + gk.RBFKernel(lengthscale_prior=GammaPrior(3.0, 6.0)), + gk.RBFKernel(lengthscale_prior=GammaPrior(2.0, 6.0)), + False, + id="gpytorch-prior", + ), + param( + _rbf_kernel(2.0, frozen=True), + _rbf_kernel(3.0, frozen=True), + False, + id="gpytorch-frozen-value", + ), + param(RBFKernel(), gk.RBFKernel(), False, id="baybe-vs-gpytorch"), + param(gk.RBFKernel(), None, False, id="gpytorch-vs-none"), + ], +) +def test_parameter_equivalence(left, right, expected): + """Parameter equivalence with kernel overrides produces the expected result.""" + p1 = NumericalContinuousParameter("p1", (0, 1), kernel_override=left) + p2 = NumericalContinuousParameter("p2", (0, 1), kernel_override=right) + assert p1.is_equivalent(p2) == expected + + @pytest.mark.parametrize( "kernel_or_factory", [ From 2357887ab8ea7821bdc7b9888f0a9cbec4f4ca93 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Thu, 24 Sep 2026 16:50:42 +0200 Subject: [PATCH 42/73] Rename kernel_override to override_kernel --- baybe/parameters/base.py | 18 +++++++-------- baybe/parameters/categorical.py | 2 +- baybe/surrogates/base.py | 2 +- .../gaussian_process/_override/parameter.py | 4 ++-- baybe/surrogates/gaussian_process/core.py | 2 +- docs/components/parameters.md | 8 +++---- tests/hypothesis_strategies/parameters.py | 20 ++++++++--------- tests/test_iterations.py | 2 +- tests/test_parameter_kernel_overrides.py | 22 +++++++++---------- tests/test_surrogate.py | 2 +- tests/validation/test_parameter_validation.py | 6 ++--- 11 files changed, 44 insertions(+), 44 deletions(-) diff --git a/baybe/parameters/base.py b/baybe/parameters/base.py index 7687f13ccd..f3cb0620d5 100644 --- a/baybe/parameters/base.py +++ b/baybe/parameters/base.py @@ -190,7 +190,7 @@ class Parameter(ABC, SerialMixin): name: str = field(validator=(instance_of(str), min_len(1))) """The name of the parameter""" - kernel_override: KernelOverride | None = field( + override_kernel: KernelOverride | None = field( default=None, converter=optional_c(Converter(_to_kernel_override, takes_self=True)), # type: ignore[misc, call-overload] kw_only=True, @@ -265,21 +265,21 @@ def is_equivalent(self, other: Parameter) -> bool: """ if type(self) is not type(other): return False - kernel_override = self.kernel_override - other_override = other.kernel_override - if isinstance(kernel_override, Kernel): + override_kernel = self.override_kernel + other_override = other.override_kernel + if isinstance(override_kernel, Kernel): # The override is owner-scoped, so rebind it to the other parameter's name - kernel_override = kernel_override._scope_to_parameter(other.name) + override_kernel = override_kernel._scope_to_parameter(other.name) elif ( - kernel_override is not None + override_kernel is not None and other_override is not None and not isinstance(other_override, Kernel) - and _is_gpytorch_kernel_equivalent(kernel_override, other_override) + and _is_gpytorch_kernel_equivalent(override_kernel, other_override) ): # GPyTorch kernels define no equality, so substitute equivalent overrides - kernel_override = other_override + override_kernel = other_override return ( - attrs.evolve(self, name=other.name, kernel_override=kernel_override) + attrs.evolve(self, name=other.name, override_kernel=override_kernel) == other ) diff --git a/baybe/parameters/categorical.py b/baybe/parameters/categorical.py index 8284c6ac28..e8b02397c0 100644 --- a/baybe/parameters/categorical.py +++ b/baybe/parameters/categorical.py @@ -88,7 +88,7 @@ class TaskParameter(CategoricalParameter): encoding: CategoricalEncoding = field(default=CategoricalEncoding.INT, init=False) # See base class. - kernel_override: None = field(init=False, default=None) + override_kernel: None = field(init=False, default=None) """Task parameters do not support parameter-specific kernel overrides.""" override_transfer_learning_mode: TransferLearningMode | None = field( diff --git a/baybe/surrogates/base.py b/baybe/surrogates/base.py index bb1d1cd1d0..a3a053ca0f 100644 --- a/baybe/surrogates/base.py +++ b/baybe/surrogates/base.py @@ -441,7 +441,7 @@ def fit( ) if not self.supports_kernel_overrides and ( - names := [p.name for p in searchspace.parameters if p.kernel_override] + names := [p.name for p in searchspace.parameters if p.override_kernel] ): warnings.warn( f"The selected surrogate model ({self.__class__.__name__}) does not " diff --git a/baybe/surrogates/gaussian_process/_override/parameter.py b/baybe/surrogates/gaussian_process/_override/parameter.py index 342ecbf0b2..f9bc5509f5 100644 --- a/baybe/surrogates/gaussian_process/_override/parameter.py +++ b/baybe/surrogates/gaussian_process/_override/parameter.py @@ -30,7 +30,7 @@ def extract_parameter_overrides( return [ (p.name, make_parameter_override_kernel(p, context.searchspace)) for p in context.searchspace.parameters - if p.kernel_override is not None + if p.override_kernel is not None ] @@ -47,7 +47,7 @@ def make_parameter_override_kernel( Returns: The GPyTorch kernel bound to the parameter's dimensions. """ - override = parameter.kernel_override + override = parameter.override_kernel assert override is not None # BayBE kernels resolve their own dimensions; raw kernels are bound manually. diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py index 7fd436b8ae..16e5d1225a 100644 --- a/baybe/surrogates/gaussian_process/core.py +++ b/baybe/surrogates/gaussian_process/core.py @@ -195,7 +195,7 @@ class GaussianProcessSurrogate(Surrogate): * :obj:`.components.kernel.KernelFactoryProtocol` * :class:`gpytorch.kernels.Kernel` - A :attr:`~baybe.parameters.base.Parameter.kernel_override` removes its parameter + A :attr:`~baybe.parameters.base.Parameter.override_kernel` removes its parameter from this kernel and contributes a separate multiplicative factor. A :attr:`~baybe.parameters.categorical.TaskParameter.override_transfer_learning_mode` replaces the task factor in the same way. When not all parameters are removed by diff --git a/docs/components/parameters.md b/docs/components/parameters.md index 41de4e6b9d..16ab637d20 100644 --- a/docs/components/parameters.md +++ b/docs/components/parameters.md @@ -304,7 +304,7 @@ Here, "surrogate kernel" refers to the ``kernel_or_factory`` of the to BayBE's default kernel when none is given. Regular parameters can replace the surrogate kernel on their computational dimensions -using {attr}`~baybe.parameters.base.Parameter.kernel_override`: +using {attr}`~baybe.parameters.base.Parameter.override_kernel`: ```python from baybe.kernels import RBFKernel @@ -313,7 +313,7 @@ from baybe.parameters import NumericalContinuousParameter x3 = NumericalContinuousParameter( name="x3", bounds=(0.0, 1.0), - kernel_override=RBFKernel(), + override_kernel=RBFKernel(), ) ``` @@ -330,7 +330,7 @@ An override can be a configured BayBE or GPyTorch kernel instance. For a BayBE k ```{admonition} Limitations :class: warning -{attr}`~baybe.parameters.base.Parameter.kernel_override` accepts kernel instances only, +{attr}`~baybe.parameters.base.Parameter.override_kernel` accepts kernel instances only, not factories. The surrogate kernel or factory of the {class}`~baybe.surrogates.gaussian_process.core.GaussianProcessSurrogate` must allow the overridden dimensions to be excluded. Otherwise, BayBE raises @@ -344,7 +344,7 @@ elsewhere, they cannot be serialized. The feature applies only to Gaussian process surrogates. Other surrogates ignore kernel overrides and emit an {class}`~baybe.exceptions.UnusedObjectWarning`. Task parameters do -not expose {attr}`~baybe.parameters.base.Parameter.kernel_override`; use +not expose {attr}`~baybe.parameters.base.Parameter.override_kernel`; use {attr}`~baybe.parameters.categorical.TaskParameter.override_transfer_learning_mode` instead. Both override mechanisms can be used in the same searchspace, in which case their kernel factors are multiplied. diff --git a/tests/hypothesis_strategies/parameters.py b/tests/hypothesis_strategies/parameters.py index 13f1cecc90..300904d980 100644 --- a/tests/hypothesis_strategies/parameters.py +++ b/tests/hypothesis_strategies/parameters.py @@ -132,13 +132,13 @@ def numerical_discrete_parameters( ) ) param_metadata = draw(measurable_metadata()) - kernel_override = draw(kernel_overrides(name)) + override_kernel = draw(kernel_overrides(name)) return NumericalDiscreteParameter( name=name, values=values, tolerance=tolerance, metadata=param_metadata, - kernel_override=kernel_override, + override_kernel=override_kernel, ) @@ -148,12 +148,12 @@ def numerical_continuous_parameters(draw: st.DrawFn): name = draw(parameter_names) bounds = draw(intervals(exclude_half_bounded=True, exclude_fully_unbounded=True)) param_metadata = draw(measurable_metadata()) - kernel_override = draw(kernel_overrides(name)) + override_kernel = draw(kernel_overrides(name)) return NumericalContinuousParameter( name=name, bounds=bounds, metadata=param_metadata, - kernel_override=kernel_override, + override_kernel=override_kernel, ) @@ -165,14 +165,14 @@ def categorical_parameters(draw: st.DrawFn): encoding = draw(st.sampled_from(CategoricalEncoding)) active_values = draw(_active_values(values)) param_metadata = draw(measurable_metadata()) - kernel_override = draw(kernel_overrides(name)) + override_kernel = draw(kernel_overrides(name)) return CategoricalParameter( name=name, values=values, encoding=encoding, active_values=active_values, metadata=param_metadata, - kernel_override=kernel_override, + override_kernel=override_kernel, ) @@ -210,7 +210,7 @@ def substance_parameters(draw: st.DrawFn): encoding = draw(st.sampled_from(encodings)) param_metadata = draw(measurable_metadata()) - kernel_override = draw(kernel_overrides(name)) + override_kernel = draw(kernel_overrides(name)) return SubstanceParameter( name=name, @@ -219,7 +219,7 @@ def substance_parameters(draw: st.DrawFn): encoding=encoding, active_values=active_values, metadata=param_metadata, - kernel_override=kernel_override, + override_kernel=override_kernel, ) @@ -231,14 +231,14 @@ def custom_parameters(draw: st.DrawFn): decorrelate = draw(decorrelations) active_values = draw(_active_values(list(data.index))) param_metadata = draw(measurable_metadata()) - kernel_override = draw(kernel_overrides(name)) + override_kernel = draw(kernel_overrides(name)) return CustomDiscreteParameter( name=name, data=data, decorrelate=decorrelate, active_values=active_values, metadata=param_metadata, - kernel_override=kernel_override, + override_kernel=override_kernel, ) diff --git a/tests/test_iterations.py b/tests/test_iterations.py index 164f4adbbe..fa88383b6f 100644 --- a/tests/test_iterations.py +++ b/tests/test_iterations.py @@ -409,7 +409,7 @@ def test_kernels(ongoing_campaign, n_iterations, batch_size): [ [ NumericalContinuousParameter("x1", (0, 1)), - NumericalContinuousParameter("x2", (0, 1), kernel_override=RBFKernel()), + NumericalContinuousParameter("x2", (0, 1), override_kernel=RBFKernel()), *( [] if mode is None diff --git a/tests/test_parameter_kernel_overrides.py b/tests/test_parameter_kernel_overrides.py index 8fbd81bb8f..221ab3d21b 100644 --- a/tests/test_parameter_kernel_overrides.py +++ b/tests/test_parameter_kernel_overrides.py @@ -75,7 +75,7 @@ def test_selector_does_not_exclude_overridden_parameters(mode): """ parameters = [ NumericalContinuousParameter("x1", (0, 1)), - NumericalContinuousParameter("x2", (0, 1), kernel_override=RBFKernel()), + NumericalContinuousParameter("x2", (0, 1), override_kernel=RBFKernel()), NumericalContinuousParameter("omitted", (0, 1)), ] expected_names = ["x1", "x2"] @@ -115,7 +115,7 @@ def test_nondefault_residual_indices(kernel_or_factory, expected_residual_dims): kernel, _ = _resolve( [ NumericalContinuousParameter("x1", (0, 1)), - NumericalContinuousParameter("x2", (0, 1), kernel_override=RBFKernel()), + NumericalContinuousParameter("x2", (0, 1), override_kernel=RBFKernel()), NumericalContinuousParameter("x3", (0, 1)), ], kernel_or_factory, @@ -159,8 +159,8 @@ def test_fitted_model_uses_parameter_kernel_overrides( ): """The fitted model uses each configured kernel on the intended dimensions.""" parameters = [ - CategoricalParameter("base", ["a", "b"], kernel_override=base_override), - evolve(override_parameter, kernel_override=RBFKernel()), + CategoricalParameter("base", ["a", "b"], override_kernel=base_override), + evolve(override_parameter, override_kernel=RBFKernel()), ] expected = {"base": gk.MaternKernel, "override": gk.RBFKernel} if task_parameter is not None: @@ -207,7 +207,7 @@ def test_valid_baybe_kernel_override(override): """Valid BayBE overrides bind all their leaves to the owning parameter.""" parameters = [ CategoricalParameter("base", ["a", "b"]), - CategoricalParameter("override", ["a", "b", "c"], kernel_override=override), + CategoricalParameter("override", ["a", "b", "c"], override_kernel=override), ] kernel, searchspace = _resolve(parameters) expected = set(searchspace.get_comp_rep_parameter_indices("override")) @@ -234,7 +234,7 @@ def test_valid_gpytorch_kernel_override(categories, override): snapshot = deepcopy(override.state_dict()) base = NumericalContinuousParameter("base", (0, 1)) kernel, searchspace = _resolve( - [base, CategoricalParameter("override", categories, kernel_override=override)] + [base, CategoricalParameter("override", categories, override_kernel=override)] ) expected = set(searchspace.get_comp_rep_parameter_indices("override")) @@ -250,7 +250,7 @@ def test_valid_gpytorch_kernel_override(categories, override): def test_gpytorch_ard_mismatch_rejected(): """A GPyTorch override with mismatched ARD size is rejected during resolution.""" override = gk.RBFKernel(ard_num_dims=2) - parameter = CategoricalParameter("p", ["a", "b", "c"], kernel_override=override) + parameter = CategoricalParameter("p", ["a", "b", "c"], override_kernel=override) with pytest.raises(IncompatibleOverrideError, match="has 3 computational"): _resolve([parameter]) @@ -294,8 +294,8 @@ def test_gpytorch_ard_mismatch_rejected(): ) def test_parameter_equivalence(left, right, expected): """Parameter equivalence with kernel overrides produces the expected result.""" - p1 = NumericalContinuousParameter("p1", (0, 1), kernel_override=left) - p2 = NumericalContinuousParameter("p2", (0, 1), kernel_override=right) + p1 = NumericalContinuousParameter("p1", (0, 1), override_kernel=left) + p2 = NumericalContinuousParameter("p2", (0, 1), override_kernel=right) assert p1.is_equivalent(p2) == expected @@ -310,7 +310,7 @@ def test_incompatible_surrogate_kernel_is_rejected(kernel_or_factory): """Surrogate kernels that cannot exclude overridden dimensions are rejected.""" parameters = [ NumericalContinuousParameter("x1", (0, 1)), - NumericalContinuousParameter("x2", (0, 1), kernel_override=RBFKernel()), + NumericalContinuousParameter("x2", (0, 1), override_kernel=RBFKernel()), ] with pytest.raises(IncompatibleOverrideError, match="can exclude these parameters"): @@ -340,7 +340,7 @@ def test_partition_validation(monkeypatch, override_kind, kernel_cls, active_dim ) if override_kind == "tl" else NumericalContinuousParameter( - "override", (0, 1), kernel_override=RBFKernel() + "override", (0, 1), override_kernel=RBFKernel() ), ] searchspace = SearchSpace.from_product(parameters) diff --git a/tests/test_surrogate.py b/tests/test_surrogate.py index 5302dde0f7..076dd745d0 100644 --- a/tests/test_surrogate.py +++ b/tests/test_surrogate.py @@ -51,7 +51,7 @@ def test_caching(patched, searchspace, objective, fake_measurements): def test_ignored_kernel_override_warning(): """Non-GP surrogates warn when ignoring a parameter kernel override.""" - parameter = NumericalDiscreteParameter("x", [0, 1], kernel_override=RBFKernel()) + parameter = NumericalDiscreteParameter("x", [0, 1], override_kernel=RBFKernel()) measurements = pd.DataFrame({"x": [0, 1], "y": [0.0, 1.0]}) with pytest.warns(UnusedObjectWarning, match=r"kernel overrides.*\['x'\]"): diff --git a/tests/validation/test_parameter_validation.py b/tests/validation/test_parameter_validation.py index 60d2286f7e..de01518e49 100644 --- a/tests/validation/test_parameter_validation.py +++ b/tests/validation/test_parameter_validation.py @@ -49,7 +49,7 @@ def test_invalid_parameter_name(name, error): partial(TaskParameter, "task", ["a", "b"]), RBFKernel(), TypeError, - "kernel_override", + "override_kernel", id="task_parameter", ), param( @@ -92,7 +92,7 @@ def test_invalid_parameter_name(name, error): def test_invalid_kernel_override(constructor, override, error, match): """Invalid parameter kernel overrides raise the expected exceptions.""" with pytest.raises(error, match=match): - constructor(kernel_override=override) + constructor(override_kernel=override) @pytest.mark.parametrize("scaled", [False, True], ids=["direct", "nested"]) @@ -101,7 +101,7 @@ def test_unsupported_kernel_override_structure(scaled): kernel = Mock(spec=Kernel) with pytest.raises(TypeError, match="Cannot traverse kernel"): NumericalContinuousParameter( - "x", (0, 1), kernel_override=ScaleKernel(kernel) if scaled else kernel + "x", (0, 1), override_kernel=ScaleKernel(kernel) if scaled else kernel ) From 135b2020e51741116b3979c608903ced2a964ddf Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Thu, 24 Sep 2026 17:13:06 +0200 Subject: [PATCH 43/73] Back override_kernel by a private attribute --- baybe/parameters/base.py | 10 ++++++++-- baybe/parameters/categorical.py | 2 +- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/baybe/parameters/base.py b/baybe/parameters/base.py index f3cb0620d5..453bfc816a 100644 --- a/baybe/parameters/base.py +++ b/baybe/parameters/base.py @@ -190,12 +190,13 @@ class Parameter(ABC, SerialMixin): name: str = field(validator=(instance_of(str), min_len(1))) """The name of the parameter""" - override_kernel: KernelOverride | None = field( + _override_kernel: KernelOverride | None = field( default=None, + alias="override_kernel", converter=optional_c(Converter(_to_kernel_override, takes_self=True)), # type: ignore[misc, call-overload] kw_only=True, ) - """An optional kernel replacing the overall kernel for this parameter.""" + """The optional kernel override, exposed via :attr:`override_kernel`.""" metadata: MeasurableMetadata = field( factory=MeasurableMetadata, @@ -229,6 +230,11 @@ def is_discrete(self) -> bool: """Boolean indicating if this is a discrete parameter.""" return isinstance(self, DiscreteParameter) + @property + def override_kernel(self) -> KernelOverride | None: + """An optional kernel replacing the overall kernel for this parameter.""" + return self._override_kernel + @property def _kind(self) -> _ParameterKind: """The kind of the parameter.""" diff --git a/baybe/parameters/categorical.py b/baybe/parameters/categorical.py index e8b02397c0..c01c646e66 100644 --- a/baybe/parameters/categorical.py +++ b/baybe/parameters/categorical.py @@ -88,7 +88,7 @@ class TaskParameter(CategoricalParameter): encoding: CategoricalEncoding = field(default=CategoricalEncoding.INT, init=False) # See base class. - override_kernel: None = field(init=False, default=None) + _override_kernel: None = field(init=False, default=None) """Task parameters do not support parameter-specific kernel overrides.""" override_transfer_learning_mode: TransferLearningMode | None = field( From bf0b64be5d413fbcfb0883cfeae6693d640b7a74 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Thu, 24 Sep 2026 17:15:40 +0200 Subject: [PATCH 44/73] Store kernel overrides unscoped --- baybe/parameters/base.py | 50 +++++++++++++----------- tests/test_parameter_kernel_overrides.py | 16 ++++++++ 2 files changed, 44 insertions(+), 22 deletions(-) diff --git a/baybe/parameters/base.py b/baybe/parameters/base.py index 453bfc816a..3eca984413 100644 --- a/baybe/parameters/base.py +++ b/baybe/parameters/base.py @@ -7,7 +7,7 @@ from abc import ABC, abstractmethod from functools import cached_property from itertools import chain -from typing import TYPE_CHECKING, Any, ClassVar, TypeAlias +from typing import TYPE_CHECKING, Any, ClassVar, TypeAlias, cast import attrs import pandas as pd @@ -117,7 +117,7 @@ def _is_gpytorch_kernel_equivalent( def _to_kernel_override(value: KernelOverride, instance: Parameter) -> KernelOverride: - """Validate a kernel override and scope BayBE kernels to their parameter. + """Validate a kernel override and store BayBE kernels unscoped. Args: value: The provided kernel override. @@ -129,10 +129,9 @@ def _to_kernel_override(value: KernelOverride, instance: Parameter) -> KernelOve TypeError: If the object is neither a BayBE nor a GPyTorch kernel. Returns: - The validated override, with BayBE kernels scoped to the parameter. + The validated override, with BayBE kernels unscoped. """ - # BayBE kernels: every basic leaf must be unscoped or scoped to the owner. The - # kernel is then rebound to the owning parameter (dropping unspecified names). + # BayBE kernels: every basic leaf must be unscoped or scoped to the owner. if isinstance(value, Kernel): from baybe.kernels.base import BasicKernel @@ -147,7 +146,13 @@ def _to_kernel_override(value: KernelOverride, instance: Parameter) -> KernelOve f"itself. Its basic kernels must specify '{names_alias}' as " f"``None`` or ({instance.name!r},)." ) - return value._scope_to_parameter(instance.name) + # NOTE: Validated BayBE kernels are stored unscoped and only scoped to the + # owning parameter when accessed via `override_kernel`. This keeps the + # stored value independent of the parameter name, so that renaming (e.g. + # via `attrs.evolve`) and `is_equivalent` work without rescoping. User + # input is still restricted to unscoped kernels or kernels scoped to the + # owning parameter. + return value._scope_to_parameter(None) # GPyTorch kernels: no explicit active dimensions allowed anywhere in the tree. # An existing GPyTorch instance implies the module is already imported. Avoid @@ -233,7 +238,9 @@ def is_discrete(self) -> bool: @property def override_kernel(self) -> KernelOverride | None: """An optional kernel replacing the overall kernel for this parameter.""" - return self._override_kernel + if isinstance(kernel := self._override_kernel, Kernel): + return kernel._scope_to_parameter(self.name) + return kernel @property def _kind(self) -> _ParameterKind: @@ -271,23 +278,22 @@ def is_equivalent(self, other: Parameter) -> bool: """ if type(self) is not type(other): return False - override_kernel = self.override_kernel - other_override = other.override_kernel - if isinstance(override_kernel, Kernel): - # The override is owner-scoped, so rebind it to the other parameter's name - override_kernel = override_kernel._scope_to_parameter(other.name) - elif ( - override_kernel is not None - and other_override is not None - and not isinstance(other_override, Kernel) - and _is_gpytorch_kernel_equivalent(override_kernel, other_override) + # Stored overrides are unscoped (see `_to_kernel_override`), so only the name + # needs to be aligned for the comparison + changes: dict[str, Any] = {} + kernel, other_kernel = self._override_kernel, other._override_kernel + if ( + kernel is not None + and not isinstance(kernel, Kernel) + and other_kernel is not None + and not isinstance(other_kernel, Kernel) + and _is_gpytorch_kernel_equivalent(kernel, other_kernel) ): # GPyTorch kernels define no equality, so substitute equivalent overrides - override_kernel = other_override - return ( - attrs.evolve(self, name=other.name, override_kernel=override_kernel) - == other - ) + # NOTE: attrs resolves all aliases during class creation + alias = cast(str, attrs.fields(Parameter)._override_kernel.alias) + changes[alias] = other_kernel + return attrs.evolve(self, name=other.name, **changes) == other @abstractmethod def summary(self) -> dict: diff --git a/tests/test_parameter_kernel_overrides.py b/tests/test_parameter_kernel_overrides.py index 221ab3d21b..bc71d7ac21 100644 --- a/tests/test_parameter_kernel_overrides.py +++ b/tests/test_parameter_kernel_overrides.py @@ -299,6 +299,22 @@ def test_parameter_equivalence(left, right, expected): assert p1.is_equivalent(p2) == expected +@pytest.mark.parametrize( + ("values", "mode", "expected"), + [ + param(["a", "b"], TransferLearningMode.INDEX_KERNEL, True, id="same"), + param(["a", "b"], TransferLearningMode.POSITIVE_INDEX_KERNEL, False, id="mode"), + param(["a", "b", "c"], TransferLearningMode.INDEX_KERNEL, False, id="values"), + ], +) +def test_task_parameter_equivalence(values, mode, expected): + """Task parameter equivalence produces the expected result.""" + mode_ref = TransferLearningMode.INDEX_KERNEL + p1 = TaskParameter("t1", ["a", "b"], override_transfer_learning_mode=mode_ref) + p2 = TaskParameter("t2", values, override_transfer_learning_mode=mode) + assert p1.is_equivalent(p2) == expected + + @pytest.mark.parametrize( "kernel_or_factory", [ From df92f4e014988407826944876f2d3ac3daa697d5 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Thu, 24 Sep 2026 17:19:59 +0200 Subject: [PATCH 45/73] Derive task parameter kernel override from transfer learning mode --- baybe/parameters/categorical.py | 26 ++++- .../gaussian_process/_override/__init__.py | 14 +-- .../_override/{core.py => kernel.py} | 88 ++++++++++++++++- .../gaussian_process/_override/parameter.py | 99 ------------------- .../gaussian_process/_override/tl.py | 62 ------------ baybe/surrogates/gaussian_process/core.py | 32 +++--- docs/components/parameters.md | 6 +- tests/test_parameter_kernel_overrides.py | 30 +++++- 8 files changed, 158 insertions(+), 199 deletions(-) rename baybe/surrogates/gaussian_process/_override/{core.py => kernel.py} (50%) delete mode 100644 baybe/surrogates/gaussian_process/_override/parameter.py delete mode 100644 baybe/surrogates/gaussian_process/_override/tl.py diff --git a/baybe/parameters/categorical.py b/baybe/parameters/categorical.py index c01c646e66..0b38e5083c 100644 --- a/baybe/parameters/categorical.py +++ b/baybe/parameters/categorical.py @@ -8,8 +8,9 @@ from attr.converters import optional as optional_c from attrs import Converter, define, field from attrs.validators import deep_iterable, instance_of, min_len -from typing_extensions import override +from typing_extensions import assert_never, override +from baybe.kernels.base import Kernel from baybe.parameters.base import _DiscreteLabelLikeParameter from baybe.parameters.enum import CategoricalEncoding, TransferLearningMode from baybe.settings import active_settings @@ -89,7 +90,7 @@ class TaskParameter(CategoricalParameter): # See base class. _override_kernel: None = field(init=False, default=None) - """Task parameters do not support parameter-specific kernel overrides.""" + """Task parameters derive their kernel override from the transfer learning mode.""" override_transfer_learning_mode: TransferLearningMode | None = field( default=None, @@ -103,6 +104,27 @@ class TaskParameter(CategoricalParameter): configured factory. """ + @override + @property + def override_kernel(self) -> Kernel | None: + """The task kernel defined by the transfer learning mode, if any.""" + from baybe.kernels.basic import IndexKernel, PositiveIndexKernel + + n_tasks, names = len(self.values), (self.name,) + match mode := self.override_transfer_learning_mode: + case None: + return None + case TransferLearningMode.POSITIVE_INDEX_KERNEL: + return PositiveIndexKernel( + num_tasks=n_tasks, rank=n_tasks, parameter_names=names + ) + case TransferLearningMode.INDEX_KERNEL: + return IndexKernel( + num_tasks=n_tasks, rank=n_tasks, parameter_names=names + ) + case _: + assert_never(mode) + # Collect leftover original slotted classes processed by `attrs.define` gc.collect() diff --git a/baybe/surrogates/gaussian_process/_override/__init__.py b/baybe/surrogates/gaussian_process/_override/__init__.py index 579dc9db3c..1fcf372d69 100644 --- a/baybe/surrogates/gaussian_process/_override/__init__.py +++ b/baybe/surrogates/gaussian_process/_override/__init__.py @@ -1,21 +1,15 @@ -"""Kernel override resolution for Gaussian process surrogates.""" +"""Component override resolution for Gaussian process surrogates.""" -from baybe.surrogates.gaussian_process._override.core import ( +from baybe.surrogates.gaussian_process._override.kernel import ( + extract_parameter_kernel_overrides, get_active_dimensions, raise_incompatible_override, reduce_kernel_spec, ) -from baybe.surrogates.gaussian_process._override.parameter import ( - extract_parameter_overrides, -) -from baybe.surrogates.gaussian_process._override.tl import ( - extract_transfer_learning_overrides, -) __all__ = [ + "extract_parameter_kernel_overrides", "get_active_dimensions", - "extract_parameter_overrides", - "extract_transfer_learning_overrides", "raise_incompatible_override", "reduce_kernel_spec", ] diff --git a/baybe/surrogates/gaussian_process/_override/core.py b/baybe/surrogates/gaussian_process/_override/kernel.py similarity index 50% rename from baybe/surrogates/gaussian_process/_override/core.py rename to baybe/surrogates/gaussian_process/_override/kernel.py index bf7d883728..b89ab32318 100644 --- a/baybe/surrogates/gaussian_process/_override/core.py +++ b/baybe/surrogates/gaussian_process/_override/kernel.py @@ -1,7 +1,8 @@ -"""Shared helpers for kernel override resolution.""" +"""Kernel override resolution for Gaussian process surrogates.""" from __future__ import annotations +from copy import deepcopy from typing import TYPE_CHECKING, NoReturn from baybe.exceptions import IncompatibleOverrideError @@ -10,10 +11,95 @@ if TYPE_CHECKING: from gpytorch.kernels import Kernel as GPyTorchKernel + from baybe.parameters.base import Parameter from baybe.searchspace import SearchSpace from baybe.surrogates.gaussian_process.components.kernel import ( KernelFactoryProtocol, ) + from baybe.surrogates.gaussian_process.core import _ModelContext + + +def extract_parameter_kernel_overrides( + context: _ModelContext, +) -> list[tuple[str, GPyTorchKernel]]: + """Extract the regular parameter-specific kernel overrides. + + Args: + context: The model context providing the search space. + + Returns: + A ``(parameter_name, kernel)`` pair for each parameter with an override. + """ + return [ + (p.name, make_parameter_override_kernel(p, context.searchspace)) + for p in context.searchspace.parameters + if p.override_kernel is not None + ] + + +def make_parameter_override_kernel( + parameter: Parameter, searchspace: SearchSpace +) -> GPyTorchKernel: + """Create the kernel factor for a parameter's override. + + Args: + parameter: The parameter carrying a non-``None`` kernel override, as ensured + by :func:`extract_parameter_kernel_overrides`. + searchspace: The search space the kernel operates on. + + Returns: + The GPyTorch kernel bound to the parameter's dimensions. + """ + override = parameter.override_kernel + assert override is not None + + # BayBE kernels resolve their own dimensions; raw kernels are bound manually. + if isinstance(override, Kernel): + return override.to_gpytorch(searchspace) + indices = searchspace.get_comp_rep_parameter_indices(parameter.name) + return bind_gpytorch_override(override, indices, parameter.name) + + +def bind_gpytorch_override( + override: GPyTorchKernel, indices: tuple[int, ...], name: str +) -> GPyTorchKernel: + """Copy a raw GPyTorch override and bind it to the given dimensions. + + Args: + override: The provided GPyTorch kernel (must not specify active dimensions). + indices: The computational column indices of the owning parameter. + name: The owning parameter name (for error messages). + + Raises: + IncompatibleOverrideError: If the kernel specifies active dimensions or an + incompatible number of ARD dimensions. + + Returns: + A copy of the kernel bound to the given dimensions. + """ + import torch + from gpytorch.kernels import Kernel as GPyTorchKernel + + for kernel in override.modules(): + if not isinstance(kernel, GPyTorchKernel): + continue + if kernel.active_dims is not None: + raise IncompatibleOverrideError( + f"The GPyTorch kernel override for parameter '{name}' must not " + f"specify 'active_dims'." + ) + if kernel.ard_num_dims not in (None, len(indices)): + raise IncompatibleOverrideError( + f"The GPyTorch kernel override for parameter '{name}' specifies " + f"{kernel.ard_num_dims} ARD dimensions, but the parameter has " + f"{len(indices)} computational dimensions." + ) + + result = deepcopy(override) + # TODO[typing]: GPyTorch annotates `active_dims` with the constructor's tuple + # type, but `register_buffer` stores a tensor at runtime. + result.active_dims = torch.tensor(indices) # pyrefly: ignore[bad-assignment] + return result def get_active_dimensions(kernel: GPyTorchKernel, searchspace: SearchSpace) -> set[int]: diff --git a/baybe/surrogates/gaussian_process/_override/parameter.py b/baybe/surrogates/gaussian_process/_override/parameter.py deleted file mode 100644 index f9bc5509f5..0000000000 --- a/baybe/surrogates/gaussian_process/_override/parameter.py +++ /dev/null @@ -1,99 +0,0 @@ -"""Resolution of parameter-specific kernel overrides.""" - -from __future__ import annotations - -from copy import deepcopy -from typing import TYPE_CHECKING - -from baybe.exceptions import IncompatibleOverrideError -from baybe.kernels.base import Kernel - -if TYPE_CHECKING: - from gpytorch.kernels import Kernel as GPyTorchKernel - - from baybe.parameters.base import Parameter - from baybe.searchspace import SearchSpace - from baybe.surrogates.gaussian_process.core import _ModelContext - - -def extract_parameter_overrides( - context: _ModelContext, -) -> list[tuple[str, GPyTorchKernel]]: - """Extract the regular parameter-specific kernel overrides. - - Args: - context: The model context providing the search space. - - Returns: - A ``(parameter_name, kernel)`` pair for each parameter with an override. - """ - return [ - (p.name, make_parameter_override_kernel(p, context.searchspace)) - for p in context.searchspace.parameters - if p.override_kernel is not None - ] - - -def make_parameter_override_kernel( - parameter: Parameter, searchspace: SearchSpace -) -> GPyTorchKernel: - """Create the kernel factor for a parameter's override. - - Args: - parameter: The parameter carrying a non-``None`` kernel override, as ensured - by :func:`extract_parameter_overrides`. - searchspace: The search space the kernel operates on. - - Returns: - The GPyTorch kernel bound to the parameter's dimensions. - """ - override = parameter.override_kernel - assert override is not None - - # BayBE kernels resolve their own dimensions; raw kernels are bound manually. - if isinstance(override, Kernel): - return override.to_gpytorch(searchspace) - indices = searchspace.get_comp_rep_parameter_indices(parameter.name) - return bind_gpytorch_override(override, indices, parameter.name) - - -def bind_gpytorch_override( - override: GPyTorchKernel, indices: tuple[int, ...], name: str -) -> GPyTorchKernel: - """Copy a raw GPyTorch override and bind it to the given dimensions. - - Args: - override: The provided GPyTorch kernel (must not specify active dimensions). - indices: The computational column indices of the owning parameter. - name: The owning parameter name (for error messages). - - Raises: - IncompatibleOverrideError: If the kernel specifies active dimensions or an - incompatible number of ARD dimensions. - - Returns: - A copy of the kernel bound to the given dimensions. - """ - import torch - from gpytorch.kernels import Kernel as GPyTorchKernel - - for kernel in override.modules(): - if not isinstance(kernel, GPyTorchKernel): - continue - if kernel.active_dims is not None: - raise IncompatibleOverrideError( - f"The GPyTorch kernel override for parameter '{name}' must not " - f"specify 'active_dims'." - ) - if kernel.ard_num_dims not in (None, len(indices)): - raise IncompatibleOverrideError( - f"The GPyTorch kernel override for parameter '{name}' specifies " - f"{kernel.ard_num_dims} ARD dimensions, but the parameter has " - f"{len(indices)} computational dimensions." - ) - - result = deepcopy(override) - # TODO[typing]: GPyTorch annotates `active_dims` with the constructor's tuple - # type, but `register_buffer` stores a tensor at runtime. - result.active_dims = torch.tensor(indices) # pyrefly: ignore[bad-assignment] - return result diff --git a/baybe/surrogates/gaussian_process/_override/tl.py b/baybe/surrogates/gaussian_process/_override/tl.py deleted file mode 100644 index d796fb553e..0000000000 --- a/baybe/surrogates/gaussian_process/_override/tl.py +++ /dev/null @@ -1,62 +0,0 @@ -"""Resolution of the transfer-learning kernel override.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING - -from typing_extensions import assert_never - -from baybe.kernels.base import Kernel -from baybe.parameters.enum import TransferLearningMode - -if TYPE_CHECKING: - from gpytorch.kernels import Kernel as GPyTorchKernel - - from baybe.surrogates.gaussian_process.core import _ModelContext - - -def extract_transfer_learning_overrides( - context: _ModelContext, -) -> list[tuple[str, GPyTorchKernel]]: - """Extract the transfer-learning kernel override, if any. - - Args: - context: The model context providing the task parameter and override mode. - - Returns: - A single ``(task_name, kernel)`` pair if an override is set, else an empty - list. - """ - task_param = context.searchspace._task_parameter - if task_param is None or context.tl_override is None: - return [] - return [(task_param.name, make_transfer_learning_override_kernel(context))] - - -def make_transfer_learning_override_kernel( - context: _ModelContext, -) -> GPyTorchKernel: - """Create the task kernel factor requested by a transfer-learning override. - - Args: - context: The model context providing the task parameter and override mode. - - Returns: - The GPyTorch task kernel for the requested transfer-learning mode. - """ - from baybe.kernels.basic import IndexKernel, PositiveIndexKernel - - task_param = context.searchspace._task_parameter - override = context.tl_override - assert task_param is not None and override is not None - n_tasks, names = context.n_tasks, (task_param.name,) - match override: - case TransferLearningMode.POSITIVE_INDEX_KERNEL: - spec: Kernel = PositiveIndexKernel( - num_tasks=n_tasks, rank=n_tasks, parameter_names=names - ) - case TransferLearningMode.INDEX_KERNEL: - spec = IndexKernel(num_tasks=n_tasks, rank=n_tasks, parameter_names=names) - case _: - assert_never(override) - return spec.to_gpytorch(context.searchspace) diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py index 16e5d1225a..0a76661518 100644 --- a/baybe/surrogates/gaussian_process/core.py +++ b/baybe/surrogates/gaussian_process/core.py @@ -27,7 +27,6 @@ from baybe.objectives.base import Objective from baybe.parameters.base import Parameter from baybe.parameters.categorical import TaskParameter -from baybe.parameters.enum import TransferLearningMode from baybe.searchspace.core import SearchSpace from baybe.surrogates.base import Surrogate from baybe.surrogates.gaussian_process import _override @@ -93,14 +92,6 @@ def task_idx(self) -> int | None: """The computational column index of the task parameter, if available.""" return self.searchspace.task_idx - @property - def tl_override(self) -> TransferLearningMode | None: - """The task parameter's transfer learning override, if any.""" - task_param = self.searchspace._task_parameter - return ( - None if task_param is None else task_param.override_transfer_learning_mode - ) - @property def is_multitask(self) -> bool: """Indicates if model is to be operated in a multi-task context.""" @@ -196,10 +187,11 @@ class GaussianProcessSurrogate(Surrogate): * :class:`gpytorch.kernels.Kernel` A :attr:`~baybe.parameters.base.Parameter.override_kernel` removes its parameter - from this kernel and contributes a separate multiplicative factor. A + from this kernel and contributes a separate multiplicative factor. For task + parameters, it is derived from :attr:`~baybe.parameters.categorical.TaskParameter.override_transfer_learning_mode` - replaces the task factor in the same way. When not all parameters are removed by - an override, the configured kernel or factory must support excluding the + and replaces the task factor in the same way. When not all parameters are removed + by an override, the configured kernel or factory must support excluding the overridden parameters; otherwise, :class:`~baybe.exceptions.IncompatibleOverrideError` is raised. If all parameters are overridden, this kernel or factory is not used. Without overrides, it is used @@ -428,8 +420,7 @@ def _resolve_kernel(self, context: _ModelContext) -> GPyTorchKernel: Returns: The resolved and partition-validated kernel. """ - overrides = _override.extract_parameter_overrides(context) - overrides += _override.extract_transfer_learning_overrides(context) + overrides = _override.extract_parameter_kernel_overrides(context) excluded_names = {name for name, _ in overrides} residual = self._resolve_residual_kernel(context, excluded_names) @@ -599,14 +590,17 @@ def _fit(self, train_x: Tensor, train_y: Tensor) -> None: context = _ModelContext(self._searchspace, self._objective, self._measurements) - # Check for custom kernel + multi-task clash (only relevant when no - # override_transfer_learning_mode is set, since the override mechanism - # handles task kernel attachment explicitly). - has_tl_override = context.tl_override is not None + # Check for custom kernel + multi-task clash (only relevant when the task + # parameter has no kernel override, since the override mechanism handles task + # kernel attachment explicitly). + task_param = self._searchspace._task_parameter + has_task_override = ( + task_param is not None and task_param.override_kernel is not None + ) if ( context.is_multitask and self._custom_kernel - and not has_tl_override + and not has_task_override and not strtobool(os.getenv("BAYBE_DISABLE_CUSTOM_KERNEL_WARNING", "False")) ): raise DeprecationError( diff --git a/docs/components/parameters.md b/docs/components/parameters.md index 16ab637d20..c235a95144 100644 --- a/docs/components/parameters.md +++ b/docs/components/parameters.md @@ -344,7 +344,7 @@ elsewhere, they cannot be serialized. The feature applies only to Gaussian process surrogates. Other surrogates ignore kernel overrides and emit an {class}`~baybe.exceptions.UnusedObjectWarning`. Task parameters do -not expose {attr}`~baybe.parameters.base.Parameter.override_kernel`; use +not support kernel overrides; use {attr}`~baybe.parameters.categorical.TaskParameter.override_transfer_learning_mode` -instead. Both override mechanisms can be used in the same searchspace, in which case -their kernel factors are multiplied. +instead. Kernel overrides of other parameters can be combined with a transfer learning +override of the task parameter, in which case their kernel factors are multiplied. diff --git a/tests/test_parameter_kernel_overrides.py b/tests/test_parameter_kernel_overrides.py index bc71d7ac21..e70f0323ec 100644 --- a/tests/test_parameter_kernel_overrides.py +++ b/tests/test_parameter_kernel_overrides.py @@ -9,14 +9,16 @@ import pytest import torch from attrs import evolve -from botorch.models.kernels.positive_index import PositiveIndexKernel +from botorch.models.kernels.positive_index import ( + PositiveIndexKernel as BoPositiveIndexKernel, +) from gpytorch import kernels as gk from gpytorch.priors import GammaPrior from pytest import param from baybe.exceptions import IncompatibleOverrideError from baybe.kernels import MaternKernel, RBFKernel -from baybe.kernels.basic import IndexKernel +from baybe.kernels.basic import IndexKernel, PositiveIndexKernel from baybe.kernels.composite import AdditiveKernel, ProductKernel, ScaleKernel from baybe.parameters import ( CategoricalParameter, @@ -169,7 +171,7 @@ def test_fitted_model_uses_parameter_kernel_overrides( gk.IndexKernel if task_parameter.override_transfer_learning_mode == TransferLearningMode.INDEX_KERNEL - else PositiveIndexKernel + else BoPositiveIndexKernel ) searchspace = SearchSpace.from_product(parameters) measurements = pd.DataFrame( @@ -315,6 +317,28 @@ def test_task_parameter_equivalence(values, mode, expected): assert p1.is_equivalent(p2) == expected +@pytest.mark.parametrize( + ("mode", "expected"), + [ + param(None, None, id="none"), + param( + TransferLearningMode.INDEX_KERNEL, + IndexKernel(num_tasks=3, rank=3, parameter_names=("task",)), + id="index", + ), + param( + TransferLearningMode.POSITIVE_INDEX_KERNEL, + PositiveIndexKernel(num_tasks=3, rank=3, parameter_names=("task",)), + id="positive-index", + ), + ], +) +def test_task_parameter_override_kernel(mode, expected): + """Task parameters derive their kernel override from the transfer learning mode.""" + task = TaskParameter("task", ["a", "b", "c"], override_transfer_learning_mode=mode) + assert task.override_kernel == expected + + @pytest.mark.parametrize( "kernel_or_factory", [ From 0c6e21e584dd9212c8f7e80c885a03f5d7aea778 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Thu, 24 Sep 2026 17:59:01 +0200 Subject: [PATCH 46/73] Regularize override kernel in iteration test The unregularized RBF override on x2 intermittently caused NaN acquisition gradients (OptimizationGradientError) with few data points, which the retry logic of run_iterations does not catch. A lengthscale prior stabilizes the fit: 6 of 60 seeds failed without it and none with it. --- tests/test_iterations.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/test_iterations.py b/tests/test_iterations.py index fa88383b6f..dbf1fcb6a2 100644 --- a/tests/test_iterations.py +++ b/tests/test_iterations.py @@ -409,7 +409,11 @@ def test_kernels(ongoing_campaign, n_iterations, batch_size): [ [ NumericalContinuousParameter("x1", (0, 1)), - NumericalContinuousParameter("x2", (0, 1), override_kernel=RBFKernel()), + NumericalContinuousParameter( + "x2", + (0, 1), + override_kernel=RBFKernel(lengthscale_prior=GammaPrior(3.0, 6.0)), + ), *( [] if mode is None From f26beb6017abb800c0628b2436374b140d8d66f5 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Mon, 5 Oct 2026 15:01:36 +0200 Subject: [PATCH 47/73] Improve docstring Co-authored-by: Alexander V. Hopp --- baybe/parameters/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/baybe/parameters/base.py b/baybe/parameters/base.py index 3eca984413..b5e4100561 100644 --- a/baybe/parameters/base.py +++ b/baybe/parameters/base.py @@ -126,7 +126,7 @@ def _to_kernel_override(value: KernelOverride, instance: Parameter) -> KernelOve Raises: ValueError: If a BayBE kernel targets a different parameter or a GPyTorch kernel specifies explicit active dimensions. - TypeError: If the object is neither a BayBE nor a GPyTorch kernel. + TypeError: If ``value`` is neither a BayBE nor a GPyTorch kernel. Returns: The validated override, with BayBE kernels unscoped. From 516ebacfd1a30adf55502a2266ffb5b1f080507a Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Mon, 5 Oct 2026 15:21:33 +0200 Subject: [PATCH 48/73] Cache override_kernel property --- baybe/parameters/base.py | 2 +- baybe/parameters/categorical.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/baybe/parameters/base.py b/baybe/parameters/base.py index b5e4100561..26224f205c 100644 --- a/baybe/parameters/base.py +++ b/baybe/parameters/base.py @@ -235,7 +235,7 @@ def is_discrete(self) -> bool: """Boolean indicating if this is a discrete parameter.""" return isinstance(self, DiscreteParameter) - @property + @cached_property def override_kernel(self) -> KernelOverride | None: """An optional kernel replacing the overall kernel for this parameter.""" if isinstance(kernel := self._override_kernel, Kernel): diff --git a/baybe/parameters/categorical.py b/baybe/parameters/categorical.py index 0b38e5083c..47b642f71e 100644 --- a/baybe/parameters/categorical.py +++ b/baybe/parameters/categorical.py @@ -105,7 +105,7 @@ class TaskParameter(CategoricalParameter): """ @override - @property + @cached_property def override_kernel(self) -> Kernel | None: """The task kernel defined by the transfer learning mode, if any.""" from baybe.kernels.basic import IndexKernel, PositiveIndexKernel From fd3b369c3b543efcf3740da7d6691847b4edb356 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Mon, 5 Oct 2026 15:21:43 +0200 Subject: [PATCH 49/73] Clarify kernel type in _reduce_base_kernels docstring --- baybe/kernels/composite.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/baybe/kernels/composite.py b/baybe/kernels/composite.py index 7aeb7e93a1..c2819fc443 100644 --- a/baybe/kernels/composite.py +++ b/baybe/kernels/composite.py @@ -26,16 +26,16 @@ def _reduce_base_kernels( kernel: AdditiveKernel | ProductKernel, name: str, searchspace: SearchSpace, / ) -> Kernel | None: - """Remove a parameter from all base kernels of a composite kernel. + """Remove a parameter from all base kernels of an additive or product kernel. Args: - kernel: The composite kernel whose base kernels are to be reduced. + kernel: The additive or product kernel whose base kernels are to be reduced. name: The name of the parameter to remove. searchspace: The search space the kernel operates on. Returns: - The reduced composite kernel, the sole remaining base kernel, or ``None`` - if no base kernel remains. + The reduced kernel of the same type, the sole remaining base kernel, or + ``None`` if no base kernel remains. """ remaining = tuple( reduced From c0405bcd7838141f8b975de1a16ea1b71779eb73 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Mon, 5 Oct 2026 15:21:50 +0200 Subject: [PATCH 50/73] Clarify parameter scoping docstring --- baybe/kernels/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/baybe/kernels/base.py b/baybe/kernels/base.py index d88dad44df..7d8fe51b4b 100644 --- a/baybe/kernels/base.py +++ b/baybe/kernels/base.py @@ -122,7 +122,7 @@ def _without_parameter( ) def _scope_to_parameter(self, name: str | None, /) -> Kernel: - """Return a copy of the kernel that acts only on the given parameter. + """Return a copy of the kernel that acts only on the named parameter. Args: name: The name of the parameter to scope the kernel to, or ``None`` to From b5da88d933849c556c4fbfb7c7455e9bf08bff34 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Mon, 5 Oct 2026 15:22:25 +0200 Subject: [PATCH 51/73] Simplify GPyTorch override binding --- .../gaussian_process/_override/kernel.py | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/baybe/surrogates/gaussian_process/_override/kernel.py b/baybe/surrogates/gaussian_process/_override/kernel.py index b89ab32318..a548d9c10c 100644 --- a/baybe/surrogates/gaussian_process/_override/kernel.py +++ b/baybe/surrogates/gaussian_process/_override/kernel.py @@ -66,28 +66,20 @@ def bind_gpytorch_override( """Copy a raw GPyTorch override and bind it to the given dimensions. Args: - override: The provided GPyTorch kernel (must not specify active dimensions). + override: The provided GPyTorch kernel. indices: The computational column indices of the owning parameter. name: The owning parameter name (for error messages). Raises: - IncompatibleOverrideError: If the kernel specifies active dimensions or an - incompatible number of ARD dimensions. + IncompatibleOverrideError: If the kernel specifies an incompatible number of + ARD dimensions. Returns: A copy of the kernel bound to the given dimensions. """ import torch - from gpytorch.kernels import Kernel as GPyTorchKernel - for kernel in override.modules(): - if not isinstance(kernel, GPyTorchKernel): - continue - if kernel.active_dims is not None: - raise IncompatibleOverrideError( - f"The GPyTorch kernel override for parameter '{name}' must not " - f"specify 'active_dims'." - ) + for kernel in (override, *override.sub_kernels()): if kernel.ard_num_dims not in (None, len(indices)): raise IncompatibleOverrideError( f"The GPyTorch kernel override for parameter '{name}' specifies " From 197868062363a8229a0e25f8432c3883ef40062b Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Mon, 5 Oct 2026 15:23:46 +0200 Subject: [PATCH 52/73] Use default iteration count in override iteration test --- tests/test_iterations.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_iterations.py b/tests/test_iterations.py index dbf1fcb6a2..e383ccce6a 100644 --- a/tests/test_iterations.py +++ b/tests/test_iterations.py @@ -428,7 +428,6 @@ def test_kernels(ongoing_campaign, n_iterations, batch_size): ], ids=["parameter_kernel_override", *(mode.name for mode in TransferLearningMode)], ) -@pytest.mark.parametrize("n_iterations", [3], ids=["i3"]) def test_parameter_kernel_override_iteration( ongoing_campaign, n_iterations, batch_size ): From 69b946070da2f013af491ea2822454085972205e Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Mon, 5 Oct 2026 15:24:29 +0200 Subject: [PATCH 53/73] Test composite surrogate kernels with parameter overrides --- tests/test_parameter_kernel_overrides.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/tests/test_parameter_kernel_overrides.py b/tests/test_parameter_kernel_overrides.py index e70f0323ec..8fb01d28b7 100644 --- a/tests/test_parameter_kernel_overrides.py +++ b/tests/test_parameter_kernel_overrides.py @@ -110,6 +110,9 @@ def test_selector_does_not_exclude_overridden_parameters(mode): (0, 2), id="callable", ), + param(MaternKernel() + RBFKernel(), (0, 2), id="sum"), + param(MaternKernel() * RBFKernel(), (0, 2), id="product"), + param(ScaleKernel(MaternKernel() + RBFKernel()), (0, 2), id="scaled-sum"), ], ) def test_nondefault_residual_indices(kernel_or_factory, expected_residual_dims): @@ -122,8 +125,14 @@ def test_nondefault_residual_indices(kernel_or_factory, expected_residual_dims): ], kernel_or_factory, ) - residual, override = kernel.kernels - assert tuple(residual.active_dims.tolist()) == expected_residual_dims + # The override is the last factor; all remaining basic kernels form the residual + override = kernel.kernels[-1] + residual_dims = { + tuple(k.active_dims.tolist()) + for k in _leaf_kernels(kernel) + if k is not override + } + assert residual_dims == {expected_residual_dims} assert tuple(override.active_dims.tolist()) == (1,) assert isinstance(override, gk.RBFKernel) From f8e0754bdd96e5680f0cac9b85c5fbe452eaa1c4 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Wed, 14 Oct 2026 13:14:55 +0200 Subject: [PATCH 54/73] Improve equation Co-authored-by: AdrianSosic --- docs/components/parameters.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/components/parameters.md b/docs/components/parameters.md index c235a95144..c0f84f296d 100644 --- a/docs/components/parameters.md +++ b/docs/components/parameters.md @@ -318,7 +318,7 @@ x3 = NumericalContinuousParameter( ``` In a search space with parameters ``x1``, ``x2``, and ``x3``, this produces -``surrogate kernel(x1, x2) * RBF(x3)``. The override does not replace a dimension inside +``surrogate_kernel(x1, x2) * RBF(x3)``. The override does not replace a dimension inside the surrogate kernel itself. It partitions the dimensions and creates a separate kernel factor. Multiple overrides produce multiple factors. For a parameter with several computational dimensions, the override acts jointly on the complete computational From b2fdfeb4ea34996c5c15a0947f72c83abfbd7f41 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Wed, 14 Oct 2026 13:15:46 +0200 Subject: [PATCH 55/73] Improve wording Co-authored-by: AdrianSosic --- docs/components/parameters.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/components/parameters.md b/docs/components/parameters.md index c0f84f296d..de69cd3f50 100644 --- a/docs/components/parameters.md +++ b/docs/components/parameters.md @@ -303,7 +303,7 @@ Here, "surrogate kernel" refers to the ``kernel_or_factory`` of the {class}`~baybe.surrogates.gaussian_process.core.GaussianProcessSurrogate`, falling back to BayBE's default kernel when none is given. -Regular parameters can replace the surrogate kernel on their computational dimensions +Most parameter types allow to replace the surrogate kernel on their computational dimensions using {attr}`~baybe.parameters.base.Parameter.override_kernel`: ```python From dfdfb5df428830d41b09d1e3f6e67b22298a94bb Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Wed, 14 Oct 2026 13:16:35 +0200 Subject: [PATCH 56/73] Improve documentation Co-authored-by: AdrianSosic --- docs/components/parameters.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/docs/components/parameters.md b/docs/components/parameters.md index de69cd3f50..6e9227a273 100644 --- a/docs/components/parameters.md +++ b/docs/components/parameters.md @@ -288,6 +288,9 @@ For details, refer to [transfer learning](../concepts/transfer_learning.md). (parameter_kernel_overrides)= ## Parameter-Specific Kernel Overrides +During a Bayesian optimization loop, the search space spanned by the user-defined +parameters will be presented to the surrogate model (e.g., a Gaussian process), +to evaluate its points in terms of their expected utility. By default, the Gaussian process models all parameters with a single surrogate kernel. A parameter can instead be assigned its own kernel, which then replaces the surrogate kernel on that parameter's computational dimensions. The remaining dimensions stay with From 5c62b083ff57d259ce8ffbc138c036587a6ed822 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Thu, 15 Oct 2026 08:10:25 +0200 Subject: [PATCH 57/73] Reuse lazy GPyTorch component check for kernel overrides --- baybe/parameters/base.py | 29 +++++------ .../gaussian_process/components/generic.py | 49 ++++++++++++++----- 2 files changed, 50 insertions(+), 28 deletions(-) diff --git a/baybe/parameters/base.py b/baybe/parameters/base.py index 26224f205c..9a2863460f 100644 --- a/baybe/parameters/base.py +++ b/baybe/parameters/base.py @@ -3,7 +3,6 @@ from __future__ import annotations import gc -import sys from abc import ABC, abstractmethod from functools import cached_property from itertools import chain @@ -155,22 +154,18 @@ def _to_kernel_override(value: KernelOverride, instance: Parameter) -> KernelOve return value._scope_to_parameter(None) # GPyTorch kernels: no explicit active dimensions allowed anywhere in the tree. - # An existing GPyTorch instance implies the module is already imported. Avoid - # importing it (and Torch) solely to validate other parameter inputs. - if sys.modules.get("gpytorch") is not None: - from gpytorch.kernels import Kernel as GPyTorchKernel - - if isinstance(value, GPyTorchKernel): - if any( - k.active_dims is not None - for k in value.modules() - if isinstance(k, GPyTorchKernel) - ): - raise ValueError( - "The GPyTorch kernel provided for the kernel override must not " - "specify 'active_dims'." - ) - return value + from baybe.surrogates.gaussian_process.components.generic import ( + GPComponentType, + _is_gpytorch_component, + ) + + if _is_gpytorch_component(value, (GPComponentType.KERNEL,)): + if any(k.active_dims is not None for k in (value, *value.sub_kernels())): + raise ValueError( + "The GPyTorch kernel provided for the kernel override must not " + "specify 'active_dims'." + ) + return value raise TypeError( f"The object provided for the kernel override of " diff --git a/baybe/surrogates/gaussian_process/components/generic.py b/baybe/surrogates/gaussian_process/components/generic.py index 57e48d9bc8..005a9513a2 100644 --- a/baybe/surrogates/gaussian_process/components/generic.py +++ b/baybe/surrogates/gaussian_process/components/generic.py @@ -4,6 +4,7 @@ import gc import sys +from collections.abc import Collection from enum import Enum from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeAlias, TypeVar @@ -83,20 +84,46 @@ def get_types(self) -> tuple[type, ...]: return tuple(types) -def _is_gpytorch_component_class(obj: Any, /) -> bool: - """Check if a class is a GPyTorch component class using lazy loading.""" - if sys.modules.get("gpytorch") is None: - return False - from gpytorch.kernels import Kernel as GPyTorchKernel - from gpytorch.likelihoods import Likelihood as GPyTorchLikelihood - from gpytorch.means import Mean as GPyTorchMean +def _is_gpytorch_component_class( + obj: Any, + /, + component_types: Collection[GPComponentType] = tuple(GPComponentType), +) -> bool: + """Check if a class is a GPyTorch component class using lazy loading. - return issubclass(obj, (GPyTorchKernel, GPyTorchMean, GPyTorchLikelihood)) + Args: + obj: The class to check. + component_types: The component types to check against. + Returns: + ``True`` if the class is a GPyTorch class of one of the given component types. + """ + # GPyTorch types are only returned if GPyTorch has already been imported + types = tuple( + t + for c in component_types + for t in c.get_types() + if not issubclass(t, BayBEGPComponent) + ) + return issubclass(obj, types) -def _is_gpytorch_component(obj: Any, /) -> TypeIs[GPyTorchGPComponent]: - """Check if an object is a GPyTorch component instance using lazy loading.""" - return _is_gpytorch_component_class(type(obj)) + +def _is_gpytorch_component( + obj: Any, + /, + component_types: Collection[GPComponentType] = tuple(GPComponentType), +) -> TypeIs[GPyTorchGPComponent]: + """Check if an object is a GPyTorch component instance using lazy loading. + + Args: + obj: The object to check. + component_types: The component types to check against. + + Returns: + ``True`` if the object is a GPyTorch instance of one of the given component + types. + """ + return _is_gpytorch_component_class(type(obj), component_types) def _validate_component(instance: Any, attribute: Attribute, value: Any) -> None: From b57e9faeac3ddff04515443747105e8bf9735ead Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Thu, 15 Oct 2026 09:05:09 +0200 Subject: [PATCH 58/73] Improve wording in kernel override comments --- baybe/parameters/base.py | 4 ++-- baybe/surrogates/gaussian_process/core.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/baybe/parameters/base.py b/baybe/parameters/base.py index 9a2863460f..83eb041d19 100644 --- a/baybe/parameters/base.py +++ b/baybe/parameters/base.py @@ -145,8 +145,8 @@ def _to_kernel_override(value: KernelOverride, instance: Parameter) -> KernelOve f"itself. Its basic kernels must specify '{names_alias}' as " f"``None`` or ({instance.name!r},)." ) - # NOTE: Validated BayBE kernels are stored unscoped and only scoped to the - # owning parameter when accessed via `override_kernel`. This keeps the + # NOTE: Validated BayBE kernels are stored unscoped and are only scoped to + # the owning parameter when accessed via `override_kernel`. This keeps the # stored value independent of the parameter name, so that renaming (e.g. # via `attrs.evolve`) and `is_equivalent` work without rescoping. User # input is still restricted to unscoped kernels or kernels scoped to the diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py index 0a76661518..2338d6c8b7 100644 --- a/baybe/surrogates/gaussian_process/core.py +++ b/baybe/surrogates/gaussian_process/core.py @@ -406,7 +406,7 @@ def _posterior(self, candidates_comp_scaled: Tensor, /) -> Posterior: return self._model.posterior(candidates_comp_scaled) def _resolve_kernel(self, context: _ModelContext) -> GPyTorchKernel: - """Resolve the GP kernel, applying parameter and transfer overrides. + """Resolve the GP kernel, applying parameter-specific overrides. The effective kernel is the surrogate kernel restricted to the non-overridden dimensions, multiplied by one factor per override. From e7b9c0ca85974a34f4126a9f3e9149afc8c30f0b Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Thu, 15 Oct 2026 09:05:31 +0200 Subject: [PATCH 59/73] Improve references in kernel override docs --- docs/components/parameters.md | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/docs/components/parameters.md b/docs/components/parameters.md index 6e9227a273..71e8ffa872 100644 --- a/docs/components/parameters.md +++ b/docs/components/parameters.md @@ -302,9 +302,10 @@ k_{\mathrm{effective}} \prod_{p \in P_{\mathrm{override}}} k_p(X_p). ``` -Here, "surrogate kernel" refers to the ``kernel_or_factory`` of the -{class}`~baybe.surrogates.gaussian_process.core.GaussianProcessSurrogate`, falling back -to BayBE's default kernel when none is given. +Here, "surrogate kernel" refers to the +{attr}`~baybe.surrogates.gaussian_process.core.GaussianProcessSurrogate.kernel_factory` +of the {class}`~baybe.surrogates.gaussian_process.core.GaussianProcessSurrogate`, +falling back to BayBE's default kernel when none is given. Most parameter types allow to replace the surrogate kernel on their computational dimensions using {attr}`~baybe.parameters.base.Parameter.override_kernel`: @@ -328,8 +329,9 @@ computational dimensions, the override acts jointly on the complete computationa representation. An override can be a configured BayBE or GPyTorch kernel instance. For a BayBE kernel, -``parameter_names`` must be ``None`` or contain exactly the owning parameter. Here, -``None`` is rebound to the owning parameter instead of acting on all dimensions. +{attr}`~baybe.kernels.base.BasicKernel.parameter_names` must be ``None`` or contain +exactly the owning parameter. Here, ``None`` is rebound to the owning parameter instead +of acting on all dimensions. ```{admonition} Limitations :class: warning From 1f77d6922697e53274edde1ef6b2f38810168841 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Thu, 15 Oct 2026 09:13:10 +0200 Subject: [PATCH 60/73] Improve override extraction signature --- baybe/surrogates/gaussian_process/_override/kernel.py | 11 +++++------ baybe/surrogates/gaussian_process/core.py | 2 +- 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/baybe/surrogates/gaussian_process/_override/kernel.py b/baybe/surrogates/gaussian_process/_override/kernel.py index a548d9c10c..a9cb7dad01 100644 --- a/baybe/surrogates/gaussian_process/_override/kernel.py +++ b/baybe/surrogates/gaussian_process/_override/kernel.py @@ -16,23 +16,22 @@ from baybe.surrogates.gaussian_process.components.kernel import ( KernelFactoryProtocol, ) - from baybe.surrogates.gaussian_process.core import _ModelContext def extract_parameter_kernel_overrides( - context: _ModelContext, + searchspace: SearchSpace, ) -> list[tuple[str, GPyTorchKernel]]: - """Extract the regular parameter-specific kernel overrides. + """Extract the parameter-specific kernel overrides. Args: - context: The model context providing the search space. + searchspace: The search space whose parameters are checked for overrides. Returns: A ``(parameter_name, kernel)`` pair for each parameter with an override. """ return [ - (p.name, make_parameter_override_kernel(p, context.searchspace)) - for p in context.searchspace.parameters + (p.name, make_parameter_override_kernel(p, searchspace)) + for p in searchspace.parameters if p.override_kernel is not None ] diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py index 2338d6c8b7..be906e8ffd 100644 --- a/baybe/surrogates/gaussian_process/core.py +++ b/baybe/surrogates/gaussian_process/core.py @@ -420,7 +420,7 @@ def _resolve_kernel(self, context: _ModelContext) -> GPyTorchKernel: Returns: The resolved and partition-validated kernel. """ - overrides = _override.extract_parameter_kernel_overrides(context) + overrides = _override.extract_parameter_kernel_overrides(context.searchspace) excluded_names = {name for name, _ in overrides} residual = self._resolve_residual_kernel(context, excluded_names) From 2e182346d6e19cd166cd63d3ca0e3535bd30d347 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Thu, 15 Oct 2026 09:36:11 +0200 Subject: [PATCH 61/73] Improve type safety of default kernel base resolution --- baybe/surrogates/gaussian_process/core.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py index be906e8ffd..b0d86ea502 100644 --- a/baybe/surrogates/gaussian_process/core.py +++ b/baybe/surrogates/gaussian_process/core.py @@ -474,7 +474,9 @@ def _resolve_residual_kernel( # Default kernel: build the (task-free) numerical base restricted to the # remaining parameters and re-add the default task kernel unless overridden. if isinstance(factory, BayBEKernelFactory): - return self._resolve_default_base(context, excluded_names) + return self._resolve_default_base( + context, excluded_names, factory.parameter_selector + ) # A fixed kernel is reduced directly by removing the excluded parameters. if isinstance(factory, PlainGPComponentFactory): @@ -500,7 +502,10 @@ def _resolve_residual_kernel( return None if spec is None else spec.to_gpytorch(searchspace) def _resolve_default_base( - self, context: _ModelContext, excluded_names: set[str] + self, + context: _ModelContext, + excluded_names: set[str], + selector: ParameterSelectorProtocol | None, ) -> GPyTorchKernel | None: """Resolve the default kernel base, excluding the given parameters. @@ -513,7 +518,6 @@ def _resolve_default_base( ) searchspace = context.searchspace - selector = getattr(self.kernel_factory, "parameter_selector", None) task_param = searchspace._task_parameter keep: ParameterSelectorProtocol = lambda parameter: ( # noqa: E731 From 6199c44b3e08803e64ca1f692390e194d3c52785 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Thu, 15 Oct 2026 09:43:58 +0200 Subject: [PATCH 62/73] Improve typing workaround comments with issue links --- baybe/parameters/base.py | 3 ++- baybe/surrogates/gaussian_process/_override/kernel.py | 6 ++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/baybe/parameters/base.py b/baybe/parameters/base.py index 83eb041d19..d8aa310705 100644 --- a/baybe/parameters/base.py +++ b/baybe/parameters/base.py @@ -285,7 +285,8 @@ def is_equivalent(self, other: Parameter) -> bool: and _is_gpytorch_kernel_equivalent(kernel, other_kernel) ): # GPyTorch kernels define no equality, so substitute equivalent overrides - # NOTE: attrs resolves all aliases during class creation + # TODO[typing]: attrs resolves all aliases during class creation, see + # https://github.com/python-attrs/attrs/issues/1540 alias = cast(str, attrs.fields(Parameter)._override_kernel.alias) changes[alias] = other_kernel return attrs.evolve(self, name=other.name, **changes) == other diff --git a/baybe/surrogates/gaussian_process/_override/kernel.py b/baybe/surrogates/gaussian_process/_override/kernel.py index a9cb7dad01..a513c5625d 100644 --- a/baybe/surrogates/gaussian_process/_override/kernel.py +++ b/baybe/surrogates/gaussian_process/_override/kernel.py @@ -88,7 +88,8 @@ def bind_gpytorch_override( result = deepcopy(override) # TODO[typing]: GPyTorch annotates `active_dims` with the constructor's tuple - # type, but `register_buffer` stores a tensor at runtime. + # type, but `register_buffer` stores a tensor at runtime, see + # https://github.com/facebook/pyrefly/issues/5208 result.active_dims = torch.tensor(indices) # pyrefly: ignore[bad-assignment] return result @@ -107,7 +108,8 @@ def get_active_dimensions(kernel: GPyTorchKernel, searchspace: SearchSpace) -> s if kernel.active_dims is not None: # TODO[typing]: GPyTorch annotates `active_dims` with the constructor's tuple - # type, but `register_buffer` stores a tensor at runtime. + # type, but `register_buffer` stores a tensor at runtime, see + # https://github.com/facebook/pyrefly/issues/5208 return set(kernel.active_dims.tolist()) # pyrefly: ignore[missing-attribute] if isinstance(kernel, (AdditiveKernel, ProductKernel)): return set().union( From 954da73035dd11ef035e44cace62f1336c2dd97c Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Thu, 15 Oct 2026 12:19:09 +0200 Subject: [PATCH 63/73] Remove branch-specific kernel override tests --- tests/test_iterations.py | 37 -- tests/test_parameter_kernel_overrides.py | 433 ------------------ tests/validation/test_parameter_validation.py | 68 --- 3 files changed, 538 deletions(-) delete mode 100644 tests/test_parameter_kernel_overrides.py diff --git a/tests/test_iterations.py b/tests/test_iterations.py index e383ccce6a..c3491f46bc 100644 --- a/tests/test_iterations.py +++ b/tests/test_iterations.py @@ -27,11 +27,6 @@ ) from baybe.kernels.composite import ScaleKernel from baybe.objectives.pareto import ParetoObjective -from baybe.parameters import ( - NumericalContinuousParameter, - TaskParameter, - TransferLearningMode, -) from baybe.priors import ( GammaPrior, HalfCauchyPrior, @@ -403,38 +398,6 @@ def test_kernels(ongoing_campaign, n_iterations, batch_size): run_iterations(ongoing_campaign, n_iterations, batch_size) -@pytest.mark.slow -@pytest.mark.parametrize( - "parameters", - [ - [ - NumericalContinuousParameter("x1", (0, 1)), - NumericalContinuousParameter( - "x2", - (0, 1), - override_kernel=RBFKernel(lengthscale_prior=GammaPrior(3.0, 6.0)), - ), - *( - [] - if mode is None - else [ - TaskParameter( - "task", ["a", "b"], override_transfer_learning_mode=mode - ) - ] - ), - ] - for mode in [None, *TransferLearningMode] - ], - ids=["parameter_kernel_override", *(mode.name for mode in TransferLearningMode)], -) -def test_parameter_kernel_override_iteration( - ongoing_campaign, n_iterations, batch_size -): - """A complete optimization iteration supports a parameter kernel override.""" - run_iterations(ongoing_campaign, n_iterations, batch_size) - - @pytest.mark.parametrize("kernel", valid_kernel_factories) @pytest.mark.parametrize("n_iterations", [3], ids=["i3"]) def test_kernel_factories(ongoing_campaign, n_iterations, batch_size): diff --git a/tests/test_parameter_kernel_overrides.py b/tests/test_parameter_kernel_overrides.py deleted file mode 100644 index 8fb01d28b7..0000000000 --- a/tests/test_parameter_kernel_overrides.py +++ /dev/null @@ -1,433 +0,0 @@ -"""Functional tests for parameter-specific kernel overrides. - -These tests resolve or fit kernels to check binding and composition behavior. -""" - -from copy import deepcopy - -import pandas as pd -import pytest -import torch -from attrs import evolve -from botorch.models.kernels.positive_index import ( - PositiveIndexKernel as BoPositiveIndexKernel, -) -from gpytorch import kernels as gk -from gpytorch.priors import GammaPrior -from pytest import param - -from baybe.exceptions import IncompatibleOverrideError -from baybe.kernels import MaternKernel, RBFKernel -from baybe.kernels.basic import IndexKernel, PositiveIndexKernel -from baybe.kernels.composite import AdditiveKernel, ProductKernel, ScaleKernel -from baybe.parameters import ( - CategoricalParameter, - NumericalContinuousParameter, - NumericalDiscreteParameter, - TaskParameter, -) -from baybe.parameters.enum import TransferLearningMode -from baybe.parameters.selectors import NameSelector -from baybe.searchspace import SearchSpace -from baybe.surrogates import GaussianProcessSurrogate -from baybe.surrogates.gaussian_process.components.kernel import ICMKernelFactory -from baybe.surrogates.gaussian_process.core import _ModelContext -from baybe.surrogates.gaussian_process.presets import BayBEKernelFactory -from baybe.targets import NumericalTarget - - -def _resolve(parameters, kernel_or_factory=None): - """Resolve the default GP kernel for the given parameters.""" - searchspace = SearchSpace.from_product(parameters) - context = _ModelContext( - searchspace, NumericalTarget("y").to_objective(), pd.DataFrame() - ) - surrogate = GaussianProcessSurrogate(kernel_or_factory=kernel_or_factory) - return surrogate._resolve_kernel(context), searchspace - - -def _leaf_kernels(kernel): - """Return the leaf kernels of a GPyTorch kernel tree.""" - children = tuple(kernel.sub_kernels()) - return ( - (kernel,) - if not children - else tuple(k for c in children for k in _leaf_kernels(c)) - ) - - -def _rbf_kernel(lengthscale, *, frozen=False): - """Create a GPyTorch RBF kernel with a given, optionally frozen, lengthscale.""" - kernel = gk.RBFKernel() - kernel.lengthscale = lengthscale - kernel.raw_lengthscale.requires_grad_(not frozen) - return kernel - - -@pytest.mark.parametrize( - "mode", - [None, *TransferLearningMode], - ids=lambda mode: mode.name if mode else "no-task", -) -def test_selector_does_not_exclude_overridden_parameters(mode): - """A selector narrows the surrogate kernel, but overrides apply regardless. - - The factory selects only ``x1``, so ``omitted`` contributes no kernel factor. - ``x2`` is likewise unselected, yet its override still yields its own factor. - """ - parameters = [ - NumericalContinuousParameter("x1", (0, 1)), - NumericalContinuousParameter("x2", (0, 1), override_kernel=RBFKernel()), - NumericalContinuousParameter("omitted", (0, 1)), - ] - expected_names = ["x1", "x2"] - if mode is not None: - parameters.append( - TaskParameter("task", ["a", "b"], override_transfer_learning_mode=mode) - ) - expected_names.append("task") - - kernel, searchspace = _resolve( - parameters, - BayBEKernelFactory(parameter_selector=NameSelector(("x1",), regex=False)), - ) - - actual_dimensions = {tuple(k.active_dims.tolist()) for k in _leaf_kernels(kernel)} - expected_dimensions = { - searchspace.get_comp_rep_parameter_indices(name) for name in expected_names - } - assert actual_dimensions == expected_dimensions - - -@pytest.mark.parametrize( - ("kernel_or_factory", "expected_residual_dims"), - [ - param(MaternKernel(), (0, 2), id="unnamed"), - param(MaternKernel(parameter_names=("x1", "x2")), (0,), id="named"), - param(ScaleKernel(MaternKernel()), (0, 2), id="scaled"), - param( - lambda space, _, __: MaternKernel(parameter_names=space.parameter_names), - (0, 2), - id="callable", - ), - param(MaternKernel() + RBFKernel(), (0, 2), id="sum"), - param(MaternKernel() * RBFKernel(), (0, 2), id="product"), - param(ScaleKernel(MaternKernel() + RBFKernel()), (0, 2), id="scaled-sum"), - ], -) -def test_nondefault_residual_indices(kernel_or_factory, expected_residual_dims): - """Removing a middle parameter preserves the original computational indices.""" - kernel, _ = _resolve( - [ - NumericalContinuousParameter("x1", (0, 1)), - NumericalContinuousParameter("x2", (0, 1), override_kernel=RBFKernel()), - NumericalContinuousParameter("x3", (0, 1)), - ], - kernel_or_factory, - ) - # The override is the last factor; all remaining basic kernels form the residual - override = kernel.kernels[-1] - residual_dims = { - tuple(k.active_dims.tolist()) - for k in _leaf_kernels(kernel) - if k is not override - } - assert residual_dims == {expected_residual_dims} - assert tuple(override.active_dims.tolist()) == (1,) - assert isinstance(override, gk.RBFKernel) - - -@pytest.mark.parametrize( - ("override_parameter", "base_override", "task_parameter"), - [ - (NumericalDiscreteParameter("override", [0, 1, 2]), None, None), - (CategoricalParameter("override", ["a", "b", "c"]), None, None), - (CategoricalParameter("override", ["a", "b"]), MaternKernel(), None), - ( - CategoricalParameter("override", ["a", "b"]), - None, - TaskParameter("task", ["a", "b"]), - ), - *[ - ( - CategoricalParameter("override", ["a", "b", "c"]), - None, - TaskParameter("task", ["a", "b"], override_transfer_learning_mode=mode), - ) - for mode in TransferLearningMode - ], - ], - ids=[ - "single-dim", - "multi-dim", - "all-overridden", - "task-without-tl-override", - *(mode.name for mode in TransferLearningMode), - ], -) -def test_fitted_model_uses_parameter_kernel_overrides( - override_parameter, base_override, task_parameter -): - """The fitted model uses each configured kernel on the intended dimensions.""" - parameters = [ - CategoricalParameter("base", ["a", "b"], override_kernel=base_override), - evolve(override_parameter, override_kernel=RBFKernel()), - ] - expected = {"base": gk.MaternKernel, "override": gk.RBFKernel} - if task_parameter is not None: - parameters.append(task_parameter) - expected["task"] = ( - gk.IndexKernel - if task_parameter.override_transfer_learning_mode - == TransferLearningMode.INDEX_KERNEL - else BoPositiveIndexKernel - ) - searchspace = SearchSpace.from_product(parameters) - measurements = pd.DataFrame( - [ - {**{p.name: p.values[i] for p in parameters}, "y": y} - for i, y in [(0, 0.0), (-1, 1.0)] - ] - ) - surrogate = GaussianProcessSurrogate() - surrogate.fit(searchspace, NumericalTarget("y").to_objective(), measurements) - assert surrogate._model is not None - kernel = surrogate._model.covar_module - leaves = _leaf_kernels(kernel) - - assert isinstance(kernel, gk.ProductKernel) - assert len(leaves) == len(expected) - assert {tuple(k.active_dims.tolist()): type(k) for k in leaves} == { - searchspace.get_comp_rep_parameter_indices(name): cls - for name, cls in expected.items() - } - assert all(k.ard_num_dims == len(k.active_dims) for k in leaves) - - -@pytest.mark.parametrize( - "override", - [ - param(RBFKernel(), id="unscoped"), - param(RBFKernel(parameter_names=("override",)), id="owner-scoped"), - param(ScaleKernel(RBFKernel()), id="scale"), - param(AdditiveKernel([RBFKernel(), MaternKernel()]), id="additive"), - param(ProductKernel([RBFKernel(), MaternKernel()]), id="product"), - ], -) -def test_valid_baybe_kernel_override(override): - """Valid BayBE overrides bind all their leaves to the owning parameter.""" - parameters = [ - CategoricalParameter("base", ["a", "b"]), - CategoricalParameter("override", ["a", "b", "c"], override_kernel=override), - ] - kernel, searchspace = _resolve(parameters) - expected = set(searchspace.get_comp_rep_parameter_indices("override")) - base_dims = set(searchspace.get_comp_rep_parameter_indices("base")) - - override_leaves = [ - k for k in _leaf_kernels(kernel) if set(k.active_dims.tolist()) != base_dims - ] - assert override_leaves - assert all(set(k.active_dims.tolist()) == expected for k in override_leaves) - - -@pytest.mark.parametrize( - ("categories", "override"), - [ - param(["a", "b"], gk.RBFKernel(), id="no-ard"), - param(["a", "b"], gk.RBFKernel(ard_num_dims=2), id="ard"), - param(["a", "b", "c"], gk.RBFKernel(ard_num_dims=3), id="multi-column-ard"), - param(["a", "b"], gk.ScaleKernel(gk.RBFKernel()), id="nested"), - ], -) -def test_valid_gpytorch_kernel_override(categories, override): - """Valid GPyTorch overrides are bound to the owner without mutating the input.""" - snapshot = deepcopy(override.state_dict()) - base = NumericalContinuousParameter("base", (0, 1)) - kernel, searchspace = _resolve( - [base, CategoricalParameter("override", categories, override_kernel=override)] - ) - expected = set(searchspace.get_comp_rep_parameter_indices("override")) - - # The override factor is the product factor acting on the override dimensions. - factor = next(k for k in kernel.kernels if set(k.active_dims.tolist()) == expected) - assert type(factor) is type(override) - - # The provided kernel is copied, not mutated. - assert override.active_dims is None - assert all(torch.equal(snapshot[k], override.state_dict()[k]) for k in snapshot) - - -def test_gpytorch_ard_mismatch_rejected(): - """A GPyTorch override with mismatched ARD size is rejected during resolution.""" - override = gk.RBFKernel(ard_num_dims=2) - parameter = CategoricalParameter("p", ["a", "b", "c"], override_kernel=override) - with pytest.raises(IncompatibleOverrideError, match="has 3 computational"): - _resolve([parameter]) - - -@pytest.mark.parametrize( - ("left", "right", "expected"), - [ - param(None, None, True, id="none"), - param(RBFKernel(), RBFKernel(), True, id="baybe"), - param(gk.MaternKernel(), gk.MaternKernel(), True, id="gpytorch"), - param( - gk.ScaleKernel(gk.RBFKernel(lengthscale_prior=GammaPrior(3.0, 6.0))), - gk.ScaleKernel(gk.RBFKernel(lengthscale_prior=GammaPrior(3.0, 6.0))), - True, - id="gpytorch-nested", - ), - param( - gk.IndexKernel(num_tasks=2), - gk.IndexKernel(num_tasks=2), - True, - id="gpytorch-random-init", - ), - param(_rbf_kernel(2.0), _rbf_kernel(3.0), True, id="gpytorch-init-value"), - param(RBFKernel(), MaternKernel(), False, id="baybe-class"), - param(gk.MaternKernel(nu=0.5), gk.MaternKernel(), False, id="gpytorch-attr"), - param( - gk.RBFKernel(lengthscale_prior=GammaPrior(3.0, 6.0)), - gk.RBFKernel(lengthscale_prior=GammaPrior(2.0, 6.0)), - False, - id="gpytorch-prior", - ), - param( - _rbf_kernel(2.0, frozen=True), - _rbf_kernel(3.0, frozen=True), - False, - id="gpytorch-frozen-value", - ), - param(RBFKernel(), gk.RBFKernel(), False, id="baybe-vs-gpytorch"), - param(gk.RBFKernel(), None, False, id="gpytorch-vs-none"), - ], -) -def test_parameter_equivalence(left, right, expected): - """Parameter equivalence with kernel overrides produces the expected result.""" - p1 = NumericalContinuousParameter("p1", (0, 1), override_kernel=left) - p2 = NumericalContinuousParameter("p2", (0, 1), override_kernel=right) - assert p1.is_equivalent(p2) == expected - - -@pytest.mark.parametrize( - ("values", "mode", "expected"), - [ - param(["a", "b"], TransferLearningMode.INDEX_KERNEL, True, id="same"), - param(["a", "b"], TransferLearningMode.POSITIVE_INDEX_KERNEL, False, id="mode"), - param(["a", "b", "c"], TransferLearningMode.INDEX_KERNEL, False, id="values"), - ], -) -def test_task_parameter_equivalence(values, mode, expected): - """Task parameter equivalence produces the expected result.""" - mode_ref = TransferLearningMode.INDEX_KERNEL - p1 = TaskParameter("t1", ["a", "b"], override_transfer_learning_mode=mode_ref) - p2 = TaskParameter("t2", values, override_transfer_learning_mode=mode) - assert p1.is_equivalent(p2) == expected - - -@pytest.mark.parametrize( - ("mode", "expected"), - [ - param(None, None, id="none"), - param( - TransferLearningMode.INDEX_KERNEL, - IndexKernel(num_tasks=3, rank=3, parameter_names=("task",)), - id="index", - ), - param( - TransferLearningMode.POSITIVE_INDEX_KERNEL, - PositiveIndexKernel(num_tasks=3, rank=3, parameter_names=("task",)), - id="positive-index", - ), - ], -) -def test_task_parameter_override_kernel(mode, expected): - """Task parameters derive their kernel override from the transfer learning mode.""" - task = TaskParameter("task", ["a", "b", "c"], override_transfer_learning_mode=mode) - assert task.override_kernel == expected - - -@pytest.mark.parametrize( - "kernel_or_factory", - [ - param(gk.MaternKernel(), id="raw-gpytorch-kernel"), - param(lambda s, o, m: gk.MaternKernel(), id="factory-returning-raw-kernel"), - ], -) -def test_incompatible_surrogate_kernel_is_rejected(kernel_or_factory): - """Surrogate kernels that cannot exclude overridden dimensions are rejected.""" - parameters = [ - NumericalContinuousParameter("x1", (0, 1)), - NumericalContinuousParameter("x2", (0, 1), override_kernel=RBFKernel()), - ] - - with pytest.raises(IncompatibleOverrideError, match="can exclude these parameters"): - _resolve(parameters, kernel_or_factory) - - -@pytest.mark.parametrize( - ("override_kind", "kernel_cls", "active_dims"), - [ - param("tl", MaternKernel, None, id="tl-unrestricted-residual"), - param("tl", MaternKernel, (0, 1), id="tl-overlapping-residual"), - param("tl", IndexKernel, None, id="unrestricted-task"), - param("tl", IndexKernel, (0, 1), id="overlapping-task"), - param("tl", IndexKernel, (), id="empty-task"), - param("regular", MaternKernel, None, id="regular-unrestricted-residual"), - param("regular", RBFKernel, None, id="unrestricted-override"), - ], -) -def test_partition_validation(monkeypatch, override_kind, kernel_cls, active_dims): - """Reject misbound regular and TL factors, matching ICM for TL partitions.""" - parameters = [ - NumericalContinuousParameter("x", (0, 1)), - TaskParameter( - "task", - ["a", "b"], - override_transfer_learning_mode=TransferLearningMode.INDEX_KERNEL, - ) - if override_kind == "tl" - else NumericalContinuousParameter( - "override", (0, 1), override_kernel=RBFKernel() - ), - ] - searchspace = SearchSpace.from_product(parameters) - objective = NumericalTarget("y").to_objective() - measurements = pd.DataFrame() - - def misbound_kernel(self, searchspace): - return gk.RBFKernel(active_dims=active_dims) - - monkeypatch.setattr(kernel_cls, "to_gpytorch", misbound_kernel) - with pytest.raises(ValueError, match="active_dims"): - GaussianProcessSurrogate(kernel_or_factory=MaternKernel())._resolve_kernel( - _ModelContext(searchspace, objective, measurements) - ) - if override_kind == "tl": - with pytest.raises(ValueError, match="active_dims"): - ICMKernelFactory( - base_kernel_or_factory=MaternKernel(parameter_names=("x",)), - task_kernel_or_factory=IndexKernel( - num_tasks=2, rank=2, parameter_names=("task",) - ), - )(searchspace, objective, measurements) - - -@pytest.mark.parametrize("as_factory", [False, True], ids=["fixed", "callable"]) -def test_raw_surrogate_kernel_without_overrides(as_factory): - """Preserve raw kernels and the factory's original searchspace without overrides.""" - raw = gk.RBFKernel(active_dims=(0,)) - searchspace = SearchSpace.from_product([NumericalContinuousParameter("x", (0, 1))]) - - def factory(received_searchspace, objective, measurements): - assert received_searchspace is searchspace - return raw - - surrogate = GaussianProcessSurrogate( - kernel_or_factory=factory if as_factory else raw - ) - context = _ModelContext( - searchspace, NumericalTarget("y").to_objective(), pd.DataFrame() - ) - assert surrogate._resolve_kernel(context) is raw diff --git a/tests/validation/test_parameter_validation.py b/tests/validation/test_parameter_validation.py index de01518e49..3e1d9f3893 100644 --- a/tests/validation/test_parameter_validation.py +++ b/tests/validation/test_parameter_validation.py @@ -1,6 +1,5 @@ """Validation tests for parameters.""" -from functools import partial from unittest.mock import Mock import numpy as np @@ -8,13 +7,9 @@ import pytest from cattrs.errors import IterableValidationError from exceptiongroup import ExceptionGroup -from gpytorch import kernels as gk from pytest import param from baybe._optional.info import CHEM_INSTALLED -from baybe.kernels import RBFKernel -from baybe.kernels.base import Kernel -from baybe.kernels.composite import ScaleKernel from baybe.parameters.categorical import ( CategoricalParameter, TaskParameter, @@ -42,69 +37,6 @@ def test_invalid_parameter_name(name, error): NumericalDiscreteParameter(name=name, values=[1, 2, 3]) -@pytest.mark.parametrize( - ("constructor", "override", "error", "match"), - [ - param( - partial(TaskParameter, "task", ["a", "b"]), - RBFKernel(), - TypeError, - "override_kernel", - id="task_parameter", - ), - param( - partial(NumericalContinuousParameter, "x", (0, 1)), - RBFKernel(parameter_names=("other",)), - ValueError, - "may only act on the parameter itself", - id="baybe_different_name", - ), - param( - partial(NumericalContinuousParameter, "x", (0, 1)), - RBFKernel(parameter_names=("x", "other")), - ValueError, - "may only act on the parameter itself", - id="baybe_multiple_names", - ), - param( - partial(NumericalContinuousParameter, "x", (0, 1)), - ScaleKernel(RBFKernel(parameter_names=("other",))), - ValueError, - "may only act on the parameter itself", - id="baybe_composite_leaf", - ), - param( - partial(NumericalContinuousParameter, "x", (0, 1)), - gk.RBFKernel(active_dims=[0]), - ValueError, - "must not specify 'active_dims'", - id="gpytorch_active_dims", - ), - param( - partial(NumericalContinuousParameter, "x", (0, 1)), - gk.ScaleKernel(gk.RBFKernel(active_dims=[0])), - ValueError, - "must not specify 'active_dims'", - id="gpytorch_nested_active_dims", - ), - ], -) -def test_invalid_kernel_override(constructor, override, error, match): - """Invalid parameter kernel overrides raise the expected exceptions.""" - with pytest.raises(error, match=match): - constructor(override_kernel=override) - - -@pytest.mark.parametrize("scaled", [False, True], ids=["direct", "nested"]) -def test_unsupported_kernel_override_structure(scaled): - """Unknown kernel structures cannot silently bypass owner-scope validation.""" - kernel = Mock(spec=Kernel) - with pytest.raises(TypeError, match="Cannot traverse kernel"): - NumericalContinuousParameter( - "x", (0, 1), override_kernel=ScaleKernel(kernel) if scaled else kernel - ) - - @pytest.mark.parametrize( ("decorrelation", "error"), [ From 3d4257a9725695ef5559f5fd769d84d69d6ab9cb Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Thu, 15 Oct 2026 12:20:09 +0200 Subject: [PATCH 64/73] Add kernel override assignment test --- tests/test_parameter_kernel_overrides.py | 161 +++++++++++++++++++++++ 1 file changed, 161 insertions(+) create mode 100644 tests/test_parameter_kernel_overrides.py diff --git a/tests/test_parameter_kernel_overrides.py b/tests/test_parameter_kernel_overrides.py new file mode 100644 index 0000000000..d460196eb8 --- /dev/null +++ b/tests/test_parameter_kernel_overrides.py @@ -0,0 +1,161 @@ +"""Tests for parameter-specific kernel overrides.""" + +from contextlib import nullcontext +from functools import partial +from unittest.mock import Mock + +import pandas as pd +import pytest +from attrs import evolve +from gpytorch import kernels as gk +from gpytorch import means as gm +from pytest import param + +from baybe.kernels import MaternKernel, RBFKernel +from baybe.kernels.base import Kernel +from baybe.kernels.composite import AdditiveKernel, ProductKernel, ScaleKernel +from baybe.parameters import ( + CategoricalParameter, + CustomDiscreteParameter, + NumericalContinuousParameter, + NumericalDiscreteParameter, + TaskParameter, +) +from baybe.parameters.base import _iter_basic_kernels + +_continuous = partial(NumericalContinuousParameter, "x", (0, 1)) +_discrete = partial(NumericalDiscreteParameter, "x", (0, 1, 2)) +_categorical = partial(CategoricalParameter, "x", ("a", "b", "c")) +_custom = partial( + CustomDiscreteParameter, + "x", + pd.DataFrame({"d1": [1.0, 2.0], "d2": [3.0, 1.0]}, index=["a", "b"]), +) +_task = partial(TaskParameter, "x", ("a", "b")) + +_SCOPE_ERROR = "may only act on the parameter itself" +_ACTIVE_DIMS_ERROR = "must not specify 'active_dims'" +_TYPE_ERROR = "must be a BayBE or GPyTorch kernel" +_TRAVERSE_ERROR = "Cannot traverse kernel" + + +@pytest.mark.parametrize( + ("constructor", "override", "error", "match"), + [ + # Valid BayBE overrides + param(_continuous, RBFKernel(), None, None, id="baybe-unscoped"), + param( + _discrete, + RBFKernel(parameter_names=("x",)), + None, + None, + id="baybe-scoped-to-own-parameter", + ), + param(_categorical, ScaleKernel(RBFKernel()), None, None, id="baybe-scaled"), + param(_custom, RBFKernel() + MaternKernel(), None, None, id="baybe-sum"), + param( + _continuous, RBFKernel() * MaternKernel(), None, None, id="baybe-product" + ), + param( + _discrete, + ScaleKernel( + AdditiveKernel([RBFKernel(), ProductKernel([RBFKernel()] * 2)]) + ), + None, + None, + id="baybe-nested", + ), + # Valid GPyTorch overrides + param(_continuous, gk.RBFKernel(), None, None, id="gpytorch-plain"), + param( + _categorical, gk.RBFKernel(ard_num_dims=3), None, None, id="gpytorch-ard" + ), + param( + _discrete, gk.ScaleKernel(gk.RBFKernel()), None, None, id="gpytorch-nested" + ), + # BayBE overrides acting on other parameters + param( + _continuous, + RBFKernel(parameter_names=("other",)), + ValueError, + _SCOPE_ERROR, + id="baybe-foreign-scope", + ), + param( + _discrete, + RBFKernel(parameter_names=("x", "other")), + ValueError, + _SCOPE_ERROR, + id="baybe-multiple-names", + ), + param( + _categorical, + ScaleKernel(RBFKernel(parameter_names=("other",))), + ValueError, + _SCOPE_ERROR, + id="baybe-nested-foreign-scope", + ), + # GPyTorch overrides with explicit active dimensions + param( + _continuous, + gk.RBFKernel(active_dims=(0,)), + ValueError, + _ACTIVE_DIMS_ERROR, + id="gpytorch-active-dims", + ), + param( + _custom, + gk.ScaleKernel(gk.RBFKernel(active_dims=(0,))), + ValueError, + _ACTIVE_DIMS_ERROR, + id="gpytorch-nested-active-dims", + ), + # Objects that are not supported kernels + param(_continuous, 1, TypeError, _TYPE_ERROR, id="no-kernel"), + param(_discrete, gm.ConstantMean(), TypeError, _TYPE_ERROR, id="gpytorch-mean"), + param( + _continuous, + Mock(spec=Kernel), + TypeError, + _TRAVERSE_ERROR, + id="unknown-kernel-type", + ), + param( + _categorical, + ScaleKernel(Mock(spec=Kernel)), + TypeError, + _TRAVERSE_ERROR, + id="unknown-nested-kernel-type", + ), + # Task parameters derive their override from the transfer learning mode + param( + _task, + RBFKernel(), + TypeError, + "unexpected keyword argument 'override_kernel'", + id="task-parameter", + ), + ], +) +def test_override_assignment(constructor, override, error, match): + """Assigning a kernel override to a parameter succeeds or raises as expected.""" + with nullcontext() if error is None else pytest.raises(error, match=match): + parameter = constructor(override_kernel=override) + + # Invalid overrides: nothing more to check + if error is not None: + return + + # GPyTorch overrides are returned unchanged + if not isinstance(override, Kernel): + assert parameter.override_kernel is override + return + + # BayBE overrides are scoped to their parameter, also after renaming + renamed = evolve(parameter, name="renamed") + for p in (parameter, renamed): + kernels = _iter_basic_kernels(p.override_kernel) + assert all(k.parameter_names == (p.name,) for k in kernels) + + # Unscoped and owner-scoped input yield the same parameter + assert parameter == constructor(override_kernel=override._scope_to_parameter("x")) From e078e97f487e30464cf3586d45898b5c6a7acb47 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Thu, 15 Oct 2026 12:21:18 +0200 Subject: [PATCH 65/73] Add parameter equivalence test --- tests/test_parameter_kernel_overrides.py | 179 +++++++++++++++++++++++ 1 file changed, 179 insertions(+) diff --git a/tests/test_parameter_kernel_overrides.py b/tests/test_parameter_kernel_overrides.py index d460196eb8..f93bd03593 100644 --- a/tests/test_parameter_kernel_overrides.py +++ b/tests/test_parameter_kernel_overrides.py @@ -9,6 +9,8 @@ from attrs import evolve from gpytorch import kernels as gk from gpytorch import means as gm +from gpytorch.constraints import Interval +from gpytorch.priors import GammaPrior as GPyTorchGammaPrior from pytest import param from baybe.kernels import MaternKernel, RBFKernel @@ -22,6 +24,8 @@ TaskParameter, ) from baybe.parameters.base import _iter_basic_kernels +from baybe.parameters.enum import TransferLearningMode +from baybe.priors import GammaPrior _continuous = partial(NumericalContinuousParameter, "x", (0, 1)) _discrete = partial(NumericalDiscreteParameter, "x", (0, 1, 2)) @@ -38,6 +42,35 @@ _TYPE_ERROR = "must be a BayBE or GPyTorch kernel" _TRAVERSE_ERROR = "Cannot traverse kernel" +_INDEX = TransferLearningMode.INDEX_KERNEL +_POSITIVE_INDEX = TransferLearningMode.POSITIVE_INDEX_KERNEL + + +def _rbf(lengthscale, *, frozen=False): + """Create a GPyTorch RBF kernel with a given, optionally frozen, lengthscale.""" + kernel = gk.RBFKernel() + kernel.lengthscale = lengthscale + kernel.raw_lengthscale.requires_grad_(not frozen) + return kernel + + +def _pair(override, other_override, other_bounds=(0, 1)): + """Create two differently named continuous parameters with given overrides.""" + return ( + NumericalContinuousParameter("p1", (0, 1), override_kernel=override), + NumericalContinuousParameter( + "p2", other_bounds, override_kernel=other_override + ), + ) + + +def _task_pair(values, mode, other_values, other_mode): + """Create two differently named task parameters.""" + return ( + TaskParameter("t1", values, override_transfer_learning_mode=mode), + TaskParameter("t2", other_values, override_transfer_learning_mode=other_mode), + ) + @pytest.mark.parametrize( ("constructor", "override", "error", "match"), @@ -159,3 +192,149 @@ def test_override_assignment(constructor, override, error, match): # Unscoped and owner-scoped input yield the same parameter assert parameter == constructor(override_kernel=override._scope_to_parameter("x")) + + +_SAME_GPYTORCH_KERNEL = gk.RBFKernel() + + +@pytest.mark.parametrize( + ("p1", "p2", "expected"), + [ + # Equivalent: only the name differs + param(*_pair(None, None), True, id="no-override"), + param(*_pair(RBFKernel(), RBFKernel()), True, id="baybe"), + param( + *_pair(RBFKernel(), RBFKernel(parameter_names=("p2",))), + True, + id="baybe-unscoped-vs-scoped", + ), + param( + *_pair( + ScaleKernel(RBFKernel() + MaternKernel()), + ScaleKernel(RBFKernel() + MaternKernel()), + ), + True, + id="baybe-composite", + ), + param(*_pair(gk.MaternKernel(), gk.MaternKernel()), True, id="gpytorch"), + param( + *_pair( + gk.ScaleKernel( + gk.RBFKernel(lengthscale_prior=GPyTorchGammaPrior(3.0, 6.0)) + ), + gk.ScaleKernel( + gk.RBFKernel(lengthscale_prior=GPyTorchGammaPrior(3.0, 6.0)) + ), + ), + True, + id="gpytorch-nested", + ), + param( + *_pair(gk.IndexKernel(num_tasks=2), gk.IndexKernel(num_tasks=2)), + True, + id="gpytorch-random-init", + ), + param(*_pair(_rbf(2.0), _rbf(3.0)), True, id="gpytorch-initial-value"), + param( + *_pair(_rbf(2.0, frozen=True), _rbf(2.0, frozen=True)), + True, + id="gpytorch-same-frozen-value", + ), + param( + *_pair(_SAME_GPYTORCH_KERNEL, _SAME_GPYTORCH_KERNEL), + True, + id="gpytorch-same-instance", + ), + param( + *_task_pair(("a", "b"), _INDEX, ("a", "b"), _INDEX), + True, + id="task-same-mode", + ), + param(*_task_pair(("a", "b"), None, ("a", "b"), None), True, id="task-no-mode"), + # Not equivalent: specification differs + param( + NumericalContinuousParameter("p1", (0, 1), override_kernel=RBFKernel()), + NumericalDiscreteParameter("p2", (0, 1), override_kernel=RBFKernel()), + False, + id="parameter-type", + ), + param( + *_pair(RBFKernel(), RBFKernel(), other_bounds=(0, 2)), + False, + id="parameter-spec", + ), + param(*_pair(RBFKernel(), None), False, id="override-vs-none"), + param(*_pair(RBFKernel(), MaternKernel()), False, id="baybe-class"), + param( + *_pair(RBFKernel(), RBFKernel(lengthscale_prior=GammaPrior(3.0, 6.0))), + False, + id="baybe-setting", + ), + param(*_pair(RBFKernel(), gk.RBFKernel()), False, id="baybe-vs-gpytorch"), + param(*_pair(gk.RBFKernel(), gk.MaternKernel()), False, id="gpytorch-class"), + param( + *_pair(gk.MaternKernel(nu=0.5), gk.MaternKernel()), + False, + id="gpytorch-attribute", + ), + param( + *_pair( + gk.RBFKernel(lengthscale_prior=GPyTorchGammaPrior(3.0, 6.0)), + gk.RBFKernel(lengthscale_prior=GPyTorchGammaPrior(2.0, 6.0)), + ), + False, + id="gpytorch-prior", + ), + param( + *_pair( + gk.RBFKernel(lengthscale_constraint=Interval(0.1, 2.0)), + gk.RBFKernel(lengthscale_constraint=Interval(0.1, 3.0)), + ), + False, + id="gpytorch-constraint", + ), + param( + *_pair(_rbf(2.0, frozen=True), _rbf(3.0, frozen=True)), + False, + id="gpytorch-frozen-value", + ), + param( + *_pair(_rbf(2.0, frozen=True), _rbf(2.0)), + False, + id="gpytorch-frozen-vs-learnable", + ), + param( + *_pair( + gk.IndexKernel(num_tasks=2, rank=1), gk.IndexKernel(num_tasks=2, rank=2) + ), + False, + id="gpytorch-shape", + ), + param( + *_pair( + gk.RBFKernel() + gk.MaternKernel(), gk.MaternKernel() + gk.RBFKernel() + ), + False, + id="gpytorch-summand-order", + ), + param( + *_task_pair(("a", "b"), _INDEX, ("a", "b"), _POSITIVE_INDEX), + False, + id="task-mode", + ), + param( + *_task_pair(("a", "b"), None, ("a", "b", "c"), None), + False, + id="task-values", + ), + param( + *_task_pair(("a", "b"), _INDEX, ("a", "b"), None), + False, + id="task-mode-vs-none", + ), + ], +) +def test_parameter_equivalence(p1, p2, expected): + """Parameter equivalence with kernel overrides produces the expected result.""" + assert p1.is_equivalent(p2) == expected + assert p2.is_equivalent(p1) == expected From 095bcd66d70a8a1cc748dfccde58f4fb1b818149 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Thu, 15 Oct 2026 12:21:59 +0200 Subject: [PATCH 66/73] Make kernel override strategy scope configurable --- tests/hypothesis_strategies/parameters.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/tests/hypothesis_strategies/parameters.py b/tests/hypothesis_strategies/parameters.py index 300904d980..4d461aa950 100644 --- a/tests/hypothesis_strategies/parameters.py +++ b/tests/hypothesis_strategies/parameters.py @@ -39,12 +39,24 @@ @st.composite -def kernel_overrides(draw: st.DrawFn, parameter_name: str): - """Generate valid kernel overrides for a parameter.""" +def kernel_overrides(draw: st.DrawFn, parameter_name: str, scoped: bool | None = None): + """Generate valid kernel overrides for a parameter. + + Args: + draw: Hypothesis draw object. + parameter_name: The name of the parameter the override belongs to. + scoped: Whether the kernel is scoped to the parameter (``True``), unscoped + (``False``), or either (``None``). + + Returns: + A kernel override or ``None``. + """ from tests.hypothesis_strategies.kernels import kernels kernel = draw(st.one_of(st.none(), kernels(parameter_names=(parameter_name,)))) - if kernel is not None and draw(st.booleans()): + if scoped is None: + scoped = draw(st.booleans()) + if kernel is not None and not scoped: kernel = kernel._scope_to_parameter(None) return kernel From e92fbb632da02b5bf41bc2e15a2b0539aedea31b Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Thu, 15 Oct 2026 12:24:34 +0200 Subject: [PATCH 67/73] Add kernel override construction test --- tests/test_parameter_kernel_overrides.py | 359 +++++++++++++++++++++++ 1 file changed, 359 insertions(+) diff --git a/tests/test_parameter_kernel_overrides.py b/tests/test_parameter_kernel_overrides.py index f93bd03593..6cfc5c3522 100644 --- a/tests/test_parameter_kernel_overrides.py +++ b/tests/test_parameter_kernel_overrides.py @@ -6,15 +6,21 @@ import pandas as pd import pytest +import torch from attrs import evolve +from botorch.models.kernels.positive_index import ( + PositiveIndexKernel as BoPositiveIndexKernel, +) from gpytorch import kernels as gk from gpytorch import means as gm from gpytorch.constraints import Interval from gpytorch.priors import GammaPrior as GPyTorchGammaPrior from pytest import param +from baybe.exceptions import IncompatibleOverrideError from baybe.kernels import MaternKernel, RBFKernel from baybe.kernels.base import Kernel +from baybe.kernels.basic import IndexKernel from baybe.kernels.composite import AdditiveKernel, ProductKernel, ScaleKernel from baybe.parameters import ( CategoricalParameter, @@ -25,7 +31,14 @@ ) from baybe.parameters.base import _iter_basic_kernels from baybe.parameters.enum import TransferLearningMode +from baybe.parameters.selectors import NameSelector from baybe.priors import GammaPrior +from baybe.searchspace import SearchSpace +from baybe.surrogates import GaussianProcessSurrogate +from baybe.surrogates.gaussian_process.components.kernel import ICMKernelFactory +from baybe.surrogates.gaussian_process.core import _ModelContext +from baybe.surrogates.gaussian_process.presets import BayBEKernelFactory +from baybe.targets import NumericalTarget _continuous = partial(NumericalContinuousParameter, "x", (0, 1)) _discrete = partial(NumericalDiscreteParameter, "x", (0, 1, 2)) @@ -72,6 +85,79 @@ def _task_pair(values, mode, other_values, other_mode): ) +def _x(name, override=None): + """Create a continuous parameter on [0, 1], optionally with a kernel override.""" + return NumericalContinuousParameter(name, (0, 1), override_kernel=override) + + +def _tl_task(mode=None): + """Create a three-valued task parameter, optionally with a TL mode.""" + return TaskParameter("task", ("a", "b", "c"), override_transfer_learning_mode=mode) + + +def _basic_kernels(kernel, columns): + """Yield each basic kernel of a GPyTorch kernel with the columns it acts on. + + Basic kernels compute on columns (e.g. RBF, Matern), as opposed to wrappers such + as Scale, Additive, or Product kernels. As in GPyTorch, ``active_dims`` of a + kernel apply to everything nested inside it. + """ + if kernel.active_dims is not None: + columns = set(kernel.active_dims.tolist()) + children = [c for c in kernel.children() if isinstance(c, gk.Kernel)] + children += [ + m for c in kernel.children() if isinstance(c, torch.nn.ModuleList) for m in c + ] + if not children: + yield kernel, columns + for child in children: + yield from _basic_kernels(child, columns) + + +def _factors_by_parameter(kernel, searchspace): + """Map the parameters each basic kernel acts on to the basic kernel classes. + + Example: ``{("x1", "x3"): [MaternKernel], ("x2",): [RBFKernel]}`` + """ + factors: dict[tuple[str, ...], list[type]] = {} + for k, names, _ in _basic_kernels_by_parameter(kernel, searchspace): + factors.setdefault(names, []).append(type(k)) + return factors + + +def _basic_kernels_by_parameter(kernel, searchspace): + """Yield each basic kernel with the parameters and columns it acts on.""" + all_columns = set(range(len(searchspace.comp_rep_columns))) + for k, columns in _basic_kernels(kernel, all_columns): + names = tuple( + p.name + for p in searchspace.parameters + if columns & set(searchspace.get_comp_rep_parameter_indices(p.name)) + ) + yield k, names, columns + + +def _matern_naming_all(searchspace, objective, measurements): + """Return a Matern kernel naming all parameters of the given search space.""" + return MaternKernel(parameter_names=searchspace.parameter_names) + + +def _unnamed_matern(searchspace, objective, measurements): + """Return a Matern kernel without parameter names.""" + return MaternKernel() + + +def _raw_matern(searchspace, objective, measurements): + """Return a raw GPyTorch kernel, which cannot exclude overridden parameters.""" + return gk.MaternKernel() + + +def _needs_full_searchspace(searchspace, objective, measurements): + """Return a kernel after accessing information unavailable on reduced spaces.""" + searchspace.task_idx + return MaternKernel() + + @pytest.mark.parametrize( ("constructor", "override", "error", "match"), [ @@ -338,3 +424,276 @@ def test_parameter_equivalence(p1, p2, expected): """Parameter equivalence with kernel overrides produces the expected result.""" assert p1.is_equivalent(p2) == expected assert p2.is_equivalent(p1) == expected + + +_X2_RBF = [_x("x1"), _x("x2", RBFKernel()), _x("x3")] +_X1_SELECTOR = BayBEKernelFactory(parameter_selector=NameSelector(("x1",), regex=False)) +_TASK_KERNEL = IndexKernel(num_tasks=3, rank=3, parameter_names=("task",)) +_EXCLUDE_ERROR = "can exclude these parameters" +_ARD_ERROR = "specifies 2 ARD dimensions, but the parameter has 3" + + +@pytest.mark.parametrize( + ("parameters", "kernel_or_factory", "expected"), + [ + # No overrides: the surrogate kernel is used unchanged + param( + [_x("x1"), _x("x2")], + None, + {("x1", "x2"): [gk.MaternKernel]}, + id="no-override-default", + ), + param( + [_x("x1"), _x("x2")], + gk.MaternKernel(), + {("x1", "x2"): [gk.MaternKernel]}, + id="no-override-gpytorch", + ), + # One override, the surrogate kernel acts on the remaining parameters + param( + _X2_RBF, + None, + {("x1", "x3"): [gk.MaternKernel], ("x2",): [gk.RBFKernel]}, + id="default", + ), + param( + _X2_RBF, + MaternKernel(), + {("x1", "x3"): [gk.MaternKernel], ("x2",): [gk.RBFKernel]}, + id="fixed", + ), + param( + _X2_RBF, + MaternKernel(parameter_names=("x1", "x2")), + {("x1",): [gk.MaternKernel], ("x2",): [gk.RBFKernel]}, + id="fixed-naming-overridden-parameter", + ), + param( + _X2_RBF, + ScaleKernel(MaternKernel()), + {("x1", "x3"): [gk.MaternKernel], ("x2",): [gk.RBFKernel]}, + id="scaled", + ), + param( + _X2_RBF, + MaternKernel() + RBFKernel(), + {("x1", "x3"): [gk.MaternKernel, gk.RBFKernel], ("x2",): [gk.RBFKernel]}, + id="sum", + ), + param( + _X2_RBF, + MaternKernel() * RBFKernel(), + {("x1", "x3"): [gk.MaternKernel, gk.RBFKernel], ("x2",): [gk.RBFKernel]}, + id="product", + ), + param( + _X2_RBF, + _matern_naming_all, + {("x1", "x3"): [gk.MaternKernel], ("x2",): [gk.RBFKernel]}, + id="factory", + ), + param( + _X2_RBF, + _X1_SELECTOR, + {("x1",): [gk.MaternKernel], ("x2",): [gk.RBFKernel]}, + id="selector", + ), + # Override shapes + param( + [ + _x("x1"), + CategoricalParameter("c", ("a", "b", "c"), override_kernel=RBFKernel()), + ], + None, + {("x1",): [gk.MaternKernel], ("c",): [gk.RBFKernel]}, + id="multi-column-parameter", + ), + param( + [_x("x1"), _x("x2", RBFKernel()), _x("x3", MaternKernel())], + None, + { + ("x1",): [gk.MaternKernel], + ("x2",): [gk.RBFKernel], + ("x3",): [gk.MaternKernel], + }, + id="multiple-overrides", + ), + param( + [_x("x1", RBFKernel()), _x("x2", MaternKernel())], + MaternKernel(), + {("x1",): [gk.RBFKernel], ("x2",): [gk.MaternKernel]}, + id="all-overridden", + ), + param( + [_x("x1"), _x("x2", gk.RBFKernel())], + None, + {("x1",): [gk.MaternKernel], ("x2",): [gk.RBFKernel]}, + id="gpytorch-override", + ), + param( + [ + _x("x1"), + CategoricalParameter( + "c", ("a", "b", "c"), override_kernel=gk.RBFKernel(ard_num_dims=3) + ), + ], + None, + {("x1",): [gk.MaternKernel], ("c",): [gk.RBFKernel]}, + id="gpytorch-override-ard", + ), + param( + [_x("x1"), _x("x2", gk.ScaleKernel(gk.RBFKernel()))], + None, + {("x1",): [gk.MaternKernel], ("x2",): [gk.RBFKernel]}, + id="gpytorch-override-nested", + ), + # Task parameters: the transfer learning mode defines the task kernel + param( + [_x("x1"), _tl_task()], + None, + {("x1",): [gk.MaternKernel], ("task",): [BoPositiveIndexKernel]}, + id="task-without-mode", + ), + param( + [_x("x1"), _tl_task(_INDEX)], + None, + {("x1",): [gk.MaternKernel], ("task",): [gk.IndexKernel]}, + id="task-index", + ), + param( + [_x("x1"), _tl_task(_POSITIVE_INDEX)], + None, + {("x1",): [gk.MaternKernel], ("task",): [BoPositiveIndexKernel]}, + id="task-positive-index", + ), + param( + [_x("x1"), _tl_task(_INDEX)], + MaternKernel(parameter_names=("x1", "task")), + {("x1",): [gk.MaternKernel], ("task",): [gk.IndexKernel]}, + id="task-index+kernel-naming-task", + ), + param( + [_x("x1"), _tl_task(_INDEX)], + MaternKernel(parameter_names=("x1",)) * _TASK_KERNEL, + {("x1",): [gk.MaternKernel], ("task",): [gk.IndexKernel]}, + id="task-index+product-with-task-kernel", + ), + param( + [_x("x1"), _tl_task(_INDEX)], + _TASK_KERNEL, + {("task",): [gk.IndexKernel]}, + id="task-index+task-only-kernel", + ), + param( + [_x("x1"), _tl_task(_INDEX)], + _unnamed_matern, + {("x1",): [gk.MaternKernel], ("task",): [gk.IndexKernel]}, + id="task-index+factory", + ), + param( + [_x("x1"), _x("x2", RBFKernel()), _tl_task(_POSITIVE_INDEX)], + None, + { + ("x1",): [gk.MaternKernel], + ("task",): [BoPositiveIndexKernel], + ("x2",): [gk.RBFKernel], + }, + id="task-positive-index+parameter-override", + ), + param( + [_x("x1"), _x("x2", RBFKernel()), _tl_task()], + _X1_SELECTOR, + { + ("x1",): [gk.MaternKernel], + ("task",): [BoPositiveIndexKernel], + ("x2",): [gk.RBFKernel], + }, + id="task-without-mode+selector", + ), + # Invalid configurations detected at kernel construction + param( + [ + _x("x1"), + CategoricalParameter( + "c", ("a", "b", "c"), override_kernel=gk.RBFKernel(ard_num_dims=2) + ), + ], + None, + (IncompatibleOverrideError, _ARD_ERROR), + id="gpytorch-ard-mismatch", + ), + param( + [ + _x("x1"), + CategoricalParameter( + "c", + ("a", "b", "c"), + override_kernel=gk.ScaleKernel(gk.RBFKernel(ard_num_dims=2)), + ), + ], + None, + (IncompatibleOverrideError, _ARD_ERROR), + id="gpytorch-nested-ard-mismatch", + ), + param( + _X2_RBF, + gk.MaternKernel(), + (IncompatibleOverrideError, _EXCLUDE_ERROR), + id="gpytorch-surrogate-kernel", + ), + param( + _X2_RBF, + _raw_matern, + (IncompatibleOverrideError, _EXCLUDE_ERROR), + id="factory-returning-gpytorch-kernel", + ), + param( + _X2_RBF, + _needs_full_searchspace, + (IncompatibleOverrideError, _EXCLUDE_ERROR), + id="factory-needing-full-searchspace", + ), + param( + [_x("x1"), _tl_task(_INDEX)], + ICMKernelFactory(task_kernel_or_factory=_TASK_KERNEL), + (IncompatibleOverrideError, _EXCLUDE_ERROR), + id="task-index+icm-factory", + ), + ], +) +def test_override_kernel_construction(parameters, kernel_or_factory, expected): + """Kernel construction with overrides yields the expected kernel or error.""" + searchspace = SearchSpace.from_product(parameters) + context = _ModelContext( + searchspace, NumericalTarget("y").to_objective(), pd.DataFrame() + ) + surrogate = GaussianProcessSurrogate(kernel_or_factory=kernel_or_factory) + + # Invalid configurations raise during construction + if isinstance(expected, tuple): + error, match = expected + with pytest.raises(error, match=match): + surrogate._resolve_kernel(context) + return + + # Each parameter is covered by the expected kernel classes + kernel = surrogate._resolve_kernel(context) + assert _factors_by_parameter(kernel, searchspace) == expected + + # User-provided GPyTorch overrides are copied, leaving the input unbound + gpytorch_overrides = { + p.name: p.override_kernel + for p in parameters + if p.override_kernel is not None and not isinstance(p.override_kernel, Kernel) + } + for override in gpytorch_overrides.values(): + assert all(k.active_dims is None for k in (override, *override.sub_kernels())) + + # BayBE-created kernels use one ARD dimension per column, user-provided GPyTorch + # kernels keep their setting + user_kernel = isinstance(kernel_or_factory, gk.Kernel) + for k, names, columns in _basic_kernels_by_parameter(kernel, searchspace): + if user_kernel or set(names) & gpytorch_overrides.keys(): + assert k.ard_num_dims in (None, len(columns)) + else: + assert k.ard_num_dims == len(columns) From de72cadb4886018809b844a83e5c7b6cafc63d4b Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Thu, 15 Oct 2026 12:25:52 +0200 Subject: [PATCH 68/73] Add fitted kernel override test --- tests/test_parameter_kernel_overrides.py | 83 ++++++++++++++++++++++++ 1 file changed, 83 insertions(+) diff --git a/tests/test_parameter_kernel_overrides.py b/tests/test_parameter_kernel_overrides.py index 6cfc5c3522..9f2f183401 100644 --- a/tests/test_parameter_kernel_overrides.py +++ b/tests/test_parameter_kernel_overrides.py @@ -1,6 +1,7 @@ """Tests for parameter-specific kernel overrides.""" from contextlib import nullcontext +from copy import deepcopy from functools import partial from unittest.mock import Mock @@ -39,6 +40,7 @@ from baybe.surrogates.gaussian_process.core import _ModelContext from baybe.surrogates.gaussian_process.presets import BayBEKernelFactory from baybe.targets import NumericalTarget +from baybe.utils.dataframe import create_fake_input _continuous = partial(NumericalContinuousParameter, "x", (0, 1)) _discrete = partial(NumericalDiscreteParameter, "x", (0, 1, 2)) @@ -697,3 +699,84 @@ def test_override_kernel_construction(parameters, kernel_or_factory, expected): assert k.ard_num_dims in (None, len(columns)) else: assert k.ard_num_dims == len(columns) + + +def _rbf_with_prior(): + """Create a regularized BayBE RBF kernel for numerically stable fits.""" + return RBFKernel(lengthscale_prior=GammaPrior(3.0, 6.0)) + + +@pytest.mark.parametrize( + ("parameters", "expected"), + [ + param( + [_x("x1"), _x("x2", _rbf_with_prior())], + {("x1",): [gk.MaternKernel], ("x2",): [gk.RBFKernel]}, + id="parameter-override", + ), + param( + [ + _x("x1"), + CategoricalParameter( + "c", ("a", "b", "c"), override_kernel=_rbf_with_prior() + ), + TaskParameter( + "task", ("a", "b"), override_transfer_learning_mode=_INDEX + ), + ], + { + ("x1",): [gk.MaternKernel], + ("c",): [gk.RBFKernel], + ("task",): [gk.IndexKernel], + }, + id="task-index+parameter-override", + ), + param( + [ + _x("x1", _rbf_with_prior()), + _x("x2", MaternKernel(lengthscale_prior=GammaPrior(3.0, 6.0))), + ], + {("x1",): [gk.RBFKernel], ("x2",): [gk.MaternKernel]}, + id="all-overridden", + ), + param( + [ + _x("x1"), + _x( + "x2", + gk.RBFKernel(lengthscale_prior=GPyTorchGammaPrior(3.0, 6.0)), + ), + ], + {("x1",): [gk.MaternKernel], ("x2",): [gk.RBFKernel]}, + id="gpytorch-override", + ), + ], +) +def test_fitted_override_kernel(parameters, expected): + """The fitted model uses the kernel constructed from the overrides.""" + searchspace = SearchSpace.from_product(parameters) + objective = NumericalTarget("y").to_objective() + measurements = create_fake_input(searchspace.parameters, objective.targets, 6) + + # Snapshot user-provided GPyTorch overrides to detect mutations by training + gpytorch_overrides = [ + p.override_kernel + for p in parameters + if p.override_kernel is not None and not isinstance(p.override_kernel, Kernel) + ] + snapshots = [deepcopy(o.state_dict()) for o in gpytorch_overrides] + + surrogate = GaussianProcessSurrogate() + surrogate.fit(searchspace, objective, measurements) + + # The model was fitted on data using the expected kernel structure + assert surrogate._model is not None + kernel = surrogate._model.covar_module + assert _factors_by_parameter(kernel, searchspace) == expected + + # Training did not touch the user-provided GPyTorch overrides + for override, snapshot in zip(gpytorch_overrides, snapshots): + assert override.active_dims is None + assert all( + torch.equal(v, override.state_dict()[k]) for k, v in snapshot.items() + ) From c8746be8a0a73299cd4da79f4476b25aeaad8359 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Thu, 15 Oct 2026 12:26:25 +0200 Subject: [PATCH 69/73] Extend multitask kernel deprecation test with transfer learning modes --- tests/test_deprecations.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/tests/test_deprecations.py b/tests/test_deprecations.py index 662581b0b9..b7228594ae 100644 --- a/tests/test_deprecations.py +++ b/tests/test_deprecations.py @@ -32,7 +32,7 @@ from baybe.objectives.desirability import DesirabilityObjective from baybe.objectives.single import SingleTargetObjective from baybe.parameters.categorical import TaskParameter -from baybe.parameters.enum import SubstanceEncoding +from baybe.parameters.enum import SubstanceEncoding, TransferLearningMode from baybe.parameters.numerical import ( NumericalDiscreteParameter, ) @@ -543,15 +543,23 @@ def test_deprecated_cache_environment_variables(monkeypatch, value: str, expecte assert Settings(restore_environment=True).cache_directory == expected +@pytest.mark.parametrize( + "mode", [None, TransferLearningMode.INDEX_KERNEL], ids=["no_mode", "mode"] +) @pytest.mark.parametrize("custom", [False, True], ids=["default", "custom"]) @pytest.mark.parametrize("env", [False, True], ids=["no_env", "env"]) @pytest.mark.parametrize("task", [False, True], ids=["no_task", "task"]) -def test_multitask_kernel_deprecation(monkeypatch, custom: bool, env: bool, task: bool): +def test_multitask_kernel_deprecation( + monkeypatch, custom: bool, env: bool, task: bool, mode: TransferLearningMode | None +): """Providing a custom kernel in a transfer learning context raises a deprecation - error unless explicitly disabled via environment variable.""" # noqa + error unless explicitly disabled via environment variable or a transfer learning + mode defines the task kernel.""" # noqa parameters = [NumericalDiscreteParameter("p", [0, 1])] if task: - parameters.append(TaskParameter("task", ["a", "b"])) + parameters.append( + TaskParameter("task", ["a", "b"], override_transfer_learning_mode=mode) + ) searchspace = SearchSpace.from_product(parameters) objective = NumericalTarget("t").to_objective() measurements = create_fake_input( @@ -564,7 +572,7 @@ def test_multitask_kernel_deprecation(monkeypatch, custom: bool, env: bool, task context = ( pytest.raises(DeprecationError) - if task and custom and not env + if task and custom and not env and mode is None else nullcontext() ) with context: From 0c10d4ef6a537c35aa0d929ea3c98ee56745efac Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Thu, 15 Oct 2026 12:27:19 +0200 Subject: [PATCH 70/73] Add kernel override iteration test --- tests/test_iterations.py | 38 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/tests/test_iterations.py b/tests/test_iterations.py index c3491f46bc..5b5e9b0715 100644 --- a/tests/test_iterations.py +++ b/tests/test_iterations.py @@ -27,6 +27,11 @@ ) from baybe.kernels.composite import ScaleKernel from baybe.objectives.pareto import ParetoObjective +from baybe.parameters import ( + NumericalContinuousParameter, + TaskParameter, + TransferLearningMode, +) from baybe.priors import ( GammaPrior, HalfCauchyPrior, @@ -398,6 +403,39 @@ def test_kernels(ongoing_campaign, n_iterations, batch_size): run_iterations(ongoing_campaign, n_iterations, batch_size) +def _make_override_parameters(mode: TransferLearningMode | None = None) -> list: + """Create parameters with a regularized kernel override and an optional task.""" + parameters = [ + NumericalContinuousParameter("x1", (0, 1)), + NumericalContinuousParameter( + "x2", + (0, 1), + override_kernel=RBFKernel(lengthscale_prior=GammaPrior(3.0, 6.0)), + ), + ] + if mode is not None: + parameters.append( + TaskParameter("task", ["a", "b"], override_transfer_learning_mode=mode) + ) + return parameters + + +@pytest.mark.slow +@pytest.mark.parametrize( + "parameters", + [ + param(_make_override_parameters(), id="no_task"), + *( + param(_make_override_parameters(mode), id=mode.name) + for mode in TransferLearningMode + ), + ], +) +def test_kernel_override_iteration(ongoing_campaign, n_iterations, batch_size): + """A full optimization loop runs with kernel overrides.""" + run_iterations(ongoing_campaign, n_iterations, batch_size) + + @pytest.mark.parametrize("kernel", valid_kernel_factories) @pytest.mark.parametrize("n_iterations", [3], ids=["i3"]) def test_kernel_factories(ongoing_campaign, n_iterations, batch_size): From 7a0fbb8e69ed13385b6fd22b98da440f3d05cb5d Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Thu, 15 Oct 2026 13:15:07 +0200 Subject: [PATCH 71/73] Move parameter equivalence test to its own file --- tests/test_parameter_equivalence.py | 192 +++++++++++++++++++++++ tests/test_parameter_kernel_overrides.py | 173 -------------------- 2 files changed, 192 insertions(+), 173 deletions(-) create mode 100644 tests/test_parameter_equivalence.py diff --git a/tests/test_parameter_equivalence.py b/tests/test_parameter_equivalence.py new file mode 100644 index 0000000000..d53b518893 --- /dev/null +++ b/tests/test_parameter_equivalence.py @@ -0,0 +1,192 @@ +"""Tests for parameter equivalence.""" + +import pytest +from gpytorch import kernels as gk +from gpytorch.constraints import Interval +from gpytorch.priors import GammaPrior as GPyTorchGammaPrior +from pytest import param + +from baybe.kernels import MaternKernel, RBFKernel +from baybe.kernels.composite import ScaleKernel +from baybe.parameters import ( + NumericalContinuousParameter, + NumericalDiscreteParameter, + TaskParameter, +) +from baybe.parameters.enum import TransferLearningMode +from baybe.priors import GammaPrior + +_INDEX = TransferLearningMode.INDEX_KERNEL +_POSITIVE_INDEX = TransferLearningMode.POSITIVE_INDEX_KERNEL + + +def _rbf(lengthscale, *, frozen=False): + """Create a GPyTorch RBF kernel with a given, optionally frozen, lengthscale.""" + kernel = gk.RBFKernel() + kernel.lengthscale = lengthscale + kernel.raw_lengthscale.requires_grad_(not frozen) + return kernel + + +def _pair(override, other_override, other_bounds=(0, 1)): + """Create two differently named continuous parameters with given overrides.""" + return ( + NumericalContinuousParameter("p1", (0, 1), override_kernel=override), + NumericalContinuousParameter( + "p2", other_bounds, override_kernel=other_override + ), + ) + + +def _task_pair(values, mode, other_values, other_mode): + """Create two differently named task parameters.""" + return ( + TaskParameter("t1", values, override_transfer_learning_mode=mode), + TaskParameter("t2", other_values, override_transfer_learning_mode=other_mode), + ) + + +_SAME_GPYTORCH_KERNEL = gk.RBFKernel() + + +@pytest.mark.parametrize( + ("p1", "p2", "expected"), + [ + # Equivalent: only the name differs + param(*_pair(None, None), True, id="no-override"), + param(*_pair(RBFKernel(), RBFKernel()), True, id="baybe"), + param( + *_pair(RBFKernel(), RBFKernel(parameter_names=("p2",))), + True, + id="baybe-unscoped-vs-scoped", + ), + param( + *_pair( + ScaleKernel(RBFKernel() + MaternKernel()), + ScaleKernel(RBFKernel() + MaternKernel()), + ), + True, + id="baybe-composite", + ), + param(*_pair(gk.MaternKernel(), gk.MaternKernel()), True, id="gpytorch"), + param( + *_pair( + gk.ScaleKernel( + gk.RBFKernel(lengthscale_prior=GPyTorchGammaPrior(3.0, 6.0)) + ), + gk.ScaleKernel( + gk.RBFKernel(lengthscale_prior=GPyTorchGammaPrior(3.0, 6.0)) + ), + ), + True, + id="gpytorch-nested", + ), + param( + *_pair(gk.IndexKernel(num_tasks=2), gk.IndexKernel(num_tasks=2)), + True, + id="gpytorch-random-init", + ), + param(*_pair(_rbf(2.0), _rbf(3.0)), True, id="gpytorch-initial-value"), + param( + *_pair(_rbf(2.0, frozen=True), _rbf(2.0, frozen=True)), + True, + id="gpytorch-same-frozen-value", + ), + param( + *_pair(_SAME_GPYTORCH_KERNEL, _SAME_GPYTORCH_KERNEL), + True, + id="gpytorch-same-instance", + ), + param( + *_task_pair(("a", "b"), _INDEX, ("a", "b"), _INDEX), + True, + id="task-same-mode", + ), + param(*_task_pair(("a", "b"), None, ("a", "b"), None), True, id="task-no-mode"), + # Not equivalent: specification differs + param( + NumericalContinuousParameter("p1", (0, 1), override_kernel=RBFKernel()), + NumericalDiscreteParameter("p2", (0, 1), override_kernel=RBFKernel()), + False, + id="parameter-type", + ), + param( + *_pair(RBFKernel(), RBFKernel(), other_bounds=(0, 2)), + False, + id="parameter-spec", + ), + param(*_pair(RBFKernel(), None), False, id="override-vs-none"), + param(*_pair(RBFKernel(), MaternKernel()), False, id="baybe-class"), + param( + *_pair(RBFKernel(), RBFKernel(lengthscale_prior=GammaPrior(3.0, 6.0))), + False, + id="baybe-setting", + ), + param(*_pair(RBFKernel(), gk.RBFKernel()), False, id="baybe-vs-gpytorch"), + param(*_pair(gk.RBFKernel(), gk.MaternKernel()), False, id="gpytorch-class"), + param( + *_pair(gk.MaternKernel(nu=0.5), gk.MaternKernel()), + False, + id="gpytorch-attribute", + ), + param( + *_pair( + gk.RBFKernel(lengthscale_prior=GPyTorchGammaPrior(3.0, 6.0)), + gk.RBFKernel(lengthscale_prior=GPyTorchGammaPrior(2.0, 6.0)), + ), + False, + id="gpytorch-prior", + ), + param( + *_pair( + gk.RBFKernel(lengthscale_constraint=Interval(0.1, 2.0)), + gk.RBFKernel(lengthscale_constraint=Interval(0.1, 3.0)), + ), + False, + id="gpytorch-constraint", + ), + param( + *_pair(_rbf(2.0, frozen=True), _rbf(3.0, frozen=True)), + False, + id="gpytorch-frozen-value", + ), + param( + *_pair(_rbf(2.0, frozen=True), _rbf(2.0)), + False, + id="gpytorch-frozen-vs-learnable", + ), + param( + *_pair( + gk.IndexKernel(num_tasks=2, rank=1), gk.IndexKernel(num_tasks=2, rank=2) + ), + False, + id="gpytorch-shape", + ), + param( + *_pair( + gk.RBFKernel() + gk.MaternKernel(), gk.MaternKernel() + gk.RBFKernel() + ), + False, + id="gpytorch-summand-order", + ), + param( + *_task_pair(("a", "b"), _INDEX, ("a", "b"), _POSITIVE_INDEX), + False, + id="task-mode", + ), + param( + *_task_pair(("a", "b"), None, ("a", "b", "c"), None), + False, + id="task-values", + ), + param( + *_task_pair(("a", "b"), _INDEX, ("a", "b"), None), + False, + id="task-mode-vs-none", + ), + ], +) +def test_parameter_equivalence(p1, p2, expected): + """Parameter equivalence with kernel overrides produces the expected result.""" + assert p1.is_equivalent(p2) == expected + assert p2.is_equivalent(p1) == expected diff --git a/tests/test_parameter_kernel_overrides.py b/tests/test_parameter_kernel_overrides.py index 9f2f183401..54f9abb811 100644 --- a/tests/test_parameter_kernel_overrides.py +++ b/tests/test_parameter_kernel_overrides.py @@ -14,7 +14,6 @@ ) from gpytorch import kernels as gk from gpytorch import means as gm -from gpytorch.constraints import Interval from gpytorch.priors import GammaPrior as GPyTorchGammaPrior from pytest import param @@ -61,32 +60,6 @@ _POSITIVE_INDEX = TransferLearningMode.POSITIVE_INDEX_KERNEL -def _rbf(lengthscale, *, frozen=False): - """Create a GPyTorch RBF kernel with a given, optionally frozen, lengthscale.""" - kernel = gk.RBFKernel() - kernel.lengthscale = lengthscale - kernel.raw_lengthscale.requires_grad_(not frozen) - return kernel - - -def _pair(override, other_override, other_bounds=(0, 1)): - """Create two differently named continuous parameters with given overrides.""" - return ( - NumericalContinuousParameter("p1", (0, 1), override_kernel=override), - NumericalContinuousParameter( - "p2", other_bounds, override_kernel=other_override - ), - ) - - -def _task_pair(values, mode, other_values, other_mode): - """Create two differently named task parameters.""" - return ( - TaskParameter("t1", values, override_transfer_learning_mode=mode), - TaskParameter("t2", other_values, override_transfer_learning_mode=other_mode), - ) - - def _x(name, override=None): """Create a continuous parameter on [0, 1], optionally with a kernel override.""" return NumericalContinuousParameter(name, (0, 1), override_kernel=override) @@ -282,152 +255,6 @@ def test_override_assignment(constructor, override, error, match): assert parameter == constructor(override_kernel=override._scope_to_parameter("x")) -_SAME_GPYTORCH_KERNEL = gk.RBFKernel() - - -@pytest.mark.parametrize( - ("p1", "p2", "expected"), - [ - # Equivalent: only the name differs - param(*_pair(None, None), True, id="no-override"), - param(*_pair(RBFKernel(), RBFKernel()), True, id="baybe"), - param( - *_pair(RBFKernel(), RBFKernel(parameter_names=("p2",))), - True, - id="baybe-unscoped-vs-scoped", - ), - param( - *_pair( - ScaleKernel(RBFKernel() + MaternKernel()), - ScaleKernel(RBFKernel() + MaternKernel()), - ), - True, - id="baybe-composite", - ), - param(*_pair(gk.MaternKernel(), gk.MaternKernel()), True, id="gpytorch"), - param( - *_pair( - gk.ScaleKernel( - gk.RBFKernel(lengthscale_prior=GPyTorchGammaPrior(3.0, 6.0)) - ), - gk.ScaleKernel( - gk.RBFKernel(lengthscale_prior=GPyTorchGammaPrior(3.0, 6.0)) - ), - ), - True, - id="gpytorch-nested", - ), - param( - *_pair(gk.IndexKernel(num_tasks=2), gk.IndexKernel(num_tasks=2)), - True, - id="gpytorch-random-init", - ), - param(*_pair(_rbf(2.0), _rbf(3.0)), True, id="gpytorch-initial-value"), - param( - *_pair(_rbf(2.0, frozen=True), _rbf(2.0, frozen=True)), - True, - id="gpytorch-same-frozen-value", - ), - param( - *_pair(_SAME_GPYTORCH_KERNEL, _SAME_GPYTORCH_KERNEL), - True, - id="gpytorch-same-instance", - ), - param( - *_task_pair(("a", "b"), _INDEX, ("a", "b"), _INDEX), - True, - id="task-same-mode", - ), - param(*_task_pair(("a", "b"), None, ("a", "b"), None), True, id="task-no-mode"), - # Not equivalent: specification differs - param( - NumericalContinuousParameter("p1", (0, 1), override_kernel=RBFKernel()), - NumericalDiscreteParameter("p2", (0, 1), override_kernel=RBFKernel()), - False, - id="parameter-type", - ), - param( - *_pair(RBFKernel(), RBFKernel(), other_bounds=(0, 2)), - False, - id="parameter-spec", - ), - param(*_pair(RBFKernel(), None), False, id="override-vs-none"), - param(*_pair(RBFKernel(), MaternKernel()), False, id="baybe-class"), - param( - *_pair(RBFKernel(), RBFKernel(lengthscale_prior=GammaPrior(3.0, 6.0))), - False, - id="baybe-setting", - ), - param(*_pair(RBFKernel(), gk.RBFKernel()), False, id="baybe-vs-gpytorch"), - param(*_pair(gk.RBFKernel(), gk.MaternKernel()), False, id="gpytorch-class"), - param( - *_pair(gk.MaternKernel(nu=0.5), gk.MaternKernel()), - False, - id="gpytorch-attribute", - ), - param( - *_pair( - gk.RBFKernel(lengthscale_prior=GPyTorchGammaPrior(3.0, 6.0)), - gk.RBFKernel(lengthscale_prior=GPyTorchGammaPrior(2.0, 6.0)), - ), - False, - id="gpytorch-prior", - ), - param( - *_pair( - gk.RBFKernel(lengthscale_constraint=Interval(0.1, 2.0)), - gk.RBFKernel(lengthscale_constraint=Interval(0.1, 3.0)), - ), - False, - id="gpytorch-constraint", - ), - param( - *_pair(_rbf(2.0, frozen=True), _rbf(3.0, frozen=True)), - False, - id="gpytorch-frozen-value", - ), - param( - *_pair(_rbf(2.0, frozen=True), _rbf(2.0)), - False, - id="gpytorch-frozen-vs-learnable", - ), - param( - *_pair( - gk.IndexKernel(num_tasks=2, rank=1), gk.IndexKernel(num_tasks=2, rank=2) - ), - False, - id="gpytorch-shape", - ), - param( - *_pair( - gk.RBFKernel() + gk.MaternKernel(), gk.MaternKernel() + gk.RBFKernel() - ), - False, - id="gpytorch-summand-order", - ), - param( - *_task_pair(("a", "b"), _INDEX, ("a", "b"), _POSITIVE_INDEX), - False, - id="task-mode", - ), - param( - *_task_pair(("a", "b"), None, ("a", "b", "c"), None), - False, - id="task-values", - ), - param( - *_task_pair(("a", "b"), _INDEX, ("a", "b"), None), - False, - id="task-mode-vs-none", - ), - ], -) -def test_parameter_equivalence(p1, p2, expected): - """Parameter equivalence with kernel overrides produces the expected result.""" - assert p1.is_equivalent(p2) == expected - assert p2.is_equivalent(p1) == expected - - _X2_RBF = [_x("x1"), _x("x2", RBFKernel()), _x("x3")] _X1_SELECTOR = BayBEKernelFactory(parameter_selector=NameSelector(("x1",), regex=False)) _TASK_KERNEL = IndexKernel(num_tasks=3, rank=3, parameter_names=("task",)) From 31a57171ba18560eae476cdeeda43b9537077eb4 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Thu, 15 Oct 2026 13:19:19 +0200 Subject: [PATCH 72/73] Consolidate kernel override test helpers --- tests/test_parameter_kernel_overrides.py | 186 +++++++++++------------ 1 file changed, 85 insertions(+), 101 deletions(-) diff --git a/tests/test_parameter_kernel_overrides.py b/tests/test_parameter_kernel_overrides.py index 54f9abb811..67120663cb 100644 --- a/tests/test_parameter_kernel_overrides.py +++ b/tests/test_parameter_kernel_overrides.py @@ -33,23 +33,18 @@ from baybe.parameters.enum import TransferLearningMode from baybe.parameters.selectors import NameSelector from baybe.priors import GammaPrior -from baybe.searchspace import SearchSpace from baybe.surrogates import GaussianProcessSurrogate from baybe.surrogates.gaussian_process.components.kernel import ICMKernelFactory from baybe.surrogates.gaussian_process.core import _ModelContext from baybe.surrogates.gaussian_process.presets import BayBEKernelFactory -from baybe.targets import NumericalTarget from baybe.utils.dataframe import create_fake_input - -_continuous = partial(NumericalContinuousParameter, "x", (0, 1)) -_discrete = partial(NumericalDiscreteParameter, "x", (0, 1, 2)) -_categorical = partial(CategoricalParameter, "x", ("a", "b", "c")) -_custom = partial( - CustomDiscreteParameter, - "x", - pd.DataFrame({"d1": [1.0, 2.0], "d2": [3.0, 1.0]}, index=["a", "b"]), +from tests.test_kernel_factories import ( + _gpytorch_returning_factory, + _matern_kernel_factory, + _unnamed_matern_factory, ) -_task = partial(TaskParameter, "x", ("a", "b")) + +_PRIOR = GammaPrior(3.0, 6.0) _SCOPE_ERROR = "may only act on the parameter itself" _ACTIVE_DIMS_ERROR = "must not specify 'active_dims'" @@ -60,71 +55,76 @@ _POSITIVE_INDEX = TransferLearningMode.POSITIVE_INDEX_KERNEL -def _x(name, override=None): +def _x(name, override_kernel=None): """Create a continuous parameter on [0, 1], optionally with a kernel override.""" - return NumericalContinuousParameter(name, (0, 1), override_kernel=override) + return NumericalContinuousParameter(name, (0, 1), override_kernel=override_kernel) + +def _task(mode=None, *, name="task", values=("a", "b", "c"), **kwargs): + """Create a task parameter, optionally with a transfer learning mode.""" + return TaskParameter(name, values, override_transfer_learning_mode=mode, **kwargs) -def _tl_task(mode=None): - """Create a three-valued task parameter, optionally with a TL mode.""" - return TaskParameter("task", ("a", "b", "c"), override_transfer_learning_mode=mode) +_continuous = partial(_x, "x") +_discrete = partial(NumericalDiscreteParameter, "x", (0, 1, 2)) +_categorical = partial(CategoricalParameter, "x", ("a", "b", "c")) +_custom = partial( + CustomDiscreteParameter, + "x", + pd.DataFrame({"d1": [1.0, 2.0], "d2": [3.0, 1.0]}, index=["a", "b"]), +) -def _basic_kernels(kernel, columns): - """Yield each basic kernel of a GPyTorch kernel with the columns it acts on. - Basic kernels compute on columns (e.g. RBF, Matern), as opposed to wrappers such - as Scale, Additive, or Product kernels. As in GPyTorch, ``active_dims`` of a - kernel apply to everything nested inside it. - """ - if kernel.active_dims is not None: - columns = set(kernel.active_dims.tolist()) - children = [c for c in kernel.children() if isinstance(c, gk.Kernel)] - children += [ - m for c in kernel.children() if isinstance(c, torch.nn.ModuleList) for m in c - ] - if not children: - yield kernel, columns - for child in children: - yield from _basic_kernels(child, columns) +def _gpytorch_overrides(parameters): + """Map parameter names to their user-provided GPyTorch kernel overrides.""" + return { + p.name: p.override_kernel + for p in parameters + if p.override_kernel is not None and not isinstance(p.override_kernel, Kernel) + } -def _factors_by_parameter(kernel, searchspace): - """Map the parameters each basic kernel acts on to the basic kernel classes. +def _basic_kernels(kernel, searchspace): + """List each basic kernel with the parameters and columns it acts on. - Example: ``{("x1", "x3"): [MaternKernel], ("x2",): [RBFKernel]}`` + Basic kernels compute on columns (e.g. RBF, Matern), as opposed to wrappers such + as Scale, Additive, or Product kernels. As in GPyTorch, ``active_dims`` of a + kernel apply to everything nested inside it. """ - factors: dict[tuple[str, ...], list[type]] = {} - for k, names, _ in _basic_kernels_by_parameter(kernel, searchspace): - factors.setdefault(names, []).append(type(k)) - return factors + def walk(k, columns): + if k.active_dims is not None: + columns = set(k.active_dims.tolist()) + children = [c for c in k.children() if isinstance(c, gk.Kernel)] + children += [ + m for c in k.children() if isinstance(c, torch.nn.ModuleList) for m in c + ] + if not children: + yield k, columns + for child in children: + yield from walk(child, columns) -def _basic_kernels_by_parameter(kernel, searchspace): - """Yield each basic kernel with the parameters and columns it acts on.""" all_columns = set(range(len(searchspace.comp_rep_columns))) - for k, columns in _basic_kernels(kernel, all_columns): + result = [] + for k, columns in walk(kernel, all_columns): names = tuple( p.name for p in searchspace.parameters if columns & set(searchspace.get_comp_rep_parameter_indices(p.name)) ) - yield k, names, columns - - -def _matern_naming_all(searchspace, objective, measurements): - """Return a Matern kernel naming all parameters of the given search space.""" - return MaternKernel(parameter_names=searchspace.parameter_names) + result.append((k, names, columns)) + return result -def _unnamed_matern(searchspace, objective, measurements): - """Return a Matern kernel without parameter names.""" - return MaternKernel() - +def _factors_by_parameter(kernel, searchspace): + """Map the parameters each basic kernel acts on to the basic kernel classes. -def _raw_matern(searchspace, objective, measurements): - """Return a raw GPyTorch kernel, which cannot exclude overridden parameters.""" - return gk.MaternKernel() + Example: ``{("x1", "x3"): [MaternKernel], ("x2",): [RBFKernel]}`` + """ + factors: dict[tuple[str, ...], list[type]] = {} + for k, names, _ in _basic_kernels(kernel, searchspace): + factors.setdefault(names, []).append(type(k)) + return factors def _needs_full_searchspace(searchspace, objective, measurements): @@ -223,7 +223,7 @@ def _needs_full_searchspace(searchspace, objective, measurements): ), # Task parameters derive their override from the transfer learning mode param( - _task, + partial(_task, name="x", values=("a", "b")), RBFKernel(), TypeError, "unexpected keyword argument 'override_kernel'", @@ -317,7 +317,7 @@ def test_override_assignment(constructor, override, error, match): ), param( _X2_RBF, - _matern_naming_all, + _matern_kernel_factory, {("x1", "x3"): [gk.MaternKernel], ("x2",): [gk.RBFKernel]}, id="factory", ), @@ -378,49 +378,49 @@ def test_override_assignment(constructor, override, error, match): ), # Task parameters: the transfer learning mode defines the task kernel param( - [_x("x1"), _tl_task()], + [_x("x1"), _task()], None, {("x1",): [gk.MaternKernel], ("task",): [BoPositiveIndexKernel]}, id="task-without-mode", ), param( - [_x("x1"), _tl_task(_INDEX)], + [_x("x1"), _task(_INDEX)], None, {("x1",): [gk.MaternKernel], ("task",): [gk.IndexKernel]}, id="task-index", ), param( - [_x("x1"), _tl_task(_POSITIVE_INDEX)], + [_x("x1"), _task(_POSITIVE_INDEX)], None, {("x1",): [gk.MaternKernel], ("task",): [BoPositiveIndexKernel]}, id="task-positive-index", ), param( - [_x("x1"), _tl_task(_INDEX)], + [_x("x1"), _task(_INDEX)], MaternKernel(parameter_names=("x1", "task")), {("x1",): [gk.MaternKernel], ("task",): [gk.IndexKernel]}, id="task-index+kernel-naming-task", ), param( - [_x("x1"), _tl_task(_INDEX)], + [_x("x1"), _task(_INDEX)], MaternKernel(parameter_names=("x1",)) * _TASK_KERNEL, {("x1",): [gk.MaternKernel], ("task",): [gk.IndexKernel]}, id="task-index+product-with-task-kernel", ), param( - [_x("x1"), _tl_task(_INDEX)], + [_x("x1"), _task(_INDEX)], _TASK_KERNEL, {("task",): [gk.IndexKernel]}, id="task-index+task-only-kernel", ), param( - [_x("x1"), _tl_task(_INDEX)], - _unnamed_matern, + [_x("x1"), _task(_INDEX)], + _unnamed_matern_factory, {("x1",): [gk.MaternKernel], ("task",): [gk.IndexKernel]}, id="task-index+factory", ), param( - [_x("x1"), _x("x2", RBFKernel()), _tl_task(_POSITIVE_INDEX)], + [_x("x1"), _x("x2", RBFKernel()), _task(_POSITIVE_INDEX)], None, { ("x1",): [gk.MaternKernel], @@ -430,7 +430,7 @@ def test_override_assignment(constructor, override, error, match): id="task-positive-index+parameter-override", ), param( - [_x("x1"), _x("x2", RBFKernel()), _tl_task()], + [_x("x1"), _x("x2", RBFKernel()), _task()], _X1_SELECTOR, { ("x1",): [gk.MaternKernel], @@ -472,7 +472,7 @@ def test_override_assignment(constructor, override, error, match): ), param( _X2_RBF, - _raw_matern, + _gpytorch_returning_factory, (IncompatibleOverrideError, _EXCLUDE_ERROR), id="factory-returning-gpytorch-kernel", ), @@ -483,19 +483,18 @@ def test_override_assignment(constructor, override, error, match): id="factory-needing-full-searchspace", ), param( - [_x("x1"), _tl_task(_INDEX)], + [_x("x1"), _task(_INDEX)], ICMKernelFactory(task_kernel_or_factory=_TASK_KERNEL), (IncompatibleOverrideError, _EXCLUDE_ERROR), id="task-index+icm-factory", ), ], ) -def test_override_kernel_construction(parameters, kernel_or_factory, expected): +def test_override_kernel_construction( + searchspace, objective, kernel_or_factory, expected +): """Kernel construction with overrides yields the expected kernel or error.""" - searchspace = SearchSpace.from_product(parameters) - context = _ModelContext( - searchspace, NumericalTarget("y").to_objective(), pd.DataFrame() - ) + context = _ModelContext(searchspace, objective, pd.DataFrame()) surrogate = GaussianProcessSurrogate(kernel_or_factory=kernel_or_factory) # Invalid configurations raise during construction @@ -510,34 +509,25 @@ def test_override_kernel_construction(parameters, kernel_or_factory, expected): assert _factors_by_parameter(kernel, searchspace) == expected # User-provided GPyTorch overrides are copied, leaving the input unbound - gpytorch_overrides = { - p.name: p.override_kernel - for p in parameters - if p.override_kernel is not None and not isinstance(p.override_kernel, Kernel) - } + gpytorch_overrides = _gpytorch_overrides(searchspace.parameters) for override in gpytorch_overrides.values(): assert all(k.active_dims is None for k in (override, *override.sub_kernels())) # BayBE-created kernels use one ARD dimension per column, user-provided GPyTorch # kernels keep their setting user_kernel = isinstance(kernel_or_factory, gk.Kernel) - for k, names, columns in _basic_kernels_by_parameter(kernel, searchspace): + for k, names, columns in _basic_kernels(kernel, searchspace): if user_kernel or set(names) & gpytorch_overrides.keys(): assert k.ard_num_dims in (None, len(columns)) else: assert k.ard_num_dims == len(columns) -def _rbf_with_prior(): - """Create a regularized BayBE RBF kernel for numerically stable fits.""" - return RBFKernel(lengthscale_prior=GammaPrior(3.0, 6.0)) - - @pytest.mark.parametrize( ("parameters", "expected"), [ param( - [_x("x1"), _x("x2", _rbf_with_prior())], + [_x("x1"), _x("x2", RBFKernel(lengthscale_prior=_PRIOR))], {("x1",): [gk.MaternKernel], ("x2",): [gk.RBFKernel]}, id="parameter-override", ), @@ -545,11 +535,11 @@ def _rbf_with_prior(): [ _x("x1"), CategoricalParameter( - "c", ("a", "b", "c"), override_kernel=_rbf_with_prior() - ), - TaskParameter( - "task", ("a", "b"), override_transfer_learning_mode=_INDEX + "c", + ("a", "b", "c"), + override_kernel=RBFKernel(lengthscale_prior=_PRIOR), ), + _task(_INDEX), ], { ("x1",): [gk.MaternKernel], @@ -560,8 +550,8 @@ def _rbf_with_prior(): ), param( [ - _x("x1", _rbf_with_prior()), - _x("x2", MaternKernel(lengthscale_prior=GammaPrior(3.0, 6.0))), + _x("x1", RBFKernel(lengthscale_prior=_PRIOR)), + _x("x2", MaternKernel(lengthscale_prior=_PRIOR)), ], {("x1",): [gk.RBFKernel], ("x2",): [gk.MaternKernel]}, id="all-overridden", @@ -579,18 +569,12 @@ def _rbf_with_prior(): ), ], ) -def test_fitted_override_kernel(parameters, expected): +def test_fitted_override_kernel(searchspace, objective, expected): """The fitted model uses the kernel constructed from the overrides.""" - searchspace = SearchSpace.from_product(parameters) - objective = NumericalTarget("y").to_objective() measurements = create_fake_input(searchspace.parameters, objective.targets, 6) # Snapshot user-provided GPyTorch overrides to detect mutations by training - gpytorch_overrides = [ - p.override_kernel - for p in parameters - if p.override_kernel is not None and not isinstance(p.override_kernel, Kernel) - ] + gpytorch_overrides = list(_gpytorch_overrides(searchspace.parameters).values()) snapshots = [deepcopy(o.state_dict()) for o in gpytorch_overrides] surrogate = GaussianProcessSurrogate() @@ -603,7 +587,7 @@ def test_fitted_override_kernel(parameters, expected): # Training did not touch the user-provided GPyTorch overrides for override, snapshot in zip(gpytorch_overrides, snapshots): - assert override.active_dims is None + assert all(k.active_dims is None for k in (override, *override.sub_kernels())) assert all( torch.equal(v, override.state_dict()[k]) for k, v in snapshot.items() ) From fe42885734b1f81ca3804e5299b3d02c0328c459 Mon Sep 17 00:00:00 2001 From: Martin Fitzner Date: Thu, 15 Oct 2026 20:20:30 +0200 Subject: [PATCH 73/73] Add TODO note --- baybe/surrogates/gaussian_process/core.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py index b0d86ea502..c1cfce0fd1 100644 --- a/baybe/surrogates/gaussian_process/core.py +++ b/baybe/surrogates/gaussian_process/core.py @@ -512,6 +512,10 @@ def _resolve_default_base( The numerical base always excludes the task; the default task kernel is re-added unless the task is itself overridden (i.e. already excluded). """ + # TODO: This method might become obsolete once it can be married with the lazy + # searchspace updates. In that case, it should be reworked so it is largely or + # fully replaced by a simple factory call. + from baybe.surrogates.gaussian_process.presets.baybe import ( _BayBENumericalKernelFactory, _BayBETaskKernelFactory,