LonghaoWang/WeatherAI

Synced from GISWLH/WeatherAI — code, tests and issues live on GitHub. No model weights are stored here: official weights stay with their authors and are downloaded on demand under their own licences.

English · 简体中文 · 日本語

WeatherAI re-implements the AI weather and climate models published in Nature, Science and their sister journals as plain, readable PyTorch nn.Modules, and checks every port against the authors' own code. Load the official weights, compare with the original, fine-tune, or start from a Lite config that runs on a laptop.

✨ Highlights

19 Oct 2026 research update: PBC probabilistic S2S calibration, GenFocal official pretrained downscaling, and gray-swan translocation diagnostics use pinned official code and published weights by default, with explicit separate lite variants. Actual execution reports distinguish pretrained GPU inference from the earlier lite-only synthetic smoke. WeatherPre defaults to the official entries. PBC and the diagnostic wrapper extend the existing 20-model zoo; they are not two additional pretrained weather forecasters.

from weatherai import PBC, PBC_lite, GenFocal, GenFocal_lite, Translocation, Translocation_lite

• 🧩 20 models, one library — Pangu, GraphCast, Aurora, GenCast, NeuralGCM, FuXi-S2S, CAS-Canglong and more, all with the same import style and plain tensors in / tensors out. No per-model environments to juggle.
• 🔀 Multi-framework — native PyTorch throughout; official JAX code is kept where it matters (NeuralGCM training and precipitation) and as the reference for GraphCast, GenCast and WeatherNext; ONNX-only releases (FuXi-ENS, FuXi-S2S) are rebuilt as PyTorch modules. Runs on CPU and CUDA GPUs.
• ✅ Verified, not just rewritten — official checkpoints strict-load for 14 models, and outputs match the official implementation (often bit-for-bit). Every claim has a script and a JSON result; GPU runs are tested on an NVIDIA Blackwell GPU through a Hugging Face ZeroGPU Space.
• 🪶 Full size and Lite — the paper-size network for real experiments, plus Lite configs (same code, small grid) for 15 models that train on a laptop CPU or a single small GPU.

🚀 Quick start

pip install git+https://github.com/GISWLH/WeatherAI.git      # or: git clone ... && pip install -e ".[dev]"

import torch
from weatherai.models import Canglong_lite, GraphCast_lite, Aurora_small

sst  = Canglong_lite()(torch.randn(1, 8, 8, 32, 64))    # 3-month SST outlook -> (1, 1, 3, 32, 64)
step = GraphCast_lite()(torch.randn(1, 4, 32, 64))      # one 6-hour step on a tiny icosahedral mesh
aurora = Aurora_small(pretrained=True)                  # official weights (~450 MB), strict-loaded into native PyTorch

🗺️ Model zoo

🔓 official weights load · 🎯 matches the official code · ⚡ tested on a CUDA GPU · 🪶 Lite config · 🆕 new. Each docs link opens a one-page summary: usage, what was verified, and what was not.

🌍 Global medium-range forecasting

🎲 Ensemble & generative

🌊 Sub-seasonal to seasonal, climate & ocean

🌀 Regional & tropical cyclones

More: all model pages · verification log · roadmap & planned models

🔬 What "verified" means

For each port, the unchanged official code (PyTorch, JAX or onnxruntime) and the native module receive the same weights and inputs, and their outputs are compared; the scripts are in scripts/ and the numbers in docs/results/. GPU runs use scripts/gpu_smoke.py on the Hugging Face Space. These are implementation checks, not forecast-skill claims; each model page lists what was not verified.

🤗 Hugging Face

• Space — LonghaoWang/weatherai-graphcast-smoke: ZeroGPU smoke tests for every model.
• Hub — LonghaoWang/WeatherAI: model card and model pages, synced from this repository by .github/workflows/sync_hf.yml.
• Official weights stay with their authors and keep their licences (Apache-2.0, MIT, CC-BY, CC-BY-NC(-ND), …): WeatherAI downloads them on demand and never redistributes them.

🧪 Tests

pip install -e ".[dev]"
pytest tests/models -q          # unit tests; parity references are stored under tests/models/*/data

📖 Citation

Please cite the original paper of each model you use (linked above) and, if WeatherAI helped, this repository (CITATION.cff):

@software{wang_weatherai,
  author = {Wang, Longhao},
  title  = {{WeatherAI}: a PyTorch model zoo of AI weather and climate models},
  url    = {https://github.com/GISWLH/WeatherAI},
  year   = {2026}
}

🤝 Contributing

Issues and pull requests are welcome. A new model is accepted when it comes with a native module, a Lite config, a parity script against the official code, unit tests and a page in docs/models/. If WeatherAI saves you time, a ⭐ helps others find it.

License

MIT — Copyright (c) 2026 Longhao Wang (GISWLH); adapted code from WeatherLearn (MIT), see NOTICE. Model weights are under their authors' licences. Research and educational code: forecast skill depends on data, training and weights, and nothing here replaces an operational forecasting system.

Editable banner artwork · Excalidraw

Demonstration captured on EHPC · 17 October 2026

New visualization of the repository's existing parity-check JSON. This displays a recorded historical check; it is not a new training or model-validation run.

Execution record and checksums.