JAX-FLUIDS：可压缩两相流的完全可微高阶计算流体动力学求解器

原文来自微信公众号“编程语言Lab”：论文精读 | JAX-FLUIDS：可压缩两相流的完全可微高阶计算流体动力学求解器
搜索关注“编程语言Lab”公众号（HW-PLLab）获取更多技术内容！
欢迎加入 编程语言社区 SIG-可微编程 参与交流讨论（加入方式：添加小助手微信 pl_lab_001，备注“加入SIG-可微编程”）。

翻译 | 瞿家港 (哈工大流体力学博士研究生)，沈明 (编程语言实验室研究员)
校对 | 海丽娟
编辑 | Hana

华为编程语言实验室从编程语言的角度面向 AI + 科学计算 进行探索研究，致力于提供更适合 AI 与科学计算融合应用开发的编程技术方案。来自慕尼黑工业大学 Adams 教授的这篇论文是最新在 CFD+ML 方向看到的比较有意思的研究工作，特此分享出来以飨读者。

论文信息

Jax-fluids: A fully-differentiable high-order computational fluid dynamics solver for compressible two-phase flows. Deniz A. Bezgin, Aaron B. Buhendwa, and Nikolaus A. Adams. 3 2022.

论文地址：https://arxiv.org/abs/2203.13760

Github：https://github.com/tumaer/JAXFLUIDS

译者注

传统的计算流体力学 (后简称 CFD)，其求解过程完全采用数值微分进行计算，参数空间完全不可微。基于机器学习的计算流体力学应用 (后简称 ML-CFD)，常常构建 ML 模型来代替 CFD 的部分功能/模块，从而得到更优的计算结果，这种情况下的 ML 模型则是可微的。

多数 ML-CFD 的应用是采用离线训练的，即，先使用传统 CFD 求解器得到训练数据，再将之导入到 ML 框架中进行训练，最后将训练好的模型耦合到传统的 CFD 求解器中。这样的流程十分繁琐，且属于完全数据驱动，与偏微分方程 (后简称 PDE) 的动力学性质无关。

如果能构建这样一个计算框架，使得 CFD 求解过程和 ML 的训练/推理过程完全融合，即能在训练过程中得到来自 CFD 求解实时的反馈，可以更有效地学习到符合该 PDE 动力学特征的模型。而且，这将使得 CFD 和 ML 模型的开发更加方便和快捷。

JAX ¹ 为我们实现这样的构想提供了方便的平台，能方便地编写 CFD 求解器，也能在 GPU 平台并行，还能提供自动微分功能，从而实现 ML 模型的训练。

相关推荐

SIG-可微编程

能微分会加速的 NumPy —— JAX

浅谈 eDSL 在科学计算和数据分析领域的发展趋势

