Merge pull request #3112 from AayushSabharwal/as/better-scoping

feat: better discover scoped parameters in parent systems

Author: ChrisRackauckas

src/systems/abstractsystem.jl

@@ -913,10 +913,24 @@ end
 """
 $(TYPEDSIGNATURES)
 
-Mark a system as completed. If a system is complete, the system will no longer
+Mark a system as completed. A completed system is a system which is done being
+defined/modified and is ready for structural analysis or other transformations.
+This allows for analyses and optimizations to be performed which require knowing
+the global structure of the system.
+        
+One property to note is that if a system is complete, the system will no longer
 namespace its subsystems or variables, i.e. `isequal(complete(sys).v.i, v.i)`.
 """
 function complete(sys::AbstractSystem; split = true)
+    if !(sys isa JumpSystem)
+        newunknowns = OrderedSet()
+        newparams = OrderedSet()
+        iv = has_iv(sys) ? get_iv(sys) : nothing
+        collect_scoped_vars!(newunknowns, newparams, sys, iv; depth = -1)
+        # don't update unknowns to not disturb `structural_simplify` order
+        # `GlobalScope`d unknowns will be picked up and added there
+        @set! sys.ps = unique!(vcat(get_ps(sys), collect(newparams)))
+    end
     if split && has_index_cache(sys)
         @set! sys.index_cache = IndexCache(sys)
         all_ps = parameters(sys)
@@ -3011,6 +3025,14 @@ function compose(sys::AbstractSystem, systems::AbstractArray; name = nameof(sys)
     if has_is_dde(sys)
         @set! sys.is_dde = _check_if_dde(equations(sys), get_iv(sys), get_systems(sys))
     end
+    newunknowns = OrderedSet()
+    newparams = OrderedSet()
+    iv = has_iv(sys) ? get_iv(sys) : nothing
+    for ssys in systems
+        collect_scoped_vars!(newunknowns, newparams, ssys, iv)
+    end
+    @set! sys.unknowns = unique!(vcat(get_unknowns(sys), collect(newunknowns)))
+    @set! sys.ps = unique!(vcat(get_ps(sys), collect(newparams)))
     return sys
 end
 function compose(syss...; name = nameof(first(syss)))

src/systems/diffeqs/odesystem.jl

@@ -323,11 +323,13 @@ function ODESystem(eqs, iv; kwargs...)
         collect_vars!(allunknowns, ps, eq.rhs, iv)
         if isdiffeq(eq)
             diffvar, _ = var_from_nested_derivative(eq.lhs)
-            isequal(iv, iv_from_nested_derivative(eq.lhs)) ||
-                throw(ArgumentError("An ODESystem can only have one independent variable."))
-            diffvar in diffvars &&
-                throw(ArgumentError("The differential variable $diffvar is not unique in the system of equations."))
-            push!(diffvars, diffvar)
+            if check_scope_depth(getmetadata(diffvar, SymScope, LocalScope()), 0)
+                isequal(iv, iv_from_nested_derivative(eq.lhs)) ||
+                    throw(ArgumentError("An ODESystem can only have one independent variable."))
+                diffvar in diffvars &&
+                    throw(ArgumentError("The differential variable $diffvar is not unique in the system of equations."))
+                push!(diffvars, diffvar)
+            end
             push!(diffeq, eq)
         else
             push!(algeeq, eq)
@@ -342,6 +344,9 @@ function ODESystem(eqs, iv; kwargs...)
             collect_vars!(allunknowns, ps, eq.rhs, iv)
         end
     end
+    for ssys in get(kwargs, :systems, ODESystem[])
+        collect_scoped_vars!(allunknowns, ps, ssys, iv)
+    end
     for v in allunknowns
         isdelay(v, iv) || continue
         collect_vars!(allunknowns, ps, arguments(v)[1], iv)

src/utils.jl

@@ -473,23 +473,56 @@ function find_derivatives!(vars, expr, f)
     return vars
 end
 
-function collect_vars!(unknowns, parameters, expr, iv; op = Differential)
+"""
+    $(TYPEDSIGNATURES)
+
+Search through equations and parameter dependencies of `sys`, where sys is at a depth of
+`depth` from the root system, looking for variables scoped to the root system. Also
+recursively searches through all subsystems of `sys`, increasing the depth if it is not
+`-1`. A depth of `-1` indicates searching for variables with `GlobalScope`.
+"""
+function collect_scoped_vars!(unknowns, parameters, sys, iv; depth = 1, op = Differential)
+    if has_eqs(sys)
+        for eq in get_eqs(sys)
+            eq.lhs isa Union{Symbolic, Number} || continue
+            collect_vars!(unknowns, parameters, eq.lhs, iv; depth, op)
+            collect_vars!(unknowns, parameters, eq.rhs, iv; depth, op)
+        end
+    end
+    if has_parameter_dependencies(sys)
+        for eq in get_parameter_dependencies(sys)
+            if eq isa Pair
+                collect_vars!(unknowns, parameters, eq[1], iv; depth, op)
+                collect_vars!(unknowns, parameters, eq[2], iv; depth, op)
+            else
+                collect_vars!(unknowns, parameters, eq.lhs, iv; depth, op)
+                collect_vars!(unknowns, parameters, eq.rhs, iv; depth, op)
+            end
+        end
+    end
+    newdepth = depth == -1 ? depth : depth + 1
+    for ssys in get_systems(sys)
+        collect_scoped_vars!(unknowns, parameters, ssys, iv; depth = newdepth, op)
+    end
+end
+
+function collect_vars!(unknowns, parameters, expr, iv; depth = 0, op = Differential)
     if issym(expr)
-        collect_var!(unknowns, parameters, expr, iv)
+        collect_var!(unknowns, parameters, expr, iv; depth)
     else
         for var in vars(expr; op)
             if iscall(var) && operation(var) isa Differential
                 var, _ = var_from_nested_derivative(var)
             end
-            collect_var!(unknowns, parameters, var, iv)
+            collect_var!(unknowns, parameters, var, iv; depth)
         end
     end
     return nothing
 end
 
-function collect_var!(unknowns, parameters, var, iv)
+function collect_var!(unknowns, parameters, var, iv; depth = 0)
     isequal(var, iv) && return nothing
-    getmetadata(var, SymScope, LocalScope()) == LocalScope() || return nothing
+    check_scope_depth(getmetadata(var, SymScope, LocalScope()), depth) || return nothing
     if iscalledparameter(var)
         callable = getcalledparameter(var)
         push!(parameters, callable)
@@ -506,6 +539,24 @@ function collect_var!(unknowns, parameters, var, iv)
     return nothing
 end
 
+"""
+    $(TYPEDSIGNATURES)
+
+Check if the given `scope` is at a depth of `depth` from the root system. Only
+returns `true` for `scope::GlobalScope` if `depth == -1`.
+"""
+function check_scope_depth(scope, depth)
+    if scope isa LocalScope
+        return depth == 0
+    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)
+    elseif scope isa GlobalScope
+        return depth == -1
+    end
+end
+
 """
 Find all the symbolic constants of some equations or terms and return them as a vector.
 """

