Revisit post-processing

Author: amontoison

src/SparseMatrixColorings.jl

@@ -47,6 +47,7 @@ include("graph.jl")
 include("forest.jl")
 include("order.jl")
 include("coloring.jl")
+include("postprocessing.jl")
 include("result.jl")
 include("matrices.jl")
 include("interface.jl")

src/coloring.jl

@@ -109,7 +109,8 @@ The optional `forced_colors` keyword argument is used to enforce predefined vert
 function star_coloring(
     g::AdjacencyGraph{T},
     vertices_in_order::AbstractVector{<:Integer},
-    postprocessing::Bool;
+    postprocessing::Bool,
+    postprocessing_minimizes::Symbol;
     forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing,
 ) where {T<:Integer}
     # Initialize data structures
@@ -168,7 +169,7 @@ function star_coloring(
     if postprocessing
         # Reuse the vector forbidden_colors to compute offsets during post-processing
         offsets = forbidden_colors
-        postprocess!(color, star_set, g, offsets)
+        postprocess!(color, star_set, g, offsets, postprocessing_minimizes)
     end
     return color, star_set
 end
@@ -273,7 +274,10 @@ If `postprocessing=true`, some colors might be replaced with `0` (the "neutral"
 > [_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::AdjacencyGraph{T}, vertices_in_order::AbstractVector{<:Integer}, postprocessing::Bool
+    g::AdjacencyGraph{T},
+    vertices_in_order::AbstractVector{<:Integer},
+    postprocessing::Bool,
+    postprocessing_minimizes::Symbol,
 ) where {T<:Integer}
     # Initialize data structures
     nv = nb_vertices(g)
@@ -345,7 +349,7 @@ function acyclic_coloring(
     if postprocessing
         # Reuse the vector forbidden_colors to compute offsets during post-processing
         offsets = forbidden_colors
-        postprocess!(color, tree_set, g, offsets)
+        postprocess!(color, tree_set, g, offsets, postprocessing_minimizes)
     end
     return color, tree_set
 end
@@ -643,150 +647,3 @@ function TreeSet(
 
     return TreeSet(reverse_bfs_orders, is_star, tree_edge_indices, nt)
 end
-
-## Postprocessing, mirrors decompression code
-
-function postprocess!(
-    color::AbstractVector{<:Integer},
-    star_or_tree_set::Union{StarSet,TreeSet},
-    g::AdjacencyGraph,
-    offsets::AbstractVector{<:Integer},
-)
-    S = pattern(g)
-    edge_to_index = edge_indices(g)
-    # flag which colors are actually used during decompression
-    nb_colors = maximum(color)
-    color_used = zeros(Bool, nb_colors)
-
-    # nonzero diagonal coefficients force the use of their respective color (there can be no neutral colors if the diagonal is fully nonzero)
-    if has_diagonal(g)
-        for i in axes(S, 1)
-            if !iszero(S[i, i])
-                color_used[color[i]] = true
-            end
-        end
-    end
-
-    if star_or_tree_set isa StarSet
-        # only the colors of the hubs are used
-        (; star, hub) = star_or_tree_set
-        nb_trivial_stars = 0
-
-        # Iterate through all non-trivial stars
-        for s in eachindex(hub)
-            h = hub[s]
-            if h > 0
-                color_used[color[h]] = true
-            else
-                nb_trivial_stars += 1
-            end
-        end
-
-        # Process the trivial stars (if any)
-        if nb_trivial_stars > 0
-            rvS = rowvals(S)
-            for j in axes(S, 2)
-                for k in nzrange(S, j)
-                    i = rvS[k]
-                    if i > j
-                        index_ij = edge_to_index[k]
-                        s = star[index_ij]
-                        h = hub[s]
-                        if h < 0
-                            h = abs(h)
-                            spoke = h == j ? i : j
-                            if color_used[color[spoke]]
-                                # Switch the hub and the spoke to possibly avoid adding one more used color
-                                hub[s] = spoke
-                            else
-                                # Keep the current hub
-                                color_used[color[h]] = true
-                            end
-                        end
-                    end
-                end
-            end
-        end
-    else
-        # only the colors of non-leaf vertices are used
-        (; reverse_bfs_orders, is_star, tree_edge_indices, nt) = star_or_tree_set
-        nb_trivial_trees = 0
-
-        # Iterate through all non-trivial trees
-        for k in 1:nt
-            # Position of the first edge in the tree
-            first = tree_edge_indices[k]
-
-            # Total number of edges in the tree
-            ne_tree = tree_edge_indices[k + 1] - first
-
-            # Check if we have more than one edge in the tree (non-trivial tree)
-            if ne_tree > 1
-                # Determine if the tree is a star
-                if is_star[k]
-                    # It is a non-trivial star and only the color of the hub is needed
-                    (_, hub) = reverse_bfs_orders[first]
-                    color_used[color[hub]] = true
-                else
-                    # It is not a star and both colors are needed during the decompression
-                    (i, j) = reverse_bfs_orders[first]
-                    color_used[color[i]] = true
-                    color_used[color[j]] = true
-                end
-            else
-                nb_trivial_trees += 1
-            end
-        end
-
-        # Process the trivial trees (if any)
-        if nb_trivial_trees > 0
-            for k in 1:nt
-                # Position of the first edge in the tree
-                first = tree_edge_indices[k]
-
-                # Total number of edges in the tree
-                ne_tree = tree_edge_indices[k + 1] - first
-
-                # Check if we have exactly one edge in the tree
-                if ne_tree == 1
-                    (i, j) = reverse_bfs_orders[first]
-                    if color_used[color[i]]
-                        # Make i the root to avoid possibly adding one more used color
-                        # Switch it with the (only) leaf
-                        reverse_bfs_orders[first] = (j, i)
-                    else
-                        # Keep j as the root
-                        color_used[color[j]] = true
-                    end
-                end
-            end
-        end
-    end
-
-    # if at least one of the colors is useless, modify the color assignments of vertices
-    if any(!, color_used)
-        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
-    end
-    return color
-end

src/interface.jl

@@ -69,11 +69,12 @@ It is passed as an argument to the main function [`coloring`](@ref).
 
 # Constructors
 
-    GreedyColoringAlgorithm{decompression}(order=NaturalOrder(); postprocessing=false)
-    GreedyColoringAlgorithm(order=NaturalOrder(); postprocessing=false, decompression=:direct)
+    GreedyColoringAlgorithm{decompression}(order=NaturalOrder(); postprocessing=false, postprocessing_minimizes=:all_colors)
+    GreedyColoringAlgorithm(order=NaturalOrder(); postprocessing=false, postprocessing_minimizes=:all_colors, decompression=:direct)
 
 - `order::Union{AbstractOrder,Tuple}`: the order in which the columns or rows are colored, which can impact the number of colors. Can also be a tuple of different orders to try out, from which the best order (the one with the lowest total number of colors) will be used.
 - `postprocessing::Bool`: whether or not the coloring will be refined by assigning the neutral color `0` to some vertices.
+- `postprocessing_minimizes::Symbol`: either `:all_colors`, `:row_colors` or `:column_colors`. The options `:row_colors` and `:column_colors` are only available for bicoloring. Otherwise, the setting defaults to `:all_colors`.
 - `decompression::Symbol`: either `:direct` or `:substitution`. Usually `:substitution` leads to fewer colors, at the cost of a more expensive coloring (and decompression). When `:substitution` is not applicable, it falls back on `:direct` decompression.
 
 !!! warning
@@ -98,27 +99,30 @@ struct GreedyColoringAlgorithm{decompression,N,O<:NTuple{N,AbstractOrder}} <:
        ADTypes.AbstractColoringAlgorithm
     orders::O
     postprocessing::Bool
+    postprocessing_minimizes::Symbol
 
     function GreedyColoringAlgorithm{decompression}(
         order_or_orders::Union{AbstractOrder,Tuple}=NaturalOrder();
         postprocessing::Bool=false,
+        postprocessing_minimizes::Symbol=:all_colors,
     ) where {decompression}
         check_valid_algorithm(decompression)
         if order_or_orders isa AbstractOrder
             orders = (order_or_orders,)
         else
             orders = order_or_orders
         end
-        return new{decompression,length(orders),typeof(orders)}(orders, postprocessing)
+        return new{decompression,length(orders),typeof(orders)}(orders, postprocessing, postprocessing_minimizes)
     end
 end
 
 function GreedyColoringAlgorithm(
     order_or_orders::Union{AbstractOrder,Tuple}=NaturalOrder();
     postprocessing::Bool=false,
     decompression::Symbol=:direct,
+    postprocessing_minimizes::Symbol=:all_colors,
 )
-    return GreedyColoringAlgorithm{decompression}(order_or_orders; postprocessing)
+    return GreedyColoringAlgorithm{decompression}(order_or_orders; postprocessing, postprocessing_minimizes)
 end
 
 ## Coloring
@@ -282,7 +286,7 @@ function _coloring(
     ag = AdjacencyGraph(A; has_diagonal=true)
     color_and_star_set_by_order = map(algo.orders) do order
         vertices_in_order = vertices(ag, order)
-        return star_coloring(ag, vertices_in_order, algo.postprocessing; forced_colors)
+        return star_coloring(ag, vertices_in_order, algo.postprocessing, :all_colors; forced_colors)
     end
     color, star_set = argmin(maximum ∘ first, color_and_star_set_by_order)
     if speed_setting isa WithResult
@@ -303,7 +307,7 @@ function _coloring(
     ag = AdjacencyGraph(A; has_diagonal=true)
     color_and_tree_set_by_order = map(algo.orders) do order
         vertices_in_order = vertices(ag, order)
-        return acyclic_coloring(ag, vertices_in_order, algo.postprocessing)
+        return acyclic_coloring(ag, vertices_in_order, algo.postprocessing, :all_colors)
     end
     color, tree_set = argmin(maximum ∘ first, color_and_tree_set_by_order)
     if speed_setting isa WithResult
@@ -327,7 +331,7 @@ function _coloring(
     outputs_by_order = map(algo.orders) do order
         vertices_in_order = vertices(ag, order)
         _color, _star_set = star_coloring(
-            ag, vertices_in_order, algo.postprocessing; forced_colors
+            ag, vertices_in_order, algo.postprocessing, algo.postprocessing_minimizes; forced_colors
         )
         (_row_color, _column_color, _symmetric_to_row, _symmetric_to_column) = remap_colors(
             eltype(ag), _color, maximum(_color), size(A)...
@@ -373,7 +377,7 @@ function _coloring(
     ag = AdjacencyGraph(A_and_Aᵀ, edge_to_index; has_diagonal=false)
     outputs_by_order = map(algo.orders) do order
         vertices_in_order = vertices(ag, order)
-        _color, _tree_set = acyclic_coloring(ag, vertices_in_order, algo.postprocessing)
+        _color, _tree_set = acyclic_coloring(ag, vertices_in_order, algo.postprocessing, algo.postprocessing_minimizes)
         (_row_color, _column_color, _symmetric_to_row, _symmetric_to_column) = remap_colors(
             eltype(ag), _color, maximum(_color), size(A)...
         )