Merge pull request #2654 from AayushSabharwal/as/fix-hierarchical-guesses

fix: fix guesses in hierarchical system, getguess for array variables

Author: ChrisRackauckas

src/systems/abstractsystem.jl

@@ -859,6 +859,11 @@ function namespace_defaults(sys)
     for (k, v) in pairs(defs))
 end
 
+function namespace_guesses(sys)
+    guess = guesses(sys)
+    Dict(unknowns(sys, k) => namespace_expr(v, sys) for (k, v) in guess)
+end
+
 function namespace_equations(sys::AbstractSystem, ivs = independent_variables(sys))
     eqs = equations(sys)
     isempty(eqs) && return Equation[]
@@ -968,7 +973,13 @@ function full_parameters(sys::AbstractSystem)
 end
 
 function guesses(sys::AbstractSystem)
-    get_guesses(sys)
+    guess = get_guesses(sys)
+    systems = get_systems(sys)
+    isempty(systems) && return guess
+    for subsys in systems
+        guess = merge(guess, namespace_guesses(subsys))
+    end
+    return guess
 end
 
 # required in `src/connectors.jl:437`

src/systems/diffeqs/odesystem.jl

@@ -354,6 +354,7 @@ function flatten(sys::ODESystem, noeqs = false)
             get_iv(sys),
             unknowns(sys),
             parameters(sys),
+            guesses = guesses(sys),
             observed = observed(sys),
             continuous_events = continuous_events(sys),
             discrete_events = discrete_events(sys),

src/variables.jl

@@ -456,7 +456,7 @@ end
 ## Guess ======================================================================
 struct VariableGuess end
 Symbolics.option_to_metadata_type(::Val{:guess}) = VariableGuess
-getguess(x::Num) = getguess(Symbolics.unwrap(x))
+getguess(x::Union{Num, Symbolics.Arr}) = getguess(Symbolics.unwrap(x))
 
 """
     getguess(x)
@@ -469,8 +469,6 @@ Create variables with a guess like this
 ```
 """
 function getguess(x)
-    p = Symbolics.getparent(x, nothing)
-    p === nothing || (x = p)
     Symbolics.getmetadata(x, VariableGuess, nothing)
 end
 

test/odesystem.jl

@@ -1121,6 +1121,18 @@ ts_vars = tearing_state.fullvars
 orig_vars = unknowns(sys)
 @test isempty(setdiff(ts_vars, orig_vars))
 
+# Guesses in hierarchical systems
+@variables x(t) y(t)
+@named sys = ODESystem(Equation[], t, [x], []; guesses = [x => 1.0])
+@named outer = ODESystem(
+    [D(y) ~ sys.x + t, 0 ~ t + y - sys.x * y], t, [y], []; systems = [sys])
+@test ModelingToolkit.guesses(outer)[sys.x] == 1.0
+outer = structural_simplify(outer)
+@test ModelingToolkit.get_guesses(outer)[sys.x] == 1.0
+prob = ODEProblem(outer, [outer.y => 2.0], (0.0, 10.0))
+int = init(prob, Rodas4())
+@test int[outer.sys.x] == 1.0
+
 # Ensure indexes of array symbolics are cached appropriately
 @variables x(t)[1:2]
 @named sys = ODESystem(Equation[], t, [x], [])

test/test_variable_metadata.jl

@@ -16,6 +16,18 @@ using ModelingToolkit
 @test hasguess(y) === true
 @test ModelingToolkit.dump_variable_metadata(y).guess == 0
 
+# Issue#2653
+@variables y[1:3] [guess = ones(3)]
+@test getguess(y) == ones(3)
+@test hasguess(y) === true
+@test ModelingToolkit.dump_variable_metadata(y).guess == ones(3)
+
+for i in 1:3
+    @test getguess(y[i]) == 1.0
+    @test hasguess(y[i]) === true
+    @test ModelingToolkit.dump_variable_metadata(y[i]).guess == 1.0
+end
+
 @variables y
 @test hasguess(y) === false
 @test !haskey(ModelingToolkit.dump_variable_metadata(y), :guess)