This is the PyTorch based source code for the paper Generative AI for fast and accurate Statistical Computation of Fluids.

GenCFD is a PyTorch-based implementation designed for training and evaluating conditional score-based diffusion models for Computational Fluid Dynamics (CFD) tasks. These generative AI models enable fast, accurate, and robust statistical computation for simulating both two-dimensional and three-dimensional turbulent fluid flows.

To set up a virtual environment and install the necessary dependencies for this project, follow these steps.

1. Create a Virtual Environment

python3 -m venv venv

2. Activate the Virtual Environment

source venv/bin/activate

3. Install dependencies: Make sure your virtual environment is active, then run:

pip install -r requirements.txt

4. Optional: Install GenCFD as a library If you would like to use GenCFD and its functions as a library, you can install it directly from the repository instead of step 3:

pip install git+https://github.com/camlab-ethz/GenCFD.git@main

Train a model from the root directory by using:

python3 -m GenCFD.train.train_gencfd \
--dataset <DATASET_NAME> \
--model_type <MODEL_NAME> \
--save_dir <DIRECTORY_PATH> \
--num_train_steps <INT>

• DATASET_NAME should be a valid dataset from the section Dataset.
• MODEL_NAME should either be PreconditionedDenoiser for the two-dimensional case, or PreconditionedDenoiser3D in the three-dimensional case.
  • 3D models have approximately 70M parameters with (64, 64, 64) resolution.
  • Recommended: GPU with 32GB memory (batch size 5) or 24GB memory (batch size 4).

After training, a JSON file with model settings is saved in the output directory for use during inference.

For a fast training, use a compiled model in parallel:

torchrun --nproc_per_node=<INT> \
-m GenCFD.train.train_gencfd \
--world_size <INT> \
--dataset <DATASET_NAME> \
--model_type <MODEL_NAME> \
--save_dir <DIRECTORY_PATH> \
--num_train_steps <INT>
--compile

The flag --world_size determines the size of the group associated with a communicator and it has to correspond to the number of processes or trainers used for parallelization. The relevant flag to set this is --nproc_per_node. Another advice for fast training is to chooce a proper number of workers which can be specified through the flag --worker.

Run inference with:

python3 -m GenCFD.eval.evaluate_gencfd \
--dataset <DATASET_NAME> \
--model_type <MODEL_NAME> \
--model_dir <DIRECTORY_PATH> \
--compute_metrics \
--monte_carlo_samples <INT> \
--visualize \
--save_gen_samples\
--save_dir <DIRECTORY_PATH>

Run inference in parallel:

torchrun --nproc_per_node=<INT> \
-m GenCFD.eval.evaluate_gencfd \
--world_size <INT> \
--dataset <DATASET_NAME> \
--model_type <MODEL_NAME> \
--model_dir <DIRECTORY_PATH> \
--compute_metrics \
--monte_carlo_samples <INT> \
--visualize \
save_gen_samples \
--save_dir <DIRECTORY_PATH>

Also here the number of models spawned should be the same for both flags --world_size and --nproc_per_node. It's also possible to compile the model when run in a parallel or a sequential setup.

• --compute_metrics: Computes evaluation metrics (e.g., mean and standard deviation) using Monte Carlo simulations.
• --visualize: Generates a single inference sample for visualization.
• --save_gen_samples: Saves randomly selected samples drawn from a uniform distribution

The number of sampling steps (--sampling_steps) for the Euler-Maruyama method should be preferably >30 for convergence.

The following table summarizes key arguments that can help optimize memory usage or fine-tune model performance.

• Action Arguments: Simply add the flag (e.g., --track_memory), no need to specify True or False.
• Boolean Flags: Requires explicit specification of either True or False (e.g. --use_mixed_precision True)

A compiled version of the model can be used through adding the flag --compile. The compiler works without any issues on the following GPUs

• NVIDIA GeForce RTX 3090
• NVIDIA GeForce RTX 4090
• NVIDIA Tesla V100-SXM2 32 GiB
• NVIDIA Tesla V100-SXM2 32 GB
• Nvidia Tesla A100

If there are some compiler warnings, you can always surpress them.

The table below provides a description of each dataset along with the corresponding flag argument for selection during training or evaluation.

We would like to extend our deepest gratitude to the Google Research team for their groundbreaking work and open-source contributions. This project builds upon their foundational models and research, which have been instrumental in advancing the development of GenCFD.

For more details, please refer to their original work: swirl-dynamics.

If you use this code or find it helpful in your research, please cite the following paper:

@misc{molinaro2025generativeaifastaccurate,
      title={Generative AI for fast and accurate statistical computation of fluids}, 
      author={Roberto Molinaro and Samuel Lanthaler and Bogdan Raonić and Tobias Rohner and Victor Armegioiu and Stephan Simonis and Dana Grund and Yannick Ramic and Zhong Yi Wan and Fei Sha and Siddhartha Mishra and Leonardo Zepeda-Núñez},
      year={2025},
      eprint={2409.18359},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2409.18359}, 
}