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 (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 | None) – The device to use for training (e.g., ‘cpu’, ‘cuda:0’).

  • n_quantities (int) – The number of quantities.

  • n_timesteps (int) – The number of timesteps.

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

  • config (dict | None) – 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 (Module) – The loss function.

  • optimizer (Optimizer) – The optimizer.

Return type:

None

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) – 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[Tensor, Tensor]

forward_from_state(state, times, params=None)#

Predict the states at several time inputs from one initial state.

Packs the branch and trunk inputs as create_dataloader does, with the parameters on the side set by params_branch. See AbstractSurrogateModel.forward_from_state.

Return type:

Tensor

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 (ndarray) – The training dataset.

  • dataset_test (ndarray) – The testing dataset.

  • dataset_val (ndarray | None) – The validation dataset.

  • timesteps (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()#

Build the legacy summed-MSE loss, optionally with mass conservation.

Deprecated: fit uses self.config.loss and does not call this method. It remains available for callers constructing the legacy loss explicitly.

Returns:

The loss function.

Return type:

callable

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

Bases: MultiONet

MultiONet variant that predicts the change from the initial state.

The branch input starts with the initial state, so adding its first N columns back onto the output leaves the network to learn the deviation.

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

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 | None) – The device to use for training (e.g., ‘cpu’, ‘cuda:0’).

  • n_quantities (int) – The number of quantities.

  • n_timesteps (int) – 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 (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=1.0, gradient_clip=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

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.

Module contents#