[Design discussion] Supporting `pyg.HeteroData` use internally in `neural-lam`

[leifdenby]

The aim of this issue to continue the discussion on whether and how to support use of pyg.HeteroData objects to represent the graph when working with neural-lam (rather than the dictionary of edge-indexes and node/edge features currently used)

To start of the discussion here are some things I have been wondering about:

1. My view is that we will eventually want to use pyg.HeteroData objects to represent the graph internally in neural-lam. This is what I did in https://github.com/leifdenby/weatherduck which was a prototype implementation trying that approach out (and also trying out fiddle-based config, which is a different conversation). It is also what @matschreiner does in https://github.com/matschreiner/equicast/tree/main?tab=readme-ov-file#data-handler. I think it simplifies the code a lot (only a single object is passed around representing the entire graph) and enables us to eventually also put the model inputs (state/forcing etc) as node features on this graph object, rather than as separate tensors being passed around (which I think is a nice design for implementing new architectures in neural-lam)
   • Is there an interest in moving towards this?
   • Have some people already started this work?
   • @SimonKamuk you and I have discussed this, have you looked at refactoring neural-lam in this way?
   • If we agree this should be done, maybe we could plan a first step in this direction to be loading the tensor-based format into a pyg.HeteroDataobject that is then passed to the torch.nn.Module-derived classes that implement the forecasting model (currently ARModel). What do people think of that?
