Merge pull request #3528 from AayushSabharwal/as/fix-delayscope

fix: fix incorrect namespacing of DelayParentScope variables in `collect_scoped_vars!`

Author: AayushSabharwal

Project.toml

@@ -151,7 +151,7 @@ StochasticDelayDiffEq = "1.8.1"
 StochasticDiffEq = "6.72.1"
 SymbolicIndexingInterface = "0.3.37"
 SymbolicUtils = "3.25.1"
-Symbolics = "6.36"
+Symbolics = "6.37"
 URIs = "1"
 UnPack = "0.1, 1.0"
 Unitful = "1.1"

src/systems/abstractsystem.jl

@@ -1085,6 +1085,7 @@ function getvar(sys::AbstractSystem, name::Symbol; namespace = does_namespacing(
 
     if has_eqs(sys)
         for eq in get_eqs(sys)
+            eq isa Equation || continue
             if eq.lhs isa AnalysisPoint && nameof(eq.rhs) == name
                 return namespace ? renamespace(sys, eq.rhs) : eq.rhs
             end
@@ -1122,9 +1123,25 @@ function _apply_to_variables(f::F, ex) where {F}
         metadata(ex))
 end
 
+"""
+Variable metadata key which contains information about scoping/namespacing of the
+variable in a hierarchical system.
+"""
 abstract type SymScope end
 
+"""
+    $(TYPEDEF)
+
+The default scope of a variable. It belongs to the system whose equations it is involved
+in and is namespaced by every level of the hierarchy.
+"""
 struct LocalScope <: SymScope end
+
+"""
+    $(TYPEDSIGNATURES)
+
+Apply `LocalScope` to `sym`.
+"""
 function LocalScope(sym::Union{Num, Symbolic, Symbolics.Arr{Num}})
     apply_to_variables(sym) do sym
         if iscall(sym) && operation(sym) === getindex
@@ -1138,9 +1155,25 @@ function LocalScope(sym::Union{Num, Symbolic, Symbolics.Arr{Num}})
     end
 end
 
+"""
+    $(TYPEDEF)
+
+Denotes that the variable does not belong to the system whose equations it is involved
+in. It is not namespaced by this system. In the immediate parent of this system, the
+scope of this variable is given by `parent`.
+
+# Fields
+
+$(TYPEDFIELDS)
+"""
 struct ParentScope <: SymScope
     parent::SymScope
 end
+"""
+    $(TYPEDSIGNATURES)
+
+Apply `ParentScope` to `sym`, with `parent` being `LocalScope`.
+"""
 function ParentScope(sym::Union{Num, Symbolic, Symbolics.Arr{Num}})
     apply_to_variables(sym) do sym
         if iscall(sym) && operation(sym) === getindex
@@ -1156,11 +1189,34 @@ function ParentScope(sym::Union{Num, Symbolic, Symbolics.Arr{Num}})
     end
 end
 
+"""
+    $(TYPEDEF)
+
+Denotes that a variable belongs to a system that is at least `N + 1` levels up in the
+hierarchy from the system whose equations it is involved in. It is namespaced by the
+first `N` parents and not namespaced by the `N+1`th parent in the hierarchy. The scope
+of the variable after this point is given by `parent`.
+
+In other words, this scope delays applying `ParentScope` by `N` levels, and applies
+`LocalScope` in the meantime.
+
+# Fields
+
+$(TYPEDFIELDS)
+"""
 struct DelayParentScope <: SymScope
     parent::SymScope
     N::Int
 end
+
+"""
+    $(TYPEDSIGNATURES)
+
+Apply `DelayParentScope` to `sym`, with a delay of `N` and `parent` being `LocalScope`.
+"""
 function DelayParentScope(sym::Union{Num, Symbolic, Symbolics.Arr{Num}}, N)
+    Base.depwarn(
+        "`DelayParentScope` is deprecated and will be removed soon", :DelayParentScope)
     apply_to_variables(sym) do sym
         if iscall(sym) && operation(sym) == getindex
             args = arguments(sym)
@@ -1174,9 +1230,29 @@ function DelayParentScope(sym::Union{Num, Symbolic, Symbolics.Arr{Num}}, N)
         end
     end
 end
+
+"""
+    $(TYPEDSIGNATURES)
+
+Apply `DelayParentScope` to `sym`, with a delay of `1` and `parent` being `LocalScope`.
+"""
 DelayParentScope(sym::Union{Num, Symbolic, Symbolics.Arr{Num}}) = DelayParentScope(sym, 1)
 
+"""
+    $(TYPEDEF)
+
+Denotes that a variable belongs to the root system in the hierarchy, regardless of which
+equations of subsystems in the hierarchy it is involved in. Variables with this scope
+are never namespaced and only added to the unknowns/parameters of a system when calling
+`complete` or `structural_simplify`.
+"""
 struct GlobalScope <: SymScope end
+
+"""
+    $(TYPEDSIGNATURES)
+
+Apply `GlobalScope` to `sym`.
+"""
 function GlobalScope(sym::Union{Num, Symbolic, Symbolics.Arr{Num}})
     apply_to_variables(sym) do sym
         if iscall(sym) && operation(sym) == getindex

src/utils.jl

@@ -535,7 +535,7 @@ recursively searches through all subsystems of `sys`, increasing the depth if it
 """
 function collect_scoped_vars!(unknowns, parameters, sys, iv; depth = 1, op = Differential)
     if has_eqs(sys)
-        for eq in get_eqs(sys)
+        for eq in equations(sys)
             eqtype_supports_collect_vars(eq) || continue
             if eq isa Equation
                 eq.lhs isa Union{Symbolic, Number} || continue
@@ -544,7 +544,7 @@ function collect_scoped_vars!(unknowns, parameters, sys, iv; depth = 1, op = Dif
         end
     end
     if has_parameter_dependencies(sys)
-        for eq in get_parameter_dependencies(sys)
+        for eq in parameter_dependencies(sys)
             if eq isa Pair
                 collect_vars!(unknowns, parameters, eq, iv; depth, op)
             else
@@ -553,17 +553,13 @@ function collect_scoped_vars!(unknowns, parameters, sys, iv; depth = 1, op = Dif
         end
     end
     if has_constraints(sys)
-        for eq in get_constraints(sys)
+        for eq in constraints(sys)
             eqtype_supports_collect_vars(eq) || continue
             collect_vars!(unknowns, parameters, eq, iv; depth, op)
         end
     end
     if has_op(sys)
-        collect_vars!(unknowns, parameters, get_op(sys), iv; depth, op)
-    end
-    newdepth = depth == -1 ? depth : depth + 1
-    for ssys in get_systems(sys)
-        collect_scoped_vars!(unknowns, parameters, ssys, iv; depth = newdepth, op)
+        collect_vars!(unknowns, parameters, objective(sys), iv; depth, op)
     end
 end
 
@@ -608,6 +604,7 @@ end
 function collect_var!(unknowns, parameters, var, iv; depth = 0)
     isequal(var, iv) && return nothing
     check_scope_depth(getmetadata(var, SymScope, LocalScope()), depth) || return nothing
+    var = setmetadata(var, SymScope, LocalScope())
     if iscalledparameter(var)
         callable = getcalledparameter(var)
         push!(parameters, callable)
@@ -636,7 +633,7 @@ function check_scope_depth(scope, depth)
     elseif scope isa ParentScope
         return depth > 0 && check_scope_depth(scope.parent, depth - 1)
     elseif scope isa DelayParentScope
-        return depth >= scope.N && check_scope_depth(scope.parent, depth - scope.N)
+        return depth >= scope.N && check_scope_depth(scope.parent, depth - scope.N - 1)
     elseif scope isa GlobalScope
         return depth == -1
     end

test/jumpsystem.jl

@@ -381,12 +381,12 @@ let
     @variables x1(t) x2(t) x3(t) x4(t) x5(t)
     x2 = ParentScope(x2)
     x3 = ParentScope(ParentScope(x3))
-    x4 = DelayParentScope(x4, 2)
+    x4 = DelayParentScope(x4)
     x5 = GlobalScope(x5)
     @parameters p1 p2 p3 p4 p5
     p2 = ParentScope(p2)
     p3 = ParentScope(ParentScope(p3))
-    p4 = DelayParentScope(p4, 2)
+    p4 = DelayParentScope(p4)
     p5 = GlobalScope(p5)
 
     j1 = ConstantRateJump(p1, [x1 ~ x1 + 1])

test/variable_scope.jl

@@ -106,12 +106,12 @@ defs = ModelingToolkit.defaults(bar)
 @variables x1(t) x2(t) x3(t) x4(t) x5(t)
 x2 = ParentScope(x2)
 x3 = ParentScope(ParentScope(x3))
-x4 = DelayParentScope(x4, 2)
+x4 = DelayParentScope(x4)
 x5 = GlobalScope(x5)
 @parameters p1 p2 p3 p4 p5
 p2 = ParentScope(p2)
 p3 = ParentScope(ParentScope(p3))
-p4 = DelayParentScope(p4, 2)
+p4 = DelayParentScope(p4)
 p5 = GlobalScope(p5)
 
 @named sys1 = ODESystem([D(x1) ~ p1, D(x2) ~ p2, D(x3) ~ p3, D(x4) ~ p4, D(x5) ~ p5], t)
@@ -126,10 +126,10 @@ p5 = GlobalScope(p5)
 sys3 = sys3 ∘ sys2
 @test length(unknowns(sys3)) == 4
 @test any(isequal(x3), unknowns(sys3))
-@test any(isequal(x4), unknowns(sys3))
+@test any(isequal(ModelingToolkit.renamespace(sys1, x4)), unknowns(sys3))
 @test length(parameters(sys3)) == 4
 @test any(isequal(p3), parameters(sys3))
-@test any(isequal(p4), parameters(sys3))
+@test any(isequal(ModelingToolkit.renamespace(sys1, p4)), parameters(sys3))
 sys4 = complete(sys3)
 @test length(unknowns(sys3)) == 4
 @test length(parameters(sys4)) == 5