Add solvable_graph basics

Author: YingboMa

src/bipartite_graph.jl

@@ -130,8 +130,12 @@ end
 ###
 ### Populate
 ###
-LightGraphs.add_edge!(g::BipartiteGraph, i::Integer, j::Integer) = add_edge!(g, BipartiteEdge(i, j))
-function LightGraphs.add_edge!(g::BipartiteGraph, edge::BipartiteEdge)
+struct NoMetadata
+end
+const NO_METADATA = NoMetadata()
+
+LightGraphs.add_edge!(g::BipartiteGraph, i::Integer, j::Integer, md=NO_METADATA) = add_edge!(g, BipartiteEdge(i, j), md)
+function LightGraphs.add_edge!(g::BipartiteGraph, edge::BipartiteEdge, md=NO_METADATA)
     @unpack fadjlist, badjlist = g
     verts = vertices(g)
     s, d = src(edge), dst(edge)
@@ -140,6 +144,9 @@ function LightGraphs.add_edge!(g::BipartiteGraph, edge::BipartiteEdge)
     index = searchsortedfirst(list, d)
     @inbounds (index <= length(list) && list[index] == d) && return false  # edge already in graph
     insert!(list, index, d)
+    if md !== NO_METADATA
+        insert!(g.metadata[s], index, md)
+    end
 
     g.ne += 1
     @inbounds list = badjlist[d]

src/systems/abstractsystem.jl

@@ -196,6 +196,7 @@ function Base.getproperty(sys::AbstractSystem, name::Symbol)
     end
 
     sts = get_states(sys)
+    @show sts
     i = findfirst(x->getname(x) == name, sts)
 
     if i !== nothing

src/systems/alias_elimination.jl

@@ -10,8 +10,9 @@ function fixpoint_sub(x, dict)
     return x
 end
 
-function substitute_aliases(diffeqs, dict)
-    lhss(diffeqs) .~ fixpoint_sub.(rhss(diffeqs), (dict,))
+function substitute_aliases(eqs, dict)
+    sub = Base.Fix2(fixpoint_sub, dict)
+    map(eq->eq.lhs ~ sub(eq.rhs), eqs)
 end
 
 # Note that we reduce parameters, too
@@ -68,8 +69,13 @@ function maybe_alias(lhs, rhs, diff_vars, iv, conservative)
     end
 end
 
-function alias_elimination(sys; conservative=true)
+function alias_elimination(sys)
     sys = flatten(sys)
+    s = get_structure(sys)
+    if !(s isa SystemStructure)
+        sys = initialize_system_structure(sys)
+        s = structure(sys)
+    end
     iv = independent_variable(sys)
     eqs = equations(sys)
     diff_vars = filter(!isnothing, map(eqs) do eq
@@ -124,7 +130,7 @@ function alias_elimination(sys; conservative=true)
     @set! sys.eqs = substitute_aliases(neweqs, Dict(subs))
     @set! sys.states = newstates
     @set! sys.observed = [observed(sys); alias_eqs]
-    return initialize_system_structure(sys)
+    return 
 end
 
 """

src/systems/diffeqs/odesystem.jl

@@ -178,6 +178,9 @@ function collect_vars!(states, parameters, expr, iv)
         collect_var!(states, parameters, expr, iv)
     else
         for var in vars(expr)
+            if istree(var) && operation(var) isa Differential
+                var, _ = var_from_nested_derivative(var)
+            end
             collect_var!(states, parameters, var, iv)
         end
     end

src/systems/systemstructure.jl

@@ -1,10 +1,9 @@
 module SystemStructures
 
 using DataStructures
-using SymbolicUtils: istree, operation
+using SymbolicUtils: istree, operation, arguments
 using ..ModelingToolkit
-import ..ModelingToolkit: isdiffeq, var_from_nested_derivative, vars!, flatten
-using SymbolicUtils: arguments
+import ..ModelingToolkit: isdiffeq, var_from_nested_derivative, vars!, flatten, value
 using ..BipartiteGraphs
 using UnPack
 using Setfield
@@ -46,7 +45,7 @@ export vartype, eqtype
 
 struct SystemStructure
     dxvar_offset::Int
-    fullvars::Vector # [xvar; dxvars; algvars]
+    fullvars::Vector # [diffvars; dervars; algvars]
     varassoc::Vector{Int}
     algeqs::BitVector
     graph::BipartiteGraph{Int,Nothing}
@@ -122,14 +121,17 @@ function init_graph(sys)
     for (i, eq) in enumerate(eqs)
         vars = OrderedSet()
         vars!(vars, eq)
+        isalgeq = true
         for var in vars
             if istree(var) && operation(var) isa Differential
+                isalgeq = false
                 diffvar = arguments(var)[1]
                 @assert !(diffvar isa Differential) "The equation [ $eq ] is not first order"
                 push!(dervars, var)
                 push!(diffvars, diffvar)
             end
         end
+        algeqs[i] = isalgeq
         varsadj[i] = vars
     end
 
@@ -142,12 +144,23 @@ function init_graph(sys)
     nvars = length(fullvars)
     idxmap = Dict(fullvars .=> 1:nvars)
     graph = BipartiteGraph(neqs, nvars)
-    solvable_graph = BipartiteGraph(neqs, nvars, metadata=Vector{Int}[])
+    solvable_graph = BipartiteGraph(neqs, nvars, metadata=map(_->Int[], 1:neqs))
 
     for (i, vs) in enumerate(varsadj)
+        eq = eqs[i]
         for v in vs
             j = get(idxmap, v, nothing)
-            j === nothing || add_edge!(graph, i, idxmap[v])
+            if j !== nothing
+                add_edge!(graph, i, idxmap[v])
+                j > algvar_offset || continue
+                D = Differential(fullvars[j])
+                c = value(expand_derivatives(D(eq.rhs - eq.lhs), false))
+                if c isa Number && c != 0
+                    # 0 here is a sentinel value for non-integer coefficients
+                    coeff = c isa Integer ? c : 0
+                    add_edge!(solvable_graph, i, j, coeff)
+                end
+            end
         end
     end
 

test/reduction.jl

@@ -68,6 +68,10 @@ eqs1 = [
 
 lorenz = name -> ODESystem(eqs1,t,name=name)
 lorenz1 = lorenz(:lorenz1)
+ss = ModelingToolkit.get_structure(initialize_system_structure(lorenz1))
+@test isequal(ss.fullvars, [x, y, z, D(x), D(y), D(z), F, u])
+@test ss.solvable_graph.fadjlist == [[7], [8], [], [8]]
+@test ss.solvable_graph.metadata == [[1], [1], [], [1]]
 lorenz2 = lorenz(:lorenz2)
 
 connected = ODESystem([s ~ a + lorenz1.x