Make method for surrogates

bofire/data_models/surrogates/bnn.py

@@ -1,11 +1,8 @@
-from typing import Literal, Optional
+from typing import Literal
 
 from bofire.data_models.kernels.api import InfiniteWidthBNNKernel
-from bofire.data_models.surrogates.single_task_gp import SingleTaskGPSurrogate
-from bofire.data_models.surrogates.trainable import Hyperconfig
+from bofire.data_models.surrogates.single_task_gp import BaseSingleTaskGPSurrogate
 
 
-class SingleTaskIBNNSurrogate(SingleTaskGPSurrogate):
+class SingleTaskIBNNSurrogate(BaseSingleTaskGPSurrogate[InfiniteWidthBNNKernel]):
     type: Literal["SingleTaskIBNNSurrogate"] = "SingleTaskIBNNSurrogate"
-    kernel: InfiniteWidthBNNKernel = InfiniteWidthBNNKernel()
-    hyperconfig: Optional[Hyperconfig] = None

bofire/data_models/surrogates/fully_bayesian.py

@@ -3,11 +3,12 @@
 from pydantic import AfterValidator, Field, field_validator, model_validator
 
 from bofire.data_models.features.api import AnyOutput, ContinuousOutput
+from bofire.data_models.surrogates.trainable import Hyperconfig
 from bofire.data_models.surrogates.trainable_botorch import TrainableBotorchSurrogate
 from bofire.data_models.types import make_unique_validator
 
 
-class FullyBayesianSingleTaskGPSurrogate(TrainableBotorchSurrogate):
+class FullyBayesianSingleTaskGPSurrogate(TrainableBotorchSurrogate[Hyperconfig]):
     type: Literal["FullyBayesianSingleTaskGPSurrogate"] = (
         "FullyBayesianSingleTaskGPSurrogate"
     )

bofire/data_models/surrogates/map_saas.py

@@ -3,14 +3,15 @@
 from pydantic import PositiveInt
 
 from bofire.data_models.features.api import AnyOutput, ContinuousOutput
+from bofire.data_models.surrogates.trainable import Hyperconfig
 from bofire.data_models.surrogates.trainable_botorch import TrainableBotorchSurrogate
 
 
 class TestSurrogate:
     pass
 
 
-class AdditiveMapSaasSingleTaskGPSurrogate(TrainableBotorchSurrogate):
+class AdditiveMapSaasSingleTaskGPSurrogate(TrainableBotorchSurrogate[Hyperconfig]):
     """Additive MAP SAAS single-task GP
 
     Maximum-a-posteriori (MAP) version of the sparse axis-aligned subspace

bofire/data_models/surrogates/mixed_single_task_gp.py

@@ -102,7 +102,9 @@ def _update_hyperparameters(
             raise ValueError(f"Kernel {hyperparameters.kernel} not known.")
 
 
-class MixedSingleTaskGPSurrogate(TrainableBotorchSurrogate):
+class MixedSingleTaskGPSurrogate(
+    TrainableBotorchSurrogate[MixedSingleTaskGPHyperconfig]
+):
     type: Literal["MixedSingleTaskGPSurrogate"] = "MixedSingleTaskGPSurrogate"
     continuous_kernel: AnyContinuousKernel = Field(
         default_factory=lambda: RBFKernel(

bofire/data_models/surrogates/multi_task_gp.py

@@ -86,7 +86,7 @@ def matern_15(ard: bool, lengthscale_prior: AnyPrior) -> MaternKernel:
             raise ValueError(f"Kernel {hyperparameters.kernel} not known.")
 
 
-class MultiTaskGPSurrogate(TrainableBotorchSurrogate):
+class MultiTaskGPSurrogate(TrainableBotorchSurrogate[MultiTaskGPHyperconfig]):
     type: Literal["MultiTaskGPSurrogate"] = "MultiTaskGPSurrogate"
     kernel: AnyKernel = Field(
         default_factory=lambda: MaternKernel(

bofire/data_models/surrogates/random_forest.py

@@ -3,10 +3,11 @@
 from pydantic import Field
 
 from bofire.data_models.features.api import AnyOutput, ContinuousOutput
+from bofire.data_models.surrogates.trainable import Hyperconfig
 from bofire.data_models.surrogates.trainable_botorch import TrainableBotorchSurrogate
 
 
-class RandomForestSurrogate(TrainableBotorchSurrogate):
+class RandomForestSurrogate(TrainableBotorchSurrogate[Hyperconfig]):
     type: Literal["RandomForestSurrogate"] = "RandomForestSurrogate"
 
     # hyperparams passed down to `RandomForestRegressor`

bofire/data_models/surrogates/robust_single_task_gp.py

@@ -15,7 +15,7 @@
 from bofire.data_models.surrogates.trainable_botorch import TrainableBotorchSurrogate
 
 
-class RobustSingleTaskGPSurrogate(TrainableBotorchSurrogate):
+class RobustSingleTaskGPSurrogate(TrainableBotorchSurrogate[SingleTaskGPHyperconfig]):
     """
     Robust Relevance Pursuit Single Task Gaussian Process Surrogate.
 

