用高效张量算法加速微分方程求解与学习,大幅降低计算开销。
Learning, Solving and Optimizing PDEs with TensorGalerkin: an efficient high-performance Galerkin assembly algorithm
- 通过张量化映射与稀疏矩阵归约,实现微分方程刚度矩阵快速组装。
- 在非结构网格上对2D/3D椭圆、抛物、双曲方程均实现显著提速与高精度。
- 适合需要高效求解、优化或物理约束学习的科学计算研究者。
我们提出一个统一的算法框架,用于具有变分结构的偏微分方程(PDEs)的数值求解、约束优化和物理信息学习。该框架基于伽辽金离散化,其高效性源于一种新型高度优化且兼容GPU的TensorGalerkin框架,用于线性系统组装(刚度矩阵和载荷向量)。TensorGalerkin通过在Python层的映射阶段对单元级操作进行张量化,再通过稀疏矩阵乘法执行全局归约,完成网格诱导稀疏图上的消息传递。映射与归约阶段在PyTorch自动微分中协同设计,使得组装图仅含$O(1)$个节点,不受单元数和局部自由度数量影响。我们通过多个基准测试验证了该$O(1)$图特性:作为高效数值求解器、端到端可微的PDE约束优化框架、以及物理信息算子学习算法。在非结构网格上的二维和三维椭圆型、抛物型、双曲型方程上,本方法在所有下游应用中均显著优于多种基线方法,展现出更高的计算效率和精度。
原文摘要 · Abstract (English)
We present a unified algorithmic framework for the numerical solution, constrained optimization, and physics-informed learning of PDEs with a variational structure. Our framework is based on a Galerkin discretization of the underlying variational forms, and its high efficiency stems from a novel highly-optimized and GPU-compliant TensorGalerkin framework for linear system assembly (stiffness matrices and load vectors). TensorGalerkin operates by tensorizing element-wise operations within a Python-level Map stage and then performs global reduction with a sparse matrix multiplication that performs message passing on the mesh-induced sparsity graph. The Map and Reduce stages are co-designed inside PyTorch's autograd so that the assembly graph contains $O(1)$ nodes regardless of how the number of elements and local DoFs scale. We validate this $O(1)$-graph property by deploying TensorGalerkin downstream as i) a highly-efficient numerical PDEs solver, ii) an end-to-end differentiable framework for PDE-constrained optimization, and iii) a physics-informed operator learning algorithm for PDEs. With multiple benchmarks, including 2D and 3D elliptic, parabolic, and hyperbolic PDEs on unstructured meshes, we demonstrate that the proposed framework provides significant computational efficiency and accuracy gains over a variety of baselines in all the targeted downstream applications.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。