From e7bd420a5e2acff21475d6be0e3c4951a62c4c47 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@juliahub.com>
Date: Tue, 18 Jun 2024 15:14:44 +0000
Subject: [PATCH] fix: fix edge cases with metadata dumping, add tests

---
 src/systems/abstractsystem.jl  |  6 +++++-
 src/variables.jl               |  4 ++++
 test/test_variable_metadata.jl | 12 +++++++-----
 3 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/src/systems/abstractsystem.jl b/src/systems/abstractsystem.jl
index d9853c8e58..7336695264 100644
--- a/src/systems/abstractsystem.jl
+++ b/src/systems/abstractsystem.jl
@@ -2663,11 +2663,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
diff --git a/src/variables.jl b/src/variables.jl
index d9c5f71cab..dc52269bff 100644
--- a/src/variables.jl
+++ b/src/variables.jl
@@ -208,6 +208,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
diff --git a/test/test_variable_metadata.jl b/test/test_variable_metadata.jl
index c715814d4c..a1b8a91d29 100644
--- a/test/test_variable_metadata.jl
+++ b/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
+    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