SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Artur P. Toshev*,1
artur.toshev@tum.de
Gianluca Galletti*,1
g.galletti@tum.de
Fabian Fritz
Stefan Adami²
Nikolaus A. Adams<sup>1,2</sup>
<sup>1</sup> Chair of Aerodynamics and Fluid Mechanics Technical University of Munich 85748 Garching b. München, Germany
$^{2}$  Munich Institute of Integrated Materials, Energy and Process Engineering Technical University of Munich 85748 Garching b. München, Germany
Machine learning has been successfully applied to grid-based PDE modeling in various scientific applications. However, learned PDE solvers based on Lagrangian particle discretizations, which are the preferred approach to problems with free surfaces or complex physics, remain largely unexplored. We present LagrangeBench, the first benchmarking suite for Lagrangian particle problems, focusing on temporal coarse-graining. In particular, our contribution is: (a) seven new fluid mechanics datasets (four in 2D and three in 3D) generated with the Smoothed Particle Hydrodynamics (SPH) method including the Taylor-Green vortex, lid-driven cavity, reverse Poiseuille flow, and dam break, each of which includes different physics like solid wall interactions or free surface, (b) efficient JAX-based API with various recent training strategies and three neighbor search routines, and (c) JAX implementation of established Graph Neural Networks (GNNs) like GNS and SEGNN with baseline results. Finally, to measure the performance of learned surrogates we go beyond established position errors and introduce physical metrics like kinetic energy MSE and Sinkhorn distance for the particle distribution. Our codebase is available under the URL: https://github.com/tumaer/lagrangebench.
Partial differential equations (PDEs) are ubiquitous in engineering and physics as they describe the temporal and spatial evolution of dynamical systems. However, for practical applications, a closed analytical solution often does not exist and thus requires numerical methods like finite elements or finite volumes. We can broadly classify such numerical methods into two distinct families: Eulerian and Lagrangian. In Eulerian schemes (grid-based or mesh-based methods), spatially fixed finite nodes, control volumes, cells, or elements are used to discretize the continuous space. In Lagrangian schemes (mesh-free or particle methods), the discretization is carried out using finite material points, often referred to as particles, which move with the local deformation of the continuum. Eulerian and Lagrangian methods offer unique schemes with different characteristics dependent on the PDEs'
underlying physics. This work exclusively focuses on the physics of fluid flows, i.e., we study the numerical solution to Navier-Stokes equations (NSEs), but offers the machine learning infrastructure for many more Lagrangian problems.
The smoothed particle hydrodynamics (SPH) method is a family of Lagrangian discretization schemes for PDEs, see, e.g., Monaghan [39, 41] and Price [52] for an introduction to SPH. In SPH, the particles move with the local flow velocity, making it a genuine Lagrangian scheme, while the spatial differential operators are approximated using a smoothing operation over neighboring particles. One of the main advantages of SPH compared to Eulerian discretization techniques is that SPH can handle large topological changes with ease since (a) no connectivity constraints between particles are required, and (b) advection is treated exactly [41]. However, SPH is typically more expensive than, e.g., a finite difference discretization on a comparable resolution due to its algorithmic complexity [16]. Despite being initially devised for applications in astrophysics [23, 35], SPH has gained traction as a versatile framework for simulating complex physics, including (but not limited to) multiphase flows [44, 38, 16, 27, 28], internal flows with complex boundaries [1], free-surface flows [4, 36, 37, 40], fluid-structure interactions [3, 68, 69], and solid mechanics [14, 25].
Eulerian Machine Learning. In recent years, the application of machine learning (ML) surrogates for PDE modeling has been a very active area of research with approaches ranging from physics-informed neural networks [54], through operator learning [34, 32, 33], to U-Net-based ones [15]. Very recent efforts towards better comparability of such methods have resulted in benchmarking projects led by academia [60] and industry [26]. However, both of these projects exclusively provide the software infrastructure for training models on problems in the Eulerian description and cannot be trivially extended to particle dynamics in 3D space. We note that convolutional neural networks (CNNs) are well suited for problems in the Eulerian description, whereas graph neural networks (GNNs) are best suited to describe points in continuous space. Given the somewhat shifted upside of research on GNNs with respect to CNNs, we conjecture that particle-based PDE surrogates are still in their early stages of development.
Lagrangian Machine Learning. Generic learned physics simulators for Lagrangian rigid and deformable bodies were developed in a series of works using graph-based representations starting with the Interaction Networks [5] and their hierarchical version [46]. There are numerous successors on this line of research mainly based on the graph networks (GN) formalism [6] which improve the network architecture [55, 56], demonstrate great performance on control tasks [30], or introduce interactions on meshes in addition to the interaction radius-based graph [50, 18]. An alternative approach to modeling fluid interactions builds on the idea of continuous convolutions over the spatial domain [65], which has recently been extended to a more physics-aware design by enforcing conservation of momentum [51]. An orthogonal approach to enforcing physical symmetries (like equivariance with respect to the Euclidean group of translation, rotations, and reflections) requires working in the spherical harmonics basis and has been originally introduced for molecular interactions [62, 7, 12], and recently also to fluid dynamics problems [64]. Last but not least, some of the learned particle-based surrogates take inspiration from classical solvers to achieve equivariance [57] or to improve model performance on engineering systems [31].
In terms of data analysis, the closest to our work are those by Li and Farimani [31] and Toshev et al. [64]. In Li and Farimani [31], a dam break and a water fall datasets are presented, the first of which is similar to our dam break, but without explicit wall particles. The Toshev et al. [64] paper uses a 3D TGV and a 3D RPF dataset of every 10th step of the numerical simulator. The fluid datasets from Sanchez-Gonzalez et al. [56] are similar in terms of their physics to the dam break from Li and Farimani [31], and also do not target temporal coarse-graining. We are the first to propose the challenging task of directly predicting every 100th numerical simulator step, on top of which we also introduce 2D/3D LDC, 2D TGV, and 2D RPF. The datasets we present are of diverse and well-studied engineering fluid mechanics systems, each of which has its unique challenges. To the best of our knowledge, there are still no established Lagrangian dynamics datasets to benchmark the performance of new models and we try to fill this gap.
To address the above shortcomings of existing research, we contribute in the following ways:
- Seven new fluid mechanics datasets generated using the SPH method. These include the 2D and 3D Taylor-Green vortex (TGV) [9, 10], 2D and 3D reverse Poiseuille flow (RPF) [17], 2D and 3D lid-driven cavity (LDC) [21], and 2D dam break (DAM) [16]. Each of
Figure 1: Time snapshots of our datasets, at the initial time (top),  $40\%$  (middle), and  $95\%$  (bottom) of the trajectory. Color temperature represents velocity magnitude. (a) Taylor Green vortex (2D and 3D), (b) Reverse Poiseuille flow (2D and 3D), (c) Lid-driven cavity (2D and 3D), (d) Dam break (2D).
these datasets captures different dynamics: TGV demonstrates the onset of turbulence, RPF has a spatially dependent external force field, LDC has static and moving wall boundaries, and DAM has a free surface.
- Efficient JAX-based API with various recent training strategies including additive random walk-type input noise [56] and the push-forward trick [13]. Also, we offer three neighbor search backends two of which are the JAX-MD implementation [58] and a more memory-efficient version thereof, and the third one being the more flexible CPU-based implementation from matscripy.2
- JAX implementation and baseline results of established GNNs including: GNS [56], EGNN [57], SEGNN [12], and an adapted PaiNN model [59].
2 Datasets
To generate the datasets introduced in this work, we use the Lagrangian SPH scheme of Adami et al. [1] to solve the weakly-compressible NSEs [45] in two- and three-dimensional space. A non-dimensional form of the governing equations reads
$$
\frac {\mathrm {d}}{\mathrm {d} t} (\rho) = - \rho (\nabla \cdot \dot {\mathbf {p}}), \tag {1}
$$
$$
\frac {\mathrm {d}}{\mathrm {d} t} (\dot {\mathbf {p}}) = - \frac {1}{\rho} \nabla p + \frac {1}{\operatorname {R e}} \nabla^ {2} \dot {\mathbf {p}} + \frac {1}{\rho} \mathbf {F}, \tag {2}
$$
where we denote with  $\rho$  the density,  $\dot{\mathbf{p}}$  velocity,  $p$  pressure, Re Reynolds number (or the inverse of the nondimensional kinematic viscosity), and  $\mathbf{F}$  an external volumetric force field. For weaklycompressible flows, density fluctuations remain small and the pressure  $p$  is coupled with the density by the barotropic equation of state  $p(\rho) = c_0^2 (\rho -\rho_0) + p_{bg}$ , where  $c_{0}$  denotes an (artificial) speed-of-sound,  $\rho_0$  a reference density, and  $p_{bg}$  a background pressure [2].
SPH is a very well-established Lagrangian approximation method of the Navier-Stokes equations. It was introduced in the 70s [23, 35] and has become a standard method for numerical fluid mechanics. The core idea of SPH is to discretize a domain with fluid particles and define the properties of the fluid at some particular locations through a truncated radial kernel interpolation over neighboring particles. By rewriting the NSE in terms of kernel interpolations of the fluid properties (i.e. velocity, pressure, density), we arrive at a system of ODEs for the particle accelerations. By integrating twice we update the velocities and positions of the particles.
$^{2}$ https://github.com/libAtoms/matscipy
To allow for the most fair runtime comparison in the baselines section, we implemented the SPH method presented in Adami et al. [1] in JAX and plan to open-source our code in the near future. Until then, we include validation results with more details on the solver parameters in Appendix B.
2.1 Overview of datasets
Our datasets are generated following a few basic design principles:
1. Each physical system should contain different physics and at the same time have a solid theoretical underpinning in fluid mechanics literature.
2. The Reynolds number is selected as large as possible to still resolve all relevant scales and at the same time to not exceed 10k particles. Above this limit and depending on model choice one would need to discuss domain decomposition and multi-GPU computing, which is beyond the scope of the current work.
3. The number of training instances was deduced on the basis of scaling runs, which we discuss in Section 3. The size of the validation and test runs was chosen to achieve a ratio of training/Validation/testing splits of  $2/1/1$ . The reason for the comparably large number of validation and testing instances is that we compute up to 20-step rollout errors which are much noisier than 1-step errors and need to be averaged over more trajectories.
Also, all of our datasets are generated by subsampling every 100th step of the SPH solver to train towards an effective particle acceleration, i.e. temporal coarse-graining. As the stepsize of the numerical solver is governed by the CFL condition (see Equation (19) in Adami et al. [1]), for which we use a speed of sound of  $10\times$  the maximum flow velocity and a CFL number of 0.25, then by taking every 100th step the fastest particle would move roughly  $2.5\times$  the average particle distance. Learning to directly predict 100 times larger steps is a nontrivial task and given that we use an interaction radius of around  $1.5\times$  the average particle distance for the connectivity graph, graph-based approaches need at least two layers to even cover the necessary receptive field.
Table 1: Datasets overview. Either the trajectory length or the number of trajectories is used for splitting into training/validation/testing datasets.
Dataset	Particle number	Trajectory length	Trajectory count	Δt [×10-3]	Δx [×10-3]	Box L×H×D [−]	Re
2D TGV	2500	126	100/50/50	40	20	1×1	100
2D RPF	3200	20k/10k/10k	1	40	25	1×2	10
2D LDC	2708	10k/5k/5k	1	40	20	1.12×1.12	100
2D DAM	5740	401	50/25/25	30	20	5.486×2.12	40k
3D TGV	8000	61	200/100/100	500	314.16	2π×2π×2π	50
3D RPF	8000	10k/5k/5k	1	100	50	1×2×0.5	10
3D LDC	8160	10k/5k/5k	1	90	41.667	1.25×1.25×0.5	100
Table 1 summarizes important physical properties of the datasets including the number of particles and trajectories, as well as the physical time step between training instances  $\Delta t$  and the average particle distance  $\Delta x$ . To generate multiple trajectories of the same system (for TGV and DAM), we randomly draw the positions of the particles and let the system relax under periodic (for TGV) and non-periodic (for DAM) boundary conditions via 1000 steps of SPH relaxation. This relaxed state is then used as the initial state of the trajectories. For the statistically stationary cases (RPF and LDC) the simulation starts with fluid velocities equal to zero and runs until equilibrium. We start collecting data for the datasets after this point. In these cases, the training/validation/testing data is obtained by splitting one very long trajectory in its first half for training, third quarter for validation, and last quarter for testing. All trajectories are generated using the same Reynolds number  $Re$ , which is also included in Table 1.
Both lid-driven cavity and dam break have solid wall boundaries, implemented using the established generalized wall boundary conditions approach [1]. This approach relies on representing the walls
with "dummy" particles and then 1) enforcing a no-slip boundary condition at the wall surface by assigning to the wall particles the opposite velocity of that of the closest fluid particles, and 2) enforcing impermeability by assigning the same pressure to the wall particles as the pressure of the closest fluid. This approach requires using multiple layers of wall particles (in our case three), but we make the ML task even more difficult by leaving only the innermost wall layer in the dataset.
In the following, we briefly describe the datasets. For more details on the datasets we refer to Appendix C, and for dataset visualizations beyond Figure 1 to Appendix D.
Decaying Taylor-Green Vortex (TGV). This problem is characterized by periodic boundary conditions on all sides and a distinct initial velocity field without external driving force, which leads to a flow with decaying kinetic energy under the influence of viscous interactions. This system was first introduced by Taylor and Green [61] in 3D to study the transition of a laminar flow to turbulence. In two dimensions, this problem has an analytical solution being the exponentially decaying velocity magnitude and kinetic energy [2]. We initialize the field with the velocity profile from [2] given in Equation (3) with  $k = 2\pi$ . In three dimensions, the case has been explored in literature in-depth, and high-fidelity reference solutions have been generated [9, 10]. Unfortunately, with our limit on particle size of around 10k particles, we cannot resolve the Reynolds numbers from the literature and we chose to work with Re=50, for which we generated our own reference solution using the JAX-Fluids solver [8]. The velocity field for the 3D case is given by Equation (4) with  $k = 1$ .
2D:  $u = -\cos (kx)\sin (ky),\qquad v = \sin (kx)\cos (ky).$  (3)
3D:  $u = \sin (kx)\cos (ky)\cos (kz),\quad v = -\cos (kx)\sin (ky)\cos (kz),\quad w = 0.$  (4)
Reverse Poiseuille flow (RPF). We chose the reverse Poiseuille flow problem [17] over the more popular Poiseuille flow, i.e. laminar channel flow, because it (a) does not require the treatment of boundary conditions as it is fully periodic and (b) has a spatially varying force field. And yet in the laminar case, we still should get the same solution as for the Poiseuille flow. However, our dataset is not fully laminar at  $Re = 10$  and some mixing can be observed, making the dynamics more diverse. The physical setup consists of a force field of magnitude 1 in the lower half of the domain and -1 in the upper half. By this constant force, the system reaches a statistically stationary state, which in contrast to the regular Poiseuille flow exhibits rotational motion in the shearing layers between the two opposed streams. The only difference between the two-dimensional and three-dimensional versions of this dataset is the addition of a third periodic dimension without any significant change in the dynamics.
Lid-driven cavity (LDC). This case introduces no-slip boundaries, including a moving no-slip boundary being the lid that drives the flow. Treating wall boundaries with learned methods is a non-trivial topic that is typically treated by either 1) adding node features corresponding to boundary distance [56] or 2) having explicit particles representing the wall [29]. The second approach has a greater generalization capability and we choose to include boundary particles in our datasets. However, instead of keeping all three layers of boundary particles required by the generalized boundary approach [1] in conjunction with the used Quintic spline kernel [45], we just keep the innermost layer by which we significantly reduce the number of particles in the system. The nominally three-dimensional version of the LDC problem with periodic boundary conditions in  $z$ -direction recovers the two-dimensional solution and learned surrogates should be able to learn the right dynamics in both formulations.
Dam break (DAM). The dam break system [16] is one of the most prominent examples of SPH because it demonstrates its capability of simulating free surfaces with a rather small amount of discretization points. In terms of classical SPH, one major difference between dam break and the previous systems is that the density cannot be simply computed by density summation due to the not full support of the smoothing kernel enforcing the use of density evolution [1]. This is a potential challenge for learned surrogates, on top of which this case is typically simulated as inviscid, i.e.  $Re = \infty$ . To stabilize the simulation we used the established "artificial viscosity" approach [42] and added on top a very small amount of physical viscosity (leading to  $Re = 40\mathrm{k}$ ) to reduce artifacts at the walls. We do not include a 3D version of this dataset as resolving all relevant features would result in exceeding our limit on the number of particles of  $10\mathrm{k}$ .
2.2 Data format and access, software stack, extensibility
The datasets are stored as HDF5 files with one file for each of the training/validation/testing splits, and one JSON metadata file per dataset. The size of the datasets ranges between 0.3-1.5 GB for the 2D ones and 1.5-2 GB for the 3D ones, resulting in 8 GB total size of all 7 datasets. This size should allow for quick model development and testing, while still covering different types of physics when training on all datasets. The datasets are hosted on Zenodo and can be accessed under the DOI: doi.org/10.5281/zenodo.10021925, [63]. A complete datasheet is provided in Appendix A.
Our benchmarking repository is designed for performance and is built on the JAX library [11] in Python, which is a high-performance numerical computing library designed for machine learning and scientific computing research. We started working on this benchmarking project before the major 2.0 release of the more popular PyTorch library [49], but even after this release introduced many of the speedup benefits of JAX, our experience shows that JAX is faster on graph machine learning tasks at the time of writing this paper. Check Appendix F for detailed speed comparison on a variety of problems. With the simple design of our JAX-based codebase and the use of the established PyTorch Dataloger, we plan to maintain the codebase in the foreseeable future. In our codebase we demonstrate how to download and use our datasets, but also how to download, preprocess, and use some of the datasets provided with the Sanchez-Gonzalez et al. [56] paper. We provide many examples of how to set up machine learning training and inference runs using yaml configuration files and plan to extend our codebase with more models and optimization strategies in the future.
Current limitations of our framework include that to use the speedup of JAX the ML models need to be compiled for the largest training/inference instance if system sizes vary. This limitation to static shapes for the just-in-time compilation is a known weakness of JAX and we note that it might be an issue to run inference on very large systems. To even allow for a varying number of particles during training we integrated the mstatsipy backend for the neighbors search because, to the best of our knowledge, there is no trivial way to achieve this functionality with the JAX-MD implementation [58]. In addition, we also improve the memory requirements of the JAX-MD neighbor search routine (at the cost of runtime) by serializing a memory-heavy vectorized distance computation over adjacent cells in the cell list. This serialization significantly reduces the memory requirements of the code and could allow for computations with larger batch sizes or larger systems, see Appendix E.
Regarding future extensions of our datasets, we plan to add 1) a multi-phase problem, e.g. Rayleigh-Taylor Instability, and 2) a multi-phase flow with surface tension, e.g. drop deformation in shear flow. By multi-phase, we refer to fluids with different material properties. Introducing such complexity will require learning rich particle embeddings. Surface tension is a phenomenon that relates to how two different media interact with each other, and this would force the learned surrogate to learn complex interactions at low-dimensional manifolds.
3Baselines
3.1 Learning problem
We define the task as the autoregressive prediction of the next state of a Lagrangian flow field. For simplicity, we adapt the notation from [56]. Given the state  $\mathbf{X}^t$  of a particle system at time  $t$ , one full trajectory of  $K + 1$  steps can be written as  $\mathbf{X}^{t_0:t_K} = (\mathbf{X}^{t_0},\dots,\mathbf{X}^{t_K})$ . Each state  $\mathbf{X}^t$  is made up of  $N$  particles, namely  $\mathbf{X}^t = (\mathbf{x}_1^t,\mathbf{x}_2^t,\ldots ,\mathbf{x}_N^t)$ , where each  $\mathbf{x}_i$  is the state vector of the  $i$ -th particle. Similarly, the system particle positions at time  $t$  are defined as  $\mathbf{P}^t = (\mathbf{p}_1^t,\mathbf{p}_2^t,\dots,\mathbf{p}_N^t)$ . The input node-wise features  $\mathbf{x}_i^t$  of particle  $i$  are extracted for every time step  $t$  based on a window of previous positions  $\mathbf{p}_i^{t - H - 1:t}$  and optionally other features. In particular, the node features can be selected from:
1. The current position  $\mathbf{p}_i^t$
2. A time sequence of  $H$  previous velocity vectors  $\dot{\mathbf{p}}_i^{t - H - 1:t}$
3. External force vector  $\mathbf{F}_i^t$  (if available).
4. The distance to the bounds  $\mathbf{B}_i^t$  (optional).
The nodes are connected based on an interaction radius of  $\sim 1.5$  times the average interparticle distance similar to what is done in [56], which results in around 10-20 one-hop neighbors. Regarding