Solving ODEs with Probabilistic Numerics

In this tutorial we solve a simple non-linear ordinary differential equation (ODE) with the probabilistic numerical ODE solvers implemented in this package.

In this tutorial, we consider a Fitzhugh-Nagumo model described by an ODE of the form

\[\begin{aligned} \dot{y}_1 &= c (y_1 - \frac{y_1^3}{3} + y_2) \\ \dot{y}_2 &= -\frac{1}{c} (y_1 - a - b y_2) \end{aligned}\]

on a time span $t \in [0, T]$, with initial value $y(0) = y_0$. In the following, we

1. define the problem with explicit choices of initial values, integration domains, and parameters,
2. solve the problem with our ODE filters, and
3. visualize the results and the corresponding uncertainties.

TL;DR: Just use DifferentialEquations.jl with the EK1 algorithm

using ProbNumDiffEq, Plots

function fitz(du, u, p, t)
    a, b, c = p
    du[1] = c * (u[1] - u[1]^3 / 3 + u[2])
    du[2] = -(1 / c) * (u[1] - a - b * u[2])
end
u0 = [-1.0; 1.0]
tspan = (0.0, 20.0)
p = (0.2, 0.2, 3.0)
prob = ODEProblem(fitz, u0, tspan, p)

sol = solve(prob, EK1())
plot(sol)

Step 1: Define the problem

First, import ProbNumDiffEq.jl

using ProbNumDiffEq

Then, set up the ODEProblem exactly as you would in DifferentialEquations.jl. Define the vector field

function fitz(du, u, p, t)
    a, b, c = p
    du[1] = c * (u[1] - u[1]^3 / 3 + u[2])
    du[2] = -(1 / c) * (u[1] - a - b * u[2])
end

and then the ODEProblem, with initial value u0, time span tspan, and parameters p

u0 = [-1.0; 1.0]
tspan = (0.0, 20.0)
p = (0.2, 0.2, 3.0)
prob = ODEProblem(fitz, u0, tspan, p)

Step 2: Solve the problem

To solve the ODE we just use DifferentialEquations.jl's solve interface, together with one of the algorithms implemented in this package. For now, let's use EK1:

sol = solve(prob, EK1())

retcode: Success
Interpolation: ODE Filter Posterior
t: 267-element Vector{Float64}:
  0.0
  0.021276864853851562
  0.055300624266158616
  0.09069833746250563
  0.13926827172051978
  0.1848666961349729
  0.24179365422080168
  0.29051815606977655
  0.34907692582629263
  0.3957155206942215
  ⋮
 19.519319840846695
 19.56250961066827
 19.609301366711644
 19.659792228999418
 19.71548688397332
 19.777100417938893
 19.846195374734496
 19.924644734480783
 20.0
u: 267-element Vector{Vector{Float64}}:
 [-1.0, 1.0]
 [-0.9783978986608123, 1.0098599972789528]
 [-0.9424079102234552, 1.025330428495128]
 [-0.9028542932163617, 1.0410170742666094]
 [-0.844534969187771, 1.061811720582789]
 [-0.7847703402736025, 1.08049702917819]
 [-0.7018976845212279, 1.1025562952597345]
 [-0.6220689316716749, 1.1201790748695197]
 [-0.5125846285785938, 1.1395948389945416]
 [-0.41230241848807325, 1.1534759756253523]
 ⋮
 [2.083114409183183, 0.910019065899965]
 [2.079628891317149, 0.8855166799786748]
 [2.074332585529352, 0.8589602035604021]
 [2.0674863702504087, 0.8303140361447457]
 [2.0590832680762823, 0.7987456120282002]
 [2.049144182012271, 0.7638744453170948]
 [2.0374964872061074, 0.7248476806136662]
 [2.023853345952432, 0.6806503850812039]
 [2.0104414549959553, 0.638318942428594]

That's it! we just computed a probabilistic numerical ODE solution!

Step 3: Analyze the solution

Let's plot the result with Plots.jl.

using Plots
plot(sol)

Looks good! Looks like the EK1 managed to solve the Fitzhugh-Nagumo problem quite well.

Plot the probabilistic error estimates

