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,17 +187,28 @@ 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))
-    # resolve symbolic parameter expressions
-    for (p, v) in pairs(varmap)
-        varmap[p] = fixpoint_sub(v, varmap)
+    varmap = canonicalize_varmap(varmap; toterm)
+    defaults = canonicalize_varmap(defaults; toterm)
+    values = Dict()
+    for var in varlist
+        var = unwrap(var)
+        val = unwrap(fixpoint_sub(fixpoint_sub(var, varmap), defaults))
+        if symbolic_type(val) === NotSymbolic()
+            values[var] = val
+        end
     end
-
-    missingvars = setdiff(varlist, collect(keys(varmap)))
+    missingvars = setdiff(varlist, collect(keys(values)))
     check && (isempty(missingvars) || throw(MissingVariablesError(missingvars)))
+    return [values[unwrap(var)] for var in varlist]
+end
 
-    out = [varmap[var] for var in varlist]
+function canonicalize_varmap(varmap; toterm = Symbolics.diff2term)
+    new_varmap = Dict()
+    for (k, v) in varmap
+        new_varmap[unwrap(k)] = unwrap(v)
+        new_varmap[toterm(unwrap(k))] = unwrap(v)
+    end
+    return new_varmap
 end
 
 @noinline function throw_missingvars(vars)

test/initial_values.jl

@@ -0,0 +1,35 @@
+using ModelingToolkit
+using ModelingToolkit: t_nounits as t, D_nounits as D, get_u0
+using SymbolicIndexingInterface: getu
+
+@variables x(t)[1:3]=[1.0, 2.0, 3.0] y(t) z(t)[1:2]
+
+@mtkbuild sys=ODESystem([D(x) ~ t * x], t) simplify=false
+@test get_u0(sys, [])[1] == [1.0, 2.0, 3.0]
+@test get_u0(sys, [x => [2.0, 3.0, 4.0]])[1] == [2.0, 3.0, 4.0]
+@test get_u0(sys, [x[1] => 2.0, x[2] => 3.0, x[3] => 4.0])[1] == [2.0, 3.0, 4.0]
+@test get_u0(sys, [2.0, 3.0, 4.0])[1] == [2.0, 3.0, 4.0]
+
+@mtkbuild sys=ODESystem([
+        D(x) ~ 3x,
+        D(y) ~ t,
+        D(z[1]) ~ z[2] + t,
+        D(z[2]) ~ y + z[1]
+    ], t) simplify=false
+
+@test_throws ModelingToolkit.MissingVariablesError get_u0(sys, [])
+getter = getu(sys, [x..., y, z...])
+@test getter(get_u0(sys, [y => 4.0, z => [5.0, 6.0]])[1]) == collect(1.0:6.0)
+@test getter(get_u0(sys, [y => 4.0, z => [3y, 4y]])[1]) == [1.0, 2.0, 3.0, 4.0, 12.0, 16.0]
+@test getter(get_u0(sys, [y => 3.0, z[1] => 3y, z[2] => 2x[1]])[1]) ==
+      [1.0, 2.0, 3.0, 3.0, 9.0, 2.0]
+
+@variables w(t)
+@parameters p1 p2
+
+@test getter(get_u0(sys, [y => 2p1, z => [3y, 2p2]], [p1 => 5.0, p2 => 6.0])[1]) ==
+      [1.0, 2.0, 3.0, 10.0, 30.0, 12.0]
+@test_throws Any getter(get_u0(sys, [y => 2w, w => 3.0, z[1] => 2p1, z[2] => 3p2]))
+@test getter(get_u0(
+    sys, [y => 2w, w => 3.0, z[1] => 2p1, z[2] => 3p2], [p1 => 3.0, p2 => 4.0])[1]) ==
+      [1.0, 2.0, 3.0, 6.0, 6.0, 12.0]

test/runtests.jl

@@ -1,6 +1,5 @@
 using SafeTestsets, Pkg, Test
 
-#=
 const GROUP = get(ENV, "GROUP", "All")
 
 function activate_extensions_env()
@@ -67,6 +66,7 @@ end
             @safetestset "Constants Test" include("constants.jl")
             @safetestset "Parameter Dependency Test" include("parameter_dependencies.jl")
             @safetestset "Generate Custom Function Test" include("generate_custom_function.jl")
+            @safetestset "Initial Values Test" include("initial_values.jl")
         end
     end
 
@@ -93,6 +93,4 @@ end
     end
 end
 
-=#
-
 @safetestset "Model Parsing Test" include("model_parsing.jl")