[Bug] Silent node feature corruption and RuntimeError in InteractionNet when trailing receiver nodes have no incoming edges

[GiGiKoneti]

Describe the Bug

InteractionNet (and PropagationNet) computes the receiver node count dynamically in __init__ as:

# neural_lam/gnn_layers.py:73
self.num_rec = edge_index[1].max() + 1
edge_index = torch.stack(
    (edge_index[0] + self.num_rec, edge_index[1]), dim=0
)

This implicitly assumes that the highest-indexed receiver node ({\text{rec}} - 1$) always has at least one incoming edge in edge_index.

If the highest-indexed receiver node (or a trailing subset of receiver nodes) has zero incoming edges (which occurs in regional LAM boundary cutoffs, unmapped boundary grid points in \rightarrow G$, irregular meshes, or hierarchical subgraphs), self.num_rec evaluates to self.num_rec < N_{\text{rec}}.

---

Impact & Failure Modes

1. Silent Feature Corruption (Worst Case):

   • In InteractionNet.forward(), node representations are concatenated as node_reps = torch.cat((rec_rep, send_rep), dim=-2).
   • Receiver nodes occupy indices [0 .. N_rec - 1], and sender nodes occupy indices [N_rec .. N_rec + N_snd - 1].
   • Because sender indices in self.edge_index[0] are offset by self.num_rec instead of N_rec, sender node 0 is looked up at index self.num_rec (which is inside rec_rep!).
   • All sender features are shifted by {\text{rec}} - \text{self.num_rec}$, causing the GNN to compute messages using the wrong physical node representations.

2. Hard Crash:

   • In aggregate(), PyG aggregates up to self.num_rec.
   • In rec_diff = self.aggr_mlp(torch.cat((rec_rep, edge_rep_aggr), dim=-1)), tensor dimensions mismatch between rec_rep ({\text{rec}}$) and edge_rep_aggr (self.num_rec), raising:

     RuntimeError: Sizes of tensors must match except in dimension -1. Expected size N_rec but got size num_rec.
     

---

Minimal Reproducible Example

import torch
from neural_lam.gnn_layers import InteractionNet

num_send = 100
num_rec = 50
hidden_dim = 16

# Receiver node 49 has no incoming edge (max receiver in edge_index is 48)
edge_index = torch.tensor([
    [0, 1, 2],    # senders
    [10, 20, 48]  # receivers
], dtype=torch.long)

gnn = InteractionNet(edge_index, input_dim=hidden_dim)

rec_rep = torch.randn((num_rec, hidden_dim))
send_rep = torch.randn((num_send, hidden_dim))
edge_rep = torch.randn((3, hidden_dim))

# Trigger forward pass
gnn(send_rep, rec_rep, edge_rep)

Output:

RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 50 but got size 49 for tensor number 1 in the list.

---

Proposed Fix

1. Allow InteractionNet and PropagationNet to accept an explicit num_rec: int | None = None in __init__:

if num_rec is not None:
    self.num_rec = int(num_rec)
elif edge_index.numel() > 0:
    self.num_rec = int(edge_index[1].max().item() + 1)
else:
    self.num_rec = 0

2. Pass num_rec explicitly from caller models (BaseGraphModel, BaseHiGraphModel, GraphLAM, GraphEFM) based on known node set sizes (num_mesh_nodes, num_grid_nodes, level_mesh_sizes).

[joeloskarsson]

This implicitly assumes that the highest-indexed receiver node ({\text{rec}} - 1$) always has at least one incoming edge in edge_index.

I guess this describes a mismatch between the implementation and graph specification (https://github.com/mllam/neural-lam/blob/main/docs/graph_storage_spec.md). I think for things like this we always need to connect it to the graph spec now.

[GiGiKoneti]

Exactly right @joeloskarsson. Section 3.1.1 of graph_storage_spec.md explicitly notes:

"it is not a requirement that all node indexes are actually present in the edge index tensors... (this is for example used when decoding to only a subset in the grid nodes when training in a limited-area setting where predictions are not required on the boundary)"

Because InteractionNet derives num_rec purely from edge_index[1].max() + 1, any valid spec-compliant graph with unmapped trailing boundary nodes violates InteractionNet's internal offset calculation.

Allowing InteractionNet and PropagationNet to accept num_rec: int | None = None (passed explicitly from known node set dimensions like num_grid_nodes and level_mesh_sizes) aligns the implementation directly with Section 3.1.1.

I'd be happy to open a PR to fix this and add a test case for unconnected trailing nodes! wt do you say ?

[zakirkg]

One thought about keeping edge_index[1].max() + 1 as fallback. since we are inferring and we cannot know for sure at initialization, would it be worth warning when num_rec isn't provided?.

This would make the assumption visible to external callers while preserving backward compatibility through num_rec: int | None = None.