PICT: Adaptive GPU Accelerated Differentiable Fluid Simulation for Machine Learning

In this paper we present PICT, our differentiable 2D and 3D fluid simulator for machine learning in the PyTorch framework with support for GPU acceleration. We implemented the PISO algorithm using custom CUDA operations for the core components and Python for the overall algorithm to achieve GPU performance while keeping the simulation easily customizable. To support geometry beyond toy examples we support spatially adaptive multi-block grids using a generalized coordinate system. This allows the user to refine the grid as necessary and align it to boundaries. The forward simulation is validated using analytical and numerical references as well as long rollouts for stability. The gradients of individual components are checked numerically, and we conducted non-trivial optimization and learning tests to verify the usability of our gradients.