[ICLR 2023] Learning Vortex Dynamics for Fluid Inference and Prediction

by Yitong Deng, Hong-Xing Yu, Jiajun Wu, and Bo Zhu.

Our paper can be found at: https://arxiv.org/abs/2301.11494.

Video results can be found at: https://yitongdeng.github.io/vortex_learning_webpage.

Environment

The environment can be installed by conda via:

conda env create -f environment.yml
conda activate vortex_env

Our code is tested on Windows 10 and Ubuntu 20.04.

Data

The 5 videos (2 synthetic and 3 real-world) used in our paper can be downloaded from Google Drive. Once downloaded, place the unzipped data folder to the project root directory.

Synthetic 1

Pretrain

First, execute the command below to pretrain the trajectory network so that the initial vortices are regularly spaced to cover the simulation domain and remains stationary.

python train.py --config configs/synthetic_1.txt --run_pretrain True

Once completed, navigate to pretrained/exp_synthetic_1/tests/ and check that the plotted dots are regularly spaced and remain roughly stationary. A file pretrained.tar shall also appear at pretrained/exp_synthetic_1/ckpts/.

Train

Then, run the command below to train.

python train.py --config configs/synthetic_1.txt

Checkpoints and testing results are written to logs/exp_synthetic_1/tests/ once every 1000 training iterations.

Results

When run on our Windows machine with AMD Ryzen Threadripper 3990X and NVIDIA RTX A6000, this is the final testing result we get:

Note that since our PyTorch code includes nondeterministic components (e.g., the CUDA grid sampler), it is expected that each training session will not generate the exact same outcome.

Synthetic 2

Pretrain

python train.py --config configs/synthetic_2.txt --run_pretrain True

Train

python train.py --config configs/synthetic_2.txt

Results

Real 1

Pretrain

python train.py --config configs/real_1.txt --run_pretrain True

Train

python train.py --config configs/real_1.txt

Results

Real 2

Pretrain

python train.py --config configs/real_2.txt --run_pretrain True

Train

python train.py --config configs/real_2.txt

Results

Real 3

Pretrain

python train.py --config configs/real_3.txt --run_pretrain True

Train

python train.py --config configs/real_3.txt

Results

Trying your own video

We assume the input is a Numpy array of shape [num_frames], 256, 256, 3, with the last dimension representing RGB pixel values between 0.0 and 1.0, located in data/[your_name_here]/imgs.npy. For fluid videos with boundaries (like in our real-world examples), it is required that a Numpy array of shape 256, 256 representing the signed distance field to the boundary be supplied in data/[your_name_here]/sdf.npy. We assume the signed distance has a unit of pixels.

For videos of higher dynamical complexity, we also encourage playing around with the number of vortex particles used. Currently, this is determined by the vorts_num_x and vorts_num_y parameters in train.py hard coded to 4, which might need to be increased as needed.

Bibliography

If you find our paper or code helpful, please consider citing:

@inproceedings{deng2023vortex,
 title={Learning Vortex Dynamics for Fluid Inference and Prediction},
 author={Yitong Deng and Hong-Xing Yu and Jiajun Wu and Bo Zhu},
 booktitle={Proceedings of the International Conference on Learning Representations},
 year={2023},
}