Code for Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos published at ICML 2025
To generate training data, we use jax_cfd for Kolmogorov flow and torch_qg for quasi-geostrophic flows.
To install in developer mode:
pip install -e .
To install otherwise:
pip install .
The codebase is heavily integrated with wandb: for each training and inference run, a new wandb experiment is created, and all config metadata and results plots are uploaded.
The training loop for the stochastic residual emulator is here. The core training algorithm for the thermalizer is found here. and here are points to inference algorithms for thermalized Kolmogorov and QG flows.
Please cite as:
@article{pedersen2025thermalizer,
doi = {10.48550/ARXIV.2503.18731},
url = {https://arxiv.org/abs/2503.18731},
author = {Pedersen, Christian and Zanna, Laure and Bruna, Joan},
title = {Thermalizer: Stable autoregressive neural emulation of spatiotemporal chaos},
publisher = {arXiv},
year = {2025},
}