allow construction of SArray derivative function

Author: ChrisRackauckas

src/ModelingToolkit.jl

@@ -23,6 +23,8 @@ Base.convert(::Type{Variable},x::Int64) = Constant(x)
 
 function caclulate_jacobian end
 
+@enum FunctionVersions ArrayFunction=1 SArrayFunction=2
+
 include("operations.jl")
 include("operators.jl")
 include("systems/diffeqs/diffeqsystem.jl")

src/systems/diffeqs/diffeqsystem.jl

@@ -36,14 +36,23 @@ function DiffEqSystem(eqs, ivs;
     DiffEqSystem(eqs, ivs, dvs, vs, ps, ivs[1].subtype, dv_name, p_name, Matrix{Expression}(undef,0,0))
 end
 
-function generate_ode_function(sys::DiffEqSystem)
+function generate_ode_function(sys::DiffEqSystem;version = ArrayFunction)
     var_exprs = [:($(sys.dvs[i].name) = u[$i]) for i in 1:length(sys.dvs)]
     param_exprs = [:($(sys.ps[i].name) = p[$i]) for i in 1:length(sys.ps)]
     sys_exprs = build_equals_expr.(sys.eqs)
-    dvar_exprs = [:(du[$i] = $(Symbol("$(sys.dvs[i].name)_$(sys.ivs[1].name)"))) for i in 1:length(sys.dvs)]
-    exprs = vcat(var_exprs,param_exprs,sys_exprs,dvar_exprs)
-    block = expr_arr_to_block(exprs)
-    :((du,u,p,t)->$(block))
+
+    if version == ArrayFunction
+      dvar_exprs = [:(du[$i] = $(Symbol("$(sys.dvs[i].name)_$(sys.ivs[1].name)"))) for i in 1:length(sys.dvs)]
+      exprs = vcat(var_exprs,param_exprs,sys_exprs,dvar_exprs)
+      block = expr_arr_to_block(exprs)
+      :((du,u,p,t)->$(block))
+    elseif version == SArrayFunction
+      dvar_exprs = [:($(Symbol("$(sys.dvs[i].name)_$(sys.ivs[1].name)"))) for i in 1:length(sys.dvs)]
+      svector_expr = :(typeof(u)($(dvar_exprs...)))
+      exprs = vcat(var_exprs,param_exprs,sys_exprs,svector_expr)
+      block = expr_arr_to_block(exprs)
+      :((u,p,t)->$(block))
+    end
 end
 
 isintermediate(eq) = eq.args[1].diff == nothing

test/system_construction.jl

@@ -15,6 +15,7 @@ eqs = [D*x ~ σ*(y-x),
        D*z ~ x*y - β*z]
 de = DiffEqSystem(eqs,[t],[x,y,z],Variable[],[σ,ρ,β])
 ModelingToolkit.generate_ode_function(de)
+ModelingToolkit.generate_ode_function(de;version=ModelingToolkit.SArrayFunction)
 jac_expr = ModelingToolkit.generate_ode_jacobian(de)
 jac = ModelingToolkit.calculate_jacobian(de)
 f = ODEFunction(de)