diff --git a/bofire/data_models/kernels/api.py b/bofire/data_models/kernels/api.py index 4f9eaaa63..975be871a 100644 --- a/bofire/data_models/kernels/api.py +++ b/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, diff --git a/bofire/data_models/kernels/continuous.py b/bofire/data_models/kernels/continuous.py index 6c68e6391..57a5da643 100644 --- a/bofire/data_models/kernels/continuous.py +++ b/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 diff --git a/bofire/data_models/surrogates/api.py b/bofire/data_models/surrogates/api.py index 3ba5749e7..eaf44e852 100644 --- a/bofire/data_models/surrogates/api.py +++ b/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: diff --git a/bofire/data_models/surrogates/quadratic.py b/bofire/data_models/surrogates/quadratic.py new file mode 100644 index 000000000..8677acfd0 --- /dev/null +++ b/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 diff --git a/bofire/kernels/mapper.py b/bofire/kernels/mapper.py index 33bd0c9c5..91f117636 100644 --- a/bofire/kernels/mapper.py +++ b/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, diff --git a/bofire/surrogates/mapper.py b/bofire/surrogates/mapper.py index 5be8d4f3e..88c7616f1 100644 --- a/bofire/surrogates/mapper.py +++ b/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, } diff --git a/tests/bofire/data_models/test_kernels.py b/tests/bofire/data_models/test_kernels.py index 1f7936e76..c2d5d7ed1 100644 --- a/tests/bofire/data_models/test_kernels.py +++ b/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", [ diff --git a/tests/bofire/surrogates/test_quadratic.py b/tests/bofire/surrogates/test_quadratic.py new file mode 100644 index 000000000..149c4989d --- /dev/null +++ b/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)