j-Wave: An open-source differentiable wave simulator

TL;DR

j-Wave is an open-source, modular, and differentiable acoustic simulation tool compatible with popular machine learning frameworks, enabling advanced scientific computing and machine learning applications in acoustics.

Contribution

It introduces a flexible, differentiable acoustic simulator that integrates with JAX and TensorFlow, supporting automatic differentiation for scientific and machine learning tasks.

Findings

Simulation accuracy comparable to k-Wave and other acoustic software

Supports automatic differentiation for optimization tasks

Modular design allows easy customization and reuse

Abstract

We present an open-source differentiable acoustic simulator, j-Wave, which can solve time-varying and time-harmonic acoustic problems. It supports automatic differentiation, which is a program transformation technique that has many applications, especially in machine learning and scientific computing. j-Wave is composed of modular components that can be easily customized and reused. At the same time, it is compatible with some of the most popular machine learning libraries, such as JAX and TensorFlow. The accuracy of the simulation results for known configurations is evaluated against the widely used k-Wave toolbox and a cohort of acoustic simulation software. j-Wave is available from https://github.com/ucl-bug/jwave.