feat: add geometry utilities for global domains (spherical coords, area weights, KNN graph)

[ayushhbasu]

Description of changes

This PR introduces foundational geometry utilities to support global forecasting capabilities. It adds spherical coordinate conversion, latitude-based area weights, and a modular KNN graph builder for global domains.

Changes:

• Adds neural_lam/geometry.py:
  • lat_lon_to_cartesian(): Converts (lon, lat) to (x, y, z) on a unit sphere
  • get_area_weights(): Computes cos(latitude) weights normalized by the mean
• Adds neural_lam/graph_builder.py:
  • build_spherical_knn_graph(): Backend-agnostic KNN graph construction on a sphere using Cartesian coordinates
• Extends BaseDatastore in neural_lam/datastores/base.py:
  • get_area_weights(): Returns a DataArray of latitude-based weights
  • get_cartesian_coords(): Returns an array of (x, y, z) coordinates

Motivation:
As part of global forecasting support (#63), reusable geometry utilities are required for area-weighted metrics and spherical graph construction. This PR establishes a modular geometry layer that future work can build upon.

Dependencies:
This PR is designed to be standalone but will benefit from the GlobalDummyDatastore introduced in #444 for testing. Tests for global-specific functionality will be added after #444 is merged.

Issue Link

Part of #63 (global forecasting support)

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 (additions or improvements to documentation)

Checklist before requesting a review

• My branch is up to date with the target branch (use pull --rebase if possible)
• I have performed a self-review of my code
• I have added docstrings for all new/modified functions and classes, clearly describing purpose, inputs, and outputs
• I have added inline comments where necessary to clarify complex logic
• I have updated the README to reflect these changes
• I have added tests to validate the new functionality
• The PR title clearly describes the change and is written in imperative form
• I have requested a reviewer and assigned the PR (if applicable)

[sadamov]

thanks @ayushhbasu I am adding @joeloskarsson here for an initial review. I see some potential overlap with our plan to move to weather-model-graphs for graph creation and coords handling. But he knows the roadmap for that better than I do. Maybe the changes here on the model side will still be required.

[princekumarlahon]

That makes sense, especially if graph creation is moving towards weather-model-graphs. Given that, do you think it would be useful to keep the model side a bit more decoupled from how graphs are constructed? For example, having a small interface/adapter so the model can work with graphs regardless of whether they come from this geometry layer or an external source. This could help keep things flexible as the graph creation workflow evolves. Happy to help explore this direction if it aligns with the roadmap.

[sadamov]

@princekumarlahon It sounds like you are talking about #384 and yes I believe that does make sense.

[ayushhbasu]

@sadamov #473 (comment)

Thank you @sadamov — I appreciate you bringing in @joeloskarsson for roadmap clarity.

I understand the concern about overlap with weather-model-graphs. This PR currently adds:

• geometry.py – lat_lon_to_cartesian() and get_area_weights() (general spherical utilities)
• datastores/base.py – get_area_weights() and get_cartesian_coords() (model-facing accessors for metrics and coordinate access)
• graph_builder.py – build_spherical_knn_graph() (potential overlap point)

The geometry and datastore changes are independent of where graph construction ultimately lives, and are required for:

• area-weighted metrics
• global coordinate handling
• boundary-free rollout

These are foundational pieces for global forecasting and will be reused across the project.

Given this, I'm happy to scope this PR down to geometry utilities and datastore methods, and remove build_spherical_knn_graph() for now if that aligns with the roadmap. The functionality would still be needed later for global mesh construction, but I'm flexible on where it should live architecturally.

For context, I've submitted a GSoC proposal (Mar 24) aligned with the Global Weather Forecasting direction, where this PR serves as the initial foundational step. I'll fully align the implementation here with the project's roadmap.

Would you recommend removing build_spherical_knn_graph() from this PR for now and deferring graph generation to weather-model-graphs?

Looking forward to your guidance, @joeloskarsson.

[arzoo0511]

Hi @ayushhbasu ,

I’ve been working on a similar direction for global forecasting support and came across this PR by @sadamov this aligns very closely with what I had in mind.

I particularly like the separation between geometry utilities and datastore-level accessors. I was also exploring a similar approach for spherical coordinate handling and modular graph construction.

Regarding the discussion on graph generation potentially moving to weather-model-graphs, I agree that keeping the geometry layer independent while deferring graph construction might be a cleaner design.

I’d be happy to contribute here instead of duplicating pr i had already raised similar pr(#561 (comment)) for example:

• helping refine the geometry utilities
• adding tests (once Add GlobalDummyDatastore for global domain testing #444 is merged)
• or helping decouple graph construction further if needed

Let me know where contributions would be most useful.

Thanks!

[Prince637-boo]

Hi @ayushhbasu, great work on these geometry utilities! I've been testing the spherical coordinate conversion logic alongside my work on weather-model-graphs with @yuvraajnarula.

Regarding the discussion with @joeloskarsson about moving graph construction to WMG: I've already implemented a 3D Plotly visualizer in WMG that handles these Cartesian (x,y,z) coordinates perfectly. I'd be happy to help verify that the graphs generated via your geometry.py utilities are fully compatible with the hierarchical visualization needs we have.

Looking forward to seeing how we can bridge this with WMG!

[ayushhbasu]

Thanks @arzoo0511 and @Prince637-boo - really appreciate the alignment and the thoughtful inputs here.

@arzoo0511 — great to hear this overlaps with your work. I’d definitely welcome collaboration. The most useful areas to contribute (especially after #444 is merged) would be:

• adding tests for geometry utilities
• validating area-weighted metrics and coordinate handling
• refining datastore interfaces if needed

On scoping: I’m planning to keep this PR focused on geometry + datastore utilities, and defer graph construction to weather-model-graphs to stay aligned with the roadmap. If helpful, further decoupling between geometry and graph construction can also be explored incrementally.

@Prince637-boo — very helpful context. Great to know WMG already supports Cartesian (x, y, z) and visualization. Once direction is confirmed, it would be valuable to validate that lat_lon_to_cartesian() outputs integrate cleanly with the WMG pipeline, and collaborate on bridging this with WMG.

More broadly, I see:

• neural-lam → consuming geometry + graph structures
• weather-model-graphs → owning graph construction

For context, this PR corresponds to the foundational geometry + datastore layer from my GSoC proposal, so I’m intentionally keeping it minimal and aligned with that direction.

I’ll wait for guidance from @joeloskarsson on whether to remove build_spherical_knn_graph() and keep this PR strictly scoped to core utilities.

Looking forward to collaborating

[Prince637-boo]

Thanks @ayushhbasu for the clear direction. I completely agree with the decoupling: keeping core geometry in Neural-LAM and graph construction in WMG makes total sense.
To support this, I have just submitted a PR in weather-model-graphs #132 introducing a Graph Validator. It specifically checks for 3D geometric consistency (L2 norm) and connectivity health. This should help ensure that the geometry utilities from your PR integrate perfectly with the WMG pipeline.
Looking forward to bridging these components!

[ayushhbasu]

Thanks @Prince637-boo — great to see this.

The Graph Validator seems like a perfect fit for validating outputs from lat_lon_to_cartesian(), especially for geometric consistency and connectivity.

Once direction is confirmed, we can validate that these coordinates pass the WMG checks and align on conventions (e.g., normalization, ordering) across both sides.

This should help cleanly bridge neural-lam with WMG.

Looking forward to collaborating

[arzoo0511]

Hi @ayushhbasu , @Prince637-boo ,

This direction makes sense to me , keeping geometry utilities in neural-lam and delegating graph construction + validation to WMG feels like the right separation.

I can focus on the following on the neural-lam side once #444 is merge @joeloskarsson :
adding unit + integration tests for geometry utilities (especially spherical ↔ Cartesian conversions)
validating area-weighted metrics and coordinate consistency
ensuring outputs from lat_lon_to_cartesian() align with WMG expectations (normalization, ordering, shape)

@Prince637-boo — the Graph Validator in WMG sounds very useful. Once things are wired up, I can help run geometry outputs through it and flag any inconsistencies early.
If helpful, I can also draft a small interface doc defining expected coordinate formats between neural-lam and WMG to avoid mismatches.

Let me know if this works
Thanks!

[Prince637-boo]

Thanks @arzoo0511! That sounds like a perfect plan. I've already implemented the Graph Validator in WMG (see PR #132: mllam/weather-model-graphs#132).

It’s specifically designed to handle the geometric consistency (L2 norm) and connectivity checks we need. I’d be happy to collaborate on the interface documentation to ensure that the outputs from lat_lon_to_cartesian() align perfectly with these WMG requirements. Let’s sync once the core geometry utilities are ready!

[sadamov]

Closing this in favour of #650, which consolidates the global forecasting / spherical domain track for v0.8.0. The geometry.py (spherical coords, cos-lat weights) and graph_builder.py (KNN graph on sphere) utilities will be folded into the consolidated PR. Thanks for the contribution @ayushhbasu, you'll be credited.