FourCastNet

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

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 proposed in Guibas-Mardani et al. [paper], [code].

Quick Links:

Pre-processed Training Data - Globus Download Link

Trained Model Weights - Globus Download Link

Trained Model Weights - Web Download Link

Version Notes:

| Release Version | Model Weights | Training Data | Normalization | |----------------------|---------------|-----------------|--------------| | Initial release v0.0.0 | FCN_weights_v0| FCN_ERA5_data_v0| stats_v0 |

Training:

The model is trained on a subset of ERA5 reanalysis data on single levels [ Hersbach 2018 ] and pressure levels [ Hersbach 2018 ] 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 Resarch Scientific Computing Center (NERSC). For convenience it is available to all via Globus at the following link.

Pre-processed Training Data

You will need a Globus account and will need to be logged in to your account in order to access the data. The full dataset that this version of FourCastNet was trained on is approximately 5TB in size.

The data directory is organized as follows:

FCN_ERA5_data_v0
│   README.md
└───train
│   │   1979.h5
│   │   1980.h5
│   │   ...
│   │   ...
│   │   2015.h5
│   
└───test
│   │   2016.h5
│   │   2017.h5
│
└───out_of_sample
│   │   2018.h5
│
└───static
│   │   orography.h5
│
└───precip
│   │   train/
│   │   test/
│   │   out_of_sample/

Precomputed stats are provided at additional and have the directory structure:

stats_v0
│   global_means.npy  
│   global_stds.npy  
│   land_sea_mask.npy  
│   latitude.npy  
│   longitude.npy  
│   time_means.npy
│   time_means_daily.h5
└───precip
│   │   time_means.npy

Training configurations can be set up in config/AFNO.yaml. The following paths need to be set by the user. These paths should point to the data and stats you downloaded in the steps above:

afno_backbone: &backbone
  <<: *FULL_FIELD
  ...
  ...
  orography: !!bool False 
  orography_path: None # provide path to orography.h5 file if set to true, 
  exp_dir:             # directory path to store training checkpoints and other output
  train_data_path:     # full path to /train/
  valid_data_path:     # full path to /test/
  inf_data_path:       # full path to /out_of_sample. Will not be used while training.
  time_means_path:     # full path to time_means.npy
  global_means_path:   # full path to global_means.npy
  global_stds_path:    # full path to global_stds.npy

An example launch script for distributed data parallel training on the slurm based HPC cluster perlmutter is provided in submit_batch.sh. Please follow the pre-training and fine-tuning procedures as described in the pre-print.

To run the precipitation diagnostic model, see the following example config:

precip: &precip
  <<: *backbone
  ...
  ...
  precip:              # full path to precipitation data files 
  time_means_path_tp:  # full path to time means for precipitation
  model_wind_path:     # full path to backbone model weights ckpt

Inference:

In order to run FourCastNet in inference mode you will need to have the following files on hand.

1. The path to the out of training sample hdf5 file. This could either be a new set of initial conditions that you downloaded from copernicus and processed yourself (see separate instructions for doing so in the next section), or it could be out_of_sample dataset hosted here. The inference script provided assumes that you are using the out_of_sample/2018.h5 file. You can modify the script to use a different h5 file that you processed yourself after downloading the raw data from Copernicus.
2. The model weights hosted at Trained Model Weights

FCN_weights_v0/
│   backbone.ckpt  
│   precip.ckpt  

3. The pre-computed normalization statistics hosted at additional. It is crucial that you use the statistics that are provided if you are using the pre-trained model weights that we have provided since these stats were used when trainig the model. The normalization statistics go hand-in-hand with the trained model weights. The stats folder contains:

stats_v0
│   global_means.npy  
│   global_stds.npy  
│   land_sea_mask.npy  
│   latitude.npy  
│   longitude.npy  
│   time_means.npy
│   time_means_daily.h5

Once you have all the file listed above you should be ready to go.

In config/AFNO.yaml, set the user defined paths

afno_backbone: &backbone
  <<: *FULL_FIELD
  ...
  ...
  orography: !!bool False 
  orography_path: None # provide path to orography.h5 file if set to true, 
  inf_data_path:       # full path to /out_of_sample. Will not be used while training.
  time_means_path:     # full path to time_means.npy
  global_means_path:   # full path to global_means.npy
  global_stds_path:    # full path to global_stds.npy

Run inference using

python inference/inference.py \
       --config=afno_backbone \
       --run_num=0 \
       --weights '/path/to/weights/backbone.ckpt' \
       --override_dir '/path/to/output/scratch/directory/ \' 

Run inference for precipitation using

python inference/inference_precip.py \
       --config=precip \
       --run_num=0 \
       --weights '/path/to/weights/precip.ckpt' \
       --override_dir '/path/to/output/scratch/directory/ \' 

Additional information on batched ensemble inference and precipitation model inference can be found at inference/README_inference.md

The outputs of the inference scripts will be written to an hdf5 file at the path specified in the --override_dir input argument. Depending on the params set in the config file, the output file will contain the computed ACC and RMSE of the forecasts and the raw forecasts of selected fields for visualization.

Inference for a custom interval

The steps will walk you through:

1. Downloading an initial condition from the (continuously expanding) ERA5 da