Add Per-Parameter Convenience Kernel Override

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

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,23 @@ def _without_parameter(
             f"Cannot remove a parameter from kernel '{self.__class__.__name__}'. "
         )
 
+    def _scope_to_parameter(self, name: str | None, /) -> Kernel:
+        """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
+                remove all scoping so that the kernel acts on all dimensions.
+
+        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 +290,10 @@ def _without_parameter(
             return self
         return evolve(self, parameter_names=remaining) if remaining else None
 
+    @override
+    def _scope_to_parameter(self, name: str | None, /) -> Kernel:
+        return evolve(self, parameter_names=None if name is None else (name,))
+
 
 @define(frozen=True)
 class CompositeKernel(Kernel, ABC):

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,39 @@
 
 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
+
+
+def _reduce_base_kernels(
+    kernel: AdditiveKernel | ProductKernel, name: str, searchspace: SearchSpace, /
+) -> Kernel | None:
+    """Remove a parameter from all base kernels of an additive or product kernel.
+
+    Args:
+        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 kernel of the same type, 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):
@@ -50,6 +81,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 _scope_to_parameter(self, name: str | None, /) -> Kernel:
+        return evolve(self, base_kernel=self.base_kernel._scope_to_parameter(name))
+
     @override
     def to_gpytorch(self, *args, **kwargs):
         import torch
@@ -76,6 +111,19 @@ 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 | None, /) -> Kernel:
+        return evolve(
+            self,
+            base_kernels=tuple(k._scope_to_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 +141,19 @@ class ProductKernel(CompositeKernel):
     )
     """The individual kernels to be multiplied."""
 
+    @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 | None, /) -> Kernel:
+        return evolve(
+            self,
+            base_kernels=tuple(k._scope_to_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))