neural_lam.models.latent.hi_graph_encoder
=========================================

.. py:module:: neural_lam.models.latent.hi_graph_encoder

.. autoapi-nested-parse::

   Latent encoder for hierarchical graphs.



Classes
-------

.. autoapisummary::

   neural_lam.models.latent.hi_graph_encoder.HiGraphLatentEncoder


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

.. py:class:: HiGraphLatentEncoder(latent_dim: int, g2m_edge_index: torch.Tensor, m2m_edge_index: neural_lam.utils.BufferList, mesh_up_edge_index: neural_lam.utils.BufferList, hidden_dim: int, intra_level_layers: int, hidden_layers: int = 1, g2m_gnn_type: str = 'InteractionNet', output_dist: str = 'isotropic')

   Bases: :py:obj:`neural_lam.models.latent.base_encoder.BaseLatentEncoder`


   Latent encoder for a hierarchical mesh: grid -> bottom mesh level via a
   g2m GNN (type set by ``g2m_gnn_type``), then propagates upward through
   mesh levels using mesh-up PropagationNets, with optional intra-level
   processing at each level. The latent distribution is read out from the
   top mesh level.

   Set up the g2m, mesh-up and intra-level GNNs and latent param map.

   :param latent_dim: Dimensionality of the latent variable at each mesh node.
   :type latent_dim: int
   :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 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 hidden_dim: Dimensionality of internal node and edge representations.
   :type hidden_dim: 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``). Mesh-up edges are not configurable;
                        they always use ``PropagationNet``.
   :type g2m_gnn_type: str
   :param output_dist: Type of output distribution: ``"isotropic"`` or ``"diagonal"``.
   :type output_dist: str


   .. py:method:: compute_dist_params(grid_rep: torch.Tensor, graph_emb: dict[str, Any], **kwargs: Any) -> torch.Tensor

      Compute distribution parameters on the top mesh level.

      :param grid_rep: Shape ``(B, num_grid_nodes, d_h)``. Grid input representation.
      :type grid_rep: torch.Tensor
      :param graph_emb: Embedded graph node and edge features, with at least 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)`` and
                        ``mesh_up``: list of ``(B, M_up[l], d_h)``.
      :type graph_emb: dict
      :param \*\*kwargs: Ignored.

      :returns: Shape ``(B, num_mesh_nodes[L], output_dim)``. Raw parameters
                of the latent distribution on the top mesh level ``L``.
      :rtype: torch.Tensor



   .. py:attribute:: g2m_gnn


   .. py:attribute:: intra_level_gnns


   .. py:attribute:: latent_param_map


   .. py:attribute:: mesh_up_gnns