Merge pull request #71 from JuliaDiffEq/myb/fix

Ensure that `generate_ode_function` always returns `Expr`

Author: YingboMa

src/systems/diffeqs/diffeqsystem.jl

@@ -40,18 +40,17 @@ 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)
-
     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))
+      :((du,u,p,t)->$(toexpr(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))
+      :((u,p,t)->$(toexpr(block)))
     end
 end
 

src/utils.jl

@@ -1,3 +1,4 @@
+using MacroTools
 function expr_arr_to_block(exprs)
   block = :(begin end)
   foreach(expr -> push!(block.args, expr), exprs)
@@ -15,3 +16,5 @@ function flatten_expr!(x)
     end
     x
 end
+
+toexpr(ex) = MacroTools.postwalk(x->x isa Union{Expression,Operation} ? Expr(x) : x, ex)

test/system_construction.jl

@@ -83,6 +83,15 @@ end
 
 ModelingToolkit.generate_nlsys_function(ns)
 
+@Var _x
+@Deriv D'~t
+@Param A B C
+eqs = [_x ~ y/C,
+       D*x ~ -A*x,
+       D*y ~ A*x - B*_x]
+de = DiffEqSystem(eqs,[t],[x,y],Variable[_x],[A,B,C])
+@test eval(ModelingToolkit.generate_ode_function(de))([0.0,0.0],[1.0,2.0],[1,2,3],0.0) ≈ -1/3
+
 # Now nonlinear system with only variables
 @Var x y z
 @Param σ ρ β