test/variable_scope.jl

@@ -101,3 +101,41 @@ bar = complete(bar)
 defs = ModelingToolkit.defaults(bar)
 @test defs[bar.p] == 2
 @test isequal(defs[bar.foo.p], bar.p)
+
+# Issue#3101
+@variables x1(t) x2(t) x3(t) x4(t) x5(t)
+x2 = ParentScope(x2)
+x3 = ParentScope(ParentScope(x3))
+x4 = DelayParentScope(x4, 2)
+x5 = GlobalScope(x5)
+@parameters p1 p2 p3 p4 p5
+p2 = ParentScope(p2)
+p3 = ParentScope(ParentScope(p3))
+p4 = DelayParentScope(p4, 2)
+p5 = GlobalScope(p5)
+
+@named sys1 = ODESystem([D(x1) ~ p1, D(x2) ~ p2, D(x3) ~ p3, D(x4) ~ p4, D(x5) ~ p5], t)
+@test isequal(x1, only(unknowns(sys1)))
+@test isequal(p1, only(parameters(sys1)))
+@named sys2 = ODESystem(Equation[], t; systems = [sys1])
+@test length(unknowns(sys2)) == 2
+@test any(isequal(x2), unknowns(sys2))
+@test length(parameters(sys2)) == 2
+@test any(isequal(p2), parameters(sys2))
+@named sys3 = ODESystem(Equation[], t)
+sys3 = sys3 ∘ sys2
+@test length(unknowns(sys3)) == 4
+@test any(isequal(x3), unknowns(sys3))
+@test any(isequal(x4), unknowns(sys3))
+@test length(parameters(sys3)) == 4
+@test any(isequal(p3), parameters(sys3))
+@test any(isequal(p4), parameters(sys3))
+sys4 = complete(sys3)
+@test length(unknowns(sys3)) == 4
+@test length(parameters(sys4)) == 5
+@test any(isequal(p5), parameters(sys4))
+sys5 = structural_simplify(sys3)
+@test length(unknowns(sys5)) == 5
+@test any(isequal(x5), unknowns(sys5))
+@test length(parameters(sys5)) == 5
+@test any(isequal(p5), parameters(sys5))