OneScience-Group/FuXi-DA

1---2license: apache-2.03language:4- en5tags:6- OneScience7- Earth Science8- Data Assimilation9- Satellite Observations10- Multimodal Fusion11- Weather Forecasting12frameworks: PyTorch13---14 15<p align="center"><strong><span style="font-size: 30px;">FuXi-DA</span></strong></p>16 17# Model Introduction18 19FuXi-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.20 21Paper: FuXi-DA: a generalized deep learning data assimilation framework for assimilating satellite observations  22https://doi.org/10.1038/s41612-025-01039-323 24# Model Description25 26The 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.27 28# Use Cases29 30| Use Case | Description |31| :---: | :--- |32| Satellite assimilation | Fuse 70-channel backgrounds with multi-time AGRI brightness temperatures. |33| Multimodal fusion | Validate multi-scale interactions among background, observation, and condition branches. |34| Local validation | Execute sampled aligned tiles under full-variable and global-coordinate protocols. |35| ModelScope/OneCode execution | Validate data, training, inference, assimilation metrics, and visualization. |36| Multi-GPU training | Validate distributed training and checkpoint workflows through `torchrun`. |37 38# Usage Instructions39 40## 1.OneCode41 42[Try intelligent, one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)43 44## 2. Download and Installation45 46```bash47hf download OneScience-Group/FuXi-DA --local-dir ./FuXi-DA48cd FuXi-DA49```50 51### Environment Dependencies52 53**Hardware Requirements**54 55- A GPU or DCU is recommended.56- A CPU can be used for connectivity validation with the default small-sample configuration.57- DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended.58 59**DCU Environment**60 61```bash62# Activate DTK and Conda first63conda create -n onescience311 python=3.11 -y64conda activate onescience31165pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai66```67 68**GPU Environment**69 70```bash71# Activate Conda first72conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=1273conda activate onescience31174pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai75```76 77### Training Data78 79This 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.80 81```bash82python scripts/fake_data.py83```84 85### Training86 87For single-GPU training, use:88 89```bash90python scripts/train.py91```92 93For multi-GPU training, use:94 95```bash96torchrun --nproc_per_node=8 --nnodes=1 --rdzv_id=1000 --rdzv_backend=c10d --max_restarts=0 --master_addr="localhost" --master_port=29500 scripts/train.py97```98 99Training 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:100 101```text102result/checkpoints/fuxi_da.pt103result/training/metrics.json104```105 106### Trained Weights107 108The paper does not provide directly loadable official model weights, and no weights are bundled under `weight/`.109 110### Inference111 112```bash113python scripts/inference.py114```115 116Inference 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:117 118```text119result/output/predictions.npz120```121 122### Evaluation and Visualization123 124```bash125python scripts/result.py126```127 128Evaluation 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:129 130```text131result/evaluation/metrics.json132result/evaluation/comparison.png133```134 135# Official OneScience Information136 137| Platform | OneScience Main Repository | Skills Repository |138| --- | --- | --- |139| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |140| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |141 142# Citation and License143 144This repository is an independent engineering reproduction of the public FuXi-DA specifications, with code licensed under the Apache License 2.0.145 146The 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.147