Bicoloring (#152)

* Implement efficient adjacency graph from bipartite (#145) * Custom adjacency graph from bipartite * Fix nb edges * More typing * Working prototype for bicoloring (coloring + decompression) (#146) * Start bicoloring result * Working bidirectional decompression * Min diff * Slight perf optim * Perf * Gettin bi * Less code * Fix tests * Remove remaining compat * Fix group typing * Avoid accessing ag.S * Add bidirectional coloring visualization (#153) * Allocation-free bidirectional decompression (#154) * Minimize diff * Bump version to 0.4.9 * Check scale for bicoloring visualization * Fix visualization tests * Improve docs * Undo view change * Typo * Add views
gdalle · Nov 8, 2024 · 0246517 · 0246517 · gdalle · Nov 8, 2024
1 parent 42622b5
commit 0246517
Show file tree

Hide file tree

Showing 14 changed files with 452 additions and 104 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "SparseMatrixColorings"
 uuid = "0a514795-09f3-496d-8182-132a7b665d35"
 authors = ["Guillaume Dalle", "Alexis Montoison"]
-version = "0.4.8"
+version = "0.4.9"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

diff --git a/docs/src/api.md b/docs/src/api.md
@@ -26,6 +26,7 @@ ConstantColoringAlgorithm
 AbstractColoringResult
 column_colors
 row_colors
+ncolors
 column_groups
 row_groups
 sparsity_pattern

diff --git a/docs/src/dev.md b/docs/src/dev.md
@@ -38,6 +38,8 @@ SparseMatrixColorings.RowColoringResult
 SparseMatrixColorings.StarSetColoringResult
 SparseMatrixColorings.TreeSetColoringResult
 SparseMatrixColorings.LinearSystemColoringResult
+SparseMatrixColorings.BicoloringResult
+SparseMatrixColorings.remap_colors
 ```
 
 ## Testing

diff --git a/ext/SparseMatrixColoringsColorsExt.jl b/ext/SparseMatrixColoringsColorsExt.jl
@@ -21,7 +21,8 @@ using SparseMatrixColorings:
     AbstractColoringResult,
     sparsity_pattern,
     column_colors,
-    row_colors
+    row_colors,
+    ncolors
 using Colors: Colorant, RGB, RGBA, distinguishable_colors
 
 const DEFAULT_BACKGROUND = RGBA(0, 0, 0, 0)
@@ -31,9 +32,6 @@ const DEFAULT_PAD = 0 # update docstring in src/images.jl when changing this def
 # Sample n distinguishable colors, excluding the background color
 default_colorscheme(n, background) = distinguishable_colors(n, background; dropseed=true)
 
-ncolors(res::AbstractColoringResult{s,:column}) where {s} = maximum(column_colors(res))
-ncolors(res::AbstractColoringResult{s,:row}) where {s} = maximum(row_colors(res))
-
 ## Top-level function that handles argument errors, eagerly promotes types and allocates output buffer
 
 function SparseMatrixColorings.show_colors(
@@ -119,4 +117,30 @@ function show_colors!(
     return out
 end
 
+function show_colors!(
+    out, res::AbstractColoringResult{s,:bidirectional}, colorscheme, scale, pad
+) where {s}
+    scale < 3 && throw(ArgumentError("`scale` has to be ≥ 3 to visualize bicoloring"))
+    ccolor_indices = mod1.(column_colors(res), length(colorscheme)) # cycle color indices if necessary
+    row_shift = maximum(column_colors(res))
+    rcolor_indices = mod1.(row_shift .+ row_colors(res), length(colorscheme)) # cycle color indices if necessary
+    ccolors = colorscheme[ccolor_indices]
+    rcolors = colorscheme[rcolor_indices]
+    pattern = sparsity_pattern(res)
+    for I in CartesianIndices(pattern)
+        if !iszero(pattern[I])
+            r, c = Tuple(I)
+            area = matrix_entry_area(I, scale, pad)
+            for i in axes(area, 1), j in axes(area, 2)
+                if j > i
+                    out[area[i, j]] = ccolors[c]
+                elseif i > j
+                    out[area[i, j]] = rcolors[r]
+                end
+            end
+        end
+    end
+    return out
+end
+
 end # module
diff --git a/src/SparseMatrixColorings.jl b/src/SparseMatrixColorings.jl
@@ -60,7 +60,7 @@ export DynamicDegreeBasedOrder, SmallestLast, IncidenceDegree, DynamicLargestFir
 export ColoringProblem, GreedyColoringAlgorithm, AbstractColoringResult
 export ConstantColoringAlgorithm
 export coloring
-export column_colors, row_colors
+export column_colors, row_colors, ncolors
 export column_groups, row_groups
 export sparsity_pattern
 export compress, decompress, decompress!, decompress_single_color!

diff --git a/src/decompression.jl b/src/decompression.jl
@@ -9,9 +9,6 @@ Compress `A` given a coloring `result` of the sparsity pattern of `A`.
 Compression means summing either the columns or the rows of `A` which share the same color.
 It is undone by calling [`decompress`](@ref) or [`decompress!`](@ref).
 
-!!! warning
-    At the moment, `:bidirectional` partitions are not implemented.
-
 # Example
 
 ```jldoctest
@@ -63,10 +60,25 @@ function compress(A, result::AbstractColoringResult{structure,:row}) where {stru
     return B
 end
 
+function compress(
+    A, result::AbstractColoringResult{structure,:bidirectional}
+) where {structure}
+    row_group = row_groups(result)
+    column_group = column_groups(result)
+    Br = stack(row_group; dims=1) do g
+        dropdims(sum(A[g, :]; dims=1); dims=1)
+    end
+    Bc = stack(column_group; dims=2) do g
+        dropdims(sum(A[:, g]; dims=2); dims=2)
+    end
+    return Br, Bc
+end
+
 """
-    decompress(B::AbstractMatrix, result::AbstractColoringResult)
+    decompress(B::AbstractMatrix, result::AbstractColoringResult{_,:column/:row})
+    decompress(Br::AbstractMatrix, Bc::AbstractMatrix, result::AbstractColoringResult{_,:bidirectional})
 
-Decompress `B` into a new matrix `A`, given a coloring `result` of the sparsity pattern of `A`.
+Decompress `B` (or the tuple `(Br,Bc)`) 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.
@@ -120,13 +132,27 @@ function decompress(B::AbstractMatrix, result::AbstractColoringResult)
     return decompress!(A, B, result)
 end
 
+function decompress(
+    Br::AbstractMatrix,
+    Bc::AbstractMatrix,
+    result::AbstractColoringResult{structure,:bidirectional},
+) where {structure}
+    A = respectful_similar(result.A, Base.promote_eltype(Br, Bc))
+    return decompress!(A, Br, Bc, result)
+end
+
 """
     decompress!(
         A::AbstractMatrix, B::AbstractMatrix,
-        result::AbstractColoringResult, [uplo=:F]
+        result::AbstractColoringResult{_,:column/:row}, [uplo=:F]
+    )
+
+    decompress!(
+        A::AbstractMatrix, Br::AbstractMatrix, Bc::AbstractMatrix
+        result::AbstractColoringResult{_,:bidirectional}
     )
 
-Decompress `B` in-place into `A`, given a coloring `result` of the sparsity pattern of `A`.
+Decompress `B` (or the tuple `(Br,Bc)`) in-place into `A`, given a coloring `result` of the sparsity pattern of `A`.
 The out-of-place alternative is [`decompress`](@ref).
 
 !!! note
@@ -632,3 +658,60 @@ function decompress!(
     end
     return A
 end
+
+## BicoloringResult
+
+function _join_compressed!(result::BicoloringResult, Br::AbstractMatrix, Bc::AbstractMatrix)
+    #=
+    Say we have an original matrix `A` of size `(n, m)` and we build an augmented matrix `A_and_Aᵀ = [zeros(n, n) Aᵀ; A zeros(m, m)]`.
+    Its first `1:n` columns have the form `[zeros(n); A[:, j]]` and its following `n+1:n+m` columns have the form `[A[i, :]; zeros(m)]`.
+    The symmetric column coloring is performed on `A_and_Aᵀ` and the column-wise compression of `A_and_Aᵀ` should return a matrix `Br_and_Bc`.
+    But in reality, `Br_and_Bc` is computed as two partial compressions: the row-wise compression `Br` (corresponding to `Aᵀ`) and the columnwise compression `Bc` (corresponding to `A`).
+    Before symmetric decompression, we must reconstruct `Br_and_Bc` from `Br` and `Bc`, knowing that the symmetric colors (those making up `Br_and_Bc`) are present in either a row of `Br`, a column of `Bc`, or both.
+    Therefore, the column indices in `Br_and_Bc` don't necessarily match with the row indices in `Br` or the column indices in `Bc` since some colors may be missing in the partial compressions.
+    The columns of the top part of `Br_and_Bc` (rows `1:n`) are the rows of `Br`, interlaced with zero columns whenever the current color hasn't been used to color any row.
+    The columns of the bottom part of `Br_and_Bc` (rows `n+1:n+m`) are the columns of `Bc`, interlaced with zero columns whenever the current color hasn't been used to color any column.
+    We use the dictionaries `col_color_ind` and `row_color_ind` to map from symmetric colors to row/column colors.
+    =#
+    (; A, col_color_ind, row_color_ind) = result
+    m, n = size(A)
+    R = Base.promote_eltype(Br, Bc)
+    if eltype(result.Br_and_Bc) == R
+        Br_and_Bc = result.Br_and_Bc
+    else
+        Br_and_Bc = similar(result.Br_and_Bc, R)
+    end
+    fill!(Br_and_Bc, zero(R))
+    for c in axes(Br_and_Bc, 2)
+        if haskey(row_color_ind, c)  # some rows were colored with symmetric color c
+            copyto!(view(Br_and_Bc, 1:n, c), view(Br, row_color_ind[c], :))
+        end
+        if haskey(col_color_ind, c)  # some columns were colored with symmetric c
+            copyto!(view(Br_and_Bc, (n + 1):(n + m), c), view(Bc, :, col_color_ind[c]))
+        end
+    end
+    return Br_and_Bc
+end
+
+function decompress!(
+    A::AbstractMatrix, Br::AbstractMatrix, Bc::AbstractMatrix, result::BicoloringResult
+)
+    m, n = size(A)
+    Br_and_Bc = _join_compressed!(result, Br, Bc)
+    A_and_Aᵀ = decompress(Br_and_Bc, result.symmetric_result)
+    copyto!(A, A_and_Aᵀ[(n + 1):(n + m), 1:n])  # original matrix in bottom left corner
+    return A
+end
+
+function decompress!(
+    A::SparseMatrixCSC, Br::AbstractMatrix, Bc::AbstractMatrix, result::BicoloringResult
+)
+    (; large_colptr, large_rowval, symmetric_result) = result
+    m, n = size(A)
+    Br_and_Bc = _join_compressed!(result, Br, Bc)
+    # pretend A is larger
+    A_and_noAᵀ = SparseMatrixCSC(m + n, m + n, large_colptr, large_rowval, A.nzval)
+    # decompress lower triangle only
+    decompress!(A_and_noAᵀ, Br_and_Bc, symmetric_result, :L)
+    return A
+end
diff --git a/src/graph.jl b/src/graph.jl
@@ -145,17 +145,13 @@ function degree(g::AdjacencyGraph, v::Integer)
 end
 
 function nb_edges(g::AdjacencyGraph)
-    S = pattern(g)
     ne = 0
-    for j in vertices(g)
-        for k in nzrange(S, j)
-            i = rowvals(S)[k]
-            if i > j
-                ne += 1
-            end
+    for v in vertices(g)
+        for u in neighbors(g, v)
+            ne += 1
         end
     end
-    return ne
+    return ne ÷ 2
 end
 
 maximum_degree(g::AdjacencyGraph) = maximum(Base.Fix1(degree, g), vertices(g))

diff --git a/src/interface.jl b/src/interface.jl
@@ -1,6 +1,6 @@
 function check_valid_problem(structure::Symbol, partition::Symbol)
     valid = (
-        (structure == :nonsymmetric && partition in (:column, :row)) ||
+        (structure == :nonsymmetric && partition in (:column, :row, :bidirectional)) ||
         (structure == :symmetric && partition == :column)
     )
     if !valid
@@ -49,7 +49,7 @@ Matrix coloring is often used in automatic differentiation, and here is the tran
 | -------- | ------- | --------------- | ---------------- | ----------- |
 | Jacobian | forward | `:nonsymmetric` | `:column`        | yes         |
 | Jacobian | reverse | `:nonsymmetric` | `:row`           | yes         |
-| Jacobian | mixed   | `:nonsymmetric` | `:bidirectional` | no          |
+| Jacobian | mixed   | `:nonsymmetric` | `:bidirectional` | yes         |
 | Hessian  | -       | `:symmetric`    | `:column`        | yes         |
 | Hessian  | -       | `:symmetric`    | `:row`           | no          |
 """
@@ -223,6 +223,37 @@ function coloring(
     return TreeSetColoringResult(A, ag, color, tree_set, decompression_eltype)
 end
 
+function coloring(
+    A::AbstractMatrix,
+    ::ColoringProblem{:nonsymmetric,:bidirectional},
+    algo::GreedyColoringAlgorithm{decompression};
+    decompression_eltype::Type{R}=Float64,
+    symmetric_pattern::Bool=false,
+) where {decompression,R}
+    m, n = size(A)
+    T = eltype(A)
+    Aᵀ = if symmetric_pattern || A isa Union{Symmetric,Hermitian}
+        A
+    else
+        transpose(A)
+    end  # TODO: fuse with next step?
+    A_and_Aᵀ = [
+        spzeros(T, n, n) SparseMatrixCSC(Aᵀ)
+        SparseMatrixCSC(A) spzeros(T, m, m)
+    ]  # TODO: slow
+    ag = AdjacencyGraph(A_and_Aᵀ)
+    if decompression == :direct
+        color, star_set = star_coloring(ag, algo.order)
+        symmetric_result = StarSetColoringResult(A_and_Aᵀ, ag, color, star_set)
+    else
+        color, tree_set = acyclic_coloring(ag, algo.order)
+        symmetric_result = TreeSetColoringResult(
+            A_and_Aᵀ, ag, color, tree_set, decompression_eltype
+        )
+    end
+    return BicoloringResult(A, ag, symmetric_result, decompression_eltype)
+end
+
 ## ADTypes interface
 
 function ADTypes.column_coloring(A::AbstractMatrix, algo::GreedyColoringAlgorithm)