README.md

---
license: other
language:
- en
- zh
tags:
- OneScience
- Earth Science
- Weather Forecasting
- Medium-Range Weather Forecasting
frameworks: PyTorch
datasets:
- OneScience/Aardvark-Weather
---

<p align="center">
  <strong><span style="font-size: 30px;">Aardvark Weather</span></strong>
</p>

# Model Overview

Aardvark Weather is an end-to-end multimodal weather forecasting model that generates global gridded forecasts and station-level predictions through an observation encoder, a global forecast processor, and a station decoder.

Paper: *End-to-end data-driven weather prediction*

https://www.nature.com/articles/s41586-025-08897-0

# Model Description

This model package reuses the official code and weights to provide the following run pipeline:

```text
Official Multimodal Sample
-> Encoder
-> Day-1 Processor
-> TAS Decoder
-> 1-day global forecast and station 2-meter temperature
```

# Use Cases

| Scenario | Description |
| :---: | :--- |
| Official Model Verification | Inspect official samples, configuration, and checkpoints. |
| Global Weather Forecasting | Output a global 1.5° gridded state of 24 variables. |
| Station Temperature Forecasting | Output 2-meter temperature at 8,719 stations. |

# Usage

## 1. OneCode

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

## 2. Manual Installation & Usage

**Hardware Requirements**

- Inference with the official weights requires an NVIDIA GPU.
- CPU can be used for resource and checkpoint inspection; running full inference on CPU is not recommended.

### Download the Model Package

```bash
hf download --model OneScience-Group/Aardvark-Weather --local-dir ./Aardvark-Weather
cd Aardvark-Weather
```

### Set Up the Runtime Environment

**DCU Environment**

```bash
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**

```bash
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
```

If the environment lacks Aardvark dependencies, refer to `official-src/environment.yml` for additional installation. The current adaptation is compatible with the `Block` parameter of the newer `timm` package.

### Data & Weights

The model package already includes the resources required for 1-day temperature inference:

```text
weights/sample_data/sample_data_final.pkl
weights/trained_model/encoder/epoch_96
weights/trained_model/processor/forecast_1/epoch_0
weights/trained_model/decoder/tas/lt_1/epoch_18
official-src/data/grid_lon_lat/
official-src/data/norm_factors/
```

To re-download from source:

```text
Official Code: https://github.com/anna-allen/aardvark-weather-public
Official Weights: https://huggingface.co/datasets/av555/aardvark-weather
```

### Training

The training entry point provides a complete pipeline with epochs, validation, early stopping, learning rate scheduling, best/latest checkpointing, and resumption. By default, the Encoder and Day-1 Processor are frozen and only the TAS Decoder is trained:

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

Paper-style end-to-end joint fine-tuning:

```bash
python scripts/train.py --train-modules all --epochs 10 --train-steps 100
```

Resuming training:

```bash
python scripts/train.py --resume weight/training/last.pth
```

The default configuration resides in `conf/config.yaml`. `--data` can point to an official-schema pickle file or a directory containing multiple `.pkl` files; multiple files are partitioned into train and validation sets in a deterministic manner. `--batch-size` concatenates multiple tasks along the existing batch dimension of the official task. When only the single official sample is included, training and validation reuse the same task — this allows end-to-end validation of the training software pipeline but does not constitute an independent validation set, nor can it provide the data diversity required to reproduce paper-level accuracy.

Training artifacts:

```text
weight/training/best.pth
weight/training/last.pth
weight/training/history.json
weight/training/train.json
```

### Inference

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

Using weights obtained from training:

```bash
python scripts/inference.py --checkpoint weight/training/best.pth
```

Inference loads the official sample and the Day-1 `tas` weights by default. Results are saved to:

```text
result/inference_one_day.json
result/prediction.pt
result/target.pt
```

### Result Inspection

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

Verified output shapes:

```text
initial_state:   [1, 121, 240, 24]
global_forecast: [1, 121, 240, 24]
station_tas:     [1, 8719]
```

Current results constitute a connectivity/end-to-end verification and do not reproduce the paper's RMSE/MAE metrics.

The result script additionally produces `result/metrics.json` and `result/comparison.png`. The `normalized_mae` and `normalized_rmse` are computed in the normalized space of the official sample and should not be directly compared to the paper's physical-unit metrics.

### Paper vs. Current Implementation I/O

| Item | Paper | Current Package |
| --- | --- | --- |
| Input | Multimodal satellite, station, ship, and radiosonde observations | Bundled official sample pickle with a field structure consistent with the official Encoder |
| Global State | `24 × 121 × 240`, 1.5° | Day-1 supported; output `[1,121,240,24]` |
| Station Output | 2 m temperature and 10 m wind, up to Day-10 | Day-1 TAS only; `[1,8719]` |
| Training | Staged pre-training followed by ~25,000 steps of end-to-end fine-tuning | Configurable full training loop; supports Decoder-only or full-model joint fine-tuning |
| Evaluation | Grid-point weighted RMSE and station MAE in physical units | MAE/RMSE in normalized sample space |

All commands should be run from the project root; `scripts/inference.py --root` converts to an absolute path. The official model internally depends on CUDA, so CPU is currently unavailable for inference. Training data must adhere to the official multimodal task dictionary schema; the current directory does not synthesize satellite or station observations — the bundled official sample serves as the default training-pipeline validator. The model package retains only official resources under `weights/` and does not carry local training artifacts under `weight/` or generated outputs under `result/`. Paper-level training still requires preparing observation data spanning the full date range and converting it into the same `.pkl` task contract.

### Real Data

Using real dates requires preparing ASCAT, AMSU-A/B, HIRS, IASI, GridSat, HadISD, ICOADS, IGRA, ERA5, topography, climatology, and their corresponding normalization statistics.

# OneScience Official 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 & License

- Official Code: https://github.com/anna-allen/aardvark-weather-public
- This directory is an independent adaptation of the official Aardvark Weather model.
- Code, weights, and data are subject to their respective official licenses and data terms.