Implement TreeSet-based acyclic decompression

src/decompression.jl

@@ -295,3 +295,102 @@ function decompress_aux!(
     end
     return A
 end
+
+function decompress_aux!(
+    A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::TreeSetColoringResult
+) where {R<:Real}
+    n = checksquare(A)
+    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)
+
+    # Recover the diagonal coefficients of A
+    for i in axes(A, 1)
+        if !iszero(S[i, i])
+            A[i, i] = B[i, color[i]]
+        end
+    end
+
+    # Recover the off-diagonal coefficients of A
+    for k in 1:ntrees
+        vertices = keys(degrees[k])
+        for vertex in vertices
+            stored_values[vertex] = zero(R)
+        end
+
+        tree = dfs_orders[k]
+        for (i, j) in tree
+            val = B[i, color[j]] - stored_values[i]
+            stored_values[j] = stored_values[j] + val
+            A[i, j] = val
+            A[j, i] = val
+        end
+    end
+    return A
+end

test/utils.jl

@@ -22,10 +22,12 @@ function test_coloring_decompression(
         B = compress(A, result)
         !isnothing(color0) && @test color == color0
         !isnothing(B0) && @test B == B0
-        @test decompress(B, result) ≈ A0
         @test decompress(B, default_result) ≈ A0
-        @test decompress!(respectful_similar(A), B, result) ≈ A0
+        @test decompress(B, result) ≈ A0
+        @test decompress(B, result) ≈ A0  # check result wasn't modified
         @test decompress!(respectful_similar(A), B, default_result) ≈ A0
+        @test decompress!(respectful_similar(A), B, result) ≈ A0
+        @test decompress!(respectful_similar(A), B, result) ≈ A0
     end
     @test all(color_vec .== Ref(color_vec[1]))
 end