ModelingToolkit.jl, An IR and Compiler for Scientific Models

Chris Rackauckas

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:

using ModelingToolkit

### Define a differential equation system

@parameters t σ ρ β
@variables x(t) y(t) z(t)
@derivatives D'~t

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], [σ,ρ,β])

### Use in DifferentialEquations.jl

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

using Plots
plot(sol,vars=(1,2,3))

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], [σ,ρ,β])
:((##524, u, p, t)->begin
          #= C:\Users\Chris Rackauckas\.julia\dev\ModelingToolkit\src\utils.jl:44 =#
          let (x, y, z, σ, ρ, β) = (u[1], u[2], u[3], p[1], p[2], p[3])
              ##524[1] = (*)(σ, (-)(y, x))
              ##524[2] = (-)((*)(x, (-)(ρ, z)), y)
              ##524[3] = (-)((*)(x, y), (*)(β, z))
          end
      end)

A special syntax in DifferentialEquations.jl for small static ODE systems uses f(u,p,t), which can be generated as well:

generate_function(de, [x,y,z], [σ,ρ,β]; version=ModelingToolkit.SArrayFunction)
:((u, p, t)->begin
          #= C:\Users\Chris Rackauckas\.julia\dev\ModelingToolkit\src\utils.jl:48 =#
          #= C:\Users\Chris Rackauckas\.julia\dev\ModelingToolkit\src\utils.jl:49 =#
          X = let (x, y, z, σ, ρ, β) = (u[1], u[2], u[3], p[1], p[2], p[3])
                  ((*)(σ, (-)(y, x)), (-)((*)(x, (-)(ρ, z)), y), (-)((*)(x, y), (*)(β, z)))
              end
          #= C:\Users\Chris Rackauckas\.julia\dev\ModelingToolkit\src\utils.jl:50 =#
          T = StaticArrays.similar_type(typeof(u), eltype(X))
          #= C:\Users\Chris Rackauckas\.julia\dev\ModelingToolkit\src\utils.jl:51 =#
          T(X)
      end)

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

jac = calculate_jacobian(de)
3×3 Array{ModelingToolkit.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)
:((##525, u, p, t)->begin
          #= C:\Users\Chris Rackauckas\.julia\dev\ModelingToolkit\src\utils.jl:44 =#
          let (x, y, z, σ, ρ, β) = (u[1], u[2], u[3], p[1], p[2], p[3])
              ##525[1] = (*)(σ, -1)
              ##525[2] = (-)(ρ, z)
              ##525[3] = y
              ##525[4] = σ
              ##525[5] = -1
              ##525[6] = x
              ##525[7] = 0
              ##525[8] = (*)(x, -1)
              ##525[9] = (*)(-1, β)
          end
      end)

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]
:((##526, u, p, gam, t)->begin
          #= C:\Users\Chris Rackauckas\.julia\dev\ModelingToolkit\src\utils.jl:44 =#
          let (x, y, z, σ, ρ, β) = (u[1], u[2], u[3], p[1], p[2], p[3])
              ##526[1] = (+)((*)(σ, gam), true)
              ##526[2] = (*)(gam, (-)(ρ, z), -1, (inv)((+)((*)(σ, gam), true)))
              ##526[3] = (*)(gam, y, -1, (inv)((+)((*)(σ, gam), true)))
              ##526[4] = (*)(gam, σ, -1)
              ##526[5] = (-)((+)(gam, true), (*)(gam, (-)(ρ, z), gam, σ, (inv)((+)((*)(σ, gam), true))))
              ##526[6] = (*)((-)((*)(gam, x, -1), (*)(gam, y, gam, σ, (inv)((+)((*)(σ, gam), true)))), (inv)((-)((+)(gam, true), (
*)(gam, (-)(ρ, z), gam, σ, (inv)((+)((*)(σ, gam), true))))))
              ##526[7] = 0
              ##526[8] = (-)((*)(x, gam), 0)
              ##526[9] = (-)((-)((+)((*)(β, gam), true), 0), (*)((-)((*)(gam, x, -1), (*)(gam, y, gam, σ, (inv)((+)((*)(σ, gam), t
