neural_lam.interaction_net

neural_lam.interaction_net#

Interaction Network layers and helper modules used by Neural-LAM.

Module Contents#

class neural_lam.interaction_net.InteractionNet(edge_index, input_dim, update_edges=True, hidden_layers=1, hidden_dim=None, edge_chunk_sizes=None, aggr_chunk_sizes=None, aggr='sum')#

Bases: torch_geometric.nn.MessagePassing

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

Initialise an InteractionNet message-passing layer.

Parameters:

  • edge_index (torch.Tensor) –

    Edge connectivity tensor in PyG format.

    • Shape: (2, M) where M is the number of 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:

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

aggregate(inputs, index, ptr, dim_size)#

Aggregate messages while also returning the per-edge values.

forward(send_rep, rec_rep, edge_rep)#

Update receiver (and optionally edge) representations via message passing.

Parameters:

  • send_rep (torch.Tensor) –

    Vector representations of sender nodes.

    • Shape: (N_send, d_h)

  • rec_rep (torch.Tensor) –

    Vector representations of receiver nodes.

    • Shape: (N_rec, d_h)

  • edge_rep (torch.Tensor) –

    Edge representations used during message passing.

    • Shape: (M, d_h)

Returns:

Updated receiver representations. If self.update_edges is True, the tuple (rec_rep, edge_rep) containing the updated receiver and edge representations is returned.

  • Shape: (N_rec, d_h) for receivers and (M, d_h) for edges.

Return type:

torch.Tensor or tuple[torch.Tensor, torch.Tensor]

message(x_j, x_i, edge_attr)#

Compute messages from node j to i using edge features.

num_rec#
update_edges = True#
class neural_lam.interaction_net.SplitMLPs(mlps, chunk_sizes)#

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 (Iterable[nn.Module]) – Sequence of MLPs to apply to each chunk.

  • chunk_sizes (Sequence[int]) – Sizes used when splitting the input along dim=-2.

Raises:

AssertionError – If the number of mlps and chunk_sizes differ.

forward(x)#

Chunk up input tensor and feed each slice through its MLP.

Parameters:

x (torch.Tensor) –

Input tensor to split and process.

  • Shape: (..., N, d) where N = sum(chunk_sizes).

Returns:

Concatenated MLP outputs assembled along the chunk dimension.

  • Shape: (..., N, d)

Return type:

torch.Tensor

chunk_sizes#
mlps#
Contents