emmi_inference.models.transolver

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

Transolver

Implementation of the Transolver model.

Module Contents

class emmi_inference.models.transolver.Transolver(dim=256, num_heads=8, depth=8, num_slices=64, mlp_expansion_factor=2, use_placeholder=True, attn_ctor=TransolverAttention, **kwargs)

Bases: emmi_inference.models.base_model.BaseModel

Implementation of the Transolver model. Reference code: https://github.com/thuml/Transolver/ Paper: https://arxiv.org/abs/2402.02366

Parameters:

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

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

  • num_heads (int) – number of heads of the multi-head self-attention. Defaults to 8.

  • depth (int) – number of Transolver blocks. Defaults to 8.

  • num_slices (int) – number of latent tokens/physics tokens used for down projecting the input points. Defaults to 64.

  • mlp_expansion_factor (int) – up projection ratio/factor of the MLP in the Transolver blocks. Defaults to 2.

  • use_placeholder (bool) – left over from the original code, no idea what the function is. Defaults to True.

  • attn_ctor (type) – constructor for the attention module. Defaults to TransolverAttention.

dim = 256
num_heads = 8
depth = 8
mlp_expansion_factor = 2
embed
blocks
forward(input_position, surface_mask_input, physics_features=None)

forward pass of the Transolver model.

Parameters:

  • input_position (torch.Tensor) – conncatenated surface and volume points coodindates.

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

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

Returns:

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

Return type:

dict[str, torch.Tensor]

kwarg_forward(input_position, surface_mask_input, physics_features=None)

forward pass of the Transolver model.

Parameters:

  • input_position (torch.Tensor) – conncatenated surface and volume points coodindates.

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

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

Returns:

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

Return type:

dict[str, torch.Tensor]