add output_scaler to mlp (#305)

Co-authored-by: Simon <simon.sung@evonik.com>

bofire/data_models/surrogates/mlp.py

@@ -1,6 +1,6 @@
 from typing import Annotated, Literal, Sequence
 
-from pydantic import Field
+from pydantic import Field, validator
 
 from bofire.data_models.surrogates.botorch import BotorchSurrogate
 from bofire.data_models.surrogates.scaler import ScalerEnum
@@ -20,3 +20,22 @@ class MLPEnsemble(BotorchSurrogate, TrainableSurrogate):
     subsample_fraction: Annotated[float, Field(gt=0.0)] = 1.0
     shuffle: bool = True
     scaler: ScalerEnum = ScalerEnum.NORMALIZE
+    output_scaler: ScalerEnum = ScalerEnum.STANDARDIZE
+
+    @validator("output_scaler")
+    def validate_output_scaler(cls, output_scaler):
+        """validates that output_scaler is a valid type
+
+        Args:
+            output_scaler (ScalerEnum): Scaler used to transform the output
+
+        Raises:
+            ValueError: when ScalerEnum.NORMALIZE is used
+
+        Returns:
+            ScalerEnum: Scaler used to transform the output
+        """
+        if output_scaler == ScalerEnum.NORMALIZE:
+            raise ValueError("Normalize is not supported as an output transform.")
+
+        return output_scaler

bofire/surrogates/mlp.py

@@ -5,11 +5,13 @@
 import torch
 import torch.nn as nn
 from botorch.models.ensemble import EnsembleModel
+from botorch.models.transforms.outcome import OutcomeTransform, Standardize
 from torch import Tensor
 from torch.utils.data import DataLoader, Dataset
 
 from bofire.data_models.enum import OutputFilteringEnum
 from bofire.data_models.surrogates.api import MLPEnsemble as DataModel
+from bofire.data_models.surrogates.scaler import ScalerEnum
 from bofire.surrogates.botorch import BotorchSurrogate
 from bofire.surrogates.single_task_gp import get_scaler
 from bofire.surrogates.trainable import TrainableSurrogate
@@ -74,7 +76,9 @@ def forward(self, x):
 
 
 class _MLPEnsemble(EnsembleModel):
-    def __init__(self, mlps: Sequence[MLP]):
+    def __init__(
+        self, mlps: Sequence[MLP], output_scaler: Optional[OutcomeTransform] = None
+    ):
         super().__init__()
         if len(mlps) == 0:
             raise ValueError("List of mlps is empty.")
@@ -84,6 +88,8 @@ def __init__(self, mlps: Sequence[MLP]):
             assert mlp.layers[0].in_features == num_in_features
             assert mlp.layers[-1].out_features == num_out_features
         self.mlps = mlps
+        if output_scaler is not None:
+            self.outcome_transform = output_scaler
         # put all models in eval mode
         for mlp in self.mlps:
             mlp.eval()
@@ -170,6 +176,7 @@ def __init__(self, data_model: DataModel, **kwargs):
         self.subsample_fraction = data_model.subsample_fraction
         self.shuffle = data_model.shuffle
         self.scaler = data_model.scaler
+        self.output_scaler = data_model.output_scaler
         super().__init__(data_model, **kwargs)
 
     _output_filtering: OutputFilteringEnum = OutputFilteringEnum.ALL
@@ -179,6 +186,11 @@ def _fit(self, X: pd.DataFrame, Y: pd.DataFrame):
         scaler = get_scaler(self.inputs, self.input_preprocessing_specs, self.scaler, X)
         transformed_X = self.inputs.transform(X, self.input_preprocessing_specs)
 
+        if self.output_scaler == ScalerEnum.STANDARDIZE:
+            output_scaler = Standardize(m=Y.shape[-1])
+        else:
+            output_scaler = None
+
         mlps = []
         subsample_size = round(self.subsample_fraction * X.shape[0])
         for _ in range(self.n_estimators):
@@ -189,7 +201,7 @@ def _fit(self, X: pd.DataFrame, Y: pd.DataFrame):
 
             dataset = RegressionDataSet(
                 X=scaler.transform(tX) if scaler is not None else tX,
-                y=ty,
+                y=output_scaler(ty)[0] if output_scaler is not None else ty,
             )
             mlp = MLP(
                 input_size=transformed_X.shape[1],
@@ -208,6 +220,6 @@ def _fit(self, X: pd.DataFrame, Y: pd.DataFrame):
                 weight_decay=self.weight_decay,
             )
             mlps.append(mlp)
-        self.model = _MLPEnsemble(mlps=mlps)
+        self.model = _MLPEnsemble(mlps, output_scaler=output_scaler)
         if scaler is not None:
             self.model.input_transform = scaler

tests/bofire/data_models/specs/surrogates.py

@@ -176,6 +176,7 @@
         "subsample_fraction": 1.0,
         "shuffle": True,
         "scaler": ScalerEnum.NORMALIZE,
+        "output_scaler": ScalerEnum.STANDARDIZE,
         "input_preprocessing_specs": {},
         "dump": None,
         "hyperconfig": None,

tests/bofire/surrogates/test_mlp.py

@@ -2,6 +2,7 @@
 import torch
 import torch.nn as nn
 from botorch.models.transforms.input import InputStandardize, Normalize
+from botorch.models.transforms.outcome import Standardize
 from pandas.testing import assert_frame_equal
 
 import bofire.surrogates.api as surrogates
@@ -163,9 +164,14 @@ def test_mlp_ensemble_forward():
 
 
 @pytest.mark.parametrize(
-    "scaler", [ScalerEnum.NORMALIZE, ScalerEnum.STANDARDIZE, ScalerEnum.IDENTITY]
+    "scaler, output_scaler",
+    [
+        [ScalerEnum.NORMALIZE, ScalerEnum.IDENTITY],
+        [ScalerEnum.STANDARDIZE, ScalerEnum.STANDARDIZE],
+        [ScalerEnum.IDENTITY, ScalerEnum.STANDARDIZE],
+    ],
 )
-def test_mlp_ensemble_fit(scaler):
+def test_mlp_ensemble_fit(scaler, output_scaler):
     bench = Himmelblau()
     samples = bench.domain.inputs.sample(10)
     experiments = bench.f(samples, return_complete=True)
@@ -175,17 +181,25 @@ def test_mlp_ensemble_fit(scaler):
         n_estimators=2,
         n_epochs=5,
         scaler=scaler,
+        output_scaler=output_scaler,
     )
     surrogate = surrogates.map(ens)
 
     surrogate.fit(experiments=experiments)
+
     if scaler == ScalerEnum.NORMALIZE:
         assert isinstance(surrogate.model.input_transform, Normalize)
     elif scaler == ScalerEnum.STANDARDIZE:
         assert isinstance(surrogate.model.input_transform, InputStandardize)
     else:
         with pytest.raises(AttributeError):
             assert surrogate.model.input_transform is None
+
+    if output_scaler == ScalerEnum.STANDARDIZE:
+        assert isinstance(surrogate.model.outcome_transform, Standardize)
+    elif output_scaler == ScalerEnum.IDENTITY:
+        assert not hasattr(surrogate.model, "outcome_transform")
+
     preds = surrogate.predict(experiments)
     dump = surrogate.dumps()
     surrogate2 = surrogates.map(ens)