jwave 0.1.2

Fast and differentiable acoustic simulations in JAX.

j-Wave: Differentiable acoustic simulations in JAX

Install | Tutorials | Documentation | Changelog

j-Wave is a library of simulators for acoustic applications. It is heavily inspired by k-Wave—a substantial portion of j-Wave is a port of k-Wave in JAX—and it's intended to be used as a collection of modular blocks that can be easily incorporated into any machine learning pipeline.

Embracing the philosophy of JAX, j-Wave is developed with the following principles in mind:

1. To be differentiable
2. To be efficient through jit compilation
3. To be easily run on GPUs
4. To be easily customizable

Install

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

Next, simply install jwave using pip:

pip install jwave

For more information, refer to the Linux installation guide.

Due to JAX's limited support on Windows, j-Wave can only be run on Windows machines using the Windows Subsystem for Linux. Please refer to the Installation 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)

Support

If you encounter any problems with the code or wish to propose new features, please feel free to open an issue. If you need general guidance, wish to discuss something, or just want to say hi, don't hesitate to leave a message in our Discord channel.

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