This repository contains the source code for the paper Benchmarking Autoregressive Conditional Diffusion Models for Turbulent Flow Simulation by Georg Kohl, Liwei Chen, and Nils Thuerey. Our work benchmarks neural operators for the prediction of turbulent flow fields from an initial condition across a range of architectures:

1. Autoregressive conditional diffusion models (ACDMs): This method relies on the DDPM approach, a class of generative models based on a parameterized Markov chain. They can be trained to learn the conditional distribution of a target variable given a conditioning. In our case, the target variable is the flow field at the next time step, and the conditioning is the flow field at the current time step, i.e., the simulation trajectory is created via autoregressive unrolling of the model. We showed that ACDMs can accurately and probabilistically predict turbulent flow fields, and that the resulting trajectories align with the statistics of the underlying physics. Furthermore, ACDMs can generalize to flow parameters beyond the training regime, and exhibit high temporal rollout stability, without compromising the quality of generated samples.

2. U-Net Variants: We consider a range of different training variants of the backbone U-Net used in ACDM. It can be trained as a classic next-step predictor, with unrolled training to improve the stability, or with additional training noise. Furthermore, we also implement PDE-Refiner based on this U-Net.

3. ResNets and Fourier Neural Operators: Our source code also integrates Dilated ResNets and Fourier Neural Operators, as established flow prediction architectures.

4. Latent-Space Transformers: Finally, we provide different implementation variants of latent-space transformers. These architectures contain an autoencoder to embed the flow states in a smaller latent space. A secondary transformer model computes updated latent spaces forward in time, that can be decoded to create flow trajectories. We train the combined architecture end-to-end.

All architectures are evaluated on three carefully designed data sets of varying difficulty:

1. Incompressible Wake Flow (Inc): This relatively simple data set contains 91 incompressible wake flow trajectories simulated with PhiFlow at Reynolds numbers 100-1000. Each sequence contains 1300 temporal snapshots of pressure and velocity.

2. Transonic Cylinder Flow (Tra): This data set contains 41 sequences of a two-dimensional transonic cylinder flow simulated with the fluid solver SU2. Each sequence has a different Mach number between 0.5 and 0.9 at Reynolds number 10000, making the flows in the transonic regime highly chaotic. All trajcetories feature 1000 temporal snapshots of pressure, density, and velocity.

3. Isotropic Turbulence (Iso): This highly complex, underdetermined data set consists of 1000 two-dimensional sequences sliced from three-dimensional, isotropic turbulence simulations. The flows were computed with DNS and the raw data is provided by the Johns Hopkins Turbulence Database. Each sequence contains 1000 temporal snapshots of pressure and velocity, including the velocity z-component.

Further information about this work can be also found at our project website. Feel free to contact us if you have questions or suggestions regarding our work or the source code provided here.

Using the button above, you can run a simple example notebook in Google Colab that demonstrates ACDM (a Google account is required). Instead, it is also possible to locally run the provided acdm-demo.ipynb file, by following the installation instructions in the next section and running the notebook inside the created conda environment afterwards.

In the following, Linux is assumed as the OS but the installation on Windows should be similar.

We recommend to install the required python packages (see requirements.yml) via a conda environment (e.g. using miniconda), but it may be possible to directly install them with pip (e.g. via venv for a separate environment) as well.

conda env create -f requirements.yml
conda activate ACDM

In the following, all commands should be run from the root directory of this source code.

The directory src/turbpred contains the general code base of this project. The src/lsim directory contains the LSiM metric that is used for evaluations. The data directory contains data generation scripts, and downloaded or generated data sets should end up there as well. The models directory contains pretrained model checkpoints once they are downloaded (see below). The runs directory contains model checkpoints as well as further log files when training models from scratch. The results directory contains the results from the sampling, evaluation, and plotting scripts. Sampled model predictions are written to this directory as compressed numpy arrays. These arrays are read by the plotting scripts, which in turn write the resulting plots to the results directory as well.

The scripts in src contain the main training, sampling, and plotting functionality. Each script contains various configuration options and architecture selections at the beginning of the file. All files should be run directly with Python according to the following pattern:

python src/training_*.py
python src/sample_models_*.py
python src/plot_*.py

Our data sets can be downloaded via any web browser, ftp, or rsync here: http://doi.org/10.14459/2024mp1734798 (rsync password: m1734798.001). Use this command to directly download all data sets at the training and evaluation resolution of $128 \times 64$ (about 146 GB):

rsync -P rsync://[email protected]/m1734798.001/128_* ./data

Once the download is complete, the data set archives can be extracted with:

unzip -o -d data "data/128_*.zip"

To separately download the individiual data subsets, replace the 128_* path in the rsync and unzip commands above with one of the following archive names:

Pretrained model weights for each data set and investigated architecture can be downloaded and extraced in the same way as described for the data above (rsync password: m1734798.001):

rsync -P rsync://[email protected]/m1734798.001/checkpoints.zip ./models
unzip -o -d models "models/checkpoints.zip"

The model checkpoints can be directly sampled on the downloaded data sets with the src/sample_models_*.py scripts.

If problems with the .zip archives occur, it is recommended to check them for corruption. This can be achieved by comparing the SHA512 hash of downloaded files with the corresponding content of the provided checksum file. The checksums for the downloaded archives can be downloaded and checked via (rsync password: m1734798.001):

rsync -P rsync://[email protected]/m1734798.001/checksums.sha512 ./data
cd data
sha512sum -c --ignore-missing checksums.sha512

If the hashes do not match, restart the download or try a different download method.

Once the data sets are downloaded or generated, model architectures can be trained from scratch using the src/training_*.py scripts. The resulting model checkpoints are identical to the downloaded models described above. The training scripts contain different various configuration options and an architecture selection at the beginning of the file.

During training, various values, statistics and plots are logged to Tensorboard, allowing for monitoring the training progress. To start Tensorboard, use the following command:

tensorboard --logdir=runs --port=6006

and open http://localhost:6006/ in your browser to inspect the logged data.

python SU2_raw/data_generation.py [Thread count] [Reynolds number] [List of Mach numbers] [List of corresponding simulation folder IDs] [Restart iteration]

python SU2_raw/data_generation.py 112 10000 0.60,0.70,0.80 0,1,2 -1

Ensure that the data set folder structure resulting from the data generation is the following to ensure the data set can be loaded correctly: data/[datasetName]/sim_[simNr]/[field]_[timestep].npz. Here datasetName is any string, but has to be adjusted accordingly when creating data set objects. The simulation folder numbers should be integers with a fixed width of six digits and increase continuously. Similarly, the timestep numbering should consist of integers with a fixed width of six digits and increase continuously. Start and end points for both can be configured when creating Dataset objects. Fields should be strings that describe the physical quantity, such as pressure, density, or velocity. Velocity components are typically stored in a single array, apart from the Iso case where the velocity z-component is stored separately as velocityZ. For example, a density snapshot at timestep zero from the Tra data set is referenced as data/128_tra/sim_000000/density_000000.npz.

src/copy_data_lowres.py can be used to downsample the generation resolution of 256x128 to the training and evaluation resolution of 128x64 for the simulated data sets. It processes all .npz data files, while creating copies of all supplementary files in the input directory. Computing mean and standard deviation statistics for the data normalization to a standard normal distribution is performed via src/compute_data_mean_std.py.

On the Inc and Tra data sets dicussed above, we observed that it is possible to train models that are unconditionally stable and do not diverge, even after 200 000 rollout steps. Key criteria for training models with such stability are outlined here.

If you use the source code, models, or data sets provided here, please consider citing our work:

@article{kohl2023_acdm,
  author = {Georg Kohl and Li{-}Wei Chen and Nils Thuerey},
  title = {Benchmarking Autoregressive Conditional Diffusion Models for Turbulent Flow Simulation},
  journal = {arXiv},
  year = {2023},
  eprint = {2309.01745},
  primaryclass = {cs},
  publisher = {arXiv},
  url = {https://doi.org/10.48550/arXiv.2309.01745},
  doi = {10.48550/arXiv.2309.01745},
  archiveprefix = {arxiv}
}

This work was supported by the ERC Consolidator Grant SpaTe (CoG-2019-863850).