From 67398491447fb4243f615b1ccd08d5e888551040 Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Wed, 2 Apr 2025 18:55:47 -0500 Subject: [PATCH 1/2] Improve the computation of TreeSet --- src/coloring.jl | 73 ++++++++++++++++++++++++++++---------------- src/decompression.jl | 39 +++++++++++++++++------ src/result.jl | 18 ++++++++--- 3 files changed, 90 insertions(+), 40 deletions(-) diff --git a/src/coloring.jl b/src/coloring.jl index bfecc0d0..a04f17ea 100644 --- a/src/coloring.jl +++ b/src/coloring.jl @@ -291,7 +291,9 @@ function acyclic_coloring( end end - tree_set = TreeSet(g, forest) + buffer = forbidden_colors + reverse_bfs_orders = first_visit_to_tree + tree_set = TreeSet(g, forest, buffer, reverse_bfs_orders) if postprocessing # Reuse the vector forbidden_colors to compute offsets during post-processing offsets = forbidden_colors @@ -372,19 +374,29 @@ Encode a set of 2-colored trees resulting from the [`acyclic_coloring`](@ref) al $TYPEDFIELDS """ struct TreeSet{T} - reverse_bfs_orders::Vector{Vector{Tuple{T,T}}} + reverse_bfs_orders::Vector{Tuple{T,T}} is_star::Vector{Bool} + num_edges_per_tree::Vector{T} end -function TreeSet(g::AdjacencyGraph{T}, forest::Forest{T}) where {T} +function TreeSet( + g::AdjacencyGraph{T}, + forest::Forest{T}, + buffer::AbstractVector{T}, + reverse_bfs_orders::Vector{Tuple{T,T}}, +) where {T} S = pattern(g) edge_to_index = edge_indices(g) nv = nb_vertices(g) nt = forest.num_trees - # dictionary that maps a tree's root to the index of the tree - roots = Dict{T,T}() - sizehint!(roots, nt) + # root_to_tree is a vector that maps a tree's root to the index of the tree + # We can recycle forest.ranks because we don't need anymore to merge trees + root_to_tree = forest.ranks + fill!(root_to_tree, zero(T)) + + # Contains the number of edges per tree + num_edges_per_tree = zeros(T, nt) # vector of dictionaries where each dictionary stores the neighbors of each vertex in a tree trees = [Dict{T,Vector{T}}() for i in 1:nt] @@ -401,13 +413,14 @@ function TreeSet(g::AdjacencyGraph{T}, forest::Forest{T}) where {T} root = find_root!(forest, index_ij) # Update roots - if !haskey(roots, root) + if iszero(root_to_tree[root]) nr += 1 - roots[root] = nr + root_to_tree[root] = nr end # index of the tree T that contains this edge - index_tree = roots[root] + index_tree = root_to_tree[root] + num_edges_per_tree[index_tree] += 1 # Update the neighbors of i in the tree T if !haskey(trees[index_tree], i) @@ -427,10 +440,7 @@ function TreeSet(g::AdjacencyGraph{T}, forest::Forest{T}) where {T} end # degrees is a vector of integers that stores the degree of each vertex in a tree - degrees = Vector{T}(undef, nv) - - # reverse breadth first (BFS) traversal order for each tree in the forest - reverse_bfs_orders = [Tuple{T,T}[] for i in 1:nt] + degrees = buffer # nvmax is the number of vertices of the biggest tree in the forest nvmax = 0 @@ -446,6 +456,10 @@ function TreeSet(g::AdjacencyGraph{T}, forest::Forest{T}) where {T} # meaning that one vertex is directly connected to all other vertices in the tree is_star = Vector{Bool}(undef, nt) + # Number of edges treated + num_edges_treated = zero(T) + + # reverse_bfs_orders contains the reverse breadth first (BFS) traversal order for each tree in the forest for k in 1:nt tree = trees[k] @@ -483,7 +497,8 @@ function TreeSet(g::AdjacencyGraph{T}, forest::Forest{T}) where {T} # Check if neighbor is the parent of the leaf or if it was a child before the tree was pruned if degrees[neighbor] != 0 # (leaf, neighbor) represents the next edge to visit during decompression - push!(reverse_bfs_orders[k], (leaf, neighbor)) + num_edges_treated += 1 + reverse_bfs_orders[num_edges_treated] = (leaf, neighbor) if bool_star # Initialize the potential hub of the star with the first parent of a leaf @@ -514,7 +529,7 @@ function TreeSet(g::AdjacencyGraph{T}, forest::Forest{T}) where {T} is_star[k] = bool_star end - return TreeSet(reverse_bfs_orders, is_star) + return TreeSet(reverse_bfs_orders, is_star, num_edges_per_tree) end ## Postprocessing, mirrors decompression code @@ -582,46 +597,52 @@ function postprocess!( end else # only the colors of non-leaf vertices are used - (; reverse_bfs_orders, is_star) = star_or_tree_set + (; reverse_bfs_orders, is_star, num_edges_per_tree) = star_or_tree_set nb_trivial_trees = 0 + # Index of the first edge in reverse_bfs_orders for the current tree + first = 1 + # Iterate through all non-trivial trees - for k in eachindex(reverse_bfs_orders) - reverse_bfs_order = reverse_bfs_orders[k] + for k in eachindex(num_edges_per_tree) + ne_tree = num_edges_per_tree[k] # Check if we have more than one edge in the tree (non-trivial tree) - if length(reverse_bfs_order) > 1 + if ne_tree > 1 # Determine if the tree is a star if is_star[k] # It is a non-trivial star and only the color of the hub is needed - (_, hub) = reverse_bfs_order[1] + (_, hub) = reverse_bfs_orders[first] color_used[color[hub]] = true else # It is not a star and both colors are needed during the decompression - (i, j) = reverse_bfs_order[1] + (i, j) = reverse_bfs_orders[first] color_used[color[i]] = true color_used[color[j]] = true end else nb_trivial_trees += 1 end + first += ne_tree end # Process the trivial trees (if any) if nb_trivial_trees > 0 - for k in eachindex(reverse_bfs_orders) - reverse_bfs_order = reverse_bfs_orders[k] + first = 1 + for k in eachindex(num_edges_per_tree) + ne_tree = num_edges_per_tree[k] # Check if we have exactly one edge in the tree - if length(reverse_bfs_order) == 1 - (i, j) = reverse_bfs_order[1] + if ne_tree == 1 + (i, j) = reverse_bfs_orders[first] if color_used[color[i]] # Make i the root to avoid possibly adding one more used color # Switch it with the (only) leaf - reverse_bfs_order[1] = (j, i) + reverse_bfs_orders[first] = (j, i) else # Keep j as the root color_used[color[j]] = true end end + first += ne_tree end end end diff --git a/src/decompression.jl b/src/decompression.jl index 64e3c517..8bebcb22 100644 --- a/src/decompression.jl +++ b/src/decompression.jl @@ -517,7 +517,7 @@ end function decompress!( A::AbstractMatrix, B::AbstractMatrix, result::TreeSetColoringResult, uplo::Symbol=:F ) - (; ag, color, reverse_bfs_orders, buffer) = result + (; ag, color, reverse_bfs_orders, num_edges_per_tree, buffer) = result (; S) = ag uplo == :F && check_same_pattern(A, S) R = eltype(A) @@ -531,24 +531,32 @@ function decompress!( # Recover the diagonal coefficients of A if has_diagonal(ag) - for i in axes(A, 1) + for i in axes(S, 1) if !iszero(S[i, i]) A[i, i] = B[i, color[i]] end end end + # Index of the first edge in reverse_bfs_orders for the current tree + first = 1 + # Recover the off-diagonal coefficients of A - for k in eachindex(reverse_bfs_orders) + for k in eachindex(num_edges_per_tree) + ne_tree = num_edges_per_tree[k] + last = first + ne_tree - 1 + # Reset the buffer to zero for all vertices in a tree (except the root) - for (vertex, _) in reverse_bfs_orders[k] + for pos in first:last + (vertex, _) = reverse_bfs_orders[pos] buffer_right_type[vertex] = zero(R) end # Reset the buffer to zero for the root vertex - (_, root) = reverse_bfs_orders[k][end] + (_, root) = reverse_bfs_orders[last] buffer_right_type[root] = zero(R) - for (i, j) in reverse_bfs_orders[k] + for pos in first:last + (i, j) = reverse_bfs_orders[pos] val = B[i, color[j]] - buffer_right_type[i] buffer_right_type[j] = buffer_right_type[j] + val @@ -559,6 +567,7 @@ function decompress!( A[j, i] = val end end + first += ne_tree end return A end @@ -573,6 +582,7 @@ function decompress!( ag, color, reverse_bfs_orders, + num_edges_per_tree, diagonal_indices, diagonal_nzind, lower_triangle_offsets, @@ -612,20 +622,28 @@ function decompress!( end end + # Index of the first edge in reverse_bfs_orders for the current tree + first = 1 + # Index of offsets in lower_triangle_offsets and upper_triangle_offsets counter = 0 # Recover the off-diagonal coefficients of A - for k in eachindex(reverse_bfs_orders) + for k in eachindex(num_edges_per_tree) + ne_tree = num_edges_per_tree[k] + last = first + ne_tree - 1 + # Reset the buffer to zero for all vertices in a tree (except the root) - for (vertex, _) in reverse_bfs_orders[k] + for pos in first:last + (vertex, _) = reverse_bfs_orders[pos] buffer_right_type[vertex] = zero(R) end # Reset the buffer to zero for the root vertex - (_, root) = reverse_bfs_orders[k][end] + (_, root) = reverse_bfs_orders[last] buffer_right_type[root] = zero(R) - for (i, j) in reverse_bfs_orders[k] + for pos in first:last + (i, j) = reverse_bfs_orders[pos] counter += 1 val = B[i, color[j]] - buffer_right_type[i] buffer_right_type[j] = buffer_right_type[j] + val @@ -665,6 +683,7 @@ function decompress!( end #! format: on end + first += ne_tree end return A end diff --git a/src/result.jl b/src/result.jl index b1ec9409..1206167e 100644 --- a/src/result.jl +++ b/src/result.jl @@ -314,7 +314,8 @@ struct TreeSetColoringResult{ ag::G color::Vector{T} group::GT - reverse_bfs_orders::Vector{Vector{Tuple{T,T}}} + reverse_bfs_orders::Vector{Tuple{T,T}} + num_edges_per_tree::Vector{T} diagonal_indices::Vector{T} diagonal_nzind::Vector{T} lower_triangle_offsets::Vector{T} @@ -329,7 +330,7 @@ function TreeSetColoringResult( tree_set::TreeSet{<:Integer}, decompression_eltype::Type{R}, ) where {T<:Integer,R} - (; reverse_bfs_orders) = tree_set + (; reverse_bfs_orders, num_edges_per_tree) = tree_set (; S) = ag nvertices = length(color) group = group_by_color(T, color) @@ -358,11 +359,18 @@ function TreeSetColoringResult( lower_triangle_offsets = Vector{T}(undef, nedges) upper_triangle_offsets = Vector{T}(undef, nedges) + # Index of the first edge in reverse_bfs_orders for the current tree + first = 1 + # Index in lower_triangle_offsets and upper_triangle_offsets index_offsets = 0 - for k in eachindex(reverse_bfs_orders) - for (leaf, neighbor) in reverse_bfs_orders[k] + for k in eachindex(num_edges_per_tree) + ne_tree = num_edges_per_tree[k] + last = first + ne_tree - 1 + + for pos in first:last + (leaf, neighbor) = reverse_bfs_orders[pos] # Update lower_triangle_offsets and upper_triangle_offsets i = leaf j = neighbor @@ -393,6 +401,7 @@ function TreeSetColoringResult( end #! format: on end + first += ne_tree end # buffer holds the sum of edge values for subtrees in a tree. @@ -405,6 +414,7 @@ function TreeSetColoringResult( color, group, reverse_bfs_orders, + num_edges_per_tree, diagonal_indices, diagonal_nzind, lower_triangle_offsets, From 408ed3b0737d33e1d6e4028aa6df7fd29307b59f Mon Sep 17 00:00:00 2001 From: Alexis Montoison Date: Wed, 2 Apr 2025 19:49:57 -0500 Subject: [PATCH 2/2] Update a comment --- src/coloring.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/coloring.jl b/src/coloring.jl index a04f17ea..5f14ece7 100644 --- a/src/coloring.jl +++ b/src/coloring.jl @@ -391,7 +391,7 @@ function TreeSet( nt = forest.num_trees # root_to_tree is a vector that maps a tree's root to the index of the tree - # We can recycle forest.ranks because we don't need anymore to merge trees + # We can recycle forest.ranks because we don't need it anymore to merge trees root_to_tree = forest.ranks fill!(root_to_tree, zero(T))