Unify generate_function

Author: HarrisonGrodin

README.md

@@ -56,7 +56,7 @@ 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
@@ -88,7 +88,7 @@ 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:
@@ -270,7 +270,7 @@ 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:

src/systems/diffeqs/diffeqsystem.jl

@@ -1,3 +1,6 @@
+export DiffEqSystem, ODEFunction
+
+
 using Base: RefValue
 
 
@@ -12,6 +15,7 @@ function Base.convert(::Type{DiffEq}, eq::Equation)
     isintermediate(eq) && throw(ArgumentError("intermediate equation received"))
     return DiffEq(eq.lhs.op, eq.lhs.args[1], eq.rhs)
 end
+Base.convert(::Type{Equation}, eq::DiffEq) = Equation(eq.D(eq.var), eq.rhs)
 Base.:(==)(a::DiffEq, b::DiffEq) = (a.D, a.var, a.rhs) == (b.D, b.var, b.rhs)
 get_args(eq::DiffEq) = Expression[eq.var, eq.rhs]
 
@@ -42,26 +46,6 @@ function DiffEqSystem(eqs, iv)
 end
 
 
-function generate_ode_function(sys::DiffEqSystem; version::FunctionVersion = ArrayFunction)
-    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
-        :((du,u,p,t) -> du .= $let_expr)
-    elseif version === SArrayFunction
-        :((u,p,t) -> begin
-            du = $let_expr
-            T = StaticArrays.similar_type(typeof(u), eltype(du))
-            T(du)
-        end)
-    end
-end
-
 function calculate_jacobian(sys::DiffEqSystem, simplify=true)
     isempty(sys.jac[]) || return sys.jac[]  # use cached Jacobian, if possible
     rhs = [eq.rhs for eq in sys.eqs]
@@ -71,6 +55,8 @@ function calculate_jacobian(sys::DiffEqSystem, simplify=true)
     return jac
 end
 
+system_eqs(sys::DiffEqSystem) = collect(Equation, sys.eqs)
+system_extras(::DiffEqSystem) = Equation[]
 system_vars(sys::DiffEqSystem) = sys.dvs
 system_params(sys::DiffEqSystem) = sys.ps
 
@@ -108,14 +94,10 @@ function generate_ode_iW(sys::DiffEqSystem, simplify=true)
 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,3 +1,6 @@
+export NonlinearSystem
+
+
 struct NonlinearSystem <: AbstractSystem
     eqs::Vector{Equation}
     vs::Vector{Variable}
@@ -9,25 +12,9 @@ function NonlinearSystem(eqs)
     NonlinearSystem(eqs, vs, ps)
 end
 
-iscalc(eq) = isequal(eq.lhs, Constant(0))
-
-function generate_nlsys_function(sys::NonlinearSystem)
-    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(iscalc, sys.eqs)
+function calculate_jacobian(sys::NonlinearSystem, simplify=true)
+    sys_eqs, calc_eqs = system_eqs(sys), filter(iscalc, sys.eqs)
     rhs = [eq.rhs for eq in sys_eqs]
 
     for calc_eq ∈ calc_eqs
@@ -39,8 +26,9 @@ function calculate_jacobian(sys::NonlinearSystem,simplify=true)
     sys_exprs
 end
 
+iscalc(eq) = !isequal(eq.lhs, Constant(0))
+
+system_eqs(sys::NonlinearSystem) = filter(!iscalc, sys.eqs)
+system_extras(sys::NonlinearSystem) = filter(eq -> isa(eq.lhs, Variable), sys.eqs)
 system_vars(sys::NonlinearSystem) = sys.vs
 system_params(sys::NonlinearSystem) = sys.ps
-
-export NonlinearSystem
-export generate_nlsys_function

src/systems/systems.jl

@@ -1,10 +1,11 @@
-export generate_jacobian
+export generate_jacobian, generate_function
 
 
 abstract type AbstractSystem end
 
-
-function system_vars end
+function system_eqs    end
+function system_extras end
+function system_vars   end
 function system_params end
 
 function generate_jacobian(sys::AbstractSystem, simplify = true)
