Botorch preset

CHANGELOG.md

@@ -17,7 +17,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   components, enabling full low-level customization
 - Configurable fitting criterion for Gaussian process hyperparameter optimization
 - Factories for all Gaussian process components
-- `CHEN`, `EDBO` and `EDBO_SMOOTHED` presets for `GaussianProcessSurrogate`
+- `BOTORCH`, `CHEN`, `EDBO` and `EDBO_SMOOTHED` presets for `GaussianProcessSurrogate`
 - `TypeSelector` and `NameSelector` classes for parameter selection in kernel factories
 - `parameter_names` attribute to basic kernels for controlling the considered parameters
 - `ParameterKind` flag enum for classifying parameters by their role and automatic

baybe/surrogates/gaussian_process/components/_gpytorch.py

@@ -0,0 +1,71 @@
+"""Custom GPyTorch components."""
+
+import torch
+from botorch.models.multitask import _compute_multitask_mean
+from botorch.models.utils.gpytorch_modules import MIN_INFERRED_NOISE_LEVEL
+from gpytorch.constraints import GreaterThan
+from gpytorch.likelihoods.hadamard_gaussian_likelihood import HadamardGaussianLikelihood
+from gpytorch.means import MultitaskMean
+from gpytorch.means.multitask_mean import Mean
+from gpytorch.priors import LogNormalPrior
+from torch import Tensor
+from torch.nn import Module
+
+
+class HadamardConstantMean(Mean):
+    """A GPyTorch mean function implementing BoTorch's multitask mean logic.
+
+    While GPyTorch already provides a :class:`~gpytorch.means.MultitaskMean` class, it
+    computes mean values for all (input, task)-pairs (where input means all parameters
+    except the task parameter), i.e. it intrinsically applies a Cartesian expansion.
+    However, for the regular transfer learning setting, we only need the means for the
+    pairs that are actually observed/requested. BoTorch subselects the relevant means
+    from the GPyTorch output in `MultiTaskGP.forward`, i.e. it uses a class-based
+    approach to define its special logic for the multitask case. In contrast, BayBE uses
+    a composition approach, which is more flexible but requires that the logic is
+    injected via a self-contained `Mean` object, which is what this class provides.
+
+    Note:
+        Analogous to GPyTorch's
+        https://github.com/cornellius-gp/gpytorch/blob/main/gpytorch/likelihoods/hadamard_gaussian_likelihood.py
+        but where the logic is applied to the mean function, i.e. we learn a different
+        (constant) mean for each task.
+    """
+
+    def __init__(self, mean_module: Module, num_tasks: int, task_feature: int):
+        super().__init__()
+        self.multitask_mean = MultitaskMean(mean_module, num_tasks=num_tasks)
+        self.task_feature = task_feature
+
+    def forward(self, x: Tensor) -> Tensor:
+        # Adapted from https://github.com/meta-pytorch/botorch/blob/e0f4f5b941b5949a4a1171bf8d4ee9f74f146f3a/botorch/models/multitask.py#L397
+
+        # Convert task feature to positive index
+        task_feature = self.task_feature % x.shape[-1]
+
+        # Split input into task and non-task components
+        x_before = x[..., :task_feature]
+        task_idcs = x[..., task_feature : task_feature + 1]
+        x_after = x[..., task_feature + 1 :]
+
+        return _compute_multitask_mean(
+            self.multitask_mean, x_before, task_idcs, x_after
+        )
+
+
+def make_botorch_multitask_likelihood(
+    num_tasks: int, task_feature: int
+) -> HadamardGaussianLikelihood:
+    """Adapted from :class:`botorch.models.multitask.MultiTaskGP`."""
+    noise_prior = LogNormalPrior(loc=-4.0, scale=1.0)
+    return HadamardGaussianLikelihood(
+        num_tasks=num_tasks,
+        batch_shape=torch.Size(),
+        noise_prior=noise_prior,
+        noise_constraint=GreaterThan(
+            MIN_INFERRED_NOISE_LEVEL,
+            transform=None,
+            initial_value=noise_prior.mode,
+        ),
+        task_feature_index=task_feature,
+    )

baybe/surrogates/gaussian_process/components/kernel.py

