Feature/ensemble map saas single task gp

bofire/data_models/surrogates/api.py

@@ -15,7 +15,7 @@
     FullyBayesianSingleTaskGPSurrogate,
 )
 from bofire.data_models.surrogates.linear import LinearSurrogate
-from bofire.data_models.surrogates.map_saas import AdditiveMapSaasSingleTaskGPSurrogate
+from bofire.data_models.surrogates.map_saas import AdditiveMapSaasSingleTaskGPSurrogate, EnsembleMapSaasSingleTaskGPSurrogate
 from bofire.data_models.surrogates.mixed_single_task_gp import (
     MixedSingleTaskGPHyperconfig,
     MixedSingleTaskGPSurrogate,
@@ -75,6 +75,7 @@
     SingleTaskIBNNSurrogate,
     PiecewiseLinearGPSurrogate,
     AdditiveMapSaasSingleTaskGPSurrogate,
+    EnsembleMapSaasSingleTaskGPSurrogate,
 ]
 
 AnyTrainableSurrogate = Union[
@@ -92,6 +93,7 @@
     TanimotoGPSurrogate,
     PiecewiseLinearGPSurrogate,
     AdditiveMapSaasSingleTaskGPSurrogate,
+    EnsembleMapSaasSingleTaskGPSurrogate,
 ]
 
 AnyRegressionSurrogate = Union[
@@ -111,6 +113,7 @@
     SingleTaskIBNNSurrogate,
     PiecewiseLinearGPSurrogate,
     AdditiveMapSaasSingleTaskGPSurrogate,
+    EnsembleMapSaasSingleTaskGPSurrogate,
 ]
 
 AnyClassificationSurrogate = ClassificationMLPEnsemble

bofire/data_models/surrogates/map_saas.py

