codes.surrogates.DeepONet package

codes.surrogates.DeepONet package#

Submodules#

codes.surrogates.DeepONet.deeponet module#

class codes.surrogates.DeepONet.deeponet.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.DeepONet.deeponet.MultiONet(device=None, n_quantities=29, n_timesteps=100, n_parameters=0, training_id=None, 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_quantities (int, optional) – The number of quantities.

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

  • n_parameters (int, optional) – The number of fixed parameters. Defaults to 0.

  • 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.

  • params_branch (-) – If True, fixed parameters are concatenated to the branch net; if False, to the trunk net.

Raises:

TypeError – Invalid configuration for MultiONet model.

create_dataloader(data, timesteps, batch_size, shuffle, dataset_params, params_in_branch, num_workers=0, pin_memory=True)#
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', multi_objective=False)#

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.

  • multi_objective (bool) – Whether multi-objective optimization is used. If True, trial.report is not used (not supported by Optuna).

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, dummy_timesteps=True, dataset_train_params=None, dataset_test_params=None, dataset_val_params=None)#

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.

  • dummy_timesteps (bool) – Whether to use dummy timesteps. Defaults to True.

Returns:

The DataLoader objects for the

training, testing, and validation data.

Return type:

tuple[DataLoader, DataLoader, DataLoader]

setup_criterion()#

Utility function to set up the loss function for training.

Returns:

The loss function.

Return type:

callable

class codes.surrogates.DeepONet.deeponet.OperatorNetwork(device=None, n_quantities=29, n_timesteps=100, training_id=None, config=None)#

Bases: AbstractSurrogateModel

Abstract class for operator networks. Child classes must implement a forward method and use the branch_net and trunk_net attributes.

Parameters:

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

  • n_quantities (int, optional) – The number of quantities.

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

Raises:

  • NotImplementedError – Child classes must implement a forward method.

  • NotImplementedError – Child classes must initialize a branch_net and trunk_net.

forward(branch_input, trunk_input)#

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]

post_init_check()#
class codes.surrogates.DeepONet.deeponet.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

codes.surrogates.DeepONet.deeponet_config module#

class codes.surrogates.DeepONet.deeponet_config.MultiONetBaseConfig(learning_rate=0.0003, regularization_factor=0.0, optimizer='adamw', momentum=0.0, scheduler='cosine', poly_power=0.9, eta_min=0.1, activation=None, loss_function=None, beta=0.0, masses=None, trunk_input_size=1, hidden_size=100, branch_hidden_layers=5, trunk_hidden_layers=5, output_factor=10, massloss_factor=0.0, params_branch=True)#

Bases: AbstractSurrogateBaseConfig

Configuration for the MultiONet surrogate model.

MultiONet is a physics-inspired architecture that splits the network into a trunk and multiple branch networks, designed to represent operator-based systems.

Attributes:

  • masses (list[float] | None) – Optional list of particle masses (if used).

  • trunk_input_size (int) – Size of the input to the trunk network.

  • hidden_size (int) – Number of neurons per hidden layer in both branches and trunk.

  • branch_hidden_layers (int) – Number of hidden layers in each branch network.

  • trunk_hidden_layers (int) – Number of hidden layers in the trunk network.

  • output_factor (int) – Multiplier for the output dimension (neurons = output_factor * num_quantities).

  • massloss_factor (float) – Additional weight for a mass conservation loss term.

  • params_branch (bool) – Flag to indicate whether parameters (if present) are passed to the branch or trunk net.

branch_hidden_layers: int = 5#
hidden_size: int = 100#
masses: list[float] | None = None#
massloss_factor: float = 0.0#
output_factor: int = 10#
params_branch: bool = True#
trunk_hidden_layers: int = 5#
trunk_input_size: int = 1#

codes.surrogates.DeepONet.don_utils module#

class codes.surrogates.DeepONet.don_utils.FlatBatchIterable(branch_t, trunk_t, target_t, batch_size, shuffle)#

Bases: IterableDataset

class codes.surrogates.DeepONet.don_utils.PreBatchedDataset(branch_inputs, trunk_inputs, targets)#

Bases: Dataset

Dataset for pre-batched data.

Parameters:

  • branch_inputs (list[Tensor]) – List of precomputed branch input batches.

  • trunk_inputs (list[Tensor]) – List of precomputed trunk input batches.

  • targets (list[Tensor]) – List of precomputed target batches.

codes.surrogates.DeepONet.don_utils.create_date_based_directory(base_dir='.', subfolder='models')#

Create a directory based on the current date (dd-mm format) inside a specified subfolder of the base directory.

Parameters:

  • base_dir – The base directory where the subfolder and date-based directory will be created.

  • subfolder – The subfolder inside the base directory to include before the date-based directory.

Returns:

The path of the created date-based directory within the specified subfolder.

codes.surrogates.DeepONet.don_utils.custom_collate_fn(batch)#

Custom collate function to ensure tensors are returned in the correct shape.

Parameters:

batch – A list of tuples from the dataset, where each tuple contains (branch_input, trunk_input, targets).

Returns:

  • branch_input: [batch_size, feature_size]

  • trunk_input: [batch_size, feature_size]

  • targets: [batch_size, feature_size]

Return type:

A tuple of tensors with correct shapes

codes.surrogates.DeepONet.don_utils.get_project_path(relative_path)#

Construct the absolute path to a project resource (data or model) based on a relative path.

Parameters:

relative_path – A relative path to the resource, e.g., “datasets/dataset100” or “models/02-28/model.pth”.

Returns:

The absolute path to the resource.

codes.surrogates.DeepONet.don_utils.list_pth_files(directory='models')#

List all .pth files in the specified directory without their extensions.

Parameters:

directory (str) – The directory to search for .pth files.

Returns:

A list of filenames without the .pth extension.

Return type:

list[str]

codes.surrogates.DeepONet.don_utils.mass_conservation_loss(masses, criterion=MSELoss(), weights=(1, 1), device=device(type='cpu'))#

Replaces the standard MSE loss with a sum of the standard MSE loss and a mass conservation loss.

Parameters:

  • masses (list) – A list of masses for the quantities.

  • criterion – The loss function to use for the standard loss.

  • weights (tuple) – A 2-tuple of weights for the standard loss and the mass conservation loss.

  • device (device) – The device to use for the loss function.

Returns:

A new loss function that includes the mass conservation loss.

codes.surrogates.DeepONet.don_utils.read_yaml_config(model_path)#
codes.surrogates.DeepONet.don_utils.save_plot_counter(filename, directory='plots')#
codes.surrogates.DeepONet.don_utils.set_random_seed(gpu_id=0)#

codes.surrogates.DeepONet.train_utils module#

codes.surrogates.DeepONet.train_utils.mass_conservation_loss(masses, criterion=MSELoss(), weights=(1, 1), device=device(type='cpu'))#

Replaces the standard MSE loss with a sum of the standard MSE loss and a mass conservation loss.

Parameters:

  • masses (list) – A list of masses for the quantities.

  • criterion – The loss function to use for the standard loss.

  • weights (tuple) – A 2-tuple of weights for the standard loss and the mass conservation loss.

  • device (device) – The device to use for the loss function.

Returns:

A new loss function that includes the mass conservation loss.

Module contents#

Contents