neural_lam.models.latent.base_encoder
=====================================

.. py:module:: neural_lam.models.latent.base_encoder

.. autoapi-nested-parse::

   Abstract base class for latent encoders.



Classes
-------

.. autoapisummary::

   neural_lam.models.latent.base_encoder.BaseLatentEncoder


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

.. py:class:: BaseLatentEncoder(latent_dim: int, output_dist: str = 'isotropic')

   Bases: :py:obj:`torch.nn.Module`


   Abstract encoder mapping an input grid representation to a Gaussian
   distribution over a latent variable defined on mesh nodes.

   Subclasses implement :meth:`compute_dist_params`, which returns the raw
   parameters used to build the output distribution. With
   ``output_dist="isotropic"`` only the mean is produced (unit variance);
   with ``output_dist="diagonal"`` both mean and a positive std are output.

   Set up output dimensionality for the chosen distribution type.

   :param latent_dim: Dimensionality of the latent variable at each mesh node.
   :type latent_dim: int
   :param output_dist: Type of output distribution: ``"isotropic"`` (mean only, unit
                       variance) or ``"diagonal"`` (mean and per-dimension std).
   :type output_dist: str


   .. py:method:: compute_dist_params(grid_rep: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor
      :abstractmethod:


      Compute raw distribution parameters from the grid representation.

      :param grid_rep: Shape ``(B, num_grid_nodes, d_h)``. Grid input representation.
      :type grid_rep: torch.Tensor
      :param \*\*kwargs: Additional inputs used by concrete encoders (e.g. graph
                         embeddings).

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



   .. py:method:: forward(grid_rep: torch.Tensor, **kwargs: Any) -> torch.distributions.Normal

      Compute the Gaussian distribution over the latent variable.

      :param grid_rep: Shape ``(B, num_grid_nodes, d_h)``. Grid input representation.
      :type grid_rep: torch.Tensor
      :param \*\*kwargs: Additional inputs forwarded to :meth:`compute_dist_params`.

      :returns: Distribution over the latent variable, with batch shape
                ``(B, num_mesh_nodes, latent_dim)``.
      :rtype: torch.distributions.Normal



   .. py:attribute:: output_dist
      :value: 'isotropic'