simulation¶

High-level simulation runner and result container.

from jbubble import run_simulation, SaveSpec

SimulationResult ¶

SimulationResult(
    ts: Array,
    state: BubbleState,
    state_dot: BubbleState,
    driving_pressure: Array,
    converged: Array,
)

Bases: Module

Output of run_simulation.

Generalises over any BubbleState subclass: when you add degrees of freedom, such as temperature, to the state, this class needs no changes.

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 and `state.R_dot` the radial velocity.
`state_dot`	`BubbleState, each field shape (N,)`	Time-derivative trajectory, \(\mathrm{d}(\text{state})/\mathrm{d}t\), that the equation of motion returns. `state_dot.R_dot` is the radial acceleration \(\ddot{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.

radial_acceleration `property` ¶

radial_acceleration

Bubble wall acceleration \(\ddot{R}(t)\) [m/s²], from the equation of motion.

radial_velocity `property` ¶

radial_velocity

Bubble wall velocity \(\dot{R}(t)\) [m/s].

radius `property` ¶

radius

Bubble wall radius \(R(t)\) [m].

run_simulation ¶

run_simulation(
    eom: EquationOfMotion,
    pulse: Pulse,
    *,
    save_spec: SaveSpec | None = None,
    state0: Any = None,
    t_max: ArrayLike | None = None,
    config: SolverConfig | None = None,
    adjoint: AbstractAdjoint | None = None,
    progress: bool = False,
) -> SimulationResult

Run a simulation and return the results in SI units.

run_simulation works under jax.jit and jax.vmap, and you can differentiate it with jax.grad. If the solver fails, for example because it reaches config.max_steps, the samples after the failure are inf and result.converged is False. Called eagerly, run_simulation also warns. Under jax.jit or jax.vmap it can't inspect the flag, so check result.converged yourself.

Parameters:

Name	Type	Description	Default
`eom`	`EquationOfMotion`	Assembled equation of motion, such as `KellerMiksis`.	required
`pulse`	`Pulse`	Driving pulse.	required
`save_spec`	`SaveSpec`	Output specification (number of samples). `None` uses `SaveSpec()`, 1024 evenly spaced time points.	`None`
`state0`	`BubbleState`	Initial state. `None` uses `eom.initial_state()`. A zero `R0` or `P_gas0`, the `BubbleState` defaults, is filled from `eom`, so `BubbleState(R=1.2 * R0)` starts at rest at 1.2 times the equilibrium radius. You can also build the state with `eom.initial_state(R=..., R_dot=...)`. For an empty cavity, set `P_gas0` to a tiny positive value, such as `1e-12`.	`None`
`t_max`	`float`	Integration end time [s]. `None` uses `pulse.t_end`.	`None`
`config`	`SolverConfig`	ODE solver settings. `None` uses `SolverConfig()`.	`None`
`adjoint`	`AbstractAdjoint`	How `jax.grad` differentiates through the solve. `None` uses `diffrax.RecursiveCheckpointAdjoint()`. Pass `diffrax.ForwardMode()` to use `jax.jacfwd`.	`None`
`progress`	`bool`	Whether to show a text progress meter. Default: `False`.	`False`

Returns:

Type	Description
`SimulationResult`	Trajectory, time derivatives, driving pressure, and convergence flag.

Warns:

Type	Description
`UserWarning`	If the solver didn't converge and `run_simulation` runs outside `jax.jit` and `jax.vmap`.

Source code in jbubble/simulation.py

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

    `run_simulation` works under `jax.jit` and `jax.vmap`, and you can
    differentiate it with `jax.grad`. If the solver fails, for example
    because it reaches `config.max_steps`, the samples after the failure
    are `inf` and `result.converged` is `False`. Called eagerly,
    `run_simulation` also warns. Under `jax.jit` or `jax.vmap` it can't
    inspect the flag, so check `result.converged` yourself.

    Parameters
    ----------
    eom : EquationOfMotion
        Assembled equation of motion, such as
        [`KellerMiksis`][jbubble.bubble.eom.KellerMiksis].
    pulse : Pulse
        Driving pulse.
    save_spec : SaveSpec, optional
        Output specification (number of samples). `None` uses
        [`SaveSpec()`][jbubble.solver.SaveSpec], 1024 evenly spaced time
        points.
    state0 : BubbleState, optional
        Initial state. `None` uses `eom.initial_state()`. A zero `R0` or
        `P_gas0`, the `BubbleState` defaults, is filled from `eom`, so
        `BubbleState(R=1.2 * R0)` starts at rest at 1.2 times the
        equilibrium radius. You can also build the state with
        `eom.initial_state(R=..., R_dot=...)`. For an empty cavity, set
        `P_gas0` to a tiny positive value, such as `1e-12`.
    t_max : float, optional
        Integration end time [s]. `None` uses `pulse.t_end`.
    config : SolverConfig, optional
        ODE solver settings. `None` uses
        [`SolverConfig()`][jbubble.solver.SolverConfig].
    adjoint : diffrax.AbstractAdjoint, optional
        How `jax.grad` differentiates through the solve. `None` uses
        `diffrax.RecursiveCheckpointAdjoint()`. Pass `diffrax.ForwardMode()`
        to use `jax.jacfwd`.
    progress : bool
        Whether to show a text progress meter. Default: `False`.

    Returns
    -------
    SimulationResult
        Trajectory, time derivatives, driving pressure, and convergence
        flag.

    Warns
    -----
    UserWarning
        If the solver didn't converge and `run_simulation` runs outside
        `jax.jit` and `jax.vmap`.
    """
    result = _simulate(
        eom,
        pulse,
        save_spec=save_spec,
        state0=state0,
        t_max=t_max,
        config=config,
        adjoint=adjoint,
        progress=progress,
    )

    converged = result.converged
    if not isinstance(converged, Tracer) and not bool(jnp.all(converged)):
        warnings.warn(
            "ODE solver did not converge, so the trajectory is incomplete: "
            "samples after the failure are inf. Check `result.converged`. "
            "If the solver reached `max_steps`, raise "
            "`SolverConfig(max_steps=...)`, or use `SolverConfig.stiff()` for "
            "a stiff problem.",
            UserWarning,
            stacklevel=2,
        )
    return result

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