djps/jwave

A JAX-based research framework for differentiable and parallelizable acoustic simulations, on CPU, GPUs and TPUs. Uses fork of jaxdf for jax=0.2.29 and jaxlib=0.1.77

j-Wave: Differentiable acoustic simulations in JAX

Install | UCL Tutorials | UCL 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, install jaxdf and then jwave using pip

pip install git+https://github.com/djps/jwave.git

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.0)
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},
}

Useful Papers

• Optimizing a DIscrete Loss (ODIL) to solve forward and inverse problems for partial differential equations using machine learning tools

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
4. jaxwell: for Maxwell's equation.