B1. Visualise solution uncertainty¶

Probabilistic ODE solvers return a posterior distribution over trajectories, not just a single mean trajectory. This example plots the posterior mean together with uncertainty bands for the state and its first few derivatives.

Set up the ODE

In [1]:

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt

import jax import jax.numpy as jnp import matplotlib.pyplot as plt

In [2]:

from probdiffeq import ivpsolve, probdiffeq

from probdiffeq import ivpsolve, probdiffeq

In [3]:

# Fail this notebook on NaN detection (to catch those in the CI)
jax.config.update("jax_debug_nans", True)

# Fail this notebook on NaN detection (to catch those in the CI) jax.config.update("jax_debug_nans", True)

In [4]:

def main():
    """Plot means and standard deviations of solvers."""

    @probdiffeq.ode
    def vf(y, /, *, t):  # noqa: ARG001
        """Evaluate the Lotka-Volterra vector field."""
        y0, y1 = y[0], y[1]

        y0_new = 0.5 * y0 - 0.05 * y0 * y1
        y1_new = -0.5 * y1 + 0.05 * y0 * y1
        return jnp.asarray([y0_new, y1_new])

    t0 = 0.0
    t1 = 2.0
    u0 = jnp.asarray([20.0, 20.0])

    # Set up a solver.

    jetexpand = probdiffeq.jetexpand_ode_padded_scan(num=3)
    tcoeffs, _ = jetexpand(vf, (u0,), t=t0)
    ssm = probdiffeq.state_space_model_blockdiag()
    iwp = ssm.prior_wiener_integrated(tcoeffs)
    ts1 = ssm.constraint_ode_ts1(vf)
    strategy = probdiffeq.strategy_smoother_fixedpoint()

    solver = probdiffeq.solver_mle(strategy=strategy, constraint=ts1)
    error = probdiffeq.error_residual_std(constraint=ts1)
    solve = ivpsolve.solve_adaptive_save_at(solver=solver, error=error)

    # Solve the ODE.

    ts = jnp.linspace(t0, t1, endpoint=True, num=50)
    sol = jax.jit(solve)(iwp, save_at=ts, dt0=0.1, atol=1e-1, rtol=1e-1)

    # Plot the solution.

    fig, axes = plt.subplots(
        nrows=3,
        ncols=len(tcoeffs),
        sharex="col",
        tight_layout=True,
        figsize=(len(sol.u.mean) * 2, 5),
    )
    _titles = ["State", "1st deriv.", "2nd deriv.", "3rd deriv."]
    for i, (u_i, std_i, ax_i) in enumerate(zip(sol.u.mean, sol.u.std, axes.T)):
        # Set up titles and axis descriptions
        title = _titles[i] if i < len(_titles) else f"{i}th deriv."
        ax_i[0].set_title(title, fontsize="medium")
        if i == 0:
            ax_i[0].set_ylabel("Prey", fontsize="medium")
            ax_i[1].set_ylabel("Predators", fontsize="medium")
            ax_i[2].set_ylabel("Std.-dev.", fontsize="medium")

        ax_i[-1].set_xlabel("Time", fontsize="medium")

        for m, std, ax in zip(u_i.T, std_i.T, ax_i):
            # Plot the mean
            ax.plot(sol.t, m)

            # Plot the standard deviation
            lower, upper = m - 3 * std, m + 3 * std
            ax.fill_between(sol.t, lower, upper, alpha=0.3)
            ax.set_xlim((jnp.amin(ts), jnp.amax(ts)))

        ax_i[2].semilogy(sol.t, std_i[:, 0], label="Prey")
        ax_i[2].semilogy(sol.t, std_i[:, 1], label="Predators")
        ax_i[2].legend(fontsize="x-small")

    fig.align_ylabels()
    plt.show()

def main(): """Plot means and standard deviations of solvers.""" @probdiffeq.ode def vf(y, /, *, t): # noqa: ARG001 """Evaluate the Lotka-Volterra vector field.""" y0, y1 = y[0], y[1] y0_new = 0.5 * y0 - 0.05 * y0 * y1 y1_new = -0.5 * y1 + 0.05 * y0 * y1 return jnp.asarray([y0_new, y1_new]) t0 = 0.0 t1 = 2.0 u0 = jnp.asarray([20.0, 20.0]) # Set up a solver. jetexpand = probdiffeq.jetexpand_ode_padded_scan(num=3) tcoeffs, _ = jetexpand(vf, (u0,), t=t0) ssm = probdiffeq.state_space_model_blockdiag() iwp = ssm.prior_wiener_integrated(tcoeffs) ts1 = ssm.constraint_ode_ts1(vf) strategy = probdiffeq.strategy_smoother_fixedpoint() solver = probdiffeq.solver_mle(strategy=strategy, constraint=ts1) error = probdiffeq.error_residual_std(constraint=ts1) solve = ivpsolve.solve_adaptive_save_at(solver=solver, error=error) # Solve the ODE. ts = jnp.linspace(t0, t1, endpoint=True, num=50) sol = jax.jit(solve)(iwp, save_at=ts, dt0=0.1, atol=1e-1, rtol=1e-1) # Plot the solution. fig, axes = plt.subplots( nrows=3, ncols=len(tcoeffs), sharex="col", tight_layout=True, figsize=(len(sol.u.mean) * 2, 5), ) _titles = ["State", "1st deriv.", "2nd deriv.", "3rd deriv."] for i, (u_i, std_i, ax_i) in enumerate(zip(sol.u.mean, sol.u.std, axes.T)): # Set up titles and axis descriptions title = _titles[i] if i < len(_titles) else f"{i}th deriv." ax_i[0].set_title(title, fontsize="medium") if i == 0: ax_i[0].set_ylabel("Prey", fontsize="medium") ax_i[1].set_ylabel("Predators", fontsize="medium") ax_i[2].set_ylabel("Std.-dev.", fontsize="medium") ax_i[-1].set_xlabel("Time", fontsize="medium") for m, std, ax in zip(u_i.T, std_i.T, ax_i): # Plot the mean ax.plot(sol.t, m) # Plot the standard deviation lower, upper = m - 3 * std, m + 3 * std ax.fill_between(sol.t, lower, upper, alpha=0.3) ax.set_xlim((jnp.amin(ts), jnp.amax(ts))) ax_i[2].semilogy(sol.t, std_i[:, 0], label="Prey") ax_i[2].semilogy(sol.t, std_i[:, 1], label="Predators") ax_i[2].legend(fontsize="x-small") fig.align_ylabels() plt.show()

In [5]:

if __name__ == "__main__":
    main()

if __name__ == "__main__": main()