GitHub

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
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,
)
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
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.9999999999999974, -9.18303665625327e-11, -7.098400431570907e-8, -0.0011795280770663706, -3.336177335561177e-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.00526081718946869
 ⋮
 4.894226425458966
 4.911340792753645
 4.928446281898732
 4.944939594798175
 4.960234245048122
 4.973988440235516
 4.9847640040260615
 4.993429967827159
 5.0
u: 622-element Vector{Vector{Float64}}:
 [1.0, 0.0, 0.0, 0.0, 0.0]
 [1.0, -2.992994094559582e-18, -4.8999999999071035e-12, -9.800000000000003e-6, -5.706217293757403e-16]
 [1.0, 5.278742785090216e-13, -4.586688442166948e-9, -0.000299831774639812, -1.003289285148396e-10]
 [0.9999999999999986, -1.1002570206373845e-11, -5.3950632953151425e-8, -0.0010283153263207937, -6.916581259942502e-9]
 [0.999999999999966, 1.4396812897648353e-11, -2.605449271805613e-7, -0.002259796577761726, -8.90804584615879e-8]
 [0.9999999999991783, -2.1028413016953296e-10, -1.2820106361283433e-6, -0.005012724656719099, -1.0782453237305222e-6]
 [0.9999999999771458, -3.803393090436643e-9, -6.7607990620882345e-6, -0.01151137096819222, -1.3743963998806765e-5]
 [0.9999999995517715, -8.150989306828036e-8, -2.9940889534892595e-5, -0.024224810325656067, -0.00013084620605558294]
 [0.9999999973704503, -4.721003199902214e-7, -7.25196493160739e-5, -0.03770126160100114, -0.0004966739294258048]
 [0.9999999908045084, -1.6735625786117297e-6, -0.00013561336631285654, -0.051556008305906835, -0.001274409462520493]
 ⋮
 [0.9929154934854335, -0.18141105675422825, -0.11882265425073626, -1.5163431809539891, -3.4956153382159423]
 [0.9892641298195628, -0.24741098253915317, -0.14613857097631908, -1.6752259094491317, -4.299855779304121]
 [0.9843681904670305, -0.3273501274488839, -0.17612287319736997, -1.8297878511006567, -5.184670149880936]
 [0.9782358877449209, -0.4186857991873351, -0.20749590382192173, -1.9735490281466135, -6.113696197852662]
 [0.9711022017485105, -0.516662109510528, -0.23866403045672485, -2.1009928690660353, -7.041041370255113]
 [0.9633302841451111, -0.6161818085693496, -0.2683184238644507, -2.2096854728828585, -7.928485939058314]
 [0.9562441865541502, -0.7019325623814641, -0.2925697693684672, -2.2902215348989943, -8.658770807849216]
 [0.9498555587710478, -0.7759506047296278, -0.31268904030120714, -2.3516395095852163, -9.268209476059964]
 [0.9445767200775522, -0.8351021819924697, -0.3282907774353968, -2.396009597037009, -9.743302870465886]

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)
Model ##MTKizedODE#309:
Equations (5):
  5 standard: see equations(##MTKizedODE#309)
Unknowns (5): see unknowns(##MTKizedODE#309)
  x₁(t)
  x₂(t)
  x₃(t)
  x₄(t)
  ⋮
Parameters (2): see parameters(##MTKizedODE#309)
  α₁
  α₂

(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#309:
Equations (5):
  5 standard: see equations(##MTKizedODE#309)
Unknowns (5): see unknowns(##MTKizedODE#309)
  x₅(t)
  x₃(t)
  x₄(t)
  x₂(t)
  ⋮
Parameters (2): see parameters(##MTKizedODE#309)
  α₁
  α₂

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/zJBUr/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/zJBUr/src/initialization.jl:22
┌ Info: Results
│   sol1_final_error = 0.0017424939900806402
│   sol1_f_evals = 2921
│   sol3_final_error = 0.07679411214033639
└   sol3_f_evals = 4613

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

[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.