JAXIGA - 一个支持GPU和可微分的等几何分析软件框架

中文摘要 AI 辅助

在本文中，我们介绍了JAXIGA，一个基于JAX的开源等几何分析库。该库将网格的构建与用于编译和微分的计算分开。这使得几何评估、组装和所支持的求解器可以组合在即时编译的程序中，并可以计算一阶反向模式导数。我们考虑了从单一逐点能量密度获得的三种数值方法。Galerkin和能量最小化方法使用相同的离散势能并具有相同的驻点，而配点法强制执行相应的强形式。收敛解关于参数的灵敏度通过求解一个伴随问题获得，其精度取决于原问题和伴随问题的容差。我们还描述了用于变分方法局部细化的分层提取，以及用于四阶问题的光滑样条空间的使用。数值结果与解析解以及独立的FEniCSx和JAX-FEM计算进行了比较。对于报告的GPU基准测试，具有36,448个未知数，组装速度比主机CPU快二十四倍。当共轭梯度求解器也在GPU上运行时，总求解时间从31.8秒减少到1.38秒，结果与直接解一致，误差为$5\times10^{-11}$。这些结果表明，求解器的选择和编译的复用对于整体计算成本的降低是重要的。

英文摘要

In this paper, we present JAXIGA, an open-source library for isogeometric analysis based on JAX. The library separates the construction of the mesh from the computations which are traced for compilation and differentiation. This allows geometry evaluation, assembly and the supported solvers to be combined in a just-in-time compiled program, for which first-order reverse-mode derivatives can be computed. We consider three numerical methods obtained from a single pointwise energy density. The Galerkin and energy-minimization methods use the same discrete potential and have the same stationary points, while collocation enforces the corresponding strong form. The sensitivities of a converged solution are obtained by solving one adjoint problem, and their accuracy depends on the tolerances of both the primal and adjoint problems. We also describe hierarchical extraction for local refinement of the variational methods and the use of smooth spline spaces for fourth-order problems. The numerical results are compared with analytical solutions and independent FEniCSx and JAX-FEM computations. For the reported GPU benchmark with 36,448 unknowns, assembly is twenty-four times faster than on the host CPU. When the conjugate-gradient solver is also run on the GPU, the total solution time is reduced from 31.8 to 1.38 seconds, and the result agrees with the direct solution to $5\times10^{-11}$. These results show that the choice of solver and the reuse of compilation are important for obtaining an overall reduction in computational cost.

发表机构

• Carnegie Mellon University（卡内基梅隆大学）

机构由 AI 辅助整理，请以论文原文为准。