The plot above looks like a standard ODE solution – but it's not! The numerical errors are just so small that we can't see them in the plot, and the probabilistic error estimates are too. We can visualize them by plotting the errors and error estimates directly:

using OrdinaryDiffEq, Statistics
reference = solve(prob, Vern9(), abstol=1e-9, reltol=1e-9, saveat=sol.t)
errors = reduce(hcat, mean.(sol.pu) .- reference.u)'
error_estimates = reduce(hcat, std.(sol.pu))'
plot(sol.t, errors, label="error", color=[1 2], xlabel="t", ylabel="err")
plot!(sol.t, zero(errors), ribbon=3error_estimates, label="error estimate",
      color=[1 2], alpha=0.2)

More about the ProbabilisticODESolution

The solution object returned by ProbNumDiffEq.jl mostly behaves just like any other ODESolution in DifferentialEquations.jl – with some added uncertainties and related functionality on top. The ProbabilisticODESolution can be indexed with

julia> sol.u[1]2-element Vector{Float64}:
 -1.0
  1.0
julia> sol.u[end]2-element Vector{Float64}:
 2.0104414549959553
 0.638318942428594
julia> sol.t[end]20.0

But since sol is a probabilistic numerical ODE solution, it contains a Gaussian distributions over solution values. The marginals of this posterior are stored in sol.pu:

julia> sol.pu[end]Gaussian{Vector{Float64},PSDMatrix{Float64, Matrix{Float64}}}(
  μ=[2.0104414549959553, 0.638318942428594],
  Σ=2x2 PSDMatrix{Float64, Matrix{Float64}}; R=[3.095013381245304e-5 9.585439719881619e-5; 4.849725631249043e-5 0.00015167434310597893; -2.6808472101866618e-5 -7.968208422450118e-5; -1.0847312089895692e-6 -2.0070824881007662e-5; 0.0 0.0; 0.0 0.0; 0.0 0.0; 0.0 0.0]
)

You can compute means, covariances, and standard deviations via Statistics.jl:

julia> using Statistics
julia> mean(sol.pu[5])2-element Vector{Float64}:
 -0.844534969187771
  1.061811720582789
julia> cov(sol.pu[5])2x2 PSDMatrix{Float64, Matrix{Float64}}
 Right square root: R=8×2 Matrix{Float64}:
 -2.8014e-6  -1.24815e-7
  0.0         2.74772e-6
  0.0         0.0
  0.0         0.0
  0.0         0.0
  0.0         0.0
  0.0         0.0
  0.0         0.0
julia> std(sol.pu[5])2-element Vector{Float64}:
 2.801399531863131e-6
 2.7505525858231918e-6

Dense output

Probabilistic numerical ODE solvers approximate the posterior distribution

\[p \Big( y(t) ~\big|~ y(0) = y_0, \{ \dot{y}(t_i) = f_\theta(y(t_i), t_i) \} \Big),\]

which describes a posterior not just for the discrete steps but for any $t$ in the continuous space $t \in [0, T]$; in classic ODE solvers, this is also known as "interpolation" or "dense output". The probabilistic solutions returned by our solvers can be interpolated as usual by treating them as functions, but they return Gaussian distributions

julia> sol(0.45)Gaussian{Vector{Float64},PSDMatrix{Float64, Matrix{Float64}}}(
  μ=[-0.27738212722925, 1.1675659399666511],
  Σ=2x2 PSDMatrix{Float64, Matrix{Float64}}; R=[-3.208354488497593e-5 -4.789969099963069e-6; 0.0 2.4206417679524815e-5; 0.0 0.0; 0.0 0.0; 0.0 0.0; 0.0 0.0; 0.0 0.0; 0.0 0.0]
)
julia> mean(sol(0.45))2-element Vector{Float64}:
 -0.27738212722925
  1.1675659399666511

Next steps

Check out one of the other tutorials:

• "Second Order ODEs and Energy Preservation" explains how to solve second-order ODEs more efficiently while also better preserving energy or other conserved quantities;
• "Solving DAEs with Probabilistic Numerics" demonstrates how to solve differential algebraic equations in a probabilistic numerical way.