Add an acyclic coloring (#36)

* Add an acyclic coloring

* Group by color

* Typo

* Noref

* Fix test

* Working indirect decompression

* Large random tests pass

---------

Co-authored-by: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>

Author: amontoison

Project.toml

@@ -6,6 +6,7 @@ version = "0.4.0"
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
+DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -15,6 +16,7 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 ADTypes = "1.2.1"
 Compat = "3.46,4.2"
 DocStringExtensions = "0.8,0.9"
+DataStructures = "0.18"
 LinearAlgebra = "<0.0.1, 1"
 Random = "<0.0.1, 1"
 SparseArrays = "<0.0.1, 1"

docs/src/api.md

@@ -66,9 +66,11 @@ SparseMatrixColorings.bipartite_graph
 SparseMatrixColorings.partial_distance2_coloring
 SparseMatrixColorings.symmetric_coefficient
 SparseMatrixColorings.star_coloring
-SparseMatrixColorings.StarSet
+SparseMatrixColorings.acyclic_coloring
 SparseMatrixColorings.group_by_color
 SparseMatrixColorings.get_matrix
+SparseMatrixColorings.StarSet
+SparseMatrixColorings.TreeSet
 ```
 
 ### Concrete coloring results
@@ -99,5 +101,7 @@ SparseMatrixColorings.matrix_versions
 ```@docs
 SparseMatrixColorings.Example
 SparseMatrixColorings.what_fig_41
+SparseMatrixColorings.what_fig_61
 SparseMatrixColorings.efficient_fig_1
+SparseMatrixColorings.efficient_fig_4
 ```

src/SparseMatrixColorings.jl

@@ -9,6 +9,7 @@ $EXPORTS
 """
 module SparseMatrixColorings
 
+using DataStructures: DisjointSets, find_root!, root_union!, num_groups
 using ADTypes: ADTypes
 using Compat: @compat, stack
 using DocStringExtensions: README, EXPORTS, SIGNATURES, TYPEDEF, TYPEDFIELDS

src/coloring.jl

@@ -76,10 +76,11 @@ The vertices are colored in a greedy fashion, following the `order` supplied.
 """
 function star_coloring(g::Graph, order::AbstractOrder)
     # Initialize data structures
-    color = zeros(Int, length(g))
-    forbidden_colors = zeros(Int, length(g))
-    first_neighbor = fill((0, 0), length(g))  # at first no neighbors have been encountered
-    treated = zeros(Int, length(g))
+    nvertices = length(g)
+    color = zeros(Int, nvertices)
+    forbidden_colors = zeros(Int, nvertices)
+    first_neighbor = fill((0, 0), nvertices)  # at first no neighbors have been encountered
+    treated = zeros(Int, nvertices)
     star = Dict{Tuple{Int,Int},Int}()
     hub = Int[]
     vertices_in_order = vertices(g, order)
@@ -119,17 +120,13 @@ function star_coloring(g::Graph, order::AbstractOrder)
 end
 
 """
-$TYPEDEF
+    StarSet
 
-Encode a set of 2-colored stars resulting from the star coloring algorithm.
+Encode a set of 2-colored stars resulting from the [`star_coloring`](@ref) algorithm.
 
 # Fields
 
 $TYPEDFIELDS
-
-# References
-
-> [_New Acyclic and Star Coloring Algorithms with Application to Computing Hessians_](https://epubs.siam.org/doi/abs/10.1137/050639879), Gebremedhin et al. (2007), Algorithm 4.1
 """
 struct StarSet
     "a mapping from edges (pair of vertices) their to star index"
@@ -262,3 +259,159 @@ function symmetric_coefficient(
         return j, color[h]
     end
 end
+
+"""
+    acyclic_coloring(g::Graph, order::AbstractOrder)
+
+Compute an acyclic coloring of all vertices in the adjacency graph `g` and return a vector of integer colors.
+
+An _acyclic coloring_ is a distance-1 coloring with the further restriction that every cycle uses at least 3 colors.
+
+The vertices are colored in a greedy fashion, following the `order` supplied.
+
+# See also
+
+- [`Graph`](@ref)
+- [`AbstractOrder`](@ref)
+
+# References
+
+> [_New Acyclic and Star Coloring Algorithms with Application to Computing Hessians_](https://epubs.siam.org/doi/abs/10.1137/050639879), Gebremedhin et al. (2007), Algorithm 3.1
+"""
+function acyclic_coloring(g::Graph, order::AbstractOrder)
+    # Initialize data structures
+    nvertices = length(g)
+    color = zeros(Int, nvertices)
+    forbidden_colors = zeros(Int, nvertices)
+    first_neighbor = fill((0, 0), nvertices)  # at first no neighbors have been encountered
+    first_visit_to_tree = fill((0, 0), nnz(g))  # Could we use less than nnz(g) values?
+    disjoint_sets = DisjointSets{Tuple{Int,Int}}()
+    vertices_in_order = vertices(g, order)
+    parent = Int[]  # Could be a vector of Bool
+
+    for v in vertices_in_order
+        for w in neighbors(g, v)
+            iszero(color[w]) && continue
+            forbidden_colors[color[w]] = v
+        end
+        for w in neighbors(g, v)
+            iszero(color[w]) && continue
+            for x in neighbors(g, w)
+                iszero(color[x]) && continue
+                if forbidden_colors[color[x]] != v
+                    _prevent_cycle!(
+                        v, w, x, color, first_visit_to_tree, forbidden_colors, disjoint_sets
+                    )
+                end
+            end
+        end
+        for i in eachindex(forbidden_colors)
+            if forbidden_colors[i] != v
+                color[v] = i
+                break
+            end
+        end
+        for w in neighbors(g, v)  # grow two-colored stars around the vertex v
+            iszero(color[w]) && continue
+            _grow_star!(v, w, color, first_neighbor, disjoint_sets, parent)
+        end
+        for w in neighbors(g, v)
+            iszero(color[w]) && continue
+            for x in neighbors(g, w)
+                (x == v || iszero(color[x])) && continue
+                if color[x] == color[v]
+                    _merge_trees!(v, w, x, disjoint_sets)  # merge trees T₁ ∋ vw and T₂ ∋ wx if T₁ != T₂
+                end
+            end
+        end
+    end
+    return color, TreeSet(disjoint_sets, parent)
+end
+
+function _prevent_cycle!(
+    # not modified
+    v::Integer,
+    w::Integer,
+    x::Integer,
+    color::AbstractVector{<:Integer},
+    # modified
+    first_visit_to_tree::AbstractVector{<:Tuple},
+    forbidden_colors::AbstractVector{<:Integer},
+    disjoint_sets::DisjointSets{<:Tuple{Int,Int}},
+)
+    wx = _sort(w, x)
+    root = find_root!(disjoint_sets, wx)  # edge wx belongs to the 2-colored tree T represented by edge "root"
+    id = disjoint_sets.intmap[root] # ID of the representative edge "root" of a two-colored tree.
+    (p, q) = first_visit_to_tree[id]
+    if p != v  # T is being visited from vertex v for the first time
+        vw = _sort(v, w)
+        first_visit_to_tree[id] = (v, w)
+    elseif q != w  # T is connected to vertex v via at least two edges
+        forbidden_colors[color[x]] = v
+    end
+    return nothing
+end
+
+function _grow_star!(
+    # not modified
+    v::Integer,
+    w::Integer,
+    color::AbstractVector{<:Integer},
+    # modified
+    first_neighbor::AbstractVector{<:Tuple},
+    disjoint_sets::DisjointSets{Tuple{Int,Int}},
+    parent::Vector{Int},
+)
+    vw = _sort(v, w)
+    push!(disjoint_sets, vw)  # Create a new tree T_{vw} consisting only of edge vw
+    push!(parent, v)  # the vertex v is the parent of the vertex w in tree T_{vw} -- parent[disjoint_sets.intmap[vw]] == v
+    (p, q) = first_neighbor[color[w]]
+    if p != v  # a neighbor of v with color[w] encountered for the first time
+        first_neighbor[color[w]] = (v, w)
+    else  # merge T_{vw} with a two-colored star being grown around v
+        vw = _sort(v, w)
+        pq = _sort(p, q)
+        root1 = find_root!(disjoint_sets, vw)
+        root2 = find_root!(disjoint_sets, pq)
+        root_union!(disjoint_sets, root1, root2)
+    end
+    return nothing
+end
+
+function _merge_trees!(
+    # not modified
+    v::Integer,
+    w::Integer,
+    x::Integer,
+    # modified
+    disjoint_sets::DisjointSets{Tuple{Int,Int}},
+)
+    vw = _sort(v, w)
+    wx = _sort(w, x)
+    root1 = find_root!(disjoint_sets, vw)
+    root2 = find_root!(disjoint_sets, wx)
+    if root1 != root2
+        root_union!(disjoint_sets, root1, root2)
+    end
+    return nothing
+end
+
+"""
+    TreeSet
+
+Encode a set of 2-colored trees resulting from the acyclic coloring algorithm.
+
+# Fields
+
+$TYPEDFIELDS
+
+# References
+
+> [_New Acyclic and Star Coloring Algorithms with Application to Computing Hessians_](https://epubs.siam.org/doi/abs/10.1137/050639879), Gebremedhin et al. (2007), Algorithm 4.1
+"""
+struct TreeSet
+    "set of 2-colored trees"
+    disjoint_sets::DisjointSets{Tuple{Int,Int}}
+    "???"  # TODO: fill this
+    parent::Vector{Int}
+end

src/decompression.jl

@@ -91,3 +91,140 @@ function decompress_aux!(
     end
     return A
 end
+
+function decompress_aux!(
+    A::AbstractMatrix{R},
+    B::AbstractMatrix{R},
+    result::AbstractColoringResult{:symmetric,:column,:substitution},
+) where {R<:Real}
+    # build T such that T * upper_nonzeros(A) = B and invert the linear system
+    # only consider the 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)
+    upper_nonzero_inds = Tuple{Int,Int}[]
+    I, J, _ = findnz(S)
+    for (i, j) in zip(I, J)
+        i >= j && push!(upper_nonzero_inds, (i, j))
+    end
+
+    T = spzeros(float(R), length(B), length(upper_nonzero_inds))
+    for (l, (i, j)) in enumerate(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))
+        if i != j
+            T[kj, l] = one(float(R))
+        end
+    end
+
+    upper_nonzeros_A = T \ vec(B)
+    for (l, (i, j)) in enumerate(upper_nonzero_inds)
+        A[i, j] = upper_nonzeros_A[l]
+        if i != j
+            A[j, i] = upper_nonzeros_A[l]
+        end
+    end
+    return A
+end
+
+#=
+function decompress_aux!(
+    A::AbstractMatrix{R},
+    B::AbstractMatrix{R},
+    result::AbstractColoringResult{:symmetric,:column,:substitution},
+) where {R<:Real}
+    @compat (; disjoint_sets, parent) = result
+
+    # to be optimized!
+    set_roots = Set{Int}()
+    ntrees = 0
+    for edge in tree_set.disjoint_sets.revmap
+        # ensure that all paths are compressed
+        root_edge = find_root!(disjoint_sets, edge)
+        root_index = tree_set.disjoint_sets.intmap[root_edge]
+
+        # we exclude trees related to diagonal coefficients
+        if (edge[1] != edge[2])
+            push!(set_roots, root_index)
+        end
+    end
+    roots = disjoint_sets.internal.parents
+
+    # DEBUG
+    println(set_roots)
+    ntrees = length(set_roots)
+    println(ntrees)
+
+    trees = [Int[] for i in 1:ntrees]
+    k = 0
+    for root in set_roots
+        k += 1
+        for (pos, val) in enumerate(roots)
+            if root == val
+                push!(trees[k], pos)
+            end
+        end
+    end
+    # for k in 1:ntrees
+    #     nedges = length(trees[k])
+    #     if nedges > 1
+    #         tree_edges = trees[k]
+    #         p = ...
+    #         trees[k] = tree_edges[p]
+    #     end
+    # end
+
+    # DEBUG
+    display(trees)
+
+    n = checksquare(A)
+    stored_values = Vector{R}(undef, n)
+    if !same_sparsity_pattern(A, S)
+        throw(DimensionMismatch("`A` and `S` must have the same sparsity pattern."))
+    end
+    A .= zero(R)
+    for i in axes(A, 1)
+        if !iszero(S[i, i])
+            A[i, i] = B[i, color[i]]
+        end
+    end
+    for tree in trees
+        nedges = length(tree)
+        if nedges == 1
+            edge_index = tree[1]
+            i, j = disjoint_sets.revmap[edge_index]
+            val = B[i, color[j]]
+            A[i, j] = val
+            A[j, i] = val
+        else
+            for edge_index in tree
+                i, j = disjoint_sets.revmap[edge_index]
+                stored_values[i] = zero(R)
+                stored_values[j] = zero(R)
+            end
+            for edge_index in tree  # edges are sorted by their distance to the root
+                i, j = disjoint_sets.revmap[edge_index]
+                parent_index = disjoint_sets.internal.parents[edge_index]
+                k, l = disjoint_sets.revmap[parent_index]
+                # k = parent[edge_index]
+                if edge_index != parent_index
+                    if i == k || i == l  # vertex i is the parent of vertex j
+                        i, j = j, i  # ensure that i always denotes a leaf vertex
+                    end
+                end
+                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
+    end
+    return A
+end
+=#