Add Per-Parameter Convenience Kernel Override

CHANGELOG.md

@@ -21,6 +21,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   number of required models instead of the number of transform outputs
 
 ### Added
+- Parameter-specific kernel overrides for composing Gaussian process kernels on
+  individual parameter dimensions
 - `coefficients` attribute for `DiscreteSumConstraint`, enabling weighted sums. Follows
   the same pattern as `ContinuousLinearConstraint.coefficients`
 - `simplex_coefficients` keyword argument to `SubspaceDiscrete.from_simplex` for
@@ -31,6 +33,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   `BayesianRecommender`
 - `Parameter.is_equivalent` method for structural parameter comparison
 - `posterior_mean_function` method to `GaussianProcessSurrogate`
+- `TaskParameter.override_transfer_learning_mode` (and the corresponding
+  `TransferLearningMode` enum) for selecting the kernel that models the task
+  correlations in transfer learning, taking precedence over the task kernel of the
+  configured kernel factory. `INDEX_KERNEL` permits arbitrary correlations while
+  `POSITIVE_INDEX_KERNEL` enforces positive ones
 
 ### Changed
 - `BOTORCH` GP preset now includes `BetaPrior(2.5, 1.5)` for the task covariance

baybe/exceptions.py

@@ -74,6 +74,14 @@ class IncompatibleArgumentError(IncompatibilityError):
     """An incompatible argument was passed to a callable."""
 
 
+class IncompatibleOverrideError(IncompatibilityError):
+    """An override conflicts with another specification."""
+
+
+class UnsupportedSearchSpaceAttributeError(AttributeError):
+    """Access to a blocked attribute on a reduced search space was attempted."""
+
+
 class NonGaussianityError(Exception):
     """An operation assuming Gaussianity is attempted on a non-Gaussian distribution."""
 

baybe/kernels/base.py

@@ -7,20 +7,20 @@
 from itertools import chain
 from typing import TYPE_CHECKING, Any
 
-from attrs import define, field
+from attrs import define, evolve, field
 from attrs.converters import optional as optional_c
 from attrs.validators import deep_iterable, instance_of
 from attrs.validators import optional as optional_v
 from typing_extensions import override
 
 from baybe.exceptions import UnmatchedAttributeError
 from baybe.priors.base import Prior
-from baybe.searchspace.core import SearchSpace
 from baybe.serialization.mixin import SerialMixin
 from baybe.settings import active_settings
 from baybe.utils.basic import 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
 
 
@@ -99,6 +99,45 @@ def to_factory(self) -> PlainKernelFactory:
 
         return PlainKernelFactory(self)
 
+    def _without_parameter(
+        self, name: str, searchspace: SearchSpace, /
+    ) -> Kernel | None:
+        """Return a copy of the kernel that no longer acts on the given parameter.
+
+        Args:
+            name: The name of the parameter to remove.
+            searchspace: The search space, used to enumerate the remaining parameter
+                names when the kernel does not explicitly specify the ones it acts on.
+
+        Raises:
+            TypeError: If the kernel structure does not support removing a single
+                parameter unambiguously.
+
+        Returns:
+            The reduced kernel, or ``None`` if removing the parameter leaves the
+            kernel with no parameters to act on.
+        """
+        raise TypeError(
+            f"Cannot remove a parameter from kernel '{self.__class__.__name__}'. "
+            f"Only basic kernels and scaled basic kernels are supported."
+        )
+
+    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
@@ -239,6 +278,22 @@ def _get_dimensions(
         )
         return active_dims, ard_num_dims
 
+    @override
+    def _without_parameter(
+        self, name: str, searchspace: SearchSpace, /
+    ) -> Kernel | None:
+        if self.parameter_names is None:
+            remaining = tuple(n for n in searchspace.parameter_names if n != name)
+        elif name in self.parameter_names:
+            remaining = tuple(n for n in self.parameter_names if n != name)
+        else:
+            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):

baybe/kernels/composite.py

@@ -1,10 +1,13 @@
 """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, field
+from attrs import define, evolve, field
 from attrs.converters import optional as optional_c
 from attrs.validators import deep_iterable, gt, instance_of, min_len
 from attrs.validators import optional as optional_v
@@ -16,6 +19,9 @@
 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):
@@ -42,6 +48,17 @@ class ScaleKernel(CompositeKernel):
     If ``False``, the output scale is frozen at its initial value and excluded from
     optimization."""
 
+    @override
+    def _without_parameter(
+        self, name: str, searchspace: SearchSpace, /
+    ) -> Kernel | None:
+        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
@@ -68,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))
@@ -85,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))

baybe/parameters/__init__.py

@@ -6,6 +6,7 @@
     CategoricalEncoding,
     CustomEncoding,
     SubstanceEncoding,
+    TransferLearningMode,
 )
 from baybe.parameters.numerical import (
     NumericalContinuousParameter,
@@ -25,4 +26,5 @@
     "SubstanceEncoding",
     "SubstanceParameter",
     "TaskParameter",
+    "TransferLearningMode",
 ]

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:

baybe/parameters/categorical.py

@@ -5,12 +5,13 @@
 
 import numpy as np
 import pandas as pd
+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 baybe.parameters.base import _DiscreteLabelLikeParameter
-from baybe.parameters.enum import CategoricalEncoding
+from baybe.parameters.enum import CategoricalEncoding, TransferLearningMode
 from baybe.settings import active_settings
 from baybe.utils.conversion import nonstring_to_tuple, sort_tuple
 from baybe.utils.validation import validate_unique_values
@@ -87,6 +88,24 @@ 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),
+    )
+    """Optional override for how the task dimension is modeled.
+
+    Only applies to :class:`.GaussianProcessSurrogate`. When ``None``, the surrogate's
+    kernel factory decides how the task dimension is treated. When set, the surrogate
+    attaches the requested task kernel to a task-free base kernel derived from the
+    configured factory.
+
+    If the configured factory does not reduce to a task-free base kernel, an
+    :class:`.IncompatibleOverrideError` is raised.
+    """
+
 
 # Collect leftover original slotted classes processed by `attrs.define`
 gc.collect()