WIP: update TearingState in alias_elimination

Author: YingboMa

src/structural_transformation/utils.jl

@@ -43,9 +43,13 @@ end
 ###
 ### Structural check
 ###
-function check_consistency(state::TearingState)
+function check_consistency(state::TearingState, ag = nothing)
     fullvars = state.fullvars
     @unpack graph, var_to_diff = state.structure
+    #n_highest_vars = count(v -> length(outneighbors(var_to_diff, v)) == 0 && !isempty(𝑑neighbors(graph, v)),
+    #                       vertices(var_to_diff))
+    #n_highest_vars = count(v -> length(outneighbors(var_to_diff, v)) == 0 && !haskey(ag, v),
+    #                       vertices(var_to_diff))
     n_highest_vars = count(v -> length(outneighbors(var_to_diff, v)) == 0,
                            vertices(var_to_diff))
     neqs = nsrcs(graph)
@@ -69,11 +73,11 @@ function check_consistency(state::TearingState)
     # details, check the equation (15) of the original paper.
     extended_graph = (@set graph.fadjlist = Vector{Int}[graph.fadjlist;
                                                         map(collect, edges(var_to_diff))])
-    extended_var_eq_matching = maximal_matching(extended_graph)
+    extended_var_eq_matching = maximal_matching(extended_graph, eq->true, v->!haskey(ag, v))
 
     unassigned_var = []
     for (vj, eq) in enumerate(extended_var_eq_matching)
-        if eq === unassigned
+        if eq === unassigned && !haskey(ag, vj)
             push!(unassigned_var, fullvars[vj])
         end
     end

src/systems/abstractsystem.jl

