codes.surrogates package

Submodules

codes.surrogates.surrogate_classes module

codes.surrogates.surrogates module

class codes.surrogates.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://robin-janssen.github.io/CODES-Benchmark/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[torch.Tensor, torch.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

torch.Tensor) -> torch.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)

Define the computation performed at every call.

Should be overridden by all subclasses. :rtype: tuple[Tensor, Tensor]

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

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[torch.Tensor, torch.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]

save(model_name, base_dir, training_id, data_params)

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.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://robin-janssen.github.io/CODES-Benchmark/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[torch.Tensor, torch.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

torch.Tensor) -> torch.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)

Define the computation performed at every call.

Should be overridden by all subclasses. :rtype: tuple[Tensor, Tensor]

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

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[torch.Tensor, torch.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]

save(model_name, base_dir, training_id, data_params)

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

class codes.surrogates.BranchNet(input_size, hidden_size, output_size, num_hidden_layers, activation=ReLU())

Bases: Module

Class that defines the branch network for the MultiONet model.

Parameters:

• input_size (int) – The input size for the network.

• hidden_size (int) – The number of hidden units in each layer.

• output_size (int) – The number of output units.

• num_hidden_layers (int) – The number of hidden layers.

forward(x)

Forward pass for the branch network.

Parameters:

x (torch.Tensor) – The input tensor.

Return type:

Tensor

class codes.surrogates.ChemDataset(raw_data, timesteps, device)

Bases: Dataset

Dataset class for the latent neural ODE model. The data is a 3D tensor with dimensions (batch, timesteps, species). The dataset also returns the timesteps for the data, as they are requred for the integration.

class codes.surrogates.Decoder(out_features, latent_features=5, width_list=None, activation=ReLU())

Bases: Module

The decoder neural network. The decoder is a simple feedforward neural network the output of which is of a higher dimension than the input. Acts as the approximate inverse of the encoder.

out_features

The number of output features.

Type:

int

latent_features

The number of latent features.

Type:

int

width_list

The width of the hidden layers.

Type:

list

activation

The activation function.

Type:

torch.nn.Module

mlp

The neural network.

Type:

torch.nn.Sequential

forward(x)

Perform a forward pass through the neural network. (“Decode” the input)

forward(x)

Perform a forward pass through the neural network.

Parameters:

x (torch.Tensor) – The input tensor.

Returns:

The output of the neural network. (“Decoded” input)

Return type:

torch.Tensor

class codes.surrogates.Encoder(in_features, latent_features=5, width_list=None, activation=ReLU())

Bases: Module

The encoder neural network. The encoder is a simple feedforward neural network the output of which is of a lower dimension than the input.

in_features

The number of input features.

Type:

int

latent_features

The number of latent features.

Type:

int

n_hidden

The number of hidden layers.

Type:

int

width_list

The width of the hidden layers.

Type:

list

activation

The activation function.

Type:

torch.nn.Module

mlp

The neural network.

Type:

torch.nn.Sequential

forward(x)

Perform a forward pass through the neural network. (“Encode” the input)

forward(x)

Perform a forward pass through the neural network.

Parameters:

x (torch.Tensor) – The input tensor.

Returns:

The output of the neural network. (“Encoded” input)

Return type:

torch.Tensor

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

Bases: AbstractSurrogateModel

create_dataloader(dataset, timesteps, batch_size, shuffle=False)

Create a DataLoader from a dataset.

Parameters:

• dataset (np.ndarray) – The dataset.

• timesteps (np.ndarray) – The timesteps.

• batch_size (int) – The batch size.

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

Returns:

The DataLoader object.

Return type:

DataLoader

epoch(data_loader, criterion, optimizer)

Perform one training epoch.

Parameters:

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

• criterion (nn.Module) – The loss function.

• optimizer (torch.optim.Optimizer) – The optimizer.

Returns:

The total loss for the training step.

Return type:

float

fit(train_loader, test_loader, epochs, position=0, description='Training FullyConnected')

Train the FullyConnected model.

Parameters:

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

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

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

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

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

Return type:

None

Returns:

None

forward(inputs)

Forward pass for the FullyConnected model.

Parameters:

• inputs (tuple[torch.Tensor, torch.Tensor]) – The input tensor and the target tensor.

• Note – The targets are not used in the forward pass but are included for compatibility with the DataLoader.

• timesteps (np.ndarray, optional) – The timesteps array.

• Note – The timesteps are not used in the forward pass but are included for compatibility with the benchmarking code.

Returns:

Output tensor of the model.

Return type:

torch.Tensor

