ModelingToolkit.jl

A lot of people are building modeling languages for their specific domains. However, while the syntax my vary greatly between these domain-specific languages (DSLs), the internals of modeling frameworks are surprisingly similar: building differential equations, calculating Jacobians, etc.

ModelingToolkit.jl is metamodeling systemitized.

After building our third modeling interface, we realized that this problem can be better approached by having a reusable internal structure which DSLs can target. This internal is ModelingToolkit.jl: an Intermediate Representation (IR) with a well-defined interface for defining system transformations and compiling to Julia functions for use in numerical libraries. Now a DSL can easily be written by simply defining the translation to ModelingToolkit.jl's primatives and querying for the mathematical quantities one needs.

Basic usage: defining differential equation systems, with performance!

Let's explore the IR itself. ModelingToolkit.jl is friendly to use, and can used as a symbolic DSL in its own right. Let's define and solve the Lorenz differential equation system using ModelingToolkit to generate the functions:

Installation

using Pkg
pkg"add ModelingToolkit DifferentialEquations"

23.4s

Julia

Example

using ModelingToolkit, DifferentialEquations, Plots

# Define a differential equation system
@parameters t σ ρ β
@variables x(t) y(t) z(t)
@derivatives D'~t

# Lorenz attractor
eqs = [D(x) ~ σ*(y-x),
       D(y) ~ x*(ρ-z)-y,
       D(z) ~ x*y - β*z]
de = ODESystem(eqs)
ode_f = ODEFunction(de, [x,y,z], [σ,ρ,β])

# Plug it in DifferentialEquations.jl

u₀ = ones(3)
tspan = (0.0,100.0)
p = [10.0,28.0,10/3]
prob = ODEProblem(ode_f,u₀,tspan,p)
sol = solve(prob,Tsit5())

plot(sol,vars=(1,2,3), lw=1)

373.3s

Julia

Generated RHS function

ModelingToolkit is a compiler for mathematical systems.

At its core, ModelingToolkit is a compiler. It's IR is its type system, and its output are Julia functions (it's a compiler for Julia code to Julia code, written in Julia).

DifferentialEquations.jl wants a function f(du,u,p,t) for defining an ODE system, which is what ModelingToolkit.jl is building.

generate_function(de, [x,y,z], [σ,ρ,β])

1.1s

Julia

Jacobian

ModelingToolkit.jl can be used to calculate the Jacobian of the differential equation system:

jac = calculate_jacobian(de)

3.3s

Julia

3×3 Array{Expression,2}: σ() * -1 σ() Constant(0) ρ() - z(t()) Constant(-1) x(t()) * -1 y(t()) x(t()) -1 * β()

It will automatically generate functions for using this Jacobian within the stiff ODE solvers for faster solving:

jac_expr = generate_jacobian(de)

1.2s

Julia

It can even do fancy linear algebra. Stiff ODE solvers need to perform an LU-factorization which is their most expensive part. But ModelingToolkit.jl can skip this operation and instead generate the analytical solution to a matrix factorization, and build a Julia function for directly computing the factorization, which is then optimized in LLVM compiler passes.

ModelingToolkit.generate_factorized_W(de)[1]

1.9s

Julia

Solving Nonlinear systems

ModelingToolkit.jl is not just for differential equations. It can be used for any mathematical target that is representable by its IR. For example, let's solve a rootfinding problem F(x)=0. What we do is define a nonlinear system and generate a function for use in NLsolve.jl

@variables x y z
@parameters σ ρ β

# Define a nonlinear system
eqs = [0 ~ σ*(y-x),
       0 ~ x*(ρ-z)-y,
       0 ~ x*y - β*z]
ns = NonlinearSystem(eqs, [x,y,z])
nlsys_func = generate_function(ns, [x,y,z], [σ,ρ,β])

1.5s

Julia

We can then tell ModelingToolkit.jl to compile this function for use in NLsolve.jl, and then numerically solve the rootfinding problem:

# Compile the function [1] = native version, [2] = in-place version
nl_f = @eval eval(nlsys_func[2])

# Make a closure over the parameters for for NLsolve.jl
f_nl(du, u) = nl_f(du, u, (10.0,26.0,2.33))

using NLsolve
nlsolve(f_nl, ones(3))

12.9s

Julia

Results of Nonlinear Solver Algorithm * Algorithm: Trust-region with dogleg and autoscaling * Starting Point: [1.0, 1.0, 1.0] * Zero: [2.2228042242798023e-10, 2.222804224296743e-10, -9.990339458476605e-11] * Inf-norm of residuals: 0.000000 * Iterations: 3 * Convergence: true * |x - x'| < 0.0e+00: false * |f(x)| < 1.0e-08: true * Function Calls (f): 4 * Jacobian Calls (df/dx): 4

Transformations on mathematical systems

