Clean up graph initialization code

Author: YingboMa

src/bipartite_graph.jl

@@ -66,11 +66,13 @@ badjlist = [[1,2,5,6],[3,4,6]]
 bg = BipartiteGraph(7, fadjlist, badjlist)
 ```
 """
-mutable struct BipartiteGraph{I<:Integer} <: LightGraphs.AbstractGraph{I}
+mutable struct BipartiteGraph{I<:Integer,M} <: LightGraphs.AbstractGraph{I}
     ne::Int
     fadjlist::Vector{Vector{I}} # `fadjlist[src] => dsts`
     badjlist::Vector{Vector{I}} # `badjlist[dst] => srcs`
+    metadata::M
 end
+BipartiteGraph(ne::Integer, fadj::AbstractVector, badj::AbstractVector) = BipartiteGraph(ne, fadj, badj, nothing)
 
 """
 ```julia
@@ -91,9 +93,13 @@ $(SIGNATURES)
 
 Build an empty `BipartiteGraph` with `nsrcs` sources and `ndsts` destinations.
 """
-BipartiteGraph(nsrcs::T, ndsts::T) where T = BipartiteGraph(0, map(_->T[], 1:nsrcs), map(_->T[], 1:ndsts))
+function BipartiteGraph(nsrcs::T, ndsts::T; metadata=nothing) where T
+    fadjlist = map(_->T[], 1:nsrcs)
+    badjlist = map(_->T[], 1:ndsts)
+    BipartiteGraph(0, fadjlist, badjlist, metadata)
+end
 
-Base.eltype(::Type{BipartiteGraph{I}}) where I = I
+Base.eltype(::Type{<:BipartiteGraph{I}}) where I = I
 Base.empty!(g::BipartiteGraph) = (foreach(empty!, g.fadjlist); foreach(empty!, g.badjlist); g.ne = 0; g)
 Base.length(::BipartiteGraph) = error("length is not well defined! Use `ne` or `nv`.")
 
@@ -106,8 +112,8 @@ LightGraphs.vertices(g::BipartiteGraph) = (𝑠vertices(g), 𝑑vertices(g))
 𝑑vertices(g::BipartiteGraph) = axes(g.badjlist, 1)
 has_𝑠vertex(g::BipartiteGraph, v::Integer) = v in 𝑠vertices(g)
 has_𝑑vertex(g::BipartiteGraph, v::Integer) = v in 𝑑vertices(g)
-𝑠neighbors(g::BipartiteGraph, i::Integer) = g.fadjlist[i]
-𝑑neighbors(g::BipartiteGraph, i::Integer) = g.badjlist[i]
+𝑠neighbors(g::BipartiteGraph, i::Integer, with_metadata::Val{M}=Val(false)) where M = M ? zip(g.fadjlist[i], g.metadata[i]) : g.fadjlist[i]
+𝑑neighbors(g::BipartiteGraph, j::Integer, with_metadata::Val{M}=Val(false)) where M = M ? zip(g.badjlist[j], (g.metadata[i][j] for i in g.badjlist[j])) : g.badjlist[j]
 LightGraphs.ne(g::BipartiteGraph) = g.ne
 LightGraphs.nv(g::BipartiteGraph) = sum(length, vertices(g))
 LightGraphs.edgetype(g::BipartiteGraph{I}) where I = BipartiteEdge{I}
@@ -170,7 +176,7 @@ Base.length(it::BipartiteEdgeIter{ALL}) = 2ne(it.g)
 
 Base.eltype(it::BipartiteEdgeIter) = edgetype(it.g)
 
-function Base.iterate(it::BipartiteEdgeIter{SRC,BipartiteGraph{T}}, state=(1, 1, SRC)) where T
+function Base.iterate(it::BipartiteEdgeIter{SRC,<:BipartiteGraph{T}}, state=(1, 1, SRC)) where T
     @unpack g = it
     neqs = nsrcs(g)
     neqs == 0 && return nothing
@@ -191,7 +197,7 @@ function Base.iterate(it::BipartiteEdgeIter{SRC,BipartiteGraph{T}}, state=(1, 1,
     return nothing
 end
 
-function Base.iterate(it::BipartiteEdgeIter{DST,BipartiteGraph{T}}, state=(1, 1, DST)) where T
+function Base.iterate(it::BipartiteEdgeIter{DST,<:BipartiteGraph{T}}, state=(1, 1, DST)) where T
     @unpack g = it
     nvars = ndsts(g)
     nvars == 0 && return nothing

src/systems/alias_elimination.jl

@@ -1,5 +1,3 @@
-export alias_elimination, flatten
-
 using SymbolicUtils: Rewriters
 
 function fixpoint_sub(x, dict)

src/systems/diffeqs/odesystem.jl

@@ -111,18 +111,31 @@ iv_from_nested_derivative(x) = missing
 vars(x::Sym) = [x]
 vars(exprs::Symbolic) = vars([exprs])
 vars(exprs) = foldl(vars!, exprs; init = Set())
+vars!(vars, eq::Equation) = (vars!(vars, eq.lhs); vars!(vars, eq.rhs); vars)
 function vars!(vars, O)
     isa(O, Sym) && return push!(vars, O)
-    !isa(O, Symbolic) && return vars
+    !istree(O) && return vars
+
+    operation(O) isa Differential && return push!(vars, O)
 
     operation(O) isa Sym && push!(vars, O)
-    for arg ∈ arguments(O)
+    for arg in arguments(O)
         vars!(vars, arg)
     end
 
     return vars
 end
 
+find_derivatives!(vars, expr::Equation, f=identity) = (find_derivatives!(vars, expr.lhs, f); find_derivatives!(vars, expr.rhs, f); vars)
+function find_derivatives!(vars, expr, f)
+    !istree(O) && return vars
+    operation(O) isa Differential && push!(vars, f(O))
+    for arg in arguments(O)
+        vars!(vars, arg)
+    end
+    return vars
+end
+
 function ODESystem(eqs, iv=nothing; kwargs...)
     # NOTE: this assumes that the order of algebric equations doesn't matter
     diffvars = OrderedSet()

src/systems/systemstructure.jl

@@ -1,5 +1,7 @@
 module SystemStructures
 
+using DataStructures
+using SymbolicUtils: istree, operation
 using ..ModelingToolkit
 import ..ModelingToolkit: isdiffeq, var_from_nested_derivative, vars!, flatten
 using SymbolicUtils: arguments
@@ -47,8 +49,8 @@ struct SystemStructure
     fullvars::Vector # [xvar; dxvars; algvars]
     varassoc::Vector{Int}
     algeqs::BitVector
-    graph::BipartiteGraph{Int}
-    solvable_graph::BipartiteGraph{Int}
+    graph::BipartiteGraph{Int,Nothing}
+    solvable_graph::BipartiteGraph{Int,Vector{Vector{Int}}}
     assign::Vector{Int}
     inv_assign::Vector{Int}
     scc::Vector{Vector{Int}}
@@ -108,55 +110,45 @@ function Base.show(io::IO, s::SystemStructure)
     show(io, S)
 end
 
-# V-nodes `[x_1, x_2, x_3, ..., dx_1, dx_2, ..., y_1, y_2, ...]` where `x`s are
-# differential variables and `y`s are algebraic variables.
-function collect_variables(sys)
-    dxvars = []
+function init_graph(sys)
+    iv = independent_variable(sys)
     eqs = equations(sys)
-    algeqs = trues(length(eqs))
-    for (i, eq) in enumerate(eqs)
-        if isdiffeq(eq)
-            algeqs[i] = false
-            lhs = eq.lhs
-            # Make sure that the LHS is a first order derivative of a var.
-            @assert !(arguments(lhs)[1] isa Differential) "The equation $eq is not first order"
+    neqs = length(eqs)
+    algeqs = trues(neqs)
+    varsadj = Vector{Any}(undef, neqs)
+    dervars = OrderedSet()
+    diffvars = OrderedSet()
 
-            push!(dxvars, lhs)
+    for (i, eq) in enumerate(eqs)
+        vars = OrderedSet()
+        vars!(vars, eq)
+        for var in vars
+            if istree(var) && operation(var) isa Differential
+                diffvar = arguments(var)[1]
+                @assert !(diffvar isa Differential) "The equation [ $eq ] is not first order"
+                push!(dervars, var)
+                push!(diffvars, diffvar)
+            end
         end
+        varsadj[i] = vars
     end
 
-    xvars = (first ∘ var_from_nested_derivative).(dxvars)
-    algvars  = setdiff(states(sys), xvars)
-    return xvars, dxvars, algvars, algeqs
-end
+    algvars  = setdiff(states(sys), diffvars)
+    fullvars = [collect(diffvars); collect(dervars); algvars]
 
-function init_graph(sys)
-    xvars, dxvars, algvars, algeqs = collect_variables(sys)
-    dxvar_offset = length(xvars)
+    dxvar_offset = length(diffvars)
     algvar_offset = 2dxvar_offset
 
-    fullvars = [xvars; dxvars; algvars]
-    eqs = equations(sys)
-    idxmap = Dict(fullvars .=> 1:length(fullvars))
-    graph = BipartiteGraph(length(eqs), length(fullvars))
-    solvable_graph = BipartiteGraph(length(eqs), length(fullvars))
+    nvars = length(fullvars)
+    idxmap = Dict(fullvars .=> 1:nvars)
+    graph = BipartiteGraph(neqs, nvars)
+    solvable_graph = BipartiteGraph(neqs, nvars, metadata=Vector{Int}[])
 
-    vs = Set()
-    for (i, eq) in enumerate(eqs)
-        # TODO: custom vars that handles D(x)
-        # TODO: add checks here
-        lhs = eq.lhs
-        if isdiffeq(eq)
-            v = lhs
-            haskey(idxmap, v) && add_edge!(graph, i, idxmap[v])
-        else
-            vars!(vs, lhs)
-        end
-        vars!(vs, eq.rhs)
+    for (i, vs) in enumerate(varsadj)
         for v in vs
-            haskey(idxmap, v) && add_edge!(graph, i, idxmap[v])
+            j = get(idxmap, v, nothing)
+            j === nothing || add_edge!(graph, i, idxmap[v])
         end
-        empty!(vs)
     end
 
     varassoc = Int[(1:dxvar_offset) .+ dxvar_offset; zeros(Int, length(fullvars) - dxvar_offset)] # variable association list