FourierFlow:
Frequency-aware Flow Matching for Generative Turbulence Modeling

1. Environment setup

conda create -n repa python=3.9 -y
conda activate repa
pip install -r requirements.txt

All our experiments are implemented on 8 x NVIDIA H800 GPUs.

2. Dataset

Dataset download

Please download Compressible N-S datasets in PDEBench and Shear Flow in The Well datasets, and set 'base-path' in train.py accordingly. (PDEBench, The Well)

3. Training

accelerate launch /you_path_to/FourierFlow/train.py --allow-tf32

Then this script will automatically create the folder in exps to save logs and checkpoints. You can adjust the following options:

• --models: [SiT-B/2, SiT-L/2, SiT-XL/2]
• --proj-coeff: Any values larger than 0
• --encoder-depth: Any values between 1 to the depth of the model
• --output-dir: Any directory that you want to save checkpoints and logs
• --exp-name: Any string name (the folder will be created under output-dir)

4. Evaluation

python eval.py

If you want to test the surrogate model and the generative model together, run the following:

python all_eval.py

If you want to test the results of each step, run:

python all_eval_step.py

We also provide the SiT-XL/2 checkpoint (trained for 4M iterations) used in the final evaluation. It will be automatically downloaded if you do not specify --ckpt.

Acknowledgement

This code is mainly built upon DiT, and SiT repositories.