Support hierarchical graphs in the HeteroData representation - #713
Open
prajwal-tech07 wants to merge 11 commits into
Open
Support hierarchical graphs in the HeteroData representation#713prajwal-tech07 wants to merge 11 commits into
prajwal-tech07 wants to merge 11 commits into
Conversation
First step of the HeteroData migration (issue mllam#385), scoped to flat (non-hierarchical) graphs. - Add neural_lam.utils.heterodata with graph_dict_to_heterodata and its inverse heterodata_to_graph_dict. Grid nodes are the "data" node type and mesh nodes the "hidden" node type, following the convention used in the reference implementations linked from the issue (leifdenby/weatherduck, matschreiner/equicast); g2m/m2m/m2g become typed edges. The grid node count is taken from the datastore rather than inferred from edge indices, and the edge-feature width is read dynamically (3 for 2D, 4 for 3D). - Wire into BaseGraphModel behind a use_heterodata flag (default False): when enabled the graph is represented as a HeteroData object and the model's graph tensors are taken from it. The extracted tensors are identical to the dict ones, so the model builds and trains identically either way, and existing checkpoints remain compatible. - Thread the flag through GraphLAM. - Tests: unit tests for the conversion (structure, exact round-trip, variable edge-feature width, datastore-provided grid count, hierarchical rejection) and a GraphLAM equivalence test asserting identical parameters, graph buffers and forward output with the flag on and off.
Rename the HeteroData node types from data/hidden to grid/mesh, matching the naming used throughout the rest of neural-lam (grid_static_features, mesh_static_features, g2m/m2m/m2g). The names remain module-level constants; the reference implementations in issue mllam#385 use data/hidden, noted in the module docstring.
Extend the GraphLAM equivalence test to run a few real optimizer steps on both the dict-based and HeteroData-based models and assert identical per-step losses and identical weights afterwards, covering the issue mllam#385 requirement to carry out training with both datastructures and confirm it proceeds identically.
Rename heterodata_to_graph_dict to graph_tensors_from_heterodata and document it as how the model obtains its graph tensors when the graph is represented as a HeteroData: every tensor is looked up on the typed node/edge stores of the object, so the HeteroData is the source of the graph data the model uses. The tensors are unchanged, so the model still builds and trains identically. Also add the changelog entry for this change.
Extend the HeteroData migration (mllam#385) to hierarchical (multi-level) graphs and the HiLAM / HiLAMParallel models. - graph_dict_to_heterodata / heterodata_to_graph_dict now handle hierarchical graphs: each mesh level is a distinct node type (mesh_0 .. mesh_{L-1}), g2m/m2g connect the grid to the bottom level, intra-level edges are (mesh_i, to, mesh_i) and inter-level edges are (mesh_i, up, mesh_{i+1}) and (mesh_{i+1}, down, mesh_i) for each of the L-1 level pairs. The round-trip through HeteroData is exact, so the model gets identical tensors. - Thread use_heterodata through BaseHiGraphModel, HiLAM and HiLAMParallel; BaseGraphModel now builds/consumes the HeteroData for hierarchical graphs too (the previous NotImplementedError guard is removed). - Tests: hierarchical builder structure + exact round-trip unit tests, and a HiLAM equivalence test asserting identical parameters, per-level graph buffers, forward output and training with and without the HeteroData datastructure.
prajwal-tech07
force-pushed
the
issue-385/heterodata-hierarchical
branch
from
July 25, 2026 13:58
49d1f6e to
5908777
Compare
Main moved the graph loading and buffer registration out of BaseGraphModel into utils.load_and_register_graph (mllam#648), so the HeteroData wiring moves there with it. It is now applied wherever a module loads a graph through that helper, which includes the new Graph-EFM step predictors, and the grid node count is read from the datastore directly.
The HeteroData wiring now lives in utils.load_and_register_graph, so pass the hierarchical flag through to the conversion helpers there instead of building the object in BaseGraphModel.
The check moved out of BaseGraphModel together with the rest of the graph loading, so test it where it now lives: loading a hierarchical graph with use_heterodata raises NotImplementedError.
This PR adds hierarchical support, so the rejection test from the flat branch becomes its positive counterpart: loading a hierarchical graph with use_heterodata builds a per-level HeteroData object whose levels match the tensors registered on the module.
Contributor
Author
|
@leifdenby @joeloskarsson review request for this one too - it stacks on #711, so that's the one to look at first. Same post-#648 merge, CI green, hierarchical re-verified live (3 levels, grid/mesh_0/mesh_1/mesh_2, 9 edge types). |
21 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Second step of the HeteroData migration (#385), extending it to hierarchical (multi-level) graphs and the
HiLAM/HiLAMParallelmodels.graph_dict_to_heterodata/heterodata_to_graph_dictnow handle hierarchical graphs. Each mesh level is a distinct node type (mesh_0…mesh_{L-1}), since levels have different node counts and their own embedders/GNNs in the model.g2m/m2gconnect the grid to the bottom level, intra-level edges are("mesh_i","to","mesh_i"), and inter-level edges are("mesh_i","up","mesh_{i+1}")and("mesh_{i+1}","down","mesh_i")for each of theL-1level pairs.use_heterodatathroughBaseHiGraphModel,HiLAMandHiLAMParallel;BaseGraphModelnow builds and consumes theHeteroDatafor hierarchical graphs as well (the previousNotImplementedErrorguard is removed).HeteroDatais exact, so the model receives identical tensors and builds/trains identically either way, and existing checkpoints remain compatible.Verified on a real 3-level graph (729/81/9 mesh nodes): the expected per-level node types and all nine edge types are present, with
m2mhavingLentries andmesh_up/mesh_downhavingL-1.Issue Link
Part of #385 (hierarchical follow-up to #711).
Type of change
Checklist before requesting a review
pullwith--rebaseoption if possible).Checklist for reviewers
Each PR comes with its own improvements and flaws. The reviewer should check the following:
Author checklist after completed review
reflecting type of change (add section where missing):
Checklist for assignee