Bug: Offset Mismatch in zero_index_m2g / zero_index_g2m when last grid node is unconnected

[GiGiKoneti]

Bug Description

When mesh_first is False (grid nodes are indexed before mesh nodes), zero_index_m2g and zero_index_g2m in neural_lam/utils.py dynamically calculate the offset for mesh nodes using the maximum connected grid node index:

• num_interior_nodes = m2g_edge_index[1].max() + 1 in zero_index_m2g
• num_grid_nodes = g2m_edge_index[0].max() + 1 in zero_index_g2m

This assumes the highest-indexed grid node is always connected to the mesh. If the highest-indexed grid nodes are unconnected (e.g., due to boundary masking or spatial selection criteria), the computed offset is smaller than the true grid size. This corrupts the zero-indexing of all subsequent mesh nodes, causing index alignment errors during encoding and decoding.

---

Minimal Reproducible Example

The following code reproduces the bug by creating a mock edge index where the highest active grid node connected in m2g is 90 (while the true grid size is 100):

import torch
from neural_lam.utils import zero_index_m2g

# True configuration: 100 grid nodes (0..99), 10 mesh nodes (100..109)
true_num_grid_nodes = 100

# Edge index where the highest active grid node connected in m2g is 90
m2g_edge_index = torch.tensor([
    [100, 101, 102, 109],  # mesh node indices
    [  0,  10,  50,  90]   # grid node indices
])

# Computes offset as 90 + 1 = 91 instead of 100
zeroed_m2g = zero_index_m2g(m2g_edge_index, mesh_static_features=[], mesh_first=False)

print(f"Original mesh indices: {m2g_edge_index[0].tolist()}")
print(f"Expected zero-indexed: {[x - true_num_grid_nodes for x in m2g_edge_index[0].tolist()]}")
print(f"Actual zero-indexed:   {zeroed_m2g[0].tolist()}")
# Actual zero-indexed:   [9, 10, 11, 18] -> Does not start at 0

---

Impact

This causes subsequent mesh node index accesses to point to the wrong locations in the mesh representation array (or crash with an IndexError if the corrupted index exceeds the array bounds), leading to silent spatial coordinate misalignment or runtime crashes.

---

Proposed Fix

Pass the true grid node count (e.g. num_grid_nodes from the datastore) directly into load_graph and use it to offset the indices rather than inferring it from the edge indices.

[joeloskarsson]

I think all grid nodes must be connected in g2m. Could you connect this to the newly introduced graph specification, and see if there is an issue with the spec (and if so, for legacy, or for newer versions), or this is an issue with the code not matching the spec?

[GiGiKoneti]

According to Section 3.1.1 (Node index space) of the newly introduced graph-on-disk spec, unconnected nodes are explicitly allowed:

"Although the integers used to label a given set of source or destination nodes must be contiguous from 0 to N-1... it is not a requirement that all node indexes are actually present in the edge index tensors, i.e. it is not a requirement that all nodes in either the source or destination set of nodes has edges connecting to them (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)"

how it connects to the codebase:

1. New-Format Graphs (Spec 0.1.0): These are zero-indexed on-save by create_graph.py and loaded without calling zero_index_m2g or zero_index_g2m. They are unaffected by this bug.
2. Legacy-Format Graphs: These still run through zero_index_m2g and zero_index_g2m on load. Using max() + 1 dynamically is a flawed assumption when legacy graphs have unconnected boundary nodes.

This is an issue of the legacy-format code not matching the specification. Since legacy graphs do not store metadata indicating the true grid size, we cannot dynamically infer the offset. We could resolve this by either dropping legacy graph support entirely or passing num_grid_nodes as an optional argument to load_graph.