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
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.806859541877657e-5
 4.5496933032029966e-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.2119452819316723e-13, -7.101167990076622e-9, -0.00037307223510401005, -2.651533791865963e-8]
 [1.0, -3.1661973720383655e-12, -1.0142857485072981e-8, -0.00044586994371389327, -2.9596714810803877e-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: 637-element Vector{Float64}:
 0.0
 1.0e-6
 1.1e-5
 6.293457247564385e-5
 0.0001717966902698356
 0.00041700284549204755
 0.001043224623113885
 0.0023835910941641687
 0.003763856534047043
 0.005191196217624695
 ⋮
 4.907553725698676
 4.924636353074359
 4.9411967014972324
 4.95668039898637
 4.970681486255768
 4.981792730511177
 4.990762595647369
 4.998031888191015
 5.0
u: 637-element Vector{Vector{Float64}}:
 [1.0, 0.0, 0.0, 0.0, 0.0]
 [1.0, -2.4320527589833776e-18, -4.8999999999373206e-12, -9.800000000000003e-6, -3.870438350247696e-16]
 [1.0, 1.779329265947601e-14, -5.929000004468873e-10, -0.00010780000000000091, -3.0178908529874883e-12]
 [0.9999999999999998, -3.705258613763331e-12, -1.9407725978901355e-8, -0.0006167588102611186, -1.3036994235702973e-9]
 [0.9999999999999896, 1.231734351558704e-11, -1.4461910374133314e-7, -0.001683607564645096, -3.158050706805358e-8]
 [0.9999999999996371, -9.147857460683334e-11, -8.520677279956897e-7, -0.0040866278858198955, -4.845676031758449e-7]
 [0.9999999999857808, -1.891233463103192e-9, -5.332756312011546e-6, -0.010223601306518774, -7.83009088217851e-6]
 [0.9999999996124844, -5.726676209454051e-8, -2.7839381844867827e-5, -0.023359192721600835, -9.45339864680247e-5]
 [0.9999999975906804, -3.4755243956250846e-7, -6.941641839844382e-5, -0.0368857940178753, -0.00037351357964688953]
 [0.9999999912816985, -1.2748040564497867e-6, -0.00013204773764628023, -0.05087372282090536, -0.0009815961712321248]
 ⋮
 [0.9901651613957098, -0.23172735115694545, -0.13990337266694583, -1.6405123591693411, -4.115166337579525]
 [0.9855712161386252, -0.3083616993192879, -0.16926127369263747, -1.795912092376441, -4.980213518414672]
 [0.9797518017566736, -0.3967501726849769, -0.20021590560182462, -1.9415667671742471, -5.894755669328316]
 [0.9728831662637859, -0.4927520327194497, -0.23129710247383958, -2.072072459058729, -6.816526114467944]
 [0.9653111301491414, -0.5912124320760092, -0.26110234215619854, -2.184252881590318, -7.704686847211488]
 [0.9582752681971104, -0.6774847635078917, -0.28584702809597784, -2.268603444057187, -8.445920248663054]
 [0.9518735971071052, -0.7524892742760094, -0.3064909010545698, -2.3332438885829054, -9.067471577304934]
 [0.9461805448181262, -0.8168252614310116, -0.32363929847919043, -2.383121166331972, -9.58619264314162]
 [0.9445584480218617, -0.834794040825908, -0.3283433240601558, -2.3962124390909865, -9.728925871084446]

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} \mathtt{x_1}\left( t \right)}{\mathrm{d}t} &= \mathtt{x_2}\left( t \right) \\ \frac{\mathrm{d} \mathtt{x_2}\left( t \right)}{\mathrm{d}t} &= \mathtt{x_1}\left( t \right) \mathtt{x_5}\left( t \right) \\ \frac{\mathrm{d} \mathtt{x_3}\left( t \right)}{\mathrm{d}t} &= \mathtt{x_4}\left( t \right) \\ \frac{\mathrm{d} \mathtt{x_4}\left( t \right)}{\mathrm{d}t} &= - \mathtt{\alpha_1} + \mathtt{x_3}\left( t \right) \mathtt{x_5}\left( t \right) \\ 0 &= \left( \mathtt{x_1}\left( t \right) \right)^{2} + \left( \mathtt{x_3}\left( t \right) \right)^{2} - \mathtt{\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} \mathtt{x_1}\left( t \right)}{\mathrm{d}t} &= \mathtt{x_2}\left( t \right) \\ \frac{\mathrm{d} \mathtt{x_2}\left( t \right)}{\mathrm{d}t} &= \mathtt{x_1}\left( t \right) \mathtt{x_5}\left( t \right) \\ \frac{\mathrm{d} \mathtt{x_3}\left( t \right)}{\mathrm{d}t} &= \mathtt{x_4}\left( t \right) \\ \frac{\mathrm{d} \mathtt{x_4}\left( t \right)}{\mathrm{d}t} &= - \mathtt{\alpha_1} + \mathtt{x_3}\left( t \right) \mathtt{x_5}\left( t \right) \\ 0 &= 2 \left( \mathtt{x_4}\left( t \right) \right)^{2} + 2 \left( \mathtt{x_2}\left( t \right) \right)^{2} + 2 \left( \mathtt{x_1}\left( t \right) \right)^{2} \mathtt{x_5}\left( t \right) + 2 \mathtt{x_3}\left( t \right) \left( - \mathtt{\alpha_1} + \mathtt{x_3}\left( t \right) \mathtt{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.010916847865436992
│   sol1_f_evals = 1085
│   sol3_final_error = 0.06751295080513564
└   sol3_f_evals = 1379

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.