[论文解读] Solving the Discretised Boltzmann Transport Equations using Neural Networks: Applications in Neutron Transport

T. R. F. Phillips, C. E. Heaney+3 more|arXiv (Cornell University)|Jan 24, 2023

Model Reduction and Neural Networks被引用 7

一句话总结

本文提出了一种新型的AI原生方法，用于求解中子输运问题中的4D离散化Boltzmann输运方程，通过利用AI库实现高效的GPU和AI硬件加速。该方法结合了空间-角度多重网格法、一种新型卷积有限元法（ConvFEM）以及非线性Galerkin-Petrov稳定化技术，实现了稳定且高阶精度的解，其性能与迎风格式相当，但分辨率更高，$k_{\text{eff}}$值也更优。

ABSTRACT

In this paper we solve the Boltzmann transport equation using AI libraries. The reason why this is attractive is because it enables one to use the highly optimised software within AI libraries, enabling one to run on different computer architectures and enables one to tap into the vast quantity of community based software that has been developed for AI and ML applications e.g. mixed arithmetic precision or model parallelism. Here we take the first steps towards developing this approach for the Boltzmann transport equation and develop the necessary methods in order to do that effectively. This includes: 1) A space-angle multigrid solution method that can extract the level of parallelism necessary to run efficiently on GPUs or new AI computers. 2) A new Convolutional Finite Element Method (ConvFEM) that greatly simplifies the implementation of high order finite elements (quadratic to quintic, say). 3) A new non-linear Petrov-Galerkin method that introduces dissipation anisotropically.

研究动机与目标

• 开发一种可扩展、硬件优化的高维Boltzmann输运方程求解方法，利用AI库实现。
• 通过利用AI软件栈，实现在GPU和新兴AI加速器上的高效并行计算。
• 通过一种新型Petrov-Galerkin方法，解决高阶有限元离散化输运方程中的数值不稳定性问题。
• 通过一种卷积、模板不变的公式（ConvFEM），简化高阶有限元的实现。
• 展示使用深度学习架构（如U-Net）求解中子输运问题多重网格解的可行性。

提出的方法

• 采用U-Net神经网络架构实现4D锯齿形多重网格循环，天然适配GPU并行计算。
• 在4个维度（2D空间，2D角度）上以因子2进行结构化网格粗化，使用标准的 prolongation 和 restriction 算子。
• 提出一种新型卷积有限元法（ConvFEM），采用固定尺寸滤波器（如5×5、7×7），实现一致且高效的高阶（二次至五次）有限元离散化。
• 应用非线性Petrov-Galerkin方法引入各向异性耗散，稳定由一阶输运方程产生的非对称、正定系统。
• 在所有网格层级上使用迎风格式，进一步稳定多重网格循环中的残差校正过程。
• 整个求解过程在AI库内实现，支持混合精度计算、模型并行化和硬件可移植性。

b Octahedron in a sphere, open

实验结果

研究问题

• RQ1 AI库能否有效用于求解中子输运问题中的高维Boltzmann输运方程？
• RQ2 基于U-Net的多重网格架构能否在GPU上高效求解4D空间-角度输运问题？
• RQ3 ConvFEM能否实现高阶有限元方法且滤波器变化极小，从而简化基于AI的实现？
• RQ4 所提出的非线性Petrov-Galerkin方法是否能稳定高阶输运解，而标准Galerkin方法会失效？
• RQ5 该AI原生框架能否通过图神经网络扩展至3D问题和非结构化网格？

主要发现

• 所提出的ConvFEM方法采用固定尺寸滤波器，实现二次至五次高阶有限元，与传统FEM相比，在AI框架中实现显著简化。
• 非线性Petrov-Galerkin方法成功稳定了离散化Boltzmann输运方程的求解，而使用中心或Bubnov-Galerkin格式则无法收敛。
• 采用二次ConvFEM滤波器的Petrov-Galerkin方法解的$k_{\text{eff}}$值高于迎风格式，表明精度更高且人工扩散更小。
• Petrov-Galerkin方法得到的标量通量解在定性上与迎风解相似，但分辨率更高，峰值通量略低，表明耗散得到有效控制。
• 该方法在GPU和CPU上均实现稳定收敛，借助AI硬件（如Cerebras 2）具备扩展至百亿亿级系统的潜力。
• 通过与其他求解流体或热传导方程的神经网络耦合，该框架可扩展至多物理场问题。

c $4\times 4$ grid on a triangle face.