eq to var dependencies

Author: isaacsas

src/ModelingToolkit.jl

@@ -119,7 +119,8 @@ export Differential, expand_derivatives, @derivatives
 export IntervalDomain, ProductDomain, ⊗, CircleDomain
 export Equation, ConstrainedEquation
 export Operation, Expression, Variable
-export independent_variable, states, parameters, equations
+export independent_variable, states, parameters, equations 
+export numstates, numparameters, numequations
 
 export calculate_jacobian, generate_jacobian, generate_function
 export calculate_tgrad, generate_tgrad

src/systems/abstractsystem.jl

@@ -203,6 +203,10 @@ function parameters(sys::AbstractSystem,args...)
     rename(x,renamespace(sys.name,newname))()
 end
 
+numstates(sys::AbstractSystem)     = length(states(sys))
+numparameters(sys::AbstractSystem) = length(parameters(sys))
+numequations(sys::AbstractSystem)  = length(equations(sys))
+
 struct AbstractSysToExpr
     sys::AbstractSystem
     states::Vector{Variable}

src/systems/dependency_graphs.jl

@@ -1,24 +1,58 @@
-
 # each system type should define extract_variables! for a single equation
 function extract_variables(eqs, vars)
     deps = Set{Variable}()
     depeqs_to_vars = Vector{Vector{Variable}}(undef,length(eqs))
 
-    for (i,eq) in enumerate(eqs)        
+    for (i,eq) in enumerate(eqs)      
         depeqs_to_vars[i] = collect(extract_variables!(deps, eq, vars))
         empty!(deps)
     end
+
+    depeqs_to_vars
+end
+
+# variables equations depend on as a vector of vectors of variables
+equation_dependencies(sys::AbstractSystem; variables=states(sys)) = extract_variables(equations(sys), variables)
+
+# modeled on LightGraphs SimpleGraph
+mutable struct BiPartiteGraph{T <: Integer}
+    ne::Int
+    fadjlist::Vector{Vector{T}}  # fadjlist[src] = [dest1,dest2,...]
+    badjlist::Vector{Vector{T}}  # badjlist[dst] = [src1,src2,...]
+end
+
+# convert equation-variable dependencies to a bipartite graph
+function asgraph(eqdeps, vtois)    
+    fadjlist = Vector{Vector{Int}}(undef, length(eqdeps))
+    for (i,dep) in enumerate(eqdeps)
+        fadjlist[i] = sort!([vtois[var] for var in dep])
+    end
+
+    badjlist = [Vector{Int}() for i = 1:length(vtois)]
+    ne = 0
+    for (eqidx,vidxs) in enumerate(fadjlist)
+        foreach(vidx -> push!(badjlist[vidx], eqidx), vidxs)
+        ne += 1
+    end
+
+    BiPartiteGraph(ne, fadjlist, badjlist)
 end
 
-equation_dependencies(sys) = extract_variables(equations(sys), states(sys))
+# could be made to directly generate graph and save memory
+function asgraph(sys::AbstractSystem; variables=nothing, variablestoids=nothing)
+    vs   = isnothing(variables) ? states(sys) : variables    
+    eqdeps = extract_variables(equations(sys), vs)
+    vtois  = isnothing(variablestoids) ? Dict(convert(Variable, v) => i for (i,v) in enumerate(vs)) : variablestoids
+    asgraph(eqdeps, vtois)
+end
 
 # map each variable to the eqs depending on it 
 function variables_to_depeqs(sys::AbstractSystem; equationdeps = nothing, statestoids = nothing)
     sts   = states(sys)
     stoi  = isnothing(statestoids) ? Dict(convert(Variable,state) => i for (i,state) in enumerate(sts)) : statestoids
 
     # map from eqs to states they depend on
-    eqdeps = isnothing(equationdeps) ? equation_dependencies(equations(sys), sts) : equationdeps
+    eqdeps = isnothing(equationdeps) ? equation_dependencies(sys) : equationdeps
 
     # reverse map and switch to integer indices of states
     dg = [Vector{Int}() for i = 1:length(sts)]

src/systems/jumps/jumpsystem.jl

@@ -128,14 +128,25 @@ end
 
 
 ### Functions to determine which states a jump depends on
-
-function extract_variables!(dep, jump::Union{ConstantRateJump,VariableRateJump}, sts)
-    foreach(st -> (convert(Variable,st) in sts) && push!(dep,st), vars(jump.rate))
+function extract_variables!(dep, jump::Union{ConstantRateJump,VariableRateJump}, variables)
+    foreach(var -> (var in variables) && push!(dep, var), vars(jump.rate))
     dep
 end
 
-function extract_variables!(dep, jump::MassActionJump, args...)    
-    foreach(st -> push!(dep,convert(Variable,st[1])), jump.reactant_stoch)    
+function extract_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
+
+    for varasop in jump.reactant_stoch
+        var = convert(Variable, varasop[1])
+        (var in variables) && push!(dep, var)
+    end
+
     dep
 end