neural_lam.models.step_predictors.graph.graph_lam
=================================================

.. py:module:: neural_lam.models.step_predictors.graph.graph_lam

.. autoapi-nested-parse::

   Graph-based LAM model with a flat mesh.



Classes
-------

.. autoapisummary::

   neural_lam.models.step_predictors.graph.graph_lam.GraphLAM


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

.. py:class:: GraphLAM(datastore: neural_lam.datastore.BaseDatastore, graph_name: str = 'multiscale', hidden_dim: int = 64, hidden_layers: int = 1, processor_layers: int = 4, mesh_aggr: str = 'sum', num_past_forcing_steps: int = 1, num_future_forcing_steps: int = 1, output_std: bool = False, output_clamping_lower: dict[str, float] | None = None, output_clamping_upper: dict[str, float] | None = None, g2m_gnn_type: str = 'InteractionNet', m2g_gnn_type: str = 'InteractionNet', **_kwargs: object)

   Bases: :py:obj:`neural_lam.models.step_predictors.graph.base.BaseGraphModel`


   Full graph-based LAM model that can be used with different
   (non-hierarchical )graphs. Mainly based on GraphCast, but the model from
   Keisler (2022) is almost identical. Used for GC-LAM and L1-LAM in
   Oskarsson et al. (2023).

   Initialize the GraphLAM model.

   :param datastore: The datastore providing grid metadata and data access.
   :type datastore: BaseDatastore
   :param graph_name: The name of the graph to load.
   :type graph_name: str, default "multiscale"
   :param hidden_dim: The dimension of the hidden representations.
   :type hidden_dim: int, default 64
   :param hidden_layers: The number of hidden layers in the MLPs.
   :type hidden_layers: int, default 1
   :param processor_layers: The number of processor layers in the GNN.
   :type processor_layers: int, default 4
   :param mesh_aggr: The aggregation method for mesh nodes.
   :type mesh_aggr: str, default "sum"
   :param num_past_forcing_steps: The number of past forcing steps to include.
   :type num_past_forcing_steps: int, default 1
   :param num_future_forcing_steps: The number of future forcing steps to include.
   :type num_future_forcing_steps: int, default 1
   :param output_std: Whether to output a predicted standard deviation.
   :type output_std: bool, default False
   :param output_clamping_lower: Lower clamping limits for state variables.
   :type output_clamping_lower: dict, optional
   :param output_clamping_upper: Upper clamping limits for state variables.
   :type output_clamping_upper: dict, optional


   .. py:method:: embedd_mesh_nodes() -> torch.Tensor

      Embed static mesh node features.

      :returns: Shape ``(num_mesh_nodes, hidden_dim)``. Embedded mesh node
                representations. Dims: ``num_mesh_nodes`` is the number of
                mesh nodes and ``hidden_dim`` is the hidden dimension.
      :rtype: torch.Tensor



   .. py:method:: get_num_mesh() -> tuple[int, int]

      Compute number of mesh nodes from loaded features,
      and number of mesh nodes that should be ignored in encoding/decoding.

      :returns: * **num_mesh_nodes** (*int*) -- The number of mesh nodes.
                * **num_ignore_mesh_nodes** (*int*) -- The number of mesh nodes to ignore.



   .. py:method:: process_step(mesh_rep: torch.Tensor) -> torch.Tensor

      Process the mesh representation through the flat message-passing
      processor (all nodes at a single level).

      :param mesh_rep: Shape ``(B, num_mesh_nodes, hidden_dim)``. Current mesh node
                       representations. Dims: ``B`` is batch size, ``num_mesh_nodes`` is
                       the number of mesh nodes, and ``hidden_dim`` is the hidden
                       dimension.
      :type mesh_rep: torch.Tensor

      :returns: Shape ``(B, num_mesh_nodes, hidden_dim)``. Updated mesh node
                representations. Dims: same as ``mesh_rep``.
      :rtype: torch.Tensor



   .. py:attribute:: m2m_edge_index
      :type:  torch.Tensor


   .. py:attribute:: m2m_embedder


   .. py:attribute:: m2m_features
      :type:  torch.Tensor


   .. py:attribute:: mesh_embedder


   .. py:attribute:: mesh_static_features
      :type:  torch.Tensor


   .. py:attribute:: processor