rue)))), (inv)((-)((+)(gam, true), (*)(gam, (-)(ρ, z), gam, σ, (inv)((+)((*)(σ, gam), true))))), (-)((*)(x, gam), 0)))
          end
      end)

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], [σ,ρ,β])
:((##528, u, p)->begin
          #= C:\Users\Chris Rackauckas\.julia\dev\ModelingToolkit\src\utils.jl:44 =#
          let (x, y, z, σ, ρ, β) = (u[1], u[2], u[3], p[1], p[2], p[3])
              ##528[1] = (*)(σ, (-)(y, x))
              ##528[2] = (-)((*)(x, (-)(ρ, z)), y)
              ##528[3] = (-)((*)(x, y), (*)(β, z))
          end
      end)

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

nl_f = @eval eval(nlsys_func)
# Make a closure over the parameters for for NLsolve.jl
f2 = (du,u) -> nl_f(du,u,(10.0,26.0,2.33))

using NLsolve
nlsolve(f2,ones(3))
Results of Nonlinear Solver Algorithm
 * Algorithm: Trust-region with dogleg and autoscaling
 * Starting Point: [1.0, 1.0, 1.0]
 * Zero: [2.2228e-10, 2.2228e-10, -9.99034e-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

Library of 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 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)
ModelingToolkit.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, ModelingToolkit.Variable[u, x, u_tt, u_t
, x_t], ModelingToolkit.Variable[], Base.RefValue{Array{ModelingToolkit.Expression,2}}(Array{Expression}(0,0)))
de1.eqs
5-element Array{ModelingToolkit.DiffEq,1}:
 ModelingToolkit.DiffEq(u_tt, 1, ((2 * u_tt(t()) + u_t(t())) + x_t(t())) + 1)
 ModelingToolkit.DiffEq(x_t, 1, x_t(t()) + 2)                                
 ModelingToolkit.DiffEq(u_t, 1, u_tt(t()))                                   
 ModelingToolkit.DiffEq(u, 1, u_t(t()))                                      
 ModelingToolkit.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.

