Make BotorchKernelFactory support parameter selection

Author: AdrianSosic

baybe/surrogates/gaussian_process/presets/botorch.py

@@ -3,61 +3,93 @@
 from __future__ import annotations
 
 import gc
-from typing import TYPE_CHECKING
+from itertools import chain
+from typing import TYPE_CHECKING, ClassVar
 
 from attrs import define
-from gpytorch.kernels import Kernel as GPyTorchKernel
 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._gpytorch import (
     make_botorch_multitask_likelihood,
 )
 from baybe.surrogates.gaussian_process.components.kernel import (
     ICMKernelFactory,
-    KernelFactoryProtocol,
+    _KernelFactory,
 )
 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(KernelFactoryProtocol):
+class BotorchKernelFactory(_KernelFactory):
     """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 __call__(
+    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,
         )
 
-        if searchspace.n_tasks == 1:
+        parameter_names = self.get_parameter_names(searchspace)
+
+        # Resolve parameter names to active dimension indices
+        active_dims: list[int] | None
+        if parameter_names is not None:
+            active_dims = list(
+                chain.from_iterable(
+                    searchspace.get_comp_rep_parameter_indices(name)
+                    for name in parameter_names
+                )
+            )
+            ard_num_dims = len(active_dims)
+        else:
+            active_dims = None
+            ard_num_dims = len(searchspace.comp_rep_columns)
+
+        # Determine if the selected parameters include a task parameter
+        task_idx = searchspace.task_idx
+        is_multitask = task_idx is not None and (
+            active_dims is None or task_idx in active_dims
+        )
+
+        if not is_multitask:
             return get_covar_module_with_dim_scaled_prior(
-                ard_num_dims=len(searchspace.comp_rep_columns), active_dims=None
+                ard_num_dims=ard_num_dims, active_dims=active_dims
             )
 
-        assert searchspace.task_idx is not None
+        assert task_idx is not None
         base_idcs = [
             idx
-            for idx in range(len(searchspace.comp_rep_columns))
-            if idx != searchspace.task_idx
+            for idx in (active_dims or range(len(searchspace.comp_rep_columns)))
+            if idx != task_idx
         ]
         base = get_covar_module_with_dim_scaled_prior(
             ard_num_dims=len(base_idcs), active_dims=base_idcs
         )
         index_kernel = PositiveIndexKernel(
             num_tasks=searchspace.n_tasks,
             rank=searchspace.n_tasks,
-            active_dims=[searchspace.task_idx],
+            active_dims=[task_idx],
         )
         return ICMKernelFactory(base, index_kernel)(searchspace, train_x, train_y)