Merge pull request #790 from SciML/mtkitize

make modelingtoolkitize more robust to weird arrays

Author: ChrisRackauckas

src/systems/diffeqs/modelingtoolkitize.jl

@@ -15,7 +15,7 @@ function modelingtoolkitize(prob::DiffEqBase.ODEProblem)
     end
 
     var(x, i) = Num(Sym{FnType{Tuple{symtype(t)}, Real}}(nameof(Variable(x, i))))
-    vars = reshape([var(:x, i)(value(t)) for i in eachindex(prob.u0)],size(prob.u0))
+    vars = ArrayInterface.restructure(prob.u0,[var(:x, i)(ModelingToolkit.value(t)) for i in eachindex(prob.u0)])
     params = p isa DiffEqBase.NullParameters ? [] :
              reshape([Num(Sym{Real}(nameof(Variable(:α, i)))) for i in eachindex(p)],size(p))
 
@@ -31,7 +31,14 @@ function modelingtoolkitize(prob::DiffEqBase.ODEProblem)
     end
 
     eqs = vcat([rhs[i] ~ lhs[i] for i in eachindex(prob.u0)]...)
-    de = ODESystem(eqs,t,vec(vars),vec(params))
+
+    params = if ndims(params) == 0
+        [params[1]]
+    else
+        Vector(vec(params))
+    end
+
+    de = ODESystem(eqs,t,Vector(vec(vars)),params)
 
     de
 end
@@ -53,7 +60,7 @@ function modelingtoolkitize(prob::DiffEqBase.SDEProblem)
         p = prob.p
     end
     var(x, i) = Num(Sym{FnType{Tuple{symtype(t)}, Real}}(nameof(Variable(x, i))))
-    vars = reshape([var(:x, i)(value(t)) for i in eachindex(prob.u0)],size(prob.u0))
+    vars = ArrayInterface.restructure(prob.u0,[var(:x, i)(ModelingToolkit.value(t)) for i in eachindex(prob.u0)])
     params = p isa DiffEqBase.NullParameters ? [] :
              reshape([Num(Sym{Real}(nameof(Variable(:α, i)))) for i in eachindex(p)],size(p))
 
@@ -83,7 +90,13 @@ function modelingtoolkitize(prob::DiffEqBase.SDEProblem)
     end
     deqs = vcat([rhs[i] ~ lhs[i] for i in eachindex(prob.u0)]...)
 
-    de = SDESystem(deqs,neqs,t,vec(vars),vec(params))
+    params = if ndims(params) == 0
+        [params[1]]
+    else
+        Vector(vec(params))
+    end
+
+    de = SDESystem(deqs,neqs,t,Vector(vec(vars)),params)
 
     de
 end

test/modelingtoolkitize.jl

@@ -1,5 +1,5 @@
 using OrdinaryDiffEq, ModelingToolkit, Test
-using GalacticOptim, Optim
+using GalacticOptim, Optim, RecursiveArrayTools
 
 N = 32
 const xyd_brusselator = range(0,stop=1,length=N)
@@ -141,3 +141,22 @@ problem = ODEProblem(SIRD_ac!, ℬ, 𝒯, 𝒫)
 sys = modelingtoolkitize(problem)
 fast_problem = ODEProblem(sys,ℬ, 𝒯, 𝒫 )
 @time solution = solve(fast_problem, Tsit5(), saveat = 1:final_time)
+
+## Issue #778
+
+r0 = [1131.340, -2282.343, 6672.423]
+v0 = [-5.64305, 4.30333, 2.42879]
+Δt = 86400.0*365
+μ = 398600.4418
+rv0 = ArrayPartition(r0,v0)
+
+function f(dy, y, μ, t)
+    r = sqrt(sum(y[1,:].^2))
+    dy[1,:] = y[2,:]
+    dy[2,:] = -μ .* y[1,:] / r^3
+end
+
+prob = ODEProblem(f, rv0, (0.0, Δt), μ)
+sol = solve(prob, Vern8())
+
+modelingtoolkitize(prob)