Accelerating wave simulations with neural dispersion correctors

TL;DR

The paper tackles the high computational cost of accurate 3-D wave simulations by addressing dispersion errors that arise on coarse grids. It introduces a neural dispersion corrector based on Fourier neural operators that learns to map low-accuracy wavefields to high-accuracy counterparts, using a small training dataset. The approach is supported by a theoretical justification showing dispersion errors depend weakly on medium properties, and it employs a memory-efficient data representation via FFT and DCT. Empirically, the method achieves about a 16x speed-up on 3-D elastic-wave problems while generalising to strongly heterogeneous media, indicating practical impact for seismic imaging and related applications.

Abstract

We present a Fourier neural operator network, designed to correct dispersion errors in numerical wave simulations. The neural dispersion corrector enables the replacement of a computationally expensive high-accuracy simulation by a less expensive low-accuracy simulation. In contrast to neural network surrogates that fully replace a wave equation, the neural dispersion corrector has only a weak dependence on the distribution of model parameters, such as wave speeds. Consequently, the network can be trained with a significantly smaller dataset, while still generalising to unseen input parameters. Following a description of the network architecture and training, we provide examples for the 3-D elastic wave equation. After training with merely 1$\,$000 examples on one GPU, the neural corrector achieves a speed-up of 16$\times$ compared to a reference spectral-element simulation and a generalisation to a broad range of strongly heterogeneous wave speed distributions.

Figures (9)

• Figure 1: Comparison of numerical wave speed $c_k$ (a) and fractional wave speed error $c_k/c$ (b) as a function of frequency, normalised with respect to the Nyquist frequency $\omega_\text{Nyquist}$. The discretisation is fixed to $dx$=100 m and $dt$=0.25. Different curves correspond to different exact wave speeds $c$. In contrast to $c_k$ itself, the fractional error only shows a weak dependence on $c$, despite its large variation from 100 to 350 m$/$s. The dependence is particularly weak in the regime where most second-order finite-difference simulations operate, i.e., at $\omega/\omega_{Nyquist}\leq 0.2$, which corresponds to $\geq$10 grid points per wavelength Moczo_2014Igel_2016.
• Figure 2: Schematic illustration of the FNO architecture used in the following numerical examples. The input consists of low-accuracy three-components wavefields in the frequency domain and static variables that describe elastic medium and source properties. These are mapped into a higher-dimensional representation through a lifting layer. Each FB consists of Fourier transforms followed by a pointwise multi-layer perceptron (MPL). After passing through four FBs, the final projection layer maps the latent representation back into six output channels, corresponding to the real and imaginary parts of the three-component frequency-domain residual between the high-accuracy wavefield $u_h$ and its low-accuracy counterpart $u_l$.
• Figure 3: Representative selection of eight of the 100 $v_\text{s}$ distributions that we used to compute $v_\text{p}$ and $\rho$, as well as the high- and low-accuracy wavefields that compose the training dataset.
• Figure 4: Comparison of low- and high-accuracy spectral-element simulations. Panel (a) shows the near-surface $v_\text{s}$ distribution used for this example. Panel (b) illustrates a close-up of the discretisation. Doubling the elements per wavelength (epw) from one to two, the number of elements increases from 1$\,$859 to 13$\,$125. Finally, panel (c) compares representative vertical-component time series from the low-accuracy wavefield $u_l$ (black) and the high-accuracy wavefield $u_h$ (blue). The goal of the neural dispersion corrector is to map $u_l$ into a close approximation of $u_h$.
• Figure 5: Summary of training and prediction performance. (a) Training and validation losses as a function of Adam optimisation epoch. (b) Mean prediction error for the testing dataset as a function of time, and (c) as a function of frequency.