neural_lam.models.latent.graph_decoder

neural_lam.models.latent.graph_decoder#

Latent decoder for flat (non-hierarchical) graphs.

Classes#

GraphLatentDecoder

Latent decoder for a flat (non-hierarchical) graph. Encodes grid into

Module Contents#

class neural_lam.models.latent.graph_decoder.GraphLatentDecoder(g2m_edge_index: torch.Tensor, m2m_edge_index: torch.Tensor, m2g_edge_index: torch.Tensor, hidden_dim: int, latent_dim: int, num_state_vars: int, m2m_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 flat (non-hierarchical) graph. Encodes grid into mesh with a g2m GNN (type set by g2m_gnn_type), processes on mesh, and reads back out to grid with an m2g GNN (type set by m2g_gnn_type). The grid representation also goes through a residual MLP that is added back to the mesh-to-grid output.

Set up the g2m, on-mesh and m2g GNNs.

Parameters:

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

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

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

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

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

  • 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]#

Fuse grid and latent reps via g2m -> m2m -> m2g.

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.

  • latent_rep (torch.Tensor) – Shape (B, num_mesh_nodes, d_h). Latent sample embedded to d_h, used as the initial mesh node representation (the receiver in the grid-to-mesh step), so all mesh processing starts from the latent.

  • 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 g2m: (B, M_g2m, d_h), m2m: (B, M_m2m, d_h) and m2g: (B, M_m2g, d_h).

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#
m2g_gnn#
m2m_gnns#