README.md

---
license: apache-2.0
language:
- en
- zh
tags:
- OneScience
- Earth science
- Weather forecasting
- Probabilistic weather forecasting
- Ensemble forecasting
- Graph neural networks
- Conditional normalization
- CRPS
- ERA5
frameworks: PyTorch
datasets:
  - OneScience/ERA5
---
<p align="center">
  <strong>
    <span style="font-size: 30px;">WeatherNext2 · FGN</span>
  </strong>
</p>


# Model Introduction

FGN (**F**unctional **G**enerative **N**etworks) was proposed by Google DeepMind and is the probabilistic global weather forecasting model in the WeatherNext series (WeatherNext2). FGN uses the fair continuous ranked probability score (fair CRPS) of the marginal distribution as its training objective. It models aleatoric uncertainty by injecting a global noise vector into conditional normalization layers and epistemic uncertainty through an ensemble of independently trained models. In this way, it captures the joint spatial structure of ensemble forecasts while optimizing only pointwise marginal objectives, and surpasses GenCast and ECMWF ENS across 15-day medium-range forecasts.

Paper:Skilful joint probabilistic weather forecasting from marginals

https://arxiv.org/abs/2506.14285

# Model Description

FGN uses a GNN encoder-processor-decoder architecture consistent with the GenCast denoiser: a sparse GNN encoder maps latitude-longitude grid inputs to a latent space on a six-times-subdivided icosahedral grid, a graph-transformer processor advances the atmospheric state on that grid, and a GNN decoder maps the latent grid back to the output grid. Each forecast samples a 32-dimensional global noise vector, embeds it with a single matrix multiplication, and injects it into all conditional LayerNorm layers (equivalent to applying a learned functional perturbation to the network parameters) as the source of ensemble spread. Under a second-order Markov assumption (using the two previous frames), the model generates forecasts autoregressively at 6-hour intervals. This repository is a minimal reproduction based on the paper and is integrated with the OneScience data loading and training workflow.

# Use Cases

| Scenario | Description |
| :---: | :--- |
| Probabilistic/ensemble medium-range weather forecasting research | Learn pointwise marginal distributions with CRPS and generate joint ensemble forecasts. |
| Uncertainty modeling research | Reproducible parameter-space noise injection (conditional normalization) and deep ensemble mechanisms. |
| Graph + Transformer latent-space model research | Encoder-processor-decoder architecture and fair CRPS objective. |
| Local quick validation | Use synthetic data to check data loading, training, inference, and result scripts. |
| ModelScope/OneCode execution | Download the model package, install dependencies, and run the scripts directly. |
| Multi-card training | Launch multi-process training with `torchrun`. |


# Usage

## 1. OneCode Usage

Use the OneCode online environment for intelligent one-click AI4S programming:

[Try intelligent one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)

## 2. Manual Installation and Usage

**Hardware Requirements**

- GPU or DCU is recommended.
- CPU can be used for imports and small-configuration connectivity validation, but full training and inference are slower.
- DCU users must install DTK beforehand. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended.


### Download the Model Package

```bash
hf download OneScience-Group/WeatherNext2 --local-dir ./WeatherNext2
cd WeatherNext2
```

### Install the Runtime Environment

**DCU Environment**

```bash
# Activate DTK and CONDA first
conda create -n onescience311 python=3.11 -y
conda activate onescience311
# uv installation is supported
pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
```

**GPU Environment**
```bash
# Activate CONDA first
conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
conda activate onescience311
# uv installation is supported
pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
```

### Training Data

The OneScience community provides ERA5 data for training (the current repository contains complete data slices subject to data-file size limits). Download it with the command below and confirm that the data path in `conf/config.yaml` is correct:

```bash
hf download --repo-type dataset OneScience-Group/ERA5 --local-dir ./data
```

For a quick workflow validation, run the synthetic data script first:

```bash
python scripts/fake_data.py
```

> Note: `scripts/fake_data.py` generates `[T, C, H, W]` data from second-order Markov inputs, forecast steps, batch size, and `grid_shape`. The current small configuration uses 6 channels and a 32×32 grid.

### Training

Single card：

```bash
python scripts/train.py
```

Multiple cards：

```bash
torchrun --nproc_per_node=8 --nnodes=1 --rdzv_id=1000 --rdzv_backend=c10d --max_restarts=0 --master_addr="localhost" --master_port=29500 scripts/train.py
```

Training outputs:

```text
data/checkpoints/model_bak.pth
data/checkpoints/trloss.npy
data/checkpoints/valoss.npy
```

### Training Weights
The `weight/` folder is reserved for model weights. Pretrained weights are not provided by default; users may train the model using the paper configuration. The paper model (per-seed latent dimension 768, 24 processor layers, an ensemble of 4 model seeds, and approximately 490 TPU-days of total computation) has no publicly released weights.

### Inference

Inference reads `data/checkpoints/model_bak.pth` by default, generates `num_members` ensemble members for each initialization time (independently sampling global noise for each), and uses the member mean as the deterministic forecast output:

```bash
python scripts/inference.py
```

Prediction results are written frame by frame to:

```text
result/output/
```

### Evaluation and Visualization

```bash
python scripts/result.py
```

Outputs include：

- `result/rmse.npy`
- `result/acc.npy`
- `result/loss.png`
- Forecast comparison plots for the specified date and variables


# Official Source and Reproduction Notes

- The paper is a Google DeepMind preprint (© 2025 Google DeepMind. All rights reserved.), and no official implementation or weights are publicly available. This repository's `model/fgn.py` is a pure PyTorch minimal reproduction that preserves the paper's encoder-processor-decoder GNN structure, conditional LayerNorm global-noise injection, and fair CRPS objective (with N=2 ensemble samples during training).
- Differences from the paper (due to the OneScience gridded data pipeline and connectivity-validation scale): the paper uses a six-times-subdivided icosahedral latent grid (approximately 40k nodes) and a 0.25° (1440×721) output grid, with approximately 180M parameters per seed. This reproduction uses a fixed regular 8-neighbor latent grid (`mesh_shape`), ERA5 gridded h5 channels as input placeholders, and a default latent dimension of 64. The number of latent-grid nodes and edge features are determined by the `mesh_shape` and `channel_weights` configurations.
- The following details are not disclosed in the paper and are assumptions in this reproduction:the specific projection from noise to each layer's scale/shift in conditional LayerNorm (currently implemented with per-layer linear projections initialized to zero so the initial behavior is standard LayerNorm); per-channel weights for the multi-task loss (currently all 1 by default); and the evaluation convention of using the ensemble mean as the deterministic forecast for probabilistic output.
- Paper-level reproduction requires the four-stage training procedure in the paper (ERA5 1°12h → 1°6h → 0.25°6h → HRES-fc0 0.25° AR fine-tuning); `conf/config.yaml` uses a small configuration for connectivity validation by default.

# Official OneScience Information

| Platform | OneScience Main Repository | Skills Repository |
| --- | --- | --- |
| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |

# Citation and License

- This repository is an independent FGN (WeatherNext2) reproduction (the model code is an original minimal implementation), with the architecture design based on the paper by Alet et al. (2025).
- Please cite:Alet, F., Price, I., El-Kadi, A., Masters, D., Markou, S., Andersson, T. R., Stott, J., Lam, R., Willson, M., Sanchez-Gonzalez, A. and Battaglia, P. Skilful joint probabilistic weather forecasting from marginals. arXiv:2506.14285, 2025.