Add BitDict

Author: YingboMa

src/structural_transformation/partial_state_selection.jl

@@ -240,7 +240,8 @@ function dummy_derivative_graph!(structure::SystemStructure, var_eq_matching, ja
     # We can eliminate variables that are not a selected state (differential
     # variables). Selected states are differentiated variables that are not
     # dummy derivatives.
-    can_eliminate = let var_to_diff = var_to_diff, dummy_derivatives_set = dummy_derivatives_set
+    can_eliminate = let var_to_diff = var_to_diff,
+        dummy_derivatives_set = dummy_derivatives_set
 
         v -> begin
             dv = var_to_diff[v]

src/structural_transformation/symbolics_tearing.jl

@@ -115,7 +115,8 @@ function solve_equation(eq, var, simplify)
     var ~ rhs
 end
 
-function substitute_vars!(graph::BipartiteGraph, subs, cache=Int[], callback! = nothing; exclude = ())
+function substitute_vars!(graph::BipartiteGraph, subs, cache = Int[], callback! = nothing;
+                          exclude = ())
     for su in subs
         su === nothing && continue
         v, v′ = su
@@ -361,7 +362,6 @@ function tearing_reassemble(state::TearingState, var_eq_matching; simplify = fal
                 add_edge!(solvable_graph, eq_idx, dx_idx)
                 add_edge!(graph, eq_idx, x_t_idx)
                 add_edge!(graph, eq_idx, dx_idx)
-
             end
             # We use this info to substitute all `D(D(x))` or `D(x_t)` except
             # the `D(D(x)) ~ x_tt` equation to `x_tt`.
@@ -382,8 +382,10 @@ function tearing_reassemble(state::TearingState, var_eq_matching; simplify = fal
             # substituted to `x_tt`.
             for idx in (ogidx, dx_idx)
                 subidx = ((idx => x_t_idx),)
-                substitute_vars!(graph, subidx, idx_buffer, sub_callback!; exclude = order_lowering_eqs)
-                substitute_vars!(solvable_graph, subidx, idx_buffer; exclude = order_lowering_eqs)
+                substitute_vars!(graph, subidx, idx_buffer, sub_callback!;
+                                 exclude = order_lowering_eqs)
+                substitute_vars!(solvable_graph, subidx, idx_buffer;
+                                 exclude = order_lowering_eqs)
             end
         end
         empty!(subinfo)
@@ -453,7 +455,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching; simplify = fal
                 push!(removed_eqs, ieq)
                 push!(removed_vars, iv)
             else
-                rhs = -b/a
+                rhs = -b / a
                 neweq = var ~ simplify ? Symbolics.simplify(rhs) : rhs
                 push!(subeqs, neweq)
                 push!(solved_equations, ieq)
@@ -497,7 +499,8 @@ function tearing_reassemble(state::TearingState, var_eq_matching; simplify = fal
 
     # Update system
     solved_variables_set = BitSet(solved_variables)
-    active_vars = setdiff!(setdiff(BitSet(1:length(fullvars)), solved_variables_set), removed_vars)
+    active_vars = setdiff!(setdiff(BitSet(1:length(fullvars)), solved_variables_set),
+                           removed_vars)
     new_var_to_diff = complete(DiffGraph(length(active_vars)))
     idx = 0
     for (v, d) in enumerate(var_to_diff)

src/systems/alias_elimination.jl

@@ -19,7 +19,8 @@ function aag_bareiss(sys::AbstractSystem)
     return aag_bareiss!(state.structure.graph, complete(state.structure.var_to_diff), mm)
 end
 
-function extreme_var(var_to_diff, v, level = nothing, ::Val{descend} = Val(true); callback = _ -> nothing) where descend
+function extreme_var(var_to_diff, v, level = nothing, ::Val{descend} = Val(true);
+                     callback = _ -> nothing) where {descend}
     g = descend ? invview(var_to_diff) : var_to_diff
     callback(v)
     while (v′ = g[v]) !== nothing
@@ -32,17 +33,20 @@ function extreme_var(var_to_diff, v, level = nothing, ::Val{descend} = Val(true)
     level === nothing ? v : (v => level)
 end
 
-function walk_to_root!(iag, v::Integer, level = 0)
+function walk_to_root!(relative_level, iag, v::Integer, level = 0)
     brs = neighbors(iag, v)
     min_var_level = v => level
     for (x, lv′) in brs
         lv = lv′ + level
-        x, lv = walk_to_root!(iag, x, lv)
+        x, lv = walk_to_root!(relative_level, iag, x, lv)
+        relative_level[x] = lv
         if min_var_level[2] > lv
             min_var_level = x => lv
         end
     end
-    return extreme_var(iag.var_to_diff, min_var_level...)
+    x, lv = extreme_var(iag.var_to_diff, min_var_level...)
+    relative_level[x] = lv
+    return x => lv
 end
 
 function alias_elimination(sys)
@@ -87,22 +91,28 @@ function alias_elimination(sys)
     Main._a[] = ag, invag
     processed = falses(nvars)
     iag = InducedAliasGraph(ag, invag, var_to_diff, processed)
+    relative_level = BitDict(nvars)
     newag = AliasGraph(nvars)
     for (v, dv) in enumerate(var_to_diff)
         processed[v] && continue
         (dv === nothing && diff_to_var[v] === nothing) && continue
 
         # TODO: use an iterator, and get a relative level vector for `processed`
         # variabels.
-        r, lv = walk_to_root!(iag, v)
+        r, lv = walk_to_root!(relative_level, iag, v)
         fill!(processed, false)
         #lv = extreme_var(var_to_diff, v, -lv, Val(false))
-        lv′ = extreme_var(var_to_diff, v, 0, Val(false))[2]
+        lv′ = extreme_var(var_to_diff, v, 0)[2]
         let
             sv = fullvars[v]
             root = fullvars[r]
-            @warn "" sv => root level = lv levelv = lv′
+            @warn "" sv=>root level=lv levelv=lv′
+            for (v, rl) in pairs(relative_level)
+                @show v, rl
+                @show fullvars[v], rl - lv, rl, lv
+            end
         end
+        empty!(relative_level)
         level_to_var = Int[r]
         v′′::Union{Nothing, Int} = v′::Int = r
         while (v′′ = var_to_diff[v′]) !== nothing
@@ -113,7 +123,7 @@ function alias_elimination(sys)
         if nlevels < (new_nlevels = length(level_to_var))
             @assert !(D isa Nothing)
             for i in (nlevels + 1):new_nlevels
-                var_to_diff[level_to_var[i-1]] = level_to_var[i]
+                var_to_diff[level_to_var[i - 1]] = level_to_var[i]
                 fullvars[level_to_var[i]] = D(fullvars[level_to_var[i - 1]])
             end
         end
@@ -147,7 +157,7 @@ function alias_elimination(sys)
     end
 
     newstates = []
-   for j in eachindex(fullvars)
+    for j in eachindex(fullvars)
         if j in keys(ag)
             _, var = ag[j]
             iszero(var) && continue
@@ -316,7 +326,9 @@ struct InducedAliasGraph
     visited::BitVector
 end
 
-InducedAliasGraph(ag, invag, var_to_diff) = InducedAliasGraph(ag, invag, var_to_diff, falses(nv(invag)))
+function InducedAliasGraph(ag, invag, var_to_diff)
+    InducedAliasGraph(ag, invag, var_to_diff, falses(nv(invag)))
+end
 
 struct IAGNeighbors
     iag::InducedAliasGraph
@@ -540,7 +552,8 @@ function locally_structure_simplify!(adj_row, pivot_var, ag, var_to_diff)
                 # loop.
                 #
                 # We're relying on `var` being produced in sorted order here.
-                nirreducible += !(alias_candidate isa Pair) || alias_var != alias_candidate[2]
+                nirreducible += !(alias_candidate isa Pair) ||
+                                alias_var != alias_candidate[2]
                 alias_candidate = new_coeff => alias_var
             end
         end

src/systems/systemstructure.jl

@@ -147,7 +147,7 @@ struct DiffChainIterator{Descend}
     v::Int
 end
 
-function Base.iterate(di::DiffChainIterator{Descend}, v=nothing) where Descend
+function Base.iterate(di::DiffChainIterator{Descend}, v = nothing) where {Descend}
     if v === nothing
         vv = di.v
         return (vv, vv)

src/utils.jl

@@ -590,3 +590,42 @@ function promote_to_concrete(vs; tofloat = true, use_union = false)
         convert.(C, vs)
     end
 end
+
+struct BitDict <: AbstractDict{Int, Int}
+    keys::Vector{Int}
+    values::Vector{Union{Nothing, Int}}
+end
+BitDict(n::Integer) = BitDict(Int[], Union{Nothing, Int}[nothing for _ in 1:n])
+struct BitDictKeySet <: AbstractSet{Int}
+    d::BitDict
+end
+
+Base.keys(d::BitDict) = BitDictKeySet(d)
+Base.in(v::Integer, s::BitDictKeySet) = s.d.values[v] !== nothing
+Base.iterate(s::BitDictKeySet, state...) = iterate(s.d.keys, state...)
+function Base.setindex!(d::BitDict, val::Integer, ind::Integer)
+    if 1 <= ind <= length(d.values) && d.values[ind] === nothing
+        push!(d.keys, ind)
+    end
+    d.values[ind] = val
+end
+function Base.getindex(d::BitDict, ind::Integer)
+    if 1 <= ind <= length(d.values) && d.values[ind] === nothing
+        return d.values[ind]
+    else
+        throw(KeyError(ind))
+    end
+end
+function Base.iterate(d::BitDict, state...)
+    r = Base.iterate(d.keys, state...)
+    r === nothing && return nothing
+    k, state = r
+    (k => d.values[k]), state
+end
+function Base.empty!(d::BitDict)
+    for v in d.keys
+        d.values[v] = nothing
+    end
+    empty!(d.keys)
+    d
+end