codes.surrogates package

codes.surrogates package#

Subpackages#

Submodules#

codes.surrogates.surrogate_classes module#

codes.surrogates.training module#

Shared training loop for the standard surrogate models.

codes.surrogates.training.fit_standard(model, train_loader, test_loader, epochs, position=0, description='Training', multi_objective=False, loss_fn=None, *, gradient_clip=None, epoch_fn=None)#

Train and restore the best checkpoint for a standard surrogate.

The model supplies config.loss, config.gradient_clip, update_epochs, and the standard optimizer, progress, checkpoint and validation methods. The optimizer returned by setup_optimizer_and_scheduler must provide train() and eval(). loss_fn may customize the batch loss as described by train_epoch(); gradient_clip overrides the configured limit. epoch_fn, if provided, receives (train_loader, criterion, optimizer) and owns the whole epoch. It cannot be combined with a batch loss or gradient clipping override.

Scheduler steps and validation checks occur once per epoch; metrics are computed at the model’s update_epochs interval. The model’s loss histories and n_epochs are populated, and its checkpoint is restored on success. Concrete fit methods can use time_execution to record fit_duration.

Return type:: None

codes.surrogates.training.train_epoch(model, data_loader, criterion, optimizer, loss_fn=None, *, gradient_clip=None)#

Train for one epoch, using a custom batch loss when supplied.

loss_fn(batch, criterion) must return a differentiable scalar. By default, the model returns predictions and targets and the criterion is applied to them. A gradient_clip override of zero disables clipping; None uses the config.

Return type:: None

Module contents#

class codes.surrogates.AbstractSurrogateModel(device=None, n_quantities=29, n_timesteps=100, n_parameters=0, training_id=None, 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 | None) – The device to run the model on. Defaults to None.
n_quantities (int) – The number of quantities. Defaults to 29.
n_timesteps (int) – The number of timesteps. Defaults to 100.
config (dict | None) – The configuration dictionary. Defaults to {}.

Attributes:

train_loss (float) – The training loss.
test_loss (float) – The test loss.
MAE (float) – The mean absolute error.
normalisation (dict) – The normalisation parameters.
train_duration (float) – The training duration.
device (str) – The device to run the model on.
n_quantities (int) – The number of quantities.
n_timesteps (int) – The number of timesteps.
L1 (nn.L1Loss) – The L1 loss function.
config (dict) – The configuration dictionary.

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_info: 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.

checkpoint(test_loss, epoch)#

If checkpointing is enabled and test_loss < self.best_test_loss, store a copy of the current weights and update best_test_loss/epoch.

Return type:: None

denormalize(data, leave_log=False, leave_norm=False)#

Denormalize the data.

Parameters:

data (Tensor | ndarray) – The data to denormalize.
leave_log (bool) – If True, do not exponentiate the data even if log10_transform is True.
leave_norm (bool) – If True, do not denormalize the data even if normalisation is applied.

Returns:

The denormalized data.

Return type:

Tensor | ndarray

abstractmethod fit(train_loader, test_loader, epochs, position, description, multi_objective)#

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.
multi_objective (bool) – Whether the training is multi-objective.

Return type:

None

abstractmethod 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]

forward_from_state(state, times, params=None)#

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

Rollout finetuning needs this to run the model on states it produced itself, without building a dataloader. Surrogates that support finetuning implement it by packing the inputs the same way as their dataloader.

Parameters:

state (Tensor) – Initial states of shape (batch, n_quantities), normalised.
times (Tensor) – Time inputs of shape (n_times,), in the dummy time coordinate the model is trained on.
params (Tensor | None) – Parameters of shape (batch, n_parameters).

Returns:

Predicted states of shape (batch, n_times, n_quantities).

Return type:

Tensor

get_checkpoint(test_loader, criterion=None)#

After training, compare the final weights with the best stored weights and keep whichever reaches the lower test loss. Both losses are the log-space L1 that validate reports, so the comparison is like for like.

Parameters:

test_loader (DataLoader) – DataLoader for computing final test loss.
criterion (Module | None) – Unused; kept for backward compatibility.

Return type:

None

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, leave_log=False, leave_norm=False)#

Evaluate the model on the given dataloader.

Parameters:

data_loader (DataLoader) – The DataLoader object containing the data the model is evaluated on.
leave_log (bool) – If True, do not exponentiate the data even if log10_transform is True.
leave_norm (bool) – If True, do not denormalize the data even if normalisation is applied.

Returns:

The predictions and targets.

Return type:

tuple[Tensor, Tensor]

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

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 | None) – 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 | None, DataLoader | None]

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_info (dict) – The data parameters.

Return type:

None

setup_checkpoint()#

Prepare tracking of the best weights seen during training. Must be called before any call to self.checkpoint(…).

