neural_lam.models.latent.base_decoder
=====================================

.. py:module:: neural_lam.models.latent.base_decoder

.. autoapi-nested-parse::

   Abstract base class for graph-based latent decoders.



Classes
-------

.. autoapisummary::

   neural_lam.models.latent.base_decoder.BaseGraphLatentDecoder


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

.. py:class:: BaseGraphLatentDecoder(hidden_dim: int, latent_dim: int, num_state_vars: int, hidden_layers: int = 1, output_std: bool = True)

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

   :param hidden_dim: Dimensionality of internal node and edge representations.
                      Latent samples are embedded to this dimensionality before
                      being fused with the grid representation.
   :type hidden_dim: int
   :param latent_dim: Dimensionality of the latent variable at each mesh node, i.e.
                      the feature dimension of the ``latent_samples`` given to
                      ``forward``.
   :type latent_dim: int
   :param num_state_vars: Number of state variables predicted at each grid node, i.e.
                          the feature dimension of the predicted mean (and std).
   :type num_state_vars: int
   :param hidden_layers: Number of hidden layers in the internal MLPs (latent embedder,
                         grid-residual MLP and output parameter map).
   :type hidden_layers: int
   :param output_std: 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.
   :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
      :abstractmethod:


      Fuse grid and latent representations and return a grid-shaped output.

      :param original_grid_rep: Shape ``(B, num_grid_nodes, d_h)``. Embedded grid input
                                features, used as the sender representation when encoding the
                                grid onto the mesh.
      :type original_grid_rep: torch.Tensor
      :param latent_rep: 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.
      :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 when
                                decoding the mesh back to the grid. This keeps a direct path
                                from the grid input to the output.
      :type residual_grid_rep: torch.Tensor
      :param graph_emb: 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.
      :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:method:: forward(grid_rep: torch.Tensor, latent_samples: torch.Tensor, graph_emb: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor | None]

      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).

      :param grid_rep: Shape ``(B, num_grid_nodes, d_h)``. Embedded grid input
                       features (states, forcing and static features).
      :type grid_rep: torch.Tensor
      :param latent_samples: 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.
      :type latent_samples: torch.Tensor
      :param graph_emb: Embedded static graph node and edge features, forwarded to
                        ``combine_with_latent``; includes at least the ``g2m``,
                        ``m2m`` and ``m2g`` edge embeddings.
      :type graph_emb: dict

      :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.



   .. py:attribute:: grid_update_mlp


   .. py:attribute:: latent_embedder


   .. py:attribute:: output_std
      :value: True



   .. py:attribute:: param_map