ModelingToolkit: define ODEs with symbolic expressions

ModelingToolkit: define ODEs with symbolic expressions#

You can also define ODE systems symbolically using ModelingToolkit.jl (MTK) and let MTK generate high-performace ODE functions.

transforming and simplifying expressions for performance
Jacobian and Hessian generation
accessing eliminated states (called observed variables)
building and connecting subsystems programatically (component-based modeling)

Radioactive decay#

Here we use the same example of decaying radioactive elements

using ModelingToolkit
using DifferentialEquations
using Plots

[ Info: Precompiling IJuliaExt [2f4121a4-3b3a-5ce6-9c5e-1f2673ce168a]

# independent variable (time) and dependent variables
@variables t c(t) RHS(t)

# parameters: decay rate
@parameters λ

# Differential operator w.r.t. time
D = Differential(t)

(::Differential) (generic function with 3 methods)

Equations in MTK use the tilde character (~) for equality. Every MTK system requires a name. The @named macro simply ensures that the symbolic name matches the name in the REPL.

eqs = [
    RHS ~ -λ * c
    D(c) ~ RHS
]

@named sys = ODESystem(eqs, t)

\[\begin{split} \begin{align} \frac{\mathrm{d} c\left( t \right)}{\mathrm{d}t} =& \mathrm{RHS}\left( t \right) \\ \mathrm{RHS}\left( t \right) =& - c\left( t \right) \lambda \end{align} \end{split}\]

structural_simplify() simplifies the two equations to one. You can also use @mtkbuild sys = ODESystem(eqs, t) to merge these two steps into one

sys = structural_simplify(sys)

\[ \begin{align} \frac{\mathrm{d} c\left( t \right)}{\mathrm{d}t} =& \mathrm{RHS}\left( t \right) \end{align} \]

Setup initial conditions, time span, parameter values, the ODEProblem, and solve the problem.

p = [λ => 1.0]
u0 = [c => 1.0]
tspan = (0.0, 2.0)
prob = ODEProblem(sys, u0, tspan, p)
sol = solve(prob)

retcode: Success
Interpolation: 3rd order Hermite
t: 8-element Vector{Float64}:
 0.0
 0.10001999200479662
 0.34208427873632274
 0.6553980290285384
 1.0312652902321524
 1.4709406498424789
 1.9659577002710475
 2.0
u: 8-element Vector{Vector{Float64}}:
 [1.0]
 [0.9048193287657775]
 [0.7102883564034526]
 [0.5192354320104405]
 [0.35655575232768816]
 [0.2297097760377979]
 [0.14002246806154703]
 [0.13533600284000216]

Visualize the solution with Plots.jl.

plot(sol, label="Exp decay")

The solution interface provides symbolic access. So you can access the results of c directly.

sol[c]

8-element Vector{Float64}:
0
9048193287657775
7102883564034526
5192354320104405
35655575232768816
2297097760377979
14002246806154703
13533600284000216

sol(0.0:0.1:2.0, idxs=c)

t: 0.0:0.1:2.0
u: 21-element Vector{Float64}:
0
9048389983836875
8225844481735297
7453905798423603
6716690052872237
6126431137597176
554633689939333
4971936192821015
45380162043196953
4138482560131354
3713271431453589
3337844113123039
30522377530459677
2789957673333738
2516591524294101
22324380097718674
20365547977113405
18672102774211408
17051216093498966
15310059539462373
13533600284000216

The eliminated term (RHS in this example) is still tracible.

plot(sol, idxs=[c, RHS], legend=:right)

The interface allows symbolic calculations.

plot(sol, idxs=[c * 1000])

Lorenz system#

We use the same Lorenz system example as above. Here we setup the initial conditions and parameters with default values.

@variables t x(t) = 1.0 y(t) = 0.0 z(t) = 0.0
@parameters (σ=10.0, ρ=28.0, β=8 / 3)

D = Differential(t)

eqs = [
    D(x) ~ σ * (y - x)
    D(y) ~ x * (ρ - z) - y
    D(z) ~ x * y - β * z
]

@mtkbuild sys = ODESystem(eqs, t)

\[\begin{split} \begin{align} \frac{\mathrm{d} x\left( t \right)}{\mathrm{d}t} =& \left( - x\left( t \right) + y\left( t \right) \right) \sigma \\ \frac{\mathrm{d} y\left( t \right)}{\mathrm{d}t} =& - y\left( t \right) + x\left( t \right) \left( - z\left( t \right) + \rho \right) \\ \frac{\mathrm{d} z\left( t \right)}{\mathrm{d}t} =& x\left( t \right) y\left( t \right) - z\left( t \right) \beta \end{align} \end{split}\]

Here we are using default values, so we pass empty arrays for initial conditions and parameter values.

tspan = (0.0, 100.0)
prob = ODEProblem(sys, [], tspan, [])
sol = solve(prob)

retcode: Success
Interpolation: 3rd order Hermite
t: 1281-element Vector{Float64}:
   0.0
   3.5678604836301404e-5
   0.0003924646531993154
   0.0032624087100077666
   0.009058076582749423
   0.016956470605311864
   0.027689959227781235
   0.04185635103821218
   0.060240410627700816
   0.0836854113984534
   0.11336499269451543
   0.14862181409827
   0.18703978025370946
   ⋮
  99.19664352825902
  99.28542618302045
  99.3530549740687
  99.4386918268877
  99.53430077629385
  99.62315909447312
  99.70299380266312
  99.7727905290976
  99.8468719463341
  99.92909223837675
  99.99676433660723
 100.0
u: 1281-element Vector{Vector{Float64}}:
 [1.0, 0.0, 0.0]
 [0.9996434557625105, 0.0009988049817849058, 1.781434788799189e-8]
 [0.9961045497425811, 0.010965399721242457, 2.1469553658389193e-6]
 [0.9693591550149778, 0.08977063252764937, 0.0001438019170127846]
 [0.924204355043198, 0.242289149116772, 0.0010461625397616113]
 [0.8800455796215916, 0.4387364900041282, 0.0034242599668253874]
 [0.8483309836977301, 0.6915629475762161, 0.008487624968655677]
 [0.8495036679850485, 1.0145426495980538, 0.018212090123888365]
 [0.9139069520070257, 1.4425599557988966, 0.03669382071284047]
 [1.0888638157267199, 2.052326562845511, 0.07402573450924263]
 [1.4608626762755585, 3.0206719763684764, 0.160039355072216]
 [2.162723385463873, 4.6333636144497445, 0.3771173678677311]
 [3.3684642325435363, 7.267693727122867, 0.9363555414291395]
 ⋮
 [2.4154751319743304, 0.9019051839955854, 23.10904521543824]
 [1.9010431518181314, 2.075631191607503, 18.469762607052317]
 [2.3313747500610864, 3.427114121484939, 15.773720459660732]
 [3.8662182925863604, 6.477675362450072, 13.7026651297458]
 [7.596364784839832, 12.931698345552782, 15.295013031866944]
 [12.999465900688321, 18.69084638217891, 25.77003643982645]
 [14.725104821648667, 12.127682492556511, 37.85997432096073]
 [10.45813774612691, 2.2355517835584817, 37.102951902956555]
 [4.858569386282909, -1.238820833292597, 30.50492838887987]
 [1.4267543691218378, -1.0945246877110832, 24.22832564643038]
 [0.2911991490440339, -0.7513530324542533, 20.18373424388144]
 [0.25815851151351493, -0.7419175275518352, 20.00966894969765]

Plot the solution with symbols instead of index numbers.

plot(sol, idxs=(x, y, z), size=(400, 400))

Non-autonomous ODEs#

Sometimes a model might have a time-variant external force, which is too complex or impossible to express it symbolically. In such situation, one could apply @register_symbolic to it to exclude it from symbolic transformations and use it numberically.

@variables t x(t) f(t)
@parameters τ
D = Differential(t)

(::Differential) (generic function with 3 methods)

Define a time-dependent random external force

value_vector = randn(10)

f_fun(t) = t >= 10 ? value_vector[end] : value_vector[Int(floor(t))+1]

f_fun (generic function with 1 method)

“Register” arbitrary Julia functions to be excluded from symbolic transformations. Just use it as-is (numberically).

@register_symbolic f_fun(t)
@named fol_external_f = ODESystem([f ~ f_fun(t), D(x) ~ (f - x) / τ], t)

\[\begin{split} \begin{align} \frac{\mathrm{d} x\left( t \right)}{\mathrm{d}t} =& \frac{ - x\left( t \right) + f\left( t \right)}{\tau} \\ f\left( t \right) =& f_{fun}\left( t \right) \end{align} \end{split}\]

sys = structural_simplify(fol_external_f)
prob = ODEProblem(sys, [x => 0.0], (0.0, 10.0), [τ => 0.75])
sol = solve(prob)
plot(sol, idxs=[x, f])

Second-order ODE systems#

ode_order_lowering(sys) automatically transforms a second-order ODE into two first-order ODEs.

using Plots
using ModelingToolkit
using DifferentialEquations

@parameters σ ρ β
@variables t x(t) y(t) z(t)
D = Differential(t)

eqs = [
    D(D(x)) ~ σ * (y - x),
    D(y) ~ x * (ρ - z) - y,
    D(z) ~ x * y - β * z
]

\[\begin{split} \begin{align} \frac{\mathrm{d}}{\mathrm{d}t} \frac{\mathrm{d} x\left( t \right)}{\mathrm{d}t} =& \left( - x\left( t \right) + y\left( t \right) \right) \sigma \\ \frac{\mathrm{d} y\left( t \right)}{\mathrm{d}t} =& - y\left( t \right) + x\left( t \right) \left( - z\left( t \right) + \rho \right) \\ \frac{\mathrm{d} z\left( t \right)}{\mathrm{d}t} =& x\left( t \right) y\left( t \right) - z\left( t \right) \beta \end{align} \end{split}\]

@mtkbuild automatically does the transformations and simplifications The same as sys = sys |> ode_order_lowering |> structural_simplify

@mtkbuild sys = ODESystem(eqs, t)

\[\begin{split} \begin{align} \frac{\mathrm{d} xˍt\left( t \right)}{\mathrm{d}t} =& \left( - x\left( t \right) + y\left( t \right) \right) \sigma \\ \frac{\mathrm{d} y\left( t \right)}{\mathrm{d}t} =& - y\left( t \right) + x\left( t \right) \left( - z\left( t \right) + \rho \right) \\ \frac{\mathrm{d} z\left( t \right)}{\mathrm{d}t} =& x\left( t \right) y\left( t \right) - z\left( t \right) \beta \\ \frac{\mathrm{d} x\left( t \right)}{\mathrm{d}t} =& xˍt\left( t \right) \end{align} \end{split}\]

Note that one needs to provide the initial condition for x’s derivative.

u0 = [
    D(x) => 2.0,
    x => 1.0,
    y => 0.0,
    z => 0.0
]

p = [
    σ => 28.0,
    ρ => 10.0,
    β => 8 / 3
]

tspan = (0.0, 100.0)
prob = ODEProblem(sys, u0, tspan, p, jac=true)
sol = solve(prob)
plot(sol, idxs=(x, y, z), label="Trajectory", size=(500, 500))

Composing systems#

By defining connection equation(s) to couple ODE systems together, we can build component-based, hierarchical models.

using Plots
using ModelingToolkit
using DifferentialEquations

@parameters σ ρ β
@variables t x(t) y(t) z(t)
D = Differential(t)

eqs = [
    D(x) ~ σ * (y - x),
    D(y) ~ x * (ρ - z) - y,
    D(z) ~ x * y - β * z
]

@named lorenz1 = ODESystem(eqs)
@named lorenz2 = ODESystem(eqs)

\[\begin{split} \begin{align} \frac{\mathrm{d} x\left( t \right)}{\mathrm{d}t} =& \left( - x\left( t \right) + y\left( t \right) \right) \sigma \\ \frac{\mathrm{d} y\left( t \right)}{\mathrm{d}t} =& - y\left( t \right) + x\left( t \right) \left( - z\left( t \right) + \rho \right) \\ \frac{\mathrm{d} z\left( t \right)}{\mathrm{d}t} =& x\left( t \right) y\left( t \right) - z\left( t \right) \beta \end{align} \end{split}\]

Define relations (connectors) between the two systems.

@variables a(t)
@parameters γ

connections = [0 ~ lorenz1.x + lorenz2.y + a * γ]

@named connLorenz = ODESystem(connections, t, [a], [γ], systems=[lorenz1, lorenz2])

\[\begin{split} \begin{align} 0 =& lorenz2_{+}y\left( t \right) + lorenz1_{+}x\left( t \right) + a\left( t \right) \gamma \\ \frac{\mathrm{d} lorenz1_{+}x\left( t \right)}{\mathrm{d}t} =& lorenz1_+\sigma \left( - lorenz1_{+}x\left( t \right) + lorenz1_{+}y\left( t \right) \right) \\ \frac{\mathrm{d} lorenz1_{+}y\left( t \right)}{\mathrm{d}t} =& - lorenz1_{+}y\left( t \right) + \left( lorenz1_+\rho - lorenz1_{+}z\left( t \right) \right) lorenz1_{+}x\left( t \right) \\ \frac{\mathrm{d} lorenz1_{+}z\left( t \right)}{\mathrm{d}t} =& - lorenz1_+\beta lorenz1_{+}z\left( t \right) + lorenz1_{+}x\left( t \right) lorenz1_{+}y\left( t \right) \\ \frac{\mathrm{d} lorenz2_{+}x\left( t \right)}{\mathrm{d}t} =& lorenz2_+\sigma \left( lorenz2_{+}y\left( t \right) - lorenz2_{+}x\left( t \right) \right) \\ \frac{\mathrm{d} lorenz2_{+}y\left( t \right)}{\mathrm{d}t} =& - lorenz2_{+}y\left( t \right) + \left( lorenz2_+\rho - lorenz2_{+}z\left( t \right) \right) lorenz2_{+}x\left( t \right) \\ \frac{\mathrm{d} lorenz2_{+}z\left( t \right)}{\mathrm{d}t} =& - lorenz2_+\beta lorenz2_{+}z\left( t \right) + lorenz2_{+}y\left( t \right) lorenz2_{+}x\left( t \right) \end{align} \end{split}\]

All state variables in the combined system

states(connLorenz)

7-element Vector{SymbolicUtils.BasicSymbolic{Real}}:
 a(t)
 lorenz1₊x(t)
 lorenz1₊y(t)
 lorenz1₊z(t)
 lorenz2₊x(t)
 lorenz2₊y(t)
 lorenz2₊z(t)

u0 = [
    lorenz1.x => 1.0, lorenz1.y => 0.0, lorenz1.z => 0.0,
    lorenz2.x => 0.0, lorenz2.y => 1.0, lorenz2.z => 0.0,
    a => 2.0
]

p = [
    lorenz1.σ => 10.0, lorenz1.ρ => 28.0, lorenz1.β => 8 / 3,
    lorenz2.σ => 10.0, lorenz2.ρ => 28.0, lorenz2.β => 8 / 3,
    γ => 2.0
]

tspan = (0.0, 100.0)
sys = connLorenz |> structural_simplify
sol = solve(ODEProblem(sys, u0, tspan, p, jac=true))
plot(sol, idxs=(a, lorenz1.x, lorenz2.x), size=(600, 600))

Convert existing functions into MTK systems#

modelingtoolkitize(prob) generates MKT systems from regular DE problems. I t can also generate analytic Jacobin functions for faster solving.

Example: DAE index reduction for the pendulum problem, which cannot be solved by regular ODE solvers.

using Plots
using ModelingToolkit
using DifferentialEquations
using LinearAlgebra

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
    # Do not write your function like this after you've learned MTK
    du[5] = x^2 + y^2 - L^2
    return nothing
end

pendulum! (generic function with 1 method)

pendulum_fun! = ODEFunction(pendulum!, mass_matrix=Diagonal([1, 1, 1, 1, 0]))
u0 = [1.0, 0.0, 0.0, 0.0, 0.0]
p = [9.8, 1.0]
tspan = (0.0, 10.0)
pendulum_prob = ODEProblem(pendulum_fun!, u0, tspan, p)

ODEProblem with uType Vector{Float64} and tType Float64. In-place: true
timespan: (0.0, 10.0)
u0: 5-element Vector{Float64}:
 1.0
 0.0
 0.0
 0.0
 0.0

Convert the ODE problem into a MTK system.

tracedSys = modelingtoolkitize(pendulum_prob)

\[\begin{split} \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} \end{split}\]

structural_simplify() and dae_index_lowering() transform the index-3 DAE into an index-0 ODE.

pendulumSys = tracedSys |> dae_index_lowering |> structural_simplify

\[\begin{split} \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} \end{split}\]

The default u0 is included in the system already so one can use an empty array [] as the initial conditions.

prob = ODAEProblem(pendulumSys, [], tspan)
sol = solve(prob, abstol=1e-8, reltol=1e-8)
plot(sol, idxs=states(tracedSys))

This notebook was generated using Literate.jl.