Use dispatch for in-place out-of-place function versions

Since anonymous functions cannot dispatch, this creates a gensym'd generic function with the two dispatches, making it simultaniously holding the oop and iip versions. Then ODEFunction just specifies which path to take. Defaults to saying it should be ran in-place.

Author: ChrisRackauckas

src/systems/diffeqs/diffeqsystem.jl

@@ -152,16 +152,18 @@ function (f::ODEToExpr)(O::Operation)
 end
 (f::ODEToExpr)(x) = convert(Expr, x)
 
-function generate_jacobian(sys::ODESystem; version::FunctionVersion = ArrayFunction)
+function generate_jacobian(sys::ODESystem; version::FunctionVersion = nothing)
+    version != nothing && @warn("version is deprecated. Both dispatches are now constructed in the same function by defualt.")
     jac = calculate_jacobian(sys)
-    return build_function(jac, sys.dvs, sys.ps, (sys.iv.name,), ODEToExpr(sys); version = version)
+    return build_function(jac, sys.dvs, sys.ps, (sys.iv.name,), ODEToExpr(sys))
 end
 
-function generate_function(sys::ODESystem, dvs, ps; version::FunctionVersion = ArrayFunction)
+function generate_function(sys::ODESystem, dvs, ps; version::FunctionVersion = nothing)
+    version != nothing && @warn("version is deprecated. Both dispatches are now constructed in the same function by defualt.")
     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); version = version)
+    return build_function(rhss, dvs′, ps′, (sys.iv.name,), ODEToExpr(sys))
 end
 
 function calculate_factorized_W(sys::ODESystem, simplify=true)
@@ -188,7 +190,8 @@ function calculate_factorized_W(sys::ODESystem, simplify=true)
     (Wfact,Wfact_t)
 end
 
-function generate_factorized_W(sys::ODESystem, simplify=true; version::FunctionVersion = ArrayFunction)
+function generate_factorized_W(sys::ODESystem, simplify=true; version::FunctionVersion = nothing)
+    version != nothing && @warn("version is deprecated. Both dispatches are now constructed in the same function by defualt.")
     (Wfact,Wfact_t) = calculate_factorized_W(sys,simplify)
 
     if version === SArrayFunction
@@ -202,8 +205,8 @@ function generate_factorized_W(sys::ODESystem, simplify=true; version::FunctionV
     end
 
     vs, ps = sys.dvs, sys.ps
-    Wfact_func   = build_function(Wfact  , vs, ps, (:gam,:t), ODEToExpr(sys); version = version, constructor=constructor)
-    Wfact_t_func = build_function(Wfact_t, vs, ps, (:gam,:t), ODEToExpr(sys); version = version, constructor=constructor)
+    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)
 
     return (Wfact_func, Wfact_t_func)
 end
@@ -215,16 +218,15 @@ Create an `ODEFunction` from the [`ODESystem`](@ref). The arguments `dvs` and `p
 are used to set the order of the dependent variable and parameter vectors,
 respectively.
 """
-function DiffEqBase.ODEFunction(sys::ODESystem, dvs, ps; version::FunctionVersion = ArrayFunction,
-                                                         jac = false, Wfact = false)
-    expr = eval(generate_function(sys, dvs, ps; version = version))
+function DiffEqBase.ODEFunction{iip}(sys::ODESystem, dvs, ps; version::FunctionVersion = nothing,
+                                                         jac = false, Wfact = false) where iip
+    version != nothing && @warn("version is deprecated. Both dispatches are now constructed in the same function by defualt.")
+    expr = eval(generate_function(sys, dvs, ps))
     jac_expr = jac ? nothing : eval(generate_jacobian(sys))
-    Wfact_expr,Wfact_t_expr = Wfact ? (nothing,nothing) : eval.(calculate_factorized_W(sys))
-    if version === ArrayFunction
-        ODEFunction{true}(eval(expr),jac=jac_expr,
-                          Wfact = Wfact_expr, Wfact_t = Wfact_t_expr)
-    elseif version === SArrayFunction
-        ODEFunction{false}(eval(expr),jac=jac_expr,
-                           Wfact = Wfact_expr, Wfact_t = Wfact_t_expr)
-    end
+    Wfact_expr,Wfact_t_expr = Wfact ? (nothing,nothing) : eval.(generate_factorized_W(sys))
+    ODEFunction{iip}(eval(expr),jac=jac_expr,
+                      Wfact = Wfact_expr, Wfact_t = Wfact_t_expr)
+end
+function DiffEqBase.ODEFunction(sys::ODESystem, args...; kwargs...)
+    ODEFunction{true}(sys, args...; kwargs...)
 end

src/utils.jl

@@ -31,26 +31,31 @@ function flatten_expr!(x)
     x
 end
 
-function build_function(rhss, vs, ps, args = (), conv = rhs -> convert(Expr, rhs); version::FunctionVersion, constructor=nothing)
+function build_function(rhss, vs, ps, args = (), conv = rhs -> convert(Expr, rhs); version::FunctionVersion=nothing, constructor=nothing)
+    version != nothing && @warn("version is deprecated. Both dispatches are now constructed in the same function by defualt.")
     var_pairs   = [(u.name, :(u[$i])) for (i, u) ∈ enumerate(vs)]
     param_pairs = [(p.name, :(p[$i])) for (i, p) ∈ enumerate(ps)]
     (ls, rs) = zip(var_pairs..., param_pairs...)
 
     var_eqs = Expr(:(=), build_expr(:tuple, ls), build_expr(:tuple, rs))
 
-    if version === ArrayFunction
-        X = gensym()
-        sys_exprs = [:($X[$i] = $(conv(rhs))) for (i, rhs) ∈ enumerate(rhss)]
-        let_expr = Expr(:let, var_eqs, build_expr(:block, sys_exprs))
-        :(($X,u,p,$(args...)) -> $let_expr)
-    elseif version === SArrayFunction
-        sys_expr = build_expr(:tuple, [conv(rhs) for rhs ∈ rhss])
-        let_expr = Expr(:let, var_eqs, sys_expr)
-        :((u,p,$(args...)) -> begin
+    fname = gensym()
+
+    X = gensym()
+    ip_sys_exprs = [:($X[$i] = $(conv(rhs))) for (i, rhs) ∈ enumerate(rhss)]
+    ip_let_expr = Expr(:let, var_eqs, build_expr(:block, ip_sys_exprs))
+
+    sys_expr = build_expr(:tuple, [conv(rhs) for rhs ∈ rhss])
+    let_expr = Expr(:let, var_eqs, sys_expr)
+    quote
+        function $fname($X,u,p,$(args...))
+            $ip_let_expr
+        end
+        function $fname(u,p,$(args...))
             X = $let_expr
             T = $(constructor === nothing ? :(StaticArrays.similar_type(typeof(u), eltype(X))) : constructor)
             T(X)
-        end)
+        end
     end
 end