Learning the elastic wave equation with Fourier neural operators

Tianze Zhang, Kristopher A. Innanen, Daniel O. Trad

Neural operators are extensions of neural networks which in supervised training learn how to map complex relationships, such as classes of PDE. Recent literature reports efforts to develop one type of these, the Fourier Neural Operator (FNO) such that it learns to create relatively general solutions to PDEs such as the Navier-Stokes equation. In this study, we seek what we believe is the first numerical instance of a Fourier Neural Operator (FNO) be trained to learn the elastic wave equation from a synthetic training data set. FNO attempts to find a manifold for elastic wave propagation. On that manifold, wave fields are represented in lower numbers of dimensions than those needed for standard solutions, and the calculations for wave propagation are correspondingly simpler. The FNO combines a linear transform, the Fourier transform, and a non-linear local activation, to produce a network with sufficient freedom to map from a general parameterization of a forward wave problem to its solution. Post-training, the FNO is observed to generate accurate elastic wave fields at approximately 10 times the speed of traditional finite difference methods on a CPU, and about a hundred times faster on a GPU.