Neural SPH: Improved Neural Modeling of Lagrangian Fluid Dynamics

Smoothed particle hydrodynamics (SPH) is omnipresent in modern engineering and scientific disciplines. SPH is a class of Lagrangian schemes that discretize fluid dynamics via finite material points that are tracked through the evolving velocity field. Due to the particle-like nature of the simulation, graph neural networks (GNNs) have emerged as appealing and successful surrogates. However, the practical utility of such GNN-based simulators relies on their ability to faithfully model physics, providing accurate and stable predictions over long time horizons - which is a notoriously hard problem. In this work, we identify particle clustering originating from tensile instabilities as one of the primary pitfalls. Based on these insights, we enhance both training and rollout inference of state-of-the-art GNN-based simulators with varying components from standard SPH solvers, including pressure, viscous, and external force components. All Neural SPH-enhanced simulators achieve better performance than the baseline GNNs, often by orders of magnitude in terms of rollout error, allowing for significantly longer rollouts and significantly better physics modeling.

Smoothed particle hydrodynamics (SPH) is a Lagrangian, mesh-free method for numerically solving differential equations, which is the prefered method to simulate problems with (a) free surfaces, (b) complex boundaries, (c) multi-phase flows, and (d) fluid-structure interactions. The method works by discretizing the domain with a set of particles at whose locations the physical quantities can then estimated by kernel interpolations. Thus, a summation over every interaction between particle \(i\) and all particles \(j\) within a given cutoff radius is carried out. Each interaction is then weighted using a radial kernel function \(W\), as illustrated below.

In Neural SPH, the approximated governing equations are the incompressible Navier-Stokes equations (NSE) which are modeled by the so-called weakly compressible NSE given below.

\[ \begin{align} \frac{d}{dt}(\rho) &= - \rho ~ (\nabla \cdot \mathbf{u}), \label{eq1}\tag{1}\\ \frac{d}{dt}(\mathbf{u}) &= - \frac{1}{\rho} \nabla p ~ + ~ \frac{\nu}{V_{ref}L_{ref}} \nabla^2 \mathbf{u} ~ + ~ \mathbf{g}, \label{eq2}\tag{2} \end{align} \]

where \(\rho\) is the fluid density, \( \mathbf{u}\) its velocity, \(p\) the pressure, \(\nu\) the kinematic viscosity, \(V_{ref}\) a reference volume, \(L_{ref}\) a reference length, and \(\mathbf{g}\) the external force. The continuity equation Eq. \eqref{eq1} and the momentum equation Eq. \eqref{eq2} are coupled via an equation of state of the form

\[p(\rho) = p_{ref} \left( \frac{\rho}{\rho_{ref}} - 1 \right). \label{eq3}\tag{3}\] For a more visual introduction to SPH, see Matthias Teschner's slides.

Both classical numerics and learned approaches to Lagrangian fluid simulations face difficulties when evolving a systems over hundreds or thousands of steps. Some of the challenges for the classical methods are the evolution of particles at free surfaces / interfaces / boundaries, as well as particle clumping. After extensively analysing the failure modes of GNN-based Lagrangian CFD surrogates, we conclude that these same challenges are present in learned models and are the primary cause of frequently observed catastrophic failures.

To qualitatively understand clumping, in the figure below we plot the histogram of the per-particle number of neighbors after 400 steps of a 2D lid-driven cavity simulation, which has regions with high particle density. We see a pronounced increase in the number of particles with 8-10 neighbors, indicating clustering artifacts.

To address the challenges of long rollouts, we propose two key components to enhance the performance of GNN-based Lagrangian fluid simulations:

1. External Force Treatment
2. SPH Relaxation

inspired by traditional SPH solvers and aim to improve the physics modeling of the learned simulators.

The learning problem formulation by Sanchez-Gonzalez et al., 2020, strives for simplicity by asking the GNN to implicitly predict the effect of both inter-particle forces and external forces. On the example of the dam break experiment, we stress that if the model learned Physics correctly, then the bias term of the last decode layer should correspond to the gravitational force. However, on a pretrained checkpoint, we observed a bias term of more than 10x smaller than the respective gravitational acceleration!

Our idea is to split the terms on the right-hand side of the momentum equation Eq. \eqref{eq2} into \([...] + \mathbf{g}\) and add the external forces back to the learned acceleration in the integrator, essentially forcing the model to learn only inter-particle forces. Implementing this idea requires special attention because we train on 100-times larger time steps than the reference SPH solver, while we only have access to instantaneous external forces. Thus, the accumulated external force over \(M=100\) SPH steps needs to be removed from the full acceleration \(\mathbf{a}\), leading to

\[\mathbf{a} = \texttt{GNN}(\textbf{X}^{t_{k-H-1}:t_k}, \mathbf{g}) + \mathbf{g}_{M}^{FD}. \label{eq4}\tag{4}\]

In addition to the external force treatment, we propose to apply an SPH relaxation after each learned solver step to improve the poor particle distribution observed in the number on neighbors figure above. Notably, our relaxation is a applied only during inference and does not require retraining the model. We propose an SPH relaxation procedure which (a) uses only the particle coordinates, i.e. no other physical quantities such as density and velocity, and (b) can handle free surfaces as well as walls discretized with one dummy particle layer. Ultimately, the relaxation consists of the \(l\) loops of the following two steps:

\[ \begin{align} \mathbf{a} &= \alpha \frac{-1}{\rho} \nabla p ~ + ~ \alpha \beta \nabla^2 \mathbf{u} , \label{eq5}\tag{5}\\ \mathbf{p} &= \mathbf{p} ~ + ~ \mathbf{a}, \label{eq6}\tag{6} \end{align} \]

where the time step and the pre-factors are hidden in the hyperparameters \(\alpha\) and \(\beta\). Furthermore, according to the weakly compressible assumption, density fluctuations should not exceed \(\sim1\%\). Due to the use of density summation, we set all \(\rho < 0.98\rho_{ref}\) to \(\rho_{ref}\) , and all \(\rho > 1.02 \rho_{ref}\) to \(1.02\rho_{ref}\). Our relaxation implementation is based on the JAX-SPH codebase, but we stress that our approach does not require a differentiably solver, as the relaxation solely operates on the particle coordinates.
Consult our full paper for various further implementation details.

We evaluate the performance of Neural SPH on all 7 LagrangeBench datasets, and both the GNS and SEGNN models. Our results are summarized in the table below, but the appendix of the paper includes various further results including: rollout visualizations, velocity and acceleration histograms, and extensive hyperparameter studies. Overall, we find that Neural SPH significantly improves the physics modeling of GNN-based Lagrangian fluid simulations, at the cost of a slightly increased computational overhead, while the Neural SPH hyperparameters are robust and easy to tune.

Dam Break

The dam break experiment is a challenging test case for fluid simulations due to the free-surface dynamics and wall interactions.

Lid Driven Cavity

The lid-driven cavity experiment is a standard benchmark for incompressible fluid simulations and is used to test the ability of the model to handle fixed and moving no-slip walls.

Reverse Poiseuille Flow

The reverse Poiseuille flow case consists of two opposed channel flows (without walls) with the lower half of the domain being forced left and the upper half being forced right. This experiment tests the ability of the model to handle spatially varying external forces.