Equivariant GNN-Based Multi-Resolution Simulator for Particle-Based Fluid Mechanics

Project description
Particle-based Computational Fluid Dynamics (CFD) has established as the preferred approach for the simulation of complex physical systems ranging from stellar collisions, having sparsely populated regions, to multi-phase flows encountered in additive manufacturing. However, given the multi-scale nature of these phenomena, such algorithms very often require high-performance computing resources due to 1) the very large number of particles required to represent the relevant physics, as well as 2) iteration steps in the order of millions. Coarse-Graining (CG) has the potential to significantly speed-up simulations by having larger particles, which in turn allows for larger time steps by the Courant–Friedrichs–Lewy condition.

Interestingly, looking at particle-based CFD methods, we observe strong algorithmic similarity to Molecular Dynamics (MD) simulations. That is why we suggest adapting the promising recent work on GNN-based CG for molecular simulations [1] by Fu et al., 2022, to problems in fluid mechanics. The first step would be exchanging the Graph Neural Network (GNN) in the dynamics prediction step from the current GNS network [2] to the equivariant SEGNN [3]. In addition to that, we tackle one particular limitation of the approach suggested in [1], namely the analytical nature of the used graph-clustering algorithm provided by METIS. Given the continuum assumption in CFD, there are no edges corresponding to the bonds in MD, forcing us to use another heuristic for the CG routine - we choose to use a neural network learning a scalar field as a refinement heuristic as done in [4]. Another common technique in MD and CFD, which has not been explored by Fu et al., is to coarse-grain towards a multi-resolution description, i.e. coarsening only regions with homogeneous properties and refining those with large variations.

Project milestones
Apply reference code of [1] to fluid mechanics data (see datasets below). This already requires some modifications, e.g. imposing solid wall boundary conditions.
Replace the dynamics GNN with SEGNN.
Implement a new coarsening and refinement routine as the one in [5] allowing different-sized particles. Decide on coarsening/refinement based on a learned scalar field as done for remeshing in [4].

Datasets
3D datasets from [2].
Our own 3D Smoothed Particle Hydrodynamics dataset containing thousands of trajectories over different systems including the Taylor-Green vortex, Reverse-Poiseuille flow, dam break, and others.

Requirements
Basics of Graph Neural Networks and JAX
Interest in Fluid Mechanics
Teamwork skills

References
[1] Simulate Time-integrated Coarse-grained Molecular Dynamics with Geometric Machine Learning, Fu et al., 2022
[2] Learning to Simulate Complex Physics with Graph Networks, Sanchez-Gonzalez et al., 2020
[3] Geometric and Physical Quantities improve E(3) Equivariant Message Passing, Brandstetter et al., 2022
[4] Learning Mesh-Based Simulation with Graph Networks, Pfaff et al., 2021
[5] A consistent multi-resolution smoothed particle hydrodynamics method, Hu et al., 2017