Implementing the discontinuous-Galerkin finite element method using graph neural networks with application to diffusion equations

Machine learning (ML) has benefited from both software and hardware advancements, leading to increasing interest in capitalising on ML throughout academia and industry. There have been efforts in the scientific computing community to leverage this development via implementing conventional partial differential equation (PDE) solvers with machine learning packages, most of which rely on structured spatial discretisation and fast convolution algorithms. However, unstructured meshes are favoured in problems with complex geometries. To bridge this gap, we propose to implement the unstructured Finite Element Method (FEM) on simplicial meshes with graph neural networks. This paper is the first to implement an unstructured mesh FEM solver using graph neural networks. All compute-intensive algorithms in the solver are represented with either convolutional or graph neural networks. Specifically, the FEM solver uses a discontinuous Galerkin formulation with an interior penalty method for spatial discretisation and a multigrid preconditioned Krylov solver as the linear solver. The multigrid method has been designed to suit the data structure within the ML package and adopts the commonly used U-Net architecture for this. A hierarchy of coarsened meshes is generated from p-multigrid and algebraic node agglomeration guided by either a space-filling curve or a smoothed aggregation algorithm. The solver is verified and assessed for solving the diffusion problems. The solver shows the theoretical convergence rate of (p+1) order. Compared with a highly optimised implementation, the solver running on GPU can reach promising throughput in terms of matrix operator evaluation at 6.8 MDOF/s. The method can easily extend to other PDEs and computing platforms beyond CPU and GPU.