Adding emulators#
In addition to providing a library of core emulators, AutoEmulate is designed to be easily extensible. This tutorial walks you through the steps of adding new emulators to the library. We cover two scenarios: adding new Gaussian Process kernels and adding entirely new models.
1. Adding Gaussian Process kernels#
Gaussian Processes (GPs) are primarily defined by their kernel functions, which determine the covariance structure of the data. AutoEmulate includes several built-in GP kernels:
Radial Basis Function (RBF)
Matern 3/2
Matern 5/2
Rational Quadratic (RQ)
Linear
You can easily create new kernels by composing any two or more of these existing kernels. For example, you might want to create a kernel that combines the RBF and Linear kernels to capture both smooth variations and linear trends in your data.
In AutoEmulate, each kernel is defined by an initialisation function that takes as inputs the number of data input features and the number of output features. Below we define a custom kernel function following this pattern.
from autoemulate.emulators.gaussian_process.kernel import rbf_kernel, linear_kernel
def rbs_plus_linear_kernel(n_features, n_outputs):
"""
Example of a custom kernel function that combines RBF and linear kernels.
"""
return rbf_kernel(n_features, n_outputs) + linear_kernel(n_features, n_outputs)
Once this function has been defined, you can create a new GP emulator class using the create_gp_subclass
function.
from autoemulate.emulators.gaussian_process.exact import GaussianProcess, create_gp_subclass
GaussianProcessRBFandLinear = create_gp_subclass(
"GaussianProcessRBFandLinear",
GaussianProcess,
# the custom kernel function goes here
covar_module_fn=rbs_plus_linear_kernel,
)
Now we can tell AutoEmulate to use the new GP class by passing it to the models
argument when initialising an AutoEmulate
object.
from autoemulate import AutoEmulate
import torch
# create some example data
x = torch.linspace(0, 1, 100).unsqueeze(-1)
y = torch.sin(2 * 3.14 * x) + 0.1 * torch.randn_like(x)
ae = AutoEmulate(x, y, models=[GaussianProcessRBFandLinear])
Comparing models: 0%| | 0.00/1.00 [00:00<?, ?model/s]
Comparing models: 100%|██████████| 1.00/1.00 [00:30<00:00, 30.6s/model]
Comparing models: 100%|██████████| 1.00/1.00 [00:30<00:00, 30.6s/model]
ae.summarise()
model_name | x_transforms | y_transforms | params | rmse_test | r2_test | r2_test_std | r2_train | r2_train_std | |
---|---|---|---|---|---|---|---|---|---|
0 | GaussianProcessRBFandLinear | [StandardizeTransform()] | [StandardizeTransform()] | {'mean_module_fn': <function linear_mean at 0x... | 0.095536 | 0.975267 | 0.00875 | 0.981082 | 0.004556 |
2. Adding new models#
It is also possible to add entirely new models to AutoEmulate. AutoEmulate has a base Emulator
class that handles most of the general functionality required for training and prediction. To implement a new emulator, one must simply subclass Emulator
and implement the abstract methods (_fit
, _predict
and is_multioutput
), get_tune_params
to enable model tuning, as well any model specific functionality and initialisations.
Since AutoEmulate supports a variety of models, there are additional Emulator
subclasses that handle specific functionality for each model type:
PytorchBackend
for PyTorch modelsSklearnBackend
for scikit-learn modelsGaussianProcess
for exact Gaussian Process implementationsEnsemble
for ensemble models
Subclassing one of these directly has slightly different requirements. For example, when subclassing PytorchBackend
or GaussianProcess
, one must implement the forward
method to define the model’s forward pass.
There are also some static methods that should be implemented to provide metadata about the model, such as is_multioutput
and get_tune_params
.
Below demonstrates adding a simple feedforward neural network (FNN) using PyTorch. The new class SimpleFNN
subclasses PytorchBackend
, which already handles fitting and prediction.
from autoemulate.core.device import TorchDeviceMixin
from autoemulate.emulators.base import PyTorchBackend
import torch.nn as nn
class SimpleFNN(PyTorchBackend):
def __init__(
self,
x,
y,
hidden_dim=64,
device = None,
):
TorchDeviceMixin.__init__(self, device=device)
nn.Module.__init__(self)
input_dim = x.shape[1]
output_dim = y.shape[1] if len(y.shape) > 1 else 1
layers = []
layers.append(nn.Linear(input_dim, hidden_dim, device=self.device))
layers.append(nn.ReLU())
layers.append(nn.Linear(hidden_dim, output_dim, device=self.device))
self.model = nn.Sequential(*layers)
self.optimizer = self.optimizer_cls(self.model.parameters(), lr=self.lr) # type: ignore[call-arg] since all optimizers include lr
self.scheduler = None
self.to(self.device)
def forward(self, x):
return self.model(x)
@staticmethod
def is_multioutput():
return True
@staticmethod
def get_tune_params():
return {
"hidden_dim": [32, 64, 128]
}
ae = AutoEmulate(x, y, models=[SimpleFNN])
Comparing models: 0%| | 0.00/1.00 [00:00<?, ?model/s]
Comparing models: 100%|██████████| 1.00/1.00 [00:01<00:00, 1.94s/model]
Comparing models: 100%|██████████| 1.00/1.00 [00:01<00:00, 1.94s/model]
ae.summarise()
model_name | x_transforms | y_transforms | params | rmse_test | r2_test | r2_test_std | r2_train | r2_train_std | |
---|---|---|---|---|---|---|---|---|---|
0 | SimpleFNN | [StandardizeTransform()] | [StandardizeTransform()] | {'hidden_dim': 64} | 0.162058 | 0.931167 | 0.02025 | 0.952293 | 0.011802 |