2. Does it make sense trying to export from weather-model-graphs to pyg.HeteroData objects and pickling these to disk?
   • If so, what would the drawbacks of these be? To me a clear drawback is that the format (pickled python objects) isn't a self-describing format, unlike what a xr.DataTree netCDF/zarr dataset could be. And it relies on torch-geometric which seems a bit overkill for something that just describes a static graph with edge and node features.
   • If we instead opt for writing xr.DataTrees to disk to represent the graph what would the pros/cons there be? Is the xr.DataTree functionality I have already drafted (Feat/xr datatree output format weather-model-graphs#47) fit-for-purpose or is there something wrong/missing?

This issue is meant as a continuation that started on #339 on bridging the gap between weather-model-graphs and neural-lam.

[Joltsy10]

The DataTree on-disk + HeteroData in-memory split still seems like the right direction to me. The on-disk format should be self-describing and human-inspectable, which DataTree over netCDF or zarr gives you. HeteroData is then purely a runtime representation that the loading layer constructs, keeping the PyG dependency out of the serialization layer entirely.

On the suggested first step, loading the existing tensor-based format into a pyg.HeteroData object before passing it to ARModel seems like a clean incremental move. It doesn't require changing the on-disk format at all, just the loading layer, and it immediately simplifies the model code since you're passing one object instead of a dict of tensors. The hierarchical case maps onto HeteroData naturally since node types and edge types are first-class concepts there.

[varunsiravuri]

hey @leifdenby I think this reads like a good incremental direction. Keeping the on-disk format self-describing (e.g. xr.DataTree / netCDF / zarr) while using pyg.HeteroData only as the in memory/runtime representation seems like a clean separation.
As a first step, converting the current tensor/dict graph representation into a HeteroData object before passing it into ARModel sounds low-risk, because it avoids changing the serialization format and lets the model interface evolve incrementally.
If that first step is aligned with the intended direction, I’d be happy to help prototype it in a PR. Would you prefer to settle the design here first, or would a small draft PR exploring just the loading-layer + ARModel interface be useful?

[yuvraajnarula]

One observation from looking at the current graph pipeline is that neural-lam already relies on PyG utilities during graph creation (e.g. from_networkx in create_graph()), but the runtime representation ends up being a dictionary of tensors loaded from *_edge_index.pt and *_features.pt files via load_graph().

Because of that, converting this dictionary into a pyg.HeteroData object at the loading layer seems like a fairly contained first step. The hierarchical structure (g2m, m2m, m2g) maps quite naturally onto typed edges in HeteroData, while keeping the on-disk format unchanged.

A possible incremental path could be:

• keep the current .pt tensor export format unchanged
• extend the graph loader to optionally assemble a pyg.HeteroData object
• pass that object to the model layer instead of a tensor dictionary
• migrate model code gradually while preserving backward compatibility

This would keep serialization independent of PyG while allowing the internal model interface to move toward a single graph object representation.

Curious whether introducing this as an optional loader pathway first would make experimentation easier without forcing an immediate refactor across the codebase.
@leifdenby
If this direction sounds reasonable, I’d also be happy to help prototype the loader-side conversion as a small exploratory PR.

[Sir-Sloth-The-Lazy]

@leifdenby One edge case worth flagging for the loading-layer conversion: the current .pt graph files use a list format for hierarchical graphs — m2m_edge_index.pt, m2m_features.pt, and mesh_features.pt are all lists of length L (one entry per mesh level). When assembling a pyg.HeteroData object from these, each level needs to become a distinct node/edge type (e.g. mesh_level_0, mesh_level_1), otherwise the hierarchical structure is flattened and lost. The loader will need to handle both the L=1 (non-hierarchical) and L>1 (hierarchical) cases without the model layer needing to know which format was on disk. This might be worth designing explicitly before the PR rather than discovering it mid-implementation.

[kshirajahere]

One thing I've been thinking about is how this interacts with the ongoing model hierarchy refactor (issue #49 / PR #208). Currently, InteractionNet stores edge_index as a registered buffer at construction time the graph topology is baked into the module's state dict. BaseGraphModel.init unpacks the graph dict and creates separate InteractionNet instances for g2m, m2m, m2g, each owning its own edge topology. This means the graph structure is frozen into the model at build time. If we move to passing a HeteroData object at forward time instead, that changes the ownership model: the graph topology would no longer be part of the nn.Module tree, but an input flowing through the forward pass. This is arguably cleaner for swapping graphs at inference time, but it creates a question about the hierarchy refactor: Where in the ForecasterModule--> Forecaster --> ARForecaster --> StepPredictor chain does the HeteroData object live? In Joel's current design for StepPredictor, the forward signature is essentially (prev_state, prev_prev_state, forcing) --> (pred_state, pred_std), pure tensor-in, tensor-out. If we follow point 1 from this issue through to its conclusion and putting state/forcing as node features on the graph object; then StepPredictor.forward would instead receive a HeteroData carrying the atmospheric state and return a HeteroData carrying the prediction. That pushes graph-structure awareness into the forward interface of every model class, which is a significant contract change. More concretely, in the AR rollout loop (ARForecaster), the current design extracts the predicted state tensor from StepPredictor, applies boundary masking (boundary_mask * border_state + interior_mask * pred_state), and feeds the result back as input for the next step. If the state lives on a HeteroData object, ARForecaster would need to know how to extract predictions from the graph, apply the boundary mask, and inject the masked result back onto the graph for the next step which couples the rollout logic to the graph schema in a way that the current tensor-based interface avoids. So the question is: should the HeteroData representation stop at the boundary of StepPredictor (i.e., StepPredictor.init receives a HeteroData and unpacks it into edge indices and features internally, but forward stays tensor-in/tensor-out), or should the graph object flow through the entire forward pass? The first keeps the refactored hierarchy's clean tensor contracts intact but limits some of the ergonomic benefits Leif described. The second is the more ambitious design but would need to be coordinated with (#49) since it changes the StepPredictor / ARForecaster interface contract @leifdenby.
I think this is worth resolving before the first incremental step, because the choice of where HeteroData stops determines whether the loader-layer approach (as suggested above) is a stepping stone toward the full vision or a temporary adapter that would later need to be unwound.

[Sir-Sloth-The-Lazy]

Great question — this is exactly the kind of design tension worth resolving early. Here's my thinking after looking at the current code:

TL;DR: I'd advocate for Option 1 — HeteroData stops at the StepPredictor boundary.

The key insight is that PR #208's hierarchy was designed with clean, tensor-based contracts at each layer:

• StepPredictor.forward(prev_state, prev_prev_state, forcing) → (pred_state, pred_std)
• ARForecaster.forward(init_states, forcing_features, boundary_states) → (prediction, pred_std)
• ForecasterModule handles training/eval orchestration

This separation is intentional — ARForecaster's boundary masking logic (boundary_mask * boundary_state + interior_mask * pred_state) works cleanly on pure tensors, and so does the state feedback loop for auto-regressive rollout. Pushing HeteroData through the full forward chain would couple the rollout logic to the graph schema, meaning ARForecaster would need to know how to extract/inject node features from/into the graph object at every AR step. That's a significant contract change that adds complexity without clear benefit at the rollout level.

Where HeteroData does add clear value is inside StepPredictor/BaseGraphModel:

Right now, BaseGraphModel.__init__ unpacks the graph dict into separate registered buffers (g2m_edge_index, m2g_edge_index, g2m_features, m2g_features, etc. — see base_graph_model.py L64-76) and creates separate InteractionNet instances for g2m, m2m, m2g, each with baked-in edge topology. A HeteroData object could cleanly replace this dictionary-of-tensors with a single structured object, simplifying the __init__ unpacking and making the internal graph representation more ergonomic.

Concretely, the approach I'd suggest:

1. Loader layer: Convert the on-disk graph format (currently flat tensors via utils.load_graph()) into a pyg.HeteroData object. This could live in utils.py or a dedicated graph-loading module.
2. BaseGraphModel.__init__ receives the HeteroData and unpacks edge indices / features from it (instead of from the current dict). InteractionNet still gets raw edge_index tensors — no change to its interface.
3. StepPredictor.forward signature stays tensor-in/tensor-out — the HeteroData is a construction-time resource, not a forward-time input. The encode-process-decode forward in BaseGraphModel continues to concatenate state features into grid nodes, embed, run GNNs, and return (pred_state, pred_std).
4. ARForecaster and ForecasterModule are untouched — they keep working with pure tensors.

This makes the loader-layer approach a genuine stepping stone: if a future design wants to push state/forcing onto the graph as node features (Leif's full vision), that change would only need to modify BaseGraphModel.forward internally, and the StepPredictor tensor contract could be preserved via an unpack_from_graph()/pack_to_graph() pattern at the BaseGraphModel boundary.

The alternative (passing HeteroData through the entire chain) would require changing Forecaster.forward, ARForecaster's rollout loop, and possibly ForecasterModule's loss computation — a much larger surface area of change with unclear benefits at those levels.

@leifdenby @joeloskarsson curious to hear your thoughts, does this align with how you envisioned the incremental path ?, just a suggestion but hey lets dicuss if you find something to be unclear😁

[prajwal-tech07]

@leifdenby I went through the current graph loading path (load_graph() → BaseGraphModel.__init__ → InteractionNet) and the model hierarchy (ARModel → BaseGraphModel → GraphLAM / BaseHiGraphModel) to map out what a HeteroData migration would actually touch. Here's a concrete first-step proposal.

Where HeteroData should stop: at BaseGraphModel.__init__

+1 to @kshirajahere's framing. The answer is clear from the code:

• ARModel.forward / ARModel.common_step works with pure state tensors + boundary masks — no graph awareness
• BaseGraphModel.__init__ is the only place that unpacks the graph dict into self.g2m_edge_index, self.m2m_features, self.mesh_static_features, etc. (lines 27-34)
• InteractionNet.__init__ receives raw edge_index tensors and does its own reindexing (line 63 of interaction_net.py)

So the natural boundary is: loader returns HeteroData → BaseGraphModel.__init__ unpacks it → everything downstream stays tensor-in/tensor-out. ARModel, StepPredictor (in #208), and the rollout loop are untouched.

Concrete sketch: load_graph() → HeteroData

# In utils.py — new function alongside existing load_graph()

def load_graph_as_heterodata(graph_dir_path, device="cpu"):
    """Load graph tensors into a pyg.HeteroData object."""
    from torch_geometric.data import HeteroData

    hierarchical, graph_dict = load_graph(graph_dir_path, device)
    data = HeteroData()

    # Node features
    data["grid"].num_nodes = graph_dict["g2m_edge_index"].max() + 1  # from g2m receiver count
    if hierarchical:
        for level_i in range(len(graph_dict["mesh_static_features"])):
            data[f"mesh_{level_i}"].x = graph_dict["mesh_static_features"][level_i]
    else:
        data["mesh"].x = graph_dict["mesh_static_features"]

    # Edge indices + features
    data["grid", "to", "mesh"].edge_index = graph_dict["g2m_edge_index"]
    data["grid", "to", "mesh"].edge_attr = graph_dict["g2m_features"]

    data["mesh", "to", "grid"].edge_index = graph_dict["m2g_edge_index"]
    data["mesh", "to", "grid"].edge_attr = graph_dict["m2g_features"]

    if hierarchical:
        for level_i in range(len(graph_dict["m2m_edge_index"])):
            data["mesh", f"m2m_{level_i}", "mesh"].edge_index = graph_dict["m2m_edge_index"][level_i]
            data["mesh", f"m2m_{level_i}", "mesh"].edge_attr = graph_dict["m2m_features"][level_i]
        for level_i in range(len(graph_dict["mesh_up_edge_index"])):
            data["mesh", f"up_{level_i}", "mesh"].edge_index = graph_dict["mesh_up_edge_index"][level_i]
            data["mesh", f"up_{level_i}", "mesh"].edge_attr = graph_dict["mesh_up_features"][level_i]
            data["mesh", f"down_{level_i}", "mesh"].edge_index = graph_dict["mesh_down_edge_index"][level_i]
            data["mesh", f"down_{level_i}", "mesh"].edge_attr = graph_dict["mesh_down_features"][level_i]
    else:
        data["mesh", "to", "mesh"].edge_index = graph_dict["m2m_edge_index"]
        data["mesh", "to", "mesh"].edge_attr = graph_dict["m2m_features"]

    data.hierarchical = hierarchical
    return data

This wraps the existing load_graph(), so the on-disk format is unchanged. BaseGraphModel.__init__ can then unpack from the HeteroData object instead of the dict — same information, structured container.

Why per-level node types for hierarchical

@Sir-Sloth-The-Lazy is right that this is the critical detail. Looking at BaseHiGraphModel.__init__:

It creates separate embedders per level (self.mesh_embedders = nn.ModuleList)
mesh_static_features is a BufferList of [N_mesh[l], d] — each level has different node count
mesh_up_edge_index / mesh_down_edge_index connect nodes across levels
So hierarchical mesh levels are genuinely distinct node types — mesh_0, mesh_1, etc. Flattening them into a single "mesh" node set in HeteroData would lose the level structure and force re-splitting later, which defeats the purpose.

Sequencing

This builds naturally on top of #384 (tensor-on-disk from wmg):

1. Supporting tensor-on-disk-format from weather-model-graphs #384 lands first — to_neural_lam() writes the current .pt format, CLI calls wmg
2. This step — load_graph_as_heterodata() added alongside existing load_graph(), BaseGraphModel` updated to accept either
3. Future — once DataTree on-disk format is ready (wmg#47), the loader reads DataTree → HeteroData, and nothing downstream changes

I can take on step 2 as a follow-up PR after #384. The scope is contained: utils.py (new loader function) + base_graph_model.py (unpack from HeteroData) — no changes to InteractionNet, ARModel, or any model forward pass.

[yuvraajnarula]

@prajwal-tech07
Clean sketch. A few additions worth considering from the icosahedral side:

1. Validation layer
We'll need tests that verify:

• dict to HeteroData conversion preserves all edges/features exactly
• Model outputs are identical with both formats (for flat and hierarchical cases)

2. Variable edge feature dimensions
Spherical graphs use 4+ edge features (len + 3D vector). Your sketch assumes fixed 3-dim edges. The loader should:

• Read feature dimension dynamically from tensor shape
• Not hardcode expectations

3. Backward compatibility testing
Add a CI step that runs existing tests with both load_graph() and load_graph_as_heterodata() to catch regressions.

4. Metadata/versioning
If we add a version tag to the graph directory later, the loader can adapt without breaking old graphs.

Happy to help with testing infrastructure and spherical graph validation once you have a draft PR.

[Sharkyii]

one thing I don't think we've talked about yet:
do we actually know if HeteroData is going to be slower than what we have now?

Indexing: Is data["mesh", "to", "grid"].edge_index noticeably slower than just grabbing graph_dict["m2g_edge_index"]? Probably not a huge deal for single lookups, but we're doing this a lot during training.

[yuvraajnarula]

one thing I don't think we've talked about yet: do we actually know if HeteroData is going to be slower than what we have now?

Indexing: Is data["mesh", "to", "grid"].edge_index noticeably slower than just grabbing graph_dict["m2g_edge_index"]? Probably not a huge deal for single lookups, but we're doing this a lot during training.

@Sharkyii
Good question. Following @leifdenby's vision of cleaner graph internals, here's the performance picture:

Overhead analysis:

• HeteroData lookups vs dict – both O(1), construction cost once per run (amortized)
• @prajwal-tech07's sketch wraps existing tensors – no extra copies, forward pass unchanged

What we should benchmark to confirm zero regression:

• Load time (once per run – not critical, but worth tracking)
• Memory usage (should be identical)
• Forward pass speed (should be identical)
• GPU transfer overhead (if any)

Optimization suggestion: If we later add state/forcing as node features (Leif's long-term vision), we can avoid separate tensor copies by keeping features on-device in HeteroData; potentially reducing transfer overhead.

@leifdenby
If this aligns with your direction, happy to help benchmark once we have a prototype.

[mandeepsingh2007]

hii @leifdenby

This is a brilliant and pragmatic approach for the immediate term! Handling the unpacking at the BaseGraphModel.init gate via load_graph is definitely the safest way to introduce HeteroData without breaking the existing regional ARModel pipelines by @prajwal-tech07

Just adding a small thought for the long-term roadmap :
We might eventually want to explore a dual-track or polymorphic architecture where the core processor uses native PyG HeteroConv layers. This would allow the engine to ingest the HeteroData object natively end-to-end without any unpacking wrappers, making the GNN entirely schema-agnostic and perfectly future-proof.