add polynomial kernel

bofire/data_models/kernels/api.py

@@ -13,6 +13,7 @@
     ContinuousKernel,
     LinearKernel,
     MaternKernel,
+    PolynomialKernel,
     RBFKernel,
 )
 from bofire.data_models.kernels.kernel import Kernel
@@ -23,6 +24,7 @@
 AnyContinuousKernel = Union[
     MaternKernel,
     LinearKernel,
+    PolynomialKernel,
     RBFKernel,
 ]
 
@@ -36,6 +38,7 @@
     ScaleKernel,
     HammondDistanceKernel,
     LinearKernel,
+    PolynomialKernel,
     MaternKernel,
     RBFKernel,
     TanimotoKernel,

bofire/data_models/kernels/continuous.py

@@ -24,3 +24,9 @@ class MaternKernel(ContinuousKernel):
 class LinearKernel(ContinuousKernel):
     type: Literal["LinearKernel"] = "LinearKernel"
     variance_prior: Optional[AnyPrior] = None
+
+
+class PolynomialKernel(ContinuousKernel):
+    type: Literal["PolynomialKernel"] = "PolynomialKernel"
+    offset_prior: Optional[AnyPrior] = None
+    power: int = 2

bofire/data_models/surrogates/api.py

@@ -15,6 +15,7 @@
         MixedSingleTaskGPSurrogate,
     )
     from bofire.data_models.surrogates.mlp import MLPEnsemble
+    from bofire.data_models.surrogates.quadratic import QuadraticSurrogate
     from bofire.data_models.surrogates.random_forest import RandomForestSurrogate
     from bofire.data_models.surrogates.single_task_gp import (
         SingleTaskGPHyperconfig,
@@ -36,6 +37,7 @@
         SaasSingleTaskGPSurrogate,
         XGBoostSurrogate,
         LinearSurrogate,
+        QuadraticSurrogate,
         TanimotoGPSurrogate,
     ]
 
@@ -47,6 +49,7 @@
         SaasSingleTaskGPSurrogate,
         XGBoostSurrogate,
         LinearSurrogate,
+        QuadraticSurrogate,
         TanimotoGPSurrogate,
     ]
 except ImportError:

bofire/data_models/surrogates/quadratic.py

@@ -0,0 +1,21 @@
+from typing import Literal
+
+from pydantic import Field
+
+from bofire.data_models.kernels.api import (
+    PolynomialKernel,
+)
+from bofire.data_models.priors.api import BOTORCH_NOISE_PRIOR, AnyPrior
+
+# from bofire.data_models.strategies.api import FactorialStrategy
+from bofire.data_models.surrogates.botorch import BotorchSurrogate
+from bofire.data_models.surrogates.scaler import ScalerEnum
+from bofire.data_models.surrogates.trainable import TrainableSurrogate
+
+
+class QuadraticSurrogate(BotorchSurrogate, TrainableSurrogate):
+    type: Literal["QuadraticSurrogate"] = "QuadraticSurrogate"
+
+    kernel: PolynomialKernel = Field(default_factory=lambda: PolynomialKernel(power=2))
+    noise_prior: AnyPrior = Field(default_factory=lambda: BOTORCH_NOISE_PRIOR())
+    scaler: ScalerEnum = ScalerEnum.NORMALIZE

bofire/kernels/mapper.py

@@ -57,6 +57,22 @@ def map_LinearKernel(
     )
 
 
