codes.surrogates.AbstractSurrogate package

Submodules

codes.surrogates.AbstractSurrogate.surrogates module

class codes.surrogates.AbstractSurrogate.surrogates.AbstractSurrogateModel(device=None, n_chemicals=29, n_timesteps=100, config=None)

Bases: ABC, Module

Abstract base class for surrogate models. This class implements the basic structure of a surrogate model and defines the methods that need to be implemented by the subclasses for it to be compatible with the benchmarking framework. For more information, see https://codes-docs.web.app/documentation.html#add_model.

Parameters:

• device (str, optional) – The device to run the model on. Defaults to None.

• n_chemicals (int, optional) – The number of chemicals. Defaults to 29.

• n_timesteps (int, optional) – The number of timesteps. Defaults to 100.

• config (dict, optional) – The configuration dictionary. Defaults to {}.

train_loss

The training loss.

Type:

float

test_loss

The test loss.

Type:

float

MAE

The mean absolute error.

Type:

float

normalisation

The normalisation parameters.

Type:

dict

train_duration

The training duration.

Type:

float

device

The device to run the model on.

Type:

str

n_chemicals

The number of chemicals.

Type:

int

n_timesteps

The number of timesteps.

Type:

int

L1

The L1 loss function.

Type:

nn.L1Loss

config

The configuration dictionary.

Type:

dict

forward(inputs

Any) -> tuple[Tensor, Tensor]: Forward pass of the model.

prepare_data(

dataset_train: np.ndarray, dataset_test: np.ndarray | None, dataset_val: np.ndarray | None, timesteps: np.ndarray, batch_size: int, shuffle: bool,

) -> tuple[DataLoader, DataLoader, DataLoader]

Gets the data loaders for training, testing, and validation.

fit(

train_loader: DataLoader, test_loader: DataLoader, epochs: int | None, position: int, description: str,

) -> None

Trains the model on the training data. Sets the train_loss and test_loss attributes.

predict(data_loader

DataLoader) -> tuple[Tensor, Tensor]: Evaluates the model on the given data loader.

save(

model_name: str, subfolder: str, training_id: str, data_params: dict,

) -> None

Saves the model to disk.

load(training_id

str, surr_name: str, model_identifier: str) -> None: Loads a trained surrogate model.

setup_progress_bar(epochs

int, position: int, description: str) -> tqdm: Helper function to set up a progress bar for training.

denormalize(data

Tensor) -> Tensor: Denormalizes the data back to the original scale.

denormalize(data)

Denormalize the data.

Parameters:

data (np.ndarray) – The data to denormalize.

Returns:

The denormalized data.

Return type:

np.ndarray

abstract fit(train_loader, test_loader, epochs, position, description)

Perform the training of the model. Sets the train_loss and test_loss attributes.

Parameters:

• train_loader (DataLoader) – The DataLoader object containing the training data.

• test_loader (DataLoader) – The DataLoader object containing the testing data.

• epochs (int) – The number of epochs to train the model for.

• position (int) – The position of the progress bar.

• description (str) – The description of the progress bar.

Return type:

None

abstract forward(inputs)

Forward pass of the model.

Parameters:

inputs (Any) – The input data as recieved from the dataloader.

Returns:

The model predictions and the targets.

Return type:

tuple[Tensor, Tensor]

classmethod get_registered_classes()

Returns the list of registered surrogate model classes.

Return type:

list[type[AbstractSurrogateModel]]

load(training_id, surr_name, model_identifier, model_dir=None)

Load a trained surrogate model.

Parameters:

• training_id (str) – The training identifier.

• surr_name (str) – The name of the surrogate model.

• model_identifier (str) – The identifier of the model (e.g., ‘main’).

Return type:

None

Returns:

None. The model is loaded in place.

predict(data_loader)

Evaluate the model on the given dataloader.

Parameters:

data_loader (DataLoader) – The DataLoader object containing the data the model is evaluated on.

Returns:

The predictions and targets.

Return type:

tuple[Tensor, Tensor]

abstract prepare_data(dataset_train, dataset_test, dataset_val, timesteps, batch_size, shuffle)

Prepare the data for training, testing, and validation. This method should return the DataLoader objects for the training, testing, and validation data.

Parameters:

• dataset_train (np.ndarray) – The training dataset.

• dataset_test (np.ndarray) – The testing dataset.

• dataset_val (np.ndarray) – The validation dataset.

• timesteps (np.ndarray) – The timesteps.

• batch_size (int) – The batch size.

• shuffle (bool) – Whether to shuffle the data.

Returns:

The DataLoader objects for the

training, testing, and validation data.

Return type:

tuple[DataLoader, DataLoader, DataLoader]

classmethod register(surrogate)

Registers a surrogate model class into the registry.

save(model_name, base_dir, training_id)

Save the model to disk.

Parameters:

• model_name (str) – The name of the model.

• subfolder (str) – The subfolder to save the model in.

• training_id (str) – The training identifier.

• data_params (dict) – The data parameters.

Return type:

None

setup_progress_bar(epochs, position, description)

Helper function to set up a progress bar for training.

Parameters:

• epochs (int) – The number of epochs.

• position (int) – The position of the progress bar.

• description (str) – The description of the progress bar.

