MindSpore AI科学计算系列（37）：Learning Mesh-Based Simulation论文分析

一． 背景

基于网格的模拟是复杂物理系统建模的核心，由于高维科学模拟的运行成本非常昂贵，网格的分辨率和强大的数值积分方法需要在准确性和效率之间取得平衡，使用AI通过几何外形快速得到流场信息变得尤为重要。在当前的高性能计算中，分为结构网格和非结构网格两类，为复杂的场景进行计算，而对于实际的工程场景，随着外形的复杂化，非结构网格逐步成为数值计算中主流的网格采用方式。由于非结构网格在其网格区域内的内部点不具有相同的毗邻单元，可以是多种形状，因而无法转换为规整的Tensor表示，从而无法使用CNN将几何和流场进行映射，需要采用GNN来对这类非结构化的数据进行处理。DeepMind提出了MeshGraphNets[1]来解决非结构网格的流场快速预测问题。

二． 网络架构

图1. MESHGRPHNETS网络架构

模型采用了Encoder-Processor-Decoder的架构进行训练，并可在推理时反复迭代生成长时间的轨迹。Encoder输入网格M(t) (M(t) = (V, E^M)),表示的是在时刻t的仿真网格有V个节点E^M条边，其转换为图，并增加了额外的World-space边。Processor通过几轮网格边和外部边之间的消息传递，更新了所有节点和边的embedding信息。Decoder提取每个节点的特征，并将其更新到网格产生下一时刻的M(t+1)。

1. Encoder

Encoder模块对当前的网格M(t)进行编码到一个多图G = (V, E^M, E^M)。其中网格点为V，网格点之间的连线为边E^M，这两个可以用来计算网格内部的动量；而对于拉格朗日系统，增加了world edge E^M，可以学习外部的动量，如（自）碰撞和接触。world edge E^M是通过空间邻近性创建的：给定网格边的最小半径为r^W，为节点i和节点j设定一个外部边，这个边符合这个不等式|x(i)-x(j)| < r^W，并排除在网格中已连接的边。

然后，我们将特征编码为图的节点和边。为了实现空间等变，位置特征作为相对边缘特征提供。我们把网格空间内的相对位置矢量u(i)(j) = u(i)-u(j) 和其范数 |u(i)(j)|编码到网格边e(i)(j)^M∈E^M；然后将world-space的相对位置矢量x(i)(j)和其范数|x(i)(j)|编码到e(i)(j)^M∈E^M和 e(i)(j)^W∈E^W；其余动量特征q(i)都被编码到了节点特征v(i)。

最终，上述特征被编码到一个长度为128的latent 向量，网络结构是MLP。

图2 网格空间和world space示意

2. Processor

Processor由L个相同的消息传递模块组成，它概括了图——网络块到多边缘集合。每个消息传递模块都包含了独立的网络参数，将e(i)(j)^M、e(i)(j)^W和v(i)更新为新的量，可以用如下式子表示：

式1

其中f^M、f^W、f^V使用了带残差连接的MLP网络。

3. Decoder and state updater

为了从t时刻的输入预测t+1时刻的状态，decoder使用了一个MLP将最终处理后的latent node特征v(i)转换为一个或多个输出特征p(i)。

我们可以把输出特征p(i)理解为q(i)的高阶导数，并把他们整合，通过一个△t=1前向欧拉积分来计算下一步的动量q(i)^(t+1)。对于一阶系统，输出p(i)积分一次来更新q(i)^(t+1) = p(i) + q(i)^t，而对于二阶系统会进行两次积分：q(i)^(t+1) = p(i) + 2q(i)^t – q^(t-1)。附加的输出特征p(i)仍被用于预测压力和应力等辅助量。最终，输出网格节点V通过q(i)^(t+1)进行更新从而产生M^(t+1)。

我们通过监督每个节点的输出特征p(i)来训练我们的动态模型，解码器对p(i)和ground truth p(i)`计算L2 loss。

三． 实验结果

论文在具有不同底层PDE的系统上评估了此模型的效果，包含布料、结构力学、不可压缩流体和可压缩流体。数据集AIRFOIL的边长范围在2*10^-4m到3.5m，同时也模拟了在轨迹过程中动态改变分辨率的网格。

表1 在不同数据集下使用MeshGraphNets与传统方法的性能和精度对比

上表为在不同的数据集下使用MeshGraphNets和使用传统方式的性能和精度对比。可以看到，对于不同的数据集，性能均有几百倍的提升，但随着仿真步数增多，产生了误差累积和结果漂移，精度也随之下降，这个问题在DeepMind 2022年ICLR的论文《Predicting Physics in Mesh-reduced Space with Temporal Attention》[2]进行了解决，具体方案将在后续的博客中进行分析。

MeshGraphNets在训练分布之外的泛化性也表现得很好，包括底层系统参数、网格形状和网格大小，这是因为在图上使用相对编码的架构已被证明非常有利于泛化。在翼型数据集，验证了模型在更陡的角度（-35°——35° VS 训练集中的-25°——25°）和更高的马赫数（0.70.9 VS训练集中的0.20.7）下，预测结果仍然是合理的且RMSE仅仅从11.5（训练）升到了12.4（更陡的角度）和13.1（更高的马赫数）。同样，在FLAGDYNAMIC变体数据集上训练了一个模型，风速和风向在轨迹之间进行变化，但在每个轨迹内保持不变；在推理时，可以自由的改变风速和风向。这表明这个模型学习的物理定律可以推断出未经训练的参数范围。

四． 总结

MeshGraphNets是一种通用的基于网格的方法，可以准确有效地对广泛的物理系统进行建模，具有良好的泛化能力，并且可以在推理时进行扩展。这种方法相对传统的求解器可实现更高效的模拟，并且由于它是可微分的，对于设计优化或最优控制任务很有用。MeshGraphNets的出现将促进研究学者们对于结构网格和非结构网格的流场快速预测提出更优和更新的网络架构，为进一步降低累积误差、提升泛化性奠定了良好的基础。

参考文献

[1] Pfaff T, Fortunato M, Sanchez-Gonzalez A, et al. Learning mesh-based simulation with graph networks[J]. arXiv preprint arXiv:2010.03409, 2020. https://arxiv.org/abs/2010.03409

[2] Han X, Gao H, Pfaff T, et al. Predicting physics in mesh-reduced space with temporal attention[J]. arXiv preprint arXiv:2201.09113, 2022. https://arxiv.org/abs/2201.0911