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.

我们提出了一种适用于深度学习框架的拉格朗日流体模拟器JAX-SPH，该模拟器基于JAX实现，并扩展了LagrangeBench项目中的代码，包括进一步的SPH算法、代码重构为Python库、求解器梯度验证，以及展示了求解逆问题和应用于Solver-in-the-Loop的梯度的实用性。

JAX-SPH：一个可微分的平滑粒子流体动力学框架