feat: un-scalarize inferred parameters, improve parameter initialization

Author: AayushSabharwal

src/systems/abstractsystem.jl

@@ -192,8 +192,7 @@ function SymbolicIndexingInterface.is_variable(sys::AbstractSystem, sym)
         ic = get_index_cache(sys)
         h = getsymbolhash(sym)
         return haskey(ic.unknown_idx, h) ||
-               haskey(ic.unknown_idx, getsymbolhash(default_toterm(sym))) ||
-               hasname(sym) && is_variable(sys, getname(sym))
+               haskey(ic.unknown_idx, getsymbolhash(default_toterm(sym)))
     else
         return any(isequal(sym), variable_symbols(sys)) ||
                hasname(sym) && is_variable(sys, getname(sym))
@@ -220,8 +219,6 @@ function SymbolicIndexingInterface.variable_index(sys::AbstractSystem, sym)
             h = getsymbolhash(default_toterm(sym))
             if haskey(ic.unknown_idx, h)
                 ic.unknown_idx[h]
-            elseif hasname(sym)
-                variable_index(sys, getname(sym))
             else
                 nothing
             end
@@ -264,8 +261,7 @@ function SymbolicIndexingInterface.is_parameter(sys::AbstractSystem, sym)
         else
             h = getsymbolhash(default_toterm(sym))
             haskey(ic.param_idx, h) || haskey(ic.discrete_idx, h) ||
-                haskey(ic.constant_idx, h) || haskey(ic.dependent_idx, h) ||
-                hasname(sym) && is_parameter(sys, getname(sym))
+                haskey(ic.constant_idx, h) || haskey(ic.dependent_idx, h)
         end
     end
     return any(isequal(sym), parameter_symbols(sys)) ||

src/systems/diffeqs/odesystem.jl

@@ -280,10 +280,23 @@ function ODESystem(eqs, iv; kwargs...)
         isdelay(v, iv) || continue
         collect_vars!(allunknowns, ps, arguments(v)[1], iv)
     end
+    new_ps = OrderedSet()
+    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))
+                push!(new_ps, par)
+            else
+                push!(new_ps, p)
+            end
+        else
+            push!(new_ps, p)
+        end
+    end
     algevars = setdiff(allunknowns, diffvars)
     # the orders here are very important!
     return ODESystem(Equation[diffeq; algeeq; compressed_eqs], iv,
-        collect(Iterators.flatten((diffvars, algevars))), collect(ps); kwargs...)
+        collect(Iterators.flatten((diffvars, algevars))), collect(new_ps); kwargs...)
 end
 
 # NOTE: equality does not check cached Jacobian

src/systems/parameter_buffer.jl

@@ -36,6 +36,12 @@ function MTKParameters(sys::AbstractSystem, p; tofloat = false, use_union = fals
         for (k, v) in p if !haskey(extra_params, unwrap(k)))
     end
 
+    for (sym, _) in p
+        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
+
     tunable_buffer = Tuple(Vector{temp.type}(undef, temp.length)
     for temp in ic.param_buffer_sizes)
     disc_buffer = Tuple(Vector{temp.type}(undef, temp.length)
@@ -48,6 +54,7 @@ function MTKParameters(sys::AbstractSystem, p; tofloat = false, use_union = fals
     for temp in ic.nonnumeric_buffer_sizes)
     function set_value(sym, val)
         h = getsymbolhash(sym)
+        done = true
         if haskey(ic.param_idx, h)
             i, j = ic.param_idx[h]
             tunable_buffer[i][j] = val
@@ -64,17 +71,24 @@ function MTKParameters(sys::AbstractSystem, p; tofloat = false, use_union = fals
             i, j = ic.nonnumeric_idx[h]
             nonnumeric_buffer[i][j] = val
         elseif !isequal(default_toterm(sym), sym)
-            set_value(default_toterm(sym), val)
+            done = set_value(default_toterm(sym), val)
         else
-            error("Symbol $sym does not have an index")
+            done = false
         end
+        return done
     end
 
     for (sym, val) in p
         sym = unwrap(sym)
         ctype = concrete_symtype(sym)
         val = convert(ctype, fixpoint_sub(val, p))
-        set_value(sym, val)
+        done = set_value(sym, val)
+        if !done && Symbolics.isarraysymbolic(sym)
+            done = all(set_value.(collect(sym), val))
+        end
+        if !done
+            error("Symbol $sym does not have an index")
+        end
     end
 
     if has_parameter_dependencies(sys) &&