ab_upt

Classes

AnchoredBranchedUPT

Implementation of the Anchored Branched UPT model.

Module Contents

class ab_upt.AnchoredBranchedUPT(config)

Bases: torch.nn.Module

Implementation of the Anchored Branched UPT model.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

config (emmi.schemas.models.AnchorBranchedUPTConfig)

data_specs
rope
encoder
geometry_blocks
pos_embed
surface_bias
volume_bias
num_perceivers = 0
physics_blocks
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:

  • geometry_position (torch.Tensor)

  • geometry_supernode_idx (torch.Tensor)

  • geometry_batch_idx (torch.Tensor)

  • condition (torch.Tensor | None)

  • geometry_attn_kwargs (dict[str, torch.Tensor])

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:

  • surface_position_all (torch.Tensor) – Tensor of shape (B, N_surface_total, D_pos)

  • volume_position_all (torch.Tensor) – Tensor of shape (B, N_volume_total, D_pos)

  • geometry_encoding (torch.Tensor) – Tensor of shape (B, N_supernodes, D_hidden)

  • physics_token_specs (list[emmi.schemas.modules.attention.anchor_attention.TokenSpec]) – List of TokenSpec defining the token specifications for the physics blocks.

  • physics_attn_kwargs (dict[str, torch.Tensor]) – Additional attention kwargs for the physics transformer blocks.

  • physics_perceiver_attn_kwargs (dict[str, torch.Tensor]) – Additional attention kwargs for the physics perceiver blocks.

  • condition (torch.Tensor | None) – Optional conditioning tensor of shape (B, D_condition)

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:

  • x_physics (torch.Tensor)

  • physics_token_specs (list[emmi.schemas.modules.attention.anchor_attention.TokenSpec])

  • surface_token_specs (list[emmi.schemas.modules.attention.anchor_attention.TokenSpec])

  • volume_token_specs (list[emmi.schemas.modules.attention.anchor_attention.TokenSpec])

  • surface_position_all (torch.Tensor)

  • volume_position_all (torch.Tensor)

  • surface_decoder_attn_kwargs (dict[str, torch.Tensor])

  • volume_decoder_attn_kwargs (dict[str, torch.Tensor])

  • condition (torch.Tensor | None)

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:

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

  • geometry_supernode_idx (torch.Tensor) – Tensor of shape (B * number of super nodes,) with indices of supernodes

  • surface_position_all (torch.Tensor) – Tensor of shape (B, N_surface_total, D_pos)

  • volume_position_all (torch.Tensor) – Tensor of shape (B, N_volume_total, D_pos)

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. # TODO: when writing the docs, we have to use consistent notation. :param geometry_position: Coordinates of the geometry mesh. Tensor of shape (B * N_geometry, D_pos), sparse tensor :param geometry_supernode_idx: Indices of the supernodes for the geometry points. Tensor of shape (B * number of super nodes,) :param geometry_batch_idx: Batch indices for the geometry points. Tensor of shape (B * N_geometry,). If None, assumes all points belong to the same batch. :param surface_anchor_position: Coordinates of the surface anchor points. Tensor of shape (B, N_surface_anchor, D_pos) :param volume_anchor_position: Coordinates of the volume anchor points. Tensor of shape (B, N_volume_anchor, D_pos) :param geometry_design_parameters: Design parameters related to the geometry to condition on. Tensor of shape (B, D_geom) :param inflow_design_parameters: Design parameters related to the inflow to condition on. Tensor of shape (B, D_inflow). :param query_surface_position: Coordinates of the query surface points. :param query_volume_position: Coordinates of the query volume points.

Parameters:

  • geometry_position (torch.Tensor)

  • geometry_supernode_idx (torch.Tensor)

  • geometry_batch_idx (torch.Tensor | None)

  • surface_anchor_position (torch.Tensor)

  • volume_anchor_position (torch.Tensor)

  • geometry_design_parameters (torch.Tensor | None)

  • inflow_design_parameters (torch.Tensor | None)

  • query_surface_position (torch.Tensor | None)

  • query_volume_position (torch.Tensor | None)

Return type:

dict[str, torch.Tensor]