Skip to content

feat: add variable-resolution latent graph builder#227

Merged
jacobbieker merged 1 commit into
openclimatefix:mainfrom
AswaniSahoo:feat/stretched-latent-graph
Jul 5, 2026
Merged

feat: add variable-resolution latent graph builder#227
jacobbieker merged 1 commit into
openclimatefix:mainfrom
AswaniSahoo:feat/stretched-latent-graph

Conversation

@AswaniSahoo

Copy link
Copy Markdown
Contributor

Pull Request

Description

Builds on #226. build_variable_resolution_mesh returns a mixed-resolution cell set, coarse over the globe and fine over the region, and this PR turns that set into the latent graph the processor message-passes over.

The problem is the seam between the two resolutions. Same-resolution neighbours come from grid_disk, but that never crosses from a coarse cell to a fine one, so on its own the fine region is a disconnected island and no information flows in or out. Each fine cell is therefore also linked to the coarse cells bordering its coarse ancestor, found with cell_to_parent and then grid_disk on that parent, keeping the coarse cells that are actually present in the mesh. Edges are bidirectional and carry [sin d, cos d] of the great-circle distance between cell centres, the same edge feature the single-resolution builder already uses.

Pure function: it takes a cell list and returns a torch_geometric Data. No model or training changes yet - wiring it into the forecaster and resizing the embedding table come in later PRs.

Related to #3

How Has This Been Tested?

Added tests/test_stretched_latent_graph.py - plain pytest, runs entirely on h3, no network:

  • at least one edge joins a coarse cell to a fine cell, so the seam is bridged
  • union-find over every edge gives a single connected component, so the fine region is not an island
  • two adjacent fine interior cells still share an edge (same-resolution case intact)
  • no self-loops
  • no duplicate edges
  • edge indices reference exactly the supplied cells
  • edge_attr is [E, 2] with each row on the unit circle

Commands:

  • pytest tests/test_stretched_latent_graph.py -q -> 7 passed
  • ruff check graph_weather/models/layers/stretched_latent_graph.py tests/test_stretched_latent_graph.py -> clean

Checklist:

  • My code follows OCF's coding style guidelines
  • I have performed a self-review of my own code
  • I have made corresponding changes to the documentation
  • I have added tests that prove my fix is effective or that my feature works
  • I have checked my code and corrected any misspellings

@jacobbieker
jacobbieker merged commit cffd0a4 into openclimatefix:main Jul 5, 2026
0 of 3 checks passed
@AswaniSahoo
AswaniSahoo deleted the feat/stretched-latent-graph branch July 5, 2026 19:02
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.

2 participants