noether.modeling.models.ab_upt

Classes

AnchoredBranchedUPT

Implementation of the Anchored Branched UPT model. Including input embedding and output projection, so this is an off-the-shelf model that can be used directly by providing the appropriate input tensors.

Module Contents

class noether.modeling.models.ab_upt.AnchoredBranchedUPT(config)

Bases: torch.nn.Module

Implementation of the Anchored Branched UPT model. Including input embedding and output projection, so this is an off-the-shelf model that can be used directly by providing the appropriate input tensors.

Parameters:

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

data_specs
rope
pos_embed
encoder
geometry_blocks
surface_bias
volume_bias
num_perceivers = 0
physics_blocks
use_geometry_branch = False
surface_decoder_blocks
volume_decoder_blocks
surface_decoder
volume_decoder
geometry_branch_forward(geometry_position, geometry_supernode_idx, geometry_batch_idx, condition, geometry_attn_kwargs)

Forward pass through the geometry branch of the model.

Parameters:
Return type:

torch.Tensor

physics_blocks_forward(surface_position_all, volume_position_all, geometry_encoding, physics_token_specs, physics_attn_kwargs, physics_perceiver_attn_kwargs, condition)

Forward pass through the physics blocks of the model. Allthough in the AB-UPT paper we only have a perceiver block a the first block, it is possible to have more perceiver blocks in the physics blocks that attend to the geometry encoding.

Parameters:
Return type:

torch.Tensor

decoder_blocks_forward(x_physics, physics_token_specs, surface_token_specs, volume_token_specs, surface_position_all, volume_position_all, surface_decoder_attn_kwargs, volume_decoder_attn_kwargs, condition)

Forward pass through the decoder blocks of the model. We have a separate decoder for surface and volume tokens.

Parameters:
Return type:

tuple[torch.Tensor, torch.Tensor]

create_rope_frequencies(geometry_position, geometry_supernode_idx, surface_position_all, volume_position_all)

Create RoPE frequencies for all relevant positions.

Parameters:
forward(geometry_position, geometry_supernode_idx, geometry_batch_idx, surface_anchor_position, volume_anchor_position, geometry_design_parameters=None, inflow_design_parameters=None, query_surface_position=None, query_volume_position=None)

Forward pass of the AB-UPT model.

Parameters:

  • geometry_position (torch.Tensor) – Coordinates of the geometry mesh. Tensor of shape (B * N_geometry, D_pos), sparse tensor

  • geometry_supernode_idx (torch.Tensor) – Indices of the supernodes for the geometry points. Tensor of shape (B * number of super nodes,)

  • geometry_batch_idx (torch.Tensor | None) – Batch indices for the geometry points. Tensor of shape (B * N_geometry,). If None, assumes all points belong to the same batch.

  • surface_anchor_position (torch.Tensor) – Coordinates of the surface anchor points. Tensor of shape (B, N_surface_anchor, D_pos)

  • volume_anchor_position (torch.Tensor) – Coordinates of the volume anchor points. Tensor of shape (B, N_volume_anchor, D_pos)

  • geometry_design_parameters (torch.Tensor | None) – Design parameters related to the geometry to condition on. Tensor of shape (B, D_geom)

  • inflow_design_parameters (torch.Tensor | None) – Design parameters related to the inflow to condition on. Tensor of shape (B, D_inflow).

  • query_surface_position (torch.Tensor | None) – Coordinates of the query surface points.

  • query_volume_position (torch.Tensor | None) – Coordinates of the query volume points.

Returns:

A dictionary containing the predictions for surface and volume fields, sliced according to the data specifications.

Return type:

dict[str, torch.Tensor]