FourCastNet Inference

This repository is forked from the official FourCastNet. For references and citations, please refer to the official repository.

In addition to the official FourCastNet, this repository contains code that performs inference on NCI version of ERA5 from project rt52.

Additional documentation for FourCastNet may be found on the FourCastNet documentation.

Setup

• Ask to join NCI project rt52 on mancini.

• Run bash setup.sh to set up the environment. This script sets up a Python virtualenv with all the required dependencies.

• The inference requires pretrained weights and input normalization statistics. The pretrained weights are at /g/data/wb00/FourCastNet/nvlab/v0/pretrained and the input normalization statistics are at /g/data/wb00/FourCastNet/nvlab/v0/data/stats.

Notebooks

We also provided several Jupyter notebooks under the notebooks directory. These notebooks demonstrate:

• Baseline inference on the out-of-distribution sample data provided by the original authors The code for baseline inference is under the inference directory.

• Inference on NCI ERA5 The code for inference on NCI ERA5 is under the inference_nci directory.

Data Location

The data used in the FourCastNet project is available at /g/data/wb00/FourCastNet/. This dataspace contains two directories: nci and nvlabs. The nci directory contains the predictions and models generated by NCI. The nvlabs folder contains the data and models provided by nvlabs on their repository. More information may be found on the FourCastNet documentation and notebook itself.

Inference on NCI ERA5

The inference algorithm on NCI ERA5 works as follows:

INPUTS:  start_time
         end_time
         prediction_length
         initial conditions (i.e. raw ERA5 data)

OUTPUTS: predictions in zarr format

let t = start_time

while t <= end_time:
  load initial conditions at t to initialize FourCastNet

  forecast from t+1 to t+1 + prediction_length

  t += t+1 + prediction_length

The inference code is tested to support both CPU and GPU. With GPU support, we use CUDA 11.7 which can be loaded via module load cuda/11.7.0 on Gadi.

A demonstration to run the inference script inference_nci/inference.py is test_inference.sh.

• To run test_inference.sh, one needs to modify checkpoint_dir and stats_dir to point to the right location.

• With a short prediction length specified in test_inference.sh, one can run the script directly from Gadi login node to get quick results in a few minutes.