emmi_inference.models.upt

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

UPT

Implementation of the UPT (Universal Physics Transformer) model.

Module Contents

class emmi_inference.models.upt.UPT(ndim, radius, dim, num_heads, depth, dec_depth, mlp_expansion_factor, norm_ctor=nn.LayerNorm, spool_pos_mode='abspos', use_physics_features=False, physics_dim=None, **kwargs)

Bases: emmi_inference.models.base_model.BaseModel

Implementation of the UPT (Universal Physics Transformer) model.

Parameters:

  • BaseModel – Base model class that contains the utilities for all models we use in this tutorial.

  • ndim (int) – rank of the input data/tensor.

  • input_dim – dimensionality of the input features.

  • output_dim – dimensionality of the output features.

  • radius (float) – radius for the supernode pooling.

  • dim (int) – hidden dimensionality of the network.

  • num_heads (int) – number of attention heads.

  • depth (int) – depth (i.e., number of layers) of the encoder/approximator.

  • dec_depth (int) – number of decoder layers for the perceiver decoder.

  • mlp_expansion_factor (int) – expansion factor for the MLP in the transformer block.

  • norm_ctor (type) – cons. Defaults to nn.LayerNorm.

  • spool_pos_mode (str) – _description_. Defaults to “abspos”.

  • use_physics_features (bool)

  • physics_dim (int | None)

encoder
blocks
pred
forward(surface_mask_query, surface_position_batch_idx, surface_position_supernode_idx, surface_position, query_position, surface_features=None)

Forward pass of the UPT model.

Parameters:

  • input_position – input coordinates of the surface points.

  • surface_mask_query (torch.Tensor) – surface mask for the query points, indicating which points are surface points.

  • surface_position_batch_idx (torch.Tensor) – batch indices for the surface positions, since the surface positions are a sparse tensor for the supernode pooling.

  • surface_position_supernode_idx (torch.Tensor) – supernode indices for the surface positions.

  • geometry_position – geometry position information.

  • query_position (torch.Tensor) – input coordinates of the query points.

  • surface_features (torch.Tensor | None) – surface features for the input points. Defaults to None.

  • surface_position (torch.Tensor)

Returns:

dictionary with the output tensors, containing the surface pressure and volume velocity.

Return type:

dict[str, torch.Tensor]