Additions for Fully Bayesian Models (#606)

* add map_saas

* implement new fully bayesian models

* make tests compatible to botorch release

* fix test

* fix exception

* fix tests

* fix tests

Author: jduerholt

bofire/data_models/surrogates/api.py

@@ -11,8 +11,11 @@
     LinearDeterministicSurrogate,
 )
 from bofire.data_models.surrogates.empirical import EmpiricalSurrogate
-from bofire.data_models.surrogates.fully_bayesian import SaasSingleTaskGPSurrogate
+from bofire.data_models.surrogates.fully_bayesian import (
+    FullyBayesianSingleTaskGPSurrogate,
+)
 from bofire.data_models.surrogates.linear import LinearSurrogate
+from bofire.data_models.surrogates.map_saas import AdditiveMapSaasSingleTaskGPSurrogate
 from bofire.data_models.surrogates.mixed_single_task_gp import (
     MixedSingleTaskGPHyperconfig,
     MixedSingleTaskGPSurrogate,
@@ -56,7 +59,7 @@
     MixedTanimotoGPSurrogate,
     ClassificationMLPEnsemble,
     RegressionMLPEnsemble,
-    SaasSingleTaskGPSurrogate,
+    FullyBayesianSingleTaskGPSurrogate,
     XGBoostSurrogate,
     LinearSurrogate,
     PolynomialSurrogate,
@@ -66,6 +69,7 @@
     MultiTaskGPSurrogate,
     SingleTaskIBNNSurrogate,
     PiecewiseLinearGPSurrogate,
+    AdditiveMapSaasSingleTaskGPSurrogate,
 ]
 
 AnyTrainableSurrogate = Union[
@@ -76,13 +80,14 @@
     MixedTanimotoGPSurrogate,
     ClassificationMLPEnsemble,
     RegressionMLPEnsemble,
-    SaasSingleTaskGPSurrogate,
+    FullyBayesianSingleTaskGPSurrogate,
     XGBoostSurrogate,
     LinearSurrogate,
     PolynomialSurrogate,
     SingleTaskIBNNSurrogate,
     TanimotoGPSurrogate,
     PiecewiseLinearGPSurrogate,
+    AdditiveMapSaasSingleTaskGPSurrogate,
 ]
 
 AnyRegressionSurrogate = Union[
@@ -93,7 +98,7 @@
     MixedSingleTaskGPSurrogate,
     MixedTanimotoGPSurrogate,
     RegressionMLPEnsemble,
-    SaasSingleTaskGPSurrogate,
+    FullyBayesianSingleTaskGPSurrogate,
     XGBoostSurrogate,
     LinearSurrogate,
     PolynomialSurrogate,
@@ -102,6 +107,7 @@
     MultiTaskGPSurrogate,
     SingleTaskIBNNSurrogate,
     PiecewiseLinearGPSurrogate,
+    AdditiveMapSaasSingleTaskGPSurrogate,
 ]
 
 AnyClassificationSurrogate = ClassificationMLPEnsemble

bofire/data_models/surrogates/botorch_surrogates.py

@@ -10,7 +10,9 @@
     LinearDeterministicSurrogate,
 )
 from bofire.data_models.surrogates.empirical import EmpiricalSurrogate