prepare_data(dataset_train, dataset_test, dataset_val, timesteps, batch_size, shuffle=True)

Prepare the data for the predict or fit methods. Note: All datasets must have shape (n_samples, n_timesteps, n_chemicals).

Parameters:

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

• dataset_test (np.ndarray) – The test data.

• dataset_val (np.ndarray, optional) – The validation data.

• timesteps (np.ndarray) – The timesteps.

• batch_size (int, optional) – The batch size.

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

Returns:

The training, test, and validation DataLoaders.

Return type:

tuple

setup_optimizer_and_scheduler(epochs)

Utility function to set up the optimizer and scheduler for training.

Parameters:

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

Returns:

The optimizer and scheduler.

Return type:

tuple (torch.optim.Optimizer, torch.optim.lr_scheduler._LRScheduler)

class codes.surrogates.FullyConnectedNet(input_size, hidden_size, output_size, num_hidden_layers, activation=ReLU())

Bases: Module

forward(inputs)

One forward pass through the network.

Parameters:

inputs (torch.Tensor) – The input tensor.

Returns:

The output tensor of the model.

Return type:

torch.Tensor

class codes.surrogates.LatentNeuralODE(device=None, n_chemicals=29, n_timesteps=100, model_config=None)

Bases: AbstractSurrogateModel

LatentNeuralODE is a class that represents a latent neural ordinary differential equation model. It includes an encoder, decoder, and neural ODE. The integrator is implemented by the torchode framework.

model

The neural network model wrapped in a ModelWrapper object.

Type:

ModelWrapper

config

The configuration for the model.

Type:

LatentNeuralODEBaseConfig

forward(inputs)

Takes whatever the dataloader outputs, performs a forward pass through the model and returns the predictions with the respective targets.

prepare_data(dataset_train, dataset_test, dataset_val, timesteps, batch_size,

shuffle): Prepares the data for training by creating a DataLoader object.

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

Fits the model to the training data. Sets the train_loss and test_loss attributes.

fit(train_loader, test_loader, epochs, position=0, description='Training LatentNeuralODE')

Fits the model to the training data. Sets the train_loss and test_loss attributes. After 10 epochs, the loss weights are renormalized to scale the individual loss terms.

Parameters:

• train_loader (DataLoader) – The data loader for the training data.

• test_loader (DataLoader) – The data loader for the test data.

• epochs (int | None) – The number of epochs to train the model. If None, uses the value from the config.

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

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

Return type:

None

forward(inputs)

Takes whatever the dataloader outputs, performs a forward pass through the model and returns the predictions with the respective targets.

Parameters:

inputs (Any) – the data from the dataloader

Returns:

predictions and targets

Return type:

tuple[torch.Tensor, torch.Tensor]

prepare_data(dataset_train, dataset_test, dataset_val, timesteps, batch_size=128, shuffle=True)

Prepares the data for training by creating a DataLoader object.

Parameters:

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

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

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

• timesteps (np.ndarray) – The array of timesteps.

• batch_size (int) – The batch size for the DataLoader.

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

Returns:

The DataLoader object containing the prepared data.

Return type:

DataLoader

class codes.surrogates.LatentPoly(device=None, n_chemicals=29, n_timesteps=100, model_config=None)

Bases: AbstractSurrogateModel

LatentPoly class for training a polynomial model on latent space trajectories. Includes an Encoder, Decoder and learnable Polynomial.

config

The configuration for the model.

Type:

LatentPolynomialBaseConfig

model

The model for the polynomial.

Type:

PolynomialModelWrapper

device

The device to use for training.

Type:

str

forward()

Perform a forward pass through the model.

prepare_data()

Prepares the data for training by creating a DataLoader object.

fit()

Fits the model to the training data. Sets the train_loss and test_loss attributes.

fit(train_loader, test_loader, epochs, position=0, description='Training LatentPoly')

Fits the model to the training data. Sets the train_loss and test_loss attributes. After 10 epochs, the loss weights are renormalized to scale the individual loss terms.

Parameters:

• train_loader (DataLoader) – The data loader for the training data.

• test_loader (DataLoader) – The data loader for the test data.

• epochs (int | None) – The number of epochs to train the model. If None, uses the value from the config.

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

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

Return type:

None

forward(inputs)

Perform a forward pass through the model.

Parameters:

inputs (torch.Tensor) – The input tensor.

Returns:

predictions and targets

Return type:

tuple[torch.Tensor, torch.Tensor]

prepare_data(dataset_train, dataset_test, dataset_val, timesteps, batch_size=128, shuffle=True)

