simulation

High-level simulation runner and result container.

from jbubble import run_simulation, SaveSpec

---

jbubble.simulation.SimulationResult

Bases: Module

Output of :func:run_simulation.

Generalises over any :class:~jbubble.bubble.BubbleState subclass: adding new degrees of freedom (e.g. temperature) to the state requires no changes here.

All array quantities are in SI units, sampled at ts.

Attributes:

Name	Type	Description
ts	(Array, shape(N))	Time points [s].
state	BubbleState, each field shape (N,)	Full state trajectory. state.R is the bubble radius, state.R_dot the radial velocity. For confined models (ConfinedBubbleState) also carries state.a and state.a_dot.
state_dot	BubbleState, each field shape (N,)	Time-derivative trajectory d(state)/dt returned by the EoM. state_dot.R_dot is the radial acceleration R̈(t).
driving_pressure	(Array, shape(N))	Applied acoustic pressure at the bubble [Pa].
converged	Array	Boolean scalar; True if the ODE solver converged successfully.

radius property

Bubble wall radius R(t) [m].

radial_velocity property

Bubble wall velocity Ṙ(t) [m/s].

radial_acceleration property

Bubble wall acceleration R̈(t) [m/s²], evaluated from the EoM RHS.

has_vessel property

True for confined-bubble models (ConfinedBubbleState).

vessel_radius property

Vessel wall radius a(t) [m], or None for unconfined models.

vessel_velocity property

Vessel wall velocity ȧ(t) [m/s], or None for unconfined models.

jbubble.simulation.run_simulation(eom, pulse, *, save_spec, state0=None, t_max=None, config=None, progress=False)

Run a simulation and return results in SI units.

Parameters:

Name	Type	Description	Default
eom	EquationOfMotion	Assembled equation of motion (e.g. KellerMiksis).	required
pulse	Pulse	Driving pulse.	required
save_spec	SaveSpec	Output specification (number of samples).	required
state0	BubbleState	Initial state. Defaults to eom.initial_state().	None
t_max	float	Integration end time [s]. None uses pulse.t_end.	None
config	SolverConfig	ODE solver settings.	None
progress	bool	Show progress meter.	False

Returns:

Type	Description
SimulationResult

Source code in jbubble/simulation.py

def run_simulation(
    eom: EquationOfMotion,
    pulse: Pulse,
    *,
    save_spec: SaveSpec,
    state0: Any = None,
    t_max: ArrayLike | None = None,
    config: SolverConfig | None = None,
    progress: bool = False,
) -> SimulationResult:
    """Run a simulation and return results in SI units.

    Parameters
    ----------
    eom : EquationOfMotion
        Assembled equation of motion (e.g. ``KellerMiksis``).
    pulse : Pulse
        Driving pulse.
    save_spec : SaveSpec
        Output specification (number of samples).
    state0 : BubbleState, optional
        Initial state.  Defaults to ``eom.initial_state()``.
    t_max : float, optional
        Integration end time [s].  ``None`` uses ``pulse.t_end``.
    config : SolverConfig, optional
        ODE solver settings.
    progress : bool
        Show progress meter.

    Returns
    -------
    SimulationResult
    """
    sol = solve_eom(
        eom,
        pulse,
        y0=state0,  # None → solve_eom calls eom.initial_state()
        t_max=t_max,
        save_spec=save_spec,
        config=config,
        progress=progress,
    )

    ts = cast(jax.Array, sol.ts)
    ys = cast(BubbleState, sol.ys)
    ys_dot: BubbleState = jax.vmap(lambda t, x: eom(t, x, pulse))(ts, ys)

    converged = diffrax.is_successful(sol.result)

    def _warn_not_converged(c: bool) -> None:
        if not c:
            warnings.warn(
                "ODE solver did not converge. "
                "Returned trajectory may be incomplete. Check `result.converged` before use.",
                UserWarning,
                stacklevel=2,
            )

    jax.debug.callback(_warn_not_converged, converged)

    return SimulationResult(
        ts=ts,
        state=ys,
        state_dot=ys_dot,
        driving_pressure=jax.vmap(pulse)(ts),
        converged=converged,
    )

---

Working with SimulationResult

SimulationResult is an Equinox module, so it is a JAX PyTree. All fields are jax.Array and can be passed through jax.jit, jax.vmap, and jax.grad.

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

# Convenience properties
ts  = result.ts                  # shape (N,) — time axis [s]
R   = result.radius              # shape (N,) — bubble radius [m]
Rd  = result.radial_velocity     # shape (N,) — wall velocity [m/s]
Rdd = result.radial_acceleration # shape (N,) — wall acceleration [m/s²]
pac = result.driving_pressure    # shape (N,) — acoustic pressure [Pa]

# Full state and derivatives
state     = result.state         # BubbleState with shape-(N,) arrays
state_dot = result.state_dot     # BubbleState time derivatives

# Solver diagnostics
ok = result.converged            # bool — False if ODE hit max_steps

Confined bubble (SphericalConfinement EoM)

When the EoM is SphericalConfinement, the result state is a ConfinedBubbleState:

from jbubble.bubble.eom import SphericalConfinement

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

print(result.has_vessel)       # True
a   = result.vessel_radius     # shape (N,) — vessel wall radius [m]
a_d = result.vessel_velocity   # shape (N,) — vessel wall velocity [m/s]