feat: scalarize equations in ODESystem, fix vars! and namespace_expr

Author: AayushSabharwal

src/systems/abstractsystem.jl

@@ -704,6 +704,9 @@ function namespace_expr(O, sys, n = nameof(sys); ivs = independent_variables(sys
             rescoped = renamespace(n, O)
             similarterm(O, operation(rescoped), renamed,
                 metadata = metadata(rescoped))
+        elseif Symbolics.isarraysymbolic(O)
+            # promote_symtype doesn't work for array symbolics
+            similarterm(O, operation(O), renamed, symtype(O), metadata = metadata(O))
         else
             similarterm(O, operation(O), renamed, metadata = metadata(O))
         end

src/systems/diffeqs/odesystem.jl

@@ -195,9 +195,8 @@ 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)
     @assert all(control -> any(isequal.(control, ps)), controls) "All controls must also be parameters."
-
+    deqs = reduce(vcat, scalarize(deqs); init = Equation[])
     iv′ = value(iv)
     ps′ = value.(ps)
     ctrl′ = value.(controls)
@@ -236,7 +235,6 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
 end
 
 function ODESystem(eqs, iv; kwargs...)
-    eqs = scalarize(eqs)
     # NOTE: this assumes that the order of algebraic equations doesn't matter
     diffvars = OrderedSet()
     allunknowns = OrderedSet()

src/utils.jl

@@ -342,22 +342,22 @@ v == Set([D(y), u])
 function vars(exprs::Symbolic; op = Differential)
     istree(exprs) ? vars([exprs]; op = op) : Set([exprs])
 end
+vars(exprs::Num; op = Differential) = vars(unwrap(exprs); op)
+vars(exprs::Symbolics.Arr; op = Differential) = vars(unwrap(exprs); op)
 vars(exprs; op = Differential) = foldl((x, y) -> vars!(x, y; op = op), exprs; init = Set())
 vars(eq::Equation; op = Differential) = vars!(Set(), eq; op = op)
 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)
+    if isvariable(O) && !(istree(O) && operation(O) === getindex)
+        return push!(vars, O)
+    end
+
     !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)
+        return vars!(vars, arr)
     end
 
     operation(O) isa op && return push!(vars, O)