In this study, we use GNN as a surrogate simulator to model granular flow behavior. Figure 3 shows an overview of the general concepts and structure of the GNN-based simulator (GNS). Consider a granular flow domain represented as particles (Figure 3a). In GNS, we represent the physical state of the granular domain at time t with a set of \(\mathbf{x}_i^t\) describing the state and properties of each particle. The GNS takes the current state of the granular flow \(\mathbf{x}_t^i \in \mathbf{X}_t\) and predicts its next state \(\mathbf{x}_{i+1}^i \in \mathbf{X}_{t+1}\) (Figure 3a). The GNS consists of two components: a parameterized function approximator \(d_\mathbf{\Theta}\) and an updater function (Figure 3b). The approximator \(d_\theta\) takes \(\mathbf{X}_t\) as an input and outputs dynamics information \(\mathbf{y}_i^t \in \mathbf{Y}_t\). The updater then computes \(\mathbf{X}_{t+1}\) using \(\mathbf{Y}_t\) and \(\mathbf{X}_t\). Figure 3c shows the details of \(d_\theta\) which consists of an encoder, a processor, and a decoder. The encoder (Figure 3c-1) takes the state of the system \(\mathbf{X}^t\) and embed it into a latent graph \(G_0=\left(\mathbf{V}_0, \mathbf{E}_0\right)\) to represent the relationship between particles, where the vertices \(\mathbf{v}_i^t \in \mathbf{V}_0\) contain latent information of the current particle state, and the edges \(\mathbf{e}_{i,j}^t \in \mathbf{E}_0\) contain latent information of the pair-wise relationship between particles. Next, the processer (Figure 3c-2) converts \(G_0\) to \(G_M\) with \(M\) stacks of message passing GNN (\(G_0\rightarrow G_1\rightarrow\cdots\rightarrow G_M\)) to compute the interaction between particles. Finally, the decoder (Figure 3c-3) extracts dynamics of the particles (\(\mathbf{Y}^t\)) from \(G_M\), such as the acceleration of the physical system. The entire simulation (Figure 3a) involves running GNS surrogate model through \(K\) timesteps predicting from the initial state \(\mathbf{X}_0\) to \(\mathbf{X}_K\) (\(\mathbf{X}_0, \mathbf{X}_1, \ldots, \mathbf{X}_K\)), updating at each step (\(\mathbf{X}_t\rightarrow\mathbf{X}_{t+1}\))

Input

The input to the GNS, \(\mathbf{x}_i^t \in \mathbf{X}^t\), is a vector consisting of the current particle position \(\mathbf{p}_i^t\), the particle velocity context \({\dot{\mathbf{p}}}_i^{\le t}\), information on boundaries \(\mathbf{b}_i^t\), and particle type embedding \({\mathbf{f}}\) (Eq. 4). \(\mathbf{x}_i^t\) will be used to construct vertex feature (\(\mathbf{v}_i^t\)) (Eq. 6).

The velocity context \({\dot{\mathbf{p}}}_i^{\le t}\) includes the current and previous particle velocities for n timesteps \(\left[{\dot{\mathbf{p}}}_i^{t-n},\cdots, {\dot{\mathbf{p}}}_i^t\right]\). We use \(n=4\) to include sufficient velocity context in the vertex feature \(\mathbf{x}_i^t\). Sanchez-Gonzalez et al. (2020) show that having \(n>1\) significantly improves the model performance. The velocities are computed using the finite difference of the position sequence (i.e., \({\dot{\mathbf{p}}}_i^t=\left(\mathbf{p}_i^t-\mathbf{p}_i^{t-1}\right)/\Delta t\)). For a 2D problem, \(\mathbf{b}_i^t\) has four components each of which indicates the distance between particles and the four walls. We normalize \(\mathbf{b}_i^t\) by the connectivity radius, which is explained in the next section, and restrict it between 1.0 to 1.0. \(\mathbf{b}_i^t\) is used to evaluate boundary interaction for a particle. \({\mathbf{f}}\) is a vector embedding describing a particle type.

In addition to \(\mathbf{x}_i^t\), we define the interaction relationship between particles \(i\) and \(j\) as \(\mathbf{r}_{i, j}^t\) using the distance and displacement of the particles in the current timestep (see Eq. 5). The former reflects the level of interaction, and the latter reflects its spatial direction. \(\mathbf{r}_{i, j}^t\) will be used to construct edge features (\(\mathbf{e}_{i,j}^t\)).

Encoder

The vertex and edge encoders (\(\varepsilon_\Theta^v\) and \(\varepsilon_\Theta^e\)) convert \(\mathbf{x}_i^t\) and \(\mathbf{r}_{i, j}^t\) into the vertex and edge feature vectors (\(\mathbf{v}_i^t\) and \(\mathbf{e}_{i,j}^t\)) (Eq. 6) and embed them into a latent graph \(G_0=\left(\mathbf{V}_0, \mathbf{E}_0\right)\), \(\mathbf{v}_i^t \in \mathbf{V}_0\), \(\mathbf{e}_{i,j}^t \in \mathbf{E}_0\).