The best weights are kept in memory per model instance, so models trained concurrently (e.g. on several devices) cannot overwrite each other’s checkpoint.

Return type:: None

setup_optimizer_and_scheduler(epochs)#

Set up optimizer and scheduler based on self.config.scheduler and self.config.optimizer. Supports “adamw”, “sgd” optimizers and “schedulefree”, “cosine”, “poly” schedulers. Patches standard optimizers so that .train() and .eval() exist as no-ops. Patches ScheduleFree optimizers to have a no-op scheduler.step(). For ScheduleFree optimizers, use lr warmup for the first 1% of epochs. For Poly scheduler, use a power decay based on self.config.poly_power. For Cosine scheduler, anneal to the minimum learning rate self.config.eta_min * self.config.learning_rate.

Parameters:: epochs (int) – The number of epochs the training will run for.
Returns:: The optimizer and scheduler instances.
Return type:: tuple[Optimizer, _LRScheduler]
Raises:: ValueError – If an unknown optimizer or scheduler is specified in the config.

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

time_pruning(current_epoch, total_epochs)#

Determine whether a trial should be pruned based on projected runtime, but only after a warmup period (10% of the total epochs).

Warmup: Do not prune if current_epoch is less than warmup_epochs. After warmup, compute the average epoch time, extrapolate the total runtime, and retrieve the threshold (runtime_threshold) from the study’s user attributes. If the projected runtime exceeds the threshold, raise an optuna.TrialPruned exception.

Parameters:

current_epoch (int) – The current epoch count.
total_epochs (int) – The planned total number of epochs.

Raises:

optuna.TrialPruned – If the projected runtime exceeds the threshold.

Return type:

None

validate(epoch, train_loader, test_loader, optimizer, progress_bar, total_epochs, multi_objective)#

Shared “validation + checkpoint” logic, to be called once per epoch in each fit().

Return type:: None

Relies on:

self.update_epochs (int)
self.train_loss (np.ndarray)
self.test_loss (np.ndarray)
self.MAE (np.ndarray)
self.optuna_trial
self.L1 (nn.L1Loss)
self.predict(…)
self.checkpoint(test_loss, epoch)

Only runs if (epoch % self.update_epochs) == 0. Main reporting metric is MAE in log10-space (i.e., Δdex). Additionally, MAE in linear space is computed.

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 (Tensor) – The input tensor.
Return type:: Tensor

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

Bases: Dataset

Dataset class for the latent neural ODE model. Returns each sample along with its timesteps and (optionally) fixed parameters.

Deprecated: use FlatSeqBatchIterable for pre-batched sequence data.

class codes.surrogates.Decoder(out_features, latent_features=5, coder_layers=3, coder_width=32, activation=ReLU(), dtype=torch.float64)#

Bases: Module

Fully connected decoder that maps the latent space back to the output.

forward(x)#

Define the computation performed at every call.

Should be overridden by all subclasses.

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.

Return type:: Tensor

class codes.surrogates.Encoder(in_features, latent_features=5, coder_layers=3, coder_width=32, activation=ReLU(), dtype=torch.float64)#

Bases: Module

Fully connected encoder that maps input features to a latent space.

forward(x)#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Return type:: Tensor

class codes.surrogates.FlatSeqBatchIterable(data_t, timesteps_t, params_t, batch_size, shuffle)#: Bases: IterableDataset

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

Bases: AbstractSurrogateModel

create_dataloader(data, timesteps, batch_size, shuffle, dataset_params, num_workers=0, pin_memory=True)#: Build CPU tensors once and yield shuffled batches each epoch. data: (n_samples, n_timesteps, n_quantities) inputs = [initial_state, time, (params?)] targets = state_at_time

epoch(data_loader, criterion, optimizer)#

Return type:: None

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

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) – 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 FullyConnected model.

Parameters:: inputs (tuple) – (x, targets) - ‘targets’ is included for a consistent interface
Return type:: Tensor
Returns:: (outputs, targets)

forward_from_state(state, times, params=None)#

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

Packs [state, time, params] as create_dataloader does. 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 | None) – 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]

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

Bases: Module

forward(inputs)#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Return type:: Tensor

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

Bases: FullyConnected

FullyConnected variant that predicts the change from the initial state.

The network output is added to the initial state, so it only has to learn the deviation from it. Everything downstream still sees absolute states.

forward(inputs)#

Forward pass for the FullyConnected model.

Parameters:: inputs (tuple) – (x, targets) - ‘targets’ is included for a consistent interface
Return type:: Tensor
Returns:: (outputs, targets)

class codes.surrogates.LatentNeuralODE(device=None, n_quantities=29, n_timesteps=100, n_parameters=0, training_id=None, config=None, dtype=torch.float32)#

Bases: AbstractSurrogateModel

