README.md

---
frameworks: PyTorch
language:
- en
license: apache-2.0
tags:
- OneScience
- Earth Science
- Weather Forecast
- Ensemble Weather Forecasting
- Diffusion Model
- GEFS
- ERA5
tasks: []
datasets:
- ERA5
---
<p align="center">
  <strong>
    <span style="font-size: 30px;">SEEDS</span>
  </strong>
</p>

# Model Introduction

SEEDS is a generative weather model released by Google in March 2024. The name stands for Scalable Ensemble Envelope Diffusion Sampler; the associated paper was published in *Science Advances*.

Paper: SEEDS: Emulation of Weather Forecast Ensembles with Diffusion Models

https://arxiv.org/abs/2306.14066

# Model Description

SEEDS is a conditional diffusion model that uses a small number of numerical weather prediction seed members to efficiently generate large forecast ensembles.

# Use Cases

| Scenario | Description |
| :---: | :--- |
| Ensemble weather forecast research | Train a conditional diffusion model on data following the project's cubed-sphere NPZ protocol and generate forecast ensembles. |
| Local quick validation | Use synthetic data to check training, inference, ensemble evaluation, and visualization. |
| ModelScope / OneCode execution | Download the standalone model package, install dependencies, and run the scripts directly. |
| Multi-GPU training | Launch PyTorch DistributedDataParallel with `torchrun`. |

# Usage Guide

## 1. OneCode Usage

Experience intelligent one-click AI4S programming through the OneCode online environment:

[Click to Experience Intelligent One-Click AI4S Programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)

## 2. Manual Installation and Usage

**Hardware Requirements**

- Training and inference require a GPU or DCU recognized by PyTorch. CPU can be used to generate synthetic data and inspect configuration, but cannot run the current training and inference scripts.
- Multi-GPU training uses the NCCL backend. Ensure that the device driver, communication libraries, and PyTorch version are compatible.
- DCU users must install DTK in advance. DTK 25.04.2 or above, or the OneScience recommended version matching your cluster, is recommended.

### Download the Model Package

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

### Install the Runtime Environment

**DCU Environment**

```bash
# Please 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
# Please 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 Introduction

The SEEDS paper uses GEFS reforecast data for training, operational GEFS members as conditioning inputs, and ERA5 as the evaluation reference. The official data and preprocessing are not bundled with this repository; prepare the NPZ files specified by `conf/config.yaml` before training.

### Generate Synthetic Data

When real data is unavailable, generate a default `6x48x48` cubed-sphere fixture. Synthetic data only validates the program flow and does not represent GEFS, ERA5, or the paper's forecast quality:

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

### Training

Single GPU:

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

Multi-GPU:

```bash
torchrun --nproc_per_node=8 scripts/train.py
```

The number of epochs is controlled by `conf/config.yaml`, and the default checkpoint is saved to `data/checkpoint/model_bak.pth`.

### Fine-tuning

To continue from an existing checkpoint, use the explicit fine-tuning flag:

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

### Training Weights

This repository provides a `weight/` directory for checkpoints trained on the official data. The weight files will be uploaded soon and are expected to be available in the near future.

### Inference

Inference reads `data/checkpoint/model_bak.pth` by default and generates ensemble members in batches controlled by `sampling.member_batch_size`:

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

Predictions, targets, and seed members are saved under `result/output/`.

### Evaluation and Visualization

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

The script computes ensemble-mean RMSE, ACC, and empirical CRPS, saves the corresponding NPY metrics, and generates `result/forecast.png`. If training loss files are available, it also writes `result/loss.png`.

# Official OneScience Resources

| 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 adaptation of the SEEDS paper and is not an official Google product.