@@ -34,3 +34,45 @@ def is_output_implemented(cls, my_type: Type[AnyOutput]) -> bool:
             bool: True if the output type is valid for the surrogate chosen, False otherwise
         """
         return isinstance(my_type, type(ContinuousOutput))
+
+class EnsembleMapSaasSingleTaskGPSurrogate(TrainableBotorchSurrogate):
+    """Instantiates an ``EnsembleMapSaasSingleTaskGP``, which is a batched
+    ensemble of ``SingleTaskGP``s with the Matern-5/2 kernel and a SAAS prior.
+    The model is intended to be trained with ``ExactMarginalLogLikelihood`` and
+    ``fit_gpytorch_mll``. Under the hood, the model is equivalent to a
+    multi-output ``BatchedMultiOutputGPyTorchModel``, but it produces a
+    ``MixtureGaussiaPosterior``, which leads to ensembling of the model outputs.
+
+    Args:
+        train_X: An `n x d` tensor of training features.
+        train_Y: An `n x 1` tensor of training observations.
+        train_Yvar: An optional `n x 1` tensor of observed measurement noise.
+        num_taus: The number of taus to use (4 if omitted). Each tau is
+            a sparsity parameter for the corresponding kernel in the ensemble.
+        taus: An optional tensor of shape `num_taus` containing the taus to use.
+            If omitted, the taus are sampled from a HalfCauchy(0.1) distribution.
+        outcome_transform: An outcome transform that is applied to the
+            training data during instantiation and to the posterior during
+            inference (that is, the `Posterior` obtained by calling
+            `.posterior` on the model will be on the original scale). We use a
+            `Standardize` transform if no `outcome_transform` is specified.
+            Pass down `None` to use no outcome transform. Note that `.train()` will
+            be called on the outcome transform during instantiation of the model.
+        input_transform: An input transform that is applied in the model's
+            forward pass.
+    """
+
+    type: Literal["EnsembleMapSaasSingleTaskGPSurrogate"] = (
+        "EnsembleMapSaasSingleTaskGPSurrogate"
+    )
+    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,7 +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.map_saas import AdditiveMapSaasSingleTaskGPSurrogate, EnsembleMapSaasSingleTaskGPSurrogate
 from bofire.surrogates.mapper import map
 from bofire.surrogates.mlp import (
     ClassificationMLPEnsemble,

bofire/surrogates/map_saas.py

@@ -1,17 +1,23 @@
 from typing import Dict, Optional
 
+import numpy as np
+import pandas as pd
 import torch
 from botorch.fit import fit_gpytorch_mll
-from botorch.models.map_saas import AdditiveMapSaasSingleTaskGP
+from botorch.models.map_saas import (
+    AdditiveMapSaasSingleTaskGP,
+    EnsembleMapSaasSingleTaskGP,
+)
 from botorch.models.transforms.input import InputTransform
-from botorch.models.transforms.outcome import OutcomeTransform
+from botorch.models.transforms.outcome import OutcomeTransform, 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.surrogates.botorch import TrainableBotorchSurrogate
+from bofire.surrogates.botorch import BotorchSurrogate, TrainableBotorchSurrogate
+from bofire.utils.torch_tools import tkwargs
 
 
 class AdditiveMapSaasSingleTaskGPSurrogate(TrainableBotorchSurrogate):
@@ -46,3 +52,54 @@ def _fit_botorch(
         )
         mll = ExactMarginalLogLikelihood(self.model.likelihood, self.model)
         fit_gpytorch_mll(mll, options=self.training_specs, max_attempts=50)
+
+
+class EnsembleMapSaasSingleTaskGPSurrogate(TrainableBotorchSurrogate):
+    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[EnsembleMapSaasSingleTaskGP] = None
+    _output_filtering: OutputFilteringEnum = OutputFilteringEnum.ALL
+    training_specs: Dict = {}
+
+    def _fit_botorch(
+        self,
+        tX: torch.Tensor,
+        tY: torch.Tensor,
+        input_transform: Optional[InputTransform] = None,
+        outcome_transform: Optional[OutcomeTransform] = None,
+        **kwargs,
+    ):
+        # EnsembleMapSaasSingleTaskGP repeats the data to create a batch dimension
+        # The outcome_transform needs to have the correct batch_shape
+        if isinstance(outcome_transform, Standardize):
+            outcome_transform = Standardize(
+                m=tY.shape[-1],
+                batch_shape=torch.Size([self.n_taus]),
+            )
+        self.model = EnsembleMapSaasSingleTaskGP(
+            train_X=tX,
+            train_Y=tY,
+            outcome_transform=outcome_transform,
+            input_transform=input_transform,
+            num_taus=self.n_taus,
+        )
+        mll = ExactMarginalLogLikelihood(self.model.likelihood, self.model)
+        fit_gpytorch_mll(mll, options=self.training_specs, max_attempts=50)
+
+    def _predict(self, transformed_X: pd.DataFrame):
+        # transform to tensor
+        X = torch.from_numpy(transformed_X.values).to(**tkwargs)
+        with torch.no_grad():
+            posterior = self.model.posterior(X=X, observation_noise=True)  # type: ignore
+
+        preds = posterior.mixture_mean.detach().numpy()
+        stds = np.sqrt(posterior.mixture_variance.detach().numpy())
+        return preds, stds

bofire/surrogates/mapper.py

@@ -12,7 +12,7 @@
 )
 from bofire.surrogates.empirical import EmpiricalSurrogate
 from bofire.surrogates.fully_bayesian import FullyBayesianSingleTaskGPSurrogate
-from bofire.surrogates.map_saas import AdditiveMapSaasSingleTaskGPSurrogate
+from bofire.surrogates.map_saas import AdditiveMapSaasSingleTaskGPSurrogate, EnsembleMapSaasSingleTaskGPSurrogate
 from bofire.surrogates.mlp import ClassificationMLPEnsemble, RegressionMLPEnsemble
 from bofire.surrogates.multi_task_gp import MultiTaskGPSurrogate
 from bofire.surrogates.random_forest import RandomForestSurrogate
@@ -68,7 +68,6 @@ def map_MixedSingleTaskGPSurrogate(
         kernel=kernel,
     )
 
-
 SURROGATE_MAP: Dict[Type[data_models.Surrogate], Type[Surrogate]] = {
     data_models.EmpiricalSurrogate: EmpiricalSurrogate,
     data_models.RandomForestSurrogate: RandomForestSurrogate,
@@ -87,6 +86,7 @@ def map_MixedSingleTaskGPSurrogate(
     data_models.PiecewiseLinearGPSurrogate: PiecewiseLinearGPSurrogate,
     data_models.CategoricalDeterministicSurrogate: CategoricalDeterministicSurrogate,
     data_models.AdditiveMapSaasSingleTaskGPSurrogate: AdditiveMapSaasSingleTaskGPSurrogate,
+    data_models.EnsembleMapSaasSingleTaskGPSurrogate: EnsembleMapSaasSingleTaskGPSurrogate,
 }
 
 

tests/bofire/data_models/specs/surrogates.py

@@ -215,6 +215,30 @@
     },
 )
 
+specs.add_valid(
+    models.EnsembleMapSaasSingleTaskGPSurrogate,
+    lambda: {
+        "inputs": Inputs(
+            features=[
+                features.valid(ContinuousInput).obj(),
+            ],
+        ).model_dump(),
+        "outputs": Outputs(
+            features=[
+                features.valid(ContinuousOutput).obj(),
+            ],
+        ).model_dump(),
+        "aggregations": None,
+        "n_taus": 4,
+        "scaler": ScalerEnum.NORMALIZE,
+        "output_scaler": ScalerEnum.STANDARDIZE,
+        "input_preprocessing_specs": {},
+        "categorical_encodings": {},
+        "hyperconfig": None,
+        "dump": None,
+    },
+)
+
 specs.add_valid(
     models.FullyBayesianSingleTaskGPSurrogate,
     lambda: {

tests/bofire/surrogates/test_map_saas.py

@@ -1,9 +1,15 @@
-from botorch.models.map_saas import AdditiveMapSaasSingleTaskGP
+from botorch.models.map_saas import (
+    AdditiveMapSaasSingleTaskGP,
+    EnsembleMapSaasSingleTaskGP,
+)
 from pandas.testing import assert_frame_equal
 
 import bofire.surrogates.api as surrogates
 from bofire.benchmarks.single import Himmelblau
-from bofire.data_models.surrogates.api import AdditiveMapSaasSingleTaskGPSurrogate
+from bofire.data_models.surrogates.api import (
+    AdditiveMapSaasSingleTaskGPSurrogate,
+    EnsembleMapSaasSingleTaskGPSurrogate,
+)
 
 
 def test_AdditiveMapSaasSingleTaskGPSurrogate():
@@ -24,3 +30,23 @@ def test_AdditiveMapSaasSingleTaskGPSurrogate():
     assert preds.shape == (10, 2)
     preds2 = gp.predict(experiments)
     assert_frame_equal(preds, preds2)
+
+
+def test_EnsembleMapSaasSingleTaskGPSurrogate():
+    bench = Himmelblau()
+    samples = bench.domain.inputs.sample(10)
+    experiments = bench.f(samples, return_complete=True)
+    data_model = EnsembleMapSaasSingleTaskGPSurrogate(
+        inputs=bench.domain.inputs,
+        outputs=bench.domain.outputs,
+    )
+    gp = surrogates.map(data_model)
+    gp.fit(experiments=experiments)
+    assert isinstance(gp.model, EnsembleMapSaasSingleTaskGP)
+    dump = gp.dumps()
+    gp2 = surrogates.map(data_model=data_model)
+    gp2.loads(dump)
+    preds = gp.predict(experiments)
+    assert preds.shape == (10, 2)
+    preds2 = gp.predict(experiments)
+    assert_frame_equal(preds, preds2)