Add adaptive meshing

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

@jacobbieker

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:

Do you have a preference between unified adaptive mesh vs separate global+regional models? Or should I just prototype the simpler one first?
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.

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:

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.
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.
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.
https://source.coop/bkr/metoffice MetOffice 10km and 2km global and regional models as well.
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.

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.

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.

@jacobbieker

Hi @jacobbieker , Following up on the RegionalForecaster in PR #221 , the next piece is boundary handling.

Hereis my plan for a BoundaryNudgingLayer:

Approach: Hybrid relaxation zone : distance-based geometric prior (centroid distance, alpha=0 at interior, alpha=1 at edges) with a small MLP learning a residual correction on top. Standard NWP relaxation concept, similar spirit to what neural-lam does with boundary_mask, but adapted for our dynamic/movable mesh.
Placement: After decode + residual, as the last operation. Neural-lam does boundary overwrite at the same point in its autoregressive loop.
Optional: Controlled by enable_nudging config flag. Default off, so the core model is unaffected.

The global context would come from a global GraphWeatherForecaster's predictions at the same observation points , the caller passes it in as global_context tensor.

Does this direction make sense for Phase 1, or would you prefer a different boundary handling approach before I implement it?

Hey, thanks for these. That approach seems fine to me

@jacobbieker

hi @jacobbieker , Small update on the regional side. Now that #223 is up, I've started on the data loader.

The approach I'm leaning toward for the RegionalDataset: sample movable sub-grids straight from the IFS hres_analysis : random bounding box, random timestep, pull that high-res crop. The coarse global field is right there in the same store too, so I can wire it in for boundary context later. Main reason I like starting here is it's one source, so the variables stay consistent and the box can sit anywhere on the globe, which is kind of the whole point of "movable". MetOffice UK 2km and ICON-D2 I'd add afterwards as a second source.

Before I go build it out, does that line up with how you were thinking about the data side? Or would you rather I start from one of the regional NWP datasets directly?

No, I think that is a perfectly fine way of starting out with this!

@jacobbieker

Hi @jacobbieker ,

Phase 1 of the movable regional prototype is pretty much done now:

