emmi_inference.models.base_model

WARNING This file is a 1:1 duplicate from the tutorial folder.

It is here to avoid installation of a tutorial as a package and keep it isolated.

Classes

BaseModel

Base class for all model we use in this tutorial.

Functions

validate_required_batch_input(func)

Decorator to validate batch input for methods that require it.

Module Contents

emmi_inference.models.base_model.validate_required_batch_input(func)

Decorator to validate batch input for methods that require it.

class emmi_inference.models.base_model.BaseModel(input_dim, output_dim, required_batch_modes, dim=None, output_projection=False, use_physics_features=False, physics_dim=None, bias_layers=True, position_projection='sincos', **kwargs)

Bases: ksuit.models.Model

Base class for all model we use in this tutorial.

Parameters:

  • Model – Base class for single models in ksuit.

  • input_dim (int) – dimension of the input features.

  • output_dim (int) – dimension of the output prediction.

  • required_batch_keys – defines the required batch keys for the model. This is used to validate the input batch.

  • dim (int | None) – dimension of the latents of the model.

  • output_projection (bool) – Boolean to indicate to use the output projection. Defaults to False.

  • use_physics_features (bool) – Boolean to indicate to use additional physics features (next to the input coordiates). Defaults to False.

  • physics_dim (int | None) – dimension of the physics features. Defaults to None.

  • bias_layers (bool) – Boolean to indicate to use bias layers. Defaults to True.

  • position_projection (str) – String to indicate the type of position projection to use. Can be “sincos” or “linear”. Defaults to “sincos”.

  • required_batch_modes (list[str])

input_dim
output_dim
physics_dim = None
use_physics_features = False
required_batch_modes
position_projection = 'sincos'
output_projection(x)

Most model implementations will have an output projection layer that maps the last latent vector into the output physics space. We have a unified projection layer that can be used in all models.

Parameters:

x (torch.Tensor) – tensor of shape (batch_size, num_points, dim) containing the features for each point.

Returns:

tensor of shape (batch_size, num_points, output_dim) containing the projected features into (normalized) physics space.

Return type:

torch.Tensor

surface_and_volume_bias(x, surface_mask)

For some of the models, the surface and volume are concatenated into a single input tesnor (e.g., Pointnet, Transolver, Transformer). For AB-UPT, we shared weight for the physics blocks. Hence, we need to indicate which points are surface and which are volume points. We do this by applying a bias (i.e., an MLP) to the surface and volume points separately. The surface mask indicates which points are surface points. This function only works for tensors where surface and volume points are concatenated along the second dimension (not for AB-UPT). Howerver, for other models, self.surface_bias and self.volume_bias can be called directly in the child class.

Parameters:

  • x (torch.Tensor) – tensor of shape (batch_size, num_points, input_dim) containing the features for each point.

  • surface_mask (torch.Tensor) – Boolean tensor of shape (batch_size, num_points) indicating which points are surface points.

Returns:

biased tensor x of shape (batch_size, num_points, input_dim) where the surface points have been processed by the surface bias and the volume points by the volume bias.

Return type:

torch.Tensor

gather_outputs(x, surface_mask)

The output gathering function is used to extract the relevant outputs from the model’s output tensor. It assumes that the output tensor has a specific structure, where the first dimension corresponds to the batch size, the second dimension corresponds to the surface/volume points, and the third dimension corresponds to the output features

The surface pressure is expected to be at index 0, the volume velocity at indices 1:4, and if the output dimension is 11, the surface wall shear stress is at indices 4:7, the volume total pressure coefficient at index 7, and the volume vorticity at indices 8:11. These last three are only available for AhmedML and DriverML datasets.

Parameters:

  • x (torch.Tensor) – output tensor from the model, shape (batch_size, num_points, output_dim)

  • surface_mask (torch.Tensor) – Indicator boolean tensor for surface points, shape (batch_size, num_points). All surface points should be True, and all volume points should be False.

Returns:

A dictionary containing the gathered outputs:

  • ”surface_pressure”: Tensor of shape (batch_size, num_surface_points, 1)

  • ”volume_velocity”: Tensor of shape (batch_size, num_volume_points, 3)

  • ”surface_wallshearstress”: Tensor of shape (batch_size, num_surface_points, 3) if output_dim is 11

  • ”volume_totalpcoeff”: Tensor of shape (batch_size, num_volume_points, 1) if output_dim is 11

  • ”volume_vorticity”: Tensor of shape (batch_size, num_volume_points, 3) if output_dim is 11

Return type:

dict[str, torch.Tensor]