neural_lam.models.latent.hi_graph_encoder

neural_lam.models.latent.hi_graph_encoder#

Latent encoder for hierarchical graphs.

Classes#

HiGraphLatentEncoder

Latent encoder for a hierarchical mesh: grid -> bottom mesh level via a

Module Contents#

class neural_lam.models.latent.hi_graph_encoder.HiGraphLatentEncoder(latent_dim: int, g2m_edge_index: torch.Tensor, m2m_edge_index: neural_lam.utils.BufferList, mesh_up_edge_index: neural_lam.utils.BufferList, hidden_dim: int, intra_level_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 for a hierarchical mesh: grid -> bottom mesh level via a g2m GNN (type set by g2m_gnn_type), then propagates upward through mesh levels using mesh-up PropagationNets, with optional intra-level processing at each level. The latent distribution is read out from the top mesh level.

Set up the g2m, mesh-up and intra-level GNNs 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 (BufferList) – Per-level edge indices of intra-level mesh edges, each of shape (2, M_m2m[l]).

  • mesh_up_edge_index (BufferList) – Per-level edge indices of upward inter-level mesh edges, each of shape (2, M_up[l]).

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

  • intra_level_layers (int) – Number of intra-level GNN layers at each mesh level; 0 disables intra-level 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). Mesh-up edges are not configurable; they always use PropagationNet.

  • 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 the top mesh level.

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: list of (B, num_mesh_nodes[l], d_h), g2m: (B, M_g2m, d_h), m2m: list of (B, M_m2m[l], d_h) and mesh_up: list of (B, M_up[l], d_h).

  • **kwargs – Ignored.

Returns:

Shape (B, num_mesh_nodes[L], output_dim). Raw parameters of the latent distribution on the top mesh level L.

Return type:

torch.Tensor

g2m_gnn#
intra_level_gnns#
latent_param_map#
mesh_up_gnns#