neural_lam.models.latent.graph_decoder
======================================

.. py:module:: neural_lam.models.latent.graph_decoder

.. autoapi-nested-parse::

   Latent decoder for flat (non-hierarchical) graphs.



Classes
-------

.. autoapisummary::

   neural_lam.models.latent.graph_decoder.GraphLatentDecoder


Module Contents
---------------

.. py:class:: 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)

   Bases: :py:obj:`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.

   :param g2m_edge_index: Shape ``(2, M_g2m)``. Edge index of grid-to-mesh edges.
   :type g2m_edge_index: torch.Tensor
   :param m2m_edge_index: Shape ``(2, M_m2m)``. Edge index of mesh-to-mesh edges.
   :type m2m_edge_index: torch.Tensor
   :param m2g_edge_index: Shape ``(2, M_m2g)``. Edge index of mesh-to-grid edges.
   :type m2g_edge_index: torch.Tensor
   :param hidden_dim: Dimensionality of internal node and edge representations.
   :type hidden_dim: int
   :param latent_dim: Dimensionality of the latent variable at each mesh node.
   :type latent_dim: int
   :param num_state_vars: Number of state variables predicted at each grid node.
   :type num_state_vars: int
   :param m2m_layers: Number of on-mesh (m2m) GNN layers; 0 disables on-mesh
                      processing.
   :type m2m_layers: int
   :param hidden_layers: Number of hidden layers in internal MLPs.
   :type hidden_layers: int
   :param g2m_gnn_type: GNN type for the grid-to-mesh step (key in
                        ``gnn_layers.GNN_TYPES``).
   :type g2m_gnn_type: str
   :param m2g_gnn_type: GNN type for the mesh-to-grid step (key in
                        ``gnn_layers.GNN_TYPES``).
   :type m2g_gnn_type: str
   :param output_std: If True, the decoder outputs both mean and std of the next-state
                      distribution; if False, only the mean.
   :type output_std: bool


   .. py:method:: combine_with_latent(original_grid_rep: torch.Tensor, latent_rep: torch.Tensor, residual_grid_rep: torch.Tensor, graph_emb: dict[str, Any]) -> torch.Tensor

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

      :param original_grid_rep: Shape ``(B, num_grid_nodes, d_h)``. Embedded grid input
                                features, used as the sender representation in the
                                grid-to-mesh step.
      :type original_grid_rep: torch.Tensor
      :param latent_rep: 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.
      :type latent_rep: torch.Tensor
      :param residual_grid_rep: Shape ``(B, num_grid_nodes, d_h)``. Residually updated grid
                                representation, used as the receiver representation in the
                                mesh-to-grid step.
      :type residual_grid_rep: torch.Tensor
      :param graph_emb: 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)``.
      :type graph_emb: dict

      :returns: Shape ``(B, num_grid_nodes, d_h)``. Combined grid
                representation, incorporating both the grid input and the
                latent sample.
      :rtype: torch.Tensor



   .. py:attribute:: g2m_gnn


   .. py:attribute:: m2g_gnn


   .. py:attribute:: m2m_gnns