distributed works with both dispatches

Author: ChrisRackauckas

src/systems/diffeqs/diffeqsystem.jl

@@ -220,21 +220,32 @@ respectively.
 function DiffEqBase.ODEFunction{iip}(sys::ODESystem, dvs, ps;
                                      version = nothing,
                                      jac = false, Wfact = false) where {iip}
-    _f = generate_function(sys, dvs, ps)
+    f_oop,f_iip = generate_function(sys, dvs, ps)
+
+    f(u,p,t) = f_oop(u,p,t)
+    f(du,u,p,t) = f_iip(du,u,p,t)
 
     if jac
-        _jac = generate_jacobian(sys, dvs, ps)
+        jac_oop,jac_iip = generate_jacobian(sys, dvs, ps)
+        _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
-        _Wfact,_Wfact_t = generate_factorized_W(sys, dvs, ps)
+        tmp_Wfact,tmp_Wfact_t = generate_factorized_W(sys, dvs, ps)
+        Wfact_oop, Wfact_iip = tmp_Wfact
+        Wfact_oop_t, Wfact_iip_t = tmp_Wfact_t
+        _Wfact(u,p,t) = Wfact_oop(u,p,t)
+        _Wfact(W,u,p,t) = Wfact_iip(W,u,p,t)
+        _Wfact_t(u,p,t) = Wfact_oop_t(u,p,t)
+        _Wfact_t(W,u,p,t) = Wfact_iip_t(W,u,p,t)
     else
         _Wfact,_Wfact_t = nothing,nothing
     end
 
-    ODEFunction{iip}(_f,jac=_jac,
+    ODEFunction{iip}(f,jac=_jac,
                       Wfact = _Wfact,
                       Wfact_t = _Wfact_t)
 end

src/utils.jl

@@ -49,15 +49,15 @@ function build_function(rhss, vs, ps = (), args = (), conv = simplified_expr; co
 
     fname = gensym(:ModelingToolkitFunction)
 
-    X = gensym()
+    X = gensym(:MTIIPVar)
     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)
 
     fargs = ps == () ? :(u,$(args...)) : :(u,p,$(args...))
-    mk_function(fargs,:(),let_expr)
+    mk_function(fargs,:(),let_expr), mk_function(:($X,$(fargs.args...)),:(),ip_let_expr)
 end
 
 is_constant(::Constant) = true

test/distributed.jl

@@ -0,0 +1,38 @@
+using Distributed
+# add processes to workspace
+addprocs(2)
+
+using ModelingToolkit
+using OrdinaryDiffEq
+
+# create the Lorenz system
+@parameters t σ ρ β
+@variables x(t) y(t) z(t)
+@derivatives D'~t
+
+eqs = [D(x) ~ σ*(y-x),
+       D(y) ~ x*(ρ-z)-y,
+       D(z) ~ x*y - β*z]
+
+de = ODESystem(eqs)
+ode_func = ODEFunction(de, [x,y,z], [σ, ρ, β])
+
+u0 = [19.,20.,50.]
+params = [16.,45.92,4]
+
+ode_prob = ODEProblem(ode_func, u0, (0., 10.),params)
+
+@everywhere begin
+
+    using DifferentialEquations
+    using ModelingToolkit
+
+    function solve_lorenz(ode_problem)
+        print(solve(ode_problem,Tsit5()))
+    end
+end
+
+solve_lorenz(ode_prob)
+
+future = @spawn solve_lorenz(ode_prob)
+fetch(future)

test/runtests.jl

@@ -5,3 +5,4 @@ using ModelingToolkit, Test
 @testset "Simplify Test" begin include("simplify.jl") end
 @testset "Direct Usage Test" begin include("direct.jl") end
 @testset "System Construction Test" begin include("system_construction.jl") end
+@testset "Distributed Test" begin include("distributed.jl") end