Add create_graph_with_wmg.py CLI using weather-model-graphs - #596
Add create_graph_with_wmg.py CLI using weather-model-graphs#596prajwal-tech07 wants to merge 12 commits into
Conversation
|
Merged Leif's |
leifdenby
left a comment
leifdenby
left a comment
There was a problem hiding this comment.
Looking good! A few suggestions for changes 🚀
| } | ||
|
|
||
|
|
||
| def _estimate_mesh_node_distance(xy): |
There was a problem hiding this comment.
maybe we could call this _estimate_grid_node_spacing? And then expose the grid-mesh-spacing ratio as a CLI arg that defaults to 3.0?
There was a problem hiding this comment.
Renamed _estimate_mesh_node_distance → _estimate_grid_node_spacing — it now returns only the average grid spacing. The ×3 multiplier is replaced by a new grid_mesh_ratio parameter (default 3.0) exposed as --grid_mesh_ratio on the CLI.
There was a problem hiding this comment.
looks great, but it would better to name it --grid_mesh_spacing_ratio I think - just be clear about what it is the ratio between :) what do you think?
There was a problem hiding this comment.
--grid_mesh_ratio → --grid_mesh_spacing_ratio
Good call — renamed the CLI arg, function parameter, and docstrings throughout to make it clear it's the ratio between mesh-node and grid-node spacing.
|
|
||
| @pytest.mark.parametrize("archetype", ["keisler", "graphcast", "hierarchical"]) | ||
| @pytest.mark.parametrize("datastore_name", DATASTORES.keys()) | ||
| def test_wmg_graph_creation(datastore_name, archetype): |
There was a problem hiding this comment.
This is great, but we could actually use the testing that #323 introduces. I should try and get that finished so that we can merge both in together :)
There was a problem hiding this comment.
Sounds like a great plan — either way works for me. Happy to adapt the tests to #323's infrastructure once it lands, or merge as-is if this PR is ready
There was a problem hiding this comment.
we can remove these tests now, no? Maybe we should instead just call the validator from #323 here?
| "plotly>=5.15.0", | ||
| "torch>=2.3.0", | ||
| "torch-geometric==2.3.1", | ||
| "torch-geometric>=2.5.3", |
There was a problem hiding this comment.
Yes, this is necessary. The weather-model-graphs[pytorch] extra requires torch-geometric>=2.5.3 in its own dependencies, so we need to allow at least that version here too.
There was a problem hiding this comment.
ah, in that case I think we maybe should reduce required version number on the weather-model-graphs side, since all we are doing with pytorch-geometric is using it to convert the networkx.DiGraph objects to torch.Tensor objects, and we already do that in the current create_graphs.py code in neural-lam with the older version. So how about we change weather-model-graphs to require torch-geometric==2.3.1 so the two are in sync? Then we don't have to change anything on the neural-lam side
There was a problem hiding this comment.
You're right — wmg only uses from_networkx which works fine with 2.3.1. I've also pushed a commit to the wmg PR branch (prajwal-tech07/weather-model-graphs@d0c693d) pinning it to ==2.3.1 there too, so both repos stay in sync.
There was a problem hiding this comment.
I will make a separate PR for this. I think this is some general CI maintenance that should be merged in before this PR.
There was a problem hiding this comment.
Noted — happy to revert the CI changes from this branch if you'd prefer to handle them in a separate PR.
leifdenby
left a comment
leifdenby
left a comment
There was a problem hiding this comment.
Thanks for this, I added additional comments on previous review to resolve:
There was a problem hiding this comment.
Ok, the CLI interface looks great now, as do the readme instructions and install works great (just tried locally on my laptop too).
I think we should also change the tests that use graph creation (i.e. the training examples, they all use the create_graph_from_datastore() function in https://github.com/mllam/neural-lam/blob/main/neural_lam/create_graph.py#L540), but I think we should use the new one that you have created that creates a graph with wmg instead.
|
Thanks for the suggestion, @leifdenby! Done in 76c7ed7. All test files that used the old
All tests pass locally (55/55 in the modified files, 72/72 in the full suite. |
|
I just completed a 200 epoch training (on a DGX Spark machine) on the DANRA test dataset using a 
I ran the following commands: So I will give this a final review, but I think this is nearly ready to merge :) |
|
Hi @leifdenby! Implementation is complete. |
7356723 to
8fe2e3e
Compare
8fe2e3e to
5055f4d
Compare
|
Rebased onto |
5055f4d to
7063f99
Compare
|
Rebased onto latest Changes since the last review:
Verification: the full test suite passes locally against the released 0.4.0 - 259 passed, including graph creation for all three archetypes, the validation test, and a full training run. pre-commit is clean. One note on the dependency: I deliberately did not use the Your earlier review is still marked as changes requested - I believe everything from it is addressed now ( |
Update pyproject.toml to install weather-model-graphs from the issue-384/to-neural-lam branch of the fork, so that CI and reviewers can test the neural-lam side against the unreleased to_neural_lam() changes before wmg PR mllam#123 is merged. Will revert to a versioned PyPI dependency once PR mllam#123 is released.
…sh_ratio, use stacked=True, add return_components comment, add README entry
Migrated all test files that used the old create_graph_from_datastore() from neural_lam.create_graph to use the new wmg-based version from neural_lam.create_graph_with_wmg instead. Changes: - test_datasets.py: Use wmg create_graph_from_datastore with archetype='keisler' - test_clamping.py: Use wmg create_graph_from_datastore with archetype='keisler' - test_plotting.py: Use wmg create_graph_from_datastore with archetype='keisler' - test_training.py: Use wmg create_graph_from_datastore with archetype='keisler' - test_plot_graph.py: Use wmg create_graph_from_datastore with keisler and hierarchical archetypes. Removed multiscale (graphcast) parametrization since the graphcast archetype produces multi-level m2m edges without up/down edges, which is not yet compatible with utils.load_graph(). Graphcast graph creation is separately tested in test_graph_creation.py.
Follow rename in weather-model-graphs (mllam/weather-model-graphs#123).
- test_graph_creation: validate create_graph_with_wmg output on disk with the graph-storage validator (docs/validate_graph.py from mllam#323) and assert the metainfo.yaml spec_version -- the end-to-end format contract for the bridge - create_graph_with_wmg: add module docstring (interrogate 100% coverage) - pyproject: merge the weather-model-graphs pin into a single [tool.uv.sources] table (rebase left two declarations, which is invalid TOML) - uv.lock: regenerate so 'uv sync --locked' matches the wmg PR mllam#123 branch pin
weather-model-graphs 0.4.0 is released on PyPI and includes save.to_torch_tensors_on_disk, so the temporary git dependency on the PR mllam#123 branch is no longer needed. - pyproject: bump to weather-model-graphs>=0.4.0 and drop the [tool.uv.sources] git pin - uv.lock: regenerate so it resolves 0.4.0 from PyPI The [pytorch] extra is intentionally not used: it requires torch-geometric>=2.5.3, which conflicts with neural-lam's torch-geometric==2.3.1 pin. It is not needed here, since neural-lam already provides torch/torch-geometric and the tensor-on-disk save path does not import torch-geometric.
7063f99 to
bba00b3
Compare
|
@leifdenby gentle ping on this one. It's green and mergeable - the bump to weather-model-graphs>=0.4.0 and the regenerated uv.lock went in on 28 July, so it now resolves 0.4.0 straight from PyPI with no fork refs. Going by the merge order we agreed: spec #323 merged, the validation test is on this branch, wmg #123 merged and v0.4.0 released - this PR is the last remaining step. The points from your April review are all addressed. Is there anything else you'd like changed before it can go in? |
| ``"hierarchical"``. | ||
| mesh_node_distance : float or None | ||
| Distance between created mesh nodes (in coordinate units). If None, | ||
| automatically estimated as ``grid_mesh_spacing_ratio * grid_spacing``. |
There was a problem hiding this comment.
what is grid_spacing here? I think it is the grid node spacing, but we should state here that we automatically try to estimate this.
There was a problem hiding this comment.
Yes, it's the grid node spacing - roughly the distance between neighbouring grid points, estimated from the datastore coordinates when mesh_node_distance isn't given. Agreed the docstring should say that outright instead of leaving it to be inferred.
I'll fold this in together with the rename once we've settled the terminology in the other thread, since the wording depends on whether we land on "distance" or "resolution". Same goes for the log line you suggested.
There was a problem hiding this comment.
let's go with distance here like you suggest :)
|
|
||
| if mesh_node_distance is None: | ||
| grid_spacing = _estimate_grid_node_spacing(xy) | ||
| mesh_node_distance = grid_spacing * grid_mesh_spacing_ratio |
There was a problem hiding this comment.
I think we should log here that because mesh_node_distance wasn't provided we estimate the grid-node spacing and use that. Maybe something like mesh_node_distance not provided so estimating grid-node spacing, with value estimated to be xx.xx, resulting mesh-node spacing will be yy.yy - what do you think?
There was a problem hiding this comment.
Yes Leif, this should be logged. It's the one number in the command the user didn't pick, so it shouldn't be silent.
One thing on the wording. Your version prints the estimated grid spacing and the resulting mesh spacing, but not the ratio between them, and the ratio is what someone would reach for if the number looks wrong. I'd rather the line explained itself:
mesh_node_distance not given; estimated grid node spacing 12.50 x grid_mesh_spacing_ratio 3.0 -> mesh_node_distance 37.50
(parameter name rather than the CLI flag, since create_graph_from_datastore gets called directly from the tests too)
The other reason I want it visible: the estimate is rougher than "estimated" suggests. It's sqrt(x_range * y_range / N), so it returns the geometric mean of dx and dy, which on an anisotropic grid lands between the two rather than on either. And since np.ptp measures n-1 intervals across n points, it comes out low by sqrt((nx-1)(ny-1)/(nx*ny)). I measured -0.2% on 500x500, -1% on 100x100, -10% on 10x10. Doesn't matter on a real domain, very visible on a small test one, which is exactly when someone is squinting at a graph wondering why it looks off.
The file has no logger yet, so I'd add loguru at info to match create_graph.py. Happy to make it a warning instead.
Wording depends on the naming thread, by the way. If we go with "distance" this becomes estimated grid node distance.
There was a problem hiding this comment.
sounds great! Let's include the ratio as well, good idea
| help="Graph archetype to create", | ||
| ) | ||
| parser.add_argument( | ||
| "--mesh_node_distance", |
There was a problem hiding this comment.
I think it is a bit confusing that we use both term "spacing" and "distance" for the mesh/grid resolution. I appreciate that wmg uses "distance", so maybe should use the same here? "resolution" would probably be an even better term. What do you think?
There was a problem hiding this comment.
Yeah, Leif - that's me mixing the two, not wmg. I had a look and there's no convention to protect either way: "spacing" doesn't appear anywhere in the codebase outside this file, and "resolution" only turns up in a comment in create_graph.py and a cartopy resolution="50m" kwarg. So we can just pick one.
I'd go with "distance". wmg's parameter is mesh_node_distance and we pass it straight through, so anything else means the flag and the thing it sets end up with different names for the same quantity. And "resolution" in a weather context usually runs the other way - higher resolution is a smaller number - so mesh_resolution=5 could be read as 5 units apart or 5x finer, where a distance of 5 can only mean one thing.
That would give mesh_node_distance (unchanged), _estimate_grid_node_distance, and grid_spacing -> grid_node_distance.
One thing I'd like your view on: --grid_mesh_spacing_ratio reads as grid:mesh, but the code is mesh_node_distance = grid_spacing * ratio, so it's mesh:grid. At the default of 3.0 the mesh nodes end up 3x further apart, not 3x closer. If we're renaming anyway, --mesh_grid_distance_ratio would match, though I don't know if churning a CLI flag is worth it.
There was a problem hiding this comment.
I'd go with "distance". wmg's parameter is
mesh_node_distanceand we pass it straight through, so anything else means the flag and the thing it sets end up with different names for the same quantity. And "resolution" in a weather context usually runs the other way - higher resolution is a smaller number - somesh_resolution=5could be read as 5 units apart or 5x finer, where a distance of 5 can only mean one thing.
I agree - good point!
One thing I'd like your view on:
--grid_mesh_spacing_ratioreads asgrid:mesh, but the code ismesh_node_distance = grid_spacing * ratio, so it's mesh:grid. At the default of 3.0 the mesh nodes end up 3x further apart, not 3x closer. If we're renaming anyway,--mesh_grid_distance_ratiowould match, though I don't know if churning a CLI flag is worth it.
Very well caught! Nice, yes rename the arg
|
|
||
| @pytest.mark.parametrize("archetype", ["keisler", "graphcast", "hierarchical"]) | ||
| @pytest.mark.parametrize("datastore_name", DATASTORES.keys()) | ||
| def test_wmg_graph_creation(datastore_name, archetype): |
There was a problem hiding this comment.
we can remove these tests now, no? Maybe we should instead just call the validator from #323 here?
| if graph_name == "hierarchical": | ||
| hierarchical = True | ||
| n_max_levels = 3 | ||
| elif graph_name == "multiscale": |
There was a problem hiding this comment.
why did you remove the multiscale option?
There was a problem hiding this comment.
I removed it because at the time it genuinely didn't work. Back then to_torch_tensors_on_disk split non-hierarchical m2m by the level edge attribute, so graphcast came out as a multi-level list. load_graph infers hierarchical = len(m2m_edge_index) > 1, so it read a flat graph as hierarchical and went looking for mesh_up_*.pt files that aren't there.
Worth being explicit that this is GC-LAM from the README's three paper graphs - it should absolutely be covered here, especially in a PR that deprecates create_graph.py.
That's no longer true. When I conformed the function to the storage spec from #323, non-hierarchical graphs started being written as a single merged m2m level, and that's what's in 0.4.0. The note I left in the fixture is just stale - I should have updated it then.
Checked against the released 0.4.0:
keisler m2m levels: 1 (8320 edges) mesh_up: no load_graph -> hierarchical=False
graphcast m2m levels: 1 (6096 edges) mesh_up: no load_graph -> hierarchical=False
hierarchical m2m levels: 3 (5512/544/40) mesh_up: yes load_graph -> hierarchical=True
So it loads fine now, and I've put multiscale back as the graphcast archetype. All three are green (6 passed).
| ### Added | ||
| - Add latent encoder/decoder modules and the `GraphEFM` (hierarchical) / `GraphEFMMultiScale` (flat) step predictors for the Graph-EFM ensemble forecasting model. [\#648](https://github.com/mllam/neural-lam/pull/648) @Sir-Sloth-The-Lazy | ||
|
|
||
| - Add `neural_lam.create_graph_with_wmg` CLI which builds graphs with |
There was a problem hiding this comment.
this is a bit verbose, can you shorten it to just a single sentence?
There was a problem hiding this comment.
Done, shortened to a single sentence.
Ah, I hadn't realised this. There won't be anything in |
Restore the multiscale (GC-LAM) case in the plot_graph fixture, using the graphcast archetype. It was dropped when non-hierarchical m2m edges were still written split by level, which load_graph misread as hierarchical. Non-hierarchical graphs are written as a single merged m2m level since to_torch_tensors_on_disk was conformed to the graph storage spec, so all three paper graph types load again. Shorten the create_graph_with_wmg changelog entry to a single sentence.
Drop the hand-rolled file presence, spec version, container type and shape assertions from test_wmg_graph_creation; the validator from mllam#323 already covers all of them, and pinning edge features to exactly 3 was stricter than the spec, which allows 3 or 4. Keep an explicit check that mesh_up/mesh_down files are present exactly when the archetype is hierarchical. The validator infers hierarchy from the graph contents, so it cannot tell whether the requested archetype was honoured.
|
On
The validator's actually already in there - it runs right after the graph is built, and that's what the I checked what would be lost first, by corrupting a valid graph and confirming the validator still fails: edge features 3 -> 2 columns, mesh features 2 -> 1, One thing I did keep: the validator infers whether a graph is hierarchical from its contents, so it can't tell whether the archetype that was requested was honoured - a hierarchical run that silently produced a flat graph would still pass. So there's now a single assert that the |
I tried it. It works, but it's a bigger change than it looks, so I'd rather not fold it into this PR. There's also a trap if we do it separately. Loosening the pin without the extra looks like a no-op, because Either way this PR doesn't need the extra. |
Ah yes, that's because we have |
3 participants
Describe your changes
Add a new CLI script
create_graph_with_wmg.pythat delegates graph creationto weather-model-graphs (wmg)
and saves the output using
wmg.save.to_neural_lam()in the tensor-on-diskformat expected by
neural_lam.utils.load_graph().This is the neural-lam side of the bridge described in #384. The wmg side
is mllam/weather-model-graphs#123,
which adds
to_neural_lam()towmg.save.What this PR does:
neural_lam/create_graph_with_wmg.pywith support for all three wmgarchetypes:
keisler(flat single-scale),graphcast(flat multiscale),and
hierarchical(Oskarsson)mesh_node_distancefrom grid spacing when not specified(Nx, Ny, 2)→(N, 2)for wmgcreate_graph.pyCLI pointing usersto the new script
weather-model-graphs[pytorch]>=0.3.0as a dependencytorch-geometricpin from==2.3.1to>=2.5.3(required bywmg's pytorch extra)
create_graph_with_wmgas a console script entry pointDependencies: Requires mllam/weather-model-graphs#123
to be merged and released first (adds
wmg.save.to_neural_lam()).Files changed (4 files, +305 −1):
neural_lam/create_graph_with_wmg.pyneural_lam/create_graph.pypyproject.tomltests/test_graph_creation.pyIssue Link
Solves #384 (neural-lam side)
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