+def map_PolynomialKernel(
+    data_model: data_models.PolynomialKernel,
+    batch_shape: torch.Size,
+    ard_num_dims: int,
+    active_dims: List[int],
+) -> gpytorch.kernels.PolynomialKernel:
+    return gpytorch.kernels.PolynomialKernel(
+        batch_shape=batch_shape,
+        active_dims=active_dims,
+        power=data_model.power,
+        offset_prior=priors.map(data_model.offset_prior)
+        if data_model.offset_prior is not None
+        else None,
+    )
+
+
 def map_AdditiveKernel(
     data_model: data_models.AdditiveKernel,
     batch_shape: torch.Size,
@@ -131,6 +147,7 @@ def map_TanimotoKernel(
     data_models.RBFKernel: map_RBFKernel,
     data_models.MaternKernel: map_MaternKernel,
     data_models.LinearKernel: map_LinearKernel,
+    data_models.PolynomialKernel: map_PolynomialKernel,
     data_models.AdditiveKernel: map_AdditiveKernel,
     data_models.MultiplicativeKernel: map_MultiplicativeKernel,
     data_models.ScaleKernel: map_ScaleKernel,

bofire/surrogates/mapper.py

@@ -19,6 +19,7 @@
     data_models.SaasSingleTaskGPSurrogate: SaasSingleTaskGPSurrogate,
     data_models.XGBoostSurrogate: XGBoostSurrogate,
     data_models.LinearSurrogate: SingleTaskGPSurrogate,
+    data_models.QuadraticSurrogate: SingleTaskGPSurrogate,
     data_models.TanimotoGPSurrogate: SingleTaskGPSurrogate,
 }
 

tests/bofire/data_models/test_kernels.py

@@ -14,6 +14,7 @@
     LinearKernel,
     MaternKernel,
     MultiplicativeKernel,
+    PolynomialKernel,
     RBFKernel,
     ScaleKernel,
     TanimotoKernel,
@@ -154,6 +155,26 @@ def test_scale_kernel():
     assert hasattr(k, "outputscale_prior") is False
 
 
+def test_poly_kernel():
+    kernel = PolynomialKernel(power=2, offset_prior=BOTORCH_SCALE_PRIOR())
+    k = kernels.map(
+        kernel,
+        batch_shape=torch.Size(),
+        ard_num_dims=10,
+        active_dims=list(range(5)),
+    )
+    assert hasattr(k, "offset_prior")
+    assert isinstance(k.offset_prior, gpytorch.priors.GammaPrior)
+    kernel = PolynomialKernel(power=2)
+    k = kernels.map(
+        kernel,
+        batch_shape=torch.Size(),
+        ard_num_dims=10,
+        active_dims=list(range(5)),
+    )
+    assert hasattr(k, "offset_prior") is False
+
+
 @pytest.mark.parametrize(
     "kernel, ard_num_dims, active_dims, expected_kernel",
     [

tests/bofire/surrogates/test_quadratic.py

@@ -0,0 +1,43 @@
+import numpy as np
+from pandas.testing import assert_frame_equal
+
+import bofire.surrogates.api as surrogates
+from bofire.data_models.domain.api import Inputs, Outputs
+from bofire.data_models.features.api import ContinuousInput, ContinuousOutput
+from bofire.data_models.kernels.api import PolynomialKernel
+from bofire.data_models.surrogates.api import QuadraticSurrogate
+
+
+def test_QuadraticSurrogate():
+    N_EXPERIMENTS = 10
+
+    inputs = Inputs(
+        features=[
+            ContinuousInput(key="a", bounds=(0, 40)),
+            ContinuousInput(key="b", bounds=(20, 60)),
+        ]
+    )
+    outputs = Outputs(features=[ContinuousOutput(key="c")])
+
+    experiments = inputs.sample(N_EXPERIMENTS)
+    experiments["c"] = (
+        experiments["a"] * 2.2
+        + experiments["b"] * -0.05
+        + experiments["b"]
+        + np.random.normal(loc=0, scale=5, size=N_EXPERIMENTS)
+    )
+    experiments["valid_c"] = 1
+
+    surrogate_data = QuadraticSurrogate(inputs=inputs, outputs=outputs)
+    surrogate = surrogates.map(surrogate_data)
+
+    assert isinstance(surrogate, surrogates.SingleTaskGPSurrogate)
+    assert isinstance(surrogate.kernel, PolynomialKernel)
+
+    # check dump
+    surrogate.fit(experiments=experiments)
+    preds = surrogate.predict(experiments)
+    dump = surrogate.dumps()
+    surrogate.loads(dump)
+    preds2 = surrogate.predict(experiments)
+    assert_frame_equal(preds, preds2)