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

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: MaxIters
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.549693054277249e-5
u: 5-element Vector{Vector{Float64}}:
 [1.0, 0.0, 0.0, 0.0, 0.0]
 [1.0, 0.0, -4.8999999999999965e-12, -9.799999999999993e-6, 0.0]
 [1.0, -5.282199999999992e-16, -5.928999999999995e-10, -0.00010779999999999991, -9.047856411508895e-11]
 [1.0, -2.2119452635682026e-13, -7.101167872169468e-9, -0.00037307223200678104, -2.6515338160403223e-8]
 [1.0, -3.1661971462360842e-12, -1.0142856375187924e-8, -0.00044586991931917005, -2.9596722490873632e-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.907726151506895
 4.925025065157816
 4.941663273699247
 4.9573070968779565
 4.97143351122974
 4.982674540352456
 4.991750117206177
 4.999103080965005
 5.0
u: 650-element Vector{Vector{Float64}}:
 [1.0, 0.0, 0.0, 0.0, 0.0]
 [1.0, -2.431486734799844e-18, -4.89999999993733e-12, -9.800000000000003e-6, -3.8700910333367305e-16]
 [1.0, 1.7789148841949144e-14, -5.929000004468161e-10, -0.00010780000000000091, -3.0176176373167573e-12]
 [0.9999999999999998, -3.704394348777798e-12, -1.9407726921277446e-8, -0.000616758825234804, -1.3035791445211927e-9]
 [0.9999999999999896, 1.2314213816938362e-11, -1.4461911744623226e-7, -0.0016836076444193256, -3.157756315501106e-8]
 [0.9999999999996368, -9.146451418516957e-11, -8.520712963690889e-7, -0.004086636443002533, -4.845251344776e-7]
 [0.9999999999857788, -1.8913631003166995e-9, -5.333137047544641e-6, -0.010223966260268558, -7.830196373585224e-6]
 [0.9999999996124179, -5.727075167382377e-8, -2.7841771139449547e-5, -0.023360195092905718, -9.45372590818924e-5]
 [0.9999999975903671, -3.4756553131252326e-7, -6.942093370365734e-5, -0.03688699364700588, -0.00037351477565057646]
 [0.9999999912808007, -1.2748138702553709e-6, -0.00013205453645642666, -0.050875032484888456, -0.0009815792521743628]
 ⋮
 [0.9901277842947906, -0.23238732189767267, -0.14016765426150804, -1.6420296707064255, -4.122644158558371]
 [0.985454398388997, -0.31021166964780994, -0.16994007637880573, -1.7993240604491696, -4.999443332563201]
 [0.9795703957592, -0.3993607636868584, -0.20110157032963066, -1.9455248396534528, -5.919444799026639]
 [0.9725784264530312, -0.49678958755533603, -0.23257516887069882, -2.077170448545791, -6.852008926699451]
 [0.9648710403998514, -0.596641727365274, -0.2627239736228932, -2.190065778576504, -7.749435266344517]
 [0.9576823802250365, -0.6844107073683618, -0.28782715608691384, -2.2750844910874544, -8.500439881334152]
 [0.9511355191005343, -0.7607545737149946, -0.3087737711720759, -2.3401583473265153, -9.130247719863329]
 [0.9453098358657869, -0.8262359332944337, -0.32617377783827084, -2.3902904820373436, -9.655686408937372]
 [0.9445663341419724, -0.8344470431825701, -0.32832063660752875, -2.396236643829112, -9.720701118104335]

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(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.010943826063200043
│   sol1_f_evals = 1085
│   sol3_final_error = 0.05928571231930675
└   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.