Modulus Checkpoints: DLWP Cubesphere

Modulus Checkpoints: DLWP Cubesphere

Logo for Modulus Checkpoints: DLWP Cubesphere
Description
DLWP Cubesphere model checkpoint package. DLWP is a light-weight deep learning model for weather prediction that uses Cube-sphere grid.
Publisher
NVIDIA
Latest Version
v0.2
Modified
September 26, 2023
Size
61.67 MB

Details

This NGC asset is a DLWP Cubesphere model checkpoint package. Model checkpoint package refers to the set of artifacts needed to run inference using pre-trained model which includes the model checkpoint, set of sample inputs, inference script.

Architecture

DLWP is a light-weight deep learning model for weather prediction that uses Cube-sphere grid. Re-implemented from: https://agupubs.onlinelibrary.wiley.com/doi/epdf/10.1029/2021MS002502.

The DLWP model can be used to predict the state of the atmosphere given a previous atmospheric state. You can infer a 320-member ensemble set of six-week forecasts at 1.4° resolution within a couple of minutes, demonstrating the potential of AI in developing near real-time digital twins for weather prediction.

Please refer Sub-Seasonal Forecasting With a Large Ensemble of Deep-Learning Weather Prediction Models paper for additional details.

Training

The model is trained on 7-channel subset of ERA5 Data that is mapped onto a cubed sphere grid with a resolution of 64x64 grid cells. The model uses years 1980-2015 for training, 2016-2017 for validation and 2018 for out of sample testing. The training scripts for the problem can be found at: https://github.com/NVIDIA/modulus-launch/tree/main/examples/weather/dlwp

How to use?

A minimal inference script is provided in this model checkpoint package to get you started easily.

To run inference on this checkpoint, you can follow the below steps (all the files needed are included in the zip file):

  1. Launch Modulus docker container

    docker run --rm --shm-size=1g --ulimit memlock=-1 --ulimit stack=67108864 --runtime nvidia -v ${PWD}:/examples -it nvcr.io/nvidia/modulus/modulus:23.09
    
  2. Download this checkpoint zip file, and unzip it

    wget 'https://api.ngc.nvidia.com/v2/models/nvidia/modulus/modulus_dlwp_cubesphere/versions/v0.2/files/dlwp_cubesphere.zip'
    unzip dlwp_cubesphere.zip
    
  3. Run the simple inference

    cd dlwp/
    python simple_inference.py
    

Yu can also use the model checkpoint from this model checkpoint package to evalue it with the earth2mip framework.

References

Sub-Seasonal Forecasting With a Large Ensemble of Deep-Learning Weather Prediction Models

Improving data-driven global weather prediction using deep convolutional neural networks on a cubed sphere

Arbitrary-Order Conservative and Consistent Remapping and a Theory of Linear Maps: Part 1

Arbitrary-Order Conservative and Consistent Remapping and a Theory of Linear Maps, Part 2