Array equations/variables support in structural_simplify

Author: YingboMa

src/systems/abstractsystem.jl

@@ -703,9 +703,9 @@ function namespace_expr(O, sys, n = nameof(sys); ivs = independent_variables(sys
             # metadata from the rescoped variable
             rescoped = renamespace(n, O)
             similarterm(O, operation(rescoped), renamed,
-                metadata = metadata(rescoped))::T
+                metadata = metadata(rescoped))
         else
-            similarterm(O, operation(O), renamed, metadata = metadata(O))::T
+            similarterm(O, operation(O), renamed, metadata = metadata(O))
         end
     elseif isvariable(O)
         renamespace(n, O)

src/systems/connectors.jl

@@ -292,18 +292,20 @@ function connection2set!(connectionsets, namespace, ss, isouter)
     end
 end
 
-function generate_connection_set(sys::AbstractSystem, find = nothing, replace = nothing)
+function generate_connection_set(
+        sys::AbstractSystem, find = nothing, replace = nothing; scalarize = false)
     connectionsets = ConnectionSet[]
     domain_csets = ConnectionSet[]
-    sys = generate_connection_set!(connectionsets, domain_csets, sys, find, replace)
+    sys = generate_connection_set!(
+        connectionsets, domain_csets, sys, find, replace, scalarize)
     csets = merge(connectionsets)
     domain_csets = merge([csets; domain_csets], true)
 
     sys, (csets, domain_csets)
 end
 
 function generate_connection_set!(connectionsets, domain_csets,
-        sys::AbstractSystem, find, replace, namespace = nothing)
+        sys::AbstractSystem, find, replace, scalarize, namespace = nothing)
     subsys = get_systems(sys)
 
     isouter = generate_isouter(sys)
@@ -325,8 +327,13 @@ function generate_connection_set!(connectionsets, domain_csets,
             end
             neweq isa AbstractArray ? append!(eqs, neweq) : push!(eqs, neweq)
         else
-            if lhs isa Number || lhs isa Symbolic
-                push!(eqs, eq) # split connections and equations
+            if lhs isa Number || lhs isa Symbolic || eltype(lhs) <: Symbolic
+                # split connections and equations
+                if eq.lhs isa AbstractArray || eq.rhs isa AbstractArray
+                    append!(eqs, Symbolics.scalarize(eq))
+                else
+                    push!(eqs, eq)
+                end
             elseif lhs isa Connection && get_systems(lhs) === :domain
                 connection2set!(domain_csets, namespace, get_systems(rhs), isouter)
             else
@@ -356,7 +363,7 @@ function generate_connection_set!(connectionsets, domain_csets,
     end
     @set! sys.systems = map(
         s -> generate_connection_set!(connectionsets, domain_csets, s,
-            find, replace,
+            find, replace, scalarize,
             renamespace(namespace, s)),
         subsys)
     @set! sys.eqs = eqs
@@ -471,8 +478,8 @@ function domain_defaults(sys, domain_csets)
 end
 
 function expand_connections(sys::AbstractSystem, find = nothing, replace = nothing;
-        debug = false, tol = 1e-10)
-    sys, (csets, domain_csets) = generate_connection_set(sys, find, replace)
+        debug = false, tol = 1e-10, scalarize = true)
+    sys, (csets, domain_csets) = generate_connection_set(sys, find, replace; scalarize)
     ceqs, instream_csets = generate_connection_equations_and_stream_connections(csets)
     _sys = expand_instream(instream_csets, sys; debug = debug, tol = tol)
     sys = flatten(sys, true)

src/systems/diffeqs/odesystem.jl

@@ -195,13 +195,13 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
         gui_metadata = nothing)
     name === nothing &&
         throw(ArgumentError("The `name` keyword must be provided. Please consider using the `@named` macro"))
-    deqs = scalarize(deqs)
+    #deqs = scalarize(deqs)
     @assert all(control -> any(isequal.(control, ps)), controls) "All controls must also be parameters."
 
-    iv′ = value(scalarize(iv))
-    ps′ = value.(scalarize(ps))
-    ctrl′ = value.(scalarize(controls))
-    dvs′ = value.(scalarize(dvs))
+    iv′ = value(iv)
+    ps′ = value.(ps)
+    ctrl′ = value.(controls)
+    dvs′ = value.(dvs)
     dvs′ = filter(x -> !isdelay(x, iv), dvs′)
 
     if !(isempty(default_u0) && isempty(default_p))

src/systems/systemstructure.jl

