IGRA-Gen: Multimodal Atmospheric Super-Resolution with Deep Generative Models

This repository contains the implementation of IGRA-Gen, a diffusion-based framework for multimodal atmospheric super-resolution and data fusion. The method leverages score-based diffusion models to generate high-resolution atmospheric states from sparse, multimodal observations.

🏗️ Installation

Clone the repository:

git clone https://github.com/dibyajyoti-chakraborty/igra-gen.git
cd igra-gen

Install dependencies:

pip install -r requirements.txt

Or install latest PyTorch version on a CUDA compatible device and the required packages.

🚀 Usage

Training

Train a diffusion model on ERA5 data using the provided script:

cd scripts
sbatch submit-cond.sh

This script trains the model on 4 nodes with 4 GPUs each for 48 hours. It is sufficient for the 13 variable model. The 69 variable model will take longer.

Sampling

Generate conditional samples using the unified sampling script:

# Unconditional sampling
python src/igra_gen/sample_cond.py --sampling_type unconditional --ens 16

# Super-resolution from coarse ERA5 grid
python src/igra_gen/sample_cond.py --sampling_type sr --ens 16

# IGRA radiosonde conditioning
python src/igra_gen/sample_cond.py --sampling_type igra --ens 16

# Multimodal fusion (SR + IGRA)
python src/igra_gen/sample_cond.py --sampling_type sr_igra --ens 16

This will save required samples in samples/. The plots in the samples/ directory were generated with the default hyperparameters. Some of the plots may vary from the original paper for stochasticity and different sampling hyperparameters. For optimal results, refer to the paper for hyperparameter tuning guidance during generation. The lr_to_sr.ipynb is with the 69 variable ERA5 model and the other notebooks are with the 13 variable model.

📚 Citation

If you use this code, please cite our paper:

@article{chakraborty2025multimodal,
  title={Multimodal atmospheric super-resolution with deep generative models},
  author={Chakraborty, Dibyajyoti and Guan, Haiwen and Stock, Jason and Arcomano, Troy and Cervone, Guido and Maulik, Romit},
  journal={Machine Learning: Earth},
  year={2025}
}