MGCNN: a learnable multigrid solver for linear PDEs on structured grids

This paper presents a learnable solver tailored to solve discretized linear partial differential equations (PDEs). This solver requires only problem-specific training data, without using specialized expertise. Its development is anchored by three core principles: (1) a multilevel hierarchy to promote rapid convergence, (2) adherence to linearity concerning the right-hand side of equations, and (3) weights sharing across different levels to facilitate adaptability to various problem sizes. Built on these foundational principles, we introduce a network adept at solving PDEs discretized on structured grids, even when faced with heterogeneous coefficients. The cornerstone of our proposed solver is the convolutional neural network (CNN), chosen for its capacity to learn from structured data and its similar computation pattern as multigrid components. To evaluate its effectiveness, the solver was trained to solve convection-diffusion equations featuring heterogeneous diffusion coefficients. The solver exhibited swift convergence to high accuracy over a range of grid sizes, extending from $31 \times 31$ to $4095 \times 4095$. Remarkably, our method outperformed the classical Geometric Multigrid (GMG) solver, demonstrating a speedup of approximately 3 to 8 times. Furthermore, we explored the solver's generalizability to untrained coefficient distributions. The findings showed consistent reliability across various other coefficient distributions, revealing that when trained on a mixed coefficient distribution, the solver is nearly as effective in generalizing to all types of coefficient distributions.