Precipitation Downscaling with Spatiotemporal Video Diffusion (STVD)

Official codebase for the NeurIPS 2024 paper Precipitation Downscaling with Spatiotemporal Video Diffusion.

STVD is a probabilistic video super-resolution method for precipitation downscaling. A deterministic downscaler first super-resolves a low-resolution precipitation sequence, and a temporally-conditioned diffusion model then refines it by modeling the residual fine-grained details. Both modules use spatio-temporal factorized attention. Evaluated on FV3GFS global atmosphere model output against six state-of-the-art baselines, STVD better captures precipitation distributions and extreme events (e.g. heavy rain).

STVD pipeline: a deterministic downscaler super-resolves the low-resolution sequence, and a conditional diffusion model captures the residual details. Blue blocks are shared between training and inference, red blocks are training-only, and green blocks are inference-only.

Repository structure

Path	Description
config.py	All training/sampling hyperparameters and data settings
trainer.py	Entry point for training
sampler.py	Entry point for sampling from a trained checkpoint
model/model.py	Model definitions (UNet, deterministic downscaler, Gaussian diffusion, trainer)
data/	Dataset/dataloader code and scripts to compute normalization statistics
fv3net.yml	Conda environment specification

Setup

conda env create -f fv3net.yml
conda activate fv3net

The environment is an exact export (Linux, CUDA GPU). If some build pins fail to resolve on your system, install the key dependencies manually: PyTorch (CUDA), xarray, zarr, ml_collections, accelerate, and einops.

Data

The model is trained on an 11-member ensemble of FV3GFS output, with coarse-resolution (c48) inputs and high-resolution (c384) precipitation targets stored as zarr archives:

• c48_precip_plus_more_ave/<member>/sfc_8xdaily_ave_coarse.zarr — coarse surface fields (inputs)
• c48_atmos_ave/<member>/atmos_8xdaily_ave_coarse.zarr — coarse atmospheric fields (inputs)
• c384_precip_ave/<member>/sfc_8xdaily_ave.zarr — high-resolution precipitation (targets)
• c384_topo/<member>/atmos_static.zarr — high-resolution topography

Point PATH at the top of data/vsrdata_ensemble.py to the directory containing these archives.

The dataloader also expects normalization statistics (min/max per channel) as pickle files under data/ensemble_c48_trainstats/ and data/ensemble_c384_trainstats/. These can be generated with the data/ensemble*trainstats.py scripts.

Usage

Training and sampling are configured through config.py (model size, diffusion steps, dataset settings, output folders, etc.).

Train:

python trainer.py

Checkpoints are written to config.results_folder every config.save_and_sample_every steps.

Sample from a trained checkpoint (set config.milestone to the checkpoint index first):

python sampler.py

Pretrained models

Trained model checkpoints will be released later.

Acknowledgements

Part of the code is adapted from lucidrains/denoising-diffusion-pytorch.

Citation

@inproceedings{srivastava2024precipitation,
  title={Precipitation Downscaling with Spatiotemporal Video Diffusion},
  author={Srivastava, Prakhar and Yang, Ruihan and Kerrigan, Gavin and Dresdner, Gideon and McGibbon, Jeremy and Bretherton, Christopher and Mandt, Stephan},
  booktitle={Advances in Neural Information Processing Systems},
  year={2024}
}