Supporting tensor-on-disk-format from weather-model-graphs

[leifdenby]

The aim of this issue is two-fold:

• Create an implementation plan for:
  • upstream changes to weather-model-graphs to support writing to the tensor-on-disk format that neural-lam currently assumes (need to create issue for this and link here)
  • changes in neural-lam to call weather-model-graphs to create graphs and write these to disk in the current tensor-on-disk format in neural-lam
• Record what changes will be future work, i.e. what is missing from the agreed implementation plan so that we can address these issues later.

This issue aims to start "bridging the gap" between weather-model-graphs and neural-lam by building on the discussion on #339

The implementation being planned in here is reliant on the tensor-on-disk format to be finalised I think, so that we can apply the validator both on the weather-model-graphs and neural-lam side.

[leifdenby]

prior work

1. existing cli interface in neural_lam.create_graph which uses networkx directly to create the graph and write tensor-on-disk format: https://github.com/mllam/neural-lam/blob/main/neural_lam/create_graph.py
2. cli interface for creating rectangular mesh graphs from @joeloskarsson and @sadamov's research branch: https://github.com/joeloskarsson/neural-lam-dev/blob/research/neural_lam/build_rectangular_graph.py
3. @prajwal-tech07's comment on #339 we have the following as a suggestion towards what could be added to neural-lam:

import weather_model_graphs as wmg

xy = datastore.get_xy("state")

# Flat graph
components = wmg.create.archetype.create_keisler_graph(
    coords=xy, return_components=True
)
wmg.save.to_neural_lam(components, "graph/keisler", hierarchical=False)

# Hierarchical graph
components = wmg.create.archetype.create_oskarsson_hierarchical_graph(
    coords=xy, return_components=True
)
wmg.save.to_neural_lam(components, "graph/hierarchical", hierarchical=True)

[leifdenby]

I've spent some time now summarising the most important differences between the three pieces of prior work above:

Aspect	Current create_graph CLI	build_rectangular_graph.py CLI	@prajwal-tech07's code snippet
Input coordinates	Uses datastore.get_xy(category="state", stacked=False) (create_graph.py:546)	Uses utils.get_stacked_lat_lons(datastore, datastore_boundary) and passes coords, coords_crs, graph_crs, and optional decode_mask into weather-model-graphs	Uses datastore.get_xy("state")
Output location control	Writes to datastore.root_path / "graph" / <name> via --name (create_graph.py:572, create_graph.py:603)	Uses --graph_name and --output_dir; saves to datastore.root_path / "graphs" / <graph_name> by default or --output_dir / <graph_name>	Hard-codes "graph/keisler" and "graph/hierarchical"
Flat Keisler graph	Not the standard non-hierarchical mode; the closest available option is --levels 1, which gives a single-level flat graph but via the in-repo construction rather than an explicit Keisler archetype	Achieved directly with --archetype keisler	Achieved directly with wmg.create.archetype.create_keisler_graph(...) and wmg.save.to_neural_lam(..., hierarchical=False)
Flat GraphCast-style multi-scale graph	This is the standard non-hierarchical path: omit --hierarchical and use multiple levels via --levels; the mesh levels are merged into one flat multi-scale graph in create_graph.py	Achieved directly with --archetype graphcast; supports --mesh_node_distance, --level_refinement_factor, and --max_num_levels	Not shown in the snippet; the snippet uses Keisler for flat graphs, not a GraphCast-style flat multi-scale archetype
Hierarchical graph	Achieved with --hierarchical; the CLI builds multiple mesh levels plus inter-level mesh_up and mesh_down edges in the in-repo implementation	Achieved directly with --archetype hierarchical; uses the Oskarsson hierarchical archetype	Achieved directly with wmg.create.archetype.create_oskarsson_hierarchical_graph(...) and wmg.save.to_neural_lam(..., hierarchical=True)
One-level graph mode	Yes, via --levels 1; covered by tests as "1level" (tests/test_graph_creation.py:16)	No dedicated one-level mode flag, but a single-level keisler graph is available via --archetype keisler	No one-level mode shown
Variable number of levels	Yes, --levels / n_max_levels limits mesh depth (create_graph.py:583, create_graph.py:247)	Yes, via --max_num_levels; also exposes --level_refinement_factor	No level parameter shown
Plotting during generation	Yes, --plot with matplotlib (create_graph.py:578)	No plotting shown	No plotting shown
Graph family / construction logic	Custom in-repo graph construction based on grid geometry, KD-trees, and hand-written connectivity (create_graph.py:158)	Delegates to weather-model-graphs archetype constructors, selected via --archetype	Delegates to weather-model-graphs archetype constructors
Neural-LAM export	Saves Neural-LAM .pt files directly in this module (create_graph.py:82)	Saves in a Neural-LAM-compatible layout using repeated wmg.save.to_pyg(...) calls rather than wmg.save.to_neural_lam(...)	Calls wmg.save.to_neural_lam(...)
Hierarchical save flag	Controlled by --hierarchical and output content	Inferred from --archetype hierarchical or kwargs["m2m_connectivity"] == "hierarchical" before saving	Explicitly passed as hierarchical=False/True
Expected file set	Produces files expected by utils.load_graph() and tests (utils.py:47, tests/test_graph_creation.py:44)	Intentionally saves a Neural-LAM-compatible file layout, including separate mesh_up and mesh_down outputs for hierarchical graphs	Intended to produce Neural-LAM format, but the exact file contract is not shown in the snippet

