Faster alias_elimination

Author: YingboMa

src/bipartite_graph.jl

@@ -314,7 +314,8 @@ function BipartiteGraph(nsrcs::T, ndsts::T, backedge::Val{B} = Val(true);
 end
 
 function Base.copy(bg::BipartiteGraph)
-    BipartiteGraph(bg.ne, copy(bg.fadjlist), copy(bg.badjlist), deepcopy(bg.metadata))
+    BipartiteGraph(bg.ne, map(copy, bg.fadjlist), map(copy, bg.badjlist),
+                   deepcopy(bg.metadata))
 end
 Base.eltype(::Type{<:BipartiteGraph{I}}) where {I} = I
 function Base.empty!(g::BipartiteGraph)

src/systems/alias_elimination.jl

@@ -47,10 +47,11 @@ 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; quick_cancel = true))
+alias_elimination(sys) = alias_elimination!(TearingState(sys))
 function alias_elimination!(state::TearingState)
     sys = state.sys
     complete!(state.structure)
+    graph_orig = copy(state.structure.graph)
     ag, mm, complete_ag, complete_mm, updated_diff_vars = alias_eliminate_graph!(state)
     isempty(ag) && return sys
 
@@ -118,10 +119,23 @@ function alias_elimination!(state::TearingState)
         old_to_new[i] = idx
     end
 
-    lineqs = BitSet(old_to_new[e] for e in mm.nzrows)
-    for (ieq, eq) in enumerate(eqs)
-        ieq in lineqs && continue
-        eqs[ieq] = substitute(eq, subs)
+    lineqs = BitSet(mm.nzrows)
+    eqs_to_update = BitSet()
+    for k in keys(ag)
+        # We need to update `D(D(x))` when we subsitute `D(x)` as well.
+        while true
+            for ieq in 𝑑neighbors(graph_orig, k)
+                ieq in lineqs && continue
+                new_eq = old_to_new[ieq]
+                new_eq < 1 && continue
+                push!(eqs_to_update, new_eq)
+            end
+            k = var_to_diff[k]
+            k === nothing && break
+        end
+    end
+    for ieq in eqs_to_update
+        eqs[ieq] = substitute(eqs[ieq], subs)
     end
 
     for old_ieq in to_expand
@@ -817,9 +831,9 @@ end
 function update_graph_neighbors!(graph, ag)
     for eq in 1:nsrcs(graph)
         set_neighbors!(graph, eq,
-                       [get(ag, n, (1, n))[2]
-                        for n in 𝑠neighbors(graph, eq)
-                        if !haskey(ag, n) || ag[n][2] != 0])
+                       Int[get(ag, n, (1, n))[2]
+                           for n in 𝑠neighbors(graph, eq)
+                           if !haskey(ag, n) || ag[n][2] != 0])
     end
     return graph
 end

src/systems/systemstructure.jl

@@ -25,33 +25,6 @@ function quick_cancel_expr(expr)
                                                                          kws...))(expr)
 end
 
-#=
-When we don't do subsitution, variable information is split into two different
-places, i.e. `states` and the right-hand-side of `observed`.
-
-eqs = [0 ~ z + x; 0 ~ y + z^2]
-states = [y, z]
-observed = [x ~ sin(y) + z]
-struct Reduced
-    var
-    expr
-    idxs
-end
-fullvars = [Reduced(x, sin(y) + z, [2, 3]), y, z]
-active_𝑣vertices = [false, true, true]
-      x   y   z
-eq1:  1       1
-eq2:      1   1
-
-      x   y   z
-eq1:      1   1
-eq2:      1   1
-
-for v in 𝑣vertices(graph); active_𝑣vertices[v] || continue
-
-end
-=#
-
 export SystemStructure, TransformationState, TearingState
 export initialize_system_structure, find_linear_equations
 export isdiffvar, isdervar, isalgvar, isdiffeq, isalgeq, algeqs