neural_lam.models.latent.hi_graph_decoder

neural_lam.models.latent.hi_graph_decoder#

Latent decoder for hierarchical graphs.

Classes#

HiGraphLatentDecoder

Latent decoder for a hierarchical mesh. The grid representation is

Module Contents#

class neural_lam.models.latent.hi_graph_decoder.HiGraphLatentDecoder(g2m_edge_index: torch.Tensor, m2m_edge_index: neural_lam.utils.BufferList, m2g_edge_index: torch.Tensor, mesh_up_edge_index: neural_lam.utils.BufferList, mesh_down_edge_index: neural_lam.utils.BufferList, hidden_dim: int, latent_dim: int, num_state_vars: int, intra_level_layers: int, hidden_layers: int = 1, g2m_gnn_type: str = 'InteractionNet', m2g_gnn_type: str = 'InteractionNet', output_std: bool = True)[source]#

Bases: neural_lam.models.latent.base_decoder.BaseGraphLatentDecoder

Latent decoder for a hierarchical mesh. The grid representation is encoded into the bottom mesh level; the message-passing then propagates up through the hierarchy (mixing in the latent at the top level), then down through the hierarchy with residual connections back to the intra-level reps from the upward pass, and finally maps back to grid via an m2g GNN (type set by m2g_gnn_type). The g2m GNN type is set by g2m_gnn_type; mesh-up edges always use InteractionNets and mesh-down edges always use PropagationNets.

Set up the g2m, m2g, mesh-up/-down and intra-level GNNs.

Parameters:

  • 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]).

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

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

  • mesh_down_edge_index (BufferList) – Per-level edge indices of downward inter-level mesh edges, each of shape (2, M_down[l]).

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

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

  • num_state_vars (int) – Number of state variables predicted at each grid node.

  • 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).

  • m2g_gnn_type (str) – GNN type for the mesh-to-grid step (key in gnn_layers.GNN_TYPES). Inter-level edges are not configurable; mesh-up edges always use InteractionNet and mesh-down edges always use PropagationNet.

  • output_std (bool) – If True, the decoder outputs both mean and std of the next-state distribution; if False, only the mean.

combine_with_latent(original_grid_rep: torch.Tensor, latent_rep: torch.Tensor, residual_grid_rep: torch.Tensor, graph_emb: dict[str, Any]) → torch.Tensor[source]#

Hierarchical up-then-down fusion of grid and latent reps.

Parameters:

  • original_grid_rep (torch.Tensor) – Shape (B, num_grid_nodes, d_h). Embedded grid input features, used as the sender representation in the grid-to-mesh step that initializes the bottom mesh level.

  • latent_rep (torch.Tensor) – Shape (B, num_mesh_nodes[L], d_h). Latent sample embedded to d_h, defined on the top mesh level L. It is used as the receiver representation of the last upward step, fusing the latent in at the top of the hierarchy.

  • residual_grid_rep (torch.Tensor) – Shape (B, num_grid_nodes, d_h). Residually updated grid representation, used as the receiver representation in the mesh-to-grid step.

  • graph_emb (dict) – Embedded static graph node and edge features, with at least the 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), mesh_up: list of (B, M_up[l], d_h), mesh_down: list of (B, M_down[l], d_h) and m2g: (B, M_m2g, d_h), where l indexes the mesh levels from bottom (0) to top (L).

Returns:

Shape (B, num_grid_nodes, d_h). Combined grid representation, incorporating both the grid input and the latent sample.

Return type:

torch.Tensor

g2m_gnn#
intra_down_gnns#
intra_up_gnns#
m2g_gnn#
mesh_down_gnns#
mesh_up_gnns#