neural_lam.models.step_predictors.graph.hierarchical
====================================================

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

.. autoapi-nested-parse::

   Base implementations for hierarchical (multi-level) graph models.



Classes
-------

.. autoapisummary::

   neural_lam.models.step_predictors.graph.hierarchical.BaseHiGraphModel


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

.. py:class:: BaseHiGraphModel(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', mesh_up_gnn_type: str = 'InteractionNet', mesh_down_gnn_type: str = 'InteractionNet')

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


   Base class for hierarchical graph models.

   Extend :class:`BaseGraphModel` with hierarchical mesh structures.


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

      Embed static mesh node features for the bottom level only;
      remaining levels are embedded at the start of ``process_step``.

      :returns: Shape ``(num_mesh_nodes[0], hidden_dim)``. Embedded bottom-level
                mesh node representations. Dims: ``num_mesh_nodes[0]`` is
                the number of nodes at level 0 and ``hidden_dim`` is the hidden
                dimension.
      :rtype: torch.Tensor



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

      Compute mesh node counts used for encoding and decoding.

      :returns: Total number of mesh nodes and the number to ignore during
                encoding/decoding.
      :rtype: tuple[int, int]



   .. py:method:: hi_processor_step(mesh_rep_levels: list[torch.Tensor], mesh_same_rep: list[torch.Tensor], mesh_up_rep: list[torch.Tensor], mesh_down_rep: list[torch.Tensor]) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]
      :abstractmethod:


      Internal processor step between mesh init and read-out.

      :param mesh_rep_levels: One tensor per level, each of shape
                              ``(B, num_mesh_nodes[l], hidden_dim)``. Node representations at
                              each hierarchy level. Dims: ``B`` is batch size,
                              ``num_mesh_nodes[l]`` is the node count at level ``l``, and
                              ``hidden_dim`` is the hidden dimension.
      :type mesh_rep_levels: list of torch.Tensor
      :param mesh_same_rep: One tensor per level, each of shape
                            ``(B, num_edges[l], hidden_dim)``. Same-level edge
                            representations. ``num_edges[l]`` is the edge count at
                            level ``l``.
      :type mesh_same_rep: list of torch.Tensor
      :param mesh_up_rep: One tensor per inter-level gap, each of shape
                          ``(B, num_edges[l], hidden_dim)``. Upward edge representations from
                          level ``l`` to ``l+1``.
      :type mesh_up_rep: list of torch.Tensor
      :param mesh_down_rep: One tensor per inter-level gap, each of shape
                            ``(B, num_edges[l], hidden_dim)``. Downward edge
                            representations from level ``l+1`` to ``l``.
      :type mesh_down_rep: list of torch.Tensor

      :returns: Updated ``(mesh_rep_levels, mesh_same_rep, mesh_up_rep,
                mesh_down_rep)`` in the same order as the inputs.
      :rtype: tuple of (list, list, list, list)



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

      Process the mesh representation across all hierarchy levels,
      implementing the full init-process-readout cycle.

      :param mesh_rep: Shape ``(B, num_mesh_nodes[0], hidden_dim)``. Bottom-level mesh
                       node representations from the encoder. Dims: ``B`` is batch
                       size, ``num_mesh_nodes[0]`` is the number of nodes at
                       level 0, and ``hidden_dim`` is the hidden dimension.
      :type mesh_rep: torch.Tensor

      :returns: Shape ``(B, num_mesh_nodes[0], hidden_dim)``. Updated bottom-level
                mesh node representations. Dims: same as ``mesh_rep``.
      :rtype: torch.Tensor



   .. py:attribute:: level_mesh_sizes


   .. py:attribute:: m2m_edge_index
      :type:  list[torch.Tensor]


   .. py:attribute:: m2m_features
      :type:  list[torch.Tensor]


   .. py:attribute:: mesh_down_edge_index
      :type:  list[torch.Tensor]


   .. py:attribute:: mesh_down_embedders


   .. py:attribute:: mesh_down_features
      :type:  list[torch.Tensor]


   .. py:attribute:: mesh_down_gnn_type
      :value: 'InteractionNet'



   .. py:attribute:: mesh_embedders


   .. py:attribute:: mesh_init_gnns


   .. py:attribute:: mesh_read_gnns


   .. py:attribute:: mesh_same_embedders


   .. py:attribute:: mesh_static_features
      :type:  list[torch.Tensor]


   .. py:attribute:: mesh_up_edge_index
      :type:  list[torch.Tensor]


   .. py:attribute:: mesh_up_embedders


   .. py:attribute:: mesh_up_features
      :type:  list[torch.Tensor]


   .. py:attribute:: mesh_up_gnn_type
      :value: 'InteractionNet'



   .. py:attribute:: num_levels