DynamicGraphBuilder (feat: add DynamicGraphBuilder for runtime bipartite graph construction #219)
RegionalForecaster (feat: add RegionalForecaster for movable regional weather prediction #221)
BoundaryNudgingLayer (feat: add BoundaryNudgingLayer for regional-global blending #223)
RegionalDataset (feat: add RegionalDataset for movable IFS sub-grid sampling #224)
the coarse global_context that feeds the nudging layer is the last bit, it's up in feat: add coarse global_context to RegionalDataset for boundary nudging #225.

One thing I want to flag honestly. I wrote a small local training script and overfit a couple of real IFS samples to check the whole chain works. It does learn (loss drops fine), but it stops improving before it fully fits even a tiny batch, and bumping the H3 resolution didn't change that floor. Not sure yet if that's a real problem or just that 6-hourly change isn't very predictable from surface fields. Haven't dug into it properly, didn't want to chase it without a clear reason.

For Phase 2 (the unified adaptive mesh with two-way coupling you mentioned), I read up on how a few other groups handle regional. The two closest to what you want look like the global stretched-grid one from MET Norway / Anemoi (single mesh, just refined over the region, no separate boundary), and OneForecast's region-refined graph with the gating mechanism (the paper from #160). Both get two-way coupling without the dual-graph boundary setup we have now.

Did you have a direction in mind? I can start moving toward one of those, orkeep building on the current dual-graph version for a bit. Whichever you prefer.

I think I would prefer the one that is the global stretched-grid from MET Norway. It feels as a single mesh, we can move the mesh around a bit easier than with the gating. So yeah, I would prefer that way of doing it.

@jacobbieker

Thanks @jacobbieker ,i also agree with stretched-grid, that direction makes sense to me.

I spent some time reading up on the Anemoi / MET Norway setup and I found that it lines up with what we already have. Since we're on H3
already, a lot of the variable resolution comes almost for free: coarse hexes over the globe, refine to a higher H3 resolution over the region, and lean onH3's parent/child links to stitch across the resolution change. The encode-process-decode core and the GNN blocks all carry straight over, which is nice. The one casualty is the BoundaryNudgingLayer, it kind of retires itself once there's a single continuous mesh with no boundary to nudge, but that feels right.

I'll keep the PRs small like before. I was thinking the first one is just the mesh-building piece: a helper that assigns H3 cells at a fine resolution inside the region and coarse outside, with tests that the coverage comes out complete and non-overlapping. No model or training changes in that one. The variable-res latent graph and the embedding table can follow as separate PRs.

Does that first step sound good, or would you rather I kick off somewhere else?

Yeah, that first one sounds good. And small PRs are great, makes it easier to review.

Quick update on the stretched-grid model (StretchedForecaster + the region-weighted loss, plus the StretchedDataset in #236), and a question before I build out the training pipeline.

I ran a sanity overfit on 2 fixed IFS samples (region-weighted MSE, res2 global / res3 region, node_dim 64, 4 processor blocks). Loss drops from ~1.2 to 0.17 and then sits flat. Decomposing that floor, the 0.17 turns out to be exactly the persistence baseline (predicting the input unchanged) - global and regional error both stay at persistence. So the model isn't learning the forecast increment, even on 2 memorized samples.

To check the signal is actually there, I overfit a graph-free per-node MLP (delta = MLP(features), out = features + delta) on the same 2 samples. That one beats persistence, 0.173 to 0.114. So the increment is partly learnable from the input frame, but the encode-process-decode graph pipeline isn't preserving it - a trivial per-node model does better than the full graph model.

My guess is the processor is over-smoothing (each observation ends up decoded from its cell's feature after global message passing), but I didn't want to keep guessing at it. Have you seen this with the graph_weather encode-process-decode stack, and is there a usual way you handle it (fewer blocks, a different decoder), or does it more likely mean I've wired the stretched encoder/decoder wrong somewhere? Script and terminal output below - it runs on the #236 branch.

I haven't seen that explicitly before, but I have seen in other papers and work that the processing grid can smooth out features quite a bit. In general, the AI weather models do tend to smooth things out. One could be to try increase the latent graph resolution to see if that makes a difference?

I would suggest double checking the wiring between the stages. But I don't have anything specific in why that might be the case.

@jacobbieker

@jacobbieker here is the script i used for


import h3
import torch
import torch.nn as nn

from graph_weather.data.regional_dataset import CORE_SURFACE
from graph_weather.data.stretched_dataset import StretchedDataset
from graph_weather.models.layers.stretched_mesh import (
    assign_points_to_mesh,
    build_variable_resolution_mesh,
)
from graph_weather.models.losses import regional_weighted_mse
from graph_weather.models.stretched_forecast import StretchedForecasterConfig

COARSE_RES, FINE_RES = 2, 3
REGION_WEIGHT = 5.0
STEPS = 100
LR = 1e-2


def region_mask(lat_lons, bbox):
    """Mark observations that land on a fine (regional) cell, as the model assigns them."""
    mesh = build_variable_resolution_mesh(bbox, COARSE_RES, FINE_RES)
    assigned = assign_points_to_mesh(lat_lons, mesh, COARSE_RES, FINE_RES)
    return torch.tensor([h3.get_resolution(c) == FINE_RES for c in assigned])


def overfit(model, samples, device):
    """Repeatedly fit the fixed samples; return the per-step mean loss."""
    model = model.to(device).train()
    opt = torch.optim.Adam(model.parameters(), lr=LR)
    losses = []
    for step in range(STEPS):
        opt.zero_grad()
        total = 0.0
        for features, lat_lons, target, bbox in samples:
            mask = region_mask(lat_lons, bbox).to(device)
            out = model(features.unsqueeze(0).to(device), lat_lons, bbox)
            loss = regional_weighted_mse(out, target.unsqueeze(0).to(device), mask, REGION_WEIGHT)
            loss.backward()
            total += loss.item()
        opt.step()
        losses.append(total / len(samples))
        if step % 20 == 0 or step == STEPS - 1:
            print(f"  step {step:3d}  loss {losses[-1]:.4f}")
    return losses


def weighted_loss(pred_fn, samples, device):
    """Mean region-weighted loss of pred_fn over the samples, without training."""
    vals = []
    for features, lat_lons, target, bbox in samples:
        mask = region_mask(lat_lons, bbox).to(device)
        pred = pred_fn(features.unsqueeze(0).to(device), lat_lons, bbox)
        tgt = target.unsqueeze(0).to(device)
        vals.append(regional_weighted_mse(pred, tgt, mask, REGION_WEIGHT).item())
    return sum(vals) / len(vals)


class PerNodeBaseline(nn.Module):
    """Predict each observation's next state from its own features alone, no graph."""

    def __init__(self, dim, hidden=64):
        """Build a per-node residual MLP at the same width as the stretched model."""
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
            nn.Linear(hidden, dim),
        )

    def forward(self, features, lat_lons, bbox):
        """Return features + a per-node learned delta; lat_lons and bbox are ignored."""
        return features + self.net(features)


def main():
    """Run the overfit, the persistence decomposition, and the per-node control."""
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dim = len(CORE_SURFACE)
    dataset = StretchedDataset(
        coarse_res=COARSE_RES,
        fine_res=FINE_RES,
        extent_deg=10.0,
        max_regional_points=200,
    )
    samples = [dataset[i] for i in range(2)]
    print(f"device={device}. 2 IFS samples, ~{samples[0][0].shape[0]} points each, {dim} vars.\n")

    print("[1] stretched model overfit:")
    torch.manual_seed(0)
    stretched = StretchedForecasterConfig(
        feature_dim=dim,
        aux_dim=0,
        coarse_res=COARSE_RES,
        fine_res=FINE_RES,
        node_dim=64,
        edge_dim=64,
        num_blocks=4,
    ).build()
    s_losses = overfit(stretched, samples, device)

    def persistence(features, lat_lons, bbox):
        return features[..., :dim]

    p_loss = weighted_loss(persistence, samples, device)

    print("\n[2] graph-free per-node baseline overfit:")
    torch.manual_seed(0)
    per_node = PerNodeBaseline(dim)
    n_losses = overfit(per_node, samples, device)

    s_gap = p_loss - s_losses[-1]
    n_gap = p_loss - n_losses[-1]
    print("\n=== summary (region-weighted MSE, standardized units) ===")
    print(f"  persistence baseline    {p_loss:.4f}")
    print(f"  stretched model         {s_losses[-1]:.4f}   (vs persistence {s_gap:+.4f})")
    print(f"  per-node MLP, no graph  {n_losses[-1]:.4f}   (vs persistence {n_gap:+.4f})")


if __name__ == "__main__":
    main()

@jacobbieker

Hi @jacobbieker Following up - I think I found it, and the fix holds up on held-out data.

The decoder was rebuilding each observation purely from its assigned cell's feature. The encoder computes per-obs features (nodes[:num_obs]) but they get discarded, and the decoder's obs nodes are seeded with torch.zeros - so the learnable delta only ever sees cell-level (smoothed) context, and an observation's own values reach the output only through the fixed residual. That is why it collapsed to persistence, independent of processor depth (I swept 1-8 blocks, all floored) or the wiring (which matches RegionalForecaster).

The fix: seed the decoder's obs nodes with the encoder's per-obs features instead of zeros - a skip connection, no new parameters. Results:

2-sample overfit: floor drops from 0.173 (= persistence) to 0.136.
Generalization (train 30 movable IFS samples, eval 8 held-out regions): persistence 0.214, baseline 0.206, skip-fix 0.171. The fix beats persistence by ~20% on held-out vs the baseline's ~4%.

The same zeros-placeholder pattern is in RegionalForecaster too. I also noticed DynamicGraphBuilder already builds a richer grid_disk(cell, 1) decoder graph that RegionalForecaster.forward discards - might be worth wiring back in.

I ran the bigger comparison on Colab since my laptop GPU is only 4GB. Happy to open a PR with the stretched fix. One honest caveat: because the encoder gives obs nodes no incoming edges, this skip path is effectively per-observation, so it does not yet prove the mesh message-passing adds value beyond a per-node model - separate thing to test. But it is a clear net improvement.

Okay, nice! Good investigation. I would say open a PR with the fix, and an issue for checking the message passing bit. Thanks for looking into this!

@jacobbieker

Thanks @jacobbieker. Follow-ups from the decoder fix are up:

Opened #238 for the message-passing question. The plan there: train the fixed stretched model and a graph-free per-node MLP on the same many-sample set and compare held-out region-weighted MSE. If the graph model wins, the mesh is earning its keep; if not, the message-passing needs work. A likely reason the mesh is under-used today is that observation nodes get no incoming edges in the encoder and the decoder connects each observation to a single cell - so the plan also includes wiring each observation to its grid_disk(cell, 1) neighbourhood (the graph DynamicGraphBuilder already builds but the forward discards).

Small follow-up PR for the identical zeros-seed bug in RegionalForecaster: (fix: seed regional decoder observation nodes with encoder features #239).

I'll run the #238 ablation on Colab (region-weighted training needs more than my 4GB card) and post the numbers here the same way as the generalization run. Does the approach above look right to you, or would you weight it differently?

Yeah, I think that approach seems fine!

@jacobbieker

@jacobbieker following up on the message-passing question (#238) with the ablation you okayed.

Ran it on held-out movable IFS regions: 30 train / 8 held-out, region-weighted MSE, 3 seeds,
node_dim 64 / 4 blocks. Persistence stays 0.2144 across runs, so the arms are comparable.

held-out region-weighted MSE (3 seeds):

persistence: 0.2144
per-node MLP (19k params): 0.1660 +/- 0.0004
mesh, single cell (2.27M params): 0.1731 +/- 0.0019
mesh + grid_disk(cell,1) (2.27M): 0.1715 +/- 0.0026

The graph-free per-node MLP wins, by 0.0055 over the best mesh variant, with about 120x fewer
parameters. Wiring each observation to its grid_disk(cell,1) neighbourhood instead of a single
cell helps only marginally (0.1731 -> 0.1715, smaller than the seed spread) and does not close
the gap. So on this test the message-passing is not adding forecasting skill - a per-observation
map does better.

Caveat: this is small scale (30 samples, 25 epochs), and the mesh is 2.27M params against the
MLP's 19k, so it may be data-starved here. Its spatial-context advantage may only show up with
more data or a multi-step target where neighbours actually matter (the grid_disk arm was still
improving at the last epoch).

Before building out the training pipeline I would rather settle this. Two directions: scale the
comparison up (more samples/epochs, maybe a multi-step target so advection matters), or treat it
as an architecture problem and rethink how observations couple to the mesh. Which would you find
more convincing, and is there a setup you would want to see?

Navigation Menu

Repository navigation

Add adaptive meshing #3

Detailed Description

Context

Possible Implementation

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Add adaptive meshing #3

Description

Detailed Description

Context

Possible Implementation

Activity

AswaniSahoo commented on May 24, 2026

jacobbieker commented on May 26, 2026

AswaniSahoo commented on May 26, 2026

AswaniSahoo commented on May 31, 2026

AswaniSahoo commented on Jun 10, 2026

jacobbieker commented on Jun 11, 2026

AswaniSahoo commented on Jun 16, 2026

jacobbieker commented on Jun 17, 2026

AswaniSahoo commented on Jun 20, 2026

jacobbieker commented on Jun 23, 2026

AswaniSahoo commented on Jun 24, 2026

jacobbieker commented on Jun 25, 2026

AswaniSahoo commented on Jul 20, 2026

AswaniSahoo commented on Jul 20, 2026

jacobbieker commented on Jul 20, 2026

AswaniSahoo commented on Jul 20, 2026

AswaniSahoo commented on Jul 21, 2026

jacobbieker commented on Jul 21, 2026

AswaniSahoo commented on Jul 26, 2026

jacobbieker commented on Jul 27, 2026

AswaniSahoo commented on Jul 27, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions