Merge pull request #804 from mohamed82008/patch-1

Define inplace and out of place methods in ODEFunctionExpr

Author: ChrisRackauckas

src/systems/diffeqs/abstractodesystem.jl

@@ -192,39 +192,53 @@ function ODEFunctionExpr{iip}(sys::AbstractODESystem, dvs = states(sys),
                                      sparse = false, simplify=false,
                                      kwargs...) where {iip}
 
-    idx = iip ? 2 : 1
-    f = generate_function(sys, dvs, ps; expression=Val{true}, kwargs...)[idx]
+    f_oop, f_iip = generate_function(sys, dvs, ps; expression=Val{true}, kwargs...)
+    fsym = gensym(:f)
+    _f = quote
+      $fsym(u,p,t) = $f_oop(u,p,t)
+      $fsym(du,u,p,t) = $f_iip(du,u,p,t)
+    end
+  
+    tgradsym = gensym(:tgrad)
     if tgrad
-        _tgrad = generate_tgrad(sys, dvs, ps;
+        tgrad_oop, tgrad_iip = generate_tgrad(sys, dvs, ps;
                                 simplify=simplify,
-                                expression=Val{true}, kwargs...)[idx]
+                                expression=Val{true}, kwargs...)
+        _tgrad = quote
+            $tgradsym(u,p,t) = $tgrad_oop(u,p,t)
+            $tgradsym(J,u,p,t) = $tgrad_iip(J,u,p,t)
+        end
     else
-        _tgrad = :nothing
+        _tgrad = :($tgradsym = nothing)
     end
 
+    jacsym = gensym(:jac)
     if jac
-        _jac = generate_jacobian(sys, dvs, ps;
+        jac_oop,jac_iip = generate_jacobian(sys, dvs, ps;
                                  sparse=sparse, simplify=simplify,
-                                 expression=Val{true}, kwargs...)[idx]
+                                 expression=Val{true}, kwargs...)
+        _jac = quote
+            $jacsym(u,p,t) = $jac_oop(u,p,t)
+            $jacsym(J,u,p,t) = $jac_iip(J,u,p,t)
+        end
     else
-        _jac = :nothing
+        _jac = :($jacsym = nothing)
     end
 
     M = calculate_massmatrix(sys)
 
     _M = (u0 === nothing || M == I) ? M : ArrayInterface.restructure(u0 .* u0',M)
 
     jp_expr = sparse ? :(similar($(get_jac(sys)[]),Float64)) : :nothing
-
     ex = quote
-        f = $f
-        tgrad = $_tgrad
-        jac = $_jac
+        $_f
+        $_tgrad
+        $_jac
         M = $_M
         ODEFunction{$iip}(
-                          f,
-                          jac = jac,
-                          tgrad = tgrad,
+                          $fsym,
+                          jac = $jacsym,
+                          tgrad = $tgradsym,
                           mass_matrix = M,
                           jac_prototype = $jp_expr,
                           syms = $(Symbol.(states(sys))),

test/odesystem.jl

@@ -35,21 +35,38 @@ generate_function(de, [x,y,z], [σ,ρ,β])
 jac_expr = generate_jacobian(de)
 jac = calculate_jacobian(de)
 jacfun = eval(jac_expr[2])
-# iip
-f = ODEFunction(de, [x,y,z], [σ,ρ,β])
-du = zeros(3)
-u  = collect(1:3)
-p  = collect(4:6)
-f(du, u, p, 0.1)
-@test du == [4, 0, -16]
-J = zeros(3, 3)
-jacfun(J, u, p, t)
-# oop
-f = ODEFunction(de, [x,y,z], [σ,ρ,β])
-du = @SArray zeros(3)
-u  = SVector(1:3...)
-p  = SVector(4:6...)
-@test f(u, p, 0.1) === @SArray [4, 0, -16]
+
+for f in [
+    ODEFunction(de, [x,y,z], [σ,ρ,β], tgrad = true, jac = true),
+    eval(ODEFunctionExpr(de, [x,y,z], [σ,ρ,β], tgrad = true, jac = true)),
+]
+    # iip
+    du = zeros(3)
+    u  = collect(1:3)
+    p  = collect(4:6)
+    f.f(du, u, p, 0.1)
+    @test du == [4, 0, -16]
+
+    # oop
+    du = @SArray zeros(3)
+    u  = SVector(1:3...)
+    p  = SVector(4:6...)
+    @test f.f(u, p, 0.1) === @SArray [4, 0, -16]
+
+    # iip vs oop
+    du = zeros(3)
+    g = similar(du)
+    J = zeros(3, 3)
+    u  = collect(1:3)
+    p  = collect(4:6)
+    f.f(du, u, p, 0.1)
+    @test du == f(u, p, 0.1)
+    f.tgrad(g, u, p, t)
+    @test g == f.tgrad(u, p, t)
+    f.jac(J, u, p, t)
+    @test J == f.jac(u, p, t)
+end
+
 
 eqs = [D(x) ~ σ*(y-x),
        D(y) ~ x*(ρ-z)-y*t,
@@ -59,6 +76,8 @@ ModelingToolkit.calculate_tgrad(de)
 
 tgrad_oop, tgrad_iip = eval.(ModelingToolkit.generate_tgrad(de))
 
+u  = SVector(1:3...)
+p  = SVector(4:6...)
 @test tgrad_oop(u,p,t) == [0.0,-u[2],0.0]
 du = zeros(3)
 tgrad_iip(du,u,p,t)