fluidnet_cxx: Accelerating Fluid Simulation with Convolutional Neural Networks. A PyTorch/ATen Implementation.

This repository is based on the paper, Accelerating Eulerian Fluid Simulation With Convolutional Networks by Jonathan Tompson, Kristofer Schlachter, Pablo Sprechmann, Ken Perlin on the accelation of fluid simulations by embedding a neural network in an existing solver for pressure prediction, replacing an expensive pressure projection linked to a Poisson equation on the pressure, which is usually solved with iterative methods (PCG or Jacobi methods). We implemented our code with PyTorch, effectively replacing all the original Torch/Lua and C++/CUDA implementation of the inviscid, incompressible fluid solver (based on the open-source fluid simulator Mantaflow, aimed at the Computer Graphics community). Find the original FluidNet repository here.

We have retaken the original FluidNet NN architecture and added different features, such as replacing upsampling with deconvolution layers, or directly replacing the complete architecture with a deeper MultiScale net which showed more accurate results at the expense of inference speed.

This work allows to compare both the code perfomace when run in a single GPU unit and the accuracy of this data-driven method in comparison with tradional mehtods (Jacobi) or other fluid simulation methods like Lattice Boltzmann Methods.

Results

Simulations of a buoyancy-driven plume flow are performed with different methods for the Poisson equation resolution. An inlet is placed at the bottom of the domain, where a lighter fluid (with density rho0) is injected with velocity v0 into a quiescent heavier fluid. Results show that some work is still needed to predict a correct plume growth rate, due probably to a poor modelling of buoyant forces by the trained model.

Resolution with ConvNet | Jacobi Method 28 iter | Jacobi Method 100 iter

Growth Rate of the plume's head for Ri=0.14

Functionalities:

• NOTE: For the moment, only 2D simulations and training are supported. 3D needs still some work.
• Full eulerian (incompressible and inviscid) fluid simulator:
  • Momentum equation resolution using a splitting algorithm:
    • Advection of velocity + External forces
    • Enforcing of non-divergence of velocity constraint through Poisson equation resolution, resulting in a pressure gradient that corrects the velocity from the previous step. Step replaced by a fully convolutional Neural Network with divergence of velocity as input and pressure as output.
  • Unconditional Stable MacCormack discretization of velocity advection algorithm.
  • Jacobi method implementation for comparison.
• Dataset:
  • Generation with FluidNet own Mantaflow sript.
  • Random insertion of objects and velocity emitters, as well as gravity forces.
  • Pre-processed into PyTorch objects
• Pre-trained models:
  • An already trained model is available in trained_models
• Training:
  • Several options for loss function:
    • MSE of pressure
    • "Physical" loss: MSE of velocity divergence (unsupervised)
    • MSE of velocity divergence after several timesteps.
  • Short term divergence loss: 8 hours training
  • Short+Long term divergence loss: ~2 days
• Inference. Two test cases:
  • Buoyant plume.
  • Rayleigh Taylor instability.
  • Launch your simulation with the available pre-trained model.
  • Comparison with Jacobi method resolution + LBM with open-sourced C++ library Palabos
• Results visualization:
  • Matplotlib
  • Paraview post-processing tool (VTK files)

Models

• Adapted FluidNet architecture
• Deep MultiScale adapted from Deep multi-scale video prediction beyond mean square error.

Requirements

• Python 3.X
• C++11
• Pytorch 0.4 (Including ATen Tensor library, exposing PyTorch library in C++)
• FluidNet own Mantaflow implementation
• PyVTK (pip install)
• (Optional) Paraview
• (Optional) OpenCV2

ATen allows to write generic code that works on both devices. More information in ATen repo. It can be called from PyTorch, using its new extension-cpp.

Installation

To install this repo:

1. Clone this repo:

https://github.com/jolibrain/fluidnet_cxx.git

2. Install Pytorch 0.4: Pytorch 0.4 NOTE: Training is done in GPUs

3. Install cpp extensions for fluid solver: C++ scripts have been written using PyTorch's backend C++ library ATen. These scripts are used for the advection part of the solver. Follow these instructions from main directory:

cd pytorch/lib/fluid/cpp
python3 setup.py install # if you want to install it on local user, use --user

Training

Dataset We use the same 2D dataset as the original FluidNet Section 1: Generating the data - Generating training data (generated with MantaFlow) for training our ConvNet.

Running the training To train the model, go to pytorch folder:

cd pytorch

The dataset file structure should be located in <dataDir> folder with the following structure:

