WeatherLearn

Implementation of the PyTorch version of the Weather Deep Learning Model Zoo.

Dependencies

python = "^3.11"
torch = ">=2.1.0"
timm = ">=0.9.10"
numpy = ">=1.23.5"
cdsapi = ">=0.6.1"

Install (editable):

pip install -e .
# or with Poetry after `poetry lock` / `poetry install`

Run unit tests (skips full Pangu() unless WEATHERLEARN_RUN_HEAVY=1):

python -m unittest discover -s tests -v

Model-zoo

Pangu-Weather

Model Architecture

Example

# Pangu
from weatherlearn.models import Pangu

import torch

if __name__ == '__main__':
    B = 1  # batch_size
    surface = torch.randn(B, 4, 721, 1440)  # B, C, Lat, Lon
    surface_mask = torch.randn(3, 721, 1440)  # topography mask, land-sea mask, soil-type mask
    upper_air = torch.randn(B, 5, 13, 721, 1440)  # B, C, Pl, Lat, Lon

    pangu_weather = Pangu()

    output_surface, output_upper_air = pangu_weather(surface, surface_mask, upper_air)

# Pangu_lite
from weatherlearn.models import Pangu_lite

import torch

if __name__ == '__main__':
    B = 1  # batch_size
    surface = torch.randn(B, 4, 721, 1440)  # B, C, Lat, Lon
    surface_mask = torch.randn(3, 721, 1440)  # topography mask, land-sea mask, soil-type mask
    upper_air = torch.randn(B, 5, 13, 721, 1440)  # B, C, Pl, Lat, Lon

    pangu_lite = Pangu_lite()

    output_surface, output_upper_air = pangu_lite(surface, surface_mask, upper_air)

References

@article{bi2023accurate,
  title={Accurate medium-range global weather forecasting with 3D neural networks},
  author={Bi, Kaifeng and Xie, Lingxi and Zhang, Hengheng and Chen, Xin and Gu, Xiaotao and Tian, Qi},
  journal={Nature},
  volume={619},
  number={7970},
  pages={533--538},
  year={2023},
  publisher={Nature Publishing Group}
}

@article{bi2022pangu,
  title={Pangu-Weather: A 3D High-Resolution Model for Fast and Accurate Global Weather Forecast},
  author={Bi, Kaifeng and Xie, Lingxi and Zhang, Hengheng and Chen, Xin and Gu, Xiaotao and Tian, Qi},
  journal={arXiv preprint arXiv:2211.02556},
  year={2022}
}

Fuxi

Model Architecture

Example

from weatherlearn.models import Fuxi

import torch

if __name__ == '__main__':
    B = 1  # batch_size
    in_chans = out_chans = 70  # number of input channels or output channels
    input = torch.randn(B, in_chans, 2, 721, 1440)  # B C T Lat Lon
    
    fuxi = Fuxi()
    # patch_size : Default: (2, 4, 4)
    # embed_dim : Default: 1536
    # num_groups : Default: 32
    # num_heads : Default: 8
    # window_size : Default: 7
    # depth : Default: 48  (U-Transformer Swin blocks; reduce for lite smoke)

    output = fuxi(input)  # B C Lat Lon

Last layer matches the paper: Linear patch FC → reshape to 720×1440 → bilinear interpolate to 721×1440. See docs/fuxi_last_layer.md (clarifies Issue #8).

References

FuXi: A cascade machine learning forecasting system for 15-day global weather forecast

Published on npj Climate and Atmospheric Science: FuXi: a cascade machine learning forecasting system for 15-day global weather forecast

by Lei Chen, Xiaohui Zhong, Feng Zhang, Yuan Cheng, Yinghui Xu, Yuan Qi, Hao Li

FengWu

Multi-modal encode–fuse–decode (arXiv:2304.02948). Default 13 levels / 69 channels (OpenEarthLab practical); paper 37 levels / 189 channels via n_levels=37.

from weatherlearn.models import FengWu, FengWu_lite
import torch

# Lite (smoke / small grid). ARCO 128x64 zarr is transposed to Lat×Lon (64, 128).
model = FengWu_lite()  # img_size=(64, 128), n_levels=13 → C=69
x = torch.randn(1, 69, 64, 128)
y = model(x)  # [1, 69, 64, 128]

# Optional uncertainty head (mean, log_var)
model_u = FengWu_lite(predict_uncertainty=True)
mean, log_var = model_u(x)

Train smoke (public GCS ERA5): python examples/fengwu/train_smoke_gcs.py

TODO

• FengWu Model (https://arxiv.org/pdf/2304.02948v1.pdf) — see weatherlearn/models/fengwu, examples/fengwu/
• FuXi Model (https://arxiv.org/pdf/2306.12873v3.pdf)
• Set a separate window_size for longitude and latitude in the Fuxi model.
• Add more unittest.
• Infer the Pangu model using the pre-trained weights provided by the official Pangu repository.