We use a two-layered 128-dimensional multi-layer perceptron (MLP) for the \(\varepsilon_\Theta^v\) and \(\varepsilon_\Theta^e\). The MLP and optimization algorithm search for the best candidate for the parameter set \(\Theta\) that estimates a proper way of representing the physical state of the particles and their relationship which will be embedded into \(G_0\).

The edge encoder \(\varepsilon_\Theta^v\) uses \(\mathbf{x}_i^t\) (Eq. 4) without the current position of the particle (\(\mathbf{p}_i^t\)), but still with its velocities (\({\dot{\mathbf{p}}}_i^{\le t}\)), since velocity governs the momentum, and the interaction dynamics is independent of the absolute position of the particles. Rubanova et al. (2022) confirmed that including position causes poorer model performance. We only use \(\mathbf{p}_i^t\) to predict the next position \(\mathbf{p}_i^{t+1}\) based on the predicted velocity \({\dot{\mathbf{p}}}_i^{t+1}\) (Eq. 9).

We consider the interaction between two particles by constructing the edges between them only if vertices are located within a certain distance called connectivity radius \(R\) (see the shaded circular area in Figure 3b). The connectivity radius is a critical hyperparameter that governs how effectively the model learns the local interaction. \(R\) should be sufficiently large to include the local interaction as edges between particles but also to capture the global dynamics of the simulation domain.

Processor

The processor performs message passing (based on Eq. 1-3) on the initial latent graph (\(G_0\)) from the encoder for M times (\(G_0\rightarrow G_1\rightarrow\cdots\rightarrow G_M\)) and returns a final updated graph \(G_M\). We use two-layered 128-dimensional MLPs for both message construction function \(\phi_{\mathbf{\Theta}_\phi}\) and vertex update function \(\gamma_{\mathbf{\Theta}_r}\), and element-wise summation for the message aggregation function \(\mathbf{\Sigma}_{j \in N\left(i\right)}\) in Eq. 1-3. We set \(M=10\) to ensure sufficient message propagation through the network. These stacks of message passing models the propagation of information through the network of particles.

Decoder

The decoder \(\delta_\Theta^v\) extracts the dynamics \(\mathbf{y}_i^t \in \mathbf{Y}^t\) of the particles from the vertices \(\mathbf{v}_i^t\) (Eq. 7) using the final graph \(G_M\). We use a two-layered 128-dimensional MLP for \(\delta_\Theta^v\) which learns to extract the relevant particle dynamics from \(G_M\).

Updater

We use the dynamics \(\mathbf{y}_i^t\) to predict the velocity and position of the particles at the next timestep (\({\dot{\mathbf{p}}}_i^{t+1}\) and \(\mathbf{p}_i^{t+1}\)) based on Euler integration (Eq. 8 and Eq. 9), which makes \(\mathbf{y}_i^t\) analogous to acceleration \({\ddot{\mathbf{p}}}_i^t\).

Based on the new particle position and velocity, we update \(\mathbf{x}_i^t \in \mathbf{X}^t\) (Eq. 5) to \(\mathbf{x}_i^{t+1} \in \mathbf{X}^{t+1}\). The updated physical state \(\mathbf{X}^{t+1}\) is then used to predict the position and velocity for the next timestep.

The updater imposes inductive biases to GNS to improve learning efficiency. GNS does not directly predict the next position from the current position and velocity (i.e., \(\mathbf{p}_i^{t+1}=GNS\left(\mathbf{p}_i^t, {\dot{\mathbf{p}}}_i^t\right)\)) which has to learn the static motion and inertial motion. Instead, it uses (1) the inertial prior (Eq. 8) where the prediction of next velocity \({\dot{\mathbf{p}}}_i^{t+1}\) should be based on the current velocity \({\dot{\mathbf{p}}}_i^t\) and (2) the static prior (Eq. 9) where the prediction of the next position \(\mathbf{p}_i^{t+1}\) should be based on the current position \(\mathbf{p}_i^t\). These make GNS to be trivial to learn static and inertial motions that is already certain and focus on learning dynamics which is uncertain. In addition, since the dynamics of particles are not controlled by their absolute position, GNS prediction can be generalizable to other geometric conditions.

GNN layers such as a IN layer can be easily implemented in PyTorch Geometric (PyG). In PyG, a GNN layer is generally implemented as a subclass of the MessagePassing class. We follow this convention and define the InteractionNetwork Class as follows

(1) Construct a message for each edge of the graph. The message is generated by concatenating the features of the edge’s two nodes and the feature of the edge itself, and transforming the concatenated vector with an MLP.

(2) Aggregate (sum up) the messages of all the incoming edges for each node.

(3) Update node features and edge features. Each edge’s new feature is the sum of its old feature and the message on the edge. Each node’s new feature is determined by its old feature and the aggregation of messages.

Let’s include the encoder, the processor and the decoder together! Before GNN layers, input features are transformed by MLP so that the expressiveness of GNN is improved without increasing GNN layers. After GNN layers, final outputs (accelerations of particles in our case) are extracted from features generated by GNN layers to meet the requirement of the task.