-from bofire.data_models.surrogates.fully_bayesian import SaasSingleTaskGPSurrogate
+from bofire.data_models.surrogates.fully_bayesian import (
+    FullyBayesianSingleTaskGPSurrogate,
+)
 from bofire.data_models.surrogates.linear import LinearSurrogate
 from bofire.data_models.surrogates.mixed_single_task_gp import (
     MixedSingleTaskGPSurrogate,
@@ -37,7 +39,7 @@
     MixedTanimotoGPSurrogate,
     RegressionMLPEnsemble,
     ClassificationMLPEnsemble,
-    SaasSingleTaskGPSurrogate,
+    FullyBayesianSingleTaskGPSurrogate,
     TanimotoGPSurrogate,
     LinearSurrogate,
     PolynomialSurrogate,

bofire/data_models/surrogates/fully_bayesian.py

@@ -1,16 +1,33 @@
-from typing import Annotated, Literal, Type
+from typing import Annotated, List, Literal, Type
 
-from pydantic import Field, field_validator
+from pydantic import AfterValidator, Field, field_validator, model_validator
 
 from bofire.data_models.features.api import AnyOutput, ContinuousOutput
 from bofire.data_models.surrogates.trainable_botorch import TrainableBotorchSurrogate
+from bofire.data_models.types import make_unique_validator
 
 
-class SaasSingleTaskGPSurrogate(TrainableBotorchSurrogate):
-    type: Literal["SaasSingleTaskGPSurrogate"] = "SaasSingleTaskGPSurrogate"
+class FullyBayesianSingleTaskGPSurrogate(TrainableBotorchSurrogate):
+    type: Literal["FullyBayesianSingleTaskGPSurrogate"] = (
+        "FullyBayesianSingleTaskGPSurrogate"
+    )
+    model_type: Literal["linear", "saas", "hvarfner"] = "saas"
     warmup_steps: Annotated[int, Field(ge=1)] = 256
     num_samples: Annotated[int, Field(ge=1)] = 128
     thinning: Annotated[int, Field(ge=1)] = 16
+    features_to_warp: Annotated[
+        List[str], AfterValidator(make_unique_validator("Features"))
+    ] = []
+
+    @model_validator(mode="after")
+    def validate_features_to_warp(self):
+        input_keys = self.inputs.get_keys()
+        for feature in self.features_to_warp:
+            if feature not in input_keys:
+                raise ValueError(
+                    f"Feature '{feature}' in features_to_warp is not a valid input key."
+                )
+        return self
 
     @field_validator("thinning")
     @classmethod

bofire/data_models/surrogates/map_saas.py

@@ -0,0 +1,36 @@
+from typing import Literal, Type
+
+from pydantic import PositiveInt
+
+from bofire.data_models.features.api import AnyOutput, ContinuousOutput
+from bofire.data_models.surrogates.trainable_botorch import TrainableBotorchSurrogate
+
+
+class TestSurrogate:
+    pass
+
+
+class AdditiveMapSaasSingleTaskGPSurrogate(TrainableBotorchSurrogate):
+    """Additive MAP SAAS single-task GP
+
+    Maximum-a-posteriori (MAP) version of the sparse axis-aligned subspace
+    `FullyBayesianSingleTaskGPSurrogate` with `model_type` equals to "saas".
+
+    Attributes:
+        n_taus (PositiveInt): Number of sub-kernels to use in the SAAS model.
+    """
+
+    type: Literal["AdditiveMapSaasSingleTaskGPSurrogate"] = (
+        "AdditiveMapSaasSingleTaskGPSurrogate"
+    )
+    n_taus: PositiveInt = 4
+
+    @classmethod
+    def is_output_implemented(cls, my_type: Type[AnyOutput]) -> bool:
+        """Abstract method to check output type for surrogate models
+        Args:
+            my_type: continuous or categorical output
+        Returns:
+            bool: True if the output type is valid for the surrogate chosen, False otherwise
+        """
+        return isinstance(my_type, type(ContinuousOutput))

bofire/surrogates/api.py

@@ -1,6 +1,7 @@
 from bofire.surrogates.botorch_surrogates import BotorchSurrogates
 from bofire.surrogates.deterministic import LinearDeterministicSurrogate
 from bofire.surrogates.empirical import EmpiricalSurrogate
+from bofire.surrogates.map_saas import AdditiveMapSaasSingleTaskGPSurrogate
 from bofire.surrogates.mapper import map
 from bofire.surrogates.mixed_single_task_gp import MixedSingleTaskGPSurrogate
 from bofire.surrogates.mixed_tanimoto_gp import MixedTanimotoGPSurrogate

bofire/surrogates/fully_bayesian.py

@@ -4,19 +4,32 @@
 import pandas as pd
 import torch
 from botorch import fit_fully_bayesian_model_nuts
-from botorch.models.fully_bayesian import SaasFullyBayesianSingleTaskGP
+from botorch.models.fully_bayesian import (
+    FullyBayesianLinearSingleTaskGP,
+    FullyBayesianSingleTaskGP,
+    SaasFullyBayesianSingleTaskGP,
+)
 from botorch.models.transforms.outcome import Standardize
 
 from bofire.data_models.enum import OutputFilteringEnum
-from bofire.data_models.surrogates.api import SaasSingleTaskGPSurrogate as DataModel
+from bofire.data_models.surrogates.api import (
+    FullyBayesianSingleTaskGPSurrogate as DataModel,
+)
 from bofire.data_models.surrogates.scaler import ScalerEnum
 from bofire.surrogates.botorch import BotorchSurrogate
 from bofire.surrogates.trainable import TrainableSurrogate
 from bofire.surrogates.utils import get_scaler
 from bofire.utils.torch_tools import tkwargs
 
 
-class SaasSingleTaskGPSurrogate(BotorchSurrogate, TrainableSurrogate):
+_model_mapper = {
+    "saas": SaasFullyBayesianSingleTaskGP,
+    "linear": FullyBayesianLinearSingleTaskGP,
+    "hvarfner": FullyBayesianSingleTaskGP,
+}
+
+
+class FullyBayesianSingleTaskGPSurrogate(BotorchSurrogate, TrainableSurrogate):
     def __init__(
         self,
         data_model: DataModel,
@@ -27,6 +40,8 @@ def __init__(
         self.thinning = data_model.thinning
         self.scaler = data_model.scaler
         self.output_scaler = data_model.output_scaler
+        self.features_to_warp = data_model.features_to_warp
+        self.model_type = data_model.model_type
         super().__init__(data_model=data_model, **kwargs)
 
     model: Optional[SaasFullyBayesianSingleTaskGP] = None
@@ -41,19 +56,40 @@ def _fit(self, X: pd.DataFrame, Y: pd.DataFrame, disable_progbar: bool = True):
             torch.from_numpy(transformed_X.values).to(**tkwargs),
             torch.from_numpy(Y.values).to(**tkwargs),
         )
+        try:
+            self.model = _model_mapper[self.model_type](
+                train_X=tX,
+                train_Y=tY,
+                outcome_transform=(
+                    Standardize(m=1)
+                    if self.output_scaler == ScalerEnum.STANDARDIZE
+                    else None
+                ),
+                input_transform=scaler,
+                use_input_warping=True if len(self.features_to_warp) > 0 else False,
+                indices_to_warp=self.inputs.get_feature_indices(
+                    self.input_preprocessing_specs, self.features_to_warp
+                )
+                if len(self.features_to_warp) > 0
+                else None,  # type: ignore
+            )
+        except TypeError:
+            # For the current release versions of BoTorch,
+            # the `use_input_warping` argument is not available
+            # we have to wait for the next release
+            self.model = _model_mapper[self.model_type](
+                train_X=tX,
+                train_Y=tY,
+                outcome_transform=(
+                    Standardize(m=1)
+                    if self.output_scaler == ScalerEnum.STANDARDIZE
+                    else None
+                ),
+                input_transform=scaler,
+            )
 
-        self.model = SaasFullyBayesianSingleTaskGP(
-            train_X=tX,
-            train_Y=tY,
-            outcome_transform=(
-                Standardize(m=1)
-                if self.output_scaler == ScalerEnum.STANDARDIZE
-                else None
-            ),
-            input_transform=scaler,
-        )
         fit_fully_bayesian_model_nuts(
-            self.model,
+            self.model,  # type: ignore
             warmup_steps=self.warmup_steps,
             num_samples=self.num_samples,
             thinning=self.thinning,

bofire/surrogates/map_saas.py

@@ -0,0 +1,57 @@
+from typing import Dict, Optional
+
+import pandas as pd
+import torch
+from botorch.fit import fit_gpytorch_mll
+from botorch.models.map_saas import AdditiveMapSaasSingleTaskGP
+from botorch.models.transforms.outcome import Standardize
+from gpytorch.mlls import ExactMarginalLogLikelihood
+
+from bofire.data_models.enum import OutputFilteringEnum
+from bofire.data_models.surrogates.api import (
+    AdditiveMapSaasSingleTaskGPSurrogate as DataModel,
+)
+from bofire.data_models.surrogates.scaler import ScalerEnum
+from bofire.surrogates.botorch import BotorchSurrogate
+from bofire.surrogates.trainable import TrainableSurrogate
+from bofire.surrogates.utils import get_scaler
+from bofire.utils.torch_tools import tkwargs
+
+
+class AdditiveMapSaasSingleTaskGPSurrogate(BotorchSurrogate, TrainableSurrogate):
+    def __init__(
+        self,
+        data_model: DataModel,
+        **kwargs,
+    ):
+        self.n_taus = data_model.n_taus
+        self.scaler = data_model.scaler
+        self.output_scaler = data_model.output_scaler
+        super().__init__(data_model=data_model, **kwargs)
+
+    model: Optional[AdditiveMapSaasSingleTaskGP] = None
+    _output_filtering: OutputFilteringEnum = OutputFilteringEnum.ALL
+    training_specs: Dict = {}
+
+    def _fit(self, X: pd.DataFrame, Y: pd.DataFrame, disable_progbar: bool = True):  # type: ignore
+        scaler = get_scaler(self.inputs, self.input_preprocessing_specs, self.scaler, X)
+        transformed_X = self.inputs.transform(X, self.input_preprocessing_specs)
+
+        tX, tY = (
+            torch.from_numpy(transformed_X.values).to(**tkwargs),
+            torch.from_numpy(Y.values).to(**tkwargs),
+        )
+
+        self.model = AdditiveMapSaasSingleTaskGP(
+            train_X=tX,
+            train_Y=tY,
+            outcome_transform=(
+                Standardize(m=1)
+                if self.output_scaler == ScalerEnum.STANDARDIZE
+                else None
+            ),
+            input_transform=scaler,
+            num_taus=self.n_taus,
+        )
+        mll = ExactMarginalLogLikelihood(self.model.likelihood, self.model)
+        fit_gpytorch_mll(mll, options=self.training_specs, max_attempts=10)

bofire/surrogates/mapper.py

@@ -6,7 +6,8 @@
     LinearDeterministicSurrogate,
 )
 from bofire.surrogates.empirical import EmpiricalSurrogate
-from bofire.surrogates.fully_bayesian import SaasSingleTaskGPSurrogate
+from bofire.surrogates.fully_bayesian import FullyBayesianSingleTaskGPSurrogate
+from bofire.surrogates.map_saas import AdditiveMapSaasSingleTaskGPSurrogate
 from bofire.surrogates.mixed_single_task_gp import MixedSingleTaskGPSurrogate
 from bofire.surrogates.mixed_tanimoto_gp import MixedTanimotoGPSurrogate
 from bofire.surrogates.mlp import ClassificationMLPEnsemble, RegressionMLPEnsemble
@@ -28,7 +29,7 @@
     data_models.MixedTanimotoGPSurrogate: MixedTanimotoGPSurrogate,
     data_models.RegressionMLPEnsemble: RegressionMLPEnsemble,
     data_models.ClassificationMLPEnsemble: ClassificationMLPEnsemble,
-    data_models.SaasSingleTaskGPSurrogate: SaasSingleTaskGPSurrogate,
+    data_models.FullyBayesianSingleTaskGPSurrogate: FullyBayesianSingleTaskGPSurrogate,
     data_models.XGBoostSurrogate: XGBoostSurrogate,
     data_models.LinearSurrogate: SingleTaskGPSurrogate,
     data_models.PolynomialSurrogate: SingleTaskGPSurrogate,
@@ -38,6 +39,7 @@
     data_models.SingleTaskIBNNSurrogate: SingleTaskGPSurrogate,
     data_models.PiecewiseLinearGPSurrogate: PiecewiseLinearGPSurrogate,
     data_models.CategoricalDeterministicSurrogate: CategoricalDeterministicSurrogate,
+    data_models.AdditiveMapSaasSingleTaskGPSurrogate: AdditiveMapSaasSingleTaskGPSurrogate,
 }