Refactor function generation

Closes #93.

Author: HarrisonGrodin

README.md

@@ -60,20 +60,11 @@ ModelingToolkit.generate_ode_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)
+          du .= 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
+      end)
+
 ```
 
 and get the generated function via:
@@ -103,19 +94,11 @@ nlsys_func = ModelingToolkit.generate_nlsys_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
+:((du, u, p)->begin
+          du .= let (y, x, z, σ, ρ, β) = (u[1], u[2], u[3], p[1], p[2], p[3])
+                  (σ * (y - x), x * (ρ - z) - y, x * y - β * z)
+              end
+      end)
 ```
 
 We can use this to build a nonlinear function for use with NLsolve.jl:
@@ -293,20 +276,12 @@ nlsys_func = ModelingToolkit.generate_nlsys_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)
+:((du, u, p)->begin
+          du .= 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
+      end)
+
 ```
 
 In addition, the Jacobian calculations take into account intermediate variables

src/systems/diffeqs/diffeqsystem.jl

@@ -43,24 +43,22 @@ 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)]
+    var_pairs   = [(u.name, :(u[$i])) for (i, u) ∈ enumerate(sys.dvs)]
+    param_pairs = [(p.name, :(p[$i])) for (i, p) ∈ enumerate(sys.ps )]
+    (ls, rs) = collect(zip(var_pairs..., param_pairs...))
+
+    var_eqs = Expr(:(=), build_expr(:tuple, ls), build_expr(:tuple, rs))
     sys_exprs = build_expr(:tuple, [convert(Expr, eq.rhs) for eq ∈ sys.eqs])
+    let_expr = Expr(:let, var_eqs, sys_exprs)
+
     if version === ArrayFunction
-        dvar_exprs = [:(du[$i] = $(Symbol("$(sys.dvs[i].name)_$(sys.iv.name)"))) for i in eachindex(sys.dvs)]
-        du_expr = :(du .= $sys_exprs)
-        exprs = vcat(var_exprs,param_exprs,du_expr)
-        block = expr_arr_to_block(exprs)
-        :((du,u,p,t)->$(toexpr(block)))
+        :((du,u,p,t) -> du .= $let_expr)
     elseif version === SArrayFunction
-        svector_expr = quote
-            du = $sys_exprs
+        :((u,p,t) -> begin
+            du = $let_expr
             T = StaticArrays.similar_type(typeof(u), eltype(du))
             T(du)
-        end
-        exprs = vcat(var_exprs,param_exprs,svector_expr)
-        block = expr_arr_to_block(exprs)
-        :((u,p,t)->$(toexpr(block)))
+        end)
     end
 end
 

src/systems/nonlinear/nonlinear_system.jl

@@ -9,20 +9,25 @@ function NonlinearSystem(eqs)
     NonlinearSystem(eqs, vs, ps)
 end
 
+iscalc(eq) = isequal(eq.lhs, Constant(0))
+
 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))
+    sys_eqs, calc_eqs = partition(iscalc, sys.eqs)
+
+    var_pairs   = [(u.name, :(u[$i])) for (i, u) ∈ enumerate(sys.vs)]
+    param_pairs = [(p.name, :(p[$i])) for (i, p) ∈ enumerate(sys.ps)]
+    calc_pairs  = [(eq.lhs.name, convert(Expr, eq.rhs)) for eq ∈ calc_eqs if isa(eq.lhs, Variable)]
+    (ls, rs) = collect(zip(var_pairs..., param_pairs..., calc_pairs...))
+
+    var_eqs = Expr(:(=), build_expr(:tuple, ls), build_expr(:tuple, rs))
+    sys_exprs = build_expr(:tuple, [convert(Expr, eq.rhs) for eq ∈ sys_eqs])
+    let_expr = Expr(:let, var_eqs, sys_exprs)
+
+    :((du,u,p) -> du .= $let_expr)
 end
 
 function calculate_jacobian(sys::NonlinearSystem,simplify=true)
-    sys_eqs, calc_eqs = partition(eq -> isequal(eq.lhs, Constant(0)), sys.eqs)
+    sys_eqs, calc_eqs = partition(iscalc, sys.eqs)
     rhs = [eq.rhs for eq in sys_eqs]
 
     for calc_eq ∈ calc_eqs

src/utils.jl

@@ -31,8 +31,6 @@ 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)
     return (xs[idxs], xs[(!).(idxs)])