WaveBench: Benchmarking Data-Driven Solvers for Linear Wave Propagation PDEs

Abstract

Wave-based imaging techniques play a critical role in diverse scientific, medical, and industrial endeavors, from discovering hidden structures beneath the Earth's surface to ultrasound diagnostics. They rely on accurate solutions to the forward and inverse problems for partial differential equations (PDEs) that govern wave propagation. Surrogate PDE solvers based on machine learning emerged as an effective approach to computing the solutions more efficiently than via classical numerical schemes. However, existing datasets for PDE surrogates offer only limited coverage of the wave propagation phenomenon. In this paper, we present WaveBench, a comprehensive collection of benchmark datasets for wave propagation PDEs. WaveBench (1) contains 24 datasets that cover a wide range of forward and inverse problems for time-harmonic and time-varying wave phenomena; (2) includes a user-friendly PyTorch environment for comparing learning-based methods; and (3) comprises reference performance and model checkpoints of popular PDE surrogates such as Fourier neural operators and U-Nets. Our evaluation on WaveBench demonstrates the impressive performance of PDE surrogates on in-distribution samples, while simultaneously unveiling their limitations on out-of-distribution samples, indicating room for future improvements. We anticipate that WaveBench will stimulate the development of accurate wave-based imaging techniques through machine learning.