Add TL override path to new GP factory

CHANGELOG.md

@@ -36,6 +36,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   inverting the specification to keep the complement
 - `DiscreteSelectionConstraint` as the condition-based filtering constraint
   (inclusion-by-default; replaces `DiscreteExcludeConstraint`)
+- `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
 
 ### 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,7 +7,7 @@
 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
@@ -99,6 +99,28 @@ 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 specified parameter.
+
+        Args:
+            name: The name of the parameter to ignore.
+            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__}'. "
+        )
+
     @abstractmethod
     def _get_dimensions(
         self, searchspace: SearchSpace
@@ -239,6 +261,18 @@ 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
+
 
 @define(frozen=True)
 class CompositeKernel(Kernel, ABC):

baybe/kernels/composite.py

@@ -4,14 +4,15 @@
 from functools import reduce
 from operator import add, mul
 
-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
 from typing_extensions import override
 
 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
@@ -42,6 +43,13 @@ 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 to_gpytorch(self, *args, **kwargs):
         import torch

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/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,18 @@ class TaskParameter(CategoricalParameter):
     encoding: CategoricalEncoding = field(default=CategoricalEncoding.INT, init=False)
     # See base class.
 
+    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.
+    """
+
 
 # Collect leftover original slotted classes processed by `attrs.define`
 gc.collect()

baybe/parameters/enum.py

@@ -174,3 +174,14 @@ class SubstanceEncoding(ParameterEncoding):
 
     WHIM = "WHIM"
     """:class:`skfp.fingerprints.WHIMFingerprint`"""
+
+
+class TransferLearningMode(Enum):
+    """Transfer learning modes for :class:`.TaskParameter`."""
+
+    INDEX_KERNEL = "INDEX_KERNEL"
+    """:class:`gpytorch.kernels.IndexKernel` for arbitrary correlations."""
+
+    POSITIVE_INDEX_KERNEL = "POSITIVE_INDEX_KERNEL"
+    """:class:`botorch.models.kernels.positive_index.PositiveIndexKernel` for positive
+    correlations."""

baybe/searchspace/core.py

@@ -16,7 +16,10 @@
 
 from baybe.constraints import validate_constraints
 from baybe.constraints.base import Constraint
-from baybe.exceptions import InfeasibilityError
+from baybe.exceptions import (
+    InfeasibilityError,
+    _UnsupportedSearchSpaceAttributeError,
+)
 from baybe.parameters import TaskParameter
 from baybe.parameters.base import Parameter
 from baybe.searchspace.continuous import SubspaceContinuous
@@ -616,7 +619,7 @@ def __getattribute__(self, name: str):
         allowed = object.__getattribute__(self, "_ALLOWED_ATTRIBUTES")
         if name in allowed:
             return object.__getattribute__(self, name)
-        raise AttributeError(
+        raise _UnsupportedSearchSpaceAttributeError(
             f"'{object.__getattribute__(self, '__class__').__name__}' does not "
             f"support attribute '{name}'. Only parameter information is available."
         )