Fix legacy graph indexing offset mismatch when boundary grid nodes are unconnected

[GiGiKoneti]

Describe your changes

This PR fixes the offset mismatch in zero_index_m2g and zero_index_g2m when mesh_first is False.
Specifically, instead of dynamically computing the number of grid nodes from the maximum active node index in the edge index (which is incorrect if some boundary grid nodes are unconnected to the GNN mesh), it computes the total number of nodes at the bottom level and subtracts the number of mesh nodes at that level:

1. max_idx is calculated over g2m, m2g, and the bottom level m2m edge indices.
2. num_grid_nodes is calculated as (max_idx + 1) - num_mesh_nodes.
3. Passes the correctly computed num_grid_nodes down to the zero-indexing helpers.

Issue Link

closes #699

Type of change

• 🐛 Bug fix (non-breaking change that fixes an issue)
• ✨ New feature (non-breaking change that adds functionality)
• 💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)
• 📖 Documentation (Addition or improvements to documentation)

Checklist before requesting a review

• My branch is up-to-date with the target branch - if not update your fork with the changes from the target branch (use pull with --rebase option if possible).
• I have performed a self-review of my code
• For any new/modified functions/classes I have added docstrings that clearly describe its purpose, expected inputs and returned values
• I have placed in-line comments to clarify the intent of any hard-to-understand passages of my code
• I have updated the README to cover introduced code changes (no README changes needed for this bug fix)
• I have added tests that prove my fix is effective or that my feature works
• I have given the PR a name that clearly describes the change, written in imperative form.
• I have requested a reviewer and an assignee (assignee is responsible for merging). This applies only if you have write access to the repo, otherwise feel free to tag a maintainer to add a reviewer and assignee.

Checklist for reviewers

• the code is readable
• the code is well tested
• the code is documented (including return types and parameters)
• the code is easy to maintain

Author checklist after completed review

• I have added a line to the CHANGELOG describing this change, in a section reflecting type of change:
  • fixes: - Fix offset mismatch in legacy graph zero-indexing... under ## [unreleased]

[Sir-Sloth-The-Lazy]

Just a small problem here, also great eye for finding this issue ! Just a ping no tests are added till now however the descirption suggests there should be one. I would suggest the exact scenario from #699 (legacy format, grid-first layout, unconnected boundary grid node) has zero coverage before or after this PR you can add a regression test here with a synthetic legacy graph that has a diconneted top-indexed grid node, mirroring the reproduction in the issue.

Also the new parameter inserted before the existing one changes the public call signature. Both zero_index_m2g and zero_index_g2m are exported from neural_lam/utils/__init__.py (public API). The PR adds num_grid_nodes: int | None = None before the existing restore: bool = False parameter. All current in-repo callers use keywords, so nothing breaks internally, but any external caller invoking these positionally (e.g. zero_index_m2g(idx, features, mesh_first, True)) would now silently bind True to num_grid_nodes instead of restore, restore would default to False while num_grid_nodes becomes truthy, corrupting the result without raising. Safer to append new optional params after restore, or make them keyword-only.

[Sir-Sloth-The-Lazy]

m2m_edge_index gets reassigned to the locally zero-indexed version (each level independently rebased to start at 0 via zero_index_edge_index) several lines before the new max_idx block runs. So by the time it's used, it no longer lives in the same raw/combined grid+mesh index space as the still-raw g2m_edge_index/m2g_edge_index, mixing them in one max() is comparing apples to oranges. In practice this doesn't blow up on typical graphs, because grid counts vastly exceed mesh counts, so the raw g2m/m2g values dominate the max and the (small, rebased) m2m term is inert. But it means the m2m term contributes nothing to the robustness the fix is meant to provide, if the mesh's topmost-indexed node were ever the one lacking a g2m/m2g edge (the same failure mode as #699, just mirrored onto the mesh side), this fix wouldn't catch it either, silently.

There's also a smaller mismatch: the comment says "we determine total number of nodes at the bottom level," but the code does max(ei.max().item() for ei in m2m_edge_index) over the entire list (all mesh levels), not just m2m_edge_index[0]. For non-hierarchical graphs this is moot, but it's inconsistent with the stated intent and themesh_static_features[0] reference right below it.

Suggested fix to flag: compute max_idx from the raw m2m_edge_index (before the zero_index_edge_index reassignment) or move the whole num_grid_nodes computation up before that reassignment and use m2m_edge_index[0] explicitly to match the "bottom level" comment.

[GiGiKoneti]

Thanks for the thorough review and catch @Sir-Sloth-The-Lazy!

1. API Signature: Updated zero_index_m2g and zero_index_g2m to place restore: bool = False before num_grid_nodes: int | None = None, preserving positional parameter order for external callers.
2. Regression Test: Added test_legacy_graph_unconnected_boundary_grid_node_zero_indexing in tests/test_graph_creation.py to cover legacy graphs (mesh_first=False) with top-indexed unconnected boundary grid nodes.

All 250 tests and pre-commit checks are passing cleanly!