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.

use_rope
encoder
pos_embed
approximator_blocks
decoder
norm
prediction_layer
compute_rope_args(geometry_batch_idx, geometry_position, geometry_supernode_idx, query_position)

Compute the RoPE frequency arguments for the geometry and query positions. If RoPE is not used, return empty dicts.

Parameters:
Return type:

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

forward(geometry_batch_idx, geometry_supernode_idx, geometry_position, query_position)

Forward pass of the UPT model.

Parameters:

  • geometry_batch_idx (torch.Tensor) – Batch indices for the geometry positions.

  • geometry_supernode_idx (torch.Tensor) – Supernode indices for the geometry positions.

  • geometry_position (torch.Tensor) – Input coordinates of the geometry mesh points.

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

Returns:

Output tensor containing the predictions at query positions.

Return type:

torch.Tensor