GrainGNN: A dynamic graph neural network for predicting 3D grain microstructure

TL;DR

GrainGNN introduces a dynamic-graph surrogate for 3D grain microstructure evolution during rapid solidification in metal additive manufacturing. It integrates a reduced graph representation with two graph-transformer LSTMs (a regressor and a classifier) to evolve grain features and topology, coupled with a graph-update procedure that preserves valid triple-junction structures and reconstructs a 3D grain orientation field. Trained on phase-field data across a broad range of meltpool conditions, GrainGNN achieves 80–90% pointwise accuracy and KS-consistent QoI distributions while delivering 150×–2000× CPU speedups over high-fidelity PF simulations, enabling simulation of thousands of grains in seconds. The method generalizes to unseen process parameters and larger domain sizes, offering a scalable, efficient tool for AM process optimization and microstructure prediction.

Abstract

We propose GrainGNN, a surrogate model for the evolution of polycrystalline grain structure under rapid solidification conditions in metal additive manufacturing. High fidelity simulations of solidification microstructures are typically performed using multicomponent partial differential equations (PDEs) with moving interfaces. The inherent randomness of the PDE initial conditions (grain seeds) necessitates ensemble simulations to predict microstructure statistics, e.g., grain size, aspect ratio, and crystallographic orientation. Currently such ensemble simulations are prohibitively expensive and surrogates are necessary. In GrainGNN, we use a dynamic graph to represent interface motion and topological changes due to grain coarsening. We use a reduced representation of the microstructure using hand-crafted features; we combine pattern finding and altering graph algorithms with two neural networks, a classifier (for topological changes) and a regressor (for interface motion). Both networks have an encoder-decoder architecture; the encoder has a multi-layer transformer long-short-term-memory architecture; the decoder is a single layer perceptron. We evaluate GrainGNN by comparing it to high-fidelity phase field simulations for in-distribution and out-of-distribution grain configurations for solidification under laser power bed fusion conditions. GrainGNN results in 80\%--90\% pointwise accuracy; and nearly identical distributions of scalar quantities of interest (QoI) between phase field and GrainGNN simulations compared using Kolmogorov-Smirnov test. GrainGNN's inference speedup (PyTorch on single x86 CPU) over a high-fidelity phase field simulation (CUDA on a single NVIDIA A100 GPU) is 150$\times$--2000$\times$ for 100-initial grain problem. Further, using GrainGNN, we model the formation of 11,600 grains in 220 seconds on a single CPU core.

Figures (13)

• Figure 1: Time evolution of the solidified microstructure of stainless steel 316L predicted by a phase field model and GrainGNN. (a) Phase field simulation. (b) GrainGNN predictions. $z_l$ is the current height of the solid-liquid interface (SLI). The domain size is $(120 \mu m, 120 \mu m, 50\mu m)$. The microstructure of 900 grains at $z_l=2\mu m$ is the input to both phase field simulation and GrainGNN. We apply a temperature field with $G=10K/\mu m$ and $R=2m/s$. The grains are growing in the $z$-direction and stopped at $z_l=50 \mu m$. The elapsed simulation time for the phase field model is 48 minutes on one Nvidia A100 GPU. The GrainGNN inference time is 19 seconds on one Intel x86 6-core CPU.
• Figure 2: A phase field simulation of epitaxial growth of 900 grains. Temperature gradient $G_z=10\ K/\mu m$ and pulling velocity $R_z=2\ m/s$ (a) Initial substrate with width $120 \mu m$ and height $2\mu m$. $\boldsymbol{\theta}_z$ is the angle between the crystal orientation and the $z$-axis. (b) Phase field microstructure after 10K time steps. The time step size is 2.41 nanoseconds. The final height of the interface is 50 $\mu m$. (c) The corresponding grain index field of (b). Each grain is assigned a unique index from 1 to 900. (d) Initial grain orientation distribution. (e) Probability distributions of grain size $d$ at $t=0$ and $t=t_H$. (f) Time evolution of the quantities of interest. $z$ represents the height of the current solid-liquid interface. As the height of the interface increases, the number of eliminated grains increases and the volume-weight misorientation $\Delta \boldsymbol{\theta}$ decreases.
• Figure 3: Graph data extraction from phase field simulations. (a) Snapshots of a phase field simulation. Each snapshot is taken when the solid-liquid interface reaches height $z_l=l\Delta z, l=0,1,2,...$ (b) $I_l$ is a 2D image of grain index field at height $z=z_l$. (c) Graph $G_l$ extracted from the image $I_l$. (d) Zoomed-in inset of a region in the image $I_l$: $j_1$ is a junction pixel with three neighboring grains, $g_1,g_2,g_3$. (e) The same region in the graph $G_l$: $G_l$ has two kinds of vertices and two kinds of edges; $j_1$ is a junction vertex corresponding to the junction pixel in (d); $g_1$ is a grain vertex; $(j_1, j_2)$ is a junction-junction edge; $(j_1, g_1)$ is a junction-grain edge. (f) A zoomed-in region in image $I_{l-1}$: $j_1$ is a junction with the same grain neighbors as $j_1$ in (d). We compare the grain neighbors to identify the same junctions in different graphs.
• Figure 4: Topological events with the update of graph connectivity. (a) A neighbor-switching event requires changes of the graph topology: an $O_E$ operation. Grains $g_3$ and $g_4$ initially are neighbors. After applying $O_E$, $g_1$ and $g_2$ are neighbors. From i@ to ii@, $j_1, j_2$ merge to a quadruple junction; from ii@ to iii@, the quadruple junction splits into two triple junctions. This event replaces edges $(j_1, j_5)$, $(j_2, j_4)$ (in red color) with the edges $(j^{\prime}_2, j_5)$, $(j^{\prime}_1, j_4)$ (in blue color). The total number of edges remains the same. (b) A grain elimination event. From i@ to ii@, a neighbor-switching event occurs on $(j_1, j_2)$; $g_3$ becomes a 2-side grain. ii@ to iii@ is an $O_G$ operation. It removes $g_3$, $j^{\prime}_1$, $j_3$, a couple of edges, and replace them with $(j^{\prime}_2, j_4)$.
• Figure 5: The LSTM architecture. We use a graph transformer LSTM with a regressor $\mathcal{R}$ and a classifier $\mathcal{C}$---both having an encoder-decoder structure. The classifier and regressor are trained separately. The junction nodes (solid circles) and grain nodes (hollow circles) have vertex features defined in \ref{['eqn:junction_feature']} and \ref{['eqn:grain_feature']}. For each target vertex (red circles), graph transformer operators in \ref{['eqn:lstm']} aggregate the features of the vertex itself with the neighboring vertices into a hidden vector. The decoder $\mathcal{R}$ then uses \ref{['eqn:r_decoder']} to transform the hidden vector to the target outputs, which are displacements of junctions, area change, and excess volume of grains. For each junction-junction edge, the model $\mathcal{C}$ concatenates the hidden vectors of the two connecting vertices and predicts the probability of the edge event (see \ref{['eqn:c_decoder']}).