neural_lam.models.latent.graph_encoder
======================================

.. py:module:: neural_lam.models.latent.graph_encoder

.. autoapi-nested-parse::

   Latent encoder for flat (non-hierarchical) graphs.



Classes
-------

.. autoapisummary::

   neural_lam.models.latent.graph_encoder.GraphLatentEncoder


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

.. py:class:: GraphLatentEncoder(latent_dim: int, g2m_edge_index: torch.Tensor, m2m_edge_index: torch.Tensor, hidden_dim: int, m2m_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 that maps grid features to mesh and outputs a Gaussian
   distribution over a latent variable on mesh nodes. Uses a flat
   (non-hierarchical) graph: one g2m GNN (type set by ``g2m_gnn_type``)
   followed by a stack of on-mesh (m2m) InteractionNet layers.

   Set up the g2m GNN, on-mesh processing stack 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: Shape ``(2, M_m2m)``. Edge index of mesh-to-mesh edges.
   :type m2m_edge_index: torch.Tensor
   :param hidden_dim: Dimensionality of internal node and edge representations.
   :type hidden_dim: 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 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 mesh from grid features.

      :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``: ``(B, num_mesh_nodes, d_h)``,
                        ``g2m``: ``(B, M_g2m, d_h)`` and ``m2m``: ``(B, M_m2m, d_h)``.
      :type graph_emb: dict
      :param \*\*kwargs: Ignored.

      :returns: Shape ``(B, num_mesh_nodes, output_dim)``. Raw parameters of
                the latent distribution.
      :rtype: torch.Tensor



   .. py:attribute:: g2m_gnn


   .. py:attribute:: latent_param_map


   .. py:attribute:: m2m_gnns