j-Wave: An open-source differentiable wave simulator

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 .

[1]  Eran Treister,et al.  Multigrid-augmented deep learning preconditioners for the Helmholtz equation , 2022, SIAM Journal on Scientific Computing.

[2]  Òscar Calderón Agudo,et al.  Stride: A flexible software platform for high-performance ultrasound computed tomography , 2021, Comput. Methods Programs Biomed..

[3]  Gaurav S. Sukhatme,et al.  Probabilistic Inference of Simulation Parameters via Parallel Differentiable Simulation , 2021, 2022 International Conference on Robotics and Automation (ICRA).

[4]  Simon R. Arridge,et al.  A research framework for writing differentiable PDE discretizations in JAX , 2021, ArXiv.

[5]  F. Vicentini,et al.  mpi4jax: Zero-copy MPI communication of JAX arrays , 2021, J. Open Source Softw..

[6]  Marco Cuturi,et al.  Efficient and Modular Implicit Differentiation , 2021, NeurIPS.

[7]  Stephan Hoyer,et al.  Machine learning–accelerated computational fluid dynamics , 2021, Proceedings of the National Academy of Sciences.

[8]  Eric F Darve,et al.  Integrating Deep Neural Networks with Full-waveform Inversion: Reparametrization, Regularization, and Uncertainty Quantification , 2020, 2012.11149.

[9]  Jan Peters,et al.  Differentiable Physics Models for Real-world Offline Model-based Reinforcement Learning , 2020, 2021 IEEE International Conference on Robotics and Automation (ICRA).

[10]  Eric F Darve,et al.  A general approach to seismic inversion with automatic differentiation , 2020, Comput. Geosci..

[11]  E. D. Cubuk,et al.  JAX, M.D. A framework for differentiable physics , 2020, NeurIPS.

[12]  Adam R. Gerlach,et al.  The Koopman Expectation: An Operator Theoretic Method for Efficient Analysis and Optimization of Uncertain Hybrid Dynamical Systems , 2020, 2008.08737.

[13]  Jaime Fern'andez del R'io,et al.  Array programming with NumPy , 2020, Nature.

[14]  Ivan Yashchuk.  Bringing PDEs to JAX with forward and reverse modes automatic differentiation , 2020, ICLR 2020.

[15]  Vladlen Koltun,et al.  Learning to Control PDEs with Differentiable Physics , 2020, ICLR.

[16]  Ali Ramadhan,et al.  Universal Differential Equations for Scientific Machine Learning , 2020, ArXiv.

[17]  Gilles Louppe,et al.  The frontier of simulation-based inference , 2019, Proceedings of the National Academy of Sciences.

[18]  Jonathan Ragan-Kelley,et al.  DiffTaichi: Differentiable Programming for Physical Simulation , 2019, ICLR.

[19]  F. Herrmann,et al.  Neural network augmented wave-equation simulation , 2019, ArXiv.

[20]  Alan Edelman,et al.  A Differentiable Programming System to Bridge Machine Learning and Scientific Computing , 2019, ArXiv.

[21]  Paris Perdikaris,et al.  Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations , 2019, J. Comput. Phys..

[22]  Ming C. Lin,et al.  Differentiable Cloth Simulation for Inverse Problems , 2019, NeurIPS.

[23]  David Duvenaud,et al.  Neural Ordinary Differential Equations , 2018, NeurIPS.

[24]  Anuj Karpatne,et al.  Physics-guided Neural Networks (PGNN): An Application in Lake Temperature Modeling , 2017, ArXiv.

[25]  K. Hynynen,et al.  A viscoelastic model for the prediction of transcranial ultrasound propagation: application for the estimation of shear acoustic properties in the human skull , 2017, Physics in medicine and biology.

[26]  Alan Edelman,et al.  Julia: A Fresh Approach to Numerical Computing , 2014, SIAM Rev..

[27]  Mose Giordano,et al.  Uncertainty propagation with functionally correlated quantities , 2016, 1610.08716.

[28]  Michael Lange,et al.  Devito: Towards a Generic Finite Difference DSL Using Symbolic Python , 2016, 2016 6th Workshop on Python for High-Performance and Scientific Computing (PyHPC).

[29]  Lindsey J. Heagy,et al.  SimPEG: An open source framework for simulation and gradient based parameter estimation in geophysical applications , 2015, Comput. Geosci..

[30]  Simon Arridge,et al.  A fast boundary element method for the scattering analysis of high-intensity focused ultrasound. , 2015, The Journal of the Acoustical Society of America.

[31]  Liangguo Dong,et al.  Full waveform inversion method using envelope objective function without low frequency data , 2014 .

[32]  Tariq Alkhalifah,et al.  Scattering-angle based filtering of the waveform inversion gradients , 2014 .

[33]  D. Christensen,et al.  Ultrasound beam simulations in inhomogeneous tissue geometries using the hybrid angular spectrum method , 2012, IEEE Transactions on Ultrasonics, Ferroelectrics and Frequency Control.

[34]  Alistair P. Rendell,et al.  Modeling nonlinear ultrasound propagation in heterogeneous media with power law absorption using a k-space pseudospectral method. , 2012, The Journal of the Acoustical Society of America.

[35]  B T Cox,et al.  k-Wave: MATLAB toolbox for the simulation and reconstruction of photoacoustic wave fields. , 2010, Journal of biomedical optics.

[36]  Jean Virieux,et al.  An overview of full-waveform inversion in exploration geophysics , 2009 .

[37]  G. Trahey,et al.  A heterogeneous nonlinear attenuating full- wave model of ultrasound , 2009, IEEE Transactions on Ultrasonics, Ferroelectrics and Frequency Control.

[38]  Alfredo Bermúdez,et al.  An optimal perfectly matched layer with unbounded absorbing function for time-harmonic acoustic scattering problems , 2007, J. Comput. Phys..

[39]  B T Cox,et al.  Fast calculation of pulsed photoacoustic fields in fluids using k-space methods. , 2005, The Journal of the Acoustical Society of America.

[40]  T. D. Mast,et al.  A k-space method for coupled first-order acoustic propagation equations. , 2002, The Journal of the Acoustical Society of America.

[41]  C. Shin,et al.  An optimal 9-point, finite-difference, frequency-space, 2-D scalar wave extrapolator , 1996 .

[42]  Edward Bedrosian,et al.  The Analytic Signal Representation of Modulated Waveforms , 1962, Proceedings of the IRE.