A benchmark to compare different deep learning models and their backbones on synthetic Navier-Stokes and real-world data from WeatherBench, published in the ICLR 2024 AI4DiffEq Workshop and in the NeurIPS 2024 Datasets and Benchmarks track:
- Comparing and Contrasting Deep Learning Weather Prediction Backbones on Navier-Stokes Dynamics (workshop paper)
- Comparing and Contrasting Deep Learning Weather Prediction Backbones on Navier-Stokes and Atmospheric Dynamics
If you find this work useful, please cite our paper
@article{karlbauer2024comparing,
title={Comparing and Contrasting Deep Learning Weather Prediction Backbones on Navier-Stokes and Atmospheric Dynamics},
author={Karlbauer, Matthias and Maddix, Danielle C and Ansari, Abdul Fatir and Han, Boran and Gupta, Gaurav and Wang, Yuyang and Stuart, Andrew and Mahoney, Michael W},
journal={arXiv preprint arXiv:2407.14129},
year={2024}
}
To install the package, first create an environment, cd into it, and install the DLWPBench package via
conda create -n dlwpbench python=3.11 -y && conda activate dlwpbench
pip install -e .
In the pip Neuraloperator package, the tucker decomposition for TFNO is not installed, so manually install the package from the source repository with
mkdir packages
cd packages
git clone https://github.com/NeuralOperator/neuraloperator
git checkout 05c01c3 # (optional) use the repository state that is compatible with checkpoints from our work
cd neuraloperator
pip install -e .
pip install -r requirements.txt
cd ../..
Moreover, install the torch-harmonics package for Spherical Fourier Neural Operators from the source repository with the following commands
cd packages
git clone https://github.com/NVIDIA/torch-harmonics.git
git checkout 13aa492
cd torch-harmonics
pip install -e .
cd ../..
To install the CUDA versions of Deep Graph Library, follow these instructions and issue
pip uninstall dgl -y
pip install dgl -f https://data.dgl.ai/wheels/cu121/repo.html
pip install dglgo -f https://data.dgl.ai/wheels-test/repo.html
Important
This DGL version requires CUDA 12.1 to be installed, e.g., following these instructions
Finally, change into the benchmark directory, which will be considered the root directory in the following, that is, cd src/dlwpbench
To generate data and run experiments in the synthetic Navier-Stokes environment, please go to the respective subdirectory and follow the steps detailed there.
To download and preprocess data and run experiments in the real-world WeatherBench environment, please go to the respective subdirectory and follow the steps detailed there.
Deep learning model repositories that are used in this study:
- HEALPix remapping: https://github.com/CognitiveModeling/dlwp-hpx
- Convolutional LSTM: https://github.com/ndrplz/ConvLSTM_pytorch/blob/master/convlstm.py
- Fourier Neural Operator: https://github.com/neuraloperator/neuraloperator
- FourCastNet: https://github.com/NVlabs/FourCastNet
- Spherical Fourier Neural Operator: https://github.com/NVIDIA/torch-harmonics
- SwinTransformer: https://github.com/microsoft/Swin-Transformer/tree/main
- Pangu-Weather: https://github.com/lizhuoq/WeatherLearn/blob/master/weatherlearn/models/pangu/pangu.py
- MeshGraphNet: https://github.com/NVIDIA/modulus/tree/main/modulus/models/meshgraphnet
- GraphCast: https://github.com/NVIDIA/modulus/tree/main/modulus/models/graphcast
See CONTRIBUTING for more information.
This project is licensed under the Apache-2.0 License.