solver¶

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:
    r"""Solve the bubble dynamics for an equation of motion.

    Integrates the initial value problem

    $$
    \frac{\mathrm{d}\,\text{state}}{\mathrm{d}t}
        = \text{eom}(t, \text{state}, \text{pulse}),
    \qquad \text{state}(0) = y_0,
    \qquad 0 \le t \le t_\text{max}.
    $$

    The solver works on the scaled state `state / eom.state_scale(y0)`, so
    the tolerances in `config` are relative to `R0` for the radius and to
    $\sqrt{P_\text{amb}/\rho_L}$ for the wall velocity, and returns the
    solution in SI units. It evaluates the equation of motion only at
    admissible states (see
    [`EquationOfMotion.is_admissible`][jbubble.bubble.eom.EquationOfMotion.is_admissible]);
    elsewhere the right-hand side is NaN, so the adaptive solver rejects
    the trial step. Rejected steps never enter the solution, so the guard
    keeps gradients finite without changing the forward solve: with an
    explicit solver, such as the default `Dopri5`, the solution is the same,
    bit for bit, as without the guard. With an implicit solver, XLA can
    round the Newton iterations differently in the last bit, so the
    solution agrees to within the solver tolerance.

    The integration doesn't raise when the solver fails (`throw=False`);
    check `diffrax.is_successful` on the solution's `result`.

    Parameters
    ----------
    eom : EquationOfMotion
        Assembled equation of motion, such as
        [`KellerMiksis`][jbubble.bubble.eom.KellerMiksis].
    pulse : Pulse
        Driving acoustic pulse. An adaptive step-size controller steps to
        every time in
        [`pulse.window_edges`][jbubble.pulse.base.Pulse.window_edges], so
        it can't step over a pulse that starts late.
    y0 : BubbleState, optional
        Initial state in SI units. `None` uses `eom.initial_state()`. A zero
        `R0` or `P_gas0`, the `BubbleState` defaults, means "unset":
        `solve_eom` fills it from the equation of motion, so
        `BubbleState(R=1.2 * R0)` starts at rest at 1.2 times the equilibrium
        radius. [`EquationOfMotion.initial_state`][jbubble.bubble.eom.EquationOfMotion.initial_state]
        with `R=` and `R_dot=` builds the same state explicitly. 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`.
    save_spec : SaveSpec, optional
        Output sampling specification. `None` uses 1024 evenly spaced time
        points.
    config : SolverConfig, optional
        Numerical integration settings. `None` uses
        [`SolverConfig()`][jbubble.solver.SolverConfig].
    adjoint : diffrax.AbstractAdjoint, optional
        How `jax.grad` differentiates through the solve. `None` uses
        `diffrax.RecursiveCheckpointAdjoint()`, which gives the exact
        gradient of the discretised solution. Pass `diffrax.ForwardMode()`
        to use `jax.jacfwd`, for example in a Levenberg-Marquardt fit.
        diffrax advises against `diffrax.BacksolveAdjoint`, whose gradients
        are approximate.
    progress : bool
        Whether to show a text progress meter. Default: `False`.

    Returns
    -------
    diffrax.Solution
        Solution object with `ts`, and `ys` in SI units.
    """
    if save_spec is None:
        save_spec = SaveSpec()
    assert isinstance(save_spec, SaveSpec)

    if config is None:
        config = SolverConfig()
    assert isinstance(config, SolverConfig)

    y0 = eom.initial_state() if y0 is None else _with_equilibrium(eom, y0)
    scale = jax.lax.stop_gradient(eom.state_scale(y0))
    z0 = jtu.tree_map(jnp.divide, y0, scale)

    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)

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

    sol = diffrax.diffeqsolve(
        diffrax.ODETerm(_guarded_vector_field),
        config.solver,
        t0=t0,
        t1=t1,
        dt0=config.dt0,
        y0=z0,
        args=(eom, pulse, jax.lax.stop_gradient(z0), scale),
        saveat=saveat,
        stepsize_controller=_step_to_window_edges(config.stepsize_controller, pulse),
        max_steps=config.max_steps,
        progress_meter=progress_meter,
        throw=False,
        adjoint=_adjoint,
    )
    ys = jtu.tree_map(jnp.multiply, sol.ys, scale)
    return eqx.tree_at(lambda s: s.ys, sol, ys)