Merge pull request #596 from SciML/mtkitize

More modelingtoolkitize tests and fix mtkitize with tuples

Author: ChrisRackauckas

src/systems/diffeqs/modelingtoolkitize.jl

@@ -7,9 +7,16 @@ function modelingtoolkitize(prob::DiffEqBase.ODEProblem)
     prob.f isa DiffEqBase.AbstractParameterizedFunction &&
                             return (prob.f.sys, prob.f.sys.states, prob.f.sys.ps)
     @parameters t
+
+    if prob.p isa Tuple || prob.p isa NamedTuple
+        p = [x for x in prob.p]
+    else
+        p = prob.p
+    end
+
     vars = reshape([Variable(:x, i)(t) for i in eachindex(prob.u0)],size(prob.u0))
-    params = prob.p isa DiffEqBase.NullParameters ? [] :
-             reshape([Variable(:α,i)() for i in eachindex(prob.p)],size(prob.p))
+    params = p isa DiffEqBase.NullParameters ? [] :
+             reshape([Variable(:α,i)() for i in eachindex(p)],size(Array(p)))
     @derivatives D'~t
 
     rhs = [D(var) for var in vars]
@@ -38,9 +45,16 @@ function modelingtoolkitize(prob::DiffEqBase.SDEProblem)
     prob.f isa DiffEqBase.AbstractParameterizedFunction &&
                             return (prob.f.sys, prob.f.sys.states, prob.f.sys.ps)
     @parameters t
+
+    if prob.p isa Tuple || prob.p isa NamedTuple
+        p = [x for x in prob.p]
+    else
+        p = prob.p
+    end
+
     vars = reshape([Variable(:x, i)(t) for i in eachindex(prob.u0)],size(prob.u0))
-    params = prob.p isa DiffEqBase.NullParameters ? [] :
-             reshape([Variable(:α,i)() for i in eachindex(prob.p)],size(prob.p))
+    params = p isa DiffEqBase.NullParameters ? [] :
+             reshape([Variable(:α,i)() for i in eachindex(p)],size(p))
     @derivatives D'~t
 
     rhs = [D(var) for var in vars]
@@ -65,7 +79,7 @@ function modelingtoolkitize(prob::DiffEqBase.SDEProblem)
         end
     end
     deqs = vcat([rhs[i] ~ lhs[i] for i in eachindex(prob.u0)]...)
-    
+
     de = SDESystem(deqs,neqs,t,vec(vars),vec(params))
 
     de

test/modelingtoolkitize.jl

@@ -0,0 +1,40 @@
+using OrdinaryDiffEq, ModelingToolkit
+const N = 32
+const xyd_brusselator = range(0,stop=1,length=N)
+brusselator_f(x, y, t) = (((x-0.3)^2 + (y-0.6)^2) <= 0.1^2) * (t >= 1.1) * 5.
+limit(a, N) = ModelingToolkit.ifelse(a == N+1, 1, ModelingToolkit.ifelse(a == 0, N, a))
+function brusselator_2d_loop(du, u, p, t)
+  A, B, alpha, dx = p
+  alpha = alpha/dx^2
+  @inbounds for I in CartesianIndices((N, N))
+    i, j = Tuple(I)
+    x, y = xyd_brusselator[I[1]], xyd_brusselator[I[2]]
+    ip1, im1, jp1, jm1 = limit(i+1, N), limit(i-1, N), limit(j+1, N), limit(j-1, N)
+    du[i,j,1] = alpha*(u[im1,j,1] + u[ip1,j,1] + u[i,jp1,1] + u[i,jm1,1] - 4u[i,j,1]) +
+                B + u[i,j,1]^2*u[i,j,2] - (A + 1)*u[i,j,1] + brusselator_f(x, y, t)
+    du[i,j,2] = alpha*(u[im1,j,2] + u[ip1,j,2] + u[i,jp1,2] + u[i,jm1,2] - 4u[i,j,2]) +
+                A*u[i,j,1] - u[i,j,1]^2*u[i,j,2]
+    end
+end
+
+# Test with tuple parameters
+p = (3.4, 1., 10., step(xyd_brusselator))
+
+function init_brusselator_2d(xyd)
+  N = length(xyd)
+  u = zeros(N, N, 2)
+  for I in CartesianIndices((N, N))
+    x = xyd[I[1]]
+    y = xyd[I[2]]
+    u[I,1] = 22*(y*(1-y))^(3/2)
+    u[I,2] = 27*(x*(1-x))^(3/2)
+  end
+  u
+end
+u0 = init_brusselator_2d(xyd_brusselator)
+
+# Test with 3-tensor inputs
+prob_ode_brusselator_2d = ODEProblem(brusselator_2d_loop,
+                                     u0,(0.,11.5),p)
+
+modelingtoolkitize(prob_ode_brusselator_2d)

test/runtests.jl

@@ -19,6 +19,7 @@ using SafeTestsets, Test
 @safetestset "ControlSystem Test" begin include("controlsystem.jl") end
 @safetestset "Build Targets Test" begin include("build_targets.jl") end
 @safetestset "Domain Test" begin include("domains.jl") end
+@safetestset "Modelingtoolkitize Test" begin include("modelingtoolkitize.jl") end
 @safetestset "Constraints Test" begin include("constraints.jl") end
 @safetestset "Reduction Test" begin include("reduction.jl") end
 @safetestset "PDE Construction Test" begin include("pde.jl") end