noether.modeling.models.upt¶

Classes¶

`UPTConfig`	Configuration for a UPT model.
`UPT`	Implementation of the UPT (Universal Physics Transformer) model.

Module Contents¶

class noether.modeling.models.upt.UPTConfig(/, **data)¶

Bases: noether.core.models.base.ModelBaseConfig, noether.core.schemas.mixins.InjectSharedFieldFromParentMixin

Configuration for a UPT model.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Parameters:: data (Any)

model_config¶: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

num_heads: int = None¶: Number of attention heads in the model.

hidden_dim: int = None¶: Hidden dimension of the model.

mlp_expansion_factor: int = None¶: Expansion factor for the MLP of the FF layers.

approximator_depth: int = None¶: Number of approximator layers.

use_rope: bool = None¶

bias: bool = None¶: Whether to use bias terms in the model’s linear layers.

supernode_pooling_config: Annotated[noether.modeling.modules.encoders.supernode_pooling.SupernodePoolingConfig, noether.core.schemas.mixins.Shared]¶

approximator_config: Annotated[noether.modeling.modules.blocks.transformer.TransformerBlockConfig, noether.core.schemas.mixins.Shared]¶

decoder_config: Annotated[noether.modeling.modules.decoders.deep_perceiver.DeepPerceiverDecoderConfig, noether.core.schemas.mixins.Shared]¶

bias_layers: bool = None¶

data_specs: noether.data.schemas.ModelDataSpecs¶

linear_output_projection_config()¶

Return type:: noether.modeling.modules.layers.linear_projection.LinearProjectionConfig

rope_frequency_config()¶

Return type:: noether.modeling.modules.layers.rope_frequency.RopeFrequencyConfig

validate_rope_usage()¶

Ensure that if use_rope is True in the main config, it is also True in the approximator_config.

Return type:: UPTConfig

pos_embedding_config()¶

Return type:: noether.modeling.modules.layers.continuous_sincos_embed.ContinuousSincosEmbeddingConfig

validate_parameters()¶

Validate validity of parameters across the model and its submodules.

Ensures that: 1. hidden_dim is divisible by num_heads in parent and all submodules with num_heads 2. hidden_dim is consistent across parent and all submodules

Return type:: UPTConfig

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

Bases: torch.nn.Module

Implementation of the UPT (Universal Physics Transformer) model.

Parameters:: config (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:

geometry_batch_idx (torch.Tensor)
geometry_position (torch.Tensor)
geometry_supernode_idx (torch.Tensor)
query_position (torch.Tensor)

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: