feat: add icosahedral mesh layout with grid-to-mesh connectivity

[yuvraajnarula]

Describe your changes

This PR adds mesh_layout='icosahedral' to support global graph generation in WMG, addressing the architectural gap identified in #71.

Key additions:

• create.mesh.layouts.icosahedral.py - Implements icosahedral mesh hierarchy using trimesh (co-developed with @mandeepsingh2007)
• Grid-to-mesh (g2m) radius-based connectivity adapted from neural-lam's prob_model_global branch (lines 224-242)
• Notebook demo in /docs/icosahedral_global_graph.ipynb showing end-to-end pipeline
• Helper utilities for spherical coordinate conversion (lat/lon ↔ cartesian)

Motivation:
Currently, WMG's mesh generation assumes rectilinear grids, which breaks for global domains (Issue #71). This PR decouples mesh layout from connectivity by introducing:

• mesh_layout='icosahedral' for quasi-uniform spherical sampling
• g2m connections using haversine-ready distance calculations
• Foundation for future spherical layouts (HEALPix, etc.)

Dependencies:

• trimesh - Lightweight mesh generation (avoiding GraphCast dependency)
• numpy, scipy (already in WMG dependencies)

Issue Link

#71

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 documentation 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)

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

Checklist for assignee

• PR is up to date with the base branch
• the tests pass
• author has added an entry to the changelog (and designated the change as added, changed or fixed)
• Once the PR is ready to be merged, squash commits and merge the PR.

[mandeepsingh2007]

@yuvraajnarula Since we are collaborating on this PR, could you please invite me as a Collaborator to your fork of weather-model-graphs?

This way, I can directly commit and push the trimesh icosahedral mesh generation script to the PR branch myself. It will make it much easier for both of us to iterate on the code without having to send files back and forth!

[yuvraajnarula]

@Raj-Taware Great point about the BallTree dependency. The 3D projection + KDTree approach is cleaner and matches what we're already doing for g2m connections.

One question:
The len attribute calculation you mentioned (2*arcsin(euclid/2)) - does this work consistently for points on the unit sphere?

[Raj-Taware]

@Raj-Taware Great point about the BallTree dependency. The 3D projection + KDTree approach is cleaner and matches what we're already doing for g2m connections.

One question: The len attribute calculation you mentioned (2*arcsin(euclid/2)) - does this work consistently for points on the unit sphere?

Hey ! great work getting trimesh up and running

Yes the formula works as long as both the coordinates are projected on a unit sphere.

If you want to generalise it the formula becomes
d = 2*R*arcsin(C/2R)

But since in create_icosahedral_mesh function you are normalising the vertices with norms, you are already creating a unit vector. and generate_icosahedral_mesh also sets the radius to 1.0 so I think the original formula can be used.

Btw when using arcsin i would recommend clipping the input to handle floating-point rounding errors. Sometimes KDTree might return a distance as something like 2.00000000000004 which will case np.arcsin(2.00000000000004/2) = np.arcsin(1.00000000000002) to throw a NaN warning.

Hope this helps !

[yuvraajnarula]

Tracking

Note : Core icosahedral work now moved to dedicated issue #77. This PR focuses on connectivity primitives that will feed into that effort.

[yuvraajnarula]

@mandeepsingh2007 Now that the PR is taking shape, wanted to check in – how does the implementation look from your end?

Specifically curious if:

• The mesh generation integration works smoothly with your trimesh code
• Any adjustments you'd recommend on my end (g2m, dispatch logic, etc.)
• Anything you think we should add/change before marking ready for review

Your mesh work is the foundation here, so your feedback matters. Thanks again for collaborating on this!

[mandeepsingh2007]

@yuvraajnarula thanks for putting all this together! The implementation is definitely taking great shape.

To answer your points:

Mesh Integration: The trimesh base generation looks smooth and is establishing the correct spherical topology as expected.

g2m & Dispatch Logic: The current logic works perfectly for establishing our baseline. However, keeping Joel's recent Slack note about networkx bottlenecks in mind, we should note that scaling this to high-resolution global grids will eventually require us to shift toward fully vectorized bulk operations (to avoid iterative graph-building overhead). But for the scope of this initial PR, your current setup is spot on.

Ready for review? I think it’s good to go! Keeping this PR tightly scoped to just establishing the lightweight icosahedral layout (without GraphCast dependencies) is the perfect first step. We can tackle the hierarchical scaling and tensor optimizations in subsequent issues.

[Joltsy10]

One thing worth flagging here. find_containing_triangle correctly computes and returns the barycentric weights, but they don't seem to be making it onto the edges in connect_nodes_across_graphs:

G_connect.edges[source_node, target_node]["vdiff"] = np.array([dlat, dlon])
# weights returned by find_containing_triangle not stored here

Storing them as a weights edge attribute would let to_pyg include them in the edge features alongside len and vdiff, which gives the model a geometric prior on each mesh node's contribution to a grid point rather than having to learn it purely from data. That's kind of the nice thing about containing_triangle over nearest_neighbours so would be good to preserve it.

G_connect.edges[source_node, target_node]["weights"] = barycentric_weight_for_this_node

That said, we could skip the weights entirely since the current LAM model operates on a learned weight basis anyway. But if we are going through the effort of computing them it seems like a shame not to pass them through to the model as well.
Might just be an oversight, wanted to flag it!

[yuvraajnarula]

@Joltsy10 Good flag
I'll add them alongside len and vdiff in the containing_triangle block. The model can then choose to use them or not, but at least the information isn't lost.

Appreciate you catching these edge cases.

[leifdenby]

Great idea to add these weights (maybe as barycentric_weight) as a node attribute. If you have done that @yuvraajnarula should we close this conversation?

[yuvraajnarula]

@mandeepsingh2007 @Joltsy10

Pushed two commits:

1. refactor: optimize mesh-to-grid connectivity with spatial indexing
2. chore: add early exit + optimized function in containing_triangle case

@Joltsy10 – your review comments on performance and weights are addressed. The containing_triangle path now uses spatial indexing and stores barycentric weights as edge attributes.

@mandeepsingh2007 – mesh integration still solid? The optimizations shouldn't affect your generation logic.

Next: Generic methods (nearest_neighbour, etc.) could also use 3D spatial indexing for icosahedral cases. I'm going to keep this PR focused on containing_triangle since that's the critical path. We can optimize the generic methods in a follow-up PR if needed.

Thoughts?

[mandeepsingh2007]

@yuvraajnarula Yes, the mesh integration is completely solid. The spatial indexing you added for the containing_triangle path sits downstream from the coordinate generation, so the underlying trimesh geometry remains perfectly intact and mathematically sound.

I also 100% agree with your approach for the next steps. Keeping this PR strictly scoped to the containing_triangle critical path and the baseline icosahedral layout is the smartest move. It prevents scope creep and makes this much easier for the mentors to review. We should definitely freeze the scope here, get this foundation merged, and tackle the generic methods in subsequent PRs.

This looks fantastic. From my end on the mesh side, this is ready to go!

[yuvraajnarula]

@leifdenby
Reverting the visualization changes from this PR to keep #76 focused on core icosahedral graph generation.

I'm opening a new issue and will track 3D visualization work there. Following the discussion in #20, the plan is:

• Add plot_3d.render_with_plotly(graph) in a separate PR
• Use level attribute for Z-axis (grid nodes default to -1)
• Batch edges with None separators for performance
• Plotly as optional dependency

Happy to iterate on the API there. Thanks for the guidance!

[Prince637-boo]

Hi @yuvraajnarula @Raj-Taware @mandeepsingh2007, great progress on the icosahedral implementation. After reviewing the latest commits, I have a few concerns regarding the global consistency:

1.Coordinate System: While Lat/Lon is 'pragmatic', using 3D Cartesian coordinates for mesh_features would avoid the pole singularity issue which is critical for GNN stability in global models.

2.Efficiency: Have we considered using torch-cluster for the spatial indexing instead of scipy? This would keep the entire pipeline on the GPU and avoid costly CPU-GPU transfers during graph construction.

3.Hierarchical Links: Does the current mesh_layout logic pre-compute the edges between different icosahedral levels? This is vital for implementing multi-scale message passing (pooling/unpooling) later on.

[Joltsy10]

@Prince637-boo

1. We did discuss on the usage of 3d coordinate over lat/lon but to follow the validator spec that is currently going on in neural-lam we ended on using lat/lon for wmg side and then applying sin/cos embedding over on the neural-lam model side.
2. The graph generation happens only once here in wmg, then the files are sent over to neural lam for training/inference. I assume you think we will be generating the graph during training but it is not so, the graph is generated once here and the only thing that changes during training are the params. We will not run into this cpu/gpu handoff.
3. yes the edges are precomputed, you can look at the visualization in feat: add 3D interactive graph visualisation with Plotly #118 to get a better idea.

[Prince637-boo]

Thanks for the clarification! It makes sense to keep Lat/Lon for compatibility with the validator spec, especially if the sin/cos embedding handles the continuity on the model side. Also, I see your point about the static nature of the graph; since it's a one-time generation, the CPU/GPU handoff isn't a bottleneck. I'll take a closer look at the hierarchy in #118 to understand how the multi-scale edges are structured.

[yuvraajnarula]

@Prince637-boo
Thanks for the thoughtful questions!

To add to what @Joltsy10 already clarified:

1. Coordinate system : We opted for lat/lon in the graph to stay validator‑compliant. The neural‑lam side can apply sin/cos embeddings (as in the prob_model_global branch), which handles the polar continuity nicely.

2. Efficiency : Since graph generation is a one‑time offline step, scipy’s KDTree is sufficient for now. Keeping it CPU‑side avoids pulling in PyTorch GPU dependencies into WMG. If we ever need to generate graphs repeatedly (e.g., dynamic graphs), torch‑cluster would be a good candidate, definitely worth keeping in mind for the future.

3. Hierarchical links : Yes, the edges between levels are pre‑computed. The hierarchical layout in #118 visualises them clearly (intra‑level in blue, up/down in red/green). For multi‑scale message passing, the structure is there to use directly.

Happy to answer any more questions!

[Prince637-boo]

Thanks @yuvraajnarula and @Joltsy10 for the detailed

1.Coordinate system: Understood. Using sin/cos embeddings on the model side is a smart way to maintain validator compliance while solving the polar discontinuity.

2.Efficiency: That makes sense. Since the graph is static and pre-computed, the CPU overhead isn't a bottleneck for training. I agree that keeping WMG lightweight without forcing GPU dependencies is the right call for now.

3.Hierarchical links: I’ve been looking at the visualization in #118, and the concentric layout is a great improvement for clarity. It’s much easier to trace the up/down edges now.

I'm currently diving deeper into the plot_3d.py logic. I'd love to help polish the global visualization further—perhaps by adding a coastline/map layer to give the icosahedral grid more geographical context. Looking forward to contributing!

[prajwal-tech07]

Hi @yuvraajnarula @mandeepsingh2007 @Joltsy10 - Leif asked me to take a look at this PR, so here's my read after going through the code and the discussion history. First off, this is really impressive work. The tangential-plane vdiff keeping 3 edge features in line with the graph-storage spec, the barycentric weights on m2g edges, and the centroid-KDTree optimization are all the right calls, and the notebooks are some of the nicest docs in the repo.

I did find a few things in the connect_nodes_across_graphs changes that I think need attention, because they affect existing (non-icosahedral) graphs too:

1. The ±360° longitude wrapping is applied unconditionally. In the generic path (neither graph icosahedral), the source KDTree is built from 3 copies of the points shifted by −360/0/+360, and vdiff/len get wrapped with (dlon + 180) % 360 - 180. But pos for existing rectilinear graphs is projected x/y in the CRS units - often meters. For e.g. a LAEA grid, a 5000 m displacement wraps to −40, silently corrupting the edge features, and the replicated KDTree adds phantom neighbours 360 m away. I think the wrapping needs to be gated on the coordinates actually being geographic lat/lon rather than applied to every graph.

2. Coordinate order. The test fixtures pass coords as [lat, lon], but everywhere else in wmg coords is xy - which for a geographic CRS with always_xy=True means [lon, lat]. The icosahedral code reads pos[0] as latitude, so one of the two conventions is silently wrong depending on how the user passes coords. Worth pinning down and asserting explicitly.

3. Which level does the grid connect to? split_graph_by_edge_attribute(m2m, "level")[0] looks up the subgraph with level == 0. For the rectilinear hierarchy level 0 is the finest (bottom) level, but in the icosahedral hierarchy level 0 is the coarsest (the 12-node base icosahedron) - so if I'm reading it right, g2m connects the grid to the coarsest level, while containing_triangle m2g uses the finest vertices. I think both should use the finest level.

4.rel_max_dist semantics changed. On main the longest-edge search only considers bottom-level edges; the new code scans edges at all scales, so for flat_multiscale graphs the coarsest edge sets the radius and g2m over-connects compared to before. If that wasn't intentional, the bottom-level-only logic should be restored (the split_on_edge_attribute_existance import got dropped in this PR).

One last thought on the rebase, since the branch now conflicts with main: #81 already added the mesh_layout argument and #92 is moving the mesh code to a mesh/layout/ subpackage where layout modules create node coordinates and connectivity modules build the edges. I think after rebasing, most of the icosahedral-specific dispatch in base.py could slot into that same split, which would also shrink the special-casing in connect_nodes_across_graphs.

None of this takes away from the core contribution - the mesh generation, containment logic and spec-aligned edge features are solid, and this is the natural next mesh layout for wmg.

[Joltsy10]

Thanks for the review, I was wanting to finish this PR off since quite a while so this is a good way to do it. All four points seem valid to me.

1. Agreed, the wrapping and the replicated KDTree assume geographic degrees and corrupt CRSs. I'll gate all the wrapping + KDTree replication behind an actual geographic-coordinate check.
2. Yeah was an oversight on our part, I'll standardize to the WMG xy convention and correct all the test fixtures.
3. That is a really good catch thanks, I'll change that to explicitly take the finest mesh instead of the level 0 hardcode
4. Agreed, ill restore the bottom-level-only longest-edge scan so multiscale graphs dont over-connect.

One heads-up: this branch predates #81's two-step coords/connectivity refactor, so I'm going to rebase the icosahedral work onto current main first (it still imports the now-removed mesh/kinds/* modules). The fixes above will land on top of that.

[yuvraajnarula]

Thanks @Joltsy10 and @prajwal-tech07. Totally agree with all four points and gating the wrapping behind a geographic check and keeping the multiscale graphs from over-connecting are definitely the right moves.

The plan to rebase onto main first to handle the #81 refactor makes perfect sense. Let me know once the rebase is in, and I can jump in to help implement these fixes or update the test fixtures for the WMG xy convention!

[Joltsy10]

thanks @yuvraajnarula! quick update, the rebase is in and i went ahead and did the first two while i was in there.

rebase onto main (#81) is done, the branch sits on top of the two step coords/connectivity refactor now and the old mesh/kinds imports are gone. one merge commit, 145 tests green.

for 1 i gated the wrapping and kdtree replication behind a geographic crs check so projected coords are left alone (checked LAEA in meters stays unwrapped and 4326 across the antimeridian still wraps). for 2 i standardised everything on the wmg xy = [lon, lat] convention, mesh nodes store [lon, lat] now, and i fixed the fixtures that were passing [lat, lon].

that leaves 3 and 4, both in connect_nodes_across_graphs / the icosahedral branch of base.py. theyre all yours if you want them, feel free to take whatever approach you see fit. ping me if anything from the rebase looks off, happy to pair.