kernels working on a given set of features

bofire/data_models/domain/features.py

@@ -593,6 +593,7 @@ def _validate_transform_specs(
         """
         # first check that the keys in the specs dict are correct also correct feature keys
         # next check that all values are of type CategoricalEncodingEnum or MolFeatures
+        checked_keys = set()
         for key, value in specs.items():
             try:
                 feat = self.get_by_key(key)
@@ -622,6 +623,21 @@ def _validate_transform_specs(
                     raise ValueError(
                         f"Forbidden transform type for feature with key {key}",
                     )
+            checked_keys.add(key)
+
+        # now check that features that must be transformed do have a transformation defined
+        for key in self.get_keys():
+            if key in checked_keys:
+                continue
+
+            feat = self.get_by_key(key)
+            if isinstance(feat, MolecularInput):
+                trx = specs.get(key)
+                if trx is None or not isinstance(trx, MolFeatures):
+                    raise ValueError(
+                        "MolecularInput features must have a input processing of type MolFeatures defined"
+                    )
+
         return specs
 
     def get_bounds(

bofire/data_models/kernels/aggregation.py

@@ -3,13 +3,13 @@
 
 from bofire.data_models.kernels.categorical import HammingDistanceKernel
 from bofire.data_models.kernels.continuous import LinearKernel, MaternKernel, RBFKernel
-from bofire.data_models.kernels.kernel import Kernel
+from bofire.data_models.kernels.kernel import AggregationKernel
 from bofire.data_models.kernels.molecular import TanimotoKernel
 from bofire.data_models.kernels.shape import WassersteinKernel
 from bofire.data_models.priors.api import AnyGeneralPrior
 
 
-class AdditiveKernel(Kernel):
+class AdditiveKernel(AggregationKernel):
     type: Literal["AdditiveKernel"] = "AdditiveKernel"
     kernels: Sequence[
         Union[
@@ -26,7 +26,7 @@ class AdditiveKernel(Kernel):
     type: Literal["AdditiveKernel"] = "AdditiveKernel"
 
 
-class MultiplicativeKernel(Kernel):
+class MultiplicativeKernel(AggregationKernel):
     type: Literal["MultiplicativeKernel"] = "MultiplicativeKernel"
     kernels: Sequence[
         Union[
@@ -42,7 +42,7 @@ class MultiplicativeKernel(Kernel):
     ]
 
 
-class ScaleKernel(Kernel):
+class ScaleKernel(AggregationKernel):
     type: Literal["ScaleKernel"] = "ScaleKernel"
     base_kernel: Union[
         RBFKernel,

bofire/data_models/kernels/api.py

@@ -17,12 +17,19 @@
     PolynomialKernel,
     RBFKernel,
 )
-from bofire.data_models.kernels.kernel import Kernel
+from bofire.data_models.kernels.kernel import AggregationKernel, ConcreteKernel, Kernel
 from bofire.data_models.kernels.molecular import MolecularKernel, TanimotoKernel
 from bofire.data_models.kernels.shape import WassersteinKernel
 
 
-AbstractKernel = Union[Kernel, CategoricalKernel, ContinuousKernel, MolecularKernel]
+AbstractKernel = Union[
+    Kernel,
+    CategoricalKernel,
+    ContinuousKernel,
+    MolecularKernel,
+    ConcreteKernel,
+    AggregationKernel,
+]
 
 AnyContinuousKernel = Union[
     MaternKernel,

bofire/data_models/kernels/categorical.py

@@ -1,9 +1,9 @@
 from typing import Literal
 
-from bofire.data_models.kernels.kernel import Kernel
+from bofire.data_models.kernels.kernel import ConcreteKernel
 
 
-class CategoricalKernel(Kernel):
+class CategoricalKernel(ConcreteKernel):
     pass
 
 

bofire/data_models/kernels/continuous.py

@@ -1,12 +1,12 @@
-from typing import Literal, Optional
+from typing import List, Literal, Optional
 
 from pydantic import PositiveInt, field_validator
 
-from bofire.data_models.kernels.kernel import Kernel
+from bofire.data_models.kernels.kernel import ConcreteKernel
 from bofire.data_models.priors.api import AnyGeneralPrior, AnyPrior
 
 
-class ContinuousKernel(Kernel):
+class ContinuousKernel(ConcreteKernel):
     pass
 
 
@@ -40,6 +40,7 @@ class PolynomialKernel(ContinuousKernel):
     power: int = 2
 
 
-class InfiniteWidthBNNKernel(Kernel):
+class InfiniteWidthBNNKernel(ContinuousKernel):
+    features: Optional[List[str]] = None
     type: Literal["InfiniteWidthBNNKernel"] = "InfiniteWidthBNNKernel"
     depth: PositiveInt = 3

bofire/data_models/kernels/kernel.py

@@ -1,5 +1,15 @@
+from typing import List, Optional
+
 from bofire.data_models.base import BaseModel
 
 
 class Kernel(BaseModel):
     type: str
+
+
+class AggregationKernel(Kernel):
+    pass
+
+
+class ConcreteKernel(Kernel):
+    features: Optional[List[str]] = None

bofire/data_models/kernels/molecular.py

@@ -1,9 +1,9 @@
 from typing import Literal
 
-from bofire.data_models.kernels.kernel import Kernel
+from bofire.data_models.kernels.kernel import ConcreteKernel
 
 
-class MolecularKernel(Kernel):
+class MolecularKernel(ConcreteKernel):
     pass
 
 

bofire/kernels/categorical.py

@@ -0,0 +1,25 @@
+import torch
+from gpytorch.kernels.kernel import Kernel
+from torch import Tensor
+
+
+class HammingKernelWithOneHots(Kernel):
+    has_lengthscale = True
+
+    def forward(
+        self,
+        x1: Tensor,
+        x2: Tensor,
+        diag: bool = False,
+        last_dim_is_batch: bool = False,
+    ) -> Tensor:
+        delta = (x1.unsqueeze(-2) - x2.unsqueeze(-3)) ** 2
+        dists = delta / self.lengthscale.unsqueeze(-2)
+        if last_dim_is_batch:
+            dists = dists.transpose(-3, -1)
+
+        dists = dists.sum(-1) / 2
+        res = torch.exp(-dists)
+        if diag:
+            res = torch.diagonal(res, dim1=-1, dim2=-2)
+        return res

bofire/kernels/mapper.py

@@ -1,22 +1,37 @@
-from typing import List
+from typing import Callable, List, Optional
 
 import gpytorch
 import torch
 from botorch.models.kernels.categorical import CategoricalKernel
+from gpytorch.constraints import GreaterThan
 from gpytorch.kernels import Kernel as GpytorchKernel
 
 import bofire.data_models.kernels.api as data_models
 import bofire.priors.api as priors
+from bofire.kernels.categorical import HammingKernelWithOneHots
 from bofire.kernels.fingerprint_kernels.tanimoto_kernel import TanimotoKernel
 from bofire.kernels.shape import WassersteinKernel
 
 
+def _compute_active_dims(
+    data_model: data_models.ConcreteKernel,
+    active_dims: List[int],
+    features_to_idx_mapper: Optional[Callable[[List[str]], List[int]]],
+) -> List[int]:
+    if data_model.features:
+        assert features_to_idx_mapper is not None
+        active_dims = features_to_idx_mapper(data_model.features)
+    return active_dims
+
+
 def map_RBFKernel(
     data_model: data_models.RBFKernel,
     batch_shape: torch.Size,
     ard_num_dims: int,
     active_dims: List[int],
+    features_to_idx_mapper: Optional[Callable[[List[str]], List[int]]],
 ) -> gpytorch.kernels.RBFKernel:
+    active_dims = _compute_active_dims(data_model, active_dims, features_to_idx_mapper)
     return gpytorch.kernels.RBFKernel(
         batch_shape=batch_shape,
         ard_num_dims=len(active_dims) if data_model.ard else None,
@@ -34,7 +49,9 @@ def map_MaternKernel(
     batch_shape: torch.Size,
     ard_num_dims: int,
     active_dims: List[int],
+    features_to_idx_mapper: Optional[Callable[[List[str]], List[int]]],
 ) -> gpytorch.kernels.MaternKernel:
+    active_dims = _compute_active_dims(data_model, active_dims, features_to_idx_mapper)
     return gpytorch.kernels.MaternKernel(
         batch_shape=batch_shape,
         ard_num_dims=len(active_dims) if data_model.ard else None,
@@ -53,6 +70,7 @@ def map_InfiniteWidthBNNKernel(
     batch_shape: torch.Size,
     ard_num_dims: int,
     active_dims: List[int],
+    features_to_idx_mapper: Optional[Callable[[List[str]], List[int]]],
 ) -> "InfiniteWidthBNNKernel":  # type: ignore # noqa: F821
     try:
         from botorch.models.kernels.infinite_width_bnn import (  # type: ignore
@@ -66,6 +84,7 @@ def map_InfiniteWidthBNNKernel(
             "requires python 3.10+.",
         )
 
+    active_dims = _compute_active_dims(data_model, active_dims, features_to_idx_mapper)
     return InfiniteWidthBNNKernel(
         batch_shape=batch_shape,
         active_dims=tuple(active_dims),
@@ -78,7 +97,9 @@ def map_LinearKernel(
     batch_shape: torch.Size,
     ard_num_dims: int,
     active_dims: List[int],
+    features_to_idx_mapper: Optional[Callable[[List[str]], List[int]]],
 ) -> gpytorch.kernels.LinearKernel:
+    active_dims = _compute_active_dims(data_model, active_dims, features_to_idx_mapper)
     return gpytorch.kernels.LinearKernel(
         batch_shape=batch_shape,
         active_dims=active_dims,
@@ -95,7 +116,9 @@ def map_PolynomialKernel(
     batch_shape: torch.Size,
     ard_num_dims: int,
     active_dims: List[int],
+    features_to_idx_mapper: Optional[Callable[[List[str]], List[int]]],
 ) -> gpytorch.kernels.PolynomialKernel:
+    active_dims = _compute_active_dims(data_model, active_dims, features_to_idx_mapper)
     return gpytorch.kernels.PolynomialKernel(
         batch_shape=batch_shape,
         active_dims=active_dims,
@@ -113,6 +136,7 @@ def map_AdditiveKernel(
     batch_shape: torch.Size,
     ard_num_dims: int,
     active_dims: List[int],
+    features_to_idx_mapper: Optional[Callable[[List[str]], List[int]]],
 ) -> gpytorch.kernels.AdditiveKernel:
     return gpytorch.kernels.AdditiveKernel(
         *[  # type: ignore
@@ -121,6 +145,7 @@ def map_AdditiveKernel(
                 batch_shape=batch_shape,
                 ard_num_dims=ard_num_dims,
                 active_dims=active_dims,
+                features_to_idx_mapper=features_to_idx_mapper,
             )
             for k in data_model.kernels
         ],
@@ -132,6 +157,7 @@ def map_MultiplicativeKernel(
     batch_shape: torch.Size,
     ard_num_dims: int,
     active_dims: List[int],
+    features_to_idx_mapper: Optional[Callable[[List[str]], List[int]]],
 ) -> gpytorch.kernels.ProductKernel:
     return gpytorch.kernels.ProductKernel(
         *[  # type: ignore
@@ -140,6 +166,7 @@ def map_MultiplicativeKernel(
                 batch_shape=batch_shape,
                 ard_num_dims=ard_num_dims,
                 active_dims=active_dims,
+                features_to_idx_mapper=features_to_idx_mapper,
             )
             for k in data_model.kernels
         ],
@@ -151,13 +178,15 @@ def map_ScaleKernel(
     batch_shape: torch.Size,
     ard_num_dims: int,
     active_dims: List[int],
+    features_to_idx_mapper: Optional[Callable[[List[str]], List[int]]],
 ) -> gpytorch.kernels.ScaleKernel:
     return gpytorch.kernels.ScaleKernel(
         base_kernel=map(
             data_model.base_kernel,
             batch_shape=batch_shape,
             ard_num_dims=ard_num_dims,
             active_dims=active_dims,
+            features_to_idx_mapper=features_to_idx_mapper,
         ),
         outputscale_prior=(
             priors.map(data_model.outputscale_prior)
@@ -172,7 +201,9 @@ def map_TanimotoKernel(
     batch_shape: torch.Size,
     ard_num_dims: int,
     active_dims: List[int],
+    features_to_idx_mapper: Optional[Callable[[List[str]], List[int]]],
 ) -> TanimotoKernel:
+    active_dims = _compute_active_dims(data_model, active_dims, features_to_idx_mapper)
     return TanimotoKernel(
         batch_shape=batch_shape,
         ard_num_dims=len(active_dims) if data_model.ard else None,
@@ -185,19 +216,41 @@ def map_HammingDistanceKernel(
     batch_shape: torch.Size,
     ard_num_dims: int,
     active_dims: List[int],
-) -> CategoricalKernel:
-    return CategoricalKernel(
-        batch_shape=batch_shape,
-        ard_num_dims=len(active_dims) if data_model.ard else None,
-        active_dims=active_dims,  # type: ignore
-    )
+    features_to_idx_mapper: Optional[Callable[[List[str]], List[int]]],
+) -> GpytorchKernel:
+    active_dims = _compute_active_dims(data_model, active_dims, features_to_idx_mapper)
+
+    with_one_hots = data_model.features is not None and len(active_dims) > 1
+    if with_one_hots and len(active_dims) == 1:
+        raise RuntimeError(
+            "only one feature for categorical kernel operating on one-hot features"
+        )
+    elif not with_one_hots and len(active_dims) > 1:
+        # this is not necessarily an issue since botorch's CategoricalKernel
+        # can work on multiple features at the same time
+        pass
+
+    if with_one_hots:
+        return HammingKernelWithOneHots(
+            batch_shape=batch_shape,
+            ard_num_dims=len(active_dims) if data_model.ard else None,
+            active_dims=active_dims,  # type: ignore
+            lengthscale_constraint=GreaterThan(1e-06),
+        )
+    else:
+        return CategoricalKernel(
+            batch_shape=batch_shape,
+            ard_num_dims=len(active_dims) if data_model.ard else None,
+            active_dims=active_dims,  # type: ignore
+        )
 
 
 def map_WassersteinKernel(
     data_model: data_models.WassersteinKernel,
     batch_shape: torch.Size,
     ard_num_dims: int,
     active_dims: List[int],
+    features_to_idx_mapper: Optional[Callable[[List[str]], List[int]]],
 ) -> WassersteinKernel:
     return WassersteinKernel(
         squared=data_model.squared,
@@ -230,10 +283,12 @@ def map(
     batch_shape: torch.Size,
     ard_num_dims: int,
     active_dims: List[int],
+    features_to_idx_mapper: Optional[Callable[[List[str]], List[int]]],
 ) -> GpytorchKernel:
     return KERNEL_MAP[data_model.__class__](
         data_model,
         batch_shape,
         ard_num_dims,
         active_dims,
+        features_to_idx_mapper,
     )