neural_lam.models.latent.base_decoder

neural_lam.models.latent.base_decoder#

Abstract base class for graph-based latent decoders.

Classes#

BaseGraphLatentDecoder

Abstract decoder mapping a grid representation plus a latent sample on

Module Contents#

class neural_lam.models.latent.base_decoder.BaseGraphLatentDecoder(hidden_dim: int, latent_dim: int, num_state_vars: int, hidden_layers: int = 1, output_std: bool = True)[source]#

Bases: torch.nn.Module

Abstract decoder mapping a grid representation plus a latent sample on mesh to the next-state increment (and optionally std) on the grid.

Subclasses implement combine_with_latent(), which fuses the latent representation with the grid representation. The resulting features are mapped to either num_state_vars outputs (mean increment only) or 2 * num_state_vars outputs (mean increment, std) depending on output_std.

Set up the latent embedder, grid-residual MLP and output param map.

Parameters:

  • hidden_dim (int) – Dimensionality of internal node and edge representations. Latent samples are embedded to this dimensionality before being fused with the grid representation.

  • latent_dim (int) – Dimensionality of the latent variable at each mesh node, i.e. the feature dimension of the latent_samples given to forward.

  • num_state_vars (int) – Number of state variables predicted at each grid node, i.e. the feature dimension of the predicted mean (and std).

  • hidden_layers (int) – Number of hidden layers in the internal MLPs (latent embedder, grid-residual MLP and output parameter map).

  • output_std (bool) – If True, the decoder outputs both mean and std of the next-state distribution (the output parameter map produces 2 * num_state_vars features per grid node); if False, only the mean.

abstract 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 representations and return a grid-shaped output.

Parameters:

  • original_grid_rep (torch.Tensor) – Shape (B, num_grid_nodes, d_h). Embedded grid input features, used as the sender representation when encoding the grid onto the mesh.

  • latent_rep (torch.Tensor) – Shape (B, num_mesh_nodes, d_h). Latent sample embedded to the internal dimensionality d_h. Where this enters the message passing is up to the concrete decoder.

  • residual_grid_rep (torch.Tensor) – Shape (B, num_grid_nodes, d_h). Residually updated grid representation, used as the receiver representation when decoding the mesh back to the grid. This keeps a direct path from the grid input to the output.

  • graph_emb (dict) – Embedded static graph node and edge features. The required entries depend on the concrete decoder, but include at least the g2m, m2m and m2g edge embeddings.

Returns:

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

Return type:

torch.Tensor

forward(grid_rep: torch.Tensor, latent_samples: torch.Tensor, graph_emb: dict[str, Any]) → tuple[torch.Tensor, torch.Tensor | None][source]#

Predict the next-state increment (and optionally std).

The latent samples are embedded to the internal dimensionality and fused with the grid representation by combine_with_latent; the result is mapped to distribution parameters. The predicted mean is an increment relative to the current state, which the caller adds onto the current state (and clamps).

Parameters:

  • grid_rep (torch.Tensor) – Shape (B, num_grid_nodes, d_h). Embedded grid input features (states, forcing and static features).

  • latent_samples (torch.Tensor) – Shape (B, num_mesh_nodes, latent_dim). Sample of the latent variable on the mesh nodes, e.g. drawn from the prior or the variational distribution.

  • graph_emb (dict) – Embedded static graph node and edge features, forwarded to combine_with_latent; includes at least the g2m, m2m and m2g edge embeddings.

Returns:

  • mean_delta (torch.Tensor) – Shape (B, num_grid_nodes, num_state_vars). Predicted increment of the next state relative to the current state.

  • pred_std (torch.Tensor or None) – Shape (B, num_grid_nodes, num_state_vars) when output_std is True, otherwise None. Predicted std of the next state, obtained from the output parameter map through a softplus.

grid_update_mlp#
latent_embedder#
output_std = True#
param_map#