A Neural Multigrid Solver for Helmholtz Equations with High Wavenumber and Heterogeneous Media

In this paper, we propose a deep learning-enhanced multigrid solver for high-frequency and heterogeneous Helmholtz equations. By applying spectral analysis, we categorize the iteration error into characteristic and non-characteristic components. We eliminate the non-characteristic components by a multigrid wave cycle, which employs carefully selected smoothers on each grid. We diminish the characteristic components by a learned phase function and the approximate solution of an advection-diffusion-reaction (ADR) equation, which is solved using another multigrid V-cycle on a coarser scale, referred to as the ADR cycle. The resulting solver, termed Wave-ADR-NS, enables the handling of error components with varying frequencies and overcomes constraints on the number of grid points per wavelength on coarse grids. Furthermore, we provide an efficient implementation using differentiable programming, making Wave-ADR-NS an end-to-end Helmholtz solver that incorporates parameters learned through a semi-supervised training. Wave-ADR-NS demonstrates robust generalization capabilities for both in-distribution and out-of-distribution velocity fields of varying difficulty. Comparative experiments with other multigrid methods validate its superior performance in solving heterogeneous 2D Helmholtz equations with wavenumbers exceeding 2000.