upt

Classes

UPT

Implementation of the UPT (Universal Physics Transformer) model.

Module Contents

class upt.UPT(model_config)

Bases: torch.nn.Module

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.

  • model_config (emmi.schemas.models.UPTConfig) – Configuration for the UPT model.

encoder
use_rope
pos_embed
approximator_blocks
decoder
norm
prediction_layer
bias_layers = False
surface_and_volume_bias(x, surface_mask)

Apply separate bias layers for surface and volume points.

Parameters:

  • x (torch.Tensor)

  • surface_mask (torch.Tensor)

Return type:

torch.Tensor

compute_rope_args(surface_position_batch_idx, surface_position, surface_position_supernode_idx, query_position)

Compute the RoPE frequency arguments for the surface_position and query_position. If we don’t uise RoPE, return empty dicts. :param surface_position_batch_idx: Batch indices for the surface positions. :param surface_position: Surface position coordinates. :param surface_position_supernode_idx: Supernode indices for the surface positions. :param query_position: Query position coordinates.

Returns:

Dictionary containing the RoPE frequency arguments.

Return type:

dict[str, torch.Tensor]

Parameters:

  • surface_position_batch_idx (torch.Tensor)

  • surface_position (torch.Tensor)

  • surface_position_supernode_idx (torch.Tensor)

  • query_position (torch.Tensor)

forward(surface_position_batch_idx, surface_position_supernode_idx, surface_position, query_position, surface_mask_query, input_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 – surface features for the input points. Defaults to None.

  • surface_position (torch.Tensor)

  • input_features (torch.Tensor | None)

Returns:

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

Return type:

dict[str, torch.Tensor]