Prepares the data for training by creating a DataLoader object.

Parameters:

• dataset (np.ndarray) – The input dataset.

• timesteps (np.ndarray) – The timesteps for the dataset.

• batch_size (int | None) – The batch size for the DataLoader. If None, the entire dataset is loaded as a single batch.

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

Returns:

The DataLoader object containing the prepared data.

Return type:

DataLoader

class codes.surrogates.ModelWrapper(config, n_chemicals)

Bases: Module

This class wraps the encoder, decoder and ODE term into a single model. It also provides the integration of the ODE term and the loss calculation.

config

The configuration for the model.

Type:

LatentNeuralODEBaseConfig

loss_weights

The weights for the loss terms.

Type:

list

encoder

The encoder neural network.

Type:

Encoder

decoder

The decoder neural network.

Type:

Decoder

ode

The neural ODE term.

Type:

ODE

forward(x, t_range)

Performs a forward pass through the model.

renormalize_loss_weights(x_true, x_pred)

Renormalizes the loss weights.

total_loss(x_true, x_pred)

Calculates the total loss.

identity_loss(x)

Calculates the identity loss (encoder -> decoder).

l2_loss(x_true, x_pred)

Calculates the L2 loss.

deriv_loss(x_true, x_pred)

Calculates the derivative loss.

deriv2_loss(x_true, x_pred)

Calculates the second derivative loss.

deriv(x)

Calculates the first derivative.

deriv2(x)

Calculates the second derivative.

classmethod deriv(x)

Calculate the numerical derivative.

Parameters:

x (torch.Tensor) – The input tensor.

Returns:

The numerical derivative.

Return type:

torch.Tensor

classmethod deriv2(x)

Calculate the numerical second derivative.

Parameters:

x (torch.Tensor) – The input tensor.

Returns:

The numerical second derivative.

Return type:

torch.Tensor

classmethod deriv2_loss(x_true, x_pred)

Difference between the curvature of the predicted and true trajectories.

Parameters:

• x_true (torch.Tensor) – The true trajectory.

• x_pred (torch.Tensor) – The predicted trajectory

Returns:

The second derivative loss.

Return type:

torch.Tensor

classmethod deriv_loss(x_true, x_pred)

Difference between the slopes of the predicted and true trajectories.

Parameters:

• x_true (torch.Tensor) – The true trajectory.

• x_pred (torch.Tensor) – The predicted trajectory

Returns:

The derivative loss.

Return type:

torch.Tensor

forward(x, t_range)

Perform a forward pass through the model. Applies the encoder to the initial state, then propagates through time in the latent space by integrating the neural ODE term. Finally, the decoder is applied to the latent state to obtain the predicted trajectory.

Parameters:

• x (torch.Tensor) – The input tensor.

• t_range (torch.Tensor) – The range of timesteps.

Returns:

The predicted trajectory.

Return type:

torch.Tensor

identity_loss(x)

Calculate the identity loss (Encoder -> Decoder).

Parameters:

x (torch.Tensor) – The input tensor.

Returns:

The identity loss.

Return type:

torch.Tensor

classmethod l2_loss(x_true, x_pred)

Calculate the L2 loss.

Parameters:

• x_true (torch.Tensor) – The true trajectory.

• x_pred (torch.Tensor) – The predicted trajectory

Returns:

The L2 loss.

Return type:

torch.Tensor

renormalize_loss_weights(x_true, x_pred)

Renormalize the loss weights based on the current loss values so that they are accurately weighted based on the provided weights. To be used once after a short burn in phase.

Parameters:

• x_true (torch.Tensor) – The true trajectory.

• x_pred (torch.Tensor) – The predicted trajectory

total_loss(x_true, x_pred)

Calculate the total loss based on the loss weights.

Parameters:

• x_true (torch.Tensor) – The true trajectory.

• x_pred (torch.Tensor) – The predicted trajectory

Returns:

The total loss.

Return type:

torch.Tensor

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

Bases: OperatorNetwork

Class that implements the MultiONet model. It differs from a standard DeepONet in that it has multiple outputs, which are obtained by splitting the outputs of branch and trunk networks and calculating the scalar product of the splits.

Parameters:

• device (str, optional) – The device to use for training (e.g., ‘cpu’, ‘cuda:0’).

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

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

• config (dict, optional) – The configuration for the model.

• information (The configuration must provide the following)

• trunk_input_size (-) – The input size for the trunk network.

• hidden_size (-) – The number of hidden units in each layer of the branch and trunk networks.

