Merge remote-tracking branch 'origin/master' into sm/symutils3

Co-authored-by: "Yingbo Ma" <[email protected]>

Author: shashi

Project.toml

@@ -1,7 +1,7 @@
 name = "ModelingToolkit"
 uuid = "961ee093-0014-501f-94e3-6117800e7a78"
 authors = ["Chris Rackauckas <[email protected]>"]
-version = "3.20.0"
+version = "3.20.1"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

src/build_function.jl

@@ -87,6 +87,16 @@ function add_integrator_header(ex, fargs, iip; X=gensym(:MTIIPVar))
   wrappedex
 end
 
+function unflatten_long_ops(op, N=4)
+    rule1 = @rule((+)((~~x)) => length(~~x) > N ?
+                 +(+((~~x)[1:N]...) + (+)((~~x)[N+1:end]...)) : nothing)
+    rule2 = @rule((*)((~~x)) => length(~~x) > N ?
+                 *(*((~~x)[1:N]...) * (*)((~~x)[N+1:end]...)) : nothing)
+
+    op = to_symbolic(op)
+    Rewriters.Fixpoint(Rewriters.Postwalk(Rewriters.Chain([rule1, rule2])))(op) |> to_mtk
+end
+
 # Scalar output
 function _build_function(target::JuliaTarget, op, args...;
                          conv = toexpr, expression = Val{true},
@@ -97,9 +107,10 @@ function _build_function(target::JuliaTarget, op, args...;
     arg_pairs = map(vars_to_pairs,zip(argnames,args))
     ls = reduce(vcat,first.(arg_pairs))
     rs = reduce(vcat,last.(arg_pairs))
-    var_eqs = Expr(:(=), ModelingToolkit.build_expr(:tuple, ls), ModelingToolkit.build_expr(:tuple, rs))
+    var_eqs = Expr(:(=), ModelingToolkit.build_expr(:tuple, ls), ModelingToolkit.build_expr(:tuple, unflatten_long_ops.(rs)))
 
     fname = gensym(:ModelingToolkitFunction)
+    op = unflatten_long_ops(op)
     out_expr = conv(op)
     let_expr = Expr(:let, var_eqs, Expr(:block, out_expr))
     bounds_block = checkbounds ? let_expr : :(@inbounds begin $let_expr end)
@@ -242,7 +253,13 @@ function _build_function(target::JuliaTarget, rhss::AbstractArray, args...;
     oidx = isnothing(outputidxs) ? (i -> i) : (i -> outputidxs[i])
     X = gensym(:MTIIPVar)
 
-	rhs_length = rhss isa SparseMatrixCSC ? length(rhss.nzval) : length(rhss)
+    if rhss isa SparseMatrixCSC
+        rhs_length = length(rhss.nzval)
+        rhss = SparseMatrixCSC(rhss.m, rhss.m, rhss.colptr, rhss.rowval, map(unflatten_long_ops, rhss.nzval))
+    else
+        rhs_length = length(rhss)
+        rhss = [unflatten_long_ops(r) for r in rhss]
+    end
 
 	if parallel isa DistributedForm
 		numworks = Distributed.nworkers()
@@ -251,6 +268,7 @@ function _build_function(target::JuliaTarget, rhss::AbstractArray, args...;
 		finalsize = rhs_length - (numworks-1)*lens
 		_rhss = vcat(reduce(vcat,[[getindex(reducevars[i],j) for j in 1:lens] for i in 1:numworks-1],init=Expr[]),
 						 [getindex(reducevars[end],j) for j in 1:finalsize])
+
     elseif parallel isa DaggerForm
 		computevars = [Variable(gensym(:MTComputeVar))() for i in axes(rhss,1)]
         reducevar = Variable(gensym(:MTReduceVar))()

src/systems/diffeqs/modelingtoolkitize.jl

@@ -7,10 +7,18 @@ 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
+
     var(x, i) = Sym{FnType{Tuple{symtype(t)}, Number}}(nameof(Variable(:x, i)))
     vars = reshape([var(:x, i)(value(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]
@@ -39,10 +47,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
     var(x, i) = Sym{FnType{Tuple{symtype(t)}, Number}}(nameof(Variable(:x, i)))
     vars = reshape([var(:x, i)(value(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]
@@ -68,7 +82,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

src/systems/diffeqs/odesystem.jl

@@ -145,5 +145,5 @@ Base.:(==)(sys1::ODESystem, sys2::ODESystem) =
 # NOTE: equality does not check cached Jacobian
 
 function rename(sys::ODESystem,name)
-    ODESystem(sys.eqs, sys.iv, sys.states, sys.ps, sys.tgrad, sys.jac, sys.Wfact, sys.Wfact_t, name, sys.systems)
+    ODESystem(sys.eqs, sys.iv, sys.states, sys.ps, sys.pins, sys.observed, sys.tgrad, sys.jac, sys.Wfact, sys.Wfact_t, name, sys.systems)
 end

test/build_function.jl

@@ -64,3 +64,10 @@ h_julia_scalar(a, b, c, d, e, g) = a[1] + b[1] + c[1] + c[2] + c[3] + d[1] + e[1
 h_str_scalar = ModelingToolkit.build_function(h_scalar, [a], [b], [c1, c2, c3], [d], [e], [g])
 h_oop_scalar = eval(h_str_scalar)
 @test h_oop_scalar(inputs...) == h_julia_scalar(inputs...)
+
+@variables z[1:100]
+@test isequal(simplify(ModelingToolkit.unflatten_long_ops(sum(z))),
+              simplify(sum(z)))
+
+@test isequal(simplify(ModelingToolkit.unflatten_long_ops(prod(z))),
+              simplify(prod(z)))

test/components.jl

@@ -149,3 +149,35 @@ jac2 = [connected1₊lorenz1₊x 0 g zeros(1,12)
         zeros(Expression,7,8) ModelingToolkit.namespace_operation.(calculate_jacobian(connected2),connected2.name,:t)]
 
 @test all(isequal.(calculate_jacobian(doublelevel),jac2))
+
+# Try higher hierarchy
+
+@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]
+
+lorenz1 = ODESystem(eqs,name=:lorenz1)
+lorenz2 = ODESystem(eqs,name=:lorenz2)
+
+@variables a(t)
+@parameters γ
+connections = [0 ~ lorenz1.x + lorenz2.y + a*γ]
+
+# Now deviate from tutorial by introducing a 2nd level of connected blocks
+# Other issue: How does scoping work? Is it ok to re-use components named lorenz1 and lorenz2, or will the hierarchy be flattened eventually
+# connected1 and 2 are pairs of Lorenz oscillators. I'd like to solve multiple pairs together.
+connected1 = ODESystem(connections,t,[a],[γ],systems=[lorenz1,lorenz2], name=:connected1)
+connected2 = ODESystem(connections,t,[a],[γ],systems=[lorenz1,lorenz2], name=:connected2)
+
+connections2ndLevel = Equation[] # Pairs are not connected
+variables2ndLevel = [] # No extra variables
+parameters2ndLevel = [] # No extra parameters
+
+totalSystem = ODESystem(connections2ndLevel,t,variables2ndLevel, parameters2ndLevel, systems = [connected1, connected2],
+                        name=:totalSystem)
+
+connections2ndLevel = [0 ~ connected1.lorenz1.x + connected2.lorenz1.x]

test/controlsystem.jl

@@ -1,15 +1,16 @@
 using ModelingToolkit
 
 @variables t x(t) v(t) u(t)
+@parameters p
 @derivatives D'~t
 
 loss = (4-x)^2 + 2v^2 + u^2
 eqs = [
     D(x) ~ v
-    D(v) ~ u^3
+    D(v) ~ p*u^3
 ]
 
-sys = ControlSystem(loss,eqs,t,[x,v],[u],[])
+sys = ControlSystem(loss,eqs,t,[x,v],[u],[p])
 dt = 0.1
 tspan = (0.0,1.0)
 runge_kutta_discretize(sys,dt,tspan)

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