Make decompression more efficient (#71)

* Make decompression more efficient

* Add skipped test for decompression into dense matrix

Author: gdalle

docs/src/dev.md

@@ -34,10 +34,10 @@ SparseMatrixColorings.TreeSet
 ## Concrete coloring results
 
 ```@docs
-SparseMatrixColorings.DefaultColoringResult
+SparseMatrixColorings.NonSymmetricColoringResult
 SparseMatrixColorings.StarSetColoringResult
 SparseMatrixColorings.TreeSetColoringResult
-SparseMatrixColorings.DirectSparseColoringResult
+SparseMatrixColorings.LinearSystemColoringResult
 ```
 
 ## Testing

src/SparseMatrixColorings.jl

@@ -20,7 +20,9 @@ using LinearAlgebra:
     Transpose,
     adjoint,
     checksquare,
+    factorize,
     issymmetric,
+    ldiv!,
     parent,
     transpose
 using Random: AbstractRNG, default_rng, randperm
@@ -45,7 +47,6 @@ include("matrices.jl")
 include("interface.jl")
 include("decompression.jl")
 include("check.jl")
-include("sparsematrixcsc.jl")
 include("examples.jl")
 
 export ColoringProblem, GreedyColoringAlgorithm, AbstractColoringResult

src/decompression.jl

@@ -7,7 +7,7 @@ Compress `A` given a coloring `result` of the sparsity pattern of `A`.
 - If `result` comes from a `:bidirectional` partition, the output is a tuple of matrices `(Br, Bc)`, where `Br` is compressed by row and `Bc` by column.
 
 Compression means summing either the columns or the rows of `A` which share the same color.
-It is undone by calling [`decompress`](@ref).
+It is undone by calling [`decompress`](@ref) or [`decompress!`](@ref).
 
 !!! warning
     At the moment, `:bidirectional` partitions are not implemented.
@@ -67,6 +67,7 @@ end
     decompress(B::AbstractMatrix, result::AbstractColoringResult)
 
 Decompress `B` into a new matrix `A`, given a coloring `result` of the sparsity pattern of `A`.
+The in-place alternative is [`decompress!`](@ref).
 
 Compression means summing either the columns or the rows of `A` which share the same color.
 It is done by calling [`compress`](@ref).
@@ -127,10 +128,14 @@ end
     )
 
 Decompress `B` in-place into `A`, given a coloring `result` of the sparsity pattern of `A`.
+The out-of-place alternative is [`decompress`](@ref).
 
 Compression means summing either the columns or the rows of `A` which share the same color.
 It is done by calling [`compress`](@ref).
 
+!!! note
+    In-place decompression is faster when `A isa SparseMatrixCSC`.
+
 # Example
 
 ```jldoctest
@@ -190,10 +195,10 @@ function decompress!(
     return decompress_aux!(A, B, result)
 end
 
+## NonSymmetricColoringResult
+
 function decompress_aux!(
-    A::AbstractMatrix{R},
-    B::AbstractMatrix{R},
-    result::AbstractColoringResult{:nonsymmetric,:column,:direct},
+    A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::NonSymmetricColoringResult{:column}
 ) where {R<:Real}
     A .= zero(R)
     S = get_matrix(result)
@@ -209,9 +214,7 @@ function decompress_aux!(
 end
 
 function decompress_aux!(
-    A::AbstractMatrix{R},
-    B::AbstractMatrix{R},
-    result::AbstractColoringResult{:nonsymmetric,:row,:direct},
+    A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::NonSymmetricColoringResult{:row}
 ) where {R<:Real}
     A .= zero(R)
     S = get_matrix(result)
@@ -227,24 +230,31 @@ function decompress_aux!(
 end
 
 function decompress_aux!(
-    A::AbstractMatrix{R},
-    B::AbstractMatrix{R},
-    result::AbstractColoringResult{:symmetric,:column,:direct},
+    A::SparseMatrixCSC{R}, B::AbstractMatrix{R}, result::NonSymmetricColoringResult{:column}
 ) where {R<:Real}
-    A .= zero(R)
-    S = get_matrix(result)
-    color = column_colors(result)
-    group = column_groups(result)
-    for ij in findall(!iszero, S)
-        i, j = Tuple(ij)
-        k, l = symmetric_coefficient(i, j, color, group, S)
-        A[i, j] = B[k, l]
+    nzA = nonzeros(A)
+    ind = result.compressed_indices
+    for i in eachindex(nzA, ind)
+        nzA[i] = B[ind[i]]
+    end
+    return A
+end
+
+function decompress_aux!(
+    A::SparseMatrixCSC{R}, B::AbstractMatrix{R}, result::NonSymmetricColoringResult{:row}
+) where {R<:Real}
+    nzA = nonzeros(A)
+    ind = result.compressed_indices
+    for i in eachindex(nzA, ind)
+        nzA[i] = B[ind[i]]
     end
     return A
 end
 
+## StarSetColoringResult
+
 function decompress_aux!(
-    A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::StarSetColoringResult{:column}
+    A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::StarSetColoringResult
 ) where {R<:Real}
     A .= zero(R)
     S = get_matrix(result)
@@ -258,117 +268,35 @@ function decompress_aux!(
 end
 
 function decompress_aux!(
-    A::AbstractMatrix{R},
-    B::AbstractMatrix{R},
-    result::AbstractColoringResult{:symmetric,:column,:substitution},
+    A::SparseMatrixCSC{R}, B::AbstractMatrix{R}, result::StarSetColoringResult
 ) where {R<:Real}
-    # build T such that T * strict_upper_nonzeros(A) = B
-    # and solve a linear least-squares problem
-    # only consider the strict upper triangle of A because of symmetry
-    # TODO: make more efficient
-    A .= zero(R)
-    S = sparse(get_matrix(result))
-    color = column_colors(result)
-
-    n = checksquare(S)
-    strict_upper_nonzero_inds = Tuple{Int,Int}[]
-    I, J, _ = findnz(S)
-    for (i, j) in zip(I, J)
-        (i < j) && push!(strict_upper_nonzero_inds, (i, j))
-        (i == j) && (A[i, i] = B[i, color[i]])
-    end
-
-    T = spzeros(float(R), length(B), length(strict_upper_nonzero_inds))
-    for (l, (i, j)) in enumerate(strict_upper_nonzero_inds)
-        ci = color[i]
-        cj = color[j]
-        ki = (ci - 1) * n + j  # A[i, j] appears in B[j, ci]
-        kj = (cj - 1) * n + i  # A[i, j] appears in B[i, cj]
-        T[ki, l] = one(float(R))
-        T[kj, l] = one(float(R))
-    end
-
-    strict_upper_nonzeros_A = T \ vec(B)
-    for (l, (i, j)) in enumerate(strict_upper_nonzero_inds)
-        A[i, j] = strict_upper_nonzeros_A[l]
-        A[j, i] = strict_upper_nonzeros_A[l]
+    nzA = nonzeros(A)
+    ind = result.compressed_indices
+    for i in eachindex(nzA, ind)
+        nzA[i] = B[ind[i]]
     end
     return A
 end
 
+## TreeSetColoringResult
+
+# TODO: add method for A::SparseMatrixCSC
+
 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
+    @compat (; degrees, dfs_orders, stored_values) = result
 
     # 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)
+    stored_values_right_type = if R == Float64
+        stored_values
+    else
+        similar(stored_values, R)
+    end
 
     # Recover the diagonal coefficients of A
     for i in axes(A, 1)
@@ -378,19 +306,43 @@ function decompress_aux!(
     end
 
     # Recover the off-diagonal coefficients of A
-    for k in 1:ntrees
+    for k in eachindex(degrees, dfs_orders)
         vertices = keys(degrees[k])
         for vertex in vertices
-            stored_values[vertex] = zero(R)
+            stored_values_right_type[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
+            val = B[i, color[j]] - stored_values_right_type[i]
+            stored_values_right_type[j] = stored_values_right_type[j] + val
             A[i, j] = val
             A[j, i] = val
         end
     end
     return A
 end
+
+## MatrixInverseColoringResult
+
+function decompress_aux!(
+    A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::LinearSystemColoringResult
+) where {R<:Real}
+    S = get_matrix(result)
+    color = column_colors(result)
+    @compat (; strict_upper_nonzero_inds, T_factorization, strict_upper_nonzeros_A) = result
+
+    # TODO: for some reason I cannot use ldiv! with a sparse QR
+    strict_upper_nonzeros_A = T_factorization \ vec(B)
+    A .= zero(R)
+    for i in axes(A, 1)
+        if !iszero(S[i, i])
+            A[i, i] = B[i, color[i]]
+        end
+    end
+    for (l, (i, j)) in enumerate(strict_upper_nonzero_inds)
+        A[i, j] = strict_upper_nonzeros_A[l]
+        A[j, i] = strict_upper_nonzeros_A[l]
+    end
+    return A
+end

src/graph.jl

@@ -16,7 +16,6 @@ struct Graph{T<:Integer}
 end
 
 Graph(S::SparseMatrixCSC) = Graph(S.colptr, S.rowval)
-Graph(S::AbstractMatrix) = Graph(sparse(S))
 
 Base.length(g::Graph) = length(g.colptr) - 1
 SparseArrays.nnz(g::Graph) = length(g.rowval)
@@ -91,7 +90,6 @@ The adjacency graph of a symmetrix matric `A ∈ ℝ^{n × n}` is `G(A) = (V, E)
 > [_What Color Is Your Jacobian? Graph Coloring for Computing Derivatives_](https://epubs.siam.org/doi/10.1137/S0036144504444711), Gebremedhin et al. (2005)
 """
 adjacency_graph(H::SparseMatrixCSC) = Graph(H - Diagonal(H))
-adjacency_graph(H::AbstractMatrix) = adjacency_graph(sparse(H))
 
 """
     bipartite_graph(J::AbstractMatrix)
@@ -109,9 +107,7 @@ The bipartite graph of a matrix `A ∈ ℝ^{m × n}` is `Gb(A) = (V₁, V₂, E)
 > [_What Color Is Your Jacobian? Graph Coloring for Computing Derivatives_](https://epubs.siam.org/doi/10.1137/S0036144504444711), Gebremedhin et al. (2005)
 """
 function bipartite_graph(J::SparseMatrixCSC)
-    g1 = Graph(transpose(J))  # rows to columns
+    g1 = Graph(sparse(transpose(J)))  # rows to columns
     g2 = Graph(J)  # columns to rows
     return BipartiteGraph(g1, g2)
 end
-
-bipartite_graph(J::AbstractMatrix) = bipartite_graph(sparse(J))