ivpsolve

Routines for estimating solutions of initial value problems.

adaptive(solver, atol=0.0001, rtol=0.01, control=None, norm_ord=None)

Make an IVP solver adaptive.

Source code in probdiffeq/ivpsolve.py

def adaptive(solver, atol=1e-4, rtol=1e-2, control=None, norm_ord=None):
    """Make an IVP solver adaptive."""
    if control is None:
        control = control_proportional_integral()

    return _AdaSolver(solver, atol=atol, rtol=rtol, control=control, norm_ord=norm_ord)

control_integral(*, clip=False, safety=0.95, factor_min=0.2, factor_max=10.0) -> _Controller

Construct an integral-controller.

Source code in probdiffeq/ivpsolve.py

def control_integral(
    *, clip=False, safety=0.95, factor_min=0.2, factor_max=10.0
) -> _Controller:
    """Construct an integral-controller."""

    def init(dt, /):
        return dt

    def apply(dt, /, error_norm, error_contraction_rate):
        error_power = error_norm ** (-1.0 / error_contraction_rate)
        scale_factor_unclipped = safety * error_power

        scale_factor_clipped_min = np.minimum(scale_factor_unclipped, factor_max)
        scale_factor = np.maximum(factor_min, scale_factor_clipped_min)
        return scale_factor * dt

    def extract(dt, /):
        return dt

    if clip:

        def clip_fun(dt, /, t, t1):
            return np.minimum(dt, t1 - t)

        return _Controller(init=init, apply=apply, extract=extract, clip=clip_fun)

    return _Controller(init=init, apply=apply, extract=extract, clip=lambda v, **_kw: v)

control_proportional_integral(*, clip: bool = False, safety=0.95, factor_min=0.2, factor_max=10.0, power_integral_unscaled=0.3, power_proportional_unscaled=0.4) -> _Controller

Construct a proportional-integral-controller with time-clipping.

Source code in probdiffeq/ivpsolve.py

def control_proportional_integral(
    *,
    clip: bool = False,
    safety=0.95,
    factor_min=0.2,
    factor_max=10.0,
    power_integral_unscaled=0.3,
    power_proportional_unscaled=0.4,
) -> _Controller:
    """Construct a proportional-integral-controller with time-clipping."""

    class PIState(containers.NamedTuple):
        dt: float
        error_norm_previously_accepted: float

    def init(dt: float, /) -> PIState:
        return PIState(dt, 1.0)

    def apply(state: PIState, /, error_norm, error_contraction_rate) -> PIState:
        dt_proposed, error_norm_prev = state
        n1 = power_integral_unscaled / error_contraction_rate
        n2 = power_proportional_unscaled / error_contraction_rate

        a1 = (1.0 / error_norm) ** n1
        a2 = (error_norm_prev / error_norm) ** n2
        scale_factor_unclipped = safety * a1 * a2

        scale_factor_clipped_min = np.minimum(scale_factor_unclipped, factor_max)
        scale_factor = np.maximum(factor_min, scale_factor_clipped_min)
        error_norm_prev = np.where(error_norm <= 1.0, error_norm, error_norm_prev)

        dt_proposed = scale_factor * dt_proposed
        return PIState(dt_proposed, error_norm_prev)

    def extract(state: PIState, /) -> float:
        dt_proposed, _error_norm_previously_accepted = state
        return dt_proposed

    if clip:

        def clip_fun(state: PIState, /, t, t1) -> PIState:
            dt_proposed, error_norm_previously_accepted = state
            dt = dt_proposed
            dt_clipped = np.minimum(dt, t1 - t)
            return PIState(dt_clipped, error_norm_previously_accepted)

        return _Controller(init=init, apply=apply, extract=extract, clip=clip_fun)

    return _Controller(init=init, apply=apply, extract=extract, clip=lambda v, **_kw: v)

dt0(vf_autonomous, initial_values, /, scale=0.01, nugget=1e-05)

Propose an initial time-step.

Source code in probdiffeq/ivpsolve.py

def dt0(vf_autonomous, initial_values, /, scale=0.01, nugget=1e-5):
    """Propose an initial time-step."""
    u0, *_ = initial_values
    f0 = vf_autonomous(*initial_values)

    norm_y0 = linalg.vector_norm(u0)
    norm_dy0 = linalg.vector_norm(f0) + nugget

    return scale * norm_y0 / norm_dy0