从自动微分到可微编程语言设计 (一）（二）（三）

# 摘要 #

背景

物理系统有偏微分方程 (PDEs) 控制，其中 Navier-Stokes 方程 ² 为描述流体流动的 PDEs，代表着具有复杂时空相互作用的非线性物理系统。PDEs 通常使用 数值方法 (数值离散-迭代求解) 进行求解，近些年机器学习 (ML) 方法为求解 PDEs 提供了一个新途径，ML 在计算流体力学 (CFD) 领域中的应用也越来越多。

动机

当前，没有一个通用的 ML-CFD 软件包可以提供以下功能：

• 提供最新的求解 CFD 的数值方法
• 提供 ML 和 CFD 之间无缝衔接
• 自动微分 (AD) 功能 (AD 提供的梯度信息能够优化 CFD 模型，因此 AD 对 ML-CFD 的研究至关重要)

提出

JAX-FLUIDS，可压缩两相流的完全可微高阶计算流体动力学 Python 求解器，可以：

• 模拟三维湍流、可压缩流动、两相流等复杂流体问题；
• 完全采用 JAX 编写，因此可以将现有的 ML 模型耦合到 CFD 求解器中；
• 支持端到端的优化，ML 模型可以使用 CFD 求解器求解过程中的梯度信息，因此 ML 模型不仅可以使用 PDE 包含的信息，还可以使用所应用的数值方法的信息。

译者注：端到端是一种解决问题的思路，与之对应的是多步骤解决问题，也就是将一个问题拆分为多个步骤分步解决，而端到端是由输入端的数据直接得到输出端的结果。

# 引言 #

经典的 CFD 方法和 ML 方法的融合需要强大的新方法达成以下目标：

• 允许数据驱动模型与 CFD 方法的无缝集成
• 整个算法 (CFD&ML) 实现端到端的自动微分

本文讨论了混合 ML 加速的 CFD 求解器所面临的挑战，并强调了如 JAX-FLUIDS 这样的新软件架构如何促进 ML-CFD 的研究。

上世纪开始，CPU 快速发展为 CFD 的发展奠定了基础，CFD 逐渐发展成为一个独立的科研研究领域。近些年来机器学习为物理科学的发展注入了新的活力，强大的自动微分框架如 TensorFlow、PyTorch 和 JAX 推动了 ML 方法在自然科学和工程科学中的发展。例如采用 ML 方法从数据发现 PDE、物理信息神经网络 (PINNs) 求解反问题等。

流体力学是一个数据丰富、计算密集的科学，这是由于 Navier-Stokes 方程的高度复杂的时空非线性本质决定的，例如湍流表现出强间歇性和非高斯性的混沌行为特征。机器学习提供了一系列新的数据驱动方法来解决流体力学中一些长期存在的问题。

用于科学计算的 ML 方法能以不同的依据进行分类：

• 按模型和训练中包含的物理先验知识的水平进行分类
  完全数据驱动模型的优点在于：实现简单、推理效率高；缺点在于：收敛性、泛化性和稳定性等不能保证，难以实施物理约束等。而经典数值方法则能保证物理约束，如今结合 ML 和经典数值方法的研究越来越多。
• 在线训练/离线训练
  目前大多 ML 模型都是离线优化的，既在 ML 模型训练完成之后，再安插到传统 CFD 软件中，用于评估下游的计算任务。

尽管离线训练相对容易，但这种方法仍存在一些缺点：

• 这些模型在训练和测试时面对的数据分布不匹配；
• 它们通常不会直接学习 PDE 的动力学先验知识；
• 流体力学求解器通常非常复杂，使用 Fortran 或 C++ 等编程语言编写，并针对 CPU 计算进行了大量优化，这与 ML 研究的一般过程相反，ML 模型通常使用 Python 进行训练，并针对 GPU 的使用进行优化。将这些 ML 模型插入到现有的 CFD 软件框架中会比较麻烦。

为了解决这些问题，研究者们开始研究直接使用 Python 编写可微分的 CFD 求解器框架，从而可以进行端到端的训练 ML 模型。ML 模型在端到端的训练过程中，可以在训练过程中感受到动态的 PDE 求解过程，并实时地看到自己的输出。目前已有的可微分科学计算框架有：

• PhiFlow ³，TUM，兼容 JAX、PyTorch、TensorFlow、Numpy，为优化和机器学习应用程序构建的开源仿真工具包。
• JAX-CFD ⁴，Google，CFD 求解器 (尚不完善，如仅支持周期性边界)，支持：
  • 空间离散支持有限体积、有限差分和谱方法
  • 时间离散：一阶格式
  • 压力求解器：共轭梯度、实 FFTs 进行快速对角化
  • 仅支持周期性边界
  • 对流：二阶 Van-Leer 格式
  • Smagorinsky 涡粘性模型
• JAX-MD ⁵，Google，端到端的分子动力学求解器。

机器学习在计算流体力学以及更广泛的计算物理领域的稳步崛起和成功需要新一代算法，这些算法允许：

• 高级编程语言的快速原型设计
• 在 CPU、GPU 和 TPU 上运行
• 将机器学习模型无缝地集成到求解器框架中
• 对数据驱动模型进行端到端优化的完全可微算法

# 物理模型 #

译者注：本节介绍了描述流体运动的控制方程 (动量方程、连续性方程、能量方程)，即描述无粘流体运动的 Euler 方程和粘性流体运动的 Navier-Stokes 方程。在 CFD 的编程范式中，尤其在有限体积法的计算程序中，习惯将微分方程写成具备统一形式的通量表示形式，如下文所介绍。

本文将关注 无粘流体的可压缩 Euler 方程，和 粘性流体的可压缩 Navier-Stokes 方程。

定义以下基本变量：处于位置 $\mathbf{X}=[x,y,z{]}^T=[x_1,x_2,x_4{]}^T$ 及时间 $t$ 的流体状态采用矢量 $\mathbf{W}=[\rho ,u,v,w,p{]}^T$ 表示。这里，

• $\rho$ 为密度
• $\mathbf{u}=[u,v,w{]}^T=[u_1,u_2,u_3{]}^T$ 为速度矢量
• $p$ 为压力

由于本文关注 可压缩流体，因此，采用 守恒变量 描述问题更加合理，即，

$$\mathbf{U}=[\rho ,\rho u,\rho v,\rho w,E{]}^T$$

其中，

• $\rho \mathbf{u}=[\rho u,\rho v,\rho w{]}^T$ 是三个空间维度上的动量
• $E=\rho e+\frac{1}{2}\rho \mathbf{u}\cdot \mathbf{u}$ 是单位体积的总能量
• $e$ 是单位质量下的内能

将可压缩 Euler 方程写成 $\mathbf{U}$ 的微分方程，

$$\frac{\partial \mathbf{U}}{\partial t}+\frac{\partial \mathbf{F}(U)}{\partial x}+\frac{\partial \mathbf{G}(U)}{\partial y}+\frac{\partial \mathbf{H}(U)}{\partial z}=0$$

其中通量 $\bm F,\bm G,\bm H $ 定义为，

$$\mathbf{F}=\left(\begin{array}{c}\rho u\\ \rho u^2+p\\ \rho uv\\ \rho uw\\ u(E+p)\end{array}\right),\mathbf{G}=\left(\begin{array}{c}\rho v\\ \rho vu\\ \rho v^2+p\\ \rho vw\\ v(E+p)\end{array}\right),\mathbf{H}=\left(\begin{array}{c}\rho w\\ \rho wu\\ \rho wv\\ \rho w^2+p\\ w(E+p)\end{array}\right)$$

可压缩 Navier-Stokes 方程也写成通量的形式：

$$\frac{\partial \mathbf{U}}{\partial t}+\frac{\partial \mathbf{F}(U)}{\partial x}+\frac{\partial \mathbf{G}(U)}{\partial y}+\frac{\partial \mathbf{H}(U)}{\partial z}=\frac{\partial {\mathbf{F}}^{\mathbf{d}}(U)}{\partial x}+\frac{\partial {\mathbf{G}}^d(U)}{\partial y}+\frac{\partial {\mathbf{H}}^d(U)}{\partial z}+\mathbf{S}(U)$$

其中耗散通量 ${\mathbf{F}}^d,{\mathbf{G}}^d,{\mathbf{H}}^d$ 定义如下：

F = ( 0 τ 11 τ 12 τ 13 ∑ i u i τ 1 i − q 1 ) , G = ( 0 τ 21 τ 22 τ 23 ∑ i u i τ 2 i − q 2 ) , H = ( 0 τ 31 τ 32 τ 33 ∑ i u i τ 3 i − q 3 ) \bm F = \begin{pmatrix} 0 \\ \tau^{11} \\ \tau^{12} \\ \tau^{13} \\ \sum_i u_i\tau^{1i}-q_1 \end{pmatrix}, \bm G = \begin{pmatrix} 0 \\ \tau^{21} \\ \tau^{22} \\ \tau^{23} \\ \sum_i u_i\tau^{2i}-q_2 \end{pmatrix}, \bm H = \begin{pmatrix} 0 \\ \tau^{31} \\ \tau^{32} \\ \tau^{33} \\ \sum_i u_i\tau^{3i}-q_3 \end{pmatrix}