ivpsolvers

Probabilistic IVP solvers.

adaptive(slvr, /, *, ssm, atol=0.0001, rtol=0.01, control=None, norm_ord=None, clip_dt: bool = False, eps: float | None = 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,
    clip_dt: bool = False,
    eps: float | None = None,
):
    """Make an IVP solver adaptive."""
    if control is None:
        control = control_proportional_integral()
    if eps is None:
        eps = 10 * np.finfo_eps(float)
    return _AdaSolver(
        slvr,
        ssm=ssm,
        atol=atol,
        rtol=rtol,
        control=control,
        norm_ord=norm_ord,
        clip_dt=clip_dt,
        eps=eps,
    )

control_integral(*, safety=0.95, factor_min=0.2, factor_max=10.0) -> _Controller[None]

Construct an integral-controller.

Source code in probdiffeq/ivpsolvers.py

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

    def init(_dt, /) -> None:
        return None

    def apply(dt, _state, /, *, 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, None

    return _Controller(init=init, apply=apply)

control_proportional_integral(*, safety=0.95, factor_min=0.2, factor_max=10.0, power_integral_unscaled=0.3, power_proportional_unscaled=0.4) -> _Controller[float]

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

Source code in probdiffeq/ivpsolvers.py

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

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

    def apply(dt: float, error_power_prev: float, /, *, error_power):
        # Equivalent: error_power = error_norm ** (-1.0 / error_contraction_rate)
        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
        return dt_proposed, error_power_prev

    return _Controller(init=init, apply=apply)

correction_slr0(vector_field, *, ssm, cubature_fun=cubature_third_order_spherical, damp: float = 0.0) -> _Correction

Zeroth-order statistical linear regression.

Source code in probdiffeq/ivpsolvers.py

def correction_slr0(
    vector_field, *, ssm, cubature_fun=cubature_third_order_spherical, damp: float = 0.0
) -> _Correction:
    """Zeroth-order statistical linear regression."""
    linearize = ssm.linearise.ode_statistical_0th(cubature_fun, damp=damp)
    return _Correction(
        ssm=ssm,
        vector_field=vector_field,
        ode_order=1,
        linearize=linearize,
        name="SLR0",
        re_linearize=True,
    )

correction_slr1(vector_field, *, ssm, cubature_fun=cubature_third_order_spherical, damp: float = 0.0) -> _Correction

First-order statistical linear regression.

Source code in probdiffeq/ivpsolvers.py

def correction_slr1(
    vector_field, *, ssm, cubature_fun=cubature_third_order_spherical, damp: float = 0.0
) -> _Correction:
    """First-order statistical linear regression."""
    linearize = ssm.linearise.ode_statistical_1st(cubature_fun, damp=damp)
    return _Correction(
        ssm=ssm,
        vector_field=vector_field,
        ode_order=1,
        linearize=linearize,
        name="SLR1",
        re_linearize=True,
    )

correction_ts0(vector_field, *, ssm, ode_order=1, damp: float = 0.0) -> _Correction

Zeroth-order Taylor linearisation.

Source code in probdiffeq/ivpsolvers.py

def correction_ts0(vector_field, *, ssm, ode_order=1, damp: float = 0.0) -> _Correction:
    """Zeroth-order Taylor linearisation."""
    linearize = ssm.linearise.ode_taylor_0th(ode_order=ode_order, damp=damp)
    return _Correction(
        name="TS0",
        vector_field=vector_field,
        ode_order=ode_order,
        ssm=ssm,
        linearize=linearize,
        re_linearize=False,
    )

correction_ts1(vector_field, *, ssm, ode_order=1, damp: float = 0.0, jvp_probes=10, jvp_probes_seed=1) -> _Correction

First-order Taylor linearisation.

Source code in probdiffeq/ivpsolvers.py

def correction_ts1(
    vector_field,
    *,
    ssm,
    ode_order=1,
    damp: float = 0.0,
    jvp_probes=10,
    jvp_probes_seed=1,
) -> _Correction:
    """First-order Taylor linearisation."""
    assert jvp_probes > 0
    linearize = ssm.linearise.ode_taylor_1st(
        ode_order=ode_order,
        damp=damp,
        jvp_probes=jvp_probes,
        jvp_probes_seed=jvp_probes_seed,
    )
    return _Correction(
        name="TS1",
        vector_field=vector_field,
        ode_order=ode_order,
        ssm=ssm,
        linearize=linearize,
        re_linearize=False,
    )

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_wiener_integrated(tcoeffs, *, ssm_fact: str, output_scale: ArrayLike | None = None, damp: float = 0.0)

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

Source code in probdiffeq/ivpsolvers.py

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

    # TODO: should the output_scale be an argument to solve()?
    # TODO: should the output scale (and all 'damp'-like factors)
    #       mirror the pytree structure of 'tcoeffs'?
    if output_scale is None:
        output_scale = np.ones_like(ssm.prototypes.output_scale())

    discretize = ssm.conditional.ibm_transitions(base_scale=output_scale)

    # Increase damping to get visually more pleasing uncertainties
    #  and more numerical robustness for
    #  high-order solvers in low precision arithmetic
    init = ssm.normal.from_tcoeffs(tcoeffs, damp=damp)
    return init, discretize, ssm

prior_wiener_integrated_discrete(ts, *args, **kwargs)

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

Source code in probdiffeq/ivpsolvers.py

def prior_wiener_integrated_discrete(ts, *args, **kwargs):
    """Compute a time-discretized, multiply-integrated Wiener process."""
    init, discretize, ssm = prior_wiener_integrated(*args, **kwargs)
    scales = np.ones_like(ssm.prototypes.output_scale())
    discretize_vmap = functools.vmap(discretize, in_axes=(0, None))
    conditionals = discretize_vmap(np.diff(ts), scales)
    return init, conditionals, ssm

solver(strategy, *, correction, prior, ssm)

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

Source code in probdiffeq/ivpsolvers.py

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

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

        u_step_from = tree_util.ravel_pytree(ssm.unravel(state.rv.mean)[0])[0]

        # Estimate the error
        transition = prior(dt, state.output_scale)
        mean = ssm.stats.mean(state.rv)
        hidden = ssm.conditional.apply(mean, transition)
        t = state.t + dt
        error, _, correction_state = correction.estimate_error(
            hidden, state.correction_state, t=t
        )

        # Do the full extrapolation step (reuse the transition)
        hidden, extra = strategy.extrapolate(
            state.rv, state.strategy_state, transition=transition
        )

        # Do the full correction step
        hidden, _, correction_state = correction.correct(hidden, correction_state, t=t)
        state = _State(
            t=t,
            rv=hidden,
            strategy_state=extra,
            correction_state=correction_state,
            output_scale=state.output_scale,
        )

        # Normalise the error
        u_proposed = tree_util.ravel_pytree(ssm.unravel(state.rv.mean)[0])[0]
        reference = np.maximum(np.abs(u_proposed), np.abs(u_step_from))
        error = _ErrorEstimate(dt * error, reference=reference)
        return error, state

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

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

Create a solver that calibrates the output scale dynamically.

Source code in probdiffeq/ivpsolvers.py

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

    def step_dynamic(state, /, *, dt, calibration):
        u_step_from = tree_util.ravel_pytree(ssm.unravel(state.rv.mean)[0])[0]

        # Estimate error and calibrate the output scale
        ones = np.ones_like(ssm.prototypes.output_scale())
        transition = prior(dt, ones)
        mean = ssm.stats.mean(state.rv)
        hidden = ssm.conditional.apply(mean, transition)

        t = state.t + dt
        error, observed, correction_state = correction.estimate_error(
            hidden, state.correction_state, t=t
        )
        output_scale = calibration.update(state.output_scale, observed=observed)

        # Do the full extrapolation with the calibrated output scale
        scale, _ = calibration.extract(output_scale)
        transition = prior(dt, scale)
        hidden, extra = strategy.extrapolate(
            state.rv, state.strategy_state, transition=transition
        )

        # Do the full correction step
        hidden, _, correction_state = correction.correct(hidden, correction_state, t=t)

        # Return solution
        state = _State(
            t=t,
            rv=hidden,
            strategy_state=extra,
            correction_state=correction_state,
            output_scale=output_scale,
        )

        # Normalise the error
        u_proposed = tree_util.ravel_pytree(ssm.unravel(state.rv.mean)[0])[0]
        reference = np.maximum(np.abs(u_proposed), np.abs(u_step_from))
        error = _ErrorEstimate(dt * error, reference=reference)
        return error, state

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

solver_mle(strategy, *, 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(strategy, *, 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, calibration):
        u_step_from = tree_util.ravel_pytree(ssm.unravel(state.rv.mean)[0])[0]

        # Estimate the error
        output_scale_prior, _calibrated = calibration.extract(state.output_scale)
        transition = prior(dt, output_scale_prior)
        mean = ssm.stats.mean(state.rv)
        mean_extra = ssm.conditional.apply(mean, transition)
        t = state.t + dt
        error, _, correction_state = correction.estimate_error(
            mean_extra, state.correction_state, t=t
        )

        # Do the full prediction step (reuse previous discretisation)
        hidden, extra = strategy.extrapolate(
            state.rv, state.strategy_state, transition=transition
        )

        # Do the full correction step
        hidden, observed, corr_state = correction.correct(hidden, correction_state, t=t)

        # Calibrate the output scale
        output_scale = calibration.update(state.output_scale, observed=observed)

        # Normalise the error

        state = _State(
            t=t,
            rv=hidden,
            strategy_state=extra,
            correction_state=corr_state,
            output_scale=output_scale,
        )
        u_proposed = tree_util.ravel_pytree(ssm.unravel(state.rv.mean)[0])[0]
        reference = np.maximum(np.abs(u_proposed), np.abs(u_step_from))
        error = _ErrorEstimate(dt * error, reference=reference)
        return error, state

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

strategy_filter(*, ssm) -> _Strategy

Construct a filter.

Source code in probdiffeq/ivpsolvers.py

def strategy_filter(*, ssm) -> _Strategy:
    """Construct a filter."""

    @containers.dataclass
    class Filter(_Strategy):
        def init(self, sol, /):
            return sol, None

        def extrapolate(self, rv, aux, /, *, transition):
            del aux
            rv = self.ssm.conditional.marginalise(rv, transition)

            return rv, None

        def extract(self, hidden_state, _extra, /):
            return hidden_state

        def interpolate(self, state_t0, state_t1, dt0, dt1, output_scale, *, prior):
            # todo: by ditching marginal_t1 and dt1, this function _extrapolates
            #  (no *inter*polation happening)
            del dt1
            marginal_t1, _ = state_t1

            hidden, extra = state_t0
            prior0 = prior(dt0, output_scale)
            hidden, extra = self.extrapolate(hidden, extra, transition=prior0)

            # Consistent state-types in interpolation result.
            interp = (hidden, extra)
            step_from = (marginal_t1, None)
            return _InterpRes(
                step_from=step_from, interpolated=interp, interp_from=interp
            )

        def interpolate_at_t1(
            self, state_t0, state_t1, dt0, dt1, output_scale, *, prior
        ):
            del prior
            del state_t0
            del dt0
            del dt1
            del output_scale
            rv, extra = state_t1
            return _InterpRes((rv, extra), (rv, extra), (rv, extra))

    return 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) -> _Strategy

Construct a fixedpoint-smoother.

Source code in probdiffeq/ivpsolvers.py

def strategy_fixedpoint(*, ssm) -> _Strategy:
    """Construct a fixedpoint-smoother."""

    @containers.dataclass
    class FixedPoint(_Strategy):
        def init(self, sol, /):
            cond = self.ssm.conditional.identity(ssm.num_derivatives + 1)
            return sol, cond

        def extrapolate(self, rv, bw0, /, *, transition):
            extrapolated, cond = self.ssm.conditional.revert(rv, transition)
            cond = self.ssm.conditional.merge(bw0, cond)
            return extrapolated, cond

        def extract(self, hidden_state, extra, /):
            return stats.MarkovSeq(init=hidden_state, conditional=extra)

        def interpolate_at_t1(
            self, state_t0, state_t1, *, dt0, dt1, output_scale, prior
        ):
            del prior
            del state_t0
            del dt0
            del dt1
            del output_scale
            rv, extra = state_t1
            cond_identity = self.ssm.conditional.identity(ssm.num_derivatives + 1)
            return _InterpRes((rv, cond_identity), (rv, extra), (rv, cond_identity))

        def interpolate(self, state_t0, state_t1, *, dt0, dt1, output_scale, prior):
            """Interpolate.

            A fixed-point smoother interpolates by

            * Extrapolating from t0 to t, which gives the "filtering" marginal
            and the backward transition from t to t0.
            * Extrapolating from t to t1, which gives another "filtering" marginal
            and the backward transition from t1 to t.
            * Applying the t1-to-t backward transition
            to compute the interpolation result.
            This intermediate result is informed about its "right-hand side" datum.

            The difference to smoother-interpolation is quite subtle:

            * The backward transition of the solution at 't'
            is merged with that at 't0'.
            The reason is that the backward transition at 't0' knows
            "how to get to the quantity of interest",
            and this is precisely what we want to interpolate.
            * Subsequent interpolations do not continue from the value at 't', but
            from a very similar value where the backward transition
            is replaced with an identity. The reason is that the interpolated solution
            becomes the new quantity of interest, and subsequent interpolations
            need to learn how to get here.
            * Subsequent solver steps do not continue from the value at 't1',
            but the value at 't1' where the backward model is replaced by
            the 't1-to-t' backward model. The reason is similar to the above:
            future steps need to know "how to get back to the quantity of interest",
            which is the interpolated solution.

            These distinctions are precisely why we need three fields
            in every interpolation result:
                the solution,
                the continue-interpolation-from-here,
                and the continue-stepping-from-here.
            All three are different for fixed point smoothers.
            (Really, I try removing one of them monthly and
            then don't understand why tests fail.)
            """
            marginal_t1, _ = state_t1
            # Extrapolate from t0 to t, and from t to t1.
            # This yields all building blocks.
            prior0 = prior(dt0, output_scale)
            extrapolated_t = self.extrapolate(*state_t0, transition=prior0)
            conditional_id = self.ssm.conditional.identity(ssm.num_derivatives + 1)
            previous_new = (extrapolated_t[0], conditional_id)

            prior1 = prior(dt1, output_scale)
            extrapolated_t1 = self.extrapolate(*previous_new, transition=prior1)

            # Marginalise from t1 to t to obtain the interpolated solution.
            conditional_t1_to_t = extrapolated_t1[1]
            rv_at_t = self.ssm.conditional.marginalise(marginal_t1, conditional_t1_to_t)

            # Return the right combination of marginals and conditionals.
            return _InterpRes(
                step_from=(marginal_t1, conditional_t1_to_t),
                interpolated=(rv_at_t, extrapolated_t[1]),
                interp_from=previous_new,
            )

    return FixedPoint(
        ssm=ssm,
        is_suitable_for_save_at=True,
        is_suitable_for_save_every_step=False,
        is_suitable_for_offgrid_marginals=False,
    )

strategy_smoother(*, ssm) -> _Strategy

Construct a smoother.

Source code in probdiffeq/ivpsolvers.py

def strategy_smoother(*, ssm) -> _Strategy:
    """Construct a smoother."""

    @containers.dataclass
    class Smoother(_Strategy):
        def init(self, sol, /):
            # Special case for implementing offgrid-marginals...
            if isinstance(sol, stats.MarkovSeq):
                rv = sol.init
                cond = sol.conditional
            else:
                rv = sol
                cond = self.ssm.conditional.identity(ssm.num_derivatives + 1)
            return rv, cond

        def extrapolate(self, rv, aux, /, *, transition):
            del aux
            return self.ssm.conditional.revert(rv, transition)

        def extract(self, hidden_state, extra, /):
            return stats.MarkovSeq(init=hidden_state, conditional=extra)

        def interpolate(self, state_t0, state_t1, *, dt0, dt1, output_scale, prior):
            """Interpolate.

            A smoother interpolates by_
            * Extrapolating from t0 to t, which gives the "filtering" marginal
            and the backward transition from t to t0.
            * Extrapolating from t to t1, which gives another "filtering" marginal
            and the backward transition from t1 to t.
            * Applying the new t1-to-t backward transition to compute the interpolation.
            This intermediate result is informed about its "right-hand side" datum.

            Subsequent interpolations continue from the value at 't'.
            Subsequent IVP solver steps continue from the value at 't1'.
            """
            # TODO: if we pass prior1 and prior2, then
            #       we don't have to pass dt0, dt1, output_scale, and prior...

            # Extrapolate from t0 to t, and from t to t1.
            prior0 = prior(dt0, output_scale)
            extrapolated_t = self.extrapolate(*state_t0, transition=prior0)
            prior1 = prior(dt1, output_scale)
            extrapolated_t1 = self.extrapolate(*extrapolated_t, transition=prior1)

            # Marginalise from t1 to t to obtain the interpolated solution.
            marginal_t1, _ = state_t1
            conditional_t1_to_t = extrapolated_t1[1]
            rv_at_t = self.ssm.conditional.marginalise(marginal_t1, conditional_t1_to_t)
            solution_at_t = (rv_at_t, extrapolated_t[1])

            # The state at t1 gets a new backward model;
            # (it must remember how to get back to t, not to t0).
            solution_at_t1 = (marginal_t1, conditional_t1_to_t)
            return _InterpRes(
                step_from=solution_at_t1,
                interpolated=solution_at_t,
                interp_from=solution_at_t,
            )

        def interpolate_at_t1(
            self, state_t0, state_t1, *, dt0, dt1, output_scale, prior
        ):
            del prior
            del state_t0
            del dt0
            del dt1
            del output_scale
            return _InterpRes(state_t1, state_t1, state_t1)

    return Smoother(
        ssm=ssm,
        is_suitable_for_save_at=False,
        is_suitable_for_save_every_step=True,
        is_suitable_for_offgrid_marginals=True,
    )