Adaptive Fourier Neural Operator (AFNO) for weather forecasting

This repository contains the code used for FourCastNet: A Global Data-driven High-resolution Weather Model using Adaptive Fourier Neural Operators

The code was developed by the authors of the preprint: Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, Pedram Hassanzadeh, Karthik Kashinath, Animashree Anandkumar

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∘ resolution. FourCastNet accurately forecasts high-resolution, fast-timescale variables such as the surface wind speed, precipitation, and atmospheric water vapor. It has important implications for planning wind energy resources, predicting extreme weather events such as tropical cyclones, extra-tropical cyclones, and atmospheric rivers. 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 variables with complex fine-scale structure, 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. We discuss how data-driven deep learning models such as FourCastNet are a valuable addition to the meteorology toolkit to aid and augment NWP models.

FourCastNet is based on the vision transformer architecture with Adaptive Fourier Neural Operator (AFNO) attention

FourCastNet.gif

Fig. 1 Comparison between the FourCastNet and the ground truth (ERA5) for \(u-10\) for different lead times.

The model is trained on a 20-channel subset of the ERA5 reanalysis data on single levels and pressure levels that is pre-processed and stored into HDF5 files. The subset of the ERA5 training data that FCN was trained on is hosted at the National Energy Research Scientific Computing Center (NERSC). For convenience it is available to all via Globus. You will need a Globus account and will need to be logged in to your account in order to access the data. You may also need the Globus Connect to transfer data. The full dataset that this version of FourCastNet was trained on is approximately 5TB in size.

Please refer to the reference paper to learn about the model architecture.

To train the model, run

Copy
Copied!
            

python train_era5.py

Progress can be monitored using MLFlow. Open a new terminal and navigate to the training directory, then run:

Copy
Copied!
            

mlflow ui -p 2458

View progress in a browser at http://127.0.0.1:2458

Data parallelism is also supported with multi-GPU runs. To launch a multi-GPU training, run

Copy
Copied!
            

mpirun -np <num_GPUs> python train_era5.py

If running inside a docker container, you may need to include the --allow-run-as-root in the multi-GPU run command.

If you find this work useful, cite it using:

Copy
Copied!
            

@article{pathak2022fourcastnet, title={Fourcastnet: A global data-driven high-resolution weather model using adaptive fourier neural operators}, author={Pathak, Jaideep and Subramanian, Shashank and Harrington, Peter and Raja, Sanjeev and Chattopadhyay, Ashesh and Mardani, Morteza and Kurth, Thorsten and Hall, David and Li, Zongyi and Azizzadenesheli, Kamyar and Hassanzadeh, Pedram and Kashinath, Karthik and Anandkumar, Animashree}, journal={arXiv preprint arXiv:2202.11214}, year={2022} }

ERA5 data was downloaded from the Copernicus Climate Change Service (C3S) Climate Data Store.

Copy
Copied!
            

Hersbach, H., Bell, B., Berrisford, P., Biavati, G., Horányi, A., Muñoz Sabater, J., Nicolas, J., Peubey, C., Radu, R., Rozum, I., Schepers, D., Simmons, A., Soci, C., Dee, D., Thépaut, J-N. (2018): ERA5 hourly data on pressure levels from 1959 to present. Copernicus Climate Change Service (C3S) Climate Data Store (CDS). 10.24381/cds.bd0915c6 Hersbach, H., Bell, B., Berrisford, P., Biavati, G., Horányi, A., Muñoz Sabater, J., Nicolas, J., Peubey, C., Radu, R., Rozum, I., Schepers, D., Simmons, A., Soci, C., Dee, D., Thépaut, J-N. (2018): ERA5 hourly data on single levels from 1959 to present. Copernicus Climate Change Service (C3S) Climate Data Store (CDS). 10.24381/cds.adbb2d47

Other references:

Adaptive Fourier Neural Operators: Efficient Token Mixers for Transformers

Previous ERA5 Data Downloader and Converter
Next Deep Learning Weather Prediction (DLWP) model for weather forecasting
© Copyright 2023, NVIDIA Modulus Team. Last updated on Jan 25, 2024.