jbubble 0.1.0

Differentiable microbubble dynamics primitives

jbubble

Differentiable microbubble dynamics in JAX.

Comparison of R(t) curves, with and without a lipid coating.

---

jbubble is a research library for simulating and fitting acoustic microbubble dynamics, built on JAX, Equinox, and diffrax. Models are designed to be composable, extendable, fully differentiable, and jit-compatible.

Developed by the Noninvasive Surgery & Biopsy Laboratory at Imperial College London.

Why jbubble?

Feature	jbubble	APECSS / MATLAB
JIT compilation	jax.jit	No
Vectorised sweeps	jax.vmap	Script loops
Gradient-based fitting	jax.grad	Finite differences
Composable physics	Mix any gas + shell + medium	Fixed combinations
Neural components	NeuralProperty, NeuralPulse	No

Installation

pip install jbubble

For development:

git clone https://github.com/imperial-nsb/jbubble.git
cd jbubble
pip install -e ".[dev]"

Requires Python ≥ 3.11. See the [installation guide] for more information.

Quick start

Using a preset

import jax
from jbubble import run_simulation, SaveSpec
from jbubble.utils.presets import lipid_bubble

eom, pulse = lipid_bubble(R0=2e-6, freq=1e6, pressure=100e3)

result = jax.jit(run_simulation)(
    eom, pulse,
    save_spec=SaveSpec(num_samples=1000),
    t_max=10e-6,
)
print(result.radius.max() / eom.R0)  # peak expansion ratio

Three presets are available: free_bubble, lipid_bubble, and thick_shell_bubble.

Composing models manually

import jax
from jbubble import run_simulation, SaveSpec
from jbubble.bubble.eom import KellerMiksis
from jbubble.bubble.gas import PolytropicGas
from jbubble.bubble.shell import NoShell
from jbubble.bubble.medium import NewtonianMedium
from jbubble.pulse import ToneBurst
from jbubble.pulse.shapes import Sine

eom = KellerMiksis(
    gas=PolytropicGas(gamma=1.4),
    shell=NoShell(sigma=0.072),
    medium=NewtonianMedium(mu=1e-3),
    R0=2e-6, P_amb=101325.0, rho_L=998.0, c_L=1500.0,
)
pulse = ToneBurst(freq=1e6, pressure=100e3, shape=Sine(), cycle_num=5)

result = jax.jit(run_simulation)(
    eom, pulse,
    save_spec=SaveSpec(num_samples=1000),
    t_max=10e-6,
)

Any gas model works with any shell model and any medium model—all combinations are valid and differentiated through automatically via jax.grad.

Key capabilities

• Composable physics — Mix and match 6 equations of motion (Rayleigh-Plesset through Gilmore), 2 gas models, 3 shell models, and 4 medium models. All combinations work automatically.

• Batch parameter sweeps — Run thousands of simulations in parallel with GridSweep, which uses jax.vmap under the hood:

  sweep = GridSweep(fn=simulate, search_space={"R0": radii, "freq": freqs})
  results = sweep.run()  # shape (len(radii), len(freqs))
  

• Gradient-based fitting — Fit model parameters to experimental data via fit_parameters and optax:

  fit_result = fit_parameters(
      make_model=lambda p: (make_eom(p), pulse),
      params0=initial_guess,
      loss_fn=my_loss,
      optimizer=optax.adam(1e-2),
  )
  

• Acoustic emissions — Compute radiated pressure at arbitrary field points from solved trajectories using IncompressibleMonopole or QuasiAcoustic emission models.

Examples

The examples/ directory contains 12 self-contained scripts:

#	Script	Description
00	Presets	Quickest start — pick a preset and run
01	Basic simulation	Assemble an EoM from components manually
02	Pulse algebra	Add, scale, and window pulse waveforms
03	Shell models	Compare no-shell, lipid, and thick-shell coatings
04	Batch sweeps	GridSweep + vmap over parameter grids
05	Parameter fitting	Gradient-based estimation of shell elasticity
06	JIT timing	Benchmark JIT compilation vs steady-state throughput
07	Cavitation regimes	Stable vs inertial cavitation physics
08	Acoustic emissions	Monopole and quasi-acoustic radiated pressure
09	Custom pulse shapes	Subclass FourierPulseShape for custom waveforms
10	Envelopes	Envelope types and their gradient compatibility
11	Gradient resonance	2D sweep + gradient descent to resonance peak

Documentation

Full documentation is available at imperial-nsb.github.io/jbubble, including:

• Quickstart guide
• Bubble models guide
• JAX tips (JIT, vmap, grad)
• API reference

Contributing

Contributions are welcome! Please see CONTRIBUTING.md for guidelines.

License

MIT License. Copyright (c) 2026 Noninvasive Surgery & Biopsy Laboratory. See LICENSE for details.