dt0_adaptive(vf, initial_values, /, t0, *, error_contraction_rate, rtol, atol)

Propose an initial time-step as a function of the tolerances.

Source code in probdiffeq/ivpsolve.py

def dt0_adaptive(vf, initial_values, /, t0, *, error_contraction_rate, rtol, atol):
    """Propose an initial time-step as a function of the tolerances."""
    # Algorithm from:
    # E. Hairer, S. P. Norsett G. Wanner,
    # Solving Ordinary Differential Equations I: Nonstiff Problems, Sec. II.4.
    # Implementation mostly copied from
    #
    # https://github.com/google/jax/blob/main/jax/experimental/ode.py
    #

    if len(initial_values) > 1:
        raise ValueError
    y0 = initial_values[0]

    f0 = vf(y0, t=t0)
    scale = atol + np.abs(y0) * rtol
    d0, d1 = linalg.vector_norm(y0), linalg.vector_norm(f0)

    dt0 = np.where((d0 < 1e-5) | (d1 < 1e-5), 1e-6, 0.01 * d0 / d1)

    y1 = y0 + dt0 * f0
    f1 = vf(y1, t=t0 + dt0)
    d2 = linalg.vector_norm((f1 - f0) / scale) / dt0

    dt1 = np.where(
        (d1 <= 1e-15) & (d2 <= 1e-15),
        np.maximum(1e-6, dt0 * 1e-3),
        (0.01 / np.maximum(d1, d2)) ** (1.0 / (error_contraction_rate + 1.0)),
    )
    return np.minimum(100.0 * dt0, dt1)

solve_adaptive_save_at(vector_field, initial_condition, save_at, adaptive_solver, dt0) -> _Solution

Solve an initial value problem and return the solution at a pre-determined grid.

This algorithm implements the method by Krämer (2024). Please consider citing it if you use it for your research. A PDF is available here and Krämer's (2024) experiments are here.

BibTex for Krämer (2024)

