JAX-SPH: A Differentiable Smoothed Particle Hydrodynamics Framework

粒子流体模拟已经成为解决 Navier-Stokes 方程的强大工具，特别是在包含复杂物理和自由表面的情况下。最近将机器学习方法加入解决此类问题的工具箱中，正在推动这种数值模拟的质量与速度之间的权衡的边界。在这项工作中，我们引领了与深度学习框架兼容的拉格朗日流体模拟器的道路，并提出了 JAX-SPH——一个在 JAX 中实现的平滑粒子流体力学（SPH）框架。JAX-SPH 建立在 LagrangeBench 项目（Toshev 等人，2023）中用于数据集生成的代码基础之上，并通过以下多种方式进行了扩展：(a) 集成了更多关键 SPH 算法，(b) 重构了代码以成为一个 Python 库，(c) 通过求解器验证了梯度，以及(d) 展示了梯度在解决反问题和“求解器中的循环”应用中的实用性。我们的代码可在 https://github.com/tumaer/jax-sph 上获得。

本论文旨在提出一种基于深度学习框架的拉格朗日流体模拟器，以解决流体力学中的复杂物理和自由表面问题。

论文提出了JAX-SPH，一个基于JAX实现的平滑粒子流体动力学（SPH）框架，该框架在LagrangeBench项目的数据集生成代码基础上进行了扩展，包括关键SPH算法的集成、代码重构、梯度验证以及梯度在求解器中的应用。

该论文的亮点包括：（1）提出了一种基于深度学习框架的流体模拟器；（2）JAX-SPH代码已在GitHub上开源；（3）论文展示了梯度在求解器中的应用，包括解决反问题和解决Solver-in-the-Loop应用；（4）该论文提供了一个新的思路和工具，以推动流体动力学领域的研究。

最近在这个领域中，还有一些相关的研究，如：（1）Deep Lagrangian Fluid Simulation with Divergence-Free Networks；（2）Fluid Net: A Physics-Based Generative Model for Fluid Flow；（3）Differentiable Fluid Simulations for Inverse Problems。