ivpsolvers

Probabilistic IVP solvers.

adaptive(slvr, /, *, ssm, atol=0.0001, rtol=0.01, control=None, norm_ord=None)

Make an IVP solver adaptive.

Source code in probdiffeq/ivpsolvers.py

def adaptive(slvr, /, *, ssm, 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(
        slvr, ssm=ssm, 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/ivpsolvers.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_power):
        # 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/ivpsolvers.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_power_previously_accepted: float

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

    def apply(state: PIState, /, *, error_power) -> PIState:
        # error_power = error_norm ** (-1.0 / error_contraction_rate)
        dt_proposed, error_power_prev = state

        a1 = error_power**power_integral_unscaled
        a2 = (error_power / error_power_prev) ** power_proportional_unscaled
        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)

        # >= 1.0 because error_power is 1/scaled_error_norm
        error_power_prev = np.where(error_power >= 1.0, error_power, error_power_prev)

        dt_proposed = scale_factor * dt_proposed
        return PIState(dt_proposed, error_power_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)

correction_slr0(*, ssm, cubature_fun=cubature_third_order_spherical) -> _Correction

Zeroth-order statistical linear regression.

Source code in probdiffeq/ivpsolvers.py

def correction_slr0(*, ssm, cubature_fun=cubature_third_order_spherical) -> _Correction:
    """Zeroth-order statistical linear regression."""
    linearize = ssm.linearise.ode_statistical_1st(cubature_fun)
    return _CorrectionSLR(ssm=ssm, ode_order=1, linearize=linearize, name="SLR0")

correction_slr1(*, ssm, cubature_fun=cubature_third_order_spherical) -> _Correction

First-order statistical linear regression.

Source code in probdiffeq/ivpsolvers.py

def correction_slr1(*, ssm, cubature_fun=cubature_third_order_spherical) -> _Correction:
    """First-order statistical linear regression."""
    linearize = ssm.linearise.ode_statistical_0th(cubature_fun)
    return _CorrectionSLR(ssm=ssm, ode_order=1, linearize=linearize, name="SLR1")

correction_ts0(*, ssm, ode_order=1) -> _Correction

Zeroth-order Taylor linearisation.

Source code in probdiffeq/ivpsolvers.py

def correction_ts0(*, ssm, ode_order=1) -> _Correction:
    """Zeroth-order Taylor linearisation."""
    linearize = ssm.linearise.ode_taylor_0th(ode_order=ode_order)
    return _CorrectionTS(name="TS0", ode_order=ode_order, ssm=ssm, linearize=linearize)

correction_ts1(*, ssm, ode_order=1) -> _Correction

First-order Taylor linearisation.

Source code in probdiffeq/ivpsolvers.py

def correction_ts1(*, ssm, ode_order=1) -> _Correction:
    """First-order Taylor linearisation."""
    linearize = ssm.linearise.ode_taylor_1st(ode_order=ode_order)
    return _CorrectionTS(name="TS1", ode_order=ode_order, ssm=ssm, linearize=linearize)

cubature_gauss_hermite(input_shape, degree=5) -> _PositiveCubatureRule