@@ -268,6 +268,7 @@ function TearingState(sys; quick_cancel = false, check = true)
     end
 
     vars = OrderedSet()
+    varsvec = []
     for (i, eq′) in enumerate(eqs)
         if eq′.lhs isa Connection
             check ? error("$(nameof(sys)) has unexpanded `connect` statements") :
@@ -282,9 +283,18 @@ function TearingState(sys; quick_cancel = false, check = true)
             eq = 0 ~ rhs - lhs
         end
         vars!(vars, eq.rhs, op = Symbolics.Operator)
+        for v in vars
+            v = scalarize(v)
+            if v isa AbstractArray
+                v = setmetadata.(v, VariableIrreducible, true)
+                append!(varsvec, v)
+            else
+                push!(varsvec, v)
+            end
+        end
         isalgeq = true
         unknownvars = []
-        for var in vars
+        for var in varsvec
             ModelingToolkit.isdelay(var, iv) && continue
             set_incidence = true
             @label ANOTHER_VAR
@@ -340,6 +350,7 @@ function TearingState(sys; quick_cancel = false, check = true)
         push!(symbolic_incidence, copy(unknownvars))
         empty!(unknownvars)
         empty!(vars)
+        empty!(varsvec)
         if isalgeq
             eqs[i] = eq
         else
@@ -350,9 +361,10 @@ function TearingState(sys; quick_cancel = false, check = true)
     # sort `fullvars` such that the mass matrix is as diagonal as possible.
     dervaridxs = collect(dervaridxs)
     sorted_fullvars = OrderedSet(fullvars[dervaridxs])
+    var_to_old_var = Dict(zip(fullvars, fullvars))
     for dervaridx in dervaridxs
         dervar = fullvars[dervaridx]
-        diffvar = lower_order_var(dervar)
+        diffvar = var_to_old_var[lower_order_var(dervar)]
         if !(diffvar in sorted_fullvars)
             push!(sorted_fullvars, diffvar)
         end

src/utils.jl

@@ -348,19 +348,22 @@ function vars!(vars, eq::Equation; op = Differential)
     (vars!(vars, eq.lhs; op = op); vars!(vars, eq.rhs; op = op); vars)
 end
 function vars!(vars, O; op = Differential)
+    !istree(O) && return vars
+    if operation(O) === (getindex)
+        arr = first(arguments(O))
+        !istree(arr) && return vars
+        operation(arr) isa op && return push!(vars, O)
+        isvariable(operation(O)) && return push!(vars, O)
+    end
+
     if isvariable(O)
         return push!(vars, O)
     end
-    !istree(O) && return vars
 
     operation(O) isa op && return push!(vars, O)
 
-    if operation(O) === (getindex) &&
-       isvariable(first(arguments(O)))
-        return push!(vars, O)
-    end
-
     isvariable(operation(O)) && push!(vars, O)
+
     for arg in arguments(O)
         vars!(vars, arg; op = op)
     end

test/odesystem.jl

@@ -517,21 +517,16 @@ eqs = [D(x) ~ x * y
 using StaticArrays
 using SymbolicUtils: term
 using SymbolicUtils.Code
-using Symbolics: unwrap, wrap
-function foo(a::Num, ms::AbstractVector)
-    a = unwrap(a)
-    ms = map(unwrap, ms)
-    wrap(term(foo, a, term(SVector, ms...)))
-end
+using Symbolics: unwrap, wrap, @register_symbolic
 foo(a, ms::AbstractVector) = a + sum(ms)
-@variables x(t) ms(t)[1:3]
-ms = collect(ms)
-eqs = [D(x) ~ foo(x, ms); D.(ms) .~ 1]
+@register_symbolic foo(a, ms::AbstractVector)
+@variables t x(t) ms(t)[1:3]
+D = Differential(t)
+eqs = [D(x) ~ foo(x, ms); D(ms) ~ ones(3)]
 @named sys = ODESystem(eqs, t, [x; ms], [])
 @named emptysys = ODESystem(Equation[], t)
-@named outersys = compose(emptysys, sys)
-outersys = complete(outersys)
-prob = ODEProblem(outersys, [sys.x => 1.0; collect(sys.ms) .=> 1:3], (0, 1.0))
+@mtkbuild outersys = compose(emptysys, sys)
+prob = ODEProblem(outersys, [sys.x => 1.0, sys.ms => 1:3], (0, 1.0))
 @test_nowarn solve(prob, Tsit5())
 
 # x/x