Multigrid-Augmented Deep Learning Preconditioners for the Helmholtz Equation

Abstract

We present a data-driven approach to iteratively solve the discrete heterogeneous Helmholtz equation at high wavenumbers. We combine multigrid ingredients with convolutional neural networks (CNNs) to form a preconditioner which is applied within a Krylov solver. Two types of preconditioners are proposed 1) U-Net as a coarse grid solver, and 2) U-Net as a deflation operator with shifted Laplacian V-cycles. The resulting CNN preconditioner can generalize over residuals and a relatively general set of wave slowness models. On top of that, we offer an encoder-solver framework where an ``encoder'' network generalizes over the medium and sends context vectors to another ``solver'' network, which generalizes over the right-hand-sides. We show that this option is more efficient than the stand-alone variant. Lastly, we suggest a mini-retraining procedure, to improve the solver after the model is known. We demonstrate the efficiency and generalization abilities of our approach on a variety of 2D problems.