Skip to content

Load flat graphs into pyg.HeteroData and wire into BaseGraphModel#711

Open
prajwal-tech07 wants to merge 3 commits into
mllam:mainfrom
prajwal-tech07:issue-385/heterodata-flat
Open

Load flat graphs into pyg.HeteroData and wire into BaseGraphModel#711
prajwal-tech07 wants to merge 3 commits into
mllam:mainfrom
prajwal-tech07:issue-385/heterodata-flat

Conversation

@prajwal-tech07

@prajwal-tech07 prajwal-tech07 commented Jul 22, 2026

Copy link
Copy Markdown
Contributor

First step of the HeteroData migration (#385), scoped to flat (non-hierarchical) graphs.

  • Adds neural_lam.utils.heterodata with graph_dict_to_heterodata and its inverse heterodata_to_graph_dict. Grid nodes are the grid node type and mesh nodes the mesh node type, matching the terminology used throughout neural-lam (grid_static_features, mesh_static_features, g2m/m2m/m2g); 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). The node/edge-type names are module-level constants; the reference implementations linked from the issue (leifdenby/weatherduck, matschreiner/equicast) use data/hidden — happy to switch if preferred.
  • Wires 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.
  • Threads the flag through GraphLAM.

Hierarchical (HiLAM) support is deliberately left as a follow-up — there is an open question on #385 about how to represent mesh levels as node/edge types.

Issue Link

Part of #385 (first implementation step; the design discussion and hierarchical follow-up remain open).

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
  • 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 (context).
  • 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

Each PR comes with its own improvements and flaws. The reviewer should check the following:

  • 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 (add section where missing):
    • added: when you have added new functionality
    • changed: when default behaviour of the code has been changed
    • fixes: when your contribution fixes a bug
    • maintenance: when your contribution is relates to repo maintenance, e.g. CI/CD or documentation

Checklist for assignee

  • PR is up to date with the base branch
  • the tests pass
  • (if the PR is not just maintenance/bugfix) the PR is assigned to the next milestone. If it is not, propose it for a future milestone.
  • author has added an entry to the changelog (and designated the change as added, changed, fixed or maintenance)
  • Once the PR is ready to be merged, squash commits and merge the PR.

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant