In this first blog post, we present Py4cast, a comprehensive framework designed to facilitate the training and comparison of AI-based weather forecasting models. Py4cast aims to streamline the experimentation process, allowing researchers to efficiently test various neural network architectures and methodologies. We demonstrate the utility of Py4cast through a series of experiments using the TITAN dataset, showcasing its capability to produce accurate weather forecasts and its flexibility in accommodating different model types.

This project, built using PyTorch and PyTorch-lightning, is a Work in Progress, intended to share ideas and design concepts with partners.

It is currently developped at Météo-France by DSM/AI Lab and CNRM/GMAP/PREV, and at Eviden.

Features

  • 7 neural network architectures : Half-Unet, U-Net, SegFormer, SwinUnetR, HiLam, GraphLam, UnetR++
  • 1 dataset with samples available on Huggingface : Titan
  • 2 training strategies : Scaled Auto-regressive steps, Differential Auto-regressive steps
  • 4 losses: Scaled RMSE, Scaled L1, Weighted MSE, Weighted L1
  • neural networks as simple torch.nn.Module
  • training with pytorchlightning
  • simple interfaces to easily add a new dataset, neural network, training strategy or loss
  • simple command line to lauch a training
  • config files to change the parameters of your dataset or neural network during training
  • experiment tracking with tensorboard and plots of forecasts with matplotlib
  • implementation of NamedTensors to tracks features and dimensions of tensors at each step of the training

Future Work

Future developments will focus on expanding the TITAN dataset, testing additional architectures such as GraphCast and Pangu, and optimizing training strategies to enhance resource efficiency. We also plan to incorporate high-resolution wave models and conduct further studies on the impact of time step duration on forecast quality.

Acknowledgments

This project started as a fork of neural-lam, a project by Joel Oskarsson, see here. Many thanks to Joel for his work!