Returns:

The progress bar.

Return type:

tqdm

Module contents

class codes.surrogates.AbstractSurrogate.AbstractSurrogateModel(device=None, n_chemicals=29, n_timesteps=100, config=None)

Bases: ABC, Module

Abstract base class for surrogate models. This class implements the basic structure of a surrogate model and defines the methods that need to be implemented by the subclasses for it to be compatible with the benchmarking framework. For more information, see https://codes-docs.web.app/documentation.html#add_model.

Parameters:

• device (str, optional) – The device to run the model on. Defaults to None.

• n_chemicals (int, optional) – The number of chemicals. Defaults to 29.

• n_timesteps (int, optional) – The number of timesteps. Defaults to 100.

• config (dict, optional) – The configuration dictionary. Defaults to {}.

train_loss

The training loss.

Type:

float

test_loss

The test loss.

Type:

float

MAE

The mean absolute error.

Type:

float

normalisation

The normalisation parameters.

Type:

dict

train_duration

The training duration.

Type:

float

device

The device to run the model on.

Type:

str

n_chemicals

The number of chemicals.

Type:

int

n_timesteps

The number of timesteps.

Type:

int

L1

The L1 loss function.

Type:

nn.L1Loss

config

The configuration dictionary.

Type:

dict

forward(inputs

Any) -> tuple[Tensor, Tensor]: Forward pass of the model.

prepare_data(

dataset_train: np.ndarray, dataset_test: np.ndarray | None, dataset_val: np.ndarray | None, timesteps: np.ndarray, batch_size: int, shuffle: bool,

) -> tuple[DataLoader, DataLoader, DataLoader]

Gets the data loaders for training, testing, and validation.

fit(

train_loader: DataLoader, test_loader: DataLoader, epochs: int | None, position: int, description: str,

) -> None

Trains the model on the training data. Sets the train_loss and test_loss attributes.

predict(data_loader

DataLoader) -> tuple[Tensor, Tensor]: Evaluates the model on the given data loader.

save(

model_name: str, subfolder: str, training_id: str, data_params: dict,

) -> None

Saves the model to disk.

load(training_id

str, surr_name: str, model_identifier: str) -> None: Loads a trained surrogate model.

setup_progress_bar(epochs

int, position: int, description: str) -> tqdm: Helper function to set up a progress bar for training.

denormalize(data

Tensor) -> Tensor: Denormalizes the data back to the original scale.

denormalize(data)

Denormalize the data.

Parameters:

data (np.ndarray) – The data to denormalize.

Returns:

The denormalized data.

Return type:

np.ndarray

abstract fit(train_loader, test_loader, epochs, position, description)

Perform the training of the model. Sets the train_loss and test_loss attributes.

Parameters:

• train_loader (DataLoader) – The DataLoader object containing the training data.

• test_loader (DataLoader) – The DataLoader object containing the testing data.

• epochs (int) – The number of epochs to train the model for.

• position (int) – The position of the progress bar.

• description (str) – The description of the progress bar.

Return type:

None

abstract forward(inputs)

Forward pass of the model.

Parameters:

inputs (Any) – The input data as recieved from the dataloader.

Returns:

The model predictions and the targets.

Return type:

tuple[Tensor, Tensor]

classmethod get_registered_classes()

Returns the list of registered surrogate model classes.

Return type:

list[type[AbstractSurrogateModel]]

load(training_id, surr_name, model_identifier, model_dir=None)

Load a trained surrogate model.

Parameters:

• training_id (str) – The training identifier.

• surr_name (str) – The name of the surrogate model.

• model_identifier (str) – The identifier of the model (e.g., ‘main’).

Return type:

None

Returns:

None. The model is loaded in place.

predict(data_loader)

Evaluate the model on the given dataloader.

Parameters:

data_loader (DataLoader) – The DataLoader object containing the data the model is evaluated on.

Returns:

The predictions and targets.

Return type:

tuple[Tensor, Tensor]

abstract prepare_data(dataset_train, dataset_test, dataset_val, timesteps, batch_size, shuffle)

Prepare the data for training, testing, and validation. This method should return the DataLoader objects for the training, testing, and validation data.

Parameters:

• dataset_train (np.ndarray) – The training dataset.

• dataset_test (np.ndarray) – The testing dataset.

• dataset_val (np.ndarray) – The validation dataset.

• timesteps (np.ndarray) – The timesteps.

• batch_size (int) – The batch size.

• shuffle (bool) – Whether to shuffle the data.

Returns:

The DataLoader objects for the

training, testing, and validation data.

Return type:

tuple[DataLoader, DataLoader, DataLoader]

classmethod register(surrogate)

Registers a surrogate model class into the registry.

save(model_name, base_dir, training_id)

Save the model to disk.

Parameters:

• model_name (str) – The name of the model.

• subfolder (str) – The subfolder to save the model in.

• training_id (str) – The training identifier.

• data_params (dict) – The data parameters.

Return type:

None

setup_progress_bar(epochs, position, description)

Helper function to set up a progress bar for training.

Parameters:

• epochs (int) – The number of epochs.

• position (int) – The position of the progress bar.

• description (str) – The description of the progress bar.

Returns:

The progress bar.

Return type:

tqdm