Skip to content

feat: add vectorized 3D cartesian utils and polar pytests#107

Open
Riteesh-NITT wants to merge 1 commit intomllam:mainfrom
Riteesh-NITT:feature/spherical-utils
Open

feat: add vectorized 3D cartesian utils and polar pytests#107
Riteesh-NITT wants to merge 1 commit intomllam:mainfrom
Riteesh-NITT:feature/spherical-utils

Conversation

@Riteesh-NITT
Copy link
Copy Markdown

Describe your changes

Summary of the changes:
This PR introduces vectorized PyTorch utilities (spherical_utils.py) to convert 2D geographic coordinates (latitude/longitude) into 3D Cartesian coordinates $(x, y, z)$. It also includes a pytest suite specifically validating mathematical stability at the North Pole singularity and across the anti-meridian (Date Line).

Note: Hi @leifdenby and @mandeepsingh2007, following up on our thread in #76/#77! I know we are waiting on the load_graph shim (Issue #385) to be merged before fully hooking this up, but I wanted to get these utilities up for early review. Let me know if the tensor shapes align with the native PyG pipeline you are building. I'll keep this as a Draft until the base branch is ready!

Motivation and context:
When mapping 2D weather grid data to a 3D icosahedral mesh, flat lat-lon grids suffer from severe spatial distortion at the poles. This utility provides the mathematical bridge to project grid nodes onto a unit sphere in 3D space, which is an architectural prerequisite for accurate Grid-to-Mesh (g2m) distance calculations and connectivity.

Dependencies:

  • torch
  • pytest (for the test suite)

Issue Link

Addresses #76 and #77

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)

@Riteesh-NITT
Copy link
Copy Markdown
Author

Hi @leifdenby and @joeloskarsson, just putting this up as a Draft! This implements the vectorized spherical-to-cartesian conversions and the polar edge-case tests we discussed in #76 and #77. Hey @mandeepsingh2007, here is the math bridge for the icosahedral mesh generation. Feel free to pull this branch locally if you need to test the grid-to-mesh (g2m) mappings with your trimesh output!

@mandeepsingh2007
Copy link
Copy Markdown

Hey @Riteesh-NITT , I just pulled the logic locally and ran some integration tests against the current NumPy g2m geometry we built in PR #76.

The outputs from your torch implementation match the current numpy pipeline outputs down to the exact floating-point precision (max difference ~1e-8).

This confirms that your PyTorch vectorization is a perfect drop-in replacement. Once we move to the HeteroData native graph loading, we can safely swap the NumPy processing out for this without breaking any downstream spatial alignments.

image

@leifdenby
Copy link
Copy Markdown
Member

since this functionality is a utility being added for icosahedral meshes, which in turn was decide in today's dev meeting has gone on the roadmap for v0.5.0, I will put this utility on the roadmap for v0.5.0 as well. Although I would prefer to just have this addition made with the functionality to use it rather than having a util that isn't used yet

@leifdenby leifdenby added this to the v0.5.0 milestone Mar 23, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants