JAX-SPH: A Differentiable Smoothed Particle Hydrodynamics Framework

Authors: artur p toshev
harish  ramachandran
jonas a erbesdobler
gianluca  galletti
johannes  brandstetter
nikolaus a adams
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-11 12:55:36

Abstract: 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 package, (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.

Main paper: 
 introduction: 
 Partial differential equations (PDEs) are the mathematical tools developed to describe natural phenomena ranging from engineering and physics to social sciences. Various numerical methods have been developed to solve these PDEs, as analytical solutions are only available for toy examples, with the most recent class of PDE solvers being the machine learning-based ones (Thuerey et al., 2021; Brunton & Kutz, 2023) . One particular class of machine learning (ML) approaches called hybrid solvers refers to combining ideas (or full algorithmic blocks) from classical numerical solvers and machine learning (Schenck & Fox, 2018; Um et al., 2020; Kochkov et al., 2021; Jagtap et al., 2022; Lienen & Günnemann, 2022; Karlbauer et al., 2022; Li & Farimani, 2022; Toshev et al., 2023b) . This has been one of the main reasons for the development of differentiable fluid mechanics solvers like PhiFlow (Holl et al., 2020) , JAX-CFD (Kochkov et al., 2021), and JAX-Fluids (Bezgin et al., 2022) . However, these three frameworks implement Eulerian, i.e., grid-based, solvers, and we lead the way to a JAX-based Lagrangian CFD solver. Eulerian solvers refer to numerical methods that discretize space into static volume elements and then track the evolution of the fluid properties at these positions, while Lagrangian solvers discretize individual material elements, which are then shifted in space following the local velocity field.Algorithmically, grid-based and particle-based methods are very different. While grid-based solvers rely on stencils akin to kernels in Convolutional Neural Networks (CNNs) (Lecun et al., 1998) , particle-based solvers rely on kernel approximations akin to Graph Neural Networks (GNNs) (Scarselli et al., 2008; Battaglia et al., 2018) operating on a radial distance-based graph. The main overhead of Lagrangian over Eulerian approaches is updating the connectivity between discretization elements at every autoregressive solver step. Even if an Eulerian scheme operates on an irregular mesh, the connectivity prescribed by this mesh can be precomputed, while by moving particles in space in Lagrangian methods, their neighbors constantly vary in time.As SPH techniques advance, various software packages have arisen in the past few years. However, most of them are designed for high-performance computing (HPC) applications and are typically implemented in low-level languages like C++ or Fortran (Crespo et al., 2015; Koschier et al., 2019; Zhang et al., 2021) , with two notable exceptions: PySPH (Ramachandran et al., 2021) in Python and juSPH (Luo et al., 2022) in Julia. Nevertheless, hardly an SPH implementation exists which is out-of-the-box compatible with modern deep learning frameworks like TensorFlow (Abadi et al., 2015) , PyTorch1 (Paszke et al., 2019) or JAX (Bradbury et al., 2018) , i.e., leverages automatic differentiation to power differentiable solvers like PhiFlow (Holl et al., 2020) . Upon identifying the lack of an ML-ready SPH solver, we choose JAX for the framework implementation for two reasons: (a) JAX tends to be faster for operations on graphs even after PyTorch 2.0 has been introduced (see Appendix F in Toshev et al. (2024a) ), and (b) we can use the cell list-based neighbor search routine of the JAX-MD library (Schoenholz & Cubuk, 2020) . We note that with our solver, we target the easier integration of SPH with ML workflows rather than developing a parallel HPC code, and we exclusively use the AD routine grad by JAX. We leave the implementation of better custom adjoints along the lines of Ma et al. (2021) ; Kidger (2022) ; Nadarajah & Jameson (2000) to future work.In this work, we significantly extend the codebase used for dataset generation within La-grangeBench2 (Toshev et al., 2024a) and demonstrate the utility of the obtained gradients in multiple ways. Our contributions are:• The addition and validation of Transport Velocity (Adami et al., 2013) , Riemann SPH (Zhang et al., 2017b) , and thermal diffusion effects (Cleary, 1998) . • Validating the accuracy of the automatic differentiation-based gradients (Griewank & Walther, 2008) over 5 solver steps. • An open-source Python package at pypi.org/project/jax-sph.• A demonstration of how to use these gradients on (a) an inverse problem over 100 SPH solver steps, and (b) using the SPH solver in a Solver-in-the-Loop fashion (Um et al., 2020) . 
sph solver: 
 In this section, we introduce the components included in our solver code. The validation of these can be found in Appendix A, for which we include various solver validation cases.Weakly compressible SPH. We follow the weakly compressible SPH approach (Monaghan, 1994; Morris et al., 1997) to evolve the dynamics of incompressible fluids. The equations governing such systems are the mass and momentum conservation equationsEQUATIONEQUATIONwith density ρ, velocity u, pressure p, Reynolds number Re, and external force f . To numerically solve these equations, SPH applies a distance-based kernel W that averages over the properties of the fluid. The default kernel in our codebase is the Quintic spline (Morris et al., 1997) , but we also include the 5th order Wendland kernel (Wendland, 1995) .The preferred way of estimating the density (Monaghan, 2005) is through density summation ρ i = j m j W (r ij |h), with m j being the mass of the adjacent particles j, h the smoothing length of the kernel, and r ij the interparticle distance. However, this approach leads to unphysically low densities at free surfaces, forcing the use of density evolution, i.e., numerically integrating the mass conservation Eq. 1. We also implement a density field reinitialization method (Zhang et al., 2017a) to mitigate errors arising from density evolution (Colagrossi & Landrini, 2003) . In weakly compressible SPH, the pressure is defined as a function of density through an equation of state, and in our code, we choose the formulation by Monaghan (1994) for standard and transport velocity formulation SPH. The equation of state used for the Riemann SPH formulation is the one proposed by Zhang et al. (2017b) .Transport velocity. The observation that tensile instabilities in standard SPH can cause particle clumping and void regions (Price, 2012) has led to the development of advanced shifting schemes like the transport velocity formulation of SPH (Adami et al., 2013; Zhang et al., 2017a) . We add the shifting velocity proposed by Adami et al. (2013) as an optional feature in our codebase.Riemann SPH. We also add the Riemann SPH solver formulation by Zhang et al. (2017b) , which is based on the idea of introducing a simple low dissipation limiter to a classical Riemann solver, resulting in decreased numerical dissipation. We construct a one-dimensional linear inter-particle Riemann problem along a unit vector pointing from particle i to particle j, which implicitly regularizes both the momentum equation and the mass conservation. We implement the dissipation limiter as proposed by Zhang et al. (2017b) .Wall boundaries. We follow the generalized wall boundary condition approach by Adami et al. (2012) to enforce different boundary conditions for the standard SPH and the transport velocity SPH formulation. This approach, in essence, implements two physical constraints: (a) impermeability of walls and (b) viscous effects at walls. For impermeability, the pressure of the adjacent fluid particles is assigned to the dummy wall particles, which leads to a zero pressure gradient in the wallnormal direction at the wall surface, and, thus prevents the penetration of fluid particles into the wall.Regarding viscosity, there are two cases to distinguish: (b1) no-slip enforces the no-slip boundary condition, i.e., the fluid must have zero velocity directly at the wall surface, and (b2) free-slip, i.e., the fluid must have a zero velocity normal to a wall, but might have any velocity tangentially. These two conditions are implemented by assigning the inverted velocity of the fluid to the wall particles, either fully for (b1) or only in the wall-normal direction for (b2). See Adami et al. (2012) for more details.For the Riemann SPH wall boundary implementation, we solve a one-sided Riemann problem as proposed by Zhang et al. (2017b) . Unlike the generalized wall boundary condition, this method avoids interpolating and then extrapolating the fluid states for the wall boundary particles. Here, these wall particles are assigned the individual Riemann states depending on the current particleparticle interaction. For a more in-depth explanation, see (Zhang et al., 2017b; Yang et al., 2020) .Thermal diffusion. Thermal diffusion in weakly compressible SPH involves the transfer of heat between neighboring particles governed by Fourier's law of heat (Cleary, 1998) . This diffusion smoothens the temperature field by applying an SPH kernel interpolation over adjacent particles to compute the rate of temperature change. As we deal with incompressible fluids, the dynamics are not influenced by thermal effects, and the pressure or velocity fields do not directly influence the temperature field. Thus, the addition of temperature does not interfere with the previously presented SPH algorithms and allows us to study thermal effects which are governed by diffusion (by our explicit diffusion modeling) and convection (as Lagrangian particles are shifted in space). For an example of a channel flow with a hot wall, see Appendix C. 
experiments: 
 Gradient validation. To test the validity of the gradients obtained through our differentiable solver, we compute analytical solver gradients via automatic differentiation and compare those to numerical gradients from finite difference schemes (Griewank & Walther, 2008) . In our setup, gradients are accumulated over 5 solver steps, preceded by 10 (forward only) warm-up steps. Epsilon for finite differences is picked as 0.001dx = 5e -5, as smaller values lead to instabilities. Fig. 1 shows the scalar gradients of kinetic energy over position changes dE kin dr when using 2-dimensional Taylor-Green vortex (TGV) (Brachet et al., 1983; 1984) and lid-driven cavity (LDC) (Ghia et al., 1982) as initial states. Inverse problem. Our first application case is an inverse problem, representing the class of inverse design and flow control problems that are successfully tackled by differentiable solvers (Holl et al., 2020) as well as differentiable learned solver surrogates (Allen et al., 2022) . The scenario involves a 2D box containing a water cube, discretized by 36 particles, which undergoes acceleration due to gravity across 100 solver steps. The task is to find the initial coordinates of the cube given its final state; see Fig. 2 . The inverse problem is formulated by computing the mean squared error of coordinates between the target final state and the end of a simulation with randomly placed initial particles. After as few as 15 gradient descent steps, we reach a state closely resembling the original one, up to some loss of information during the deformation of the water cube while interacting with the wall.Solver-in-the-loop. As a second experiment to showcase the solver differentiability, we adapt the popular "Solver-in-the-Loop" (SitL) (Um et al., 2020) training scheme to SPH. Initially developed to tackle spatial coarsening on grids, SitL interleaves a solver operating on a coarse spatial and/or temporal discretization with a learnable correction function. The solver manages coarse, low-frequency components, while the learnable function adjusts high-frequency details. Due to the inherent difficulties of spatial coarse-graining in particle systems, our objective is to implement SitL only for temporal coarsening. This significantly differs from the original application of SitL and mandates a series of design changes to the original architecture, mainly related to the normalization and training procedure, see Appendix B.  
conclusion: 
 We have developed JAX-SPH, a framework for simulating Lagrangian fluid problems, which can be easily integrated into design/control problems as well as hybrid solver approaches. By building our code on the high-performance library JAX and validating the simulation results of our solver, we offer a fast and reliable SPH solver in Python. With our work, we hope to accelerate the development of more hybrid Lagrangian solvers, e.g., Toshev et al. (2024b) , and we leave the addition of more SPH algorithms and simulation cases to future work. One particularly exciting future direction is developing foundation models for PDEs that can operate on both Eulerian and Lagrangian data (Alkin et al., 2024) , potentially combined with encoded symmetries (Toshev et al., 2023a) . 
 
 Appendix: A.1 TAYLOR GREEN VORTEX 2D Fig. 4 shows the absolute velocities of each particle at the start and at the end of the Taylor Green vortex (Brachet et al., 1983; 1984) simulation at Re = 100. Here, the transport velocity SPH formulation is used. In Fig. 5 , one can see the decay of the maximum velocity and the kinetic energy over time for the different methods, standard SPH (SPH), transport velocity formulation SPH (SPH + tvf), and Riemann SPH (Riemann). The following figures compare the different methods for a 2D lid-driven cavity (Ghia et al., 1982) simulation at Re = 100. The reference data for the velocity profiles, i.e., black dots for U (velocity in x direction at a vertical cut through the middle of the cavity) and black squares for V (velocity in y direction at a horizontal cut through the middle of the cavity), are from Ghia et al. (1982) . The following Fig. 8 shows the nondimensionalized pressure of a Riemann SPH (Zhang et al., 2017b) dam break (Colagrossi & Landrini, 2003) Solver-in-the-Loop was trained with LagrangeBench (Toshev et al., 2024a) , with a custom RPF 2D dataset with 20-step temporal coarsening. For SitL, n sitl=3. Both the corrector model and the reference GNS (Sanchez-Gonzalez et al., 2020) are message-passing networks with 10 layers and 64 latent dimensions. The starting learning rate is set to 1e -3 for both, and noise std is set to 1e -5 for SitL and to 1e -3 for GNS. Table 1 shows the LagrangeBench performance metrics on these models. Best models are picked based on the MSE 20 loss on the validation set.Solver only GNS SitL MSE 51.7e -7 6.7e -9 3.3e -9 MSE 20 7.9e -6 1.9e -7 1.3e -7 MSE E kin 0.13 2.8e -4 7.4e -5 Sinkhorn 3.4e -7 3.7e -8 9.3e -9Table 1 : LagrangeBench metrics on the RPF 2D dataset over 20 steps.Step 0Step 1000Step 3000 For similar work in progress, we refer the reader to TorchSPH https://github.com/wi-re/pytorchSPH. Compared to our JAX-SPH, this code uses PyTorch, different SPH algorithms, different experiments, and does not validate the solver or the gradients.2 https://github.com/tumaer/lagrangebench/blob/main/notebooks/data gen.ipynb. 


Decision: The paper introduces a differentiable lagrangian fluid simulator based on smooth particle hydrodynamics (SPH), written in JAX. The authors extend the SPH code in LagrangeBench to be differentiable, demonstrating some applications of the said approach. The paper is a useful contribution to the community given the growing interest in SPH codes. There is limited comparisons to other SPH methods but nevertheless it is a well written paper and should be accepted for the poster session.

Reviews: review: title: title: Good contribution to the field of differentiable lagrangian-based solvers
main_review: main_review: **Summary:** The paper introduces a differentiable lagrangian fluid simulator based on smooth particle hydrodynamics (SPH), written in JAX. The authors extend the SPH code in LagrangeBench to be differentiable, demonstrating some applications of the said approach. 

**Strengths:**

1. The paper seems well-motivated and addresses the need for differentiable Lagrangian CFD solvers.
2. Extending the differentiability of SPH solvers is a novel contribution with examples demonstrating the usability of the software.
3. Incorporation of various features such as Riemann SPH, Transport velocity, Wall boundaries, and Thermal diffusion. 

**Areas of Improvement:**

1. The work can significantly benefit from a performance analysis with other available open-source SPH solvers, which questions the suitability for HPC applications.
2. Unsucessful evaluation of the original purpose of "Solver-in-the-loop" (SitL), i.e., spatial coarsening, concerns the applicability of methods to other use cases.
3. There has been a significant discussion about using continuous vs discrete adjoints [ma2021, kidger2022, nadarajah2000], which can have significant performance impacts. Evaluating the solvers and implementing the said adjoints could strengthen the technical aspects.

Moreover, I would be interested in more details of what specifics were done to make the solver differentiable, which can serve as a roadmap for other scientists implementing such programs in niche scientific disciplines.

[ma2021] Ma, Yingbo, et al. "A comparison of automatic differentiation and continuous sensitivity analysis for derivatives of differential equation solutions." 2021 IEEE High Performance Extreme Computing Conference (HPEC). IEEE, 2021. 
 
[kidger2022] Kidger, Patrick. "On neural differential equations." arXiv preprint arXiv:2202.02435 (2022).  

[nadarajah2000]: Nadarajah, Siva, and Antony Jameson. "A comparison of the continuous and discrete adjoint approach to automatic aerodynamic optimization." 38th Aerospace sciences meeting and exhibit. 2000.