tonyzzr / jwave

A JAX-based research framework for differentiable and parallelizable acoustic simulations, on CPU, GPUs and TPUs

j-Wave: Differentiable acoustic simulations in JAX

Install | Tutorials | Documentation | Changelog

j-Wave is a library of simulators for acoustic applications. Is heavily inspired by k-Wave (a big portion of j-Wave is a port of k-Wave in JAX), and its intented to be used as a collection of modular blocks that can be easily included into any machine learning pipeline.

Following the phylosophy of JAX, j-Wave is developed with the following principles in mind

1. To be differntiable
2. To be fast via jit compilation
3. Easy to run on GPUs
4. Easy to customize

Install

Follow the instructions to install Jax with CUDA support if you want to use your GPU.

Then, simply install jwave using pip

pip install git+https://github.com/ucl-bug/jwave.git

For more details, see the Linux install guide.

Because JAX has limited support on Windows, j-Wave can be run on Windows machines only using the Windows Subsystem for Linux. See the Install on Windows guide for more details.

Example

This example simulates an acoustic initial value problem, which is often used as a simple model for photoacoustic acquisitions:

from jax import jit

from jwave import FourierSeries
from jwave.acoustics.time_varying import simulate_wave_propagation
from jwave.geometry import Domain, Medium, TimeAxis
from jwave.utils import load_image_to_numpy

# Simulation parameters
N, dx = (256, 256), (0.1e-3, 0.1e-3)
domain = Domain(N, dx)
medium = Medium(domain=domain, sound_speed=1500.)
time_axis = TimeAxis.from_medium(medium, cfl=0.3, t_end=.8e-05)

# Initial pressure field
p0 = load_image_to_numpy("docs/assets/images/jwave.png", image_size=N)/255.
p0 = FourierSeries(p0, domain)

# Compile and run the simulation
@jit
def solver(medium, p0):
  return simulate_wave_propagation(medium, time_axis, p0=p0)

pressure = solver(medium, p0)

Citation

If you use jwave for your research, please consider citing it as:

@article{stanziola2022jwave,
    author={Stanziola, Antonio and Arridge, Simon R. and Cox, Ben T. and Treeby, Bradley E.},
    title = {j-Wave: An open-source differentiable wave simulator},
    publisher = {arXiv},
    year = {2022},
}

Related Projects

1. ADSeismic.jl: a finite difference acoustic simulator with support for AD and JIT compilation in Julia.
2. stride: a general optimisation framework for medical ultrasound tomography.
3. k-wave-python: A python interface to k-wave GPU accelerated binaries