Skip to content

sig-gis/neural-operators-weather-downscaling

 
 

Repository files navigation

On the effectiveness of neural operators at zero-shot weather downscaling

Overview

Overall, our work investigates the zero-shot downscaling potential of neural operators. To summarize, our contributions are:

  1. We provide a comparative analysis based on two challenging weather downscaling problems, between various neural operator and non-neural-operator methods with large upsampling factors (e.g., 8x and 15x) and fine grid resolutions (e.g., 2 km × 2 km wind speed).

  2. We examine whether neural operator layers provide unique advantages when testing downscaling models on upsampling factors higher than those seen during training, i.e., zero-shot downscaling. Our results instead show the surprising success of an approach that combines a powerful transformer-based model with a parameter-free interpolation step at zero-shot weather downscaling.

  3. We find that this Swin-Transformer-based approach mostly outperforms all neural operator models in terms of average error metrics, whereas an enhanced super-resolution generative adversarial network (ESRGAN)-based approach is better than most models in capturing the physics of the system, and suggests their use in future work as strong baselines. However, these approaches still do not capture variations at smaller spatial scales well, including the physical characteristics of turbulence in the HR data. This suggests a potential for improvement in transformer or GAN-based methods and neural-operator-based methods for zero-shot weather downscaling.

Framework Overview

Overview of the neural operator and non-neural-operator zero-shot weather downscaling approaches. We show 5x to 15x zero-shot downscaling as an example.

This codebase builds upon:

Setup

1. Clone the Repository

git clone https://github.com/NREL/neural-operators-weather-downscaling.git
cd neural-operators-weather-downscaling

2. Create the Conda Environment

conda env create -f environment.yml
conda activate downscaling

Usage

1. Download the Dataset

We work on two weather downscaling problems. Download the appropriate dataset(s) for the experiments:

  • ERA5 (Low Resolution) → ERA5 (High Resolution)
    We use the public dataset released by the SuperBench paper by Pu Ren et al., 2023 :
    🔗 ERA5-to-ERA5 dataset

  • ERA5 (Low Resolution) → WTK (High Resolution)
    We have released the dataset used in these experiments here:
    🔗 ERA5-to-WTK dataset
    See the link for more details.

After downloading, we structure the folder like this:

../datasets/
├── era_to_wtk/
│   ├── region1/
│   └── region2/
└── era5/

2. Model Training

Baseline models are trained using the script located at:

scripts/baseline/train.sh → calls scripts/baseline/train.py
Models trained:

  • SRCNN
  • EDSR
  • ESRGAN
  • SwinIR

Neural operator models are trained using the script located at:

scripts/neuraloperator/train.sh → calls scripts/neuraloperator/train.py
Models trained:

  • FNO
  • DFNO
  • DUNO
  • DAFNO
  • DCNO

Run training using:

sbatch scripts/baseline/train.sh         # for baseline models
sbatch scripts/neuraloperator/train.sh   # for neural operator models

For descriptions of all the baseline and neural operator models, please refer to our paper.

Note: We configure .sh scripts for SLURM-based scheduling and all models are trained on a single GPU node.

The train.sh scripts contain hyperparameter details including references to configuration files.

Model implementations are located in:

  • src/baseline_models/ — for SRCNN, EDSR, ESRGAN, SwinIR
  • src/neuraloperators/neuralop/models/ — for FNO, DFNO, DUNO, DAFNO, DCNO

3. Model Evaluation

Refer to scripts/baseline/eval.sh and scripts/neuraloperator/eval.sh for evaluating baseline and neural operator models, respectively. Each eval.sh script internally calls its corresponding Python script—scripts/baseline/eval.py or scripts/neuraloperator/eval.py.

To perform zero-shot evaluation, specific zero-shot arguments are included when calling these Python scripts.
Please refer to the respective .sh files for details on how these arguments are configured.

Run evaluation using:

sbatch scripts/baseline/eval.sh         # for baseline models
sbatch scripts/neuraloperator/eval.sh   # for neural operator models

4. Result Analysis

The result_analysis/ folder contains utilities for visualizing model outputs and energy spectra of the downscaled outputs.

  • plot_outputs.py: Visualizes sample downscaled outputs for comparison across models.
  • energy_spectrum.py: Computes and compares the kinetic energy spectrum plots for each model which shows a distribution of energy across various wavenumbers.

For detailed results and analysis, please refer to our paper.

Citation

If you use this codebase or ideas from this project in your work, please cite:

@article{sinha2025effectiveness,
  title={On the effectiveness of neural operators at zero-shot weather downscaling},
  author={Sinha, Saumya and Benton, Brandon and Emami, Patrick},
  journal={Environmental Data Science},
  volume={4},
  pages={e21},
  year={2025},
  publisher={Cambridge University Press}
}

About

No description, website, or topics provided.

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages

  • Python 82.3%
  • Cuda 11.7%
  • C++ 3.8%
  • Shell 1.5%
  • Other 0.7%