noether.modeling.models.upt

Classes

UPT

Implementation of the UPT (Universal Physics Transformer) model.

Module Contents

class noether.modeling.models.upt.UPT(config)

Bases: torch.nn.Module

Implementation of the UPT (Universal Physics Transformer) model.

Parameters:

config (noether.core.schemas.models.UPTConfig) – Configuration for the UPT model. See UPTConfig for details.

encoder
use_rope
pos_embed
approximator_blocks
decoder
norm
prediction_layer
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 use RoPE, return empty dicts.

Parameters:
Return type:

tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]

forward(surface_position_batch_idx, surface_position_supernode_idx, surface_position, query_position)

Forward pass of the UPT model.

Parameters:

  • 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.

  • surface_position (torch.Tensor) – Input coordinates of the surface points.

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

Returns:

Output tensor containing the predictions for the surface and volume fields, sliced according to the data specifications.

Return type:

torch.Tensor