codes.surrogates.AbstractSurrogate package

Submodules

codes.surrogates.AbstractSurrogate.abstract_config module

class codes.surrogates.AbstractSurrogate.abstract_config.AbstractSurrogateBaseConfig(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)

Bases: object

Base configuration for the AbstractSurrogate model.

This class defines shared attributes and methods for surrogate models.

Attributes:

• learning_rate (float) – Learning rate for the optimizer.

• regularization_factor (float) – Regularization coefficient, applied as weight decay.

• optimizer (str) – Type of optimizer to use. Supported options: adamw, sgd.

• momentum (float) – Momentum factor for the optimizer (used only if optimizer == “sgd”).

• scheduler (str) – Type of learning rate scheduler to use. - “schedulefree”: Use schedulefree optimizer. - “cosine”: Use cosine annealing scheduler. - “poly”: Use polynomial decay scheduler.

• poly_power (float) – Power for polynomial decay scheduler (used only if scheduler == “poly”).

• eta_min (float) – Minimum learning rate of the cosine annealing scheduler, as a fraction of learning_rate (used only if scheduler == “cosine”).

• activation (nn.Module | None) – Activation module used in the model.

• loss_function (nn.Module | None) – Loss function used for training.

• beta (float) – Beta parameter forwarded to SmoothL1Loss (used only if loss_function is nn.SmoothL1Loss). Defaults to 1.0, PyTorch’s default.

• gradient_clip (float) – Maximum gradient norm. Set to 0 to disable clipping.

activation: Module | None = None

beta: float = 1.0

eta_min: float = 0.1

gradient_clip: float = 0.0

learning_rate: float = 0.0003

property loss: Module

Returns the loss function to be used for training.

If the loss function is nn.SmoothL1Loss, it returns an instance with the specified beta. Otherwise, it returns the loss function as is.

loss_function: Module | None = None

momentum: float = 0.0

optimizer: str = 'adamw'

poly_power: float = 0.9

regularization_factor: float = 0.0

scheduler: str = 'cosine'

codes.surrogates.AbstractSurrogate.abstract_surrogate module

class codes.surrogates.AbstractSurrogate.abstract_surrogate.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.

Module contents

class codes.surrogates.AbstractSurrogate.AbstractSurrogateBaseConfig(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)

Bases: object

Base configuration for the AbstractSurrogate model.

This class defines shared attributes and methods for surrogate models.

Attributes:

• learning_rate (float) – Learning rate for the optimizer.

• regularization_factor (float) – Regularization coefficient, applied as weight decay.

• optimizer (str) – Type of optimizer to use. Supported options: adamw, sgd.

• momentum (float) – Momentum factor for the optimizer (used only if optimizer == “sgd”).

• scheduler (str) – Type of learning rate scheduler to use. - “schedulefree”: Use schedulefree optimizer. - “cosine”: Use cosine annealing scheduler. - “poly”: Use polynomial decay scheduler.

• poly_power (float) – Power for polynomial decay scheduler (used only if scheduler == “poly”).

• eta_min (float) – Minimum learning rate of the cosine annealing scheduler, as a fraction of learning_rate (used only if scheduler == “cosine”).

• activation (nn.Module | None) – Activation module used in the model.

• loss_function (nn.Module | None) – Loss function used for training.

• beta (float) – Beta parameter forwarded to SmoothL1Loss (used only if loss_function is nn.SmoothL1Loss). Defaults to 1.0, PyTorch’s default.

• gradient_clip (float) – Maximum gradient norm. Set to 0 to disable clipping.

activation: Module | None = None

beta: float = 1.0

eta_min: float = 0.1

gradient_clip: float = 0.0

learning_rate: float = 0.0003

property loss: Module

Returns the loss function to be used for training.

If the loss function is nn.SmoothL1Loss, it returns an instance with the specified beta. Otherwise, it returns the loss function as is.

loss_function: Module | None = None

momentum: float = 0.0

optimizer: str = 'adamw'

poly_power: float = 0.9

regularization_factor: float = 0.0

scheduler: str = 'cosine'

class codes.surrogates.AbstractSurrogate.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.