bofire/data_models/surrogates/shape.py

@@ -89,7 +89,9 @@ def _update_hyperparameters(
             raise ValueError(f"Kernel {hyperparameters.kernel} not known.")
 
 
-class PiecewiseLinearGPSurrogate(TrainableBotorchSurrogate):
+class PiecewiseLinearGPSurrogate(
+    TrainableBotorchSurrogate[PiecewiseLinearGPSurrogateHyperconfig]
+):
     """GP surrogate that is based on a `WassersteinKernel` for modeling functions
     that take a monotonically increasing piecewise linear function as input. The
     computation of the covariance between the piecewise linears is done by the

bofire/data_models/surrogates/single_task_gp.py

@@ -1,4 +1,4 @@
-from typing import Literal, Optional, Type, Union
+from typing import Generic, Literal, Optional, Type, TypeVar, Union
 
 import pandas as pd
 from pydantic import Field
@@ -134,19 +134,14 @@ def matern_15(
             surrogate_data.kernel = base_kernel
 
 
-class SingleTaskGPSurrogate(TrainableBotorchSurrogate):
-    type: Literal["SingleTaskGPSurrogate"] = "SingleTaskGPSurrogate"
+T = TypeVar("T")
 
-    kernel: AnyKernel = Field(
-        default_factory=lambda: RBFKernel(
-            ard=True,
-            lengthscale_prior=HVARFNER_LENGTHSCALE_PRIOR(),
-        )
-    )
+
+class BaseSingleTaskGPSurrogate(
+    TrainableBotorchSurrogate[SingleTaskGPHyperconfig], Generic[T]
+):
+    kernel: T
     noise_prior: AnyPrior = Field(default_factory=lambda: HVARFNER_NOISE_PRIOR())
-    hyperconfig: Optional[SingleTaskGPHyperconfig] = Field(
-        default_factory=lambda: SingleTaskGPHyperconfig(),
-    )
 
     @classmethod
     def is_output_implemented(cls, my_type: Type[AnyOutput]) -> bool:
@@ -157,3 +152,16 @@ 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 SingleTaskGPSurrogate(BaseSingleTaskGPSurrogate[AnyKernel]):
+    type: Literal["SingleTaskGPSurrogate"] = "SingleTaskGPSurrogate"
+    kernel: AnyKernel = Field(
+        default_factory=lambda: RBFKernel(
+            ard=True,
+            lengthscale_prior=HVARFNER_LENGTHSCALE_PRIOR(),
+        )
+    )
+    hyperconfig: Optional[SingleTaskGPHyperconfig] = Field(
+        default_factory=lambda: SingleTaskGPHyperconfig(),
+    )

bofire/data_models/surrogates/tanimoto_gp.py

@@ -16,10 +16,11 @@
     AnyPrior,
 )
 from bofire.data_models.surrogates.scaler import ScalerEnum
+from bofire.data_models.surrogates.single_task_gp import SingleTaskGPHyperconfig
 from bofire.data_models.surrogates.trainable_botorch import TrainableBotorchSurrogate
 
 
-class TanimotoGPSurrogate(TrainableBotorchSurrogate):
+class TanimotoGPSurrogate(TrainableBotorchSurrogate[SingleTaskGPHyperconfig]):
     type: Literal["TanimotoGPSurrogate"] = "TanimotoGPSurrogate"
 
     kernel: AnyKernel = Field(

bofire/data_models/surrogates/trainable.py

@@ -1,5 +1,5 @@
 import warnings
-from typing import Annotated, List, Literal, Optional, Union
+from typing import Annotated, Generic, List, Literal, Optional, TypeVar, Union
 
 import pandas as pd
 from pydantic import Field, field_validator, model_validator
@@ -92,8 +92,11 @@ def _update_hyperparameters(surrogate_data, hyperparameters: pd.Series):
         )
 
 
-class TrainableSurrogate(BaseModel):
-    hyperconfig: Optional[Hyperconfig] = None
+T = TypeVar("T", bound=Hyperconfig)
+
+
+class TrainableSurrogate(BaseModel, Generic[T]):
+    hyperconfig: Optional[T] = None
     aggregations: Optional[Annotated[List[AnyAggregation], Field(min_length=1)]] = None
 
     @model_validator(mode="after")

bofire/data_models/surrogates/trainable_botorch.py

@@ -1,11 +1,16 @@
+from typing import Generic, TypeVar
+
 from pydantic import field_validator
 
 from bofire.data_models.surrogates.botorch import BotorchSurrogate
 from bofire.data_models.surrogates.scaler import ScalerEnum
-from bofire.data_models.surrogates.trainable import TrainableSurrogate
+from bofire.data_models.surrogates.trainable import Hyperconfig, TrainableSurrogate
+
+
+T = TypeVar("T", bound=Hyperconfig)
 
 
-class TrainableBotorchSurrogate(BotorchSurrogate, TrainableSurrogate):
+class TrainableBotorchSurrogate(BotorchSurrogate, TrainableSurrogate[T], Generic[T]):
     scaler: ScalerEnum = ScalerEnum.NORMALIZE
     output_scaler: ScalerEnum = ScalerEnum.STANDARDIZE
 

bofire/surrogates/deterministic.py

@@ -1,14 +1,20 @@
+from typing import Dict, Optional, cast
+
 import torch
 from botorch.models.deterministic import AffineDeterministicModel
+from typing_extensions import Self
 
+from bofire.data_models.domain.features import Inputs, Outputs
 from bofire.data_models.enum import CategoricalEncodingEnum
 from bofire.data_models.surrogates.api import (
     CategoricalDeterministicSurrogate as CategoricalDeterministicSurrogateDataModel,
 )
 from bofire.data_models.surrogates.api import (
     LinearDeterministicSurrogate as LinearDeterministicSurrogateDataModel,
 )
+from bofire.data_models.types import InputTransformSpecs
 from bofire.surrogates.botorch import BotorchSurrogate
+from bofire.surrogates.model_utils import make_surrogate
 from bofire.utils.torch_tools import get_NumericToCategorical_input_transform, tkwargs
 
 
@@ -30,6 +36,28 @@ def __init__(
             .unsqueeze(-1),
         )
 
+    @classmethod
+    def make(
+        cls,
+        inputs: Inputs,
+        outputs: Outputs,
+        coefficients: Dict[str, float],
+        intercept: float,
+        input_preprocessing_specs: Optional[InputTransformSpecs] = None,
+        dump: str | None = None,
+        categorical_encodings: Optional[InputTransformSpecs] = None,
+    ) -> Self:
+        """
+        Factory method to create an EmpiricalSurrogate from a data model.
+        Args:
+            # document parameters
+        Returns:
+            LinearDeterministicSurrogate: A new instance.
+        """
+        return cast(
+            Self, make_surrogate(cls, LinearDeterministicSurrogateDataModel, locals())
+        )
+
 
 class CategoricalDeterministicSurrogate(BotorchSurrogate):
     def __init__(
@@ -53,3 +81,25 @@ def __init__(
             inputs=self.inputs,
             transform_specs={self.inputs[0].key: CategoricalEncodingEnum.ONE_HOT},
         )
+
+    @classmethod
+    def make(
+        cls,
+        inputs: Inputs,
+        outputs: Outputs,
+        mapping: Dict[str, float],
+        input_preprocessing_specs: Optional[InputTransformSpecs] = None,
+        dump: str | None = None,
+        categorical_encodings: Optional[InputTransformSpecs] = None,
+    ) -> Self:
+        """
+        Factory method to create an EmpiricalSurrogate from a data model.
+        Args:
+            # document parameters
+        Returns:
+            CategoricalDeterministicSurrogate: A new instance.
+        """
+        return cast(
+            Self,
+            make_surrogate(cls, CategoricalDeterministicSurrogateDataModel, locals()),
+        )

bofire/surrogates/empirical.py

@@ -1,13 +1,17 @@
 import base64
 import io
 import warnings
-from typing import Optional
+from typing import Optional, cast
 
 import torch
 from botorch.models.deterministic import DeterministicModel
+from typing_extensions import Self
 
+from bofire.data_models.domain.features import Inputs, Outputs
 from bofire.data_models.surrogates.api import EmpiricalSurrogate as DataModel
+from bofire.data_models.types import InputTransformSpecs
 from bofire.surrogates.botorch import BotorchSurrogate
+from bofire.surrogates.model_utils import make_surrogate
 
 
 class EmpiricalSurrogate(BotorchSurrogate):
@@ -28,6 +32,24 @@ def __init__(
 
     model: Optional[DeterministicModel] = None
 
+    @classmethod
+    def make(
+        cls,
+        inputs: Inputs,
+        outputs: Outputs,
+        input_preprocessing_specs: Optional[InputTransformSpecs] = None,
+        dump: str | None = None,
+        categorical_encodings: Optional[InputTransformSpecs] = None,
+    ) -> Self:
+        """
+        Factory method to create an EmpiricalSurrogate from a data model.
+        Args:
+            # document parameters
+        Returns:
+            EmpiricalSurrogate: A new instance.
+        """
+        return cast(Self, make_surrogate(cls, DataModel, locals()))
+
     def _dumps(self) -> str:
         """Dumps the actual model to a string via pickle as this is not directly json serializable."""
         with warnings.catch_warnings(record=True) as w:

bofire/surrogates/fully_bayesian.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional
+from typing import Dict, List, Literal, Optional, cast
 
 import numpy as np
 import pandas as pd
@@ -11,12 +11,18 @@
 )
 from botorch.models.transforms.input import InputTransform
 from botorch.models.transforms.outcome import OutcomeTransform
+from typing_extensions import Self
 
+from bofire.data_models.domain.features import Inputs, Outputs
 from bofire.data_models.enum import OutputFilteringEnum
 from bofire.data_models.surrogates.api import (
     FullyBayesianSingleTaskGPSurrogate as DataModel,
 )
+from bofire.data_models.surrogates.scaler import ScalerEnum
+from bofire.data_models.surrogates.trainable import AnyAggregation, Hyperconfig
+from bofire.data_models.types import InputTransformSpecs
 from bofire.surrogates.botorch import TrainableBotorchSurrogate
+from bofire.surrogates.model_utils import make_surrogate
 from bofire.utils.torch_tools import tkwargs
 
 
@@ -84,3 +90,44 @@ def _predict(self, transformed_X: pd.DataFrame):
         preds = posterior.mixture_mean.detach().numpy()
         stds = np.sqrt(posterior.mixture_variance.detach().numpy())
         return preds, stds
+
+    @classmethod
+    def make(
+        cls,
+        inputs: Inputs,
+        outputs: Outputs,
+        hyperconfig: Optional[Hyperconfig] = None,
+        aggregations: Optional[List[AnyAggregation]] = None,
+        input_preprocessing_specs: Optional[InputTransformSpecs] = None,
+        dump: Optional[str] = None,
+        categorical_encodings: Optional[InputTransformSpecs] = None,
+        scaler: ScalerEnum = ScalerEnum.NORMALIZE,
+        output_scaler: ScalerEnum = ScalerEnum.STANDARDIZE,
+        model_type: Literal["linear", "saas", "hvarfner"] = "saas",
+        warmup_steps: int = 256,
+        num_samples: int = 128,
+        thinning: int = 16,
+        features_to_warp: Optional[List[str]] = None,
+    ) -> Self:
+        """
+        Factory method to create a SingleTaskGPSurrogate from a data model.
+        Args:
+            inputs: Inputs
+            outputs: Outputs
+            hyperconfig: Hyperconfig | None
+            aggregations: List[AnyAggregation] | None
+            type: Literal['FullyBayesianSingleTaskGPSurrogate']
+            input_preprocessing_specs: InputTransformSpecs
+            dump: str | None
+            categorical_encodings: InputTransformSpecs
+            scaler: ScalerEnum
+            output_scaler: ScalerEnum
+            model_type: Literal['linear', 'saas', 'hvarfner']
+            warmup_steps: int
+            num_samples: int
+            thinning: int
+            features_to_warp: List[str]
+        Returns:
+            SingleTaskGPSurrogate: A new instance.
+        """
+        return cast(Self, make_surrogate(cls, DataModel, locals()))