@@ -17,3 +18,27 @@ function generate_jacobian(sys::AbstractSystem, simplify = true)
     block = expr_arr_to_block(exprs)
     :((J,u,p,t) -> $(block))
 end
+
+function generate_function(sys::AbstractSystem; version::FunctionVersion = ArrayFunction)
+    sys_eqs, calc_eqs = system_eqs(sys), system_extras(sys)
+    vs, ps = system_vars(sys), system_params(sys)
+
+    var_pairs   = [(u.name, :(u[$i])) for (i, u) ∈ enumerate(vs)]
+    param_pairs = [(p.name, :(p[$i])) for (i, p) ∈ enumerate(ps)]
+    calc_pairs  = [(eq.lhs.name, convert(Expr, eq.rhs)) for eq ∈ calc_eqs]
+    (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)
+
+    if version === ArrayFunction
+        :((du,u,p,t) -> du .= $let_expr)
+    elseif version === SArrayFunction
+        :((u,p,t) -> begin
+            du = $let_expr
+            T = StaticArrays.similar_type(typeof(u), eltype(du))
+            T(du)
+        end)
+    end
+end

test/system_construction.jl

@@ -12,8 +12,8 @@ eqs = [D(x) ~ σ*(y-x),
        D(y) ~ x*(ρ-z)-y,
        D(z) ~ x*y - β*z]
 de = DiffEqSystem(eqs,t,[x,y,z],[σ,ρ,β])
-ModelingToolkit.generate_ode_function(de)
-ModelingToolkit.generate_ode_function(de;version=ModelingToolkit.SArrayFunction)
+generate_function(de)
+generate_function(de;version=ModelingToolkit.SArrayFunction)
 jac_expr = generate_jacobian(de)
 jac = ModelingToolkit.calculate_jacobian(de)
 f = ODEFunction(de)
@@ -39,7 +39,7 @@ test_vars_extraction(de, de2)
            D(y) ~ x*(ρ-z)-y,
            D(z) ~ x*y - β*z]
     de = DiffEqSystem(eqs,[t],[x,y,z],[σ,ρ,β])
-    ModelingToolkit.generate_ode_function(de)
+    generate_function(de)
 
     #=
     ```julia
@@ -83,7 +83,7 @@ eqs = [D(x) ~ σ*a,
        D(y) ~ x*(ρ-z)-y,
        D(z) ~ x*y - β*z]
 de = DiffEqSystem(eqs,t,[x,y,z],[σ,ρ,β])
-ModelingToolkit.generate_ode_function(de)
+generate_function(de)
 jac = ModelingToolkit.calculate_jacobian(de)
 f = ODEFunction(de)
 
@@ -99,7 +99,7 @@ for el in (:vs, :ps)
     @test names2 == names
 end
 
-ModelingToolkit.generate_nlsys_function(ns)
+generate_function(ns)
 
 @Deriv D'~t
 @Param A B C
@@ -110,7 +110,7 @@ de = DiffEqSystem(eqs,t,[x,y],[A,B,C])
 test_vars_extraction(de, DiffEqSystem(eqs,t))
 test_vars_extraction(de, DiffEqSystem(eqs))
 @test begin
-    f = eval(ModelingToolkit.generate_ode_function(de))
+    f = eval(generate_function(de))
     du = [0.0,0.0]
     f(du, [1.0,2.0], [1,2,3], 0.0)
     du ≈ [-1, -1/3]
@@ -137,7 +137,7 @@ jac = ModelingToolkit.calculate_jacobian(ns)
     @test jac[3,2] == x
     @test jac[3,3] == -1 * β
 end
-nlsys_func = ModelingToolkit.generate_nlsys_function(ns)
+nlsys_func = generate_function(ns)
 jac_func = generate_jacobian(ns)
 f = @eval eval(nlsys_func)
 
@@ -148,6 +148,6 @@ eqs = [a ~ y-x,
        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)
 jac = ModelingToolkit.calculate_jacobian(ns)
 jac = generate_jacobian(ns)