The reason for using ModelingToolkit is not just for defining performant Julia functions for solving systems, but also for performing mathematical transformations which may be required in order to numerically solve the system. For example, let's solve a third order ODE. The way this is done is by transforming the third order ODE into a first order ODE, and then solving the resulting ODE. This transformation is given by the ode_order_lowering(de) function.

@derivatives D3'''~t
@derivatives D2''~t
@variables u(t), x(t)
eqs = [D3(u) ~ 2(D2(u)) + D(u) + D(x) + 1
       D2(x) ~ D(x) + 2]
de = ODESystem(eqs)
de1 = ode_order_lowering(de)

1.7s

Julia

ODESystem(ModelingToolkit.DiffEq[DiffEq(u_tt, 1, ((2 * u_tt(t()) + u_t(t())) + x_t(t())) + 1), DiffEq(x_t, 1, x_t(t()) + 2), DiffEq(u_t, 1, u_tt(t())), DiffEq(u, 1, u_t(t())), DiffEq(x, 1, x_t(t()))], t, Variable[u, x, u_tt, u_t, x_t], Variable[], RefValue{Array{Expression,2}}(Array{Expression}(undef,0,0)), RefValue{Array{Expression,2}}(Array{Expression}(undef,0,0)), RefValue{Array{Expression,2}}(Array{Expression}(undef,0,0)))

de1.eqs

1.0s

Julia

5-element Array{DiffEq,1}: DiffEq(u_tt, 1, ((2 * u_tt(t()) + u_t(t())) + x_t(t())) + 1) DiffEq(x_t, 1, x_t(t()) + 2) DiffEq(u_t, 1, u_tt(t())) DiffEq(u, 1, u_t(t())) DiffEq(x, 1, x_t(t()))

This has generated a system of 5 first order ODE systems which can now be used in the ODE solvers.

Generated DE and Linear Algebra

Let's take a look at how to extend ModelingToolkit.jl in new directions. Let's define a Jacobian just by using the derivative primatives by hand:

@parameters t σ ρ β
@variables x(t) y(t) z(t)
@derivatives D'~t Dx'~x Dy'~y Dz'~z
eqs = [D(x) ~ σ*(y-x),
       D(y) ~ x*(ρ-z)-y,
       D(z) ~ x*y - β*z]
J = [Dx(eqs[1].rhs) Dy(eqs[1].rhs) Dz(eqs[1].rhs)
 Dx(eqs[2].rhs) Dy(eqs[2].rhs) Dz(eqs[2].rhs)
 Dx(eqs[3].rhs) Dy(eqs[3].rhs) Dz(eqs[3].rhs)]

1.2s

Julia

3×3 Array{Operation,2}: Differential(σ() * (y(t()) - x(t()))) … Differential(σ() * (y(t()) - x(t()))) Differential(x(t()) * (ρ() - z(t())) - y(t())) Differential(x(t()) * (ρ() - z(t())) - y(t())) Differential(x(t()) * y(t()) - β() * z(t())) Differential(x(t()) * y(t()) - β() * z(t()))

Notice that this writes the derivatives in a "lazy" manner. If we want to actually compute the derivatives, we can expand out those expressions by expand_derivatives.(J)

J = expand_derivatives.(J)

0.5s

Julia

3×3 Array{Expression,2}: σ() * -1 σ() Constant(0) ρ() - z(t()) Constant(-1) x(t()) * -1 y(t()) x(t()) -1 * β()

Here's the magic of ModelingToolkit.jl: Julia treats ModelingToolkit expressions like a Number, and so generic numerical functions are directly usable on ModelingToolkit expressions! Let's compute the LU-factorization of this Jacobian we defined using Julia's Base linear algebra library.

using LinearAlgebra
luJ = lu(J)

1.2s

Julia

LU{Expression,Array{Expression,2}} L factor: 3×3 Array{Expression,2}: Constant(1) … Constant(0) (ρ() - z(t())) * inv(σ() * -1) identity(0) y(t()) * inv(σ() * -1) Constant(1) U factor: 3×3 Array{Expression,2}: σ() * -1 … Constant(0) identity(0) x(t()) * -1 - ((ρ() - z(t())) * inv(σ() * -1)) * 0 identity(0) (-1 * β() - (y(t()) * inv(σ() * -1)) * 0) - ((x(t()) - (y(t()) * inv(σ() * -1)) * σ()) * inv(-1 - ((ρ() - z(t())) * inv(σ() * -1)) * σ())) * (x(t()) * -1 - ((ρ() - z(t())) * inv(σ() * -1)) * 0)

luJ.L

0.2s

Julia

3×3 Array{Expression,2}: Constant(1) … Constant(0) (ρ() - z(t())) * inv(σ() * -1) identity(0) y(t()) * inv(σ() * -1) Constant(1)

and the inverse?

invJ = inv(J)

2.0s

Julia

Thus ModelingToolkit.jl can utilize existing numerical code on symbolic codes

Let's follow this thread a little deeper.

Automatically convert old numerical codes to symbolic ones

Let's take someone's code written to numerically solve the Lorenz equation:

function lorenz(du,u,p,t)
 du[1] = p[1]*(u[2]-u[1])
 du[2] = u[1]*(p[2]-u[3]) - u[2]
 du[3] = u[1]*u[2] - p[3]*u[3]
end

0.6s

Julia

lorenz (generic function with 1 method)

Since ModelingToolkit can trace generic numerical functions in Julia, let's trace it with Operations. When we do this, it'll spit out a symbolic representation of their numerical code:

u = [x,y,z]
du = similar(u)
p = [σ,ρ,β]
lorenz(du,u,p,t)
du

0.9s

Julia

3-element Array{Operation,1}: σ() * (y(t()) - x(t())) x(t()) * (ρ() - z(t())) - y(t()) x(t()) * y(t()) - β() * z(t())

We can then perform symbolic manipulations on their numerical code, and build a new numerical code that optimizes/fixes their original function!

J = [Dx(du[1]) Dy(du[1]) Dz(du[1])
     Dx(du[2]) Dy(du[2]) Dz(du[2])
     Dx(du[3]) Dy(du[3]) Dz(du[3])]
J = expand_derivatives.(J)

0.2s

Julia

3×3 Array{Expression,2}: σ() * -1 σ() Constant(0) ρ() - z(t()) Constant(-1) x(t()) * -1 y(t()) x(t()) -1 * β()

Automated Sparsity Detection

In many cases one has to speed up large modeling frameworks by taking into account sparsity. While ModelingToolkit.jl can be used to compute Jacobians, we can write a standard Julia function in order to get a spase matrix of expressions which automatically detects and utilizes the sparsity of their function.

using Pkg
pkg"add SparseArrays"

42.1s

Julia

import SparseArrays
function SparseArrays.SparseMatrixCSC(M::Matrix{T}) where {T<:ModelingToolkit.Expression}
    idxs = findall(!iszero, M)
    I = [i[1] for i in idxs]
    J = [i[2] for i in idxs]
    V = [M[i] for i in idxs]
    return SparseArrays.sparse(I, J, V, size(M)...)
end
sJ = SparseArrays.SparseMatrixCSC(J)

1.6s

Julia

3×3 SparseMatrixCSC{Expression,Int64} with 8 stored entries: [1, 1] = σ() * -1 [2, 1] = ρ() - z(t()) [3, 1] = y(t()) [1, 2] = σ() [2, 2] = Constant(-1) [3, 2] = x(t()) [2, 3] = x(t()) * -1 [3, 3] = -1 * β()

Dependent Variables, Functions, Chain Rule

"Variables" are overloaded. When you are solving a differential equation, the variable u(t) is actually a function of time. In order to handle these kinds of variables in a mathematically correct and extensible manner, the ModelingToolkit IR actually treats variables as functions, and constant variables are simply 0-ary functions (t()).

We can utilize this idea to have parameters that are also functions. For example, we can have a parameter σ which acts as a function of 1 argument, and then utilize this function within our differential equations:

@parameters σ(..)
eqs = [D(x) ~ σ(t-1)*(y-x),
       D(y) ~ x*(σ(t^2)-z)-y,
       D(z) ~ x*y - β*z]

0.9s

Julia

3-element Array{Equation,1}: Equation(Differential(x(t())), σ(t() - 1) * (y(t()) - x(t()))) Equation(Differential(y(t())), x(t()) * (σ(t() ^ 2) - z(t())) - y(t())) Equation(Differential(z(t())), x(t()) * y(t()) - β() * z(t()))

Notice that when we calculate the derivative with respect to t, the chain rule is automatically handled:

@derivatives Dₜ'~t
Dₜ(x*(σ(t^2)-z)-y)
expand_derivatives(Dₜ(x*(σ(t^2)-z)-y))

0.5s

Julia

(σ(t() ^ 2) - z(t())) * Differential(x(t())) + x(t()) * (Differential(σ(t() ^ 2)) + -1 * Differential(z(t()))) + -1 * Differential(y(t()))

Hackability: Extend directly from the language

ModelingToolkit.jl is written in Julia, and thus it can be directly extended from Julia itself. Let's define a normal Julia function and call it with a variable:

_f(x) = 2x + x^2
_f(x)

0.2s

Julia

2 * x(t()) + x(t()) ^ 2

Recall that when we do that, it will automatically trace this function and then build a symbolic expression. But what if we wanted our function to be a primative in the symbolic framework? This can be done by registering the function.

f(x) = 2x + x^2
@register f(x)

2.3s

Julia

f (generic function with 2 methods)

Now this function is a new primitive:

f(x)

0.2s

Julia

f(x(t()))

and we can now define derivatives of our function:

function ModelingToolkit.derivative(::typeof(f), args::NTuple{1,Any}, ::Val{1})
    2 + 2args[1]
end
expand_derivatives(Dx(f(x)))

0.2s

Julia

2 + 2 * x(t())