A sharp-interface immersed-boundary incompressible flow solver based on JAX

Sharp-interface immersed-boundary (SIIB) flow solvers have seen great success in the past. The rapid development in machine learning (ML) methods and GPU hardware calls for the migration of such solvers to the GPU for higher computational capacity and tighter integration with ML frameworks. In this work, we develop an SIIB solver purely in JAX, a Python library providing just-in-time (JIT) compilation, auto differentiation, and distributed GPU parallelization. JAX enables developers to write backend-agnostic high-level code that runs on multiple GPUs and is fully differentiable. Immersed bodies are represented using triangular surface meshes. The ghost cell method is implemented for the immersed boundary, the movement of which is treated to enable flow-structure interaction simulations. The fractional step method is implemented to solve the incompressible Navier-Stokes equation. The geometric multigrid method is implemented to solve the pressure Poisson equation. The solver is verified on canonical flow problems. Its performance is benchmarked against the CPU counterpart and other GPU solvers. More features such as finite element bodies and end-to-end optimization capability are planned for future development.