Merge pull request #3711 from AayushSabharwal/as/var-of-var-of-iv

fix: fix namespacing of variables of variables

Author: AayushSabharwal

src/independent_variables.jl

@@ -7,8 +7,12 @@ Define one or more independent variables. For example:
     @variables x(t)
 """
 macro independent_variables(ts...)
-    :(@parameters $(ts...)) |> esc # TODO: treat independent variables separately from variables and parameters
+    Symbolics._parse_vars(:independent_variables,
+        Real,
+        ts,
+        toiv) |> esc
 end
 
-toiv(s::Symbolic) = setmetadata(s, MTKVariableTypeCtx, PARAMETER)
+toiv(s::Symbolic) = GlobalScope(setmetadata(s, MTKVariableTypeCtx, PARAMETER))
+toiv(s::Symbolics.Arr) = wrap(toiv(value(s)))
 toiv(s::Num) = Num(toiv(value(s)))

src/systems/abstractsystem.jl

@@ -1169,7 +1169,9 @@ function is_array_of_symbolics(x)
 end
 
 function namespace_expr(
-        O, sys, n = nameof(sys); ivs = independent_variables(sys))
+        O, sys, n = (sys === nothing ? nothing : nameof(sys));
+        ivs = sys === nothing ? nothing : independent_variables(sys))
+    sys === nothing && return O
     O = unwrap(O)
     # Exceptions for arrays of symbolic and Ref of a symbolic, the latter
     # of which shows up in broadcasts
@@ -1538,9 +1540,9 @@ function defaults_and_guesses(sys::AbstractSystem)
     merge(guesses(sys), defaults(sys))
 end
 
-unknowns(sys::Union{AbstractSystem, Nothing}, v) = renamespace(sys, v)
+unknowns(sys::Union{AbstractSystem, Nothing}, v) = namespace_expr(v, sys)
 for vType in [Symbolics.Arr, Symbolics.Symbolic{<:AbstractArray}]
-    @eval unknowns(sys::AbstractSystem, v::$vType) = renamespace(sys, v)
+    @eval unknowns(sys::AbstractSystem, v::$vType) = namespace_expr(v, sys)
     @eval parameters(sys::AbstractSystem, v::$vType) = toparam(unknowns(sys, v))
 end
 parameters(sys::Union{AbstractSystem, Nothing}, v) = toparam(unknowns(sys, v))

src/systems/codegen_utils.jl

@@ -123,7 +123,11 @@ function isdelay(var, iv)
     if iscall(var) && !ModelingToolkit.isoperator(var, Symbolics.Operator)
         args = arguments(var)
         length(args) == 1 || return false
-        isequal(args[1], iv) || return true
+        arg = args[1]
+        isequal(arg, iv) && return false
+        iscall(arg) || return true
+        issym(operation(arg)) && !iscalledparameter(arg) && return false
+        return true
     end
     return false
 end

test/namespacing.jl

@@ -1,5 +1,6 @@
 using ModelingToolkit
-using ModelingToolkit: t_nounits as t, D_nounits as D, iscomplete, does_namespacing
+using ModelingToolkit: t_nounits as t, D_nounits as D, iscomplete, does_namespacing,
+                       renamespace
 
 @variables x(t)
 @parameters p
@@ -24,3 +25,23 @@ nsys = toggle_namespacing(sys, false)
 
 @test_throws ["namespacing", "inner"] System(
     Equation[], t; systems = [nsys], name = :a)
+
+@testset "Variables of variables" begin
+    @variables x(t) y(x)
+    @named inner = System([D(x) ~ x, y ~ 2x + 1], t)
+    @test issetequal(unknowns(inner), [x, y])
+    ss = mtkcompile(inner)
+    @test isequal(only(unknowns(ss)), x)
+    @test isequal(only(observed(ss)), y ~ 2x + 1)
+
+    @named sys = System(Equation[], t; systems = [inner])
+    xx, yy = let sys = inner
+        xx = renamespace(sys, x)
+        yy = only(@variables y(xx))
+        xx, renamespace(sys, yy)
+    end
+    @test issetequal(unknowns(sys), [xx, yy])
+    ss = mtkcompile(sys)
+    @test isequal(only(unknowns(ss)), xx)
+    @test isequal(only(observed(ss)), yy ~ 2xx + 1)
+end