Learning, Solving and Optimizing PDEs with TensorGalerkin: an efficient high-performance Galerkin assembly algorithm
This work addresses the computational bottleneck in PDE analysis for researchers and engineers in scientific computing, offering a unified framework for multiple applications, though it is incremental as it builds on existing Galerkin methods with optimizations.
The authors tackled the problem of efficiently solving, optimizing, and learning PDEs with variational structures by developing TensorGalerkin, a high-performance Galerkin assembly algorithm, which demonstrated significant computational efficiency and accuracy gains over baselines in benchmarks including 2D and 3D PDEs on unstructured meshes.
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. It can be seamlessly employed 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.