Surrogate

class aido.surrogate.NoiseAdder(n_time_steps: int, betas=(0.0001, 0.02))

Bases: Module

forward(x, t)

x: (B, C, H, W) t: (B, 1) z: (B, C, H, W) x_t: (B, C, H, W)

class aido.surrogate.Surrogate(num_parameters: int, num_context: int, num_targets: int, num_reconstructed: int, initial_means: List[float32], initial_stds: List[float32], n_time_steps: int = 50, betas: Tuple[float, float] = (0.0001, 0.02))

Bases: Module

Surrogate model class and the surrogate model training function, given a dataset consisting of events. The surrogate model itself can be very simple. It is just a feed-forward model but used as a diffusion model.

betas

Tuple containing the start and end beta values for the diffusion process.

Type:

Tuple[float]

t_is

Tensor containing time steps normalized by the number of time steps.

Type:

torch.Tensor

forward(parameters, context, reconstructed, time_step)

Forward pass of the model. Concatenates the input features and passes them through the network.

to(device=None)

Moves the model and its buffers to the specified device.

create_noisy_input(x)

Adds noise to a tensor for the diffusion process.

sample_forward(parameters, context)

Samples from the model in a forward pass using the diffusion process.

train_model(surrogate_dataset, batch_size, n_epochs, lr)

Trains the surrogate diffusion model using the provided dataset.

apply_model_in_batches(dataset, batch_size, oversample=1)

Applies the model to the dataset in batches and returns the results.

__init__(num_parameters: int, num_context: int, num_targets: int, num_reconstructed: int, initial_means: List[float32], initial_stds: List[float32], n_time_steps: int = 50, betas: Tuple[float, float] = (0.0001, 0.02))

Initializes the surrogate model.

Parameters:

• num_parameters (int) – Number of input parameters.

• num_context (int) – Number of context variables.

• num_reconstructed (int) – Number of reconstructed variables.

• n_time_steps (int, optional) – Number of time steps for the DDPM schedule. Defaults to 50. Setting it higher might lead to divergence towards infinity which will register as NaN when reaching float32 accuracy ~= 2e9.

• betas (Tuple[float], optional) – Tuple containing the start and end values for the beta schedule. Defaults to (1e-4, 0.02).

apply_model_in_batches(dataset: SurrogateDataset, batch_size: int, oversample: int = 1) → Tensor

Apply the model to the given dataset in batches.

Parameters:

• dataset (SurrogateDataset) – The dataset to apply the model to.

• batch_size (int) – The size of each batch.

• oversample (int, optional) – The number of times to oversample the dataset, by default 1.

Returns:

torch.Tensor – The surrogate model’s predictions.

Notes

In most cases, the resulting tensor with sampled data is not of importance, as the main value lies in the trained model weights.

create_noisy_input(x: Tensor, scale: float = 1.0) → Tuple[Tensor, Tensor, Tensor]

Add gaussian noise to a tensor.

Scale the noise with ‘scale’, by default the noise is N(0, 1).

Parameters:

• x (torch.Tensor) – The input tensor to which noise will be added.

• scale (float, default=1.0) – The scale factor for the noise.

Returns:

Tuple[torch.Tensor, torch.Tensor, torch.Tensor] –

• The noisy tensor

• The noise added to the input tensor

• The time steps used for generating the noise

forward(parameters: Tensor, context: Tensor, targets: Tensor, reconstructed: Tensor, time_step: Tensor)

When sampling forward, ‘parameters’ has only one entry, therefore it is broadcast to the shape of ‘context’.

to(device: str | None = None) → Self

Move whole model and data to device

Parameters:

device (str | None)

Returns:

Self

train_model(surrogate_dataset: SurrogateDataset, batch_size: int, n_epochs: int, lr: float) → float

Train the Surrogate Diffusion model.

The training loop includes noise addition as part of the diffusion process.

Parameters:

• surrogate_dataset (SurrogateDataset) – The dataset containing the training data.

• batch_size (int) – The size of each training batch.

• n_epochs (int) – The number of training epochs.

• lr (float) – The learning rate for the optimizer.

Returns:

float – The final training loss value.

class aido.surrogate.SurrogateDataset(input_df: DataFrame, parameter_key: str = 'Parameters', context_key: str = 'Context', target_key: str = 'Targets', reconstructed_key: str = 'Reconstructed', device: str = 'cpu', means: List[float32] | None = None, stds: List[float32] | None = None, normalize_parameters: bool = False)

Bases: Dataset

Dataset class for the Surrogate model

TODO: Accommodate for discrete parameters

__init__(input_df: DataFrame, parameter_key: str = 'Parameters', context_key: str = 'Context', target_key: str = 'Targets', reconstructed_key: str = 'Reconstructed', device: str = 'cpu', means: List[float32] | None = None, stds: List[float32] | None = None, normalize_parameters: bool = False)

Initializes the Surrogate model with the provided DataFrame and keys. All inputs must be unnormalized and will be normalized internally.

Parameters:

• input_df (pd.DataFrame) – The input DataFrame containing the data.

• parameter_key (str, optional) – The key for the parameters column in the DataFrame. Defaults to “Parameters”.

• context_key (str, optional) – The key for the context column in the DataFrame. Defaults to “Context”.

• target_key (str, optional) – The key for the target column in the DataFrame. Defaults to “Targets”.

• reconstructed_key (str, optional) – The key for the reconstruction loss column in the DataFrame. Defaults to “Loss”.

• device (str) – Torch device. Defaults to ‘cuda’ if available, else ‘cpu’.

• means (List[np.float32], optional) – Predefined means for normalization. Defaults to None.

• stds (List[np.float32], optional) – Predefined standard deviations for normalization. Defaults to None.

• normalize_parameters (bool, optional) – Whether to normalize parameters. Defaults to False.

filter_infs_and_nans(df: DataFrame) → DataFrame

Removes all events that contain infs or nans.

normalize_features(target: Tensor | ndarray, index: int) → Tensor | ndarray

Normalize a feature using stored means and standard deviations.

Parameters:

• target (torch.Tensor or np.ndarray) – The feature to normalize.

• index (int) – Index indicating the feature type: - 0: Parameters - 1: Context - 2: Targets

Returns:

torch.Tensor or np.ndarray – The normalized feature.

unnormalize_features(target: Tensor | ndarray, index: int) → Tensor | ndarray

Convert normalized features back to their original scale.

Parameters:

• target (torch.Tensor or np.ndarray) – The normalized features to convert back.

• index (int) – Index indicating the feature type: - 0: Parameters - 1: Context - 2: Targets

Returns:

torch.Tensor or np.ndarray – The unnormalized features in their original scale.

aido.surrogate.ddpm_schedules(beta1: float, beta2: float, n_time_steps: int) → dict[str, Tensor]

No-index:

Returns pre-computed schedules for DDPM sampling, training process.