LatentNeuralODE represents a latent neural ODE model. It includes an encoder, decoder, and neural ODE. Fixed parameters can be injected either into the encoder or later into the ODE network, controlled by config.encode_params.

Parameters:

device (str | None) – Device for training (e.g. ‘cpu’, ‘cuda:0’).
n_quantities (int) – Number of quantities.
n_timesteps (int) – Number of timesteps.
n_parameters (int) – Number of fixed parameters (default 0).
config (dict | None) – Configuration for the model.

create_dataloader(data, timesteps, batch_size, shuffle, dataset_params, num_workers=0, pin_memory=True)#

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

Train the LatentNeuralODE model.

Parameters:

train_loader (DataLoader) – The data loader for the training data.
test_loader (DataLoader) – The data loader for the test data.
epochs (int) – The positive 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

forward(inputs)#: Forward pass through the model. Expects inputs to be either (data, timesteps) or (data, timesteps, params).

prepare_data(dataset_train, dataset_test, dataset_val, timesteps, batch_size=128, 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 | None) – 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 | None, DataLoader | None]

training_loss(batch, criterion)#: Compute the trajectory and auxiliary losses for one sequence batch.

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

Bases: AbstractSurrogateModel

LatentPoly class for training a polynomial model on latent space trajectories.

This model includes an encoder, decoder, and a learnable polynomial applied on the latent space. The architecture is chosen based on the version flag in the configuration.

Attributes:

config (LatentPolynomialBaseConfig) – The configuration for the model.
model (PolynomialModelWrapper) – The wrapped model (encoder, decoder, polynomial).
device (str) – Device for training.

create_dataloader(data, timesteps, batch_size, shuffle, dataset_params, num_workers=0, pin_memory=True)#

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

Train the LatentPoly model.

Parameters:

train_loader (DataLoader) – The data loader for the training data.
test_loader (DataLoader) – The data loader for the test data.
epochs (int) – The positive 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

forward(inputs)#

Perform a forward pass through the model.

Parameters:: inputs (tuple) – Tuple containing the input tensor and timesteps. If fixed parameters are provided, the tuple is (data, timesteps, params).
Returns:: (Predictions, Targets)
Return type:: tuple[Tensor, Tensor]

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 | None) – 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]

training_loss(batch, criterion)#: Compute the trajectory and auxiliary losses for one sequence batch.

class codes.surrogates.ModelWrapper(config, n_quantities, n_parameters=0, n_timesteps=101, dtype=torch.float64, use_pid=False, adjoint_type='autodiff')#

Bases: Module

Wraps the encoder, decoder, and neural ODE in three distinct modes:

No parameters (n_parameters=0) - Encoder: input = state_dim - ODE: latent_dim -> latent_dim (the solver always evolves the latent state) - Decoder: latent_dim -> output dimensions
encode_params=True - Encoder: input = state_dim + param_dim - ODE: latent_dim -> latent_dim - Decoder: latent_dim -> output dimensions
encode_params=False - Encoder: input = state_dim - Base ODE: (latent_dim + param_dim) -> latent_dim - Wrapped in ODEWithParams so that solver state = latent_dim - Decoder: latent_dim -> output dimensions

first_derivative(x)#

forward(x0, t_range, params=None)#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

identity_loss(x_true, params=None)#

Calculate the identity loss (Encoder -> Decoder) on the initial state x0.

Parameters:

x_true (Tensor) – The full trajectory (batch, timesteps, features).
params (Tensor) – Fixed parameters (batch, n_parameters).

Returns:

The identity loss on x0.

Return type:

Tensor

second_derivative(x)#

total_loss(x_true, x_pred, params=None, criterion=MSELoss())#: Total loss: weighted sum of trajectory reconstruction, identity, first derivative, and second derivative losses. All terms remain in the computation graph.

class codes.surrogates.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 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.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.ODE(input_shape, output_shape, activation, ode_layers, ode_width, tanh_reg, dtype=torch.float64)#

Bases: Module

Neural ODE module defining the function for latent dynamics.

forward(t, x)#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class codes.surrogates.Polynomial(degree, dimension)#

Bases: Module

Learnable polynomial model.

Attributes:

degree (int) – Degree of the polynomial.
dimension (int) – Dimension of the in- and output.
coef (nn.Linear) – Linear layer representing polynomial coefficients.
t_matrix (Tensor) – Time matrix for polynomial evaluation.

forward(t)#

Evaluate the polynomial at given timesteps.

Parameters:: t (Tensor) – Time tensor.
Returns:: Evaluated polynomial.
Return type:: 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 (Tensor) – The input tensor.
Return type:: Tensor

codes.surrogates package

Contents

codes.surrogates package#

Subpackages#

Submodules#

codes.surrogates.surrogate_classes module#

codes.surrogates.training module#

Module contents#