Merge pull request #377 from isaacsas/update-get-variables

update get_variables! for Equations

Author: ChrisRackauckas

src/systems/dependency_graphs.jl

@@ -2,11 +2,12 @@
 # each system type should define extract_variables! for a single equation
 function equation_dependencies(sys::AbstractSystem; variables=states(sys))
     eqs  = equations(sys)
-    deps = Set{Variable}()
+    deps = Set{Operation}()
     depeqs_to_vars = Vector{Vector{Variable}}(undef,length(eqs))
 
     for (i,eq) in enumerate(eqs)      
-        depeqs_to_vars[i] = collect(get_variables!(deps, eq, variables))
+        get_variables!(deps, eq, variables)
+        depeqs_to_vars[i] = [convert(Variable,v) for v in deps]
         empty!(deps)
     end
 
@@ -57,11 +58,11 @@ function variable_dependencies(sys::AbstractSystem; variables=states(sys), varia
     eqs   = equations(sys)
     vtois = isnothing(variablestoids) ? Dict(convert(Variable, v) => i for (i,v) in enumerate(variables)) : variablestoids
 
-    deps = Set{Variable}()
+    deps = Set{Operation}()
     badjlist = Vector{Vector{Int}}(undef, length(eqs))    
     for (eidx,eq) in enumerate(eqs)
         modified_states!(deps, eq, variables)
-        badjlist[eidx] = sort!([vtois[var] for var in deps])
+        badjlist[eidx] = sort!([vtois[convert(Variable,var)] for var in deps])
         empty!(deps)
     end
 

src/systems/jumps/jumpsystem.jl

@@ -139,39 +139,26 @@ end
 
 
 ### Functions to determine which states a jump depends on
-function get_variables!(dep, jump::Union{ConstantRateJump,VariableRateJump}, variables)
-    foreach(var -> (var in variables) && push!(dep, var), vars(jump.rate))
-    dep
-end
+get_variables!(dep, jump::Union{ConstantRateJump,VariableRateJump}, variables) = get_variables!(dep, jump.rate, variables)
 
 function get_variables!(dep, jump::MassActionJump, variables)
-    jsr = jump.scaled_rates
-
-    if jsr isa Variable
-        (jsr in variables) && push!(dep, jsr)
-    elseif jsr isa Operation
-        foreach(var -> (var in variables) && push!(dep, var),  vars(jsr))            
-    end
-
+    get_variables!(dep, jump.scaled_rates, variables)
     for varasop in jump.reactant_stoch
-        var = convert(Variable, varasop[1])
-        (var in variables) && push!(dep, var)
+        (varasop[1].op in variables) && push!(dep, varasop[1])
     end
-
     dep
 end
 
 ### Functions to determine which states are modified by a given jump
 function modified_states!(mstates, jump::Union{ConstantRateJump,VariableRateJump}, sts)
     for eq in jump.affect!
-        st = convert(Variable, eq.lhs)
-        (st in sts) && push!(mstates, st)
+        st = eq.lhs 
+        (st.op in sts) && push!(mstates, st)
     end
 end
 
 function modified_states!(mstates, jump::MassActionJump, sts)
     for (state,stoich) in jump.net_stoch
-        st = convert(Variable, state)
-        (st in sts) && push!(mstates, st)
+        (state.op in sts) && push!(mstates, state)
     end
 end

src/utils.jl

@@ -88,15 +88,25 @@ get_variables(O::Operation)
 
 Returns the variables in the Operation
 """
-get_variables(e::Constant, vars = Operation[]) = vars
-function get_variables(e::Operation, vars = Operation[])
+get_variables!(vars, e::Constant, varlist=nothing) = vars
+get_variables(e::Constant, varlist=nothing) = get_variables!(Operation[], e, varlist)
+
+function get_variables!(vars, e::Operation, varlist=nothing)
     if is_singleton(e)
-        push!(vars, e)
+      (isnothing(varlist) ? true : (e.op in varlist)) && push!(vars, e)
     else
-        foreach(x -> get_variables(x, vars), e.args)
+        foreach(x -> get_variables!(vars, x, varlist), e.args)
     end
     return unique(vars)
 end
+get_variables(e::Operation, varlist=nothing) = get_variables!(Operation[], e, varlist)
+
+function get_variables!(vars, e::Equation, varlist=nothing)
+  get_variables!(vars, e.rhs, varlist)
+end
+get_variables(e::Equation, varlist=nothing) = get_variables!(Operation[],e,varlist)
+
+modified_states!(mstates, e::Equation, statelist=nothing) = get_variables!(mstates, e.lhs, statelist)
 
 # variable substitution
 """
@@ -119,4 +129,4 @@ function _substitute(expr, dict::Dict)
     simplify(SymbolicUtils.substitute(expr, dict))
 end
 
-@deprecate substitute_expr!(expr,s) substitute(expr,s)
+@deprecate substitute_expr!(expr,s) substitute(expr,s)

test/dep_graphs.jl

@@ -14,7 +14,7 @@ rs = ReactionSystem(rxs, t, [S,I,R], [k1,k2])
 
 
 ################################# 
-#      testing for Jumps 
+#  testing for Jumps / all dgs
 #################################
 js = convert(JumpSystem, rs)
 S  = convert(Variable,S); I = convert(Variable,I); R = convert(Variable,R)
@@ -71,4 +71,32 @@ end
 dg4 = varvar_dependencies(depsbg,deps2)
 @test dg == dg4
 
-   
+#####################################
+#       testing for ODE/SDEs
+#####################################
+os = convert(ODESystem, rs)
+deps = equation_dependencies(os)
+eq_sdeps  = [[S,I], [S,I], [S,I,R]]
+@test all(i -> isequal(Set(eq_sdeps[i]),Set(deps[i])), 1:length(deps))
+
+sdes = convert(SDESystem, rs)
+deps = equation_dependencies(sdes)
+@test all(i -> isequal(Set(eq_sdeps[i]),Set(deps[i])), 1:length(deps))
+
+deps = variable_dependencies(os)
+s_eqdeps = [[1],[2],[3]]
+@test deps.fadjlist == s_eqdeps
+
+#####################################
+#       testing for nonlin sys
+#####################################
+@variables x y z
+@parameters σ ρ β
+
+eqs = [0 ~ σ*(y-x),
+       0 ~ ρ-y,
+       0 ~ y - β*z]
+ns = NonlinearSystem(eqs, [x,y,z],[σ,ρ,β])
+deps = equation_dependencies(ns)
+eq_sdeps = [[x,y],[y],[y,z]]
+@test all(i -> isequal(Set(deps[i]),Set(convert.(Variable,eq_sdeps[i]))), 1:length(deps))

test/jumpsystem.jl

@@ -118,11 +118,14 @@ jprob = JumpProblem(js3, dprob, Direct())
 m3 = getmean(jprob,Nsims)
 @test abs(m-m3)/m < .01
 
-# maj jump test with dep graphs
+# maj jump test with various dep graphs
 js3b = JumpSystem([maj1,maj2], t, [S,I,R], [β,γ])
 jprobb = JumpProblem(js3b, dprob, NRM())
 m4 = getmean(jprobb,Nsims)
 @test abs(m-m4)/m < .01
+jprobc = JumpProblem(js3b, dprob, RSSA())
+m4 = getmean(jprobc,Nsims)
+@test abs(m-m4)/m < .01
 
 # mass action jump tests for other reaction types (zero order, decay)
 maj1 = MassActionJump(2.0, [0 => 1], [S => 1])