@@ -1035,12 +1035,12 @@ function structural_simplify(sys::AbstractSystem, io = nothing; simplify = false
     has_io = io !== nothing
     has_io && markio!(state, io...)
     state, input_idxs = inputs_to_parameters!(state, io)
-    sys = alias_elimination!(state)
+    sys, ag = alias_elimination!(state)
     # TODO: avoid construct `TearingState` again.
-    state = TearingState(sys)
-    has_io && markio!(state, io..., check = false)
-    check_consistency(state)
-    find_solvables!(state; kwargs...)
+    #state = TearingState(sys)
+    #has_io && markio!(state, io..., check = false)
+    check_consistency(state, ag)
+    #find_solvables!(state; kwargs...)
     sys = dummy_derivative(sys, state; simplify)
     fullstates = [map(eq -> eq.lhs, observed(sys)); states(sys)]
     @set! sys.observed = topsort_equations(observed(sys), fullstates)

src/systems/alias_elimination.jl

@@ -47,7 +47,7 @@ function extreme_var(var_to_diff, v, level = nothing, ::Val{descend} = Val(true)
     level === nothing ? v : (v => level)
 end
 
-alias_elimination(sys) = alias_elimination!(TearingState(sys))
+alias_elimination(sys) = alias_elimination!(TearingState(sys))[1]
 function alias_elimination!(state::TearingState)
     sys = state.sys
     complete!(state.structure)
@@ -56,7 +56,7 @@ function alias_elimination!(state::TearingState)
     isempty(ag) && return sys
 
     fullvars = state.fullvars
-    @unpack var_to_diff, graph = state.structure
+    @unpack var_to_diff, graph, solvable_graph = state.structure
 
     if !isempty(updated_diff_vars)
         has_iv(sys) ||
@@ -105,19 +105,36 @@ function alias_elimination!(state::TearingState)
         end
     end
     deleteat!(eqs, sort!(dels))
-    old_to_new = Vector{Int}(undef, nsrcs(graph))
+    old_to_new_eq = Vector{Int}(undef, nsrcs(graph))
     idx = 0
     cursor = 1
     ndels = length(dels)
-    for i in eachindex(old_to_new)
+    for i in eachindex(old_to_new_eq)
         if cursor <= ndels && i == dels[cursor]
             cursor += 1
-            old_to_new[i] = -1
+            old_to_new_eq[i] = -1
             continue
         end
         idx += 1
-        old_to_new[i] = idx
+        old_to_new_eq[i] = idx
     end
+    n_new_eqs = idx
+
+    old_to_new_var = Vector{Int}(undef, ndsts(graph))
+    idx = 0
+    for i in eachindex(old_to_new_var)
+        if haskey(ag, i)
+            old_to_new_var[i] = -1
+        else
+            idx += 1
+            old_to_new_var[i] = idx
+        end
+    end
+    n_new_vars = idx
+    #for d in dels
+    #    set_neighbors!(graph, d, ())
+    #    set_neighbors!(solvable_graph, d, ())
+    #end
 
     lineqs = BitSet(mm.nzrows)
     eqs_to_update = BitSet()
@@ -126,7 +143,7 @@ function alias_elimination!(state::TearingState)
         while true
             for ieq in 𝑑neighbors(graph_orig, k)
                 ieq in lineqs && continue
-                new_eq = old_to_new[ieq]
+                new_eq = old_to_new_eq[ieq]
                 new_eq < 1 && continue
                 push!(eqs_to_update, new_eq)
             end
@@ -139,7 +156,7 @@ function alias_elimination!(state::TearingState)
     end
 
     for old_ieq in to_expand
-        ieq = old_to_new[old_ieq]
+        ieq = old_to_new_eq[old_ieq]
         eqs[ieq] = expand_derivatives(eqs[ieq])
     end
 
@@ -150,12 +167,53 @@ function alias_elimination!(state::TearingState)
             diff_to_var[j] === nothing && push!(newstates, fullvars[j])
         end
     end
+    #=
+    new_graph = BipartiteGraph(n_new_eqs, ndsts(graph))
+    new_solvable_graph = BipartiteGraph(n_new_eqs, ndsts(graph))
+    new_eq_to_diff = DiffGraph(n_new_eqs)
+    eq_to_diff = state.structure.eq_to_diff
+    for (i, ieq) in enumerate(old_to_new_eq)
+        ieq > 0 || continue
+        set_neighbors!(new_graph, ieq, 𝑠neighbors(graph, i))
+        set_neighbors!(new_solvable_graph, ieq, 𝑠neighbors(solvable_graph, i))
+        new_eq_to_diff[ieq] = eq_to_diff[i]
+    end
+    state.structure.graph = new_graph
+    state.structure.solvable_graph = new_solvable_graph
+    state.structure.eq_to_diff = new_eq_to_diff
+    @show length(new_eq_to_diff), nsrcs(new_graph), nsrcs(new_solvable_graph), length(eqs)
+    =#
+
+    new_graph = BipartiteGraph(n_new_eqs, n_new_vars)
+    new_solvable_graph = BipartiteGraph(n_new_eqs, n_new_vars)
+    new_eq_to_diff = DiffGraph(n_new_eqs)
+    eq_to_diff = state.structure.eq_to_diff
+    new_var_to_diff = DiffGraph(n_new_vars)
+    var_to_diff = state.structure.var_to_diff
+    for (i, ieq) in enumerate(old_to_new_eq)
+        ieq > 0 || continue
+        set_neighbors!(new_graph, ieq, [old_to_new_var[v] for v in 𝑠neighbors(graph, i) if old_to_new_var[v] > 0])
+        set_neighbors!(new_solvable_graph, ieq, [old_to_new_var[v] for v in 𝑠neighbors(solvable_graph, i) if old_to_new_var[v] > 0])
+        new_eq_to_diff[ieq] = eq_to_diff[i]
+    end
+    new_fullvars = Vector{Any}(undef, n_new_vars)
+    for (i, iv) in enumerate(old_to_new_var)
+        iv > 0 || continue
+        new_var_to_diff[iv] = var_to_diff[i]
+        new_fullvars[iv] = fullvars[i]
+    end
+    state.structure.graph = new_graph
+    state.structure.solvable_graph = new_solvable_graph
+    state.structure.eq_to_diff = complete(new_eq_to_diff)
+    state.structure.var_to_diff = complete(new_var_to_diff)
+    state.fullvars = new_fullvars
 
     sys = state.sys
     @set! sys.eqs = eqs
     @set! sys.states = newstates
     @set! sys.observed = [observed(sys); obs]
-    return invalidate_cache!(sys)
+    state.sys = sys
+    return invalidate_cache!(sys), ag
 end
 
 """