⛈️ GeoDES (Geospatial Diffusion-Based Evolution Synthesis)

Generating high-fidelity, physically consistent synthetic weather events for research.

Welcome to GeoDES! This project uses Denoising Diffusion Probabilistic Models (DDPM) to create realistic synthetic examples of extratropical cyclones.

Project goals:

• Escape the "Mean State": Demonstrate that generative models can produce sharp, realistic weather features rather than blurred averages.
• Tail Convergence: Accurately reproduce the statistical distribution of wind speeds, especially at the extreme tail (high-impact storms).
• Physical Consistency: Ensure that generated variables (Wind, Pressure, Humidity) maintain realistic correlations.

Current status:

• Ready to run in the single climate variable case (e.g., generate only SLP or only 500hpa wind magnitude)
  • Validated to match the Power Spectral Density (PSD) of ERA5 data
  • Outperforms ClimaX and Climatology baselines in distributional metrics (FVD and KVD)
• Debugged in the multivariate case, with training and evaluations ongoing

Installation

conda create --name geodes python==3.11.11

pip install numpy torch==2.6.0 torchvision torchaudio xarray diffusers==0.32.2 tqdm transformers==4.50.0 pandas matplotlib notebook accelerate==1.5.2 opencv-python==4.11.0.86 einops==0.8.1 wandb scipy scikit-learn

Training code

Image pretraining phase:

python train_2d.py --train --epochs 1 --lr 1e-9 --dataset <dataset> --name debug2d

Video training phase:
Single-GPU:

python train_3d.py --train --epochs 1 --lr 1e-9 --dataset <dataset> --img_model debug2d

Multi-GPU:

CUDA_VISIBLE_DEVICES=0,1 accelerate launch --num_processes <num_GPUs> train_3d.py \
    --train --epochs 1 --lr 1e-9 --dataset <dataset> --img_model debug2d

Stable diffusion (older approach for baseline):

accelerate launch train_svd.py --dataset <dataset>  --output_dir <output_dir> \
    --per_gpu_batch_size=16 --gradient_accumulation_steps=1 --max_train_steps=500 \
    --channels=1 --width=32 --height=32 --checkpointing_steps=500 \
    --checkpoints_total_limit=1 --learning_rate=1e-5 --lr_warmup_steps=0 --seed=123 \
    --validation_steps=100 --num_frames=8 --mixed_precision="fp16" 

Attributions

• SVD code is from https://github.com/wangqiang9/SVD_Xtend/

Contact

Please reach out to Sonia Cromp (cromp@wisc.edu, socromp.github.io) with any questions or comments (or postdoc opportunities :))!