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

TL;DR

This work introduces Wave-ADR-NS, a neural multigrid solver for high-frequency, heterogeneous Helmholtz equations. By separating errors into non-characteristic and characteristic components, it uses a wave cycle with learnable smoothing on fine and mid grids to remove the former, and an ADR cycle with a learned phase function $\tau$ to suppress the latter, all implemented in a differentiable, matrix-free framework trained in a semi-supervised manner. Key contributions include a spectral-analysis–informed wave cycle, an ADR-based characteristic-cycle, neural learning of smoothing parameters $\alpha$ and phase $\tau$, and comprehensive 2D experiments demonstrating superior efficiency and robust generalization to in-distribution and out-of-distribution velocity fields for wavenumbers up to $>2000$. The approach offers a scalable, end-to-end Helmholtz solver with strong transfer capabilities and potential applicability beyond 2D, including further development toward wavenumber-independent performance and 3D extensions.

Abstract

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.

Figures (18)

• Figure 1: Eigenvalues of the error propagation matrix in \ref{['eq:sjac']} for the damped Jacobi method on nested mesh grids.
• Figure 1: Examples of slowness models converted from the CIFAR-10 dataset with different resolutions
• Figure 2: Eigenvalues of the error propagation matrix of the Chebyshev semi-iteration method on nested mesh grids.
• Figure 2: (a) Example of a slowness model from CIFAR-10 dataset, (b) Solution of the heterogeneous Helmholtz equation with a point source when $\omega=40\pi$ and $N=256$, (c) The module of Fourier transform of the error after one wave cycle iteration, (d) Solution of the eikonal equation, (e) Gradient field of $\tau$, (f) The module of Fourier transform of $e^{-i\omega\tau}$.
• Figure 3: (a) Solution of Helmholtz equation in homogeneous media with a point source and absorbing boundary condition. (b) The module of the error in Fourier space between one wave cycle iterative solution and the reference solution.