$\mathcal{H}$-HIGNN: A Scalable Graph Neural Network Framework with Hierarchical Matrix Acceleration for Simulation of Large-Scale Particulate Suspensions

TL;DR

This work addresses the computational bottleneck of simulating hydrodynamic interactions in large particulate suspensions by integrating Hierarchical Matrix ($\mathcal{H}$-matrix) techniques into the Hydrodynamic Interaction Graph Neural Network (HIGNN). It models the mobility tensor $\mathbf{M}(\mathbf{X})$ as a sum of one-, two-, and three-body HI terms and learns the two surrogates $\mathbf{h}_{\Theta_2}$ and $\mathbf{g}_{\Theta_3}$ on graph edges and 3-cliques to predict particle velocities via a forward pass, achieving transferability across particle counts and external forces. The key contribution is the $\6H$-HIGNN framework, which uses cluster-tree block partitioning, adaptive cross approximation (ACA), and a GPU-accelerated parallel implementation to reduce the dominant $O(N^2)$ two-body matrix-vector multiplications to $O(N\log N)$, validated up to $N=10^7$ particles with near-ideal strong scaling. The results demonstrate quasi-linear scaling, substantial speedups, and applicability to large-scale suspensions of both rigid particles and flexible filaments, while maintaining accuracy controlled by ACA tolerance $\zeta$ and leaf size $C_{\text{leaf}}$. Overall, this work enables practical, high-fidelity simulations of complex suspensions on modest hardware, with implications for industrial and biological systems and potential extensions to active matter.

Abstract

We present a fast and scalable framework, leveraging graph neural networks (GNNs) and hierarchical matrix ($\mathcal{H}$-matrix) techniques, for simulating large-scale particulate suspensions, which have broader impacts across science and engineering. The framework draws on the Hydrodynamic Interaction Graph Neural Network (HIGNN) that employs GNNs to model the mobility tensor governing particle motion under hydrodynamic interactions (HIs) and external forces. HIGNN offers several advantages: it effectively captures both short- and long-range HIs and their many-body nature; it realizes a substantial speedup over traditional methodologies, by requiring only a forward pass through its neural networks at each time step; it provides explainability beyond black-box neural network models, through direct correspondence between graph connectivity and physical interactions; and it demonstrates transferability across different systems, irrespective of particles' number, concentration, configuration, or external forces. While HIGNN provides significant speedup, the quadratic scaling of its overall prediction cost (with respect to the total number of particles), due to intrinsically slow-decaying two-body HIs, limits its scalability. To achieve superior efficiency across all scales, in the present work we integrate $\mathcal{H}$-matrix techniques into HIGNN, reducing the prediction cost scaling to quasi-linear. Through comprehensive evaluations, we validate $\mathcal{H}$-HIGNN's accuracy, and demonstrate its quasi-linear scalability and superior computational efficiency. It requires only minimal computing resources; for example, a single mid-range GPU is sufficient for a system containing 10 million particles. Finally, we demonstrate $\mathcal{H}$-HIGNN's ability to efficiently simulate practically relevant large-scale suspensions of both particles and flexible filaments.

Figures (15)

• Figure 1: Main components of the HIGNN framework.
• Figure 2: Geometry-based partitioning using bounding boxes, illustrated in 2D. The black dots represent particles. From left to right, the initial bounding box is first divided into two boxes along the $x$-direction, and then each of the new bounding boxes is divided along the $y$-direction.
• Figure 3: The diameter and distance of two clusters: (a) definition, (b) approximation by bounding box.
• Figure 4: Example of cluster tree and block partitioning with admissible and inadmissible blocks.
• Figure 5: The accuracy of $\mathcal{H}$-HIGNN evaluated by the relative error defined in Eq. \ref{['Eq:error_equation']} with different $\zeta$ but fixed $C_\text{leaf} = 50$.

Theorems & Definitions (2)

• Definition 1
• Definition 2