Source code for probnum.diffeq._perturbsolve_ivp

"""Perturbation-based probabilistic ODE solver."""

from __future__ import annotations

from typing import Callable, Optional

import numpy as np
import scipy.integrate

from probnum import problems
from probnum.diffeq import perturbed, stepsize
from probnum.typing import ArrayLike, FloatLike

__all__ = ["perturbsolve_ivp"]


METHODS = {
    "RK45": scipy.integrate.RK45,
    "RK23": scipy.integrate.RK23,
}
"""Implemented Scipy solvers."""


# aliases for (otherwise too-)long lines
lognorm = perturbed.step.PerturbedStepSolver.construct_with_lognormal_perturbation
uniform = perturbed.step.PerturbedStepSolver.construct_with_uniform_perturbation
PERTURBS = {
    "step-lognormal": lognorm,
    "step-uniform": uniform,
}
"""Implemented perturbation-styles."""


# This interface function is allowed to have many input arguments.
# Having many input arguments implies having many local arguments,
# so we need to disable both here.
# pylint: disable="too-many-arguments,too-many-locals"
[docs]def perturbsolve_ivp( f: Callable, t0: FloatLike, tmax: FloatLike, y0: ArrayLike, rng: np.random.Generator, method: str = "RK45", perturb: str = "step-lognormal", noise_scale: FloatLike = 10.0, adaptive: bool = True, atol: FloatLike = 1e-6, rtol: FloatLike = 1e-3, step: Optional[FloatLike] = None, time_stops: Optional[ArrayLike] = None, ) -> perturbed.step.PerturbedStepSolution: """Solve an initial value problem with a perturbation-based ODE solver. Parameters ---------- f ODE vector field. t0 Initial time point. tmax Final time point. y0 Initial value. rng Random number generator. method Integration method to use. The following are available (docs adapted from https://docs.scipy.org/doc/scipy/\ reference/generated/scipy.integrate.solve_ivp.html): * `RK45` (default): Explicit Runge-Kutta method of order 5(4) [2]_. The error is controlled assuming accuracy of the fourth-order method, but steps are taken using the fifth-order accurate formula (local extrapolation is done). A quartic interpolation polynomial is used for the dense output [3]_. Can be applied in the complex domain. * `RK23`: Explicit Runge-Kutta method of order 3(2) [4]_. The error is controlled assuming accuracy of the second-order method, but steps are taken using the third-order accurate formula (local extrapolation is done). A cubic Hermite polynomial is used for the dense output. Can be applied in the complex domain. Other integrators are not supported currently. perturb Which perturbation style to use. Currently, the following methods are supported: * `step-lognormal`: Perturbed-step (aka random time-step numerical integration) method with lognormally distributed perturbation of the step-size [1]_. * `step-uniform`: Perturbed-step (aka random time-step numerical integration) method with uniformly distributed perturbation of the step-size [1]_. Other integrators are not supported currently. noise_scale Scale of the perturbation. Optional. Default is 10.0. The magnitude of this parameter significantly impacts the width of the posterior. adaptive Whether to use adaptive steps or not. Default is `True`. atol Absolute tolerance of the adaptive step-size selection scheme. Optional. Default is ``1e-6``. rtol Relative tolerance of the adaptive step-size selection scheme. Optional. Default is ``1e-3``. step Step size. If atol and rtol are not specified, this step-size is used for a fixed-step ODE solver. If they are specified, this only affects the first step. Optional. Default is None, in which case the first step is chosen as prescribed by :meth:`propose_firststep`. time_stops Time-points through which the solver must step. Optional. Default is None. Raises ------ ValueError If the 'method' string does not correspond to a supported method. ValueError If the 'perturb' string does not correspond to a supported perturbation style. Returns ------- perturbed.step.PerturbedStepSolution ODE Solution of the perturbed-step-solver. Examples -------- >>> from probnum.diffeq import perturbsolve_ivp >>> import numpy as np Solve a simple logistic ODE with fixed steps. Per default, `perturbsolve_ivp` uses a perturbed-step solver with lognormal perturbation. >>> rng = np.random.default_rng(seed=2) >>> >>> def f(t, x): ... return 4*x*(1-x) >>> >>> y0 = np.array([0.15]) >>> t0, tmax = 0., 1.5 >>> solution = perturbsolve_ivp( ... f, t0, tmax, y0, rng=rng, step=0.25, method="RK23", adaptive=False ... ) >>> print(np.round(solution.states.mean, 3)) [[0.15 ] [0.325] [0.56 ] [0.772] [0.893] [0.964] [0.989]] Each solution is the result of a randomly-perturbed call of an underlying Runge-Kutta solver. Therefore, if you call it again, the output will be different: >>> other_solution = perturbsolve_ivp( ... f, t0, tmax, y0, rng=rng, step=0.25, method="RK23", adaptive=False ... ) >>> print(np.round(other_solution.states.mean, 3)) [[0.15 ] [0.319] [0.57 ] [0.785] [0.908] [0.968] [0.989]] Other methods, such as `RK45` (as well as other perturbation styles) are easily accessible. Let us solve the same equation, with an adaptive RK45 solver that uses uniformly perturbed steps. >>> solution = perturbsolve_ivp( ... f, t0, tmax, y0, rng=rng, atol=1e-5, rtol=1e-6, ... method="RK45", perturb="step-uniform", adaptive=True ... ) >>> print(np.round(solution.states.mean, 3)) [[0.15 ] [0.152] [0.167] [0.26 ] [0.431] [0.646] [0.849] [0.883] [0.915] [0.953] [0.976] [0.986]] References ---------- .. [1] Abdulle, A. and Garegnani, G.. Random time step probabilistic methods for uncertainty quantification in chaotic and geometric numerical integration. Statistics and Computing. 2020. .. [2] J. R. Dormand, P. J. Prince.. A family of embedded Runge-Kutta formulae. Journal of Computational and Applied Mathematics, Vol. 6, No. 1, pp. 19-26, 1980. .. [3] L. W. Shampine. Some Practical Runge-Kutta Formulas. Mathematics of Computation, Vol. 46, No. 173, pp. 135-150, 1986. .. [4] P. Bogacki, L.F. Shampine. A 3(2) Pair of Runge-Kutta Formulas. Appl. Math. Lett. Vol. 2, No. 4. pp. 321-325, 1989. """ ivp = problems.InitialValueProblem(t0=t0, tmax=tmax, y0=np.asarray(y0), f=f) steprule = _create_steprule(adaptive, atol, ivp, rtol, step) if method not in METHODS: raise ValueError( f"Parameter method='{method}' is not supported. " f"Possible values are {list(METHODS.keys())}." ) if perturb not in PERTURBS: raise ValueError( f"Parameter perturb='{perturb}' is not supported. " f"Possible values are {list(PERTURBS.keys())}." ) wrapped_scipy_solver = perturbed.scipy_wrapper.WrappedScipyRungeKutta( METHODS[method], steprule=steprule ) perturbed_solver = PERTURBS[perturb]( rng=rng, solver=wrapped_scipy_solver, noise_scale=noise_scale, ) return perturbed_solver.solve(ivp=ivp, stop_at=time_stops)
def _create_steprule(adaptive, atol, ivp, rtol, step): if adaptive is True: if atol is None or rtol is None: raise ValueError( "Please provide absolute and relative tolerance for adaptive steps." ) firststep = step if step is not None else stepsize.propose_firststep(ivp) steprule = stepsize.AdaptiveSteps(firststep=firststep, atol=atol, rtol=rtol) else: steprule = stepsize.ConstantSteps(step) return steprule