add vrj dep graph tests

Author: isaacsas

src/systems/dependency_graphs.jl

@@ -36,8 +36,8 @@ equation_dependencies(jumpsys)
 equation_dependencies(jumpsys, variables = parameters(jumpsys))
 ```
 """
-function equation_dependencies(sys::AbstractSystem; variables = unknowns(sys))
-    eqs = equations(sys)
+function equation_dependencies(sys::AbstractSystem; variables = unknowns(sys),
+        eqs = equations(sys))
     deps = Set()
     depeqs_to_vars = Vector{Vector}(undef, length(eqs))
 
@@ -114,8 +114,9 @@ digr = asgraph(jumpsys)
 ```
 """
 function asgraph(sys::AbstractSystem; variables = unknowns(sys),
-        variablestoids = Dict(v => i for (i, v) in enumerate(variables)))
-    asgraph(equation_dependencies(sys, variables = variables), variablestoids)
+        variablestoids = Dict(v => i for (i, v) in enumerate(variables)),
+        eqs = equations(sys))
+    asgraph(equation_dependencies(sys; variables, eqs), variablestoids)
 end
 
 """
@@ -141,8 +142,8 @@ variable_dependencies(jumpsys)
 ```
 """
 function variable_dependencies(sys::AbstractSystem; variables = unknowns(sys),
-        variablestoids = nothing)
-    eqs = equations(sys)
+        variablestoids = nothing, eqs = equations(sys))
+
     vtois = isnothing(variablestoids) ? Dict(v => i for (i, v) in enumerate(variables)) :
             variablestoids
 
@@ -193,8 +194,8 @@ dg = asdigraph(digr, jumpsys)
 ```
 """
 function asdigraph(g::BipartiteGraph, sys::AbstractSystem; variables = unknowns(sys),
-        equationsfirst = true)
-    neqs = length(equations(sys))
+        equationsfirst = true, eqs = equations(sys))
+    neqs = length(eqs)
     nvars = length(variables)
     fadjlist = deepcopy(g.fadjlist)
     badjlist = deepcopy(g.badjlist)

src/systems/jumps/jumpsystem.jl

@@ -502,10 +502,12 @@ function JumpProcesses.JumpProblem(js::JumpSystem, prob,
         error("Use continuous problems such as an ODEProblem or a SDEProblem with VariableRateJumps")
     jset = JumpSet(Tuple(vrjs), Tuple(crjs), nothing, majs)
 
+    # dep graphs are only for constant rate jumps
+    nonvrjs = ArrayPartition(eqs.x[1], eqs.x[2])
     if needs_vartojumps_map(aggregator) || needs_depgraph(aggregator) ||
        (aggregator isa JumpProcesses.NullAggregator)
-        jdeps = asgraph(js)
-        vdeps = variable_dependencies(js)
+        jdeps = asgraph(js; eqs = nonvrjs)
+        vdeps = variable_dependencies(js; eqs = nonvrjs)
         vtoj = jdeps.badjlist
         jtov = vdeps.badjlist
         jtoj = needs_depgraph(aggregator) ? eqeq_dependencies(jdeps, vdeps).fadjlist :

test/dep_graphs.jl

@@ -1,5 +1,5 @@
 using Test
-using ModelingToolkit, Graphs, JumpProcesses
+using ModelingToolkit, Graphs, JumpProcesses, RecursiveArrayTools
 using ModelingToolkit: t_nounits as t, D_nounits as D
 import ModelingToolkit: value
 
@@ -72,6 +72,80 @@ end
 dg4 = varvar_dependencies(depsbg, deps2)
 @test dg == dg4
 
+# testing when ignoring VariableRateJumps
+let
+    @parameters k1 k2
+    @variables S(t) I(t) R(t)
+    j₁ = MassActionJump(k1, [0 => 1], [S => 1])
+    j₂ = MassActionJump(k1, [S => 1], [S => -1])
+    j₃ = MassActionJump(k2, [S => 1, I => 1], [S => -1, I => 1])
+    j₄ = MassActionJump(k2, [S => 2, R => 1], [R => -1])
+    j₅ = ConstantRateJump(k1 * I, [R ~ R + 1])
+    j₆ = VariableRateJump(k1 * k2 / (1 + t) * S, [S ~ S - 1, R ~ R + 1])
+    eqs = [j₁, j₂, j₃, j₄, j₅, j₆]
+    @named js = JumpSystem(eqs, t, [S, I, R], [k1, k2])
+    S = value(S);
+    I = value(I);
+    R = value(R);
+    k1 = value(k1);
+    k2 = value(k2);
+    # eq to vars they depend on
+    eq_sdeps = [Variable[], [S], [S, I], [S, R], [I]]
+    eq_sidepsf = [Int[], [1], [1, 2], [1, 3], [2]]
+    eq_sidepsb = [[2, 3, 4], [3, 5], [4]]
+
+    # filter out vrjs in making graphs
+    eqs = ArrayPartition(equations(js).x[1], equations(js).x[2])
+    deps = equation_dependencies(js; eqs)
+    @test length(deps) == length(eq_sdeps)
+    @test all(i -> isequal(Set(eq_sdeps[i]), Set(deps[i])), 1:length(eqs))
+    depsbg = asgraph(js; eqs)
+    @test depsbg.fadjlist == eq_sidepsf
+    @test depsbg.badjlist == eq_sidepsb
+
+    # eq to params they depend on
+    eq_pdeps = [[k1], [k1], [k2], [k2], [k1]]
+    eq_pidepsf = [[1], [1], [2], [2], [1]]
+    eq_pidepsb = [[1, 2, 5], [3, 4]]
+    deps = equation_dependencies(js; variables = parameters(js), eqs)
+    @test length(deps) == length(eq_pdeps)
+    @test all(i -> isequal(Set(eq_pdeps[i]), Set(deps[i])), 1:length(eqs))
+    depsbg2 = asgraph(js; variables = parameters(js), eqs)
+    @test depsbg2.fadjlist == eq_pidepsf
+    @test depsbg2.badjlist == eq_pidepsb
+
+    # var to eqs that modify them
+    s_eqdepsf = [[1, 2, 3], [3], [4, 5]]
+    s_eqdepsb = [[1], [1], [1, 2], [3], [3]]
+    ne = 6
+    bg = BipartiteGraph(ne, s_eqdepsf, s_eqdepsb)
+    deps2 = variable_dependencies(js; eqs)
+    @test isequal(bg, deps2)
+
+    # eq to eqs that depend on them
+    eq_eqdeps = [[2, 3, 4], [2, 3, 4], [2, 3, 4, 5], [4], [4], [2, 3, 4]]
+    dg = SimpleDiGraph(5)
+    for (eqidx, eqdeps) in enumerate(eq_eqdeps)
+        for eqdepidx in eqdeps
+            add_edge!(dg, eqidx, eqdepidx)
+        end
+    end
+    dg3 = eqeq_dependencies(depsbg, deps2)
+    @test dg == dg3
+
+    # var to vars that depend on them
+    var_vardeps = [[1, 2, 3], [1, 2, 3], [3]]
+    ne = 7
+    dg = SimpleDiGraph(3)
+    for (vidx, vdeps) in enumerate(var_vardeps)
+        for vdepidx in vdeps
+            add_edge!(dg, vidx, vdepidx)
+        end
+    end
+    dg4 = varvar_dependencies(depsbg, deps2)
+    @test dg == dg4
+end
+
 #####################################
 #       testing for ODE/SDEs
 #####################################