OneScience-Group/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:

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

Use Cases

Usage

1. OneCode

Click to experience intelligent one-click AI4S programming

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

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

Set Up the Runtime Environment

DCU Environment

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

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:

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:

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:

python scripts/train.py

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

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

Resuming training:

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:

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

Inference

python scripts/inference.py

Using weights obtained from training:

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:

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

Result Inspection

python scripts/result.py

Verified output shapes:

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

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

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.