Fix square mesh assumption in create_graph to support non-square domains by osten-antonio · Pull Request #373 · mllam/neural-lam

osten-antonio · 2026-03-10T18:34:48Z

Describe your changes

create_graph previously generated square mesh graphs regardless of the data grid's physical aspect ratio, by passing the same value n for both x and y dimensions of mk_2d_graph. This caused two concrete problems, as stated in issue #371 :

Uneven mesh spacing: for a 200x30 domain (6.7:1 aspect ratio), the original 27x27 mesh packs y-nodes 6.7x more densely than x-nodes. This means the model sees different effective resolutions in each direction — weather features moving diagonally will appear different depending on their orientation relative to the mesh.
Redundant nodes: some of the rows of the smaller dimension in the original mesh cover the same grid points, wasting computation with no benefit to model accuracy.

Changes:

At each mesh level, n_x and n_y are now derived independently by scaling the larger dimension proportionally to the physical
domain extent (x_extent / y_extent), assigning n_larger to whichever axis has greater physical extent.
mesh_dims tracks (n_x, n_y) per level, replacing the square assumption in the merge loop reshape.
dm is now computed from the mesh cell diagonal (sqrt(dx**2 + dy**2)) rather than x-direction spacing only. This would also guarantee full grid coverage with zero orphan node, which was a problem when migrating to a non square resolution.

No new dependencies.

Issue Link

Closes #371

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.

…n create_graph

kshirajahere · 2026-03-12T07:37:01Z

    for lev in range(1, mesh_levels + 1):
-        n = int(nleaf / (nx**lev))
-        g = mk_2d_graph(xy, n, n)
+        n_larger = int(nleaf / (nx**lev))


This rounding scheme is not compatible with the existing [1::3, 1::3] parent-selection logic below. For a rectangular grid like xy.shape == (100, 600), this produces level dims (14, 81) and then (4, 27), but sampling the previous level yields (5, 27) parents, so create_graph(..., n_max_levels=2) crashes at line 399 with ValueError: cannot reshape array of size 270 into shape (108,2). The coarse dims need to be derived from the sampled previous level, or the relabel/compose logic needs to change.

kshirajahere · 2026-03-12T07:42:33Z

                        assert r.shape[1] == d_features
+
+
+def test_graph_creation_non_square_aspect_ratio():


This test only covers the n_max_levels=1 case, so it misses the actual breakage introduced by the new sizing logic. Please add at least one rectangular multiscale case (n_max_levels=2) and one square-grid regression check for unchanged g2m connectivity, otherwise both failures above still pass CI.

kshirajahere · 2026-03-12T07:42:35Z

-    dm = np.sqrt(
-        np.sum((vm.data("pos")[(0, 1, 0)] - vm.data("pos")[(0, 0, 0)]) ** 2)
-    )
+    dx = x_extent / mesh_dims[0][0]  # cell width at finest mesh level


This changes g2m connectivity for every existing square graph, not just non-square ones. Previously the radius was based on one mesh-node spacing; here it becomes the cell diagonal, which is sqrt(2) larger on square meshes. On an 81x81 square grid, g2m edges jump from 9009 on main to 17757 here. That is a topology regression unless it was explicitly intended and rebase lined everywhere.

kshirajahere

The aspect-ratio fix is not safe to merge as written. It introduces a hard failure for non-square multiscale graphs because the rounded per-level mesh dimensions no longer match the existing 3x downsampling/relabeling assumption in the non-hierarchical path. A concrete repro is create_graph(..., xy.shape=(100, 600), n_max_levels=2), which now fails at create_graph.py:399 with a reshape error.

It also changes g2m topology globally by switching the coverage radius from single-axis neighbor spacing to the mesh-cell diagonal. On a square 81x81 grid, that increases g2m edges from 9009 on main to 17757 in this PR. The added test only exercises the one-level rectangular case, so neither regression is caught.

I would block this until:

Multi-level rectangular meshes are generated with level sizes that are consistent with the parent mapping logic.
Square-grid g2m connectivity is preserved, or the topology change is made explicit and covered with updated tests.

…d edge preservation

osten-antonio · 2026-03-13T09:39:30Z

Hi @kshirajahere , thanks for the review! I have addressed both issue; fixed n_max_levels = 2 crash by using the previous level's coarse dimension, and on square grids I reverted to using the original dm formula. I have also added the tests covering multi-level rectangular mesh, and square grid g2m edge count preservation. Let me know if there is anything else to address :D

joeloskarsson · 2026-03-22T17:48:48Z

