SlowGuess/NeurIPS_2023_OCR huggingface dataset:
the edge features, they can be for example the displacement vectors  $\mathbf{d}_{ij}^{t}$  and/or distances  $d_{ij}^{t}$ . The availability of some features is dataset-dependent. For example, the external force vectors are only relevant for reverse Poiseuille and dam break.
For every time step  $t$ , the model will autoregressively predict the next step particle positions  $\mathbf{P}^{(t + 1)}$  from the system state  $\mathbf{X}^t$  by indirectly estimating acceleration  $\ddot{\mathbf{P}}$  or velocity  $\dot{\mathbf{P}}$ , or by directly inferring the next position  $\mathbf{P}$ . When using indirect predictions, the next positions  $\mathbf{P}^{(t + 1)}$  are computed by semi-implicit Euler integration from acceleration or forward Euler from velocity. Finally, the next nodes' state  $\mathbf{X}^{(t + 1)}$  is computed through feature extraction from the updated positions  $\mathbf{P}^{t - H:t + 1}$ .
3.2 Baseline models
We provide a diverse set of graph neural networks [24, 22, 6] adapted for our learning task: GNS [56], which is a popular model for engineering particle problems; EGNN [57], which has shown promising results on N-body dynamics; SEGNN [12], which performs well on similar physical problems [64]; PaiNN [59], which is another simple yet powerful equivariant model from the molecular property-predictions literature. Our choice of GNNs is motivated by their natural extendibility to point cloud data, with relational information between nodes and local neighborhood interactions. Other models not strictly based on Message Passing, such as PointNet++ [53] and Graph UNets [19], have also been considered, but go beyond the scope of the current work. Another non-MP alternative is Graph Transformers [67], but we found them not well-suited for our problems as they would not learn local interactions, which are essential for scaling to larger fluid systems.
Graph Network-based Simulator (GNS). The GNS model [56] is a popular learned surrogate for physical particle-based simulations. It is based on the encoder-processor-decoder architecture [6], where the processor consists of multiple graph network blocks [5]. Even though the architecture is fairly simple, as it employs only fully connected layers and layer norm building blocks, GNS has been proven to be capable of modeling both 2D and 3D particle systems. However, performance on long rollouts is unstable and strongly depends on the choice of Gaussian noise to perturb the inputs.
$\mathbf{E}(\mathfrak{n})$ -equivariant Graph Neural Network (EGNN).  $\mathrm{E}(\mathrm{n})$ -equivariant GNNs are a class of GNNs equivariant with respect to isometries of the  $n$ -dimensional Euclidean space, namely rotations, translations, and reflections. While other popular equivariant models such as Tensor Field Networks [62], NequIP [7], and SEGNN [12] rely on expensive Clebsch-Gordan tensor product operations to compute equivariant features, the EGNN model [57] is an instance of  $\mathrm{E}(n)$  equivariant networks that do not require complex operations to achieve higher order interactions; instead, it treats scalar and vector features separately: the scalars are updated similarly to a traditional GNN, while the vectors (positions) are updated following a layer-stepping method akin to a numerical integrator.
Steerable E(3)-equivariant Graph Neural Network (SEGNN). SEGNN [12] is a general implementation of an E(3) equivariant graph neural network. SEGNN layers are directly conditioned on steerable attributes for both nodes and edges. The main building block is the steerable MLP: a stack of learnable linear Clebsch-Gordan tensor products interleaved with gated non-linearities [66]. SEGNN layers are message passing layers [22] where Steerable MLPs replace the traditional non-equivariant MLPs for both message and node update functions. The flexibility of the SEGNN framework makes the model well suited for a wide array of physical problems, where vectorial and/or tensorial features can be modeled E(3)-equivariantly. Additionally, the steerable node and edge attributes provide a large flexibility for feature engineering choices. SEGNN is designed to work with 3D data and to make it work on the 2D datasets we set the third dimension to zero; this probably results in a less efficient code, but still preserves equivariance.
Polarizable atom interaction Neural Network (PaiNN). PaiNN [59] is yet another variant of E(3) equivariant GNN. Similarly to EGNN, PaiNN models equivariant interactions directly in Cartesian space, meaning separately for scalars and vectors, and does not require tensor products with Clebsch-Gordan coefficients. Because PaiNN was originally designed for molecular property prediction and wasn't intended for general-purpose applications, we had to extend its functionalities in two straightforward yet vital ways: (a) In the original implementation the input vectors are set to the zero vector  $\mathbf{0} \in \mathbb{R}^{F \times 3}$  (where  $F$  is the latent representation dimension) since there is no directional node information available initially for the problems considered in the paper. We add the possibility of
passing non-zero input vectors, which are lifted to  $F$ -dimensional tensors by an embedding layer. (b) We include the Gated readout block [59] to predict vectorial features, e.g. acceleration.
Model naming. The models are named using the scheme " $\{model\} - \{MP layers\} - \{latent size\}"$ . SEGNNs have one additional tag, the superscript "  $L = l$  ", where  $l$  is the maximum tensor product order considered in the steerable MLP layers.
3.3 Error measures
We measure the performance of the models in three aspects when evaluating on the test datasets:
1. Mean-squared error (MSE) of particle positions.  $\mathrm{MSE}_n$  is the  $n$ -step average rollout loss.
2. Sinkhorn distance as an optimal transport distance measure between particle distributions. Lower values indicate that the particle distribution is closer to the reference one.
3. **Kinetic energy error  $E_{kin}$  ( $= 0.5\sum_{i}m_{i}v_{i}^{2}$ ) as a global measure of physical behavior. Especially relevant for statistically stationary datasets (e.g. RPF, LDC) and decaying flows (e.g. TGV).  $E_{kin}$  is a scalar value describing all particles at a given time instance, and the  $MSE_{E_{kin}}$  is the MSE between the rollout and dataset  $E_{kin}$  evolution.
3.4 Data scaling
Before we discuss the performance of different GNN baselines, we investigate the complexity of our datasets. One simple way to achieve that is by training one and the same model on different portions of each of the datasets. We chose to train on  $1\%$ ,  $5\%$ ,  $10\%$ ,  $20\%$ ,  $50\%$ , and  $100\%$  of our training datasets, and expect to see a saturation of model performance if more data does not add more information to the GNN training. We chose to work with the GNS model as it is probably the fastest (given its simplicity) and most popular one. The number of layers was set to 10 and the latent dimension to 64, which in our experience gives reasonable results. The results are summarized in Figure 2.
Figure 2: Scaling evaluation on all datasets. The  $x$ -axis shows the amount of available data, and the  $y$ -axis shows the position MSE loss values. The model is GNS-10-64 trained for 1M steps (with 40k steps early stopping). Every mark represents a new GNS instance trained with a different amount of data.
We started exploring data scaling with the 2D RPF dataset and that is why it has more training instances than most other datasets (2D RPF and 2D DAM have  $20\mathrm{k}$  training frames, whereas all others have  $10\mathrm{k}$ ). It is 2D RPF from which we saw that  $10\mathrm{k}$  samples are enough and that is why most other datasets have roughly  $10\mathrm{k}$  training samples. Only dam break seems to keep improving with more data, probably due to the free surface which demonstrates different wave structures, but we decided to restrict the dataset size to twice as much as the others.
3.5 Baseline performance
We benchmark each dataset by training every model with different hyperparameters. We observed that higher tensor product orders in SEGNN lead to better performance, but we restrict ourselves to  $L = 1$  as higher values are increasingly expensive (see Table 3 for runtimes of  $L = \{1,2\}$ ). Table 2 shows the best model performance we obtained on all the datasets, and thus sets the state-of-the-art on our datasets. We specify the number of layers and hidden dimensions in the same table and refer to Appendix G for a more complete description of the hyperparameters. Also, we refer to Appendix H for a more complete version of Table 2 including all experiments we conducted. Regarding the effect of the training strategies, we performed a comparison study and summarized the results in Appendix I.
An unfortunate result was that we could not make EGNN and our PaiNN version work well on our datasets. EGNNs are unstable in any configurations we tried, and no model converges to reasonable results. It was also reported in other papers [12] that EGNN gradients are often unstable and produce unfavorable results, especially when predicting the force vectors on larger systems. We leave further investigations to future work. To the best of our knowledge, we are the first to extend PaiNN to general vectorial inputs and consequently apply it to engineering data. We aimed to get similar results to SEGNN with  $L = 1$ , as in this special case both models only make use of  $L = \{0,1\}$  features. However, the loss we got was on average two orders of magnitude worse than with GNS or SEGNN possibly because, even with our changes, PaiNN is still designed for atomistic systems, which are in a multitude of ways different from our problem. We include our JAX implementation of both EGNN and PaiNN anyway, together with GNS and SEGNN, in the LagrangeBench code repository as working templates for future research. In Appendix G we include more details on what we tried with EGNN and PaiNN.
Looking at Table 2, we see that GNS performs best in most cases with boundary conditions (LDC and DAM) and SEGNN in most pure fluid problems. This result is consistent for smaller and bigger models (see Appendix H) and can be interpreted in two ways: 1) equivariance is not the right inductive bias for problems with boundary conditions, or 2) representing boundaries with one layer of dummy particles is not optimal for equivariant models. We leave this investigation to future work.
Table 2: Baseline results of the best performing models. During training, the best model weights are tracked and saved based on the  $\mathrm{MSE}_{20}$  loss on the validation dataset. Reported metrics come from the best checkpoint and are averaged on the full test dataset, and over 3 different random seeds. The intervals are the standard deviation over the seeds.
Dataset	Model	MSE5	MSE20	Sinkhorn	MSEEkin
2D TGV	SEGNN-10-64	2.4e-7±5.8e-9	4.4e-6±5.8e-9	2.1e-7±2.8e-8	4.5e-7±8.3e-8
2D RPF	GNS-10-128	1.1e-7±1.2e-9	3.3e-6±1.2e-9	1.4e-7±2.7e-8	1.7e-5±4.4e-7
2D LDC	GNS-10-128	6.4e-7±1.4e-8	1.4e-5±1.4e-8	1.0e-6±1.2e-7	3.7e-3±5.2e-3
2D DAM	GNS-10-128	1.3e-6±7.2e-8	3.3e-5±7.2e-8	1.4e-5±2.0e-6	1.3e-4±1.7e-5
3D TGV	SEGNN-10-64	1.7e-4±3.9e-6	5.2e-3±3.9e-6	6.4e-6±1.5e-6	2.7e-2±4.4e-3
3D RPF	SEGNN-10-64	3.0e-7±2.9e-9	1.8e-5±2.9e-9	2.9e-7±1.6e-8	3.5e-6±1.5e-6
3D LDC	GNS-10-128	7.4e-7±2.1e-8	4.0e-5±2.1e-8	6.0e-7±1.7e-7	2.6e-8±5.0e-9
Further, we compare the performance of some instances of our baseline models in terms of inference speed and memory usage in Table 3. It is known that Clebsch-Gordan tensor products are slower than matrix-vector multiplications in fully connected layers, and we see this in the GNS-SEGNN comparison. We didn't expect EGNNs to be noticeably slower than GNS. We noticed in profiling that a large portion of EGNN runtime is spent in computing spatial periodic displacements and shifts (JAX-MD space functions, implemented with modulo operations). Interestingly, when using non-periodic displacement and shift functions the inference time drops to  $3.49\mathrm{ms}$  for the 2D case (3.2K particles) and to  $15.1\mathrm{ms}$  for the 3D case (8.1K particles), which is closer to the expected runtime. On the other hand, PaiNN is quite fast to run despite having a significant memory usage likely caused by its design: PaiNN applies a  $3\times$  uplifting to the latent embeddings, which are then
Table 3: Inference time and memory usage of the models on 2D and 3D cases (forward only). Runtimes are averaged over 1000 forward passes on one system, RAM is reported as the max memory usage through the rollout. Evaluated on a single Nvidia A6000 48GB GPU.
Model	Params	Inference [ms]	Memory [MB]
2D (3.2K)	3D (8.1K)	2D (3.2K)	3D (8.1K)
GNS-5-64	161K	2.05	8.63	1121	1889
GNS-10-64	307K	3.89	16.4	1121	1889
GNS-10-128	1.2M	6.66	32.0	1377	2913
SEGNN-5-64L=1	183K	15.1	81.2	1377	2913
SEGNN-10-64L=1	360K	29.7	161	1379	4963
SEGNN-10-64L=2	397K	86.0	470	1893	9189
EGNN-5-128	663K	50.0	206	1377	4961
PaiNN-5-128	1.0M	9.09	54.2	3041	17505
split into three different representations. This uplifting requires a large matrix multiplication and we suspect that this causes the large memory footprint of PaiNN.
In our work, we propose LagrangeBench, a novel JAX-based framework to develop and evaluate machine learning models on Lagrangian systems. To our knowledge, this is the first formalized benchmarking tool for engineering particle systems, and we consider it a necessary step in the future development of Lagrangian machine learning solvers. We provide a diverse set of 2D and 3D datasets with a solid foundation from engineering fluid simulations. Finally, we include a collection of baseline models originating from multiple machine learning fields and adapt them to our problem setting.
There are many possibilities for extensions of our framework, the most crucial of which is probably implementing a multi-GPU parallelization based on domain decomposition with load balancing. This would allow for training and inference of much bigger systems potentially up to tens of millions of particles with models like Allegro [47, 48]. There are also alternatives to the given yaml configuration files in order to automate hyperparameter tuning, which might speed up the development cycles of new machine learning models.
The authors thank Dr. Johannes Brandstetter for his great support in the initial phase of the project including SEGNN reimplementation, push-forward trick, and data analysis. Also, thanks to Christopher Zöller for fruitful discussions on the dataset generation.
A.T. and G.G. developed the codebase and ran the experiments together. A.T. led the project and focused on data loading, preprocessing, neighbor search, and reimplemented GNS. G.G. implemented the training pipeline, added the SEGNN, PaiNN, and EGNN models, and engineered the software package. F.F. and S.A. developed the first version of the in-house SPH solver. A.T. extended this solver and generated the datasets advised by F.F. and S.A. in terms of problem selection, literature references, and dataset analysis. N.A. supervised the project from conception to design of experiments and analysis of the results. A.T., G.G., F.F., and N.A. contributed to the manuscript.
[1] S. Adami, X. Hu, and N. A. Adams. A generalized wall boundary condition for smoothed particle hydrodynamics. Journal of Computational Physics, 231(21):7057-7075, 2012.
[2] S. Adami, X. Hu, and N. A. Adams. A transport-velocity formulation for smoothed particle hydrodynamics. Journal of Computational Physics, 241:292–307, 2013.
[3] C. Antoci, M. Gallati, and S. Sibilla. Numerical simulation of fluid-structure interaction by sph. Computers & structures, 85(11-14):879-890, 2007.
[4] M. Antuono, A. Colagrossi, S. Marrone, and D. Molteni. Free-surface flows solved by means of sph schemes with numerical diffusive terms. Computer Physics Communications, 181(3): 532-549, 2010.
[5] P. Battaglia, R. Pascanu, M. Lai, D. Jimenez Rezende, et al. Interaction networks for learning about objects, relations and physics. Advances in neural information processing systems, 29, 2016.
[6] P. W. Battaglia, J. B. Hamrick, V. Bapst, A. Sanchez-Gonzalez, V. Zambaldi, M. Malinowski, A. Tacchetti, D. Raposo, A. Santoro, R. Faulkner, et al. Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv:1806.01261, 2018.
[7] S. Batzner, A. Musaelian, L. Sun, M. Geiger, J. P. Mailoa, M. Kornbluth, N. Molinari, T. E. Smidt, and B. Kozinsky. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature communications, 13(1):2453, 2022.
[8] D. A. Bezgin, A. B. Buhendwa, and N. A. Adams. Jax-fluids: A fully-differentiable high-order computational fluid dynamics solver for compressible two-phase flows. Computer Physics Communications, page 108527, 2022. ISSN 0010-4655.
[9] M. E. Brachet, D. I. Meiron, S. A. Orszag, B. Nickel, R. H. Morf, and U. Frisch. Small-scale structure of the taylor-green vortex. Journal of Fluid Mechanics, 130:411-452, 1983.
[10] M. E. Brachet, D. Meiron, S. Orszag, B. Nickel, R. Morf, and U. Frisch. The taylor-green vortex and fully developed turbulence. Journal of Statistical Physics, 34(5-6):1049-1063, 1984.
[11] J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milne, and Q. Zhang. JAX: composable transformations of Python+NumPy programs, 2018.
[12] J. Brandstetter, R. Hesselink, E. van der Pol, E. J. Bekkers, and M. Welling. Geometric and physical quantities improve e(3) equivariant message passing. In ICLR, 2022.
[13] J. Brandstetter, D. E. Worrall, and M. Welling. Message passing neural PDE solvers. In *ICLR*, 2022.
A Datasheet
Datasheet adapted from "Datasheets for Datasets" [20]. Our datasets are stored on Zenodo and can be found under the following DOI: 10.5281/zenodo.10021925 [63].
The datasets can be used to train autoregressive particle-dynamics surrogates. To the best of our knowledge, these are some of the first Lagrangian fluid dynamics datasets coming from the fluid mechanics literature.
The datasets were created by Artur Toshev on behalf of the Chair of Aerodynamics and Fluid Mechanics at the Technical University of Munich.
No grant.
The dataset instances are simulation trajectories of SPH-based fluid system evolutions in time. We generated the datasets by recording every 100th simulation step of the solver to create a challenging temporal coarse-graining benchmark for Lagrangian fluid dynamics.
See Table 1 in main paper.
We provide all datasets that were used for training our machine-learning models and generating the benchmarking results.
Each of our 7 datasets consists of 4 main files: "train.h5", "valid.h5", "test.h5", "metadata.json", the first three of which are  $2/1/1$  splits of the SPH simulation trajectories for training/validation/testing, and the metadata file contains the parameters used to generate the dataset. Each data split is an HDF5 file with the first dictionary keys being the trajectory number as a string ("00000", "00001", ..., and each trajectory containing two single precision arrays named "position" (shape: [# time steps, # particles, dimension]) and "particle_type" (shape: [# particles]).
In addition, the datasets that have external forces, i.e. RPF and DAM, have a fifth file force.py, which defines the force acting on a particle based on its coordinates.
- Is there a label or target associated with each instance? If so, please provide a description.
Yes. The target in autoregressive dynamics prediction is the next instance in time.
- Is any information missing from individual instances? If so, please provide a description, explaining why this information is missing (e.g., because it was unavailable). This does not include intentionally removed information, but might include, e.g., redacted text.
- Are relationships between individual instances made explicit (e.g., users' movie ratings, social network links)? If so, please describe how these relationships are made explicit. N/A.
- Are there recommended data splits (e.g., training, development/Validation, testing)? If so, please provide a description of these splits, explaining the rationale behind them. The provided 2/1/1 split (see the fourth question of this subsection) is designed to provide enough validation and testing instances for long rollout evaluation.
- Are there any errors, sources of noise, or redundancies in the dataset? If so, please provide a description.
The datasets directly represent the dynamics of the described ground-truth SPH solver.
- Is the dataset self-contained, or does it link to or otherwise rely on external resources (e.g., websites, tweets, other datasets)? If it links to or relies on external resources, a) are there guarantees that they will exist, and remain constant, over time; b) are there official archival versions of the complete dataset (i.e., including the external resources as they existed at the time the dataset was created); c) are there any restrictions (e.g., licenses, fees) associated with any of the external resources that might apply to a dataset consumer? Please provide descriptions of all external resources and any restrictions associated with them, as well as links or other access points, as appropriate.
- Does the dataset contain data that might be considered confidential (e.g., data that is protected by legal privilege or by doctor-patient confidentiality, data that includes the content of individuals' non-public communications)? If so, please provide a description. No.
- Does the dataset contain data that, if viewed directly, might be offensive, insulting, threatening, or might otherwise cause anxiety? If so, please describe why. No.
- Any other comments? No.
The data was generated by performing SPH simulations.
- What mechanisms or procedures were used to collect the data (e.g., hardware apparatus(es) or sensor(s), manual human curation, software program(s), software API(s))? How were these mechanisms or procedures validated?
The simulations were performed using our in-house JAX-based (in Python) SPH solver. The postprocessing was done by means of Python scripts.
See the third question of the previous section on "Composition".
Artur Toshev set up and ran the numerical simulations, and performed the postprocessing. See "Acknowledgements" and "Author Contributions" for further people involved in the dataset design. There were no crowdworkers.
The datasets were generated between July 2022 and May 2023.
- Were any ethical review processes conducted (e.g., by an institutional review board)? If so, please provide a description of these review processes, including the outcomes, as well as a link or other access point to any supporting documentation. No.
- Any other comments?
We subsampled the full solver solution at every 100th step to achieve a temporal coarse-graining problem.
We used the full trajectories to validate the SPH solver but never used them to train machine learning models.
- Has the dataset been used for any tasks already? If so, please provide a description. Not beyond this paper.
- Is there a repository that links to any or all papers or systems that use the dataset? If so, please provide a link or other access point.
We plan on listing papers that use our datasets in the LagrangeBench Github repository: https://github.com/tumaer/lagrangebench.
Any task related to learning the Navier-Stokes dynamics based on particle discretizations.
- Are there tasks for which the dataset should not be used? If so, please provide a description. No.
Yes, the dataset is freely and publicly available and accessible.
- How will the dataset be distributed (e.g., tarball on website, API, GitHub)? Does the dataset have a digital object identifier (DOI)?
The datasets will be distributed as zip files through Zenodo under the DOI doi.org/10.5281/zenodo.10021925.
The dataset is on Zenodo as of October 2023.
The dataset is licensed under "Creative Commons Attribution 4.0 International".
- Do any export controls or other regulatory restrictions apply to the dataset or to individual instances? If so, please describe these restrictions, and provide a link or other access point to, or otherwise reproduce, any supporting documentation. No.
- Any other comments? No.
The dataset will be hosted on Zenodo.