Graph neural network for multitask prediction of rheological and microstructural behavior in suspensions

Fast prediction of suspension rheology is fundamental for optimizing process efficiency and performance in numerous industrial settings. However, traditional simulations are computationally demanding due to explicit evaluation of contact networks and stress tensors in dense regimes approaching shear thickening and jamming. This study presents a microstructure-informed multitask learning framework based on the graph neural network (GNN) that learns an implicit mapping between particle configurations and emergent microstructural and rheological properties of suspensions. This model simultaneously predicts particle pressure $Π$, viscosity $η$, and friction coordination $Z_μ$, in a dynamic steady-state, without explicit knowledge of interparticle forces. Here, semi-dilute to dense suspension systems in 2D were simulated across a wide range of shear stresses $σ$, spanning continuous, discontinuous shear thickening, and shear-jamming conditions. The trained models demonstrated high correlation coefficients ($R^2$ = 0.99) with narrow mean absolute error for packing fractions up to $ϕ\le ϕ_J^μ$ for all predictive targets. However, prediction scatter increases near jamming conditions, attributed to inherent fluctuations in suspension behavior as the critical packing fraction is approached, yet predictions remain in excellent agreement, closely following the trend of the simulated flow curves across stress evolution. Once trained, the model can infer rheological responses directly from structural topology, avoiding explicit stress evaluation during prediction. The approach yields computationally efficient mesoscale surrogates for accelerated simulation with potential for real-time exploration of particulate suspension behavior.

💡 Research Summary

The authors present a novel data‑driven framework that predicts the rheological and microstructural properties of dense non‑Brownian suspensions using a graph neural network (GNN). Instead of relying on computationally intensive lubrication‑discrete element method (LF‑DEM) simulations that explicitly calculate hydrodynamic and contact forces, the proposed approach extracts only the geometric configuration of particles at a single steady‑state snapshot and encodes it as a graph: particles become nodes (with a one‑hot encoded radius) and pairwise relationships become edges carrying dimensionless gap, relative distance, and directional information (x, y components, sine and cosine of the inter‑particle vector).

A multitask learning (MTL) strategy is employed so that a single Deep Graph Convolutional Network (DeepGCN) predicts three correlated outputs simultaneously: relative viscosity η_r, particle pressure Π, and frictional coordination number Z_µ. The multitask setting leverages the intrinsic physical coupling among these quantities, reduces over‑fitting risk, and improves data efficiency compared with training separate single‑task models.

The training data consist of 380 graphs generated from 2‑D LF‑DEM simulations covering packing fractions ϕ = 0.70–0.80 and shear stresses σ/σ₀ = 0.5–200 (σ₀ is the characteristic stress set by the critical load model). For each (ϕ, σ) pair a snapshot is taken after the system reaches a dynamic steady state, and the three target quantities are extracted directly from that snapshot rather than averaging over time. The dataset is split into five folds; each fold is used once as a test set while the remaining four folds train the network, ensuring unbiased performance estimates.

The DeepGCN architecture comprises multiple residual graph convolutional layers (ResGCN). Residual connections mitigate the vanishing‑gradient problem, allowing the network to be deep enough to capture complex, non‑linear relationships. After each convolution, node features are batch‑normalized and passed through a ReLU activation. The final linear head outputs the three scalar predictions. Training uses a weighted mean‑squared‑error loss across tasks, Adam optimizer, and early‑stopping based on validation loss.

Results show excellent agreement with the LF‑DEM reference: the coefficient of determination R² reaches 0.99 for all three outputs, and mean absolute errors remain below a few percent of the respective scales for packing fractions up to the frictional jamming limit ϕ_J^µ. Near the jamming transition, prediction scatter increases, reflecting the intrinsic fluctuations of the suspension, but the overall trends of viscosity thickening, pressure rise, and coordination number growth are faithfully reproduced. Inference is extremely fast—sub‑millisecond per graph on a standard CPU—representing a speed‑up of several orders of magnitude over full LF‑DEM simulations.

The study highlights several key insights: (1) the microstructural topology alone encodes sufficient information to infer macroscopic rheology, eliminating the need for explicit force calculations; (2) multitask learning exploits the physical coupling among viscosity, pressure, and coordination number, yielding more robust models; (3) residual GCNs provide a stable training platform for deep graph architectures in soft‑matter applications. Limitations include the restriction to two‑dimensional systems, a binary size distribution, and the absence of direct experimental validation.

In conclusion, the paper demonstrates that a graph‑based, multitask deep learning model can serve as a high‑fidelity surrogate for suspension rheology, enabling rapid “what‑if” exploration of process parameters and potentially supporting real‑time digital twins in industrial settings. Future work should extend the methodology to three dimensions, incorporate broader particle size/shape distributions, and integrate experimental data to further validate and generalize the approach.