@article{krämer2024adaptive,
    title={Adaptive Probabilistic {ODE} Solvers Without
    Adaptive Memory Requirements},
    author={Kr{\"a}mer, Nicholas},
    year={2024},
    eprint={2410.10530},
    archivePrefix={arXiv},
    url={https://arxiv.org/abs/2410.10530},
}

Source code in probdiffeq/ivpsolve.py

def solve_adaptive_save_at(
    vector_field, initial_condition, save_at, adaptive_solver, dt0
) -> _Solution:
    r"""Solve an initial value problem and return the solution at a pre-determined grid.

    This algorithm implements the method by Krämer (2024).
    Please consider citing it if you use it for your research.
    A PDF is available [here](https://arxiv.org/abs/2410.10530)
    and Krämer's (2024) experiments are
    [here](https://github.com/pnkraemer/code-adaptive-prob-ode-solvers).


    ??? note "BibTex for Krämer (2024)"
        ```bibtex
        @article{krämer2024adaptive,
            title={Adaptive Probabilistic {ODE} Solvers Without
            Adaptive Memory Requirements},
            author={Kr{\"a}mer, Nicholas},
            year={2024},
            eprint={2410.10530},
            archivePrefix={arXiv},
            url={https://arxiv.org/abs/2410.10530},
        }
        ```

    """
    if not adaptive_solver.solver.is_suitable_for_save_at:
        msg = (
            f"Strategy {adaptive_solver.solver} should not "
            f"be used in solve_adaptive_save_at. "
        )
        warnings.warn(msg, stacklevel=1)

    (_t, solution_save_at), _, num_steps = _solve_adaptive_save_at(
        tree_util.Partial(vector_field),
        save_at[0],
        initial_condition,
        save_at=save_at[1:],
        adaptive_solver=adaptive_solver,
        dt0=dt0,
    )

    # I think the user expects the initial condition to be part of the state
    # (as well as marginals), so we compute those things here
    posterior_t0, *_ = initial_condition
    posterior_save_at, output_scale = solution_save_at
    _tmp = _userfriendly_output(posterior=posterior_save_at, posterior_t0=posterior_t0)
    marginals, posterior = _tmp
    u = impl.stats.qoi(marginals)
    return _Solution(
        t=save_at,
        u=u,
        marginals=marginals,
        posterior=posterior,
        output_scale=output_scale,
        num_steps=num_steps,
    )

solve_adaptive_save_every_step(vector_field, initial_condition, t0, t1, adaptive_solver, dt0) -> _Solution

Solve an initial value problem and save every step.

This function uses a native-Python while loop.

Warning

Not JITable, not reverse-mode-differentiable.

Source code in probdiffeq/ivpsolve.py

def solve_adaptive_save_every_step(
    vector_field, initial_condition, t0, t1, adaptive_solver, dt0
) -> _Solution:
    """Solve an initial value problem and save every step.

    This function uses a native-Python while loop.

    !!! warning
        Not JITable, not reverse-mode-differentiable.
    """
    if not adaptive_solver.solver.is_suitable_for_save_every_step:
        msg = (
            f"Strategy {adaptive_solver.solver} should not "
            f"be used in solve_adaptive_save_every_step."
        )
        warnings.warn(msg, stacklevel=1)

    generator = _solution_generator(
        tree_util.Partial(vector_field),
        t0,
        initial_condition,
        t1=t1,
        adaptive_solver=adaptive_solver,
        dt0=dt0,
    )
    tmp = tree_array_util.tree_stack(list(generator))
    (t, solution_every_step), _dt, num_steps = tmp

    # I think the user expects the initial time-point to be part of the grid
    # (Even though t0 is not computed by this function)
    t = np.concatenate((np.atleast_1d(t0), t))

    # I think the user expects marginals, so we compute them here
    posterior_t0, *_ = initial_condition
    posterior, output_scale = solution_every_step
    _tmp = _userfriendly_output(posterior=posterior, posterior_t0=posterior_t0)
    marginals, posterior = _tmp

    u = impl.stats.qoi(marginals)
    return _Solution(
        t=t,
        u=u,
        marginals=marginals,
        posterior=posterior,
        output_scale=output_scale,
        num_steps=num_steps,
    )

solve_adaptive_terminal_values(vector_field, initial_condition, t0, t1, adaptive_solver, dt0) -> _Solution

Simulate the terminal values of an initial value problem.

Source code in probdiffeq/ivpsolve.py

def solve_adaptive_terminal_values(
    vector_field, initial_condition, t0, t1, adaptive_solver, dt0
) -> _Solution:
    """Simulate the terminal values of an initial value problem."""
    save_at = np.asarray([t1])
    (_t, solution_save_at), _, num_steps = _solve_adaptive_save_at(
        tree_util.Partial(vector_field),
        t0,
        initial_condition,
        save_at=save_at,
        adaptive_solver=adaptive_solver,
        dt0=dt0,
    )
    # "squeeze"-type functionality (there is only a single state!)
    squeeze_fun = functools.partial(np.squeeze_along_axis, axis=0)
    solution_save_at = tree_util.tree_map(squeeze_fun, solution_save_at)
    num_steps = tree_util.tree_map(squeeze_fun, num_steps)

    # I think the user expects marginals, so we compute them here
    posterior, output_scale = solution_save_at
    marginals = posterior.init if isinstance(posterior, stats.MarkovSeq) else posterior
    u = impl.stats.qoi(marginals)
    return _Solution(
        t=t1,
        u=u,
        marginals=marginals,
        posterior=posterior,
        output_scale=output_scale,
        num_steps=num_steps,
    )

solve_fixed_grid(vector_field, initial_condition, grid, solver) -> _Solution

Solve an initial value problem on a fixed, pre-determined grid.

Source code in probdiffeq/ivpsolve.py

def solve_fixed_grid(vector_field, initial_condition, grid, solver) -> _Solution:
    """Solve an initial value problem on a fixed, pre-determined grid."""
    # Compute the solution

    def body_fn(s, dt):
        _error, s_new = solver.step(state=s, vector_field=vector_field, dt=dt)
        return s_new, s_new

    t0 = grid[0]
    state0 = solver.init(t0, initial_condition)
    _, result_state = control_flow.scan(body_fn, init=state0, xs=np.diff(grid))
    _t, (posterior, output_scale) = solver.extract(result_state)

    # I think the user expects marginals, so we compute them here
    posterior_t0, *_ = initial_condition
    _tmp = _userfriendly_output(posterior=posterior, posterior_t0=posterior_t0)
    marginals, posterior = _tmp

    u = impl.stats.qoi(marginals)
    return _Solution(
        t=grid,
        u=u,
        marginals=marginals,
        posterior=posterior,
        output_scale=output_scale,
        num_steps=np.arange(1.0, len(grid)),
    )