Based on this we should design what the create_graph CLI should look like. I would err on the side of it working like the CLI on the dev-branch, but putting all functionality wrt to writing to the neural-lam tensor-on-disk format in weather-model-graphs instead of in the CLI script itself.

[Joltsy10]

Quick question on scope. Is global icosahedral graph support something we want to fold into this discussion, or is that tracked separately in #348? Asking because some of the to_neural_lam() design decisions made here, particularly around edge feature conventions and mesh node feature semantics, will directly affect whether an icosahedral graph can use the same path without a separate code path. Given that global support is being targeted fairly soon it seems worth aligning early rather than discovering conflicts mid-implementation.

On the implementation plan itself, one thing worth including in the to_neural_lam() spec is a node_id_map companion file alongside the .pt outputs. The wmg node labels are discarded during from_networkx() conversion with no record kept, which blocks latent space visualization and round-trip validation as discussed in #20 and #339. It is a small addition to the save step but the window to bake it into the spec is now, before the implementation solidifies.

[leifdenby]

Quick question on scope. Is global icosahedral graph support something we want to fold into this discussion, or is that tracked separately in #348?

@Joltsy10 I would view this issue primarily to be around how to design the new CLI for calling weather-model-graphs to create tensor-on-disk graph representations. If your question relates to that CLI needing more command-line arguments (for example to be able to specify that the graph created is for global, spherical, forecasting) or the call arguments to weather-model-graphs needing to be different then we should discuss that here. Although, I would expect that would be a future work aspect of this CLI, as-in we will support creating rectilinear graphs first. Once we have that working we should extend the cli to support calling weather-model-graphs to create other graph topologies. In terms of extending the tensor-on-disk format, getting wmg to write to that and then supporting it in the CLI I would expect the following steps:

• issue + PR describing what needs to change to the tensor-on-disk format and the validator. We have the issue discussing this, but discussion is still ongoing I think
• issue + PR extending weather-model-graphs to match the new tensor-on-disk format. This might first make sense when we have the format nailed down, but you are welcome to create an issue on the wmg repo if you like and have some ideas already :)
• issue + PR extending the CLI being discussed here to create graphs. This would be a new issue building on the one we're discussing here.

Does this make sense?

[varunsiravuri]

@leifdenby I think it makes sense to treat this issue as the planning thread and keep the first implementation target narrow: rectilinear graphs first, with broader topology support captured as future work.
From the comparison above, I agree the CLI should probably look closer to the dev-branch interface, while pushing the Neural-LAM tensor-on-disk writing logic down into weather-model-graphs rather than keeping that logic in the CLI layer.
If that sounds right, would it be helpful to split out a follow-up issue for one concrete implementation step once the plan is settled, for example the weather-model-graphs side needed to emit the current Neural-LAM tensor-on-disk format? I’d be happy to help on that once scope is fixed.

[yuvraajnarula]

@leifdenby
It might help to think of the bridge here as three layers:

• graph construction : handled by weather-model-graphs (archetypes like keisler / hierarchical)
• serialization : writing the Neural-LAM tensor-on-disk format
• execution interface : the create_graph CLI orchestrating the above

If the tensor serialization lives inside weather-model-graphs (e.g. wmg.save.to_neural_lam(...)) and the CLI only coordinates creation + saving, the format contract stays centralized and the CLI can remain thin.

Since both repos depend on the same tensor-on-disk spec, running the validator both after WMG export and before utils.load_graph() in neural-lam might help catch format drift early.

Also wondering if defining a small format/version tag alongside the tensor outputs would make future format evolution easier without breaking compatibility across the two repos.

[Sir-Sloth-The-Lazy]

@leifdenby — one thing worth explicitly capturing in this implementation plan, building on the open discussion in PR #323: the current spec still has an inconsistency in where normalization happens. mesh_features are pre-normalized at save time (divided by max_pos), while edge features are normalized at load time in utils.load_graph (divided by longest_edge). You noted on March 12 that this is deferred future work on the neural-lam side. For wmg.save.to_neural_lam to produce correct files, the implementer needs a clear answer: should it write raw or pre-normalized features? It might be worth flagging this explicitly as a blocking dependency in the implementation plan so it gets resolved before the wmg side starts writing files.

[prajwal-tech07]

