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.

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
Non-trivial mass matrix: 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,
)

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
Non-trivial mass matrix: 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 4th (Rodas6P = 5th) order "free" stiffness-aware interpolation
t: 3-element Vector{Float64}:
 0.0
 1.0e-6
 0.00012036000786391582
u: 3-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]
 [0.9999999999999976, -9.18303665625327e-11, -7.098400431570907e-8, -0.0011795280770663706, -3.3361773355611767e-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: 622-element Vector{Float64}:
 0.0
 1.0e-6
 3.059507904487559e-5
 0.00010493013533892183
 0.00023059148752658964
 0.00051150251599221
 0.00117462969063216
 0.002471919421168622
 0.0038470675125998333
 0.005260817088876566
 ⋮
 4.8922543410906085
 4.909381193749295
 4.926485578617025
 4.943083940106332
 4.9585125763118
 4.972428674793375
 4.983385462838629
 4.992206484066411
 5.0
u: 622-element Vector{Vector{Float64}}:
 [1.0, 0.0, 0.0, 0.0, 0.0]
 [1.0, -2.9929940654790135e-18, -4.8999999999071035e-12, -9.800000000000003e-6, -5.706217204585046e-16]
 [1.0, 5.278742751208898e-13, -4.586688442167038e-9, -0.00029983177463981206, -1.0032892693779432e-10]
 [0.9999999999999986, -1.1002570125058682e-11, -5.395063295315021e-8, -0.001028315326320793, -6.916581149841972e-9]
 [0.999999999999966, 1.4396812775675609e-11, -2.605449271805619e-7, -0.0022597965777617253, -8.908045703415268e-8]
 [0.9999999999991783, -2.102841273506073e-10, -1.2820106361283378e-6, -0.005012724656719099, -1.0782453063726296e-6]
 [0.9999999999771458, -3.803393026577135e-9, -6.760799062088254e-6, -0.011511370968192219, -1.3743963776990392e-5]
 [0.9999999995517715, -8.150989178577761e-8, -2.9940889534892412e-5, -0.024224810325656067, -0.0001308462039415099]
 [0.9999999973704503, -4.7210031229162706e-7, -7.251964931607653e-5, -0.03770126160100114, -0.0004966739213981555]
 [0.9999999908045091, -1.6735624224887852e-6, -0.00013561336112672556, -0.051556007320104034, -0.0012744093685905177]
 ⋮
 [0.9932666629562176, -0.17465042935745456, -0.11585049250241208, -1.497815151769782, -3.408244949955014]
 [0.9897409779743751, -0.239168206441898, -0.14287335990085748, -1.6572252610230434, -4.203798206047013]
 [0.9850003703074962, -0.3174473559017723, -0.17255222769019696, -1.8123259577702995, -5.079480109798372]
 [0.9790024959669662, -0.40769394112966373, -0.20384826429589092, -1.9576737599168494, -6.0060372667597735]
 [0.9719811903870471, -0.5049976503243307, -0.23505865369910028, -2.08696494257465, -6.9344046128888595]
 [0.9642811832135773, -0.6043780479142012, -0.26488074994508204, -2.1976664963535626, -7.826627516718407]
 [0.9572025210590386, -0.690628941061858, -0.28941898240213015, -2.2801528573369936, -8.56537002913494]
 [0.9507964728795579, -0.7652949723527499, -0.30981621178897595, -2.343143773869247, -9.183121966132624]
 [0.9445766403502521, -0.8351892752986719, -0.3282910291361937, -2.3959852451645682, -9.74587222358895]

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

plot(sol)

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)

Model ##MTKizedODE#417:
Equations (5):
  5 standard: see equations(##MTKizedODE#417)
Unknowns (5): see unknowns(##MTKizedODE#417)
  x₁(t)
  x₂(t)
  x₃(t)
  x₄(t)
  ⋮
Parameters (2): see parameters(##MTKizedODE#417)
  α₁
  α₂

(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))

Model ##MTKizedODE#417:
Equations (5):
  5 standard: see equations(##MTKizedODE#417)
Unknowns (5): see unknowns(##MTKizedODE#417)
  x₅(t)
  x₃(t)
  x₄(t)
  x₂(t)
  ⋮
Parameters (2): see parameters(##MTKizedODE#417)
  α₁
  α₂

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(), abstol=1e-7)
sol1 = solve(prob_index1, EK1(), abstol=1e-7)

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
perm = [findfirst(x -> x === s, unknowns(traced_sys)) for s in unknowns(simplified_sys)]
sol3_final_error = norm(sol3.u[end][perm] - truesol.u[end])
sol3_f_evals     = sol3.stats.nf
@info "Results" sol1_final_error sol1_f_evals sol3_final_error sol3_f_evals

┌ Warning: Initialization system is overdetermined. 1 equations for 0 unknowns. Initialization will default to using least squares. `SCCNonlinearProblem` can only be used for initialization of fully determined systems and hence will not be used here.
│
│
│ To suppress this warning, pass `warn_initialize_determined = false`. To turn this warning into an error, pass `fully_determined = true`.
└ @ ModelingToolkit ~/.julia/packages/ModelingToolkit/wVorF/src/initialization.jl:22
┌ Warning: Initialization system is overdetermined. 4 equations for 2 unknowns. Initialization will default to using least squares. `SCCNonlinearProblem` can only be used for initialization of fully determined systems and hence will not be used here.
│
│
│ To suppress this warning, pass `warn_initialize_determined = false`. To turn this warning into an error, pass `fully_determined = true`.
└ @ ModelingToolkit ~/.julia/packages/ModelingToolkit/wVorF/src/initialization.jl:22
┌ Info: Results
│   sol1_final_error = 0.0031806557760521697
│   sol1_f_evals = 2909
│   sol3_final_error = 0.07656544386494461
└   sol3_f_evals = 4511

The error for the index-1 DAE solve is quite a bit 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