fix: avoid scalarizing params in structural_simplify, variable defaults in get_u0

Author: AayushSabharwal

src/systems/diffeqs/abstractodesystem.jl

@@ -820,19 +820,6 @@ 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) &&
-           Symbolics.shape(unwrap(k)) !== Symbolics.Unknown()
-            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)
     else

src/systems/systemstructure.jl

@@ -285,6 +285,12 @@ function TearingState(sys; quick_cancel = false, check = true)
         end
         vars!(vars, eq.rhs, op = Symbolics.Operator)
         for v in vars
+            _var, _ = var_from_nested_derivative(v)
+            any(isequal(_var), ivs) && continue
+            if isparameter(_var) ||
+               (istree(_var) && isparameter(operation(_var)) || isconstant(_var))
+                continue
+            end
             v = scalarize(v)
             if v isa AbstractArray
                 v = setmetadata.(v, VariableIrreducible, true)

src/variables.jl

@@ -187,13 +187,24 @@ end
 
 function _varmap_to_vars(varmap::Dict, varlist; defaults = Dict(), check = false,
         toterm = Symbolics.diff2term, initialization_phase = false)
-    varmap = merge(defaults, varmap) # prefers the `varmap`
-    varmap = Dict(toterm(value(k)) => value(varmap[k]) for k in keys(varmap))
+    _varmap = merge(defaults, varmap) # prefers the `varmap`
+    varmap = Dict()
+    for (k, v) in _varmap
+        varmap[value(k)] = value(v)
+        varmap[toterm(value(k))] = value(v)
+    end
     # resolve symbolic parameter expressions
     for (p, v) in pairs(varmap)
-        varmap[p] = fixpoint_sub(v, varmap)
+        varmap[p] = fixpoint_sub(unwrap(v), varmap)
+    end
+    for var in varlist
+        var = value(var)
+        haskey(varmap, var) && continue
+        val = fixpoint_sub(unwrap(var), varmap)
+        if symbolic_type(val) === NotSymbolic()
+            varmap[var] = val
+        end
     end
-
     missingvars = setdiff(varlist, collect(keys(varmap)))
     check && (isempty(missingvars) || throw(MissingVariablesError(missingvars)))