neural_lam.models.latent.graph_encoder

neural_lam.models.latent.graph_encoder#

Latent encoder for flat (non-hierarchical) graphs.

Classes#

GraphLatentEncoder

Latent encoder that maps grid features to mesh and outputs a Gaussian

Module Contents#

class neural_lam.models.latent.graph_encoder.GraphLatentEncoder(latent_dim: int, g2m_edge_index: torch.Tensor, m2m_edge_index: torch.Tensor, hidden_dim: int, m2m_layers: int, hidden_layers: int = 1, g2m_gnn_type: str = 'InteractionNet', output_dist: str = 'isotropic')[source]#

Bases: neural_lam.models.latent.base_encoder.BaseLatentEncoder

Latent encoder that maps grid features to mesh and outputs a Gaussian distribution over a latent variable on mesh nodes. Uses a flat (non-hierarchical) graph: one g2m GNN (type set by g2m_gnn_type) followed by a stack of on-mesh (m2m) InteractionNet layers.

Set up the g2m GNN, on-mesh processing stack and latent param map.

Parameters:

  • latent_dim (int) – Dimensionality of the latent variable at each mesh node.

  • g2m_edge_index (torch.Tensor) – Shape (2, M_g2m). Edge index of grid-to-mesh edges.

  • m2m_edge_index (torch.Tensor) – Shape (2, M_m2m). Edge index of mesh-to-mesh edges.

  • hidden_dim (int) – Dimensionality of internal node and edge representations.

  • m2m_layers (int) – Number of on-mesh (m2m) GNN layers; 0 disables on-mesh processing.

  • hidden_layers (int) – Number of hidden layers in internal MLPs.

  • g2m_gnn_type (str) – GNN type for the grid-to-mesh step (key in gnn_layers.GNN_TYPES).

  • output_dist (str) – Type of output distribution: "isotropic" or "diagonal".

compute_dist_params(grid_rep: torch.Tensor, graph_emb: dict[str, Any], **kwargs: Any) → torch.Tensor[source]#

Compute distribution parameters on mesh from grid features.

Parameters:

  • grid_rep (torch.Tensor) – Shape (B, num_grid_nodes, d_h). Grid input representation.

  • graph_emb (dict) – Embedded graph node and edge features, with at least entries mesh: (B, num_mesh_nodes, d_h), g2m: (B, M_g2m, d_h) and m2m: (B, M_m2m, d_h).

  • **kwargs – Ignored.

Returns:

Shape (B, num_mesh_nodes, output_dim). Raw parameters of the latent distribution.

Return type:

torch.Tensor

g2m_gnn#
latent_param_map#
m2m_gnns#