.
└── dataDir
    └── dataset
        ├── te
        └── tr

Precise the location of the dataset in pytorch/config.yaml writing the folder location at dataDir (use absolute paths). Precise also dataset (name of the dataset), and output folder modelDirwhere the trained model and loss logs will be stored and the model name modelFilename.

Run the training :

python3 fluid_net_train.py

For a given dataset, a pre-processing operation must be performed to save it as PyTorch objects, easily loaded when training. This is done automatically if no preprocessing log is detected. This process can take some time but it is necessary only once per dataset.

Training can be stopped using Ctrl+C and then resumed by running:

python3 fluid_net_train.py --resume

You can also monitor the loss during training by running in /pytorch

python3 plot_loss.py <modelDir> #For total training and validation losses
#or
python3 plot_5loss.py <modelDir> #For each of the losses (e.g: L1(div) and L2(div))

It is also possible to load the saved model and print its output fields and compare it to targets (Pressure, Velocity, Divergence and Errors):

python3 print_output.py <modelDir> <modelFilename>
#example:
python3 print_output.py data/model_pLoss_L2 convModel

Training options

You can set the following options for training from the terminal command line:

• -h : displays help message
• --trainingConf : YAML config file for training. Default = config.yaml.
• --modelDir : Output folder location for trained model. When resuming, reads from this location.
• --modelFilename : Model name.
• --dataDir : Dataset location.
• --resume : Resumes training from checkpoint in modelDir
• --bsz : Batch size for training.
• --maxEpochs : Maximum number training epochs.
• --noShuffle : Remove dataset shuffle when training.
• --lr : Learning rate.
• --numWorkers : Number of parallel workers for dataset loading.
• --outMode : Training debug options. Prints or shows validation dataset. save = saves plots to disk show = shows plots in window during training none = do nothing

The rest of the training parameters are set in the trainingConf file, by default config.yaml.

Parameters in the YAML config file are copied into a python dictionary and saved as two separated dictionaries in modelDir, one conf dictionary for parameters related to training (batch size, maximum number of epochs) and one mconf dictionary for parameters related to the model (inputs, losses, scaling options etc)

Test

Run the buoyant plume test case by running:

cd pytorch
python3 plume.py --modelDir <modelDir> --modelFilename <modelFilename> --outputFolder <outputFolder>

with:

• <modelDir> : folder with trained model.
• <modelFilename> : Trained model name.
• <outputFolder> : Folder for saving simulation results.

You can also stop the simulation (Ctrl+C) and restart it afterwards:

python3 plume.py --restartSim

Test options

• -h : displays help message
• --simConf : YAML config file for simulation. Default = plumeConfig.yaml.
• --trainingConf : YAML config file for training. Default = config.yaml.
• --modelDir : Trained model location.
• --modelFilename : Model name.
• --outputFolder : Location of output results.
• --restartSim : Restart simulation from checkpoint in <outputFolder>.

Check plumeConfig.yaml to see how the configuation file for the simulation is organized.

Modifying the NN architecture

If you want to try your own architecture, you only have to follow these simple rules:

• Write your model in a separate script and save it inside pytorch/lib.
• Open model.py and import your own script as a module. Go to class FluidNet here.
• Ideally, as with the Multi-Scale Net example, you should just have to precise the number of channels from the input, and add your net forward pass as in the multicale example here

Extending the cpp code:

The cpp code, written with ATen library, can be compiled, tested and run on its own. You will need OpenCV2 to visualize output of the pressure and velocity fields, as matplotlib is unfortunately not available in cpp!

Test

First, generate the test data from FluidNet Section 3. Limitations of the current system - Unit Testing and write the location of your folder in:

solver_cpp/test/test_fluid.cpp
#define DATA <path_to_data>

Run the following commands:

cd solver_cpp/
mkdir build_test
cd build_test
cmake .. -DFLUID_TEST=ON # Default is OFF
./test/fluidnet_sim

This will test every routine of the solver (advection, divergence calculation, velocity update, adding of gravity and buoyancy, linear system resolution with Jacobi method). These tests are taken from FluidNet and compare outputs of Manta to ours, except for advection when there is no Manta equivalent. In that case, we compare to the original FluidNet advection.

Run

cd solver_cpp/
mkdir build
cd build
cmake .. -DFLUID_TEST=OFF # Default is OFF
./simulate/fluidnet_sim

Output images will be written in build folder, and can be converted into gif using ImageMagick.

NOTE: For the moment, only 2D simulations and training are supported, as bugs are still found for the 3D advection.