@leifdenby Thanks for the detailed comparison. Here's a concrete CLI proposal that closes the gaps in my original snippet (missing GraphCast/multi-scale, level params, plotting, one-level mode):

Proposed CLI

python -m neural_lam.create_graph \
  --datastore <config_path> \
  --archetype keisler|graphcast|hierarchical \
  --graph_name <name> \
  --output_dir <path>              # optional, defaults to datastore.root_path / "graphs" / <graph_name>
  --mesh_node_distance <float>     # forwarded to wmg (graphcast/hierarchical)
  --max_num_levels <int>           # forwarded to wmg (graphcast/hierarchical)
  --level_refinement_factor <int>  # forwarded to wmg (graphcast/hierarchical)
  --plot                           # optional post-creation visualization

Archetype mapping covers all rows from the comparison table:

Use case	CLI invocation
Flat Keisler	--archetype keisler
Flat GraphCast multi-scale	--archetype graphcast --max_num_levels 3
Hierarchical	--archetype hierarchical --max_num_levels 4
One-level (flat)	--archetype keisler (single-level by definition)

CLI internals (thin wrapper)

def create_graph(args):
    datastore = init_datastore(args.datastore)
    xy = datastore.get_xy("state")

    # Resolve CRS + decode_mask from datastore (mirrors dev-branch approach)
    coords_crs = datastore.coords_projection
    graph_crs = datastore.coords_projection
    decode_mask = ~datastore.boundary_mask if hasattr(datastore, 'boundary_mask') else None

    archetype_fn = {
        "keisler": wmg.create.archetype.create_keisler_graph,
        "graphcast": wmg.create.archetype.create_graphcast_graph,
        "hierarchical": wmg.create.archetype.create_oskarsson_hierarchical_graph,
    }[args.archetype]

    # Forward relevant kwargs to wmg
    wmg_kwargs = {k: v for k, v in vars(args).items()
                  if k in ("mesh_node_distance", "max_num_levels", "level_refinement_factor")
                  and v is not None}

    components = archetype_fn(
        coords=xy, coords_crs=coords_crs, graph_crs=graph_crs,
        decode_mask=decode_mask, return_components=True, **wmg_kwargs
    )

    hierarchical = (args.archetype == "hierarchical")
    wmg.save.to_neural_lam(components, args.output_dir, hierarchical=hierarchical)

    if args.plot:
        # Post-creation visualization using wmg.visualise
        ...

No tensor manipulation in the CLI — all serialization logic lives in wmg.save.to_neural_lam().

On normalization (@Sir-Sloth-The-Lazy's question)

to_neural_lam() should write raw (unnormalized) features. Current create_graph.py writes raw values and load_graph() normalizes at load time (edge features ÷ longest_edge). Keeping this convention means the wmg save step stays simple and neural-lam owns its normalization policy.

What can proceed now (in parallel with #323)

Track	Work	Where
1	Implement to_neural_lam() based on my #339 draft, validate against #323 spec (also addresses wmg#88)	wmg
2	Refactor CLI as above, replacing create_graph.py	neural-lam
3	End-to-end test: archetype -> to_neural_lam() -> load_graph() -> model forward with MEPS fixture	both

Ready to open draft PRs for tracks 1 and 2 once you confirm this direction.

[yuvraajnarula]

@prajwal-tech07
Your CLI proposal looks great. From the icosahedral graph side, I'd like to help ensure it works seamlessly for global forecasting.

Specifically, I can:

1. Test the CLI with icosahedral graphs generated, once to_neural_lam() is ready
2. Add validation tests that run the graph through the entire pipeline: WMG ->CLI -> loader -> model forward
3. Document any icosahedral-specific requirements for the tensor format (e.g., 3D mesh features, barycentric weights)

@leifdenby
Happy to coordinate with @prajwal-tech07 on this. I can have a test PR ready once the CLI implementation settles.

[manthansubhash01]

Hello @leifdenby,
I went through the issue I implement a solution I think that would solve the issue at best of my knowledge please kindly review it and provide feedback if needed any changes

[leifdenby]

Happy to coordinate with @prajwal-tech07 on this. I can have a test PR ready once the CLI implementation settles.

@prajwal-tech07 and @yuvraajnarula thank for your contributions to the discussion here! Your CLI suggestion for the neural-lam side looks great @prajwal-tech07 - let's go with that 👍 Do you or @yuvraajnarula want to take on implementing that? One caveat: we decided at the dev meeting today to keep the old cli script so let's:

• add the new wmg-based CLI script as create_graph_with_wmg.py and
• add a warning to the current create_graph.py to say that this old way of creating graphs will be deprecated and they should use the new CLI going forward.

@manthansubhash01 thank you for making a PR already, but your PR mixes together the current and new cli functionality being discussed here. I hope it ok I close your PR since you made it before we had finished deciding the plan of action here :)