JAX-CFD：利用机器学习加速CFD计算

谷歌JAX-CFD论文：《Machine learning–accelerated computational fluid dynamics》

摘要：CFD在许多物理现象（如天气、气候、空气动力学和等离子体物理学）模拟中起着至关重要的作用。Navier-Stokes 方程能够很好地描述流体流动，但受到解析最小时空特征的计算成本的限制，大规模地求解这些方程仍然令人望而生畏，这导致需要在精度和可处理性之间进行不利的权衡。这里我们使用端到端深度学习来改进计算流体动力学中的近似以对二维湍流进行模拟。对于湍流的直接数值模拟和大涡模拟，我们的结果与基准求解器一样精确，每个空间维度的分辨率提高了 8 到 10 倍，计算速度提高了 40 到 80 倍。我们的方法在长时间的模拟过程中保持稳定，并推广到受训流程之外的强制函数和雷诺数，这与黑盒机器学习方法形成了鲜明对比。我们的方法举例说明了科学计算如何利用机器学习和硬件加速器在不牺牲准确性或泛化性的情况下改进模拟。

模拟由非线性偏微分方程（PDE）描述的复杂物理系统是工程和物理科学的核心，其应用范围从天气和气候以及车辆或发动机的工程设计到野火和等离子体物理。尽管运动方程和基本物理定律之间存在直接联系，但不可能在这些重要问题所需的规模上直接进行数值模拟。这个基本问题几十年来一直阻碍着科学计算的进展，其原因是精确的模拟必须解析最小的时空尺度。

一个典型的例子，如湍流流动为天气、气候和空气动力学模拟的基础，最小的涡的尺寸是很小的，如弦长2米的飞机，最小长度尺度（Kolomogorov尺度）为  。计算流体动力学（CFD）的经典方法，如有限差分法、有限体积法、有限元法和谱分析法，只有在网格上的物理场平滑变化时才是准确的，因此网格必须解析最小的物理特征以保证收敛。对于湍流流动，求解最小流动特征的要求意味着计算成本与  成比例，其中  ，  和  为典型的速度和长度尺度，  为运动粘度。  增加10倍会导致计算成本增加1000倍。因此，对于气候和天气等，直接的数值模拟（DNS）是不可能的。

相反，传统的做法是使用平滑的Navier-Stokes方程，其允许更粗的计算网格，同时牺牲计算精度，如雷诺平均Navier-Stokes及大涡模拟（LES）。例如，目前最先进的LES，其网格大小为  到  百万，已被用于内燃机、燃气涡轮发动机和涡轮机械的设计。尽管过去二十年来LES取得了可喜的进展，但能准确模拟的内容仍有严重的局限性。这主要是由于LES对亚网格尺度（SGS）模型的一阶依赖性，特别是对于速率控制尺度未被解析的流动。

在这里，我们介绍了一种计算非线性偏微分方程解的精确时间演化的方法，同时使用比传统方法更粗的网格来获得相同的精度。这是一种数值求解器，它不对未解析的自由度进行平均值，而是使用离散方程，在未解析的网格上给出逐点精确解。我们使用机器学习（ML）来发现这些算法，通过用学习到的替代方法取代传统求解器中受分辨率损失影响最大的部分。如图1A所示，对于湍流的二维DNS模拟，我们的算法在保持精度的同时，在每个维度上使用了10倍的粗分辨率，相对于类似精度的先进数值方法，计算时间缩短了约80倍。该模型学习如何插值解的局部特征，因此可以准确地泛化到不同的流动条件，如不同的驱动因素，甚至不同的雷诺数（图1B）。我们还将该方法应用于湍流的高分辨率LES模拟，并显示出类似的性能提升，在Re = 100000的LES模拟中使用8×粗网格保持逐点精度，能实现约40倍的计算加速。

 

图1为我们的方法和结果概述。(A) 基准（直接模拟）和ML加速求解器的精度与计算成本的关系。X轴对应的是点上的精度，显示了模拟与真值高度相关的时间，Y轴为在一个张量处理单元（TPU）核心上进行一个模拟时间单位所需的计算时间。每个点都有相应的空间网格大小的注释；(B) 说明性的训练和验证例子，显示了模型的强大泛化泛化能力。(C) LI模型的单个时间步长结构，卷积神经网控制标准数值求解器的对流计算中的学习近似值。    和    分别为对流和速度分量。

”

最近有大量的工作使用ML来改善湍流模拟。其中一个主要系列的方法基于与高分辨率DNS的一致性，使用ML来拟合经典湍流模型的封闭性。虽然这些新模型可能比传统的湍流模型更精确，但并没有实现计算费用的降低。另一个主要的方向是使用"纯" ML，旨在使用基于深度神经网络的近似取代整个Navier-Stokes模拟。纯ML方法可以非常有效，其避免了传统方法稳定性所需的严格时间步长约束。然而这些模型不包括基本物理学，它们通常不能强制执行硬性约束，如动量守恒和不可压缩性等。虽然这些模型通常对来自训练分布的数据表现良好，但它们往往在泛化方面存在困难。例如，当遇到新的边界条件时，它们的表现更差。我们相信结合最佳ML和传统数值方法的"混合"方法更有前途。例如，ML可以代替或加速在一些模拟方法中使用的迭代求解而不降低精度。在这里，我们重点关注混合模型，它使用ML来校正廉价的、为充分解析的模拟中的误差。这些模型借鉴了粗粒度模拟的优势，并且由于减小了网格尺寸，计算速度通常比纯数值模拟快得多。

在这项工作中，我们设计了一些算法，通过使用性能更好的机器学习替代物取代受分辨率损失影响最大的部分，在更粗的网格上准确地求解方程。我们使用数据驱动的离散化将微分算子高精度地插值到粗网格上（图1C）。我们在一个标准的数值方法中训练模型，将基础偏微分方程作为可微分程序来求解，将神经网络和数值方法编写在一个支持反向模式自动微分的框架JAX中。这允许对整个算法进行端到端的梯度优化，类似于之前在密度泛函理论、分子动力学方面的工作。我们推导出的方法是特定于方程的，需要高分辨率的真实模拟作为训练数据。由于偏微分方程的动力学是局部的，因此可以在小的区域上进行高分辨率的模拟。这些模型在长时间模拟中能够保持稳定，并且具有良好的鲁棒性和可预测的泛化特性，其中在小区域上训练的模型能够在更大的区域上产生精确的模拟结构。与纯ML基准的比较表明，泛化来自于方法模型中固有的物理约束。

 

JAX-CFD是一个实验研究项目，旨在探索机器学习、自动微分和硬件加速器（GPU/TPU）在计算流体动力学中的潜力。其主要利用JAX实现。JAX是谷歌开发的深度学习框架，坊间传言JAX是Tensorflow的替代品。JAX-CFD网址：https://github.com/google/jax-cfd

”

（完毕）