dataymeric/ArchesWeatherSR

Super-Resolving Coarse-Resolution Weather Forecasts with Flow Matching

A flow matching–based generative super-resolution model for global weather forecasts. Takes coarse-resolution (1.5°) forecast trajectories and generates stochastic high-resolution (0.25°) outputs, recovering fine-scale variability while preserving large-scale structure.

Model description

ArchesWeatherSR is formulated as a stochastic inverse problem using flow matching. It learns the residual between the bicubically interpolated coarse field and the true ERA5 analysis at 0.25°, concentrating model capacity on fine-scale structure. At inference, the residual is sampled and added back to the interpolated field. The backbone is a 3D Swin U-Net Transformer shared with ArchesWeather & ArchesWeatherGen.

Usage

Install the package and download the model:

git clone https://github.com/dataymeric/ArchesWeatherSR.git
cd ArchesWeatherSR
uv sync
hf download dataymeric/ArchesWeatherSR --local-dir runs/archesweathersr

Run inference:

python train.py mode=test ++name=archesweathersr

Or load programmatically:

from geoarches.lightning_modules import load_module

sr_model, cfg = load_module("runs/archesweathersr")
sr_model = sr_model.cuda().eval()

# sample a super-resolved state from a batch
samples = sr_model.sample(batch)

Training

• Training data: ERA5 reanalysis (WeatherBench2), 1979–2018 (train), 2019 (val), 2020 (test)
• Variables: 6 upper-air (z, u, v, t, q, w) × 13 levels + 4 surface (t2m, msl, u10, v10)
• Optimizer: AdamW, lr=3×10⁻⁴, β=(0.9, 0.98), wd=0.01
• Training steps: 75,000
• Hardware: 4 × A100 80GB, ~40h