| --- |
| license: apache-2.0 |
| language: |
| - en |
| - zh |
| tags: |
| - OneScience |
| - Earth science |
| - Weather forecasting |
| - Medium-range weather forecasting |
| - Observation-driven |
| - Graph neural networks |
| - ERA5 |
| frameworks: PyTorch |
| datasets: |
| - OneScience/ERA5 |
| --- |
| <p align="center"> |
| <strong> |
| <span style="font-size: 30px;">GraphDOP</span> |
| </strong> |
| </p> |
| |
|
|
| # Model Introduction |
|
|
| GraphDOP (Graph-based Direct Observation Prediction) was proposed by ECMWF. It is an end-to-end observation-driven (AI-DOP) weather forecasting model based on graph neural networks (GNNs). The model uses only Earth-system observations (polar-orbiting/geostationary satellite brightness temperatures, occultation bending angles, scatterometer backscatter, radar altimetry, radiosonde and conventional surface observations, and so on) as inputs and training targets. It does not use any physics-based reanalysis field and can produce medium-range forecasts beyond 5 days. |
|
|
| Paper:GraphDOP: Towards skilful data-driven medium-range weather forecasts learnt and initialised directly from observations |
|
|
| https://arxiv.org/abs/2412.15687 |
|
|
| # Model Description |
|
|
| GraphDOP uses an encoder-processor-decoder architecture: the GNN encoder maps observations in the input window to an O96 (approximately 1°) latent grid according to spatial proximity, the Transformer processor advances the atmospheric state in latent space, and the GNN decoder maps the latent grid back to target observation locations and outputs forecasts channel by channel. The training objective is channel-weighted mean squared error (WMSE). 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 | |
| | :---: | :--- | |
| | Observation-driven medium-range weather forecasting research | Learn atmospheric-state representations directly from observations and forecast future windows. | |
| | Graph + Transformer latent-space model research | Reproducible encoder-processor-decoder architecture and WMSE 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/GraphDOP --local-dir ./GraphDOP |
| cd GraphDOP |
| ``` |
|
|
| ### 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 the model input/output windows and `grid_shape`. Because ERA5Datapipe supports only regular grids, this project approximates the irregular Level-1 observations in the paper with six observation-grid channels. |
|
|
| ### 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 (1024 latent channels, an O96 latent grid, and 70k steps on 64×H100 GPUs) has no publicly released weights. |
|
|
| ### Inference |
|
|
| Inference reads `data/checkpoints/model_bak.pth`: |
|
|
| ```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 an ECMWF AI-DOP preprint, and no official implementation is publicly available (the source is based on PyTorch Geometric and builds dynamic observation graphs per batch). This repository's `model/graphdop.py` is a pure PyTorch minimal reproduction that preserves the paper's encoder-processor-decoder GNN structure and WMSE objective. |
| - Differences from the paper (due to limitations of the OneScience gridded data pipeline): the paper consumes irregular raw Level-1 observations, with each observation connected to its nearest latent-grid neighbors using a dynamically constructed graph. This reproduction uses ERA5 gridded h5 channels as a placeholder observation field and fixes the latent grid to a regular 8-neighbor graph. The edge features (azimuth and Haversine distance) match the paper. Training augmentations such as random observation dropout (25% for satellites/50% for conventional observations) are not implemented. |
| - `conf/config.yaml` uses a small configuration (32×32 grid, 8×8 latent grid, latent_dim=64) for connectivity validation by default; paper-level reproduction requires an O96 latent grid, latent_dim=1024, and a larger data scale. |
| - The following details are not disclosed in the paper and are assumptions in this reproduction:per-channel weights `w_{c,i}` (currently all 1 by default), graph-construction details, and the specific implementation of windowed attention in the processor. |
|
|
| # 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 GraphDOP reproduction (the model code is an original minimal implementation), with the architecture design based on the paper by Alexe et al. (2024). |
| - Please cite:Alexe, M., E. Boucher, P. Lean, E. Pinnington, P. Laloyaux, A. McNally et al. GraphDOP: Towards skilful data-driven medium-range weather forecasts learnt and initialised directly from observations. arXiv:2412.15687, 2024. |
|
|