Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
79 changes: 77 additions & 2 deletions src/coloring.jl
Original file line number Diff line number Diff line change
Expand Up @@ -328,7 +328,7 @@ function acyclic_coloring(g::Graph, order::AbstractOrder)
end
end
end
return color, TreeSet(forest)
return color, TreeSet(forest, nvertices)
end

function _prevent_cycle!(
Expand Down Expand Up @@ -408,5 +408,80 @@ $TYPEDFIELDS
"""
struct TreeSet
"a forest of two-colored trees"
forest::DisjointSets{Tuple{Int,Int}}
trees::Vector{Vector{Tuple{Int,Int}}}
"???"
nodes::Vector{Vector{Int}}
"???"
stored_values::Vector{Float64}
end

function TreeSet(forest::DisjointSets, nvertices::Int)
# forest is a structure DisjointSets from DataStructures.jl
# - forest.intmap: a dictionary that maps an edge (i, j) to an integer k
# - forest.revmap: a dictionary that does the reverse of intmap, mapping an integer k to an edge (i, j)
# - forest.internal.ngroups: the number of trees in the forest
ntrees = forest.internal.ngroups

# vector of trees where each tree contains the indices of its edges
trees = [Int[] for i in 1:ntrees]

# dictionary that maps a tree's root to the index of the tree
roots = Dict{Int,Int}()

k = 0
for edge in forest.revmap
root_edge = find_root!(forest, edge)
root = forest.intmap[root_edge]
if !haskey(roots, root)
k += 1
roots[root] = k
end
index_tree = roots[root]
push!(trees[index_tree], forest.intmap[edge])
end

# vector of dictionaries where each dictionary stores the degree of each vertex in a tree
degrees = [Dict{Int,Int}() for k in 1:ntrees]
for k in 1:ntrees
tree = trees[k]
degree = degrees[k]
for edge_index in tree
i, j = forest.revmap[edge_index]
!haskey(degree, i) && (degree[i] = 0)
!haskey(degree, j) && (degree[j] = 0)
degree[i] += 1
degree[j] += 1
end
end

# depth-first search (DFS) traversal order for each tree in the forest
dfs_orders = [Vector{Tuple{Int,Int}}() for k in 1:ntrees]
for k in 1:ntrees
tree = trees[k]
degree = degrees[k]
while sum(values(degree)) != 0
for (t, edge_index) in enumerate(tree)
if edge_index != 0
i, j = forest.revmap[edge_index]
if (degree[i] == 1) || (degree[j] == 1) # leaf vertex
if degree[i] > degree[j] # vertex i is the parent of vertex j
i, j = j, i # ensure that i always denotes a leaf vertex
end
degree[i] -= 1 # decrease the degree of vertex i
degree[j] -= 1 # decrease the degree of vertex j
tree[t] = 0 # remove the edge (i,j)
push!(dfs_orders[k], (i, j))
end
end
end
end
end

# stored_values holds the sum of edge values for subtrees in a tree.
# For each vertex i, stored_values[i] is the sum of edge values in the subtree rooted at i.
stored_values = Vector{Float64}(undef, nvertices)

nodes = [[vertex for vertex in keys(degrees[k])] for k = 1:ntrees]

return TreeSet(dfs_orders, nodes, stored_values)
end
72 changes: 4 additions & 68 deletions src/decompression.jl
Original file line number Diff line number Diff line change
Expand Up @@ -303,72 +303,8 @@ function decompress_aux!(
A .= zero(R)
S = get_matrix(result)
color = column_colors(result)

# forest is a structure DisjointSets from DataStructures.jl
# - forest.intmap: a dictionary that maps an edge (i, j) to an integer k
# - forest.revmap: a dictionary that does the reverse of intmap, mapping an integer k to an edge (i, j)
# - forest.internal.ngroups: the number of trees in the forest
forest = result.tree_set.forest
ntrees = forest.internal.ngroups

# vector of trees where each tree contains the indices of its edges
trees = [Int[] for i in 1:ntrees]

# dictionary that maps a tree's root to the index of the tree
roots = Dict{Int,Int}()

k = 0
for edge in forest.revmap
root_edge = find_root!(forest, edge)
root = forest.intmap[root_edge]
if !haskey(roots, root)
k += 1
roots[root] = k
end
index_tree = roots[root]
push!(trees[index_tree], forest.intmap[edge])
end

# vector of dictionaries where each dictionary stores the degree of each vertex in a tree
degrees = [Dict{Int,Int}() for k in 1:ntrees]
for k in 1:ntrees
tree = trees[k]
degree = degrees[k]
for edge_index in tree
i, j = forest.revmap[edge_index]
!haskey(degree, i) && (degree[i] = 0)
!haskey(degree, j) && (degree[j] = 0)
degree[i] += 1
degree[j] += 1
end
end

# depth-first search (DFS) traversal order for each tree in the forest
dfs_orders = [Vector{Tuple{Int,Int}}() for k in 1:ntrees]
for k in 1:ntrees
tree = trees[k]
degree = degrees[k]
while sum(values(degree)) != 0
for (t, edge_index) in enumerate(tree)
if edge_index != 0
i, j = forest.revmap[edge_index]
if (degree[i] == 1) || (degree[j] == 1) # leaf vertex
if degree[i] > degree[j] # vertex i is the parent of vertex j
i, j = j, i # ensure that i always denotes a leaf vertex
end
degree[i] -= 1 # decrease the degree of vertex i
degree[j] -= 1 # decrease the degree of vertex j
tree[t] = 0 # remove the edge (i,j)
push!(dfs_orders[k], (i, j))
end
end
end
end
end

# stored_values holds the sum of edge values for subtrees in a tree.
# For each vertex i, stored_values[i] is the sum of edge values in the subtree rooted at i.
stored_values = Vector{R}(undef, n)
@compat (; trees, nodes, stored_values) = result.tree_set
ntrees = length(trees)

# Recover the diagonal coefficients of A
for i in axes(A, 1)
Expand All @@ -379,12 +315,12 @@ function decompress_aux!(

# Recover the off-diagonal coefficients of A
for k in 1:ntrees
vertices = keys(degrees[k])
vertices = nodes[k]
for vertex in vertices
stored_values[vertex] = zero(R)
end

tree = dfs_orders[k]
tree = trees[k]
for (i, j) in tree
val = B[i, color[j]] - stored_values[i]
stored_values[j] = stored_values[j] + val
Expand Down