Solving DAEs with Probabilistic Numerics

ProbNumDiffEq.jl provides probabilistic numerical solvers for differential algebraic equations (DAEs). Currently, we recommend using the semi-implicit EK1 algorithm.

Note

For a more general tutorial on DAEs check out the DifferentialEquations.jl DAE tutorial.

Solving mass-matrix DAEs with the EK1

First, define the DAE (here the ROBER problem) as an ODE problem with singular mass matrix:

using ProbNumDiffEq, Plots, LinearAlgebra, OrdinaryDiffEq, ModelingToolkit, LinearAlgebra

function rober(du, u, p, t)
    y₁, y₂, y₃ = u
    k₁, k₂, k₃ = p
    du[1] = -k₁ * y₁ + k₃ * y₂ * y₃
    du[2] = k₁ * y₁ - k₃ * y₂ * y₃ - k₂ * y₂^2
    du[3] = y₁ + y₂ + y₃ - 1
    nothing
end
M = [1 0 0
     0 1 0
     0 0 0]
f = ODEFunction(rober, mass_matrix=M)
prob_mm = ODEProblem(f, [1.0, 0.0, 0.0], (0.0, 1e5), (0.04, 3e7, 1e4))
ODEProblem with uType Vector{Float64} and tType Float64. In-place: true
timespan: (0.0, 100000.0)
u0: 3-element Vector{Float64}:
 1.0
 0.0
 0.0

We can solve this problem directly with the EK1:

sol = solve(prob_mm, EK1(), reltol=1e-8, abstol=1e-8)
plot(
    sol,
    xscale=:log10,
    tspan=(1e-6, 1e5),
    layout=(3, 1),
    legend=false,
    ylabel=["u₁(t)" "u₂(t)" "u₃(t)"],
    xlabel=["" "" "t"],
    denseplot=false,
)
Example block output

Looks good!

Solving an Index-3 DAE directly

The following is based on the "Automatic Index Reduction of DAEs" tutorial by ModelingToolkit.jl, which demonstrates how the classic Rodas4 solver fails to solve a DAE due to the fact that it is of index 3; which is why ModelingToolkit's automatic index reduction is so useful.

It turns out that our probabilistic numerical solvers can directly solve the index-3 DAE!

First, define the pendulum problem as in the tutorial:

function pendulum!(du, u, p, t)
    x, dx, y, dy, T = u
    g, L = p
    du[1] = dx
    du[2] = T * x
    du[3] = dy
    du[4] = T * y - g
    du[5] = x^2 + y^2 - L^2
end
pendulum_fun! = ODEFunction(pendulum!, mass_matrix=Diagonal([1, 1, 1, 1, 0]))
u0 = [1.0, 0, 0, 0, 0];
p = [9.8, 1];
tspan = (0, 5.0);
pendulum_prob = ODEProblem(pendulum_fun!, u0, tspan, p)
ODEProblem with uType Vector{Float64} and tType Float64. In-place: true
timespan: (0.0, 5.0)
u0: 5-element Vector{Float64}:
 1.0
 0.0
 0.0
 0.0
 0.0

We can try to solve it directly with one of the classic mass-matrix DAE solvers from OrdinaryDiffEq.jl:

solve(pendulum_prob, Rodas4())
retcode: Unstable
Interpolation: specialized 3rd order "free" stiffness-aware interpolation
t: 5-element Vector{Float64}:
 0.0
 1.0e-6
 1.1e-5
 3.806859510273279e-5
 4.367846401990628e-5
u: 5-element Vector{Vector{Float64}}:
 [1.0, 0.0, 0.0, 0.0, 0.0]
 [1.0, 0.0, -4.899999999999997e-12, -9.799999999999994e-6, 0.0]
 [1.0, -5.282199999999992e-16, -5.928999999999995e-10, -0.00010779999999999988, -9.04785641150889e-11]
 [1.0, -2.2119452635682024e-13, -7.1011678721694636e-9, -0.000373072232006781, -2.6515338160403213e-8]
 [1.0, -3.0396474329579032e-12, -9.348260273777398e-9, -0.00042804894739508116, -3.896506115963402e-6]

It does not work! This is because of the index of the DAE; see for example this explanation from the tutorial.

Does this also hold for the EK1 solver? Let's find out:

sol = solve(pendulum_prob, EK1())
retcode: Success
Interpolation: ODE Filter Posterior
t: 650-element Vector{Float64}:
 0.0
 1.0e-6
 1.1e-5
 6.293457400357088e-5
 0.0001717966984100629
 0.00041700371867394907
 0.0010432618632924343
 0.0023836933769503994
 0.00376397894518298
 0.00519132985680733
 ⋮
 4.907142111362168
 4.92444553498383
 4.9411116199936265
 4.956803287770239
 4.97098674548783
 4.9822968246165
 4.99143218593388
 4.998834745818967
 5.0
u: 650-element Vector{Vector{Float64}}:
 [1.0, 0.0, 0.0, 0.0, 0.0]
 [1.0, -2.4314864246071144e-18, -4.899999999937332e-12, -9.800000000000003e-6, -3.870090693653406e-16]
 [1.0, 1.778914640672942e-14, -5.929000004468074e-10, -0.00010780000000000091, -3.017617371928309e-12]
 [0.9999999999999998, -3.704393921873193e-12, -1.940772692127836e-8, -0.0006167588252348035, -1.3035790293759336e-9]
 [0.9999999999999896, 1.2314212502343228e-11, -1.4461911744623022e-7, -0.0016836076444193238, -3.157756035951345e-8]
 [0.9999999999996368, -9.146450513208143e-11, -8.520712963690976e-7, -0.004086636443002532, -4.845250915236145e-7]
 [0.9999999999857788, -1.8913629380116342e-9, -5.333137047544597e-6, -0.01022396626026856, -7.830195678929709e-6]
 [0.9999999996124179, -5.727074650413097e-8, -2.7841771139450103e-5, -0.023360195092905718, -9.453725069261283e-5]
 [0.9999999975903671, -3.4756550064331693e-7, -6.942093370364985e-5, -0.03688699364700587, -0.00037351474250189485]
 [0.9999999912808006, -1.274813756871893e-6, -0.0001320545364564354, -0.05087503248488846, -0.0009815791650569213]
 ⋮
 [0.9902628022671714, -0.230012218970361, -0.13921057057397462, -1.6366422719422729, -4.0945422713871436]
 [0.9856333315526842, -0.3073677975051621, -0.16889919052551622, -1.7941346919319325, -4.968908564858052]
 [0.9797897702820494, -0.39617679785706794, -0.20003001727924735, -1.9407717381750471, -5.888064423522993]
 [0.972827745793681, -0.49345240804801094, -0.23153008794499258, -2.073030808278413, -6.821514288820706]
 [0.9651366055329357, -0.593329573335952, -0.2617467158743893, -2.186589633729141, -7.72111506539165]
 [0.957940002319977, -0.6813737682784675, -0.28696857832357914, -2.2722969869869667, -8.475874423432991]
 [0.9513764315907425, -0.7580420941403236, -0.3080306846666298, -2.337925795788078, -9.10937542703355]
 [0.9455309355620407, -0.8238465818165532, -0.325532285560256, -2.388489111563977, -9.638156836814543]
 [0.944566419026587, -0.8345069039015303, -0.3283203924231612, -2.3962208016104234, -9.722642388563708]

Nope! The EK1 is able to solve the index-3 DAE directly. Pretty cool!

plot(sol)
Example block output

Is index-reduction still worth it?

The point of the "Automatic Index Reduction of DAEs" tutorial is to demonstrate ModelingToolkit's utility for automatic index reduction, which enables the classic implicit Runge-Kutta solvers such as Rodas5 to solve this DAE. Let's see if that still helps in this context here.

First, modelingtoolkitize the problem:

traced_sys = modelingtoolkitize(pendulum_prob)

\[ \begin{align} \frac{\mathrm{d} x_1\left( t \right)}{\mathrm{d}t} &= x_2\left( t \right) \\ \frac{\mathrm{d} x_2\left( t \right)}{\mathrm{d}t} &= x_1\left( t \right) x_5\left( t \right) \\ \frac{\mathrm{d} x_3\left( t \right)}{\mathrm{d}t} &= x_4\left( t \right) \\ \frac{\mathrm{d} x_4\left( t \right)}{\mathrm{d}t} &= - \alpha_1 + x_3\left( t \right) x_5\left( t \right) \\ 0 &= \left( x_1\left( t \right) \right)^{2} + \left( x_3\left( t \right) \right)^{2} - \alpha_2^{2} \end{align} \]

(how cool is this latex output ?!?)

Next, lower the DAE index and simplify it with ModelingToolkit's dae_index_lowering and structural_simplify:

simplified_sys = structural_simplify(dae_index_lowering(traced_sys))

\[ \begin{align} \frac{\mathrm{d} x_1\left( t \right)}{\mathrm{d}t} &= x_2\left( t \right) \\ \frac{\mathrm{d} x_2\left( t \right)}{\mathrm{d}t} &= x_1\left( t \right) x_5\left( t \right) \\ \frac{\mathrm{d} x_3\left( t \right)}{\mathrm{d}t} &= x_4\left( t \right) \\ \frac{\mathrm{d} x_4\left( t \right)}{\mathrm{d}t} &= - \alpha_1 + x_3\left( t \right) x_5\left( t \right) \\ 0 &= 2 \left( x_4\left( t \right) \right)^{2} + 2 \left( x_2\left( t \right) \right)^{2} + 2 \left( x_1\left( t \right) \right)^{2} x_5\left( t \right) + 2 x_3\left( t \right) \left( - \alpha_1 + x_3\left( t \right) x_5\left( t \right) \right) \end{align} \]

Let's build two different ODE problems, and check how well we can solve each:

prob_index3 = ODEProblem(complete(traced_sys), Pair[], tspan)
prob_index1 = ODEProblem(simplified_sys, Pair[], tspan)

sol3 = solve(prob_index3, EK1())
sol1 = solve(prob_index1, EK1())

truesol = solve(prob_index1, Rodas4(), abstol=1e-10, reltol=1e-10)

sol1_final_error = norm(sol1.u[end] - truesol.u[end])
sol1_f_evals     = sol1.stats.nf
sol3_final_error = norm(sol3.u[end] - truesol.u[end])
sol3_f_evals     = sol3.stats.nf
@info "Results" sol1_final_error sol1_f_evals sol3_final_error sol3_f_evals
┌ Info: Results
│   sol1_final_error = 0.010930749098469738
│   sol1_f_evals = 1085
│   sol3_final_error = 0.06122741667364596
└   sol3_f_evals = 1409

The error for the index-1 DAE solve is much lower. Thus it seems that, even if the index-3 DAE could also be solved directly, index lowering might still be beneficial when solving DAEs with the EK1!

References

[1]
N. Bosch, F. Tronarp and P. Hennig. Pick-and-Mix Information Operators for Probabilistic ODE Solvers. In: Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, Vol. 151 of Proceedings of Machine Learning Research, edited by G. Camps-Valls, F. J. Ruiz and I. Valera (PMLR, 28–30 Mar 2022); pp. 10015–10027.