Optimize remap_colors (#197)

* Avoid the call to remap_colors in postprocess!

* Optimize remap_colors

* Apply suggestions from code review

* Update src/coloring.jl

* Update coloring.jl

Author: amontoison

src/coloring.jl

@@ -123,7 +123,8 @@ function star_coloring(g::AdjacencyGraph, order::AbstractOrder; postprocessing::
     end
     star_set = StarSet(star, hub, nb_spokes)
     if postprocessing
-        postprocess!(color, star_set, g)
+        # Reuse the vector forbidden_colors to compute offsets during post-processing
+        postprocess!(color, star_set, g, forbidden_colors)
     end
     return color, star_set
 end
@@ -351,7 +352,8 @@ function acyclic_coloring(g::AdjacencyGraph, order::AbstractOrder; postprocessin
     end
     tree_set = TreeSet(forest, nb_vertices(g))
     if postprocessing
-        postprocess!(color, tree_set, g)
+        # Reuse the vector forbidden_colors to compute offsets during post-processing
+        postprocess!(color, tree_set, g, forbidden_colors)
     end
     return color, tree_set
 end
@@ -551,6 +553,7 @@ function postprocess!(
     color::AbstractVector{<:Integer},
     star_or_tree_set::Union{StarSet,TreeSet},
     g::AdjacencyGraph,
+    offsets::Vector{Int},
 )
     (; S) = g
     # flag which colors are actually used during decompression
@@ -642,15 +645,28 @@ function postprocess!(
 
     # if at least one of the colors is useless, modify the color assignments of vertices
     if any(!, color_used)
-        # assign the neutral color to every vertex with a useless color
+        num_colors_useless = 0
+
+        # determine what are the useless colors and compute the offsets
+        for ci in 1:nb_colors
+            if color_used[ci]
+                offsets[ci] = num_colors_useless
+            else
+                num_colors_useless += 1
+            end
+        end
+
+        # assign the neutral color to every vertex with a useless color and remap the colors
         for i in eachindex(color)
             ci = color[i]
             if !color_used[ci]
+                # assign the neutral color
                 color[i] = 0
+            else
+                # remap the color to not have any gap
+                color[i] -= offsets[ci]
             end
         end
-        # remap colors to decrease the highest one by filling gaps
-        color .= remap_colors(color)[1]
     end
     return color
 end

src/decompression.jl

@@ -711,9 +711,9 @@ function _join_compressed!(result::BicoloringResult, Br::AbstractMatrix, Bc::Abs
     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.
+    We use the vectors `symmetric_to_row` and `symmetric_to_column` to map from symmetric colors to row and column colors.
     =#
-    (; A, col_color_ind, row_color_ind) = result
+    (; A, symmetric_to_column, symmetric_to_row) = result
     m, n = size(A)
     R = Base.promote_eltype(Br, Bc)
     if eltype(result.Br_and_Bc) == R
@@ -723,11 +723,13 @@ function _join_compressed!(result::BicoloringResult, Br::AbstractMatrix, Bc::Abs
     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], :))
+        if symmetric_to_row[c] > 0  # some rows were colored with the symmetric color c
+            copyto!(view(Br_and_Bc, 1:n, c), view(Br, symmetric_to_row[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]))
+        if symmetric_to_column[c] > 0  # some columns were colored with the symmetric color c
+            copyto!(
+                view(Br_and_Bc, (n + 1):(n + m), c), view(Bc, :, symmetric_to_column[c])
+            )
         end
     end
     return Br_and_Bc

src/result.jl

@@ -456,22 +456,55 @@ end
 ## Bicoloring result
 
 """
-    remap_colors(color::Vector{Int})
+    remap_colors(color::Vector{Int}, num_sym_colors::Int, m::Int, n::Int)
 
-Renumber the colors in `color` using their index in the vector `sort(unique(color))`, so that the non-zero colors are forced to go from `1` to some `cmax` contiguously.
+Return a tuple `(row_color, column_color, symmetric_to_row, symmetric_to_column)` such that `row_color` and `column_color` are vectors containing the renumbered colors for rows and columns.
+`symmetric_to_row` and `symmetric_to_column` are vectors that map symmetric colors to row and column colors.
 
-Return a tuple `(remapped_colors, color_to_ind)` such that `remapped_colors` is a vector containing the renumbered colors and `color_to_ind` is a dictionary giving the translation between old and new color numberings.
+For all vertex indices `i` between `1` and `m` we have:
 
-For all vertex indices `i` we have:
+    row_color[i] = symmetric_to_row[color[n+i]]
 
-    remapped_color[i] = color_to_ind[color[i]]
+For all vertex indices `j` between `1` and `n` we have:
+
+    column_color[j] = symmetric_to_column[color[j]]
 """
-function remap_colors(color::Vector{Int})
-    sorted_colors = sort!(unique(color))
-    offset = sorted_colors[1] == 0 ? 1 : 0
-    color_to_ind = Dict(c => i - offset for (i, c) in enumerate(sorted_colors))
-    remapped_colors = [color_to_ind[c] for c in color]
-    return remapped_colors, color_to_ind
+function remap_colors(color::Vector{Int}, num_sym_colors::Int, m::Int, n::Int)
+    # Map symmetric colors to column colors
+    symmetric_to_column = zeros(Int, num_sym_colors)
+    column_color = zeros(Int, n)
+
+    counter = 0
+    for j in 1:n
+        cj = color[j]
+        if cj > 0
+            # First time that we encounter this column color
+            if symmetric_to_column[cj] == 0
+                counter += 1
+                symmetric_to_column[cj] = counter
+            end
+            column_color[j] = symmetric_to_column[cj]
+        end
+    end
+
+    # Map symmetric colors to row colors
+    symmetric_to_row = zeros(Int, num_sym_colors)
+    row_color = zeros(Int, m)
+
+    counter = 0
+    for i in (n + 1):(n + m)
+        ci = color[i]
+        if ci > 0
+            # First time that we encounter this row color
+            if symmetric_to_row[ci] == 0
+                counter += 1
+                symmetric_to_row[ci] = counter
+            end
+            row_color[i - n] = symmetric_to_row[ci]
+        end
+    end
+
+    return row_color, column_color, symmetric_to_row, symmetric_to_column
 end
 
 """
@@ -509,10 +542,10 @@ struct BicoloringResult{
     row_group::V
     "result for the coloring of the symmetric 2 x 2 block matrix"
     symmetric_result::SR
-    "column color to index"
-    col_color_ind::Dict{Int,Int}
-    "row color to index"
-    row_color_ind::Dict{Int,Int}
+    "maps symmetric colors to column colors"
+    symmetric_to_column::Vector{Int}
+    "maps symmetric colors to row colors"
+    symmetric_to_row::Vector{Int}
     "combination of `Br` and `Bc` (almost a concatenation up to color remapping)"
     Br_and_Bc::Matrix{R}
     "CSC storage of `A_and_noAᵀ - `colptr`"
@@ -535,11 +568,13 @@ function BicoloringResult(
 ) where {R}
     m, n = size(A)
     symmetric_color = column_colors(symmetric_result)
-    column_color, col_color_ind = remap_colors(symmetric_color[1:n])
-    row_color, row_color_ind = remap_colors(symmetric_color[(n + 1):(n + m)])
+    num_sym_colors = maximum(symmetric_color)
+    row_color, column_color, symmetric_to_row, symmetric_to_column = remap_colors(
+        symmetric_color, num_sym_colors, m, n
+    )
     column_group = group_by_color(column_color)
     row_group = group_by_color(row_color)
-    Br_and_Bc = Matrix{R}(undef, n + m, maximum(column_colors(symmetric_result)))
+    Br_and_Bc = Matrix{R}(undef, n + m, num_sym_colors)
     large_colptr = copy(ag.S.colptr)
     large_colptr[(n + 2):end] .= large_colptr[n + 1]  # last few columns are empty
     large_rowval = ag.S.rowval[1:(end ÷ 2)]  # forget the second half of nonzeros
@@ -551,8 +586,8 @@ function BicoloringResult(
         column_group,
         row_group,
         symmetric_result,
-        col_color_ind,
-        row_color_ind,
+        symmetric_to_column,
+        symmetric_to_row,
         Br_and_Bc,
         large_colptr,
         large_rowval,