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实现的平滑粒子流体动力学（SPH）框架。JAX-SPH以LagrangeBench项目（Toshev等，2023）中的数据集生成代码为基础，并在以下方面进行了拓展：(a) 集成更多关键SPH算法，(b) 将代码重构为Python库，(c) 验证求解器中的梯度计算，(d) 展示梯度在逆问题求解及“求解器-在-循环”应用中的实用性。我们的代码已开源至https://github.com/tumaer/jax-sph。