update collecting vars with scoping

Author: isaacsas

src/systems/abstractsystem.jl

@@ -923,15 +923,14 @@ 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, flatten = 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
+    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)))
+    
     if flatten
         eqs = equations(sys)
         if eqs isa AbstractArray && eltype(eqs) <: Equation

src/systems/jumps/jumpsystem.jl

@@ -217,6 +217,7 @@ function JumpSystem(eqs, iv, unknowns, ps;
 end
 
 ##### MTK dispatches for JumpSystems #####
+eqtype_supports_collect_vars(j::MassActionJump) = true
 function collect_vars!(unknowns, parameters, j::MassActionJump, iv; depth = 0, 
         op = Differential)    
     collect_vars!(unknowns, parameters, j.scaled_rates, iv; depth, op)
@@ -228,6 +229,7 @@ function collect_vars!(unknowns, parameters, j::MassActionJump, iv; depth = 0,
     return nothing
 end
 
+eqtype_supports_collect_vars(j::Union{ConstantRateJump,VariableRateJump}) = true
 function collect_vars!(unknowns, parameters, j::Union{ConstantRateJump,VariableRateJump}, 
         iv; depth = 0, op = Differential)
     collect_vars!(unknowns, parameters, j.rate, iv; depth, op)

test/jumpsystem.jl

@@ -376,3 +376,49 @@ let
     @test issetequal(us, [x1, x5, x2])
     @test issetequal(ps, [p4, p5])
 end
+
+# scoping tests
+let
+    @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)
+    
+    j1 = ConstantRateJump(p1, [x1 ~ x1 + 1])
+    j2 = MassActionJump(p2, [x2 => 1], [x3 => -1])
+    j3 = VariableRateJump(p3, [x3 ~ x3 + 1, x4 ~ x4 + 1])
+    j4 = MassActionJump(p4*p5, [x1 => 1, x5 => 1], [x1 => -1, x5 => -1, x2 => 1])
+    @named js = JumpSystem([j1, j2, j3, j4], t, [x1, x2, x3, x4, x5], [p1, p2, p3, p4, p5])
+
+    us = Set(); ps = Set()
+    iv = t
+    MT.collect_scoped_vars!(us, ps, js, iv)
+    @test issetequal(us, [x2])
+    @test issetequal(ps, [p2])
+
+    empty!.((us,ps))
+    MT.collect_scoped_vars!(us, ps, js, iv; depth = 0)
+    @test issetequal(us, [x1])
+    @test issetequal(ps, [p1])
+
+    empty!.((us,ps))
+    MT.collect_scoped_vars!(us, ps, js, iv; depth = 1)
+    @test issetequal(us, [x2])
+    @test issetequal(ps, [p2])
+
+    empty!.((us,ps))
+    MT.collect_scoped_vars!(us, ps, js, iv; depth = 2)
+    @test issetequal(us, [x3, x4])
+    @test issetequal(ps, [p3, p4])
+
+    empty!.((us,ps))
+    MT.collect_scoped_vars!(us, ps, js, iv; depth = -1)
+    @test issetequal(us, [x5])
+    @test issetequal(ps, [p5])
+end