(Statistician's) Gauss-Hermite cubature.

The number of cubature points is prod(input_shape)**degree.

Source code in probdiffeq/ivpsolvers.py

def cubature_gauss_hermite(input_shape, degree=5) -> _PositiveCubatureRule:
    """(Statistician's) Gauss-Hermite cubature.

    The number of cubature points is `prod(input_shape)**degree`.
    """
    assert len(input_shape) == 1
    (dim,) = input_shape

    # Roots of the probabilist/statistician's Hermite polynomials (in Numpy...)
    _roots = special.roots_hermitenorm(n=degree, mu=True)
    pts, weights, sum_of_weights = _roots
    weights = weights / sum_of_weights

    # Transform into jax arrays and take square root of weights
    pts = np.asarray(pts)
    weights_sqrtm = np.sqrt(np.asarray(weights))

    # Build a tensor grid and return class
    tensor_pts = _tensor_points(pts, d=dim)
    tensor_weights_sqrtm = _tensor_weights(weights_sqrtm, d=dim)
    return _PositiveCubatureRule(points=tensor_pts, weights_sqrtm=tensor_weights_sqrtm)

cubature_third_order_spherical(input_shape) -> _PositiveCubatureRule

Third-order spherical cubature integration.

Source code in probdiffeq/ivpsolvers.py

def cubature_third_order_spherical(input_shape) -> _PositiveCubatureRule:
    """Third-order spherical cubature integration."""
    assert len(input_shape) <= 1
    if len(input_shape) == 1:
        (d,) = input_shape
        points_mat, weights_sqrtm = _third_order_spherical_params(d=d)
        return _PositiveCubatureRule(points=points_mat, weights_sqrtm=weights_sqrtm)

    # If input_shape == (), compute weights via input_shape=(1,)
    # and 'squeeze' the points.
    points_mat, weights_sqrtm = _third_order_spherical_params(d=1)
    (S, _) = points_mat.shape
    points = np.reshape(points_mat, (S,))
    return _PositiveCubatureRule(points=points, weights_sqrtm=weights_sqrtm)

cubature_unscented_transform(input_shape, r=1.0) -> _PositiveCubatureRule

Unscented transform.

Source code in probdiffeq/ivpsolvers.py

def cubature_unscented_transform(input_shape, r=1.0) -> _PositiveCubatureRule:
    """Unscented transform."""
    assert len(input_shape) <= 1
    if len(input_shape) == 1:
        (d,) = input_shape
        points_mat, weights_sqrtm = _unscented_transform_params(d=d, r=r)
        return _PositiveCubatureRule(points=points_mat, weights_sqrtm=weights_sqrtm)

    # If input_shape == (), compute weights via input_shape=(1,)
    # and 'squeeze' the points.
    points_mat, weights_sqrtm = _unscented_transform_params(d=1, r=r)
    (S, _) = points_mat.shape
    points = np.reshape(points_mat, (S,))
    return _PositiveCubatureRule(points=points, weights_sqrtm=weights_sqrtm)

prior_ibm(tcoeffs, *, ssm_fact: str, output_scale=None)

Construct an adaptive(/continuous-time), multiply-integrated Wiener process.

Source code in probdiffeq/ivpsolvers.py

def prior_ibm(tcoeffs, *, ssm_fact: str, output_scale=None):
    """Construct an adaptive(/continuous-time), multiply-integrated Wiener process."""
    ssm = impl.choose(ssm_fact, tcoeffs_like=tcoeffs)

    output_scale_user = output_scale or np.ones_like(ssm.prototypes.output_scale())
    discretize = ssm.conditional.ibm_transitions(output_scale=output_scale_user)

    output_scale_calib = np.ones_like(ssm.prototypes.output_scale())
    prior = _MarkovProcess(tcoeffs, output_scale_calib, discretize=discretize)
    return prior, ssm

prior_ibm_discrete(ts, *, tcoeffs_like, ssm_fact: str, output_scale=None)

Compute a time-discretized, multiply-integrated Wiener process.

Source code in probdiffeq/ivpsolvers.py

def prior_ibm_discrete(ts, *, tcoeffs_like, ssm_fact: str, output_scale=None):
    """Compute a time-discretized, multiply-integrated Wiener process."""
    prior, ssm = prior_ibm(tcoeffs_like, output_scale=output_scale, ssm_fact=ssm_fact)
    transitions, (p, p_inv) = functools.vmap(prior.discretize)(np.diff(ts))

    preconditioner_apply_vmap = functools.vmap(ssm.conditional.preconditioner_apply)
    conditionals = preconditioner_apply_vmap(transitions, p, p_inv)

    output_scale = np.ones_like(ssm.prototypes.output_scale())
    init = ssm.normal.standard(len(tcoeffs_like), output_scale=output_scale)
    return stats.MarkovSeq(init, conditionals), ssm

solver(extrapolation, /, *, correction, prior, ssm)

Create a solver that does not calibrate the output scale automatically.

Source code in probdiffeq/ivpsolvers.py

def solver(extrapolation, /, *, correction, prior, ssm):
    """Create a solver that does not calibrate the output scale automatically."""

    def step(state: _State, *, vector_field, dt, calibration):
        del calibration  # unused

        prior_discretized = prior.discretize(dt)
        hidden, extra = extrapolation.begin(
            state.hidden, state.aux_extra, prior_discretized=prior_discretized
        )
        t = state.t + dt
        error, _, corr = correction.estimate_error(
            hidden, vector_field=vector_field, t=t
        )

        hidden, extra = extrapolation.complete(
            hidden, extra, output_scale=state.output_scale
        )
        hidden, corr = correction.complete(hidden, corr)

        # Extract and return solution
        state = _State(
            t=t, hidden=hidden, aux_extra=extra, output_scale=state.output_scale
        )
        return dt * error, state

    return _ProbabilisticSolver(
        ssm=ssm,
        prior=prior,
        extrapolation=extrapolation,
        correction=correction,
        calibration=_calibration_none(),
        step_implementation=step,
        name="Probabilistic solver",
        requires_rescaling=False,
    )

solver_dynamic(extrapolation, *, correction, prior, ssm)

Create a solver that calibrates the output scale dynamically.

Source code in probdiffeq/ivpsolvers.py

def solver_dynamic(extrapolation, *, correction, prior, ssm):
    """Create a solver that calibrates the output scale dynamically."""

    def step_dynamic(state, /, *, dt, vector_field, calibration):
        prior_discretized = prior.discretize(dt)
        hidden, extra = extrapolation.begin(
            state.hidden, state.aux_extra, prior_discretized=prior_discretized
        )
        t = state.t + dt
        error, observed, corr = correction.estimate_error(
            hidden, vector_field=vector_field, t=t
        )

        output_scale = calibration.update(state.output_scale, observed=observed)

        prior_, _calibrated = calibration.extract(output_scale)
        hidden, extra = extrapolation.complete(hidden, extra, output_scale=prior_)
        hidden, corr = correction.complete(hidden, corr)

        # Return solution
        state = _State(t=t, hidden=hidden, aux_extra=extra, output_scale=output_scale)
        return dt * error, state

    return _ProbabilisticSolver(
        prior=prior,
        ssm=ssm,
        extrapolation=extrapolation,
        correction=correction,
        calibration=_calibration_most_recent(ssm=ssm),
        name="Dynamic probabilistic solver",
        step_implementation=step_dynamic,
        requires_rescaling=False,
    )

solver_mle(extrapolation, /, *, correction, prior, ssm)

Create a solver that calibrates the output scale via maximum-likelihood.

Warning: needs to be combined with a call to stats.calibrate() after solving if the MLE-calibration shall be used.

Source code in probdiffeq/ivpsolvers.py

def solver_mle(extrapolation, /, *, correction, prior, ssm):
    """Create a solver that calibrates the output scale via maximum-likelihood.

    Warning: needs to be combined with a call to stats.calibrate()
    after solving if the MLE-calibration shall be *used*.
    """

    def step_mle(state, /, *, dt, vector_field, calibration):
        output_scale_prior, _calibrated = calibration.extract(state.output_scale)

        prior_discretized = prior.discretize(dt)
        hidden, extra = extrapolation.begin(
            state.hidden, state.aux_extra, prior_discretized=prior_discretized
        )
        t = state.t + dt
        error, _, corr = correction.estimate_error(
            hidden, vector_field=vector_field, t=t
        )

        hidden, extra = extrapolation.complete(
            hidden, extra, output_scale=output_scale_prior
        )
        hidden, observed = correction.complete(hidden, corr)

        output_scale = calibration.update(state.output_scale, observed=observed)
        state = _State(t=t, hidden=hidden, aux_extra=extra, output_scale=output_scale)
        return dt * error, state

    return _ProbabilisticSolver(
        ssm=ssm,
        name="Probabilistic solver with MLE calibration",
        prior=prior,
        calibration=_calibration_running_mean(ssm=ssm),
        step_implementation=step_mle,
        extrapolation=extrapolation,
        correction=correction,
        requires_rescaling=True,
    )

strategy_filter(*, ssm)

Construct a filter.

Source code in probdiffeq/ivpsolvers.py

def strategy_filter(*, ssm):
    """Construct a filter."""
    return _ExtraImplFilter(
        name="Filter",
        ssm=ssm,
        is_suitable_for_save_at=True,
        is_suitable_for_save_every_step=True,
        is_suitable_for_offgrid_marginals=True,
    )

strategy_fixedpoint(*, ssm)

Construct a fixedpoint-smoother.

Source code in probdiffeq/ivpsolvers.py

def strategy_fixedpoint(*, ssm):
    """Construct a fixedpoint-smoother."""
    return _ExtraImplFixedPoint(
        name="Fixed-point smoother",
        ssm=ssm,
        is_suitable_for_save_at=True,
        is_suitable_for_save_every_step=False,
        is_suitable_for_offgrid_marginals=False,
    )

strategy_smoother(*, ssm)

Construct a smoother.

Source code in probdiffeq/ivpsolvers.py

def strategy_smoother(*, ssm):
    """Construct a smoother."""
    return _ExtraImplSmoother(
        name="Smoother",
        ssm=ssm,
        is_suitable_for_save_at=False,
        is_suitable_for_save_every_step=True,
        is_suitable_for_offgrid_marginals=True,
    )