Merge pull request #94 from JuliaDiffEq/hg/refactor/systems

Refactor system functions

Author: ChrisRackauckas

README.md

@@ -56,24 +56,16 @@ This can then generate the function. For example, we can see the
 generated code via:
 
 ```julia
-ModelingToolkit.generate_ode_function(de)
+generate_function(de)
 
 ## Which returns:
-:((du, u, p, t)->begin
-                x = u[1]
-                y = u[2]
-                z = u[3]
-                σ = p[1]
-                ρ = p[2]
-                β = p[3]
-                x_t = σ * (y - x)
-                y_t = x * (ρ - z) - y
-                z_t = x * y - β * z
-                du[1] = x_t
-                du[2] = y_t
-                du[3] = z_t
-            end
-        end)
+:((##363, u, p, t)->begin
+          let (x, y, z, σ, ρ, β) = (u[1], u[2], u[3], p[1], p[2], p[3])
+              ##363[1] = σ * (y - x)
+              ##363[2] = x * (ρ - z) - y
+              ##363[3] = x * y - β * z
+          end
+      end)
 ```
 
 and get the generated function via:
@@ -97,25 +89,19 @@ eqs = [0 ~ σ*(y-x),
        0 ~ x*(ρ-z)-y,
        0 ~ x*y - β*z]
 ns = NonlinearSystem(eqs)
-nlsys_func = ModelingToolkit.generate_nlsys_function(ns)
+nlsys_func = generate_function(ns)
 ```
 
 which generates:
 
 ```julia
-(du, u, p)->begin  # C:\Users\Chris\.julia\v0.6\ModelingToolkit\src\systems.jl, line 51:
-        begin  # C:\Users\Chris\.julia\v0.6\ModelingToolkit\src\utils.jl, line 2:
-            y = u[1]
-            x = u[2]
-            z = u[3]
-            σ = p[1]
-            ρ = p[2]
-            β = p[3]
-            resid[1] = σ * (y - x)
-            resid[2] = x * (ρ - z) - y
-            resid[3] = x * y - β * z
-        end
-    end
+:((##364, u, p)->begin
+          let (x, z, y, ρ, σ, β) = (u[1], u[2], u[3], p[1], p[2], p[3])
+              ##364[1] = σ * (y - x)
+              ##364[2] = x * (ρ - z) - y
+              ##364[3] = x * y - β * z
+          end
+      end)
 ```
 
 We can use this to build a nonlinear function for use with NLsolve.jl:
@@ -287,26 +273,19 @@ eqs = [0 ~ σ*a,
        0 ~ x*(ρ-z)-y,
        0 ~ x*y - β*z]
 ns = NonlinearSystem(eqs,[x,y,z],[σ,ρ,β])
-nlsys_func = ModelingToolkit.generate_nlsys_function(ns)
+nlsys_func = generate_function(ns)
 ```
 
 expands to:
 
 ```julia
-:((du, u, p)->begin  # C:\Users\Chris\.julia\v0.6\ModelingToolkit\src\systems.jl, line 85:
-            begin  # C:\Users\Chris\.julia\v0.6\ModelingToolkit\src\utils.jl, line 2:
-                x = u[1]
-                y = u[2]
-                z = u[3]
-                σ = p[1]
-                ρ = p[2]
-                β = p[3]
-                a = y - x
-                resid[1] = σ * a
-                resid[2] = x * (ρ - z) - y
-                resid[3] = x * y - β * z
-            end
-        end)
+:((##365, u, p)->begin
+          let (x, y, z, σ, ρ, β) = (u[1], u[2], u[3], p[1], p[2], p[3])
+              ##365[1] = σ * (y - x)
+              ##365[2] = x * (ρ - z) - y
+              ##365[3] = x * y - β * z
+          end
+      end)
 ```
 
 In addition, the Jacobian calculations take into account intermediate variables

src/ModelingToolkit.jl

@@ -9,7 +9,6 @@ import MacroTools: splitdef, combinedef
 abstract type Expression <: Number end
 abstract type AbstractOperation <: Expression end
 abstract type AbstractComponent <: Expression end
-abstract type AbstractSystem end
 
 include("variables.jl")
 
@@ -25,6 +24,7 @@ function caclulate_jacobian end
 include("operations.jl")
 include("differentials.jl")
 include("equations.jl")
+include("systems/systems.jl")
 include("systems/diffeqs/diffeqsystem.jl")
 include("systems/diffeqs/first_order_transform.jl")
 include("systems/nonlinear/nonlinear_system.jl")

src/systems/diffeqs/diffeqsystem.jl

@@ -1,3 +1,6 @@
+export DiffEqSystem, ODEFunction
+
+
 using Base: RefValue
 
 
@@ -42,34 +45,7 @@ function DiffEqSystem(eqs, iv)
 end
 
 
-function generate_ode_function(sys::DiffEqSystem; version::FunctionVersion = ArrayFunction)
-    var_exprs = [:($(sys.dvs[i].name) = u[$i]) for i in eachindex(sys.dvs)]
-    param_exprs = [:($(sys.ps[i].name) = p[$i]) for i in eachindex(sys.ps)]
-    sys_exprs = build_equals_expr.(sys.eqs)
-    if version === ArrayFunction
-        dvar_exprs = [:(du[$i] = $(Symbol("$(sys.dvs[i].name)_$(sys.iv.name)"))) for i in eachindex(sys.dvs)]
-        exprs = vcat(var_exprs,param_exprs,sys_exprs,dvar_exprs)
-        block = expr_arr_to_block(exprs)
-        :((du,u,p,t)->$(toexpr(block)))
-    elseif version === SArrayFunction
-        dvar_exprs = [:($(Symbol("$(sys.dvs[i].name)_$(sys.iv.name)"))) for i in eachindex(sys.dvs)]
-        svector_expr = quote
-            E = eltype(tuple($(dvar_exprs...)))
-            T = StaticArrays.similar_type(typeof(u), E)
-            T($(dvar_exprs...))
-        end
-        exprs = vcat(var_exprs,param_exprs,sys_exprs,svector_expr)
-        block = expr_arr_to_block(exprs)
-        :((u,p,t)->$(toexpr(block)))
-    end
-end
-
-function build_equals_expr(eq::DiffEq)
-    lhs = Symbol(eq.var.name, :_, eq.D.x.name)
-    return :($lhs = $(convert(Expr, eq.rhs)))
-end
-
-function calculate_jacobian(sys::DiffEqSystem, simplify=true)
+function calculate_jacobian(sys::DiffEqSystem)
     isempty(sys.jac[]) || return sys.jac[]  # use cached Jacobian, if possible
     rhs = [eq.rhs for eq in sys.eqs]
 
@@ -78,20 +54,19 @@ function calculate_jacobian(sys::DiffEqSystem, simplify=true)
     return jac
 end
 
-function generate_ode_jacobian(sys::DiffEqSystem, simplify=true)
-    var_exprs = [:($(sys.dvs[i].name) = u[$i]) for i in eachindex(sys.dvs)]
-    param_exprs = [:($(sys.ps[i].name) = p[$i]) for i in eachindex(sys.ps)]
-    jac = calculate_jacobian(sys, simplify)
-    jac_exprs = [:(J[$i,$j] = $(convert(Expr, jac[i,j]))) for i in 1:size(jac,1), j in 1:size(jac,2)]
-    exprs = vcat(var_exprs,param_exprs,vec(jac_exprs))
-    block = expr_arr_to_block(exprs)
-    :((J,u,p,t)->$(block))
+function generate_jacobian(sys::DiffEqSystem; version::FunctionVersion = ArrayFunction)
+    jac = calculate_jacobian(sys)
+    return build_function(jac, sys.dvs, sys.ps, (sys.iv.name,); version = version)
+end
+
+function generate_function(sys::DiffEqSystem; version::FunctionVersion = ArrayFunction)
+    rhss = [eq.rhs for eq ∈ sys.eqs]
+    return build_function(rhss, sys.dvs, sys.ps, (sys.iv.name,); version = version)
 end
 
-function generate_ode_iW(sys::DiffEqSystem, simplify=true)
-    var_exprs = [:($(sys.dvs[i].name) = u[$i]) for i in eachindex(sys.dvs)]
-    param_exprs = [:($(sys.ps[i].name) = p[$i]) for i in eachindex(sys.ps)]
-    jac = calculate_jacobian(sys, simplify)
+
+function generate_ode_iW(sys::DiffEqSystem, simplify=true; version::FunctionVersion = ArrayFunction)
+    jac = calculate_jacobian(sys)
 
     gam = Parameter(:gam)
 
@@ -110,25 +85,18 @@ function generate_ode_iW(sys::DiffEqSystem, simplify=true)
         iW_t = simplify_constants.(iW_t)
     end
 
-    iW_exprs = [:(iW[$i,$j] = $(convert(Expr, iW[i,j]))) for i in 1:size(iW,1), j in 1:size(iW,2)]
-    exprs = vcat(var_exprs,param_exprs,vec(iW_exprs))
-    block = expr_arr_to_block(exprs)
+    vs, ps = sys.dvs, sys.ps
+    iW_func   = build_function(iW  , vs, ps, (:gam,:t); version = version)
+    iW_t_func = build_function(iW_t, vs, ps, (:gam,:t); version = version)
 
-    iW_t_exprs = [:(iW[$i,$j] = $(convert(Expr, iW_t[i,j]))) for i in 1:size(iW_t,1), j in 1:size(iW_t,2)]
-    exprs = vcat(var_exprs,param_exprs,vec(iW_t_exprs))
-    block2 = expr_arr_to_block(exprs)
-    :((iW,u,p,gam,t)->$(block)),:((iW,u,p,gam,t)->$(block2))
+    return (iW_func, iW_t_func)
 end
 
 function DiffEqBase.ODEFunction(sys::DiffEqSystem; version::FunctionVersion = ArrayFunction)
-    expr = generate_ode_function(sys; version = version)
+    expr = generate_function(sys; version = version)
     if version === ArrayFunction
         ODEFunction{true}(eval(expr))
     elseif version === SArrayFunction
         ODEFunction{false}(eval(expr))
     end
 end
-
-
-export DiffEqSystem, ODEFunction
-export generate_ode_function

src/systems/nonlinear/nonlinear_system.jl

@@ -1,5 +1,18 @@
+export NonlinearSystem
+
+
+struct NLEq
+    rhs::Expression
+end
+function Base.convert(::Type{NLEq}, eq::Equation)
+    isequal(eq.lhs, Constant(0)) || return NLEq(eq.rhs - eq.lhs)
+    return NLEq(eq.rhs)
+end
+Base.:(==)(a::NLEq, b::NLEq) = a.rhs == b.rhs
+get_args(eq::NLEq) = Expression[eq.rhs]
+
 struct NonlinearSystem <: AbstractSystem
-    eqs::Vector{Equation}
+    eqs::Vector{NLEq}
     vs::Vector{Variable}
     ps::Vector{Variable}
 end
@@ -9,40 +22,19 @@ function NonlinearSystem(eqs)
     NonlinearSystem(eqs, vs, ps)
 end
 
-function generate_nlsys_function(sys::NonlinearSystem)
-    var_exprs = [:($(sys.vs[i].name) = u[$i]) for i in 1:length(sys.vs)]
-    param_exprs = [:($(sys.ps[i].name) = p[$i]) for i in 1:length(sys.ps)]
-    sys_eqs, calc_eqs = partition(eq -> isequal(eq.lhs, Constant(0)), sys.eqs)
-    calc_exprs = [:($(eq.lhs.name) = $(eq.rhs)) for eq in calc_eqs if isa(eq.lhs, Variable)]
-    sys_exprs = [:($(Symbol("resid[$i]")) = $(sys_eqs[i].rhs)) for i in eachindex(sys_eqs)]
 
-    exprs = vcat(var_exprs,param_exprs,calc_exprs,sys_exprs)
-    block = expr_arr_to_block(exprs)
-    :((du,u,p)->$(block))
+function calculate_jacobian(sys::NonlinearSystem)
+    rhs = [eq.rhs for eq in sys.eqs]
+    jac = expand_derivatives.(calculate_jacobian(rhs, sys.vs))
+    return jac
 end
 
-function calculate_jacobian(sys::NonlinearSystem,simplify=true)
-    sys_eqs, calc_eqs = partition(eq -> isequal(eq.lhs, Constant(0)), sys.eqs)
-    rhs = [eq.rhs for eq in sys_eqs]
-
-    for calc_eq ∈ calc_eqs
-        find_replace!.(rhs, calc_eq.lhs, calc_eq.rhs)
-    end
-
-    sys_exprs = calculate_jacobian(rhs,sys.vs)
-    sys_exprs = Expression[expand_derivatives(expr) for expr in sys_exprs]
-    sys_exprs
+function generate_jacobian(sys::NonlinearSystem; version::FunctionVersion = ArrayFunction)
+    jac = calculate_jacobian(sys)
+    return build_function(jac, sys.vs, sys.ps; version = version)
 end
 
-function generate_nlsys_jacobian(sys::NonlinearSystem,simplify=true)
-    var_exprs = [:($(sys.vs[i].name) = u[$i]) for i in 1:length(sys.vs)]
-    param_exprs = [:($(sys.ps[i].name) = p[$i]) for i in 1:length(sys.ps)]
-    jac = calculate_jacobian(sys,simplify)
-    jac_exprs = [:(J[$i,$j] = $(convert(Expr, jac[i,j]))) for i in 1:size(jac,1), j in 1:size(jac,2)]
-    exprs = vcat(var_exprs,param_exprs,vec(jac_exprs))
-    block = expr_arr_to_block(exprs)
-    :((J,u,p,t)->$(block))
+function generate_function(sys::NonlinearSystem; version::FunctionVersion = ArrayFunction)
+    rhss = [eq.rhs for eq ∈ sys.eqs]
+    return build_function(rhss, sys.vs, sys.ps; version = version)
 end
-
-export NonlinearSystem
-export generate_nlsys_function

src/systems/systems.jl

@@ -0,0 +1,30 @@
+export generate_jacobian, generate_function
+
+
+abstract type AbstractSystem end
+
+function generate_jacobian end
+function generate_function end
+
+function build_function(rhss, vs, ps, args = (); version::FunctionVersion)
+    var_pairs   = [(u.name, :(u[$i])) for (i, u) ∈ enumerate(vs)]
+    param_pairs = [(p.name, :(p[$i])) for (i, p) ∈ enumerate(ps)]
+    (ls, rs) = zip(var_pairs..., param_pairs...)
+
+    var_eqs = Expr(:(=), build_expr(:tuple, ls), build_expr(:tuple, rs))
+
+    if version === ArrayFunction
+        X = gensym()
+        sys_exprs = [:($X[$i] = $(convert(Expr, rhs))) for (i, rhs) ∈ enumerate(rhss)]
+        let_expr = Expr(:let, var_eqs, build_expr(:block, sys_exprs))
+        :(($X,u,p,$(args...)) -> $let_expr)
+    elseif version === SArrayFunction
+        sys_expr = build_expr(:tuple, [convert(Expr, rhs) for rhs ∈ rhss])
+        let_expr = Expr(:let, var_eqs, sys_expr)
+        :((u,p,$(args...)) -> begin
+            X = $let_expr
+            T = StaticArrays.similar_type(typeof(u), eltype(X))
+            T(X)
+        end)
+    end
+end

src/utils.jl

@@ -31,12 +31,9 @@ function flatten_expr!(x)
     x
 end
 
-toexpr(ex) = MacroTools.postwalk(x -> isa(x, Expression) ? convert(Expr, x) : x, ex)
-
 function partition(f, xs)
     idxs = map(f, xs)
-    not_idxs = eachindex(xs) .∉ (idxs,)
-    return (xs[idxs], xs[not_idxs])
+    return (xs[idxs], xs[(!).(idxs)])
 end
 
 is_constant(::Constant) = true