OneScience-Group/FuXi-DA

FuXi-DA

Model Introduction

FuXi-DA fuses weather background fields and satellite observations into improved global atmospheric analyses for deep-learning data assimilation and forecast initialization. It learns background errors, observation biases, and analysis increments in a unified feature space to improve subsequent forecasts with multi-time geostationary satellite data.

Paper: FuXi-DA: a generalized deep learning data assimilation framework for assimilating satellite observations
https://doi.org/10.1038/s41612-025-01039-3

Model Description

The method was proposed by teams from the Shanghai Artificial Intelligence Laboratory, Fudan University, and collaborating institutions. The paper trains on ERA5, FuXi background fields, and Fengyun-4B AGRI observations. FuXi-DA performs multi-scale fusion with background, observation, and condition encoders and uses frozen-FuXi forecast errors to supervise assimilation. The model supports 0.25-degree global background correction and geostationary multi-channel satellite assimilation.

Use Cases

Use Case	Description
Satellite assimilation	Fuse 70-channel backgrounds with multi-time AGRI brightness temperatures.
Multimodal fusion	Validate multi-scale interactions among background, observation, and condition branches.
Local validation	Execute sampled aligned tiles under full-variable and global-coordinate protocols.
ModelScope/OneCode execution	Validate data, training, inference, assimilation metrics, and visualization.
Multi-GPU training	Validate distributed training and checkpoint workflows through torchrun.

Usage Instructions

1.OneCode

Try intelligent, one-click AI4S programming

2. Download and Installation

hf download OneScience-Group/FuXi-DA --local-dir ./FuXi-DA
cd FuXi-DA

Environment Dependencies

Hardware Requirements

• A GPU or DCU is recommended.
• A CPU can be used for connectivity validation with the default small-sample configuration.
• DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended.

DCU Environment

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

GPU Environment

# 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
pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai

Training Data

This repository uses a small synthetic dataset for engineering validation, including 70-variable global backgrounds, eight AGRI times, 15 observation channels, six-hour relationships, and the real [70,721,1440] background and [8,15,640,640] observation protocols. Procedural full-coordinate fields and a few aligned original-coordinate tiles preserve variables, channels, and spatiotemporal specifications while reducing executed samples and coverage. These data validate the multi-branch U-Net, feature fusion, training, inference, and evaluation only and do not represent official ERA5 or AGRI distributions and scale.

python scripts/fake_data.py

Training

For single-GPU training, use:

python scripts/train.py

For multi-GPU training, use:

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 combines latitude-weighted analysis L1 with multi-lead supervision from a frozen forecast proxy. The default reduces tile coverage, network width, iterations, and forecast steps while recording the paper's 6,000 iterations and ten forecast leads. Training artifacts are saved to:

result/checkpoints/fuxi_da.pt
result/training/metrics.json

Trained Weights

The paper does not provide directly loadable official model weights, and no weights are bundled under weight/.

Inference

python scripts/inference.py

Inference loads the local checkpoint and takes a 70-channel background tile and corresponding multi-time AGRI observations. Background, observation, and condition branches generate an increment that is added to the background to form the analysis. Outputs retain targets, backgrounds, observations, analyses, coordinates, and coverage metadata and explicitly indicate incomplete global coverage. Inference results are saved to:

result/output/predictions.npz

Evaluation and Visualization

python scripts/result.py

Evaluation compares latitude-weighted RMSE for the raw background, correction baseline, and FuXi-DA analysis and saves Z, T, U, V, R, and surface-variable groups. It also reports forecast-step errors and localization of analysis increments from a single-observation perturbation and generates a comparison of analysis schemes and variable-group errors. Synthetic tile results validate engineering only and do not represent complete global paper performance. Evaluation results are saved to:

result/evaluation/metrics.json
result/evaluation/comparison.png

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 engineering reproduction of the public FuXi-DA specifications, with code licensed under the Apache License 2.0.

The original paper is licensed under CC BY-NC-ND 4.0; the paper, official model weights, ERA5 data, and Fengyun-4B AGRI data remain subject to the licenses and terms of their respective projects.