Introduction
FourCastNet (https://github.com/NVlabs/FourCastNet), short for Fourier ForeCasting Neural Network, is a global data-driven weather forecasting model that provides accurate short to medium-range global predictions at 0.25 degree resolution, which corresponds to a spatial resolution of roughly 30 km × 30 km near the equator and a global grid size of 720 × 1440 pixels. The purpose of FourCastNet is to accurately forecast high-resolution, fast-timescale climate variables such as surface wind speed, precipitation, and atmospheric water vapour. FourCastNet is about 45,000 times faster than traditional NWP models on a node-hour basis. This orders of magnitude speedup, along with the unprecedented accuracy of FourCastNet at high resolution, enables the inexpensive generation of extremely large ensemble forecasts. This approach dramatically improves probabilistic weather forecasting.
Research Prospect
Many studies have attempted to develop data-driven deep learning (DL) models for weather forecasting, either trained on climate model outputs, general circulation models (GCM), reanalysis products, or a blend of climate model outputs and reanalysis products. Some challenges researchers face when trying to develop DL models for weather forecasting include: lack of high-quality observational data sets, limited understanding of how the atmospheric dynamics interact with each other, and difficulties in incorporating external factors into the models. Data-driven models have great potential to improve weather predictions by overcoming model biases present in NWP models and by enabling the generation of large ensembles at low computational cost for probabilistic forecasting and data assimilation.
The FourCastNet team at NVLabs published a journal paper in 2022 with details of their work and results (https://arxiv.org/pdf/2202.11214.pdf) of using a Fourier-based neural network forecasting model. Their results show that FourCastNet matches the forecasting accuracy of the ECMWF Integrated Forecasting System (IFS), a state-of-the-art Numerical Weather Prediction (NWP) model, at short lead times for large-scale variables, while outperforming IFS for small-scale variables, including precipitation. FourCastNet generates a week-long forecast in less than 2 seconds, orders of magnitude faster than IFS.
The speed of FourCastNet enables the creation of rapid and inexpensive large-ensemble forecasts with thousands of ensemble-members for improving probabilistic forecasting. FourCastNet has shown impressive results on predicting various variables on ERA5. While FourCastNet enjoys quasi-linear time and memory complexity in sequence length compared to quadratic complexity in vanilla transformers, training FourCastNet on ERA5 from scratch still requires large amount of compute resources which is expensive or even inaccessible to average researchers.
Wind and Precipitation prediction in FourCastNet notebook
Figure 1 below shows the multilayer transformer architecture which utilises the Adaptive Fourier Neural Operator with shared MLP and frequency soft-thresholding for spatial token mixing. The input frame (720x1440 lat-lon grid) is first projected to a 2D grid (h×w) of patches, where each patch has a small size (p×p×c). Each patch is then embedded in a higher dimensional space with high number of latent channels (AFNO layers) and position embedding is added to form a sequence of tokens. Each layer receives an input tensor of patches 𝑋∈ℝh*w*c and performs spatial mixing followed by channel mixing. The spatial mixing happens in the Fourier domain.
Figure 1: FourCastNet Model Architecture
The notebook replicates the prediction for wind and precipitation variables found in the paper: https://arxiv.org/pdf/2202.11214.pdf. Figure 2 is an illustrative example of the prediction of atmospheric rivers using FourCastNet. It is compared against the ground truth presented in the ERA5 dataset. The animation presents results of the prediction at different time steps 6 hours apart for up to 90 hours. The code to achieve this prediction can be found in the notebook. A local repository of FourCastNet code and our code for using on NCI infrastructure and data are available at this GitHub repository: https://github.com/nci/FourCastNet.
Figure 2: Prediction of Atmospheric Rivers vs Ground Truth
The FourCastNet model uses the ERA5 dataset. The notebook starts by implementing normalised statistics prepared by NVLabs to replicate the results in the paper https://arxiv.org/pdf/2202.11214.pdf. This is followed by running the pertained model on the ERA5 dataset. The raw ERA5 datasets used in this inference notebook are located in our ERA5 data collection in "/g/data/rt52/era5/single-levels/reanalysis".
FourCastNet Demo
You may find more information about running FourCastNet on Gadi using our example.