Linear Algebra... for free?

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)]
3×3 Array{ModelingToolkit.Operation,2}:
          (D'~x(t()))(σ() * (y(t()) - x(t())))  …           (D'~z(t()))(σ() * (y(t()) - x(t())))
 (D'~x(t()))(x(t()) * (ρ() - z(t())) - y(t()))     (D'~z(t()))(x(t()) * (ρ() - z(t())) - y(t()))
   (D'~x(t()))(x(t()) * y(t()) - β() * z(t()))       (D'~z(t()))(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:

J = expand_derivatives.(J)
3×3 Array{ModelingToolkit.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)
LinearAlgebra.LU{ModelingToolkit.Expression,Array{ModelingToolkit.Expression,2}}
L factor:
3×3 Array{ModelingToolkit.Expression,2}:
                    Constant(1)  …  Constant(0)
 (ρ() - z(t())) * inv(σ() * -1)     identity(0)
         y(t()) * inv(σ() * -1)     Constant(1)
U factor:
3×3 Array{ModelingToolkit.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
3×3 Array{ModelingToolkit.Expression,2}:
                    Constant(1)  …  Constant(0)
 (ρ() - z(t())) * inv(σ() * -1)     identity(0)
         y(t()) * inv(σ() * -1)     Constant(1)

and the inverse?

invJ = inv(J)
3×3 Array{ModelingToolkit.Operation,2}:
 (σ() * -1) \ ((identity(true) - identity(0) * (((-1 * β() - (y(t()) * inv(σ() * -1)) * identity(0)) - ((x(t()) - (y(t()) * inv(σ(
) * -1)) * σ()) * inv(identity(-1) - ((ρ() - z(t())) * inv(σ() * -1)) * σ())) * (x(t()) * -1 - ((ρ() - z(t())) * inv(σ() * -1)) * 
identity(0))) \ ((identity(0) - (y(t()) * inv(σ() * -1)) * identity(true)) - ((x(t()) - (y(t()) * inv(σ() * -1)) * σ()) * inv(iden
tity(-1) - ((ρ() - z(t())) * inv(σ() * -1)) * σ())) * (identity(0) - ((ρ() - z(t())) * inv(σ() * -1)) * identity(true))))) - σ() *
 ((identity(-1) - ((ρ() - z(t())) * inv(σ() * -1)) * σ()) \ ((identity(0) - ((ρ() - z(t())) * inv(σ() * -1)) * identity(true)) - (
x(t()) * -1 - ((ρ() - z(t())) * inv(σ() * -1)) * identity(0)) * (((-1 * β() - (y(t()) * inv(σ() * -1)) * identity(0)) - ((x(t()) -
 (y(t()) * inv(σ() * -1)) * σ()) * inv(identity(-1) - ((ρ() - z(t())) * inv(σ() * -1)) * σ())) * (x(t()) * -1 - ((ρ() - z(t())) * 
inv(σ() * -1)) * identity(0))) \ ((identity(0) - (y(t()) * inv(σ() * -1)) * identity(true)) - ((x(t()) - (y(t()) * inv(σ() * -1)) 
* σ()) * inv(identity(-1) - ((ρ() - z(t())) * inv(σ() * -1)) * σ())) * (identity(0) - ((ρ() - z(t())) * inv(σ() * -1)) * identity(
true)))))))  …  (σ() * -1) \ ((identity(0) - identity(0) * (((-1 * β() - (y(t()) * inv(σ() * -1)) * identity(0)) - ((x(t()) - (y(t
()) * inv(σ() * -1)) * σ()) * inv(identity(-1) - ((ρ() - z(t())) * inv(σ() * -1)) * σ())) * (x(t()) * -1 - ((ρ() - z(t())) * inv(σ
() * -1)) * identity(0))) \ ((identity(true) - (y(t()) * inv(σ() * -1)) * identity(0)) - ((x(t()) - (y(t()) * inv(σ() * -1)) * σ()
) * inv(identity(-1) - ((ρ() - z(t())) * inv(σ() * -1)) * σ())) * (identity(0) - ((ρ() - z(t())) * inv(σ() * -1)) * identity(0))))
) - σ() * ((identity(-1) - ((ρ() - z(t())) * inv(σ() * -1)) * σ()) \ ((identity(0) - ((ρ() - z(t())) * inv(σ() * -1)) * identity(0
)) - (x(t()) * -1 - ((ρ() - z(t())) * inv(σ() * -1)) * identity(0)) * (((-1 * β() - (y(t()) * inv(σ() * -1)) * identity(0)) - ((x(
t()) - (y(t()) * inv(σ() * -1)) * σ()) * inv(identity(-1) - ((ρ() - z(t())) * inv(σ() * -1)) * σ())) * (x(t()) * -1 - ((ρ() - z(t(
))) * inv(σ() * -1)) * identity(0))) \ ((identity(true) - (y(t()) * inv(σ() * -1)) * identity(0)) - ((x(t()) - (y(t()) * inv(σ() *
 -1)) * σ()) * inv(identity(-1) - ((ρ() - z(t())) * inv(σ() * -1)) * σ())) * (identity(0) - ((ρ() - z(t())) * inv(σ() * -1)) * ide
ntity(0)))))))
                                                                                                                                  
                                                                                                                                  
                                                                                                                                  
                                                                                                                                  
    (identity(-1) - ((ρ() - z(t())) * inv(σ() * -1)) * σ()) \ ((identity(0) - ((ρ() - z(t())) * inv(σ() * -1)) * identity(true)) -
 (x(t()) * -1 - ((ρ() - z(t())) * inv(σ() * -1)) * identity(0)) * (((-1 * β() - (y(t()) * inv(σ() * -1)) * identity(0)) - ((x(t())
 - (y(t()) * inv(σ() * -1)) * σ()) * inv(identity(-1) - ((ρ() - z(t())) * inv(σ() * -1)) * σ())) * (x(t()) * -1 - ((ρ() - z(t())) 
* inv(σ() * -1)) * identity(0))) \ ((identity(0) - (y(t()) * inv(σ() * -1)) * identity(true)) - ((x(t()) - (y(t()) * inv(σ() * -1)
) * σ()) * inv(identity(-1) - ((ρ() - z(t())) * inv(σ() * -1)) * σ())) * (identity(0) - ((ρ() - z(t())) * inv(σ() * -1)) * identit
y(true)))))                                                                                                                       
                                                                                                                                  
                                                                                                                                  
                                                                                                                                  
             (identity(-1) - ((ρ() - z(t())) * inv(σ() * -1)) * σ()) \ ((identity(0) - ((ρ() - z(t())) * inv(σ() * -1)) * identity
(0)) - (x(t()) * -1 - ((ρ() - z(t())) * inv(σ() * -1)) * identity(0)) * (((-1 * β() - (y(t()) * inv(σ() * -1)) * identity(0)) - ((
x(t()) - (y(t()) * inv(σ() * -1)) * σ()) * inv(identity(-1) - ((ρ() - z(t())) * inv(σ() * -1)) * σ())) * (x(t()) * -1 - ((ρ() - z(
t())) * inv(σ() * -1)) * identity(0))) \ ((identity(true) - (y(t()) * inv(σ() * -1)) * identity(0)) - ((x(t()) - (y(t()) * inv(σ()
 * -1)) * σ()) * inv(identity(-1) - ((ρ() - z(t())) * inv(σ() * -1)) * σ())) * (identity(0) - ((ρ() - z(t())) * inv(σ() * -1)) * i
dentity(0)))))
                                                                                                                                  
                                                                                                                                  
                                                                                                                                  
                                                                                                                                  
                                                                                                                                  
                                                                     ((-1 * β() - (y(t()) * inv(σ() * -1)) * identity(0)) - ((x(t(
)) - (y(t()) * inv(σ() * -1)) * σ()) * inv(identity(-1) - ((ρ() - z(t())) * inv(σ() * -1)) * σ())) * (x(t()) * -1 - ((ρ() - z(t())
) * inv(σ() * -1)) * identity(0))) \ ((identity(0) - (y(t()) * inv(σ() * -1)) * identity(true)) - ((x(t()) - (y(t()) * inv(σ() * -
1)) * σ()) * inv(identity(-1) - ((ρ() - z(t())) * inv(σ() * -1)) * σ())) * (identity(0) - ((ρ() - z(t())) * inv(σ() * -1)) * ident
ity(true)))                                                                                                                       
                                                                                                                                  
                                                                                                                                  
                                                                                                                                  
                                                                                                                                  
                                                                           ((-1 * β() - (y(t()) * inv(σ() * -1)) * identity(0)) - 
((x(t()) - (y(t()) * inv(σ() * -1)) * σ()) * inv(identity(-1) - ((ρ() - z(t())) * inv(σ() * -1)) * σ())) * (x(t()) * -1 - ((ρ() - 
z(t())) * inv(σ() * -1)) * identity(0))) \ ((identity(true) - (y(t()) * inv(σ() * -1)) * identity(0)) - ((x(t()) - (y(t()) * inv(σ
() * -1)) * σ()) * inv(identity(-1) - ((ρ() - z(t())) * inv(σ() * -1)) * σ())) * (identity(0) - ((ρ() - z(t())) * inv(σ() * -1)) *
 identity(0)))

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

Let's follow this thread a little deeper.

Automatically convert numerical codes to symbolic

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
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)
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)
3×3 Array{ModelingToolkit.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 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_IJ_sorted!(I, J, V, size(M)...)
end
sJ = SparseMatrixCSC(J)
3×3 SparseArrays.SparseMatrixCSC{ModelingToolkit.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]
3-element Array{ModelingToolkit.Equation,1}:
 ModelingToolkit.Equation((D'~t())(x(t())), σ(t() - 1) * (y(t()) - x(t())))         
 ModelingToolkit.Equation((D'~t())(y(t())), x(t()) * (σ(t() ^ 2) - z(t())) - y(t()))
 ModelingToolkit.Equation((D'~t())(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))
(σ(t() ^ 2) - z(t())) * (D'~t())(x(t())) + x(t()) * ((D'~t())(σ(t() ^ 2)) + -1 * (D'~t())(z(t()))) + -1 * (D'~t())(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)
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)
f (generic function with 2 methods)

Now this function is a new primitive:

f(x)
Main.WeaveSandBox14.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)))
2 + 2 * x(t())

Use Case: PuMaS.jl

Let's look at PuMaS.jl, a software for simulating and estimating nonlinear mixed effects models in pharmacometrics. In these models, a user brings in Electronic Health Record (EHR) data to define patient covariates, and then a population model defines fixed effects and random effects. These are collated together to give dynamical parameters which define a differential equation, and then derived variables from the differential equation are observed. For example, let's take a look at a multiple response model:

using PuMaS

subject = process_nmtran(example_nmtran_data("event_data/data23"),
                         [], [:ev1,:cp,:periph,:resp])[1]


mrm = @model begin
    @param   θ  VectorDomain(12)
    @random  η ~ MvNormal(Matrix{Float64}(I, 11, 11))

    @pre begin
        Ka1     = θ[1]
        CL      = θ[2]*exp(η[1])
        Vc      = θ[3]*exp(η[2])
        Q       = θ[4]*exp(η[3])
        Vp      = θ[5]*exp(η[4])
        Kin     = θ[6]*exp(η[5])
        Kout    = θ[7]*exp(η[6])
        IC50    = θ[8]*exp(η[7])
        IMAX    = θ[9]*exp(η[8])
        γ       = θ[10]*exp(η[9])
        Vmax    = θ[11]*exp(η[10])
        Km      = θ[12]*exp(η[11])
    end

    @init begin
        Resp = θ[6]/θ[7]
    end

    @dynamics begin
        Ev1'    = -Ka1*Ev1
        Cent'   =  Ka1*Ev1 - (CL+Vmax/(Km+(Cent/Vc))+Q)*(Cent/Vc)  + Q*(Periph/Vp)
        Periph' =  Q*(Cent/Vc)  - Q*(Periph/Vp)
        Resp'   =  Kin*(1-(IMAX*(Cent/Vc)^γ/(IC50^γ+(Cent/Vc)^γ)))  - Kout*Resp
    end

    @derived begin
        ev1    = Ev1
        cp     = Cent / θ[3]
        periph = Periph
        resp   = Resp
    end
end
PuMaSModel
  Parameters: θ
  Random effects: η
  Covariates: 
  Dynamical variables: Ev1, Cent, Periph, Resp
  Derived: ev1, cp, periph, resp
  Observed: ev1, cp, periph, resp

Notice that

@dynamics begin
  Ev1'    = -Ka1*Ev1
  Cent'   =  Ka1*Ev1 - (CL+Vmax/(Km+(Cent/Vc))+Q)*(Cent/Vc)  + Q*(Periph/Vp)
  Periph' =  Q*(Cent/Vc)  - Q*(Periph/Vp)
  Resp'   =  Kin*(1-(IMAX*(Cent/Vc)^γ/(IC50^γ+(Cent/Vc)^γ)))  - Kout*Resp
end

is implemented by translating this differential equation system into ModelingToolkit IR, and then the functionality shown above is used to get the Julia functions for numerical simulation.

We can now set parameters and simulate the system for a given patient:

param = (θ = [
              1, # Ka1  Absorption rate constant 1 (1/time)
              1, # CL   Clearance (volume/time)
              20, # Vc   Central volume (volume)
              2, # Q    Inter-compartmental clearance (volume/time)
              10, # Vp   Peripheral volume of distribution (volume)
              10, # Kin  Response in rate constant (1/time)
              2, # Kout Response out rate constant (1/time)
              2, # IC50 Concentration for 50% of max inhibition (mass/volume)
              1, # IMAX Maximum inhibition
              1, # γ    Emax model sigmoidicity
              0, # Vmax Maximum reaction velocity (mass/time)
              2  # Km   Michaelis constant (mass/volume)
              ],)
randeffs = (η = zeros(11),)
sim = simobs(mrm, subject, param, randeffs)
PuMaS.SimulatedObservations{PuMaS.Subject{NamedTuple{(:ev1, :cp, :periph, :resp),NTuple{4,Array{Union{Missing, Float64},1}}},Nothi
ng,Array{PuMaS.Event{Float64,Float64,Float64,Float64,Float64,Float64,Int64},1},Array{Float64,1}},Array{Float64,1},NamedTuple{(:ev1
, :cp, :periph, :resp),NTuple{4,Array{Float64,1}}}}(Subject
  ID: 1
  Events: 4
  Observables: ev1: (n=2401), cp: (n=2401), periph: (n=2401), resp: (n=2401)
, [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9  …  239.1, 239.2, 239.3, 239.4, 239.5, 239.6, 239.7, 239.8, 239.9, 240.0], (ev
1 = [100.0, 90.4837, 81.8731, 74.0818, 67.032, 60.6531, 54.8812, 49.6585, 44.9329, 40.657  …  -3.08978e-43, -1.91753e-43, -9.03806
e-44, -4.86243e-45, 6.48019e-44, 1.18612e-43, 1.56569e-43, 1.78672e-43, 1.84921e-43, 1.75316e-43], cp = [0.0, 0.472219, 0.892564, 
1.26616, 1.59765, 1.89122, 2.15065, 2.37936, 2.58045, 2.75668  …  0.0286307, 0.0285409, 0.0284515, 0.0283623, 0.0282733, 0.0281847
, 0.0280963, 0.0280083, 0.0279205, 0.027833], periph = [0.0, 0.0478097, 0.182936, 0.393897, 0.670403, 1.00324, 1.38417, 1.80581, 2
.26161, 2.74568  …  0.339601, 0.338536, 0.337474, 0.336416, 0.335362, 0.33431, 0.333262, 0.332218, 0.331176, 0.330139], resp = [5.
0, 4.90282, 4.68845, 4.42607, 4.15099, 3.88189, 3.62839, 3.39518, 3.18393, 2.99469  …  4.92832, 4.92855, 4.92877, 4.92899, 4.92921
, 4.92943, 4.92964, 4.92986, 4.93008, 4.93029]))

The solution to this model gives the concentration profile for each subject. For example:

plot(sim)

PuMaS's features can then be used to learn the fixed and random effects from data, and from this know how each individual's drug matabolism will react differently. Using these results, we can predict optimal drug doses specific to an individual's biology. This mechanism and the software validation is currently undergoing clinical trials at Johns Hopkins Medical Center.

Appendix

This tutorial is part of the DiffEqTutorials.jl repository, found at: https://github.com/JuliaDiffEq/DiffEqTutorials.jl

To locally run this tutorial, do the following commands:

using DiffEqTutorials
DiffEqTutorials.weave_file("ode_extras","ModelingToolkit.jmd")

Computer Information:

Julia Version 1.1.0
Commit 80516ca202 (2019-01-21 21:24 UTC)
Platform Info:
  OS: Windows (x86_64-w64-mingw32)
  CPU: Intel(R) Core(TM) i5-7200U CPU @ 2.50GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-6.0.1 (ORCJIT, skylake)

Package Information:

Status `~\.julia\environments\v1.1\Project.toml`
[28f2ccd6-bb30-5033-b560-165f7b14dc2f] ApproxFun 0.11.1
[c52e3926-4ff0-5f6e-af25-54175e0327b1] Atom 0.8.5
[aae01518-5342-5314-be14-df237901396f] BandedMatrices 0.9.1
[6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf] BenchmarkTools 0.4.2
[336ed68f-0bac-5ca0-87d4-7b16caf5d00b] CSV 0.4.3
[49dc2e85-a5d0-5ad3-a950-438e2897f1b9] Calculus 0.4.1
[a93c6f00-e57d-5684-b7b6-d8193f3e46c0] DataFrames 0.18.2
[176a2513-9382-508e-b587-15b5aca5cd00] DataInterpolations 0.0.0
[bcd4f6db-9728-5f36-b5f7-82caef46ccdb] DelayDiffEq 5.2.0+
[39dd38d3-220a-591b-8e3c-4c3a8c710a94] Dierckx 0.4.1
[2b5f629d-d688-5b77-993f-72d75c75574e] DiffEqBase 5.7.0
[459566f4-90b8-5000-8ac3-15dfb0a30def] DiffEqCallbacks 2.5.2+
[f3b72e0c-5b89-59e1-b016-84e28bfd966d] DiffEqDevTools 2.8.0
[01453d9d-ee7c-5054-8395-0335cb756afa] DiffEqDiffTools 0.8.1
[aae7a2af-3d4f-5e19-a356-7da93b79d9d0] DiffEqFlux 0.5.0
[c894b116-72e5-5b58-be3c-e6d8d4ac2b12] DiffEqJump 6.1.1+
[78ddff82-25fc-5f2b-89aa-309469cbf16f] DiffEqMonteCarlo 0.14.0
[77a26b50-5914-5dd7-bc55-306e6241c503] DiffEqNoiseProcess 3.2.0
[9fdde737-9c7f-55bf-ade8-46b3f136cc48] DiffEqOperators 3.5.0+
[1130ab10-4a5a-5621-a13d-e4788d82bd4c] DiffEqParamEstim 1.6.0
[a077e3f3-b75c-5d7f-a0c6-6bc4c8ec64a9] DiffEqProblemLibrary 4.1.0
[41bf760c-e81c-5289-8e54-58b1f1f8abe2] DiffEqSensitivity 3.2.2
[6d1b261a-3be8-11e9-3f2f-0b112a9a8436] DiffEqTutorials 0.1.0
[163ba53b-c6d8-5494-b064-1a9d43ac40c5] DiffResults 0.0.4
[0c46a032-eb83-5123-abaf-570d42b7fbaa] DifferentialEquations 6.3.0
[31c24e10-a181-5473-b8eb-7969acd0382f] Distributions 0.18.0
[e30172f5-a6a5-5a46-863b-614d45cd2de4] Documenter 0.22.3
[4f61f5a4-77b1-5117-aa51-3ab5ef4ef0cd] FFTViews 0.2.0
[7a1cc6ca-52ef-59f5-83cd-3a7055c09341] FFTW 0.2.4
[53c48c17-4a7d-5ca2-90c5-79b7896eea93] FixedPointNumbers 0.5.3
[587475ba-b771-5e3f-ad9e-33799f191a9c] Flux 0.8.3+
[f6369f11-7733-5829-9624-2563aa707210] ForwardDiff 0.10.3
[7073ff75-c697-5162-941a-fcdaad2a7d2a] IJulia 1.18.1
[c601a237-2ae4-5e1e-952c-7a85b0c7eef1] Interact 0.10.2
[b6b21f68-93f8-5de0-b562-5493be1d77c9] Ipopt 0.5.4
[1c8ee90f-4401-5389-894e-7a04a3dc0f4d] IterableTables 0.11.0
[4076af6c-e467-56ae-b986-b466b2749572] JuMP 0.19.0
[e5e0dc1b-0480-54bc-9374-aad01c23163d] Juno 0.7.0
[5ab0869b-81aa-558d-bb23-cbf5423bbe9b] KernelDensity 0.5.1
[2ee39098-c373-598a-b85f-a56591580800] LabelledArrays 0.7.1
[eb30cadb-4394-5ae3-aed4-317e484a6458] MLDatasets 0.3.0
[eff96d63-e80a-5855-80a2-b1b0885c5ab7] Measurements 2.0.0
[961ee093-0014-501f-94e3-6117800e7a78] ModelingToolkit 0.2.0+
[0987c9cc-fe09-11e8-30f0-b96dd679fdca] MonteCarloMeasurements 0.1.4
[76087f3c-5699-56af-9a33-bf431cd00edd] NLopt 0.5.1
[2774e3e8-f4cf-5e23-947b-6d7e65073b56] NLsolve 4.0.0
[c030b06c-0b6d-57c2-b091-7029874bd033] ODE 2.4.0
[09606e27-ecf5-54fc-bb29-004bd9f985bf] ODEInterfaceDiffEq 3.2.0
[961449f6-2712-11e9-291b-8fec89bf7786] Olas 0.1.0
[429524aa-4258-5aef-a3af-852621145aeb] Optim 0.18.1
[1dea7af3-3e70-54e6-95c3-0bf5283fa5ed] OrdinaryDiffEq 5.6.0
[2dcacdae-9679-587a-88bb-8b444fb7085b] ParallelDataTransfer 0.5.0
[65888b18-ceab-5e60-b2b9-181511a3b968] ParameterizedFunctions 4.1.1+
[91a5bcdd-55d7-5caf-9e0b-520d859cae80] Plots 0.24.0
[4ce49e30-a246-11e8-284f-6b410f95840d] ProgressLogging 0.1.0
[71ad9d73-34c4-5ce9-b7b1-f7bd31ac38ba] PuMaS 0.0.0
[1fd47b50-473d-5c70-9696-f719f8f3bcdc] QuadGK 2.0.3
[731186ca-8d62-57ce-b412-fbd966d074cd] RecursiveArrayTools 0.20.0
[295af30f-e4ad-537b-8983-00126c2a3abe] Revise 2.1.3
[c4c386cf-5103-5370-be45-f3a111cca3b8] Rsvg 0.2.3
[05bca326-078c-5bf0-a5bf-ce7c7982d7fd] SimpleDiffEq 0.3.1+
[60ddc479-9b66-56df-82fc-76a74619b69c] StatPlots 0.9.2
[90137ffa-7385-5640-81b9-e52037218182] StaticArrays 0.10.3
[2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91] StatsBase 0.30.0
[789caeaf-c7a9-5a7d-9973-96adeb23e2a0] StochasticDiffEq 6.2.0
[c3572dad-4567-51f8-b174-8c6c989267f4] Sundials 3.4.0
[9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c] Tracker 0.2.0
[84d833dd-6860-57f9-a1a7-6da5db126cff] TransformVariables 0.3.2
[1986cc42-f94f-5a68-af5c-568840ba703d] Unitful 0.15.0
[44d3d7a6-8a23-5bf8-98c5-b353f8df5ec9] Weave 0.9.0
[0f1e0344-ec1d-5b48-a673-e5cf874b6c29] WebIO 0.8.1
[e88e6eb3-aa80-5325-afca-941959d7151f] Zygote 0.3.0