• branch_hidden_layers (-) – The number of hidden layers in the branch network.

• trunk_hidden_layers (-) – The number of hidden layers in the trunk network.

• output_factor (-) – The factor by which the number of outputs is multiplied.

• learning_rate (-) – The learning rate for the optimizer.

• schedule (-) – Whether to use a learning rate schedule.

• regularization_factor (-) – The regularization factor for the optimizer.

• masses (-) – The masses for mass conservation loss.

• massloss_factor (-) – The factor for the mass conservation loss.

Raises:

TypeError – Invalid configuration for MultiONet model.

create_dataloader(data, timesteps, batch_size, shuffle=False)

Create a DataLoader for the given data.

Parameters:

• data (np.ndarray) – The data to load. Must have shape (n_samples, n_timesteps, n_chemicals).

• timesteps (np.ndarray) – The timesteps.

• batch_size (int, optional) – The batch size.

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

epoch(data_loader, criterion, optimizer)

Perform one training epoch.

Parameters:

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

• criterion (nn.Module) – The loss function.

• optimizer (torch.optim.Optimizer) – The optimizer.

Returns:

The total loss for the training step.

Return type:

float

fit(train_loader, test_loader, epochs, position=0, description='Training DeepONet')

Train the MultiONet model.

Parameters:

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

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

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

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

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

Return type:

None

Returns:

None. The training loss, test loss, and MAE are stored in the model.

forward(inputs)

Forward pass for the MultiONet model.

Parameters:

inputs (tuple) – The input tuple containing branch_input, trunk_input, and targets.

Returns:

The model outputs and the targets.

Return type:

tuple

prepare_data(dataset_train, dataset_test, dataset_val, timesteps, batch_size, shuffle=True)

Prepare the data for the predict or fit methods. Note: All datasets must have shape (n_samples, n_timesteps, n_chemicals).

Parameters:

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

• dataset_test (np.ndarray) – The test data.

• dataset_val (np.ndarray, optional) – The validation data.

• timesteps (np.ndarray) – The timesteps.

• batch_size (int, optional) – The batch size.

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

Returns:

The training, test, and validation DataLoaders.

Return type:

tuple

setup_criterion()

Utility function to set up the loss function for training.

Returns:

The loss function.

Return type:

callable

setup_optimizer_and_scheduler(epochs)

Utility function to set up the optimizer and scheduler for training.

Parameters:

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

Returns:

The optimizer and scheduler.

Return type:

tuple (torch.optim.Optimizer, torch.optim.lr_scheduler._LRScheduler)

class codes.surrogates.ODE(input_shape, output_shape, activation, n_hidden, layer_width, tanh_reg)

Bases: Module

The neural ODE term. The term itself is a simple feedforward neural network, a scaled tanh function is applied to the output if tanh_reg is set to True.

tanh_reg

Whether to apply a tanh regularization to the output.

Type:

bool

reg_factor

The regularization factor.

Type:

torch.Tensor

activation

The activation function.

Type:

torch.nn.Module

mlp

The neural network.

Type:

torch.nn.Sequential

forward(t, x)

Perform a forward pass through the neural network.

forward(t, x)

The forward pass through the neural network.

Parameters:

• t (torch.Tensor) – The time tensor.

• x (torch.Tensor) – The input tensor.

Returns:

The output of the neural network.

Return type:

torch.Tensor

class codes.surrogates.Polynomial(degree, dimension)

Bases: Module

Polynomial class with learnable parameters derived from nn.Module.

degree

the degree of the polynomial

Type:

int

dimension

The dimension of the in- and output variables

Type:

int

coef

The linear layer for the polynomial coefficients

Type:

nn.Linear

t_matrix

The matrix of time values

Type:

torch.Tensor

forward()

Evaluate the polynomial at the given timesteps.

_prepare_t()

Prepare the time values in matrix form for the polynomial.

forward(t)

Evaluate the polynomial at the given timesteps.

Parameters:

t (torch.Tensor) – The input tensor.

Returns:

The evaluated polynomial.

Return type:

torch.Tensor

class codes.surrogates.TrunkNet(input_size, hidden_size, output_size, num_hidden_layers, activation=ReLU())

Bases: Module

Class that defines the trunk network for the MultiONet model.

Parameters:

• input_size (int) – The input size for the network.

• hidden_size (int) – The number of hidden units in each layer.

• output_size (int) – The number of output units.

• num_hidden_layers (int) – The number of hidden layers.

forward(x)

Forward pass for the trunk network.

Parameters:

x (torch.Tensor) – The input tensor.

Return type:

Tensor