• [2026.6.17] 🙌 Training and evaluation code is now available! See below for details about how to launch training runs and evaluate models.
• [2026.6.07] 👀 The ClimateSuite dataset has been publicly released! Training and evaluation code will be released soon.
• [2026.2.20] 🔥 Our work has been accepted to CVPR 2026!
• [2025.12.02] 🎉 Our paper is available! Dataset and code will be released soon. Please feel free to watch 👀 this repository for updates.

• Highlights
• Main Results
• Environment Setup
• Data Preparation
• Training & Validating
• License
• Acknowledgements
• Citation

Spatiotemporal Pyramid Flows (SPFs) are a new class of flow matching approaches to efficiently generate samples of future climate trajectories at different timescales.

SPF divides generation into stages, each beginning with DiT denoising and followed by either a spatiotemporal transition (green) or a spatial-only transition (orange). Spatiotemporal transitions funnel into a timestep for the selected target period and upsample the latent in both space and time, while spatial transitions upsample only in space. This sequence of denoising and stage transitions continues until the final stage, which outputs clean samples at the target period and timescale.

We introduce a new dataset for climate emulation called ClimateSuite which we use to train a scaled version of SPF. ClimateSuite, comprises more than 33,000 simulation-years of climate data spanning 276 state-of-the-art simulations from 10 ESMs and 39 stratospheric aerosol injection (SAI) simulations.

Comparison of ClimateSuite to existing climate-scale datasets.

Climate model and scenario breakdown in ClimateSuite.

We demonstrate that SPFs:

• obtain superior accuracy and inference efficiency compared to strong deterministic baselines, pre-trained models, and flow matching approaches on ClimateBench.
• achieve good generalization to emissions and intervention scenarios across climate models when trained on ClimateSuite.
• obtain further improved performance on ClimateBench after fine-tuning a model pre-trained on ClimateSuite.

ClimateBench test metrics on held-out scenario (SSP2-4.5).

Effect of scale and ClimateSuite pre-training on ClimateBench performance.

Yearly metrics on held-out scenarios across climate models in ClimateSuite.

The maintained code path supports SPF plus the included PyramidFlow, PixelFlow, and single-timescale SPF baselines. We recommend using uv for a clean local environment:

./scripts/create_uv_env.sh
source .venv/bin/activate

The setup script uses pyproject.toml and uv.lock, installs CUDA-enabled PyTorch, and runs an import smoke test. If uv is not already on your PATH, install it first:

curl -LsSf https://astral.sh/uv/install.sh | sh

Training and evaluation configs load the private Hugging Face dataset jirvin16/ClimateSuite and extract the required archive-sharded Zarr files automatically. Authenticate before the first run:

huggingface-cli login

An HF_TOKEN environment variable with dataset access also works in non-interactive jobs.

Use a scratch or high-throughput filesystem for the dataset cache; otherwise the download may fill your home directory.

export CLIMATESUITE_CACHE_DIR=/path/to/fast/climatesuite_hf_cache

You can also override the cache per Hydra command:

datamodule.hf_dataset_cache_dir=/path/to/fast/climatesuite_hf_cache/

Evaluation metrics use an on-disk cache to avoid holding all predictions in memory. Configure it with either SPF_EVAL_CACHE_DIR or eval_cache_dir=....

export SPF_EVAL_CACHE_DIR=/path/to/fast/eval_cache

Use logger=none for local runs that should not write to WandB.

Train SPF-4:

HYDRA_FULL_ERROR=1 uv run python emulator/run.py \
  experiment=spf/climatebench_spf4_train.yaml \
  datamodule.hf_dataset_cache_dir=/path/to/fast/climatesuite_hf_cache/ \
  logger=none

Train the ClimateSuite SPF configs:

HYDRA_FULL_ERROR=1 uv run python emulator/run.py experiment=spf/climatesuite_spf4_train.yaml logger=none
HYDRA_FULL_ERROR=1 uv run python emulator/run.py experiment=spf/climatesuite_spf12_train.yaml logger=none

Baseline configs are available under experiment=baseline/....

• Pyramid Flows The model we built upon.
• ClimateSet The codebase and dataset we built upon.

• This project is released under the Apache 2.0 license as found in the LICENSE file.

If you find our paper and code useful in your research, please consider giving a star ⭐ and citation 📝.

@article{irvin2025spatiotemporal,
  title={Spatiotemporal Pyramid Flow Matching for Climate Emulation},
  author={Irvin, Jeremy Andrew and Han, Jiaqi and Wang, Zikui and Alharbi, Abdulaziz and Zhao, Yufei and Bayarsaikhan, Nomin-Erdene and Visioni, Daniele and Ng, Andrew Y. and Watson-Parris, Duncan},
  journal={arXiv preprint arXiv:2512.02268},
  year={2025}
}