solver

Low-level ODE integration primitives. Most users should prefer run_simulation from jbubble.simulation, which wraps these with post-processing. Use solve_eom directly when you need access to the raw diffrax.Solution object.

from jbubble import solve_eom, SaveSpec, SolverConfig

---

jbubble.solver.SaveSpec

Bases: Module

Specification for ODE output sampling.

Parameters:

Name	Type	Description	Default
num_samples	int	Number of evenly-spaced time points to record. Default: 1024.	required

jbubble.solver.SolverConfig

Bases: Module

Numerical integration settings for :func:solve_eom.

Fields

solver : diffrax.AbstractSolver ODE solver. Default: Kvaerno5(). stepsize_controller : diffrax.AbstractStepSizeController Step-size controller. Default: PIDController(rtol=1e-3, atol=1e-6). dt0 : float Initial step size [s]. Default: 1e-9. max_steps : int Maximum solver steps per integration. Default: 50 000.

jbubble.solver.solve_eom(eom, pulse, *, y0=None, t_max=None, save_spec=None, config=None, adjoint=None, progress=False)

Solve bubble dynamics for an EquationOfMotion.

Parameters:

Name	Type	Description	Default
eom	EquationOfMotion	Assembled equation of motion (e.g. KellerMiksis).	required
pulse	Pulse	Driving acoustic pulse.	required
y0	BubbleState	Initial state (dimensionless). If None, derived from eom.initial_state().	None
t_max	float	Integration end time [s]. If None, derived from pulse.t_end.	None
save_spec	SaveSpec	Output sampling specification. Default: 1024 evenly-spaced time points.	None
config	SolverConfig	Numerical integration settings. Default: Kvaerno5 with PIDController(rtol=1e-3, atol=1e-6).	None
adjoint	AbstractAdjoint	Adjoint method for gradient computation. Default: diffrax built-in (RecursiveCheckpointAdjoint). For gradient-based fitting through an explicit solver use diffrax.BacksolveAdjoint().	None
progress	bool	Show a text progress meter.	False

Returns:

Type	Description
Solution	Solution object with ts and ys.

Source code in jbubble/solver.py

def solve_eom(
    eom: EquationOfMotion,
    pulse: Pulse,
    *,
    y0: Any = None,
    t_max: ArrayLike | None = None,
    save_spec: SaveSpec | None = None,
    config: SolverConfig | None = None,
    adjoint: diffrax.AbstractAdjoint | None = None,
    progress: bool = False,
) -> diffrax.Solution:
    """Solve bubble dynamics for an ``EquationOfMotion``.

    Parameters
    ----------
    eom : EquationOfMotion
        Assembled equation of motion (e.g. ``KellerMiksis``).
    pulse : Pulse
        Driving acoustic pulse.
    y0 : BubbleState, optional
        Initial state (dimensionless).  If *None*, derived from
        ``eom.initial_state()``.
    t_max : float, optional
        Integration end time [s].  If *None*, derived from ``pulse.t_end``.
    save_spec : SaveSpec, optional
        Output sampling specification.  Default: 1024 evenly-spaced
        time points.
    config : SolverConfig, optional
        Numerical integration settings.  Default: Kvaerno5 with
        PIDController(rtol=1e-3, atol=1e-6).
    adjoint : diffrax.AbstractAdjoint, optional
        Adjoint method for gradient computation.  Default: diffrax built-in
        (``RecursiveCheckpointAdjoint``).  For gradient-based fitting through
        an explicit solver use ``diffrax.BacksolveAdjoint()``.
    progress : bool
        Show a text progress meter.

    Returns
    -------
    diffrax.Solution
        Solution object with ``ts`` and ``ys``.
    """
    if save_spec is None:
        save_spec = SaveSpec(num_samples=1024)

    if config is None:
        config = SolverConfig()

    if y0 is None:
        y0 = eom.initial_state()

    t0 = jnp.asarray(0.0)
    t1 = jnp.asarray(pulse.t_end if t_max is None else t_max)
    saveat = save_spec.build(t0, t1)

    def ode_func(t, state, args):
        eom_model, pulse_model = args
        return eom_model(t, state, pulse_model)

    term = diffrax.ODETerm(ode_func)
    progress_meter = (
        diffrax.TextProgressMeter() if progress else diffrax.NoProgressMeter()
    )
    _adjoint = adjoint if adjoint is not None else diffrax.RecursiveCheckpointAdjoint()

    return diffrax.diffeqsolve(
        term,
        config.solver,
        t0=t0,
        t1=t1,
        dt0=config.dt0,
        y0=y0,
        args=(eom, pulse),
        saveat=saveat,
        stepsize_controller=config.stepsize_controller,
        max_steps=config.max_steps,
        progress_meter=progress_meter,
        throw=False,
        adjoint=_adjoint,
    )

---

Solver choice and tolerances

The default solver is Kvaerno5 (an implicit 5th-order Runge–Kutta method) with a PID step-size controller at relative tolerance \(10^{-3}\) and absolute tolerance \(10^{-6}\). This is appropriate for most bubble dynamics simulations.

For highly stiff problems (e.g. very small bubbles, extreme driving pressures, or large shear moduli in the medium), tighten the tolerances:

import diffrax
from jbubble import SolverConfig

config = SolverConfig(
    stepsize_controller=diffrax.PIDController(rtol=1e-5, atol=1e-8),
    max_steps=50_000,
)

For gradient-based fitting, consider using the RecursiveCheckpointAdjoint to reduce memory usage during backpropagation:

import diffrax

result = fit_parameters(
    ...,
    adjoint=diffrax.RecursiveCheckpointAdjoint(),
)