Sorry for not looking at this earlier, but I would not put effort into such improvements to create_graph.py, as this functionality already exists and is much more robust in weather-model-graphs. The plan is to drop create_graph.py in its current form once #83 is complete.

kshirajahere

Thanks for addressing the earlier blockers.
The multiscale reshape crash concern and square-grid g2m topology regression were both handled, and the added tests for those paths are useful.
I am requesting some more changes please address them.

kshirajahere · 2026-03-22T19:12:56Z

+    dy = y_extent / mesh_dims[0][1]  # cell height at finest mesh level
+
+    # compare using mesh_dims because dx dy can cause rounding issue
+    if mesh_dims[0][0] == mesh_dims[0][1]:


This condition should be based on spacing equivalence (for example dx vs dy with tolerance), not only n_x == n_y. Rounded equal dims can happen on non-square domains and cause coverage regressions.

kshirajahere · 2026-03-22T19:13:16Z

+            # assign n_larger to whichever direction has greater physical extent
+            if x_extent >= y_extent:
+                n_x = n_larger
+                n_y = max(1, round(n_larger * y_extent / x_extent))


Please guard degenerate coordinate extents before division. A zero extent on one axis should fail fast with a clear error or use a safe fallback.

kshirajahere · 2026-03-22T19:16:49Z

Also add changelog entry

osten-antonio · 2026-03-24T09:43:24Z

Hi @kshirajahere thanks for the rereview! I think I will follow @joeloskarsson's advice seeing that there is much more development on weather-model-graph end with its multiple layout types.

sadamov · 2026-04-28T03:30:40Z

Thanks @osten-antonio. Closing per your own follow-up - going with @joeloskarsson's earlier suggestion to drop changes to create_graph.py in favour of the weather-model-graphs migration tracked in #384 / #596. If any of these fixes still apply on the wmg side once #596 lands, please port them there.

osten-antonio added 2 commits March 11, 2026 01:11

fix: adjust mesh dimensions to account for non-square aspect ratios i…

e7fadaa

…n create_graph

test: add test for non-square mesh aspect ratio in graph creation

67789e4

kshirajahere reviewed Mar 12, 2026

View reviewed changes

kshirajahere suggested changes Mar 12, 2026

View reviewed changes

sadamov requested a review from joeloskarsson March 12, 2026 20:29

sadamov added the enhancement New feature or request label Mar 12, 2026

osten-antonio added 2 commits March 13, 2026 16:17

fix: derive coarse mesh dims from prev slice and preserve square grid dm

a15deaf

test: add test for multiscale non-square mesh creation and square gri…

6943ca3

…d edge preservation

osten-antonio requested a review from kshirajahere March 13, 2026 09:40

Merge branch 'main' into fix/non-square-mesh-aspect-ratio

b460037

sadamov mentioned this pull request Mar 16, 2026

create_graph always generates square mesh grids, ignoring data grid aspect ratio #371

Open

Merge branch 'main' into fix/non-square-mesh-aspect-ratio

073eb21

This comment was marked as duplicate.

Sign in to view

kshirajahere suggested changes Mar 22, 2026

View reviewed changes

sadamov closed this Apr 28, 2026

		assert r.shape[1] == d_features


		def test_graph_creation_non_square_aspect_ratio():

Conversation

osten-antonio commented Mar 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Describe your changes

Issue Link

Type of change

Checklist before requesting a review

Checklist for reviewers

Author checklist after completed review

Checklist for assignee

Uh oh!

kshirajahere Mar 12, 2026

Choose a reason for hiding this comment

Uh oh!

kshirajahere Mar 12, 2026

Choose a reason for hiding this comment

Uh oh!

kshirajahere Mar 12, 2026

Choose a reason for hiding this comment

Uh oh!

kshirajahere left a comment

Choose a reason for hiding this comment

Uh oh!

osten-antonio commented Mar 13, 2026

Uh oh!

joeloskarsson commented Mar 22, 2026

Uh oh!

This comment was marked as duplicate.

Uh oh!

kshirajahere left a comment

Choose a reason for hiding this comment

Uh oh!

kshirajahere Mar 22, 2026

Choose a reason for hiding this comment

Uh oh!

kshirajahere Mar 22, 2026

Choose a reason for hiding this comment

Uh oh!

kshirajahere commented Mar 22, 2026

Uh oh!

osten-antonio commented Mar 24, 2026

Uh oh!

sadamov commented Apr 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

osten-antonio commented Mar 10, 2026 •

edited

Loading