keep expression output

Author: ChrisRackauckas

src/systems/diffeqs/diffeqsystem.jl

@@ -160,16 +160,16 @@ function (f::ODEToExpr)(O::Operation)
 end
 (f::ODEToExpr)(x) = convert(Expr, x)
 
-function generate_jacobian(sys::ODESystem, dvs = sys.dvs, ps = sys.ps)
+function generate_jacobian(sys::ODESystem, dvs = sys.dvs, ps = sys.ps, expression = Val{true})
     jac = calculate_jacobian(sys)
-    return build_function(jac, dvs, ps, (sys.iv.name,), ODEToExpr(sys))
+    return build_function(jac, dvs, ps, (sys.iv.name,), ODEToExpr(sys), expression)
 end
 
-function generate_function(sys::ODESystem, dvs = sys.dvs, ps = sys.ps)
+function generate_function(sys::ODESystem, dvs = sys.dvs, ps = sys.ps, expression = Val{true})
     rhss = [deq.rhs for deq ∈ sys.eqs]
     dvs′ = [clean(dv) for dv ∈ dvs]
     ps′ = [clean(p) for p ∈ ps]
-    return build_function(rhss, dvs′, ps′, (sys.iv.name,), ODEToExpr(sys))
+    return build_function(rhss, dvs′, ps′, (sys.iv.name,), ODEToExpr(sys), expression)
 end
 
 function calculate_factorized_W(sys::ODESystem, simplify=true)
@@ -196,16 +196,16 @@ function calculate_factorized_W(sys::ODESystem, simplify=true)
     (Wfact,Wfact_t)
 end
 
-function generate_factorized_W(sys::ODESystem, vs = sys.dvs, ps = sys.ps, simplify=true)
+function generate_factorized_W(sys::ODESystem, vs = sys.dvs, ps = sys.ps, simplify=true, expression = Val{true})
     (Wfact,Wfact_t) = calculate_factorized_W(sys,simplify)
     siz = size(Wfact)
     constructor = :(x -> begin
                         A = SMatrix{$siz...}(x)
                         StaticArrays.LU(LowerTriangular( SMatrix{$siz...}(UnitLowerTriangular(A)) ), UpperTriangular(A), SVector(ntuple(n->n, max($siz...))))
                     end)
 
-    Wfact_func   = build_function(Wfact  , vs, ps, (:gam,:t), ODEToExpr(sys);constructor=constructor)
-    Wfact_t_func = build_function(Wfact_t, vs, ps, (:gam,:t), ODEToExpr(sys);constructor=constructor)
+    Wfact_func   = build_function(Wfact  , vs, ps, (:gam,:t), ODEToExpr(sys), expression;constructor=constructor)
+    Wfact_t_func = build_function(Wfact_t, vs, ps, (:gam,:t), ODEToExpr(sys), expression;constructor=constructor)
 
     return (Wfact_func, Wfact_t_func)
 end
@@ -220,21 +220,21 @@ respectively.
 function DiffEqBase.ODEFunction{iip}(sys::ODESystem, dvs, ps;
                                      version = nothing,
                                      jac = false, Wfact = false) where {iip}
-    f_oop,f_iip = generate_function(sys, dvs, ps)
+    f_oop,f_iip = generate_function(sys, dvs, ps, Val{false})
 
     f(u,p,t) = f_oop(u,p,t)
     f(du,u,p,t) = f_iip(du,u,p,t)
 
     if jac
-        jac_oop,jac_iip = generate_jacobian(sys, dvs, ps)
+        jac_oop,jac_iip = generate_jacobian(sys, dvs, ps, Val{false})
         _jac(u,p,t) = jac_oop(u,p,t)
         _jac(J,u,p,t) = jac_iip(J,u,p,t)
     else
         _jac = nothing
     end
 
     if Wfact
-        tmp_Wfact,tmp_Wfact_t = generate_factorized_W(sys, dvs, ps)
+        tmp_Wfact,tmp_Wfact_t = generate_factorized_W(sys, dvs, ps, Val{false})
         Wfact_oop, Wfact_iip = tmp_Wfact
         Wfact_oop_t, Wfact_iip_t = tmp_Wfact_t
         _Wfact(u,p,t) = Wfact_oop(u,p,t)

src/utils.jl

@@ -38,7 +38,7 @@ mk_function(args, kwargs, body) =
         GG.RuntimeFn{Args, Kwargs, Body}()
     end
 
-function build_function(rhss, vs, ps = (), args = (), conv = simplified_expr; constructor=nothing)
+function build_function(rhss, vs, ps = (), args = (), conv = simplified_expr, expression = Val{true}; constructor=nothing)
     _vs = map(x-> x isa Operation ? x.op : x, vs)
     _ps = map(x-> x isa Operation ? x.op : x, ps)
     var_pairs   = [(u.name, :(u[$i])) for (i, u) ∈ enumerate(_vs)]
@@ -57,7 +57,24 @@ function build_function(rhss, vs, ps = (), args = (), conv = simplified_expr; co
     let_expr = Expr(:let, var_eqs, sys_expr)
 
     fargs = ps == () ? :(u,$(args...)) : :(u,p,$(args...))
-    mk_function(fargs,:(),let_expr), mk_function(:($X,$(fargs.args...)),:(),ip_let_expr)
+
+    if expression == Val{true}
+        return quote
+            function $fname($X,$(fargs.args...))
+                $ip_let_expr
+                nothing
+            end
+            function $fname($(fargs.args...))
+                X = $let_expr
+                T = promote_type(map(typeof,X)...)
+                convert.(T,X)
+                construct = $(constructor === nothing ? :(u isa ModelingToolkit.StaticArrays.StaticArray ? ModelingToolkit.StaticArrays.similar_type(typeof(u), eltype(X)) : x->(du=similar(u, T, $(size(rhss)...)); vec(du) .= x; du)) : constructor)
+                construct(X)
+            end
+        end
+    else
+        return mk_function(fargs,:(),let_expr), mk_function(:($X,$(fargs.args...)),:(),ip_let_expr)
+    end
 end
 
 is_constant(::Constant) = true