Merge pull request #2836 from AayushSabharwal/as/more-metadata

fix: fix edge cases with metadata dumping, add tests

Author: ChrisRackauckas

src/systems/abstractsystem.jl

@@ -2776,11 +2776,15 @@ ModelingToolkit.dump_unknowns(sys)
 See also: [`ModelingToolkit.dump_variable_metadata`](@ref), [`ModelingToolkit.dump_parameters`](@ref)
 """
 function dump_unknowns(sys::AbstractSystem)
-    defs = defaults(sys)
+    defs = varmap_with_toterm(defaults(sys))
+    gs = varmap_with_toterm(guesses(sys))
     map(dump_variable_metadata.(unknowns(sys))) do meta
         if haskey(defs, meta.var)
             meta = merge(meta, (; default = defs[meta.var]))
         end
+        if haskey(gs, meta.var)
+            meta = merge(meta, (; guess = gs[meta.var]))
+        end
         meta
     end
 end

src/variables.jl

@@ -30,8 +30,13 @@ ModelingToolkit.dump_variable_metadata(p)
 function dump_variable_metadata(var)
     uvar = unwrap(var)
     vartype, name = get(uvar.metadata, VariableSource, (:unknown, :unknown))
-    shape = Symbolics.shape(var)
-    if shape == ()
+    type = symtype(uvar)
+    if type <: AbstractArray
+        shape = Symbolics.shape(var)
+        if shape == ()
+            shape = nothing
+        end
+    else
         shape = nothing
     end
     unit = get(uvar.metadata, VariableUnit, nothing)
@@ -208,6 +213,10 @@ function _varmap_to_vars(varmap::Dict, varlist; defaults = Dict(), check = false
     return [values[unwrap(var)] for var in varlist]
 end
 
+function varmap_with_toterm(varmap; toterm = Symbolics.diff2term)
+    return merge(todict(varmap), Dict(toterm(unwrap(k)) => v for (k, v) in varmap))
+end
+
 function canonicalize_varmap(varmap; toterm = Symbolics.diff2term)
     new_varmap = Dict()
     for (k, v) in varmap

test/test_variable_metadata.jl

@@ -131,13 +131,15 @@ sp = Set(p)
 
 @test_nowarn show(stdout, "text/plain", sys)
 
-# Defaults overridden by system, parameter dependencies
-@variables x(t) = 1.0
+# Defaults, guesses overridden by system, parameter dependencies
+@variables x(t)=1.0 y(t) [guess = 1.0]
 @parameters p=2.0 q
-@named sys = ODESystem(Equation[], t, [x], [p]; defaults = Dict(x => 2.0, p => 3.0),
-    parameter_dependencies = [q => 2p])
-x_meta = ModelingToolkit.dump_unknowns(sys)[]
-@test x_meta.default == 2.0
+@named sys = ODESystem(Equation[], t, [x, y], [p]; defaults = Dict(x => 2.0, p => 3.0),
+    guesses = Dict(y => 2.0), parameter_dependencies = [q => 2p])
+unks_meta = ModelingToolkit.dump_unknowns(sys)
+unks_meta = Dict([ModelingToolkit.getname(meta.var) => meta for meta in unks_meta])
+@test unks_meta[:x].default == 2.0
+@test unks_meta[:y].guess == 2.0
 params_meta = ModelingToolkit.dump_parameters(sys)
 params_meta = Dict([ModelingToolkit.getname(meta.var) => meta for meta in params_meta])
 @test params_meta[:p].default == 3.0