File size: 2,219 Bytes
82b3be0
 
793db94
 
 
 
 
 
 
82b3be0
793db94
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
---
license: bsd-3-clause
tags:
- weather
- super-resolution
- flow-matching
- era5
- atmospheric-science
library_name: diffusers
---

# Super-Resolving Coarse-Resolution Weather Forecasts with Flow Matching

[![arXiv](https://img.shields.io/badge/arXiv-2604.00897-b31b1b.svg)](https://arxiv.org/abs/2604.00897)
[![GitHub](https://img.shields.io/badge/GitHub-ArchesWeatherSR-black)](https://github.com/dataymeric/ArchesWeatherSR)

![ArchesWeatherSR overview](https://raw.githubusercontent.com/dataymeric/ArchesWeatherSR/main/assets/archesweathersr_overview.png)

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](https://github.com/INRIA/geoarches).

## Usage

Install the package and download the model:

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

Run inference:

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

Or load programmatically:

```python
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