neural_lam.gnn_layers#

Interaction Network and PropagationNet GNN layers used by Neural-LAM.

Attributes#

GNN_TYPES

Classes#

`InteractionNet`	Implementation of a generic Interaction Network,
`PropagationNet`	Alternative version of InteractionNet that incentivizes the propagation
`SplitMLPs`	Module that feeds chunks of input through different MLPs.

Functions#

get_gnn_class(→ type[torch_geometric.nn.MessagePassing])

Look up a GNN class by name.

Module Contents#

class neural_lam.gnn_layers.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')[source]#

Bases: torch_geometric.nn.MessagePassing

Implementation of a generic Interaction Network, from Battaglia et al. (2016)

Create a new InteractionNet.

Parameters:

edge_index (torch.Tensor) – Edge connectivity tensor in PyG format. Shape (2, num_edges).
input_dim (int) – Dimensionality of both node and edge input representations.
update_edges (bool, optional) – If True, compute and return updated edge representations in addition to node representations. Default is True.
hidden_layers (int, optional) – Number of hidden layers in each MLP. Default is 1.
hidden_dim (int or None, optional) – Width of hidden layers. If None, defaults to input_dim.
edge_chunk_sizes (list[int] or None, optional) – Chunk sizes for splitting edge representations across separate MLPs. None uses a single shared MLP.
aggr_chunk_sizes (list[int] or None, optional) – Chunk sizes for splitting aggregated node representations across separate MLPs. None uses a single shared MLP.
aggr ({"sum", "mean"}, optional) – Message aggregation method. Default is "sum".

Raises:

ValueError – If aggr is not one of "sum" or "mean".

aggregate(inputs: torch.Tensor, index: torch.Tensor, ptr: torch.Tensor | None, dim_size: int | None) → tuple[torch.Tensor, torch.Tensor][source]#

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.

forward(send_rep: torch.Tensor, rec_rep: torch.Tensor, edge_rep: torch.Tensor) → torch.Tensor | tuple[torch.Tensor, torch.Tensor][source]#

Apply the interaction network to update receiver node representations, and optionally edge representations.

Parameters:

send_rep (torch.Tensor) – Sender node representations. Shape (num_send, input_dim).
rec_rep (torch.Tensor) – Receiver node representations. Shape (num_rec, input_dim).
edge_rep (torch.Tensor) – Edge representations. Shape (num_edges, input_dim).

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.

message(x_j: torch.Tensor, x_i: torch.Tensor, edge_attr: torch.Tensor) → torch.Tensor[source]#: Compute messages from node j to i.

node_residual_target(rec_rep: torch.Tensor, edge_rep_aggr: torch.Tensor) → torch.Tensor[source]#: Return the base tensor for the node residual connection. InteractionNet uses the original receiver representation.

num_rec#

update_edges = True#

class neural_lam.gnn_layers.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')[source]#

Bases: InteractionNet

Alternative version of InteractionNet that incentivizes the propagation of information from sender nodes to receivers.

Initialise the PropagationNet layer.

Parameters share the meaning of InteractionNet.__init__; see that class for the full description. The propagation variant overrides aggr defaults internally to favour stability of the propagation residual.

message(x_j: torch.Tensor, x_i: torch.Tensor, edge_attr: torch.Tensor) → torch.Tensor[source]#: Compute messages from node j to node i.

node_residual_target(rec_rep: torch.Tensor, edge_rep_aggr: torch.Tensor) → torch.Tensor[source]#: Return the base tensor for the node residual connection. PropagationNet uses the aggregated edge messages, propagating sender information to receiver nodes.

class neural_lam.gnn_layers.SplitMLPs(mlps: list[torch.nn.Module], chunk_sizes: list[int])[source]#

Bases: 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.

Parameters:

mlps (list of nn.Module) – Sequence of MLPs to apply to each chunk.
chunk_sizes (list of int) – Sizes used when splitting the input along dim=-2.

Raises:

AssertionError – If the number of mlps and chunk_sizes differ.

forward(x: torch.Tensor) → torch.Tensor[source]#

Split input along dim -2 and feed each chunk through its MLP.

Parameters:: x (torch.Tensor) – Shape (..., N, d). Input tensor where N = sum(chunk_sizes) and d is the feature dimension.
Returns:: Shape (..., N, d). Concatenated outputs from all MLPs.
Return type:: torch.Tensor

chunk_sizes#

mlps#

neural_lam.gnn_layers.get_gnn_class(gnn_type: str) → type[torch_geometric.nn.MessagePassing][source]#

Look up a GNN class by name.

Parameters:: gnn_type (str) – One of the keys in GNN_TYPES (currently “InteractionNet” or “PropagationNet”)
Returns:: The corresponding GNN class.
Return type:: type[pyg.nn.MessagePassing]

neural_lam.gnn_layers.GNN_TYPES#