codes.surrogates.LatentPolynomial package

codes.surrogates.LatentPolynomial package#

Submodules#

codes.surrogates.LatentPolynomial.latent_poly module#

class codes.surrogates.LatentPolynomial.latent_poly.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_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, DataLoader]

training_loss(batch, criterion)#

Compute the trajectory and auxiliary losses for one sequence batch.

class codes.surrogates.LatentPolynomial.latent_poly.OldDecoder(out_features, latent_features=5, layers_factor=8, activation=ReLU())#

Bases: Module

Old decoder network corresponding to the fixed 4–2–1 encoder.

Parameters:

  • out_features (int) – Number of output features.

  • latent_features (int) – Dimension of the latent space.

  • layers_factor (int) – Factor to scale the base widths [4, 2, 1].

  • activation (Module) – Activation function.

forward(x)#

Forward pass through the old decoder.

Return type:

Tensor

class codes.surrogates.LatentPolynomial.latent_poly.OldEncoder(in_features, latent_features=5, layers_factor=8, activation=ReLU())#

Bases: Module

Old encoder network implementing a fixed 4–2–1 structure scaled by layers_factor.

Parameters:

  • in_features (int) – Number of input features.

  • latent_features (int) – Dimension of the latent space.

  • layers_factor (int) – Factor to scale the base widths [4, 2, 1].

  • activation (Module) – Activation function.

forward(x)#

Forward pass through the old encoder.

Return type:

Tensor

class codes.surrogates.LatentPolynomial.latent_poly.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.LatentPolynomial.latent_poly.PolynomialModelWrapper(config, device, n_parameters=0, n_timesteps=101)#

Bases: Module

Wraps the Encoder, Decoder, and learnable Polynomial into a single model.

If config.coeff_network is True, fixed parameters are not concatenated to the encoder input; instead, a coefficient network predicts the polynomial coefficients based on the parameters. If False, fixed parameters are concatenated to the encoder input and the polynomial coefficients are learned directly.

Attributes:

  • config (LatentPolynomialBaseConfig) – Model configuration.

  • loss_weights (list[float]) – Weights for the loss terms.

  • device (str) – Device for training.

  • encoder (Module) – The encoder network.

  • decoder (Module) – The decoder network.

  • poly (Polynomial) – The polynomial module.

  • coefficient_net (Module | None) – The coefficient network (if config.coeff_network is True).

first_derivative(x)#
forward(x, t_range, params=None)#

Forward pass through the model.

Parameters:

  • x (Tensor) – Input tensor of shape (batch, timesteps, n_quantities).

  • t_range (Tensor) – Time range tensor.

  • params (Tensor | None) – Fixed parameters of shape (batch, n_parameters), if provided.

Returns:

Predicted trajectory.

Return type:

Tensor

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.

codes.surrogates.LatentPolynomial.latent_poly_config module#

class codes.surrogates.LatentPolynomial.latent_poly_config.LatentPolynomialBaseConfig(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, model_version='v2', latent_features=5, degree=2, coder_hidden=4, layers_factor=8, coder_layers=3, coder_width=32, coeff_network=False, coeff_width=32, coeff_layers=4)#

Bases: AbstractSurrogateBaseConfig

Standard model config for LatentPolynomial with versioning.

Attributes:

  • model_version (str) – “v1” for the old fixed 4-2-1 structure; “v2” for the new FCNN design.

  • latent_features (int) – Dimension of the latent space.

  • degree (int) – Degree of the learnable polynomial.

  • coder_hidden (int) – (v1 only) Base hidden size for fixed structure.

  • layers_factor (int) – (v1 only) Factor multiplied with coder_hidden to determine layer widths.

  • coder_layers (int) – (v2 only) Number of hidden layers in the encoder/decoder.

  • coder_width (int) – (v2 only) Number of neurons in every hidden layer in the encoder/decoder.

  • coeff_network (bool) – Whether to use a coefficient network for polynomial coefficients.

  • coeff_width (int) – Width of the coefficient network (if used).

  • coeff_layers (int) – Number of layers in the coefficient network (if used).

coder_hidden: int = 4#
coder_layers: int = 3#
coder_width: int = 32#
coeff_layers: int = 4#
coeff_network: bool = False#
coeff_width: int = 32#
degree: int = 2#
latent_features: int = 5#
layers_factor: int = 8#
model_version: str = 'v2'#

Module contents#