JAX-SPH: A Differentiable Smoothed Particle Hydrodynamics Framework
Artur P. Toshev Harish Ramachandran Jonas A. Erbesdobler Gianluca Galletti Johannes Brandstetter Nikolaus A. Adams
Particle-based fluid simulations have emerged as a powerful tool for solving
the Navier-Stokes equations, especially in cases that include intricate physics
and free surfaces. The recent addition of machine learning methods to the
toolbox for solving such problems is pushing the boundary of the quality vs.
speed tradeoff of such numerical simulations. In this work, we lead the way to
Lagrangian fluid simulators compatible with deep learning frameworks, and
propose JAX-SPH - a Smoothed Particle Hydrodynamics (SPH) framework implemented
in JAX. JAX-SPH builds on the code for dataset generation from the
LagrangeBench project (Toshev et al., 2023) and extends this code in multiple
ways: (a) integration of further key SPH algorithms, (b) restructuring the code
toward a Python library, (c) verification of the gradients through the solver,
and (d) demonstration of the utility of the gradients for solving inverse
problems as well as a Solver-in-the-Loop application. Our code is available at
https://github.com/tumaer/jax-sph.