Add adaptive meshing

[jacobbieker]

Detailed Description

Being able to get more detailed output around areas of interest is quite important, such as around the UK or some other country of interest.

Context

Possible Implementation

Uses the adaptive mesh from the original DeepMind paper, and train on 0.25 degree GFS data downsampled in some areas to 1 degree. Or something like that

[AswaniSahoo]

Hi @jacobbieker,

Went through all the DMI regional NWP datasets on Source Cooperative and dug into the codebase to figure out what we'd need to change. Here's where I'm at.

The data

There are 14 datasets across 4 model families. The coverage is actually quite diverse geographically — HARMONIE over Greenland/Iceland, MEPS over Scandinavia, AROME over France and five overseas territories (Caribbean, French Guiana, Indian Ocean, New Caledonia, Polynesia), plus Alaska HRRR and Hawaii NAM.

The grid formats vary a lot: HARMONIE and Alaska use 2D rotated grids, MEPS is Lambert conformal, AROME uses regular lat/lon, and Hawaii NAM is flattened to 1D. This is actually a good thing for graph_weather since the model already takes arbitrary lat/lon coordinates, so the grid projection differences don't matter.

The tradeoff across datasets is between temporal depth and variable richness. MEPS postprocessed has the longest time series (2.4 years, 21k timesteps) but only 18 surface variables. The AROME overseas domains have 68 variables on 19 pressure levels plus 12 height levels, but only 40 days of data. MEPS raw sits in the middle with 14 months and 223 variables. No single dataset gives us everything, so training will probably need to pull from multiple datasets.

I can put together a more detailed write-up of the variable lists and grid dimensions per dataset if that would be useful.

Architecture

Your design to keep everything accepting arbitrary lat/lon inputs is what makes this feasible — the main extension is making the bipartite graph construction dynamic so it rebuilds when the regional coordinates change between batches.

You mentioned the adaptive mesh from the original DeepMind paper in this issue. That approach (refining the H3 latent mesh to higher resolution over the regional area) would give the cleanest physics since information propagates naturally across resolution boundaries. The main cost is rebuilding a variable-resolution graph every time the region moves.

An alternative starting point: keep the global model on its static res 2 mesh and run a second GNN on the regional grid that takes boundary conditions from the global forecast through a relaxation zone at the edges. Less elegant but gets us to something testable faster, and the data pipeline is simpler when we're sampling different spatial crops from different datasets each step. We can always move toward the unified adaptive mesh once we know where the actual bottlenecks are.

Two questions that would help me figure out where to start:

1. Do you have a preference between unified adaptive mesh vs separate global+regional models? Or should I just prototype the simpler one first?
2. Should we plan for two-way coupling (regional feeding back into global) or is one-way nesting enough for now?

Happy to coordinate here or on a branch.

[jacobbieker]

For 1. I prefer a unified adaptive mesh, but prototyping with simpler one first sounds good.
2. It should be 2-way, hence the unified adaptive mesh might work better longer-term. But whichever you want to start with sounds good!

Also to note, for the datasets, there are some other collections I've built that might be relevant:

1. https://source.coop/bkr/geos is GEOS aerosol model, but only the 15 minutely data. Its 0.25 resolution, and global though, so for trying to learn sub-hourly information that could be good. There is a few years of data there.
2. https://source.coop/bkr/icon ICON global and EU model. As well as raw files from some of their newer models. There is also the Dynamical ICON-EU Icechunk/Zarr that could be helpful/useful.
3. https://source.coop/bkr/ifs IFS Analysis going back to 2016, 6 hourly, but could be nice global model to use, the 0.25 degree resampled version is still being built so is behind by quite a bit, same with the 1 degree version, but the native resolution is ready to go.
4. https://source.coop/bkr/metoffice MetOffice 10km and 2km global and regional models as well.
5. https://source.coop/bkr/silam-dust 0.1 degree global dust forecasts, more similar to the GEOS data, but also a potentially interesting one.

Not saying you need to use all or any of the above! But they do provide different resolutions, coverage, and such to complement the other datasets you mentioned.

[AswaniSahoo]

Thanks Jacob. Will start with the simpler dual-graph prototype and keep the design compatible with migrating to the unified adaptive mesh later.

Thanks for the additional dataset links. IFS Analysis and ICON EU look particularly useful. I'll explore those and the rest over the coming days.

Starting with a codebase deep-dive into the encoder/decoder graph construction to figure out exactly where to make things dynamic. Will push to a branch once I have something worth looking at.

[AswaniSahoo]

Thanks for confirming the unified adaptive mesh and two-way coupling direction. I'll prototype with the simpler dual-graph approach first to get something running, designed so we can migrate to the unified mesh without rewriting everything.

Went through the datasets you shared. The IFS native res store (hres_analysis.icechunk) is the big find - 10 years of 6-hourly global data at ~9km, 25 pressure levels, 43 variables. That covers global pre-training on its own. For regional targets, ICON D2 gives us 2km hourly over Germany and MetOffice UK is 2km with 33 vertical levels. The GEOS 15-minute aerosol data at 0.25 deg is interesting for sub-hourly work down the line.

Also opened #218 - AssimilatorEncoder was using deprecated H3 v3 calls that crash under the pinned h3==4.3.1. Six name substitutions plus a numpy 2.x fix in tests.

Next I'm working on making the encoder's bipartite graph construction dynamic so it rebuilds per-batch when the regional coordinates change. That's the core piece needed before we can feed in any of these regional datasets.