neural_lam.gnn_layers ===================== .. py:module:: neural_lam.gnn_layers .. autoapi-nested-parse:: Interaction Network and PropagationNet GNN layers used by Neural-LAM. Attributes ---------- .. autoapisummary:: neural_lam.gnn_layers.GNN_TYPES Classes ------- .. autoapisummary:: neural_lam.gnn_layers.InteractionNet neural_lam.gnn_layers.PropagationNet neural_lam.gnn_layers.SplitMLPs Functions --------- .. autoapisummary:: neural_lam.gnn_layers.get_gnn_class Module Contents --------------- .. py:class:: InteractionNet(edge_index: torch.Tensor, input_dim: int, update_edges: bool = True, hidden_layers: int = 1, hidden_dim: int | None = None, edge_chunk_sizes: list[int] | None = None, aggr_chunk_sizes: list[int] | None = None, aggr: str = 'sum') Bases: :py:obj:`torch_geometric.nn.MessagePassing` Implementation of a generic Interaction Network, from Battaglia et al. (2016) Create a new InteractionNet. :param edge_index: Edge connectivity tensor in PyG format. Shape ``(2, num_edges)``. :type edge_index: torch.Tensor :param input_dim: Dimensionality of both node and edge input representations. :type input_dim: int :param update_edges: If ``True``, compute and return updated edge representations in addition to node representations. Default is ``True``. :type update_edges: bool, optional :param hidden_layers: Number of hidden layers in each MLP. Default is ``1``. :type hidden_layers: int, optional :param hidden_dim: Width of hidden layers. If ``None``, defaults to ``input_dim``. :type hidden_dim: int or None, optional :param edge_chunk_sizes: Chunk sizes for splitting edge representations across separate MLPs. ``None`` uses a single shared MLP. :type edge_chunk_sizes: list[int] or None, optional :param aggr_chunk_sizes: Chunk sizes for splitting aggregated node representations across separate MLPs. ``None`` uses a single shared MLP. :type aggr_chunk_sizes: list[int] or None, optional :param aggr: Message aggregation method. Default is ``"sum"``. :type aggr: {"sum", "mean"}, optional :raises ValueError: If ``aggr`` is not one of ``"sum"`` or ``"mean"``. .. py:method:: aggregate(inputs: torch.Tensor, index: torch.Tensor, ptr: torch.Tensor | None, dim_size: int | None) -> tuple[torch.Tensor, torch.Tensor] Overridden aggregation function to: * return both aggregated and per-edge messages, * only aggregate to the number of receiver nodes (``self.num_rec``) rather than to ``dim_size``. .. py:method:: forward(send_rep: torch.Tensor, rec_rep: torch.Tensor, edge_rep: torch.Tensor) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] Apply the interaction network to update receiver node representations, and optionally edge representations. :param send_rep: Sender node representations. Shape ``(num_send, input_dim)``. :type send_rep: torch.Tensor :param rec_rep: Receiver node representations. Shape ``(num_rec, input_dim)``. :type rec_rep: torch.Tensor :param edge_rep: Edge representations. Shape ``(num_edges, input_dim)``. :type edge_rep: torch.Tensor :returns: * **rec_rep** (*torch.Tensor*) -- Updated receiver node representations. Shape ``(num_rec, input_dim)``. * **edge_rep** (*torch.Tensor*) -- Updated edge representations. Shape ``(num_edges, input_dim)``. Only returned when ``update_edges=True``. .. py:method:: message(x_j: torch.Tensor, x_i: torch.Tensor, edge_attr: torch.Tensor) -> torch.Tensor Compute messages from node ``j`` to ``i``. .. py:method:: node_residual_target(rec_rep: torch.Tensor, edge_rep_aggr: torch.Tensor) -> torch.Tensor Return the base tensor for the node residual connection. InteractionNet uses the original receiver representation. .. py:attribute:: num_rec .. py:attribute:: update_edges :value: True .. py:class:: PropagationNet(edge_index: torch.Tensor, input_dim: int, update_edges: bool = True, hidden_layers: int = 1, hidden_dim: int | None = None, edge_chunk_sizes: list[int] | None = None, aggr_chunk_sizes: list[int] | None = None, aggr: str = 'sum') Bases: :py:obj:`InteractionNet` Alternative version of InteractionNet that incentivizes the propagation of information from sender nodes to receivers. Initialise the :class:`PropagationNet` layer. Parameters share the meaning of :class:`InteractionNet.__init__`; see that class for the full description. The propagation variant overrides ``aggr`` defaults internally to favour stability of the propagation residual. .. py:method:: message(x_j: torch.Tensor, x_i: torch.Tensor, edge_attr: torch.Tensor) -> torch.Tensor Compute messages from node j to node i. .. py:method:: node_residual_target(rec_rep: torch.Tensor, edge_rep_aggr: torch.Tensor) -> torch.Tensor Return the base tensor for the node residual connection. PropagationNet uses the aggregated edge messages, propagating sender information to receiver nodes. .. py:class:: SplitMLPs(mlps: list[torch.nn.Module], chunk_sizes: list[int]) Bases: :py:obj:`torch.nn.Module` Module that feeds chunks of input through different MLPs. Split up input along dim -2 using given chunk sizes and feeds each chunk through separate MLPs. Create a module that dispatches chunks of the input to separate MLPs. :param mlps: Sequence of MLPs to apply to each chunk. :type mlps: list of nn.Module :param chunk_sizes: Sizes used when splitting the input along ``dim=-2``. :type chunk_sizes: list of int :raises AssertionError: If the number of ``mlps`` and ``chunk_sizes`` differ. .. py:method:: forward(x: torch.Tensor) -> torch.Tensor Split input along dim -2 and feed each chunk through its MLP. :param x: Shape ``(..., N, d)``. Input tensor where ``N = sum(chunk_sizes)`` and ``d`` is the feature dimension. :type x: torch.Tensor :returns: Shape ``(..., N, d)``. Concatenated outputs from all MLPs. :rtype: torch.Tensor .. py:attribute:: chunk_sizes .. py:attribute:: mlps .. py:function:: get_gnn_class(gnn_type: str) -> type[torch_geometric.nn.MessagePassing] Look up a GNN class by name. :param gnn_type: One of the keys in GNN_TYPES (currently "InteractionNet" or "PropagationNet") :type gnn_type: str :returns: The corresponding GNN class. :rtype: type[pyg.nn.MessagePassing] .. py:data:: GNN_TYPES