@@ -70,14 +70,10 @@ def __attrs_post_init__(self):
                 f"they actually use the selected parameter names."
             )
 
-    def get_parameter_names(self, searchspace: SearchSpace) -> tuple[str, ...] | None:
+    def get_parameter_names(self, searchspace: SearchSpace) -> tuple[str, ...]:
         """Get the names of the parameters to be considered by the kernel."""
-        if self.parameter_selector is None:
-            return None
-
-        return tuple(
-            p.name for p in searchspace.parameters if self.parameter_selector(p)
-        )
+        selector = self.parameter_selector or (lambda _: True)
+        return tuple(p.name for p in searchspace.parameters if selector(p))
 
     def _validate_parameter_kinds(self, parameters: Iterable[Parameter]) -> None:
         """Validate that the given parameters are supported by the factory.
@@ -102,7 +98,7 @@ def _validate_parameter_kinds(self, parameters: Iterable[Parameter]) -> None:
     @override
     def __call__(
         self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
-    ) -> Kernel:
+    ) -> Kernel | GPyTorchKernel:
         """Construct the kernel, validating parameter kinds before construction."""
         if self.parameter_selector is not None:
             params = [p for p in searchspace.parameters if self.parameter_selector(p)]
@@ -115,7 +111,7 @@ def __call__(
     @abstractmethod
     def _make(
         self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
-    ) -> Kernel:
+    ) -> Kernel | GPyTorchKernel:
         """Construct the kernel."""
 
 
@@ -127,7 +123,7 @@ class _MetaKernelFactory(KernelFactoryProtocol, ABC):
     @abstractmethod
     def __call__(
         self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
-    ) -> Kernel: ...
+    ) -> Kernel | GPyTorchKernel: ...
 
 
 @define
@@ -170,11 +166,40 @@ def _default_task_kernel_factory(self) -> KernelFactoryProtocol:
     @override
     def __call__(
         self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
-    ) -> Kernel:
+    ) -> Kernel | GPyTorchKernel:
+        if searchspace.task_idx is None:
+            raise IncompatibleSearchSpaceError(
+                f"'{type(self).__name__}' can only be used with a searchspace that "
+                f"contains a '{TaskParameter.__name__}'."
+            )
+
         base_kernel = self.base_kernel_factory(searchspace, train_x, train_y)
         task_kernel = self.task_kernel_factory(searchspace, train_x, train_y)
         if isinstance(base_kernel, Kernel):
             base_kernel = base_kernel.to_gpytorch(searchspace)
         if isinstance(task_kernel, Kernel):
             task_kernel = task_kernel.to_gpytorch(searchspace)
+
+        # Ensure correct partitioning between base and task kernels active dimensions
+        all_idcs = set(range(len(searchspace.comp_rep_columns)))
+        allowed_task_idcs = {searchspace.task_idx}
+        allowed_base_idcs = all_idcs - allowed_task_idcs
+        base_idcs = (
+            set(d.tolist()) if (d := base_kernel.active_dims) is not None else all_idcs
+        )
+        task_idcs = (
+            set(d.tolist()) if (d := task_kernel.active_dims) is not None else all_idcs
+        )
+
+        if not base_idcs <= allowed_base_idcs:
+            raise ValueError(
+                f"The base kernel's 'active_dims' {base_idcs} must be a subset of "
+                f"the non-task indices {allowed_base_idcs}."
+            )
+        if task_idcs != allowed_task_idcs:
+            raise ValueError(
+                f"The task kernel's 'active_dims' {task_idcs} does not match "
+                f"the task index {allowed_task_idcs}."
+            )
+
         return base_kernel * task_kernel

baybe/surrogates/gaussian_process/core.py

@@ -183,8 +183,8 @@ class GaussianProcessSurrogate(Surrogate):
         * :class:`gpytorch.likelihoods.Likelihood`
     """
 
-    criterion_factory: FitCriterionFactoryProtocol = field(
-        alias="criterion_or_factory",
+    fit_criterion_factory: FitCriterionFactoryProtocol = field(
+        alias="fit_criterion_or_factory",
         factory=BayBEFitCriterionFactory,
         converter=partial(  # type: ignore[misc]
             to_component_factory, component_type=GPComponentType.CRITERION
@@ -215,7 +215,9 @@ def from_preset(
         likelihood_or_factory: LikelihoodFactoryProtocol
         | GPyTorchLikelihood
         | None = None,
-        criterion_or_factory: FitCriterion | FitCriterionFactoryProtocol | None = None,
+        fit_criterion_or_factory: FitCriterion
+        | FitCriterionFactoryProtocol
+        | None = None,
     ) -> Self:
         """Create a Gaussian process surrogate from one of the defined presets."""
         preset = GaussianProcessPreset(preset)
@@ -228,9 +230,13 @@ def from_preset(
         kernel = kernel_or_factory or getattr(module, "KERNEL_FACTORY")
         mean = mean_or_factory or getattr(module, "MEAN_FACTORY")
         likelihood = likelihood_or_factory or getattr(module, "LIKELIHOOD_FACTORY")
-        criterion = criterion_or_factory or getattr(module, "FIT_CRITERION_FACTORY")
+        fit_criterion = fit_criterion_or_factory or getattr(
+            module, "FIT_CRITERION_FACTORY"
+        )
 
-        return cls(kernel, mean, likelihood, criterion)
+        gp = cls(kernel, mean, likelihood, fit_criterion)
+        gp._custom_kernel = False  # preset are first-party features
+        return gp
 
     @override
     def to_botorch(self) -> GPyTorchModel:
@@ -301,7 +307,7 @@ def _fit(self, train_x: Tensor, train_y: Tensor) -> None:
         likelihood = self.likelihood_factory(context.searchspace, train_x, train_y)
 
         ### Criterion
-        criterion = self.criterion_factory(context.searchspace, train_x, train_y)
+        criterion = self.fit_criterion_factory(context.searchspace, train_x, train_y)
 
         ### Model construction and fitting
         self._model = botorch.models.SingleTaskGP(
@@ -323,7 +329,7 @@ def __str__(self) -> str:
             to_string("Mean factory", self.mean_factory, single_line=True),
             to_string("Likelihood factory", self.likelihood_factory, single_line=True),
             to_string(
-                "Fit criterion factory", self.criterion_factory, single_line=True
+                "Fit criterion factory", self.fit_criterion_factory, single_line=True
             ),
         ]
         return to_string(super().__str__(), *fields)

baybe/surrogates/gaussian_process/presets/__init__.py

@@ -11,6 +11,13 @@
     BayBEMeanFactory,
 )
 
+# BoTorch preset
+from baybe.surrogates.gaussian_process.presets.botorch import (
+    BotorchKernelFactory,
+    BotorchLikelihoodFactory,
+    BotorchMeanFactory,
+)
+
 # Chen preset
 from baybe.surrogates.gaussian_process.presets.chen import (
     CHENFitCriterionFactory,
@@ -45,6 +52,10 @@
     "BayBEKernelFactory",
     "BayBELikelihoodFactory",
     "BayBEMeanFactory",
+    # BoTorch preset
+    "BotorchKernelFactory",
+    "BotorchLikelihoodFactory",
+    "BotorchMeanFactory",
     # Chen preset
     "CHENFitCriterionFactory",
     "CHENKernelFactory",

baybe/surrogates/gaussian_process/presets/botorch.py

@@ -0,0 +1,140 @@
+"""BoTorch preset for Gaussian process surrogates."""
+
+from __future__ import annotations
+
+import gc
+from itertools import chain
+from typing import TYPE_CHECKING, ClassVar
+
+from attrs import define
+from typing_extensions import override
+
+from baybe.kernels.base import Kernel
+from baybe.parameters.enum import _ParameterKind
+from baybe.searchspace.core import SearchSpace
+from baybe.surrogates.gaussian_process.components import LikelihoodFactoryProtocol
+from baybe.surrogates.gaussian_process.components.fit_criterion import (
+    FitCriterion,
+    PlainFitCriterionFactory,
+)
+from baybe.surrogates.gaussian_process.components.kernel import (
+    ICMKernelFactory,
+    _PureKernelFactory,
+)
+from baybe.surrogates.gaussian_process.components.mean import MeanFactoryProtocol
+
+if TYPE_CHECKING:
+    from gpytorch.kernels import Kernel as GPyTorchKernel
+    from gpytorch.likelihoods import Likelihood as GPyTorchLikelihood
+    from gpytorch.means import Mean as GPyTorchMean
+    from torch import Tensor
+
+
+@define
+class BotorchKernelFactory(_PureKernelFactory):
+    """A factory providing BoTorch kernels."""
+
+    _uses_parameter_names: ClassVar[bool] = True
+    # See base class.
+
+    _supported_parameter_kinds: ClassVar[_ParameterKind] = (
+        _ParameterKind.REGULAR | _ParameterKind.TASK
+    )
+    # See base class.
+
+    @override
+    def _make(
+        self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
+    ) -> Kernel | GPyTorchKernel:
+        from botorch.models.kernels.positive_index import PositiveIndexKernel
+        from botorch.models.utils.gpytorch_modules import (
+            get_covar_module_with_dim_scaled_prior,
+        )
+
+        parameter_names = self.get_parameter_names(searchspace)
+
+        # For regular parameters, resolve parameter names to active dimension indices
+        active_dims = list(
+            chain.from_iterable(
+                searchspace.get_comp_rep_parameter_indices(name)
+                for name in parameter_names
+                if searchspace.get_parameters_by_name([name])[0]._kind
+                is _ParameterKind.REGULAR
+            )
+        )
+        ard_num_dims = len(active_dims)
+
+        # Create the base kernel for the regular parameters
+        base_kernel = get_covar_module_with_dim_scaled_prior(
+            ard_num_dims=ard_num_dims, active_dims=active_dims
+        )
+
+        # Single-task case
+        if (task_idx := searchspace.task_idx) is None:
+            return base_kernel
+
+        index_kernel = PositiveIndexKernel(
+            num_tasks=searchspace.n_tasks,
+            rank=searchspace.n_tasks,
+            active_dims=[task_idx],
+        )
+        return ICMKernelFactory(base_kernel, index_kernel)(
+            searchspace, train_x, train_y
+        )
+
+
+class BotorchMeanFactory(MeanFactoryProtocol):
+    """A factory providing BoTorch mean functions."""
+
+    @override
+    def __call__(
+        self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
+    ) -> GPyTorchMean:
+        from gpytorch.means import ConstantMean
+
+        from baybe.surrogates.gaussian_process.components._gpytorch import (
+            HadamardConstantMean,
+        )
+
+        if searchspace.n_tasks == 1:
+            return ConstantMean()
+
+        assert searchspace.task_idx is not None
+        return HadamardConstantMean(
+            ConstantMean(), searchspace.n_tasks, searchspace.task_idx
+        )
+
+
+class BotorchLikelihoodFactory(LikelihoodFactoryProtocol):
+    """A factory providing BoTorch likelihoods."""
+
+    @override
+    def __call__(
+        self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor
+    ) -> GPyTorchLikelihood:
+
+        if searchspace.n_tasks == 1:
+            from botorch.models.utils.gpytorch_modules import (
+                get_gaussian_likelihood_with_lognormal_prior,
+            )
+
+            return get_gaussian_likelihood_with_lognormal_prior()
+
+        from baybe.surrogates.gaussian_process.components._gpytorch import (
+            make_botorch_multitask_likelihood,
+        )
+
+        assert searchspace.task_idx is not None
+        return make_botorch_multitask_likelihood(
+            num_tasks=searchspace.n_tasks, task_feature=searchspace.task_idx
+        )
+
+
+# Collect leftover original slotted classes processed by `attrs.define`
+gc.collect()
+
+# Aliases for generic preset imports
+KERNEL_FACTORY = BotorchKernelFactory()
+MEAN_FACTORY = BotorchMeanFactory()
+LIKELIHOOD_FACTORY = BotorchLikelihoodFactory()
+FIT_CRITERION_FACTORY = PlainFitCriterionFactory(FitCriterion.MARGINAL_LOG_LIKELIHOOD)

baybe/surrogates/gaussian_process/presets/core.py

@@ -11,6 +11,9 @@ class GaussianProcessPreset(Enum):
     BAYBE = "BAYBE"
     """The default BayBE settings of the Gaussian process surrogate class."""
 
+    BOTORCH = "BOTORCH"
+    """The BoTorch settings."""
+
     CHEN = "CHEN"
     """The adaptive kernel hyperprior settings proposed by :cite:p:`Chen2026`."""