neural_lam.models.latent.hi_graph_decoder
=========================================

.. py:module:: neural_lam.models.latent.hi_graph_decoder

.. autoapi-nested-parse::

   Latent decoder for hierarchical graphs.



Classes
-------

.. autoapisummary::

   neural_lam.models.latent.hi_graph_decoder.HiGraphLatentDecoder


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

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

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

   :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: Per-level edge indices of intra-level mesh edges, each of shape
                          ``(2, M_m2m[l])``.
   :type m2m_edge_index: BufferList
   :param m2g_edge_index: Shape ``(2, M_m2g)``. Edge index of mesh-to-grid edges.
   :type m2g_edge_index: torch.Tensor
   :param mesh_up_edge_index: Per-level edge indices of upward inter-level mesh edges, each of
                              shape ``(2, M_up[l])``.
   :type mesh_up_edge_index: BufferList
   :param mesh_down_edge_index: Per-level edge indices of downward inter-level mesh edges, each
                                of shape ``(2, M_down[l])``.
   :type mesh_down_edge_index: BufferList
   :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 intra_level_layers: Number of intra-level GNN layers at each mesh level; 0 disables
                              intra-level processing.
   :type intra_level_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``). Inter-level edges are not
                        configurable; mesh-up edges always use ``InteractionNet`` and
                        mesh-down edges always use ``PropagationNet``.
   :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

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

      :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 that initializes the bottom mesh level.
      :type original_grid_rep: torch.Tensor
      :param latent_rep: 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.
      :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
                        ``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``).
      :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:: intra_down_gnns


   .. py:attribute:: intra_up_gnns


   .. py:attribute:: m2g_gnn


   .. py:attribute:: mesh_down_gnns


   .. py:attribute:: mesh_up_gnns