feat: flatten equations to avoid scalarizing array arguments

Author: AayushSabharwal

src/systems/diffeqs/abstractodesystem.jl

@@ -807,6 +807,17 @@ function get_u0(sys, u0map, parammap = nothing; symbolic_u0 = false)
         defs = mergedefaults(defs, parammap, ps)
     end
     defs = mergedefaults(defs, u0map, dvs)
+    for (k, v) in defs
+        if Symbolics.isarraysymbolic(k)
+            ks = scalarize(k)
+            length(ks) == length(v) || error("$k has default value $v with unmatched size")
+            for (kk, vv) in zip(ks, v)
+                if !haskey(defs, kk)
+                    defs[kk] = vv
+                end
+            end
+        end
+    end
 
     if symbolic_u0
         u0 = varmap_to_vars(u0map, dvs; defaults = defs, tofloat = false, use_union = false)
@@ -1415,3 +1426,19 @@ function isisomorphic(sys1::AbstractODESystem, sys2::AbstractODESystem)
     end
     return false
 end
+
+function flatten_equations(eqs)
+    mapreduce(vcat, eqs; init = Equation[]) do eq
+        islhsarr = eq.lhs isa AbstractArray || Symbolics.isarraysymbolic(eq.lhs)
+        isrhsarr = eq.rhs isa AbstractArray || Symbolics.isarraysymbolic(eq.rhs)
+        if islhsarr || isrhsarr
+            islhsarr && isrhsarr ||
+                error("LHS ($(eq.lhs)) and RHS ($(eq.rhs)) must either both be array expressions or both scalar")
+            size(eq.lhs) == size(eq.rhs) ||
+                error("Size of LHS ($(eq.lhs)) and RHS ($(eq.rhs)) must match: got $(size(eq.lhs)) and $(size(eq.rhs))")
+            return collect(eq.lhs) .~ collect(eq.rhs)
+        else
+            eq
+        end
+    end
+end

src/systems/diffeqs/odesystem.jl

@@ -202,7 +202,19 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
     name === nothing &&
         throw(ArgumentError("The `name` keyword must be provided. Please consider using the `@named` macro"))
     @assert all(control -> any(isequal.(control, ps)), controls) "All controls must also be parameters."
-    deqs = reduce(vcat, scalarize(deqs); init = Equation[])
+    deqs = mapreduce(vcat, deqs; init = Equation[]) do eq
+        islhsarr = eq.lhs isa AbstractArray || Symbolics.isarraysymbolic(eq.lhs)
+        isrhsarr = eq.rhs isa AbstractArray || Symbolics.isarraysymbolic(eq.rhs)
+        if islhsarr || isrhsarr
+            islhsarr && isrhsarr ||
+                error("LHS ($(eq.lhs)) and RHS ($(eq.rhs)) must either both be array expressions or both scalar")
+            size(eq.lhs) == size(eq.rhs) ||
+                error("Size of LHS ($(eq.lhs)) and RHS ($(eq.rhs)) must match: got $(size(eq.lhs)) and $(size(eq.rhs))")
+            return collect(eq.lhs) .~ collect(eq.rhs)
+        else
+            eq
+        end
+    end
     iv′ = value(iv)
     ps′ = value.(ps)
     ctrl′ = value.(controls)
@@ -284,7 +296,8 @@ function ODESystem(eqs, iv; kwargs...)
     for p in ps
         if istree(p) && operation(p) === getindex
             par = arguments(p)[begin]
-            if Symbolics.shape(Symbolics.unwrap(par)) !== Symbolics.Unknown() && all(par[i] in ps for i in eachindex(par))
+            if Symbolics.shape(Symbolics.unwrap(par)) !== Symbolics.Unknown() &&
+               all(par[i] in ps for i in eachindex(par))
                 push!(new_ps, par)
             else
                 push!(new_ps, p)

src/systems/diffeqs/sdesystem.jl

@@ -171,7 +171,14 @@ function SDESystem(deqs::AbstractVector{<:Equation}, neqs::AbstractArray, iv, dv
         gui_metadata = nothing)
     name === nothing &&
         throw(ArgumentError("The `name` keyword must be provided. Please consider using the `@named` macro"))
-    deqs = scalarize(deqs)
+    deqs = flatten_equations(deqs)
+    neqs = mapreduce(vcat, neqs) do expr
+        if expr isa AbstractArray || Symbolics.isarraysymbolic(expr)
+            collect(expr)
+        else
+            expr
+        end
+    end
     iv′ = value(iv)
     dvs′ = value.(dvs)
     ps′ = value.(ps)

src/systems/jumps/jumpsystem.jl

@@ -151,7 +151,7 @@ function JumpSystem(eqs, iv, unknowns, ps;
         kwargs...)
     name === nothing &&
         throw(ArgumentError("The `name` keyword must be provided. Please consider using the `@named` macro"))
-    eqs = scalarize(eqs)
+    eqs = flatten_equations(eqs)
     sysnames = nameof.(systems)
     if length(unique(sysnames)) != length(sysnames)
         throw(ArgumentError("System names must be unique."))

src/systems/parameter_buffer.jl

@@ -37,7 +37,8 @@ function MTKParameters(sys::AbstractSystem, p; tofloat = false, use_union = fals
     end
 
     for (sym, _) in p
-        if istree(sym) && operation(sym) === getindex && is_parameter(sys, arguments(sym)[begin])
+        if istree(sym) && operation(sym) === getindex &&
+           is_parameter(sys, arguments(sym)[begin])
             # error("Scalarized parameter values are not supported. Instead of `[p[1] => 1.0, p[2] => 2.0]` use `[p => [1.0, 2.0]]`")
         end
     end

test/odesystem.jl

@@ -520,15 +520,29 @@ using SymbolicUtils.Code
 using Symbolics: unwrap, wrap, @register_symbolic
 foo(a, ms::AbstractVector) = a + sum(ms)
 @register_symbolic foo(a, ms::AbstractVector)
-@variables t x(t) ms(t)[1:3]
-D = Differential(t)
+@variables x(t) ms(t)[1:3]
 eqs = [D(x) ~ foo(x, ms); D(ms) ~ ones(3)]
 @named sys = ODESystem(eqs, t, [x; ms], [])
 @named emptysys = ODESystem(Equation[], t)
 @mtkbuild outersys = compose(emptysys, sys)
 prob = ODEProblem(outersys, [sys.x => 1.0, sys.ms => 1:3], (0, 1.0))
 @test_nowarn solve(prob, Tsit5())
 
+# array equations
+bar(x, p) = p * x
+@register_array_symbolic bar(x::AbstractVector, p::AbstractMatrix) begin
+    size = size(x)
+    eltype = promote_type(eltype(x), eltype(p))
+end
+@parameters p[1:3, 1:3]
+eqs = [D(x) ~ foo(x, ms); D(ms) ~ bar(ms, p)]
+@named sys = ODESystem(eqs, t)
+@named emptysys = ODESystem(Equation[], t)
+@mtkbuild outersys = compose(emptysys, sys)
+prob = ODEProblem(
+    outersys, [sys.x => 1.0, sys.ms => 1:3], (0.0, 1.0), [sys.p => ones(3, 3)])
+@test_nowarn solve(prob, Tsit5())
+
 # x/x
 @variables x(t)
 @named sys = ODESystem([D(x) ~ x / x], t)