neural_lam.utils.networks
=========================

.. py:module:: neural_lam.utils.networks

.. autoapi-nested-parse::

   Constructors for neural-network building blocks.



Functions
---------

.. autoapisummary::

   neural_lam.utils.networks.make_gnn_seq
   neural_lam.utils.networks.make_mlp


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

.. py:function:: make_gnn_seq(edge_index: torch.Tensor, num_gnn_layers: int, hidden_layers: int, hidden_dim: int, gnn_type: str = 'InteractionNet') -> torch.nn.Module

   Build a sequential stack of GNN layers that propagates both node and
   edge representations.

   All layer types share the ``(send, rec, edge) -> (rec, edge)``
   interface, so the stack can be applied as a single module.

   :param edge_index: Shape ``(2, M)``. Edge index of the edges that the GNN layers
                      operate on.
   :type edge_index: torch.Tensor
   :param num_gnn_layers: Number of stacked GNN layers; must be at least 1. Callers that
                          want a no-op stage (e.g. zero intra-level layers) should skip
                          building and applying the stack rather than calling this with 0.
   :type num_gnn_layers: int
   :param hidden_layers: Number of hidden layers in the MLPs of each GNN layer.
   :type hidden_layers: int
   :param hidden_dim: Dimensionality of node and edge representations.
   :type hidden_dim: int
   :param gnn_type: GNN layer type, any key in ``gnn_layers.GNN_TYPES``.
   :type gnn_type: str

   :returns: Sequential module mapping ``(mesh_rep, edge_rep)`` to updated
             ``(mesh_rep, edge_rep)``.
   :rtype: pyg.nn.Sequential

   :raises ValueError: If ``num_gnn_layers`` is less than 1.


.. py:function:: make_mlp(blueprint: list[int], layer_norm: bool = True) -> torch.nn.Sequential

   Construct a multilayer perceptron from a blueprint of layer widths.

   :param blueprint: Sequence of layer dimensions where ``blueprint[0]`` is the input size,
                     ``blueprint[-1]`` is the output size, the intermediate entries specify
                     the hidden layer widths, and ``len(blueprint) - 2`` is the number of
                     hidden layers.
   :type blueprint: list[int]
   :param layer_norm: If ``True``, append a ``LayerNorm`` to the output as in GraphCast.
   :type layer_norm: bool, optional

   :returns: Sequential module implementing the specified MLP.
   :rtype: torch.nn.Sequential