Merge pull request #2603 from AayushSabharwal/as/fix-namespacing-defaults

fix: fix variable namespacing issues

Author: ChrisRackauckas

src/systems/abstractsystem.jl

@@ -746,38 +746,64 @@ end
 abstract type SymScope end
 
 struct LocalScope <: SymScope end
-function LocalScope(sym::Union{Num, Symbolic})
+function LocalScope(sym::Union{Num, Symbolic, Symbolics.Arr{Num}})
     apply_to_variables(sym) do sym
-        setmetadata(sym, SymScope, LocalScope())
+        if istree(sym) && operation(sym) === getindex
+            args = arguments(sym)
+            a1 = setmetadata(args[1], SymScope, LocalScope())
+            similarterm(sym, operation(sym), [a1, args[2:end]...])
+        else
+            setmetadata(sym, SymScope, LocalScope())
+        end
     end
 end
 
 struct ParentScope <: SymScope
     parent::SymScope
 end
-function ParentScope(sym::Union{Num, Symbolic})
+function ParentScope(sym::Union{Num, Symbolic, Symbolics.Arr{Num}})
     apply_to_variables(sym) do sym
-        setmetadata(sym, SymScope,
-            ParentScope(getmetadata(value(sym), SymScope, LocalScope())))
+        if istree(sym) && operation(sym) === getindex
+            args = arguments(sym)
+            a1 = setmetadata(args[1], SymScope,
+                ParentScope(getmetadata(value(args[1]), SymScope, LocalScope())))
+            similarterm(sym, operation(sym), [a1, args[2:end]...])
+        else
+            setmetadata(sym, SymScope,
+                ParentScope(getmetadata(value(sym), SymScope, LocalScope())))
+        end
     end
 end
 
 struct DelayParentScope <: SymScope
     parent::SymScope
     N::Int
 end
-function DelayParentScope(sym::Union{Num, Symbolic}, N)
+function DelayParentScope(sym::Union{Num, Symbolic, Symbolics.Arr{Num}}, N)
     apply_to_variables(sym) do sym
-        setmetadata(sym, SymScope,
-            DelayParentScope(getmetadata(value(sym), SymScope, LocalScope()), N))
+        if istree(sym) && operation(sym) == getindex
+            args = arguments(sym)
+            a1 = setmetadata(args[1], SymScope,
+                DelayParentScope(getmetadata(value(args[1]), SymScope, LocalScope()), N))
+            similarterm(sym, operation(sym), [a1, args[2:end]...])
+        else
+            setmetadata(sym, SymScope,
+                DelayParentScope(getmetadata(value(sym), SymScope, LocalScope()), N))
+        end
     end
 end
-DelayParentScope(sym::Union{Num, Symbolic}) = DelayParentScope(sym, 1)
+DelayParentScope(sym::Union{Num, Symbolic, Symbolics.Arr{Num}}) = DelayParentScope(sym, 1)
 
 struct GlobalScope <: SymScope end
-function GlobalScope(sym::Union{Num, Symbolic})
+function GlobalScope(sym::Union{Num, Symbolic, Symbolics.Arr{Num}})
     apply_to_variables(sym) do sym
-        setmetadata(sym, SymScope, GlobalScope())
+        if istree(sym) && operation(sym) == getindex
+            args = arguments(sym)
+            a1 = setmetadata(args[1], SymScope, GlobalScope())
+            similarterm(sym, operation(sym), [a1, args[2:end]...])
+        else
+            setmetadata(sym, SymScope, GlobalScope())
+        end
     end
 end
 
@@ -793,6 +819,11 @@ function renamespace(sys, x)
             return similarterm(x, operation(x),
                 Any[renamespace(sys, only(arguments(x)))])::T
         end
+        if istree(x) && operation(x) === getindex
+            args = arguments(x)
+            return similarterm(
+                x, operation(x), vcat(renamespace(sys, args[1]), args[2:end]))::T
+        end
         let scope = getmetadata(x, SymScope, LocalScope())
             if scope isa LocalScope
                 rename(x, renamespace(getname(sys), getname(x)))::T
@@ -849,7 +880,8 @@ function namespace_assignment(eq::Assignment, sys)
     Assignment(_lhs, _rhs)
 end
 
-function namespace_expr(O, sys, n = nameof(sys); ivs = independent_variables(sys))
+function namespace_expr(
+        O, sys, n = nameof(sys); ivs = independent_variables(sys))
     O = unwrap(O)
     if any(isequal(O), ivs)
         return O
@@ -1500,8 +1532,7 @@ function default_to_parentscope(v)
     uv isa Symbolic || return v
     apply_to_variables(v) do sym
         if !hasmetadata(uv, SymScope)
-            setmetadata(sym, SymScope,
-                ParentScope(getmetadata(value(sym), SymScope, LocalScope())))
+            ParentScope(sym)
         else
             sym
         end

test/input_output_handling.jl

@@ -20,9 +20,9 @@ end
 @named sys = ODESystem([D(x) ~ -x + u], t) # both u and x are unbound
 @named sys1 = ODESystem([D(x) ~ -x + v[1] + v[2]], t) # both v and x are unbound
 @named sys2 = ODESystem([D(x) ~ -sys.x], t, systems = [sys]) # this binds sys.x in the context of sys2, sys2.x is still unbound
-@named sys21 = ODESystem([D(x) ~ -sys.x], t, systems = [sys1]) # this binds sys.x in the context of sys2, sys2.x is still unbound
+@named sys21 = ODESystem([D(x) ~ -sys1.x], t, systems = [sys1]) # this binds sys.x in the context of sys2, sys2.x is still unbound
 @named sys3 = ODESystem([D(x) ~ -sys.x + sys.u], t, systems = [sys]) # This binds both sys.x and sys.u
-@named sys31 = ODESystem([D(x) ~ -sys.x + sys1.v[1]], t, systems = [sys1]) # This binds both sys.x and sys1.v[1]
+@named sys31 = ODESystem([D(x) ~ -sys1.x + sys1.v[1]], t, systems = [sys1]) # This binds both sys.x and sys1.v[1]
 
 @named sys4 = ODESystem([D(x) ~ -sys.x, u ~ sys.u], t, systems = [sys]) # This binds both sys.x and sys3.u, this system is one layer deeper than the previous. u is directly forwarded to sys.u, and in this case sys.u is bound while u is not
 
@@ -43,7 +43,7 @@ end
 @test is_bound(sys2, sys.x)
 @test !is_bound(sys2, sys.u)
 @test !is_bound(sys2, sys2.sys.u)
-@test is_bound(sys21, sys.x)
+@test is_bound(sys21, sys1.x)
 @test !is_bound(sys21, sys1.v[1])
 @test !is_bound(sys21, sys1.v[2])
 @test is_bound(sys31, sys1.v[1])

test/odesystem.jl

@@ -1066,3 +1066,57 @@ prob = SteadyStateProblem(sys, u0, p)
 @test prob isa SteadyStateProblem
 prob = SteadyStateProblem(ODEProblem(sys, u0, (0.0, 10.0), p))
 @test prob isa SteadyStateProblem
+
+# Issue#2344
+using ModelingToolkitStandardLibrary.Blocks
+
+function FML2(; name)
+    @parameters begin
+        k2[1:1] = [1.0]
+    end
+    systems = @named begin
+        constant = Constant(k = k2[1])
+    end
+    @variables begin
+        x(t) = 0
+    end
+    eqs = [
+        D(x) ~ constant.output.u + k2[1]
+    ]
+    ODESystem(eqs, t; systems, name)
+end
+
+@mtkbuild model = FML2()
+
+@test isequal(ModelingToolkit.defaults(model)[model.constant.k], model.k2[1])
+@test_nowarn ODEProblem(model, [], (0.0, 10.0))
+
+# Issue#2477
+function RealExpression(; name, y)
+    vars = @variables begin
+        u(t)
+    end
+    eqns = [
+        u ~ y
+    ]
+    sys = ODESystem(eqns, t, vars, []; name)
+end
+
+function RealExpressionSystem(; name)
+    vars = @variables begin
+        x(t)
+        z(t)[1:1]
+    end # doing a collect on z doesn't work either. 
+    @named e1 = RealExpression(y = x) # This works perfectly. 
+    @named e2 = RealExpression(y = z[1]) # This bugs. However, `full_equations(e2)` works as expected. 
+    systems = [e1, e2]
+    ODESystem(Equation[], t, Iterators.flatten(vars), []; systems, name)
+end
+
+@named sys = RealExpressionSystem()
+sys = complete(sys)
+@test Set(equations(sys)) == Set([sys.e1.u ~ sys.x, sys.e2.u ~ sys.z[1]])
+tearing_state = TearingState(expand_connections(sys))
+ts_vars = tearing_state.fullvars
+orig_vars = unknowns(sys)
+@test isempty(setdiff(ts_vars, orig_vars))

test/variable_scope.jl

@@ -73,3 +73,13 @@ ps = ModelingToolkit.getname.(parameters(level3))
 @test isequal(ps[4], :level2₊level0₊d)
 @test isequal(ps[5], :level1₊level0₊e)
 @test isequal(ps[6], :f)
+
+# Issue@2252
+# Tests from PR#2354
+@parameters xx[1:2]
+arr_p = [ParentScope(xx[1]), xx[2]]
+arr0 = ODESystem(Equation[], t, [], arr_p; name = :arr0)
+arr1 = ODESystem(Equation[], t, [], []; name = :arr1) ∘ arr0
+arr_ps = ModelingToolkit.getname.(parameters(arr1))
+@test isequal(arr_ps[1], Symbol("xx"))
+@test isequal(arr_ps[2], Symbol("arr0₊xx"))