Remove generated derivative variables

Author: HarrisonGrodin

src/systems/diffeqs/diffeqsystem.jl

@@ -45,30 +45,25 @@ 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)
+    sys_exprs = build_expr(:tuple, [convert(Expr, eq.rhs) for eq ∈ 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)
+        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)))
     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...))
+            du = $sys_exprs
+            T = StaticArrays.similar_type(typeof(u), eltype(du))
+            T(du)
         end
-        exprs = vcat(var_exprs,param_exprs,sys_exprs,svector_expr)
+        exprs = vcat(var_exprs,param_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)
     isempty(sys.jac[]) || return sys.jac[]  # use cached Jacobian, if possible
     rhs = [eq.rhs for eq in sys.eqs]

test/system_construction.jl

@@ -109,7 +109,12 @@ eqs = [D(x) ~ -A*x,
 de = DiffEqSystem(eqs,t,[x,y],[A,B,C])
 test_vars_extraction(de, DiffEqSystem(eqs,t))
 test_vars_extraction(de, DiffEqSystem(eqs))
-@test eval(ModelingToolkit.generate_ode_function(de))([0.0,0.0],[1.0,2.0],[1,2,3],0.0) ≈ -1/3
+@test begin
+    f = eval(ModelingToolkit.generate_ode_function(de))
+    du = [0.0,0.0]
+    f(du, [1.0,2.0], [1,2,3], 0.0)
+    du ≈ [-1, -1/3]
+end
 
 # Now nonlinear system with only variables
 @Unknown x y z