Revisit the post-processing

src/coloring.jl

@@ -81,7 +81,7 @@ end
 """
     star_coloring(
         g::AdjacencyGraph, vertices_in_order::AbstractVector, postprocessing::Bool;
-        forced_colors::Union{AbstractVector,Nothing}=nothing
+        postprocessing_minimizes::Symbol=:all_colors, forced_colors::Union{AbstractVector,Nothing}=nothing
     )
 
 Compute a star coloring of all vertices in the adjacency graph `g` and return a tuple `(color, star_set)`, where
@@ -110,6 +110,7 @@ function star_coloring(
     g::AdjacencyGraph{T},
     vertices_in_order::AbstractVector{<:Integer},
     postprocessing::Bool;
+    postprocessing_minimizes::Symbol=:all_colors,
     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
@@ -250,7 +251,8 @@ struct StarSet{T}
 end
 
 """
-    acyclic_coloring(g::AdjacencyGraph, vertices_in_order::AbstractVector, postprocessing::Bool)
+    acyclic_coloring(g::AdjacencyGraph, vertices_in_order::AbstractVector, postprocessing::Bool;
+                     postprocessing_minimizes::Symbol=:all_colors)
 
 Compute an acyclic coloring of all vertices in the adjacency graph `g` and return a tuple `(color, tree_set)`, where
 
@@ -273,7 +275,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=:all_colors,
 ) where {T<:Integer}
     # Initialize data structures
     nv = nb_vertices(g)
@@ -345,7 +350,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

src/graph.jl

@@ -227,6 +227,7 @@ The adjacency graph of a symmetric matrix `A ∈ ℝ^{n × n}` is `G(A) = (V, E)
 struct AdjacencyGraph{T<:Integer,augmented_graph}
     S::SparsityPatternCSC{T}
     edge_to_index::Vector{T}
+    original_size::Tuple{Int,Int}
 end
 
 Base.eltype(::AdjacencyGraph{T}) where {T} = T
@@ -235,15 +236,16 @@ function AdjacencyGraph(
     S::SparsityPatternCSC{T},
     edge_to_index::Vector{T}=build_edge_to_index(S);
     augmented_graph::Bool=false,
+    original_size::Tuple{Int,Int}=size(S),    
 ) where {T}
-    return AdjacencyGraph{T,augmented_graph}(S, edge_to_index)
+    return AdjacencyGraph{T,augmented_graph}(S, edge_to_index, original_size)
 end
 
-function AdjacencyGraph(A::SparseMatrixCSC; augmented_graph::Bool=false)
+function AdjacencyGraph(A::SparseMatrixCSC; augmented_graph::Bool=false, original_size::Tuple{Int,Int}=size(A))
     return AdjacencyGraph(SparsityPatternCSC(A); augmented_graph)
 end
 
-function AdjacencyGraph(A::AbstractMatrix; augmented_graph::Bool=false)
+function AdjacencyGraph(A::AbstractMatrix; augmented_graph::Bool=false, original_size::Tuple{Int,Int}=size(A))
     return AdjacencyGraph(SparseMatrixCSC(A); augmented_graph)
 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::Bool`: whether or not the coloring will be refined by assigning the neutral color `0` to some vertices. This option does not affect row or column colorings.
+- `postprocessing_minimizes::Symbol`: which number of distinct colors is heuristically minimized by postprocessing, either `:all_colors`, `:row_colors` or `:column_colors`. This option only affects bidirectional colorings.
 - `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
@@ -279,7 +283,7 @@ function _coloring(
     symmetric_pattern::Bool;
     forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing,
 )
-    ag = AdjacencyGraph(A; augmented_graph=false)
+    ag = AdjacencyGraph(A; augmented_graph=false, original_size=size(A))
     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)
@@ -300,7 +304,7 @@ function _coloring(
     decompression_eltype::Type{R},
     symmetric_pattern::Bool,
 ) where {R}
-    ag = AdjacencyGraph(A; augmented_graph=false)
+    ag = AdjacencyGraph(A; augmented_graph=false, original_size=size(A))
     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)
@@ -323,11 +327,12 @@ function _coloring(
     forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing,
 ) where {R}
     A_and_Aᵀ, edge_to_index = bidirectional_pattern(A; symmetric_pattern)
-    ag = AdjacencyGraph(A_and_Aᵀ, edge_to_index; augmented_graph=true)
+    ag = AdjacencyGraph(A_and_Aᵀ, edge_to_index; augmented_graph=true, original_size=size(A))
+    postprocessing_minimizes = algo.postprocessing_minimizes
     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; postprocessing_minimizes, forced_colors
         )
         (_row_color, _column_color, _symmetric_to_row, _symmetric_to_column) = remap_colors(
             eltype(ag), _color, maximum(_color), size(A)...
@@ -370,10 +375,11 @@ function _coloring(
     symmetric_pattern::Bool,
 ) where {R}
     A_and_Aᵀ, edge_to_index = bidirectional_pattern(A; symmetric_pattern)
-    ag = AdjacencyGraph(A_and_Aᵀ, edge_to_index; augmented_graph=true)
+    ag = AdjacencyGraph(A_and_Aᵀ, edge_to_index; augmented_graph=true, original_size=size(A))
+    postprocessing_minimizes = algo.postprocessing_minimizes
     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; postprocessing_minimizes)
         (_row_color, _column_color, _symmetric_to_row, _symmetric_to_column) = remap_colors(
             eltype(ag), _color, maximum(_color), size(A)...
         )