Implement ConstantColoringAlgorithm (#127)

* Implement ConstantColoringAlgorithm

* Improve docs and add ADTypes

Author: gdalle

docs/src/api.md

@@ -17,6 +17,7 @@ SparseMatrixColorings
 coloring
 ColoringProblem
 GreedyColoringAlgorithm
+ConstantColoringAlgorithm
 ```
 
 ## Result analysis

src/SparseMatrixColorings.jl

@@ -49,12 +49,14 @@ include("coloring.jl")
 include("result.jl")
 include("matrices.jl")
 include("interface.jl")
+include("constant.jl")
 include("decompression.jl")
 include("check.jl")
 include("examples.jl")
 
 export NaturalOrder, RandomOrder, LargestFirst
 export ColoringProblem, GreedyColoringAlgorithm, AbstractColoringResult
+export ConstantColoringAlgorithm
 export coloring
 export column_colors, row_colors
 export column_groups, row_groups

src/constant.jl

@@ -0,0 +1,130 @@
+"""
+    ConstantColoringAlgorithm{partition}  <: ADTypes.AbstractColoringAlgorithm
+
+Coloring algorithm which always returns the same precomputed vector of colors.
+Useful when the optimal coloring of a matrix can be determined a priori due to its specific structure (e.g. banded).
+
+It is passed as an argument to the main function [`coloring`](@ref), but will only work if the associated `problem` has `:nonsymmetric` structure.
+Indeed, for symmetric coloring problems, we need more than just the vector of colors to allow fast decompression.
+
+# Constructors
+
+    ConstantColoringAlgorithm{partition}(matrix_template, color)
+    ConstantColoringAlgorithm(matrix_template, color; partition=:column)
+
+- `partition::Symbol`: either `:row` or `:column`.
+- `matrix_template::AbstractMatrix`: matrix for which the vector of colors was precomputed (the algorithm will only accept matrices of the exact same size).
+- `color::Vector{Int}`: vector of integer colors, one for each row or column (depending on `partition`).
+
+!!! warning
+    The second constructor (based on keyword arguments) is type-unstable.
+
+We do not necessarily verify consistency between the matrix template and the vector of colors, this is the responsibility of the user.
+
+# Example
+
+```jldoctest
+julia> using SparseMatrixColorings, LinearAlgebra
+
+julia> matrix_template = Diagonal(ones(Bool, 5))
+5×5 Diagonal{Bool, Vector{Bool}}:
+ 1  ⋅  ⋅  ⋅  ⋅
+ ⋅  1  ⋅  ⋅  ⋅
+ ⋅  ⋅  1  ⋅  ⋅
+ ⋅  ⋅  ⋅  1  ⋅
+ ⋅  ⋅  ⋅  ⋅  1
+
+julia> color = ones(Int, 5)  # coloring a Diagonal is trivial
+5-element Vector{Int64}:
+ 1
+ 1
+ 1
+ 1
+ 1
+
+julia> problem = ColoringProblem(; structure=:nonsymmetric, partition=:column);
+
+julia> algo = ConstantColoringAlgorithm(matrix_template, color; partition=:column);
+
+julia> result = coloring(similar(matrix_template), problem, algo);
+
+julia> column_colors(result)
+5-element Vector{Int64}:
+ 1
+ 1
+ 1
+ 1
+ 1
+```
+
+# ADTypes coloring interface
+
+`ConstantColoringAlgorithm` is a subtype of [`ADTypes.AbstractColoringAlgorithm`](@extref ADTypes.AbstractColoringAlgorithm), which means the following methods are also applicable (although they will error if the kind of coloring demanded not consistent):
+
+- [`ADTypes.column_coloring`](@extref ADTypes.column_coloring)
+- [`ADTypes.row_coloring`](@extref ADTypes.row_coloring)
+"""
+struct ConstantColoringAlgorithm{
+    partition,M<:AbstractMatrix,R<:AbstractColoringResult{:nonsymmetric,partition,:direct}
+} <: ADTypes.AbstractColoringAlgorithm
+    matrix_template::M
+    color::Vector{Int}
+    result::R
+end
+
+function ConstantColoringAlgorithm{:column}(
+    matrix_template::AbstractMatrix, color::Vector{Int}
+)
+    S = convert(SparseMatrixCSC, matrix_template)
+    result = ColumnColoringResult(S, color)
+    M, R = typeof(matrix_template), typeof(result)
+    return ConstantColoringAlgorithm{:column,M,R}(matrix_template, color, result)
+end
+
+function ConstantColoringAlgorithm{:row}(
+    matrix_template::AbstractMatrix, color::Vector{Int}
+)
+    S = convert(SparseMatrixCSC, matrix_template)
+    result = RowColoringResult(S, color)
+    M, R = typeof(matrix_template), typeof(result)
+    return ConstantColoringAlgorithm{:row,M,R}(matrix_template, color, result)
+end
+
+function ConstantColoringAlgorithm(
+    matrix_template::AbstractMatrix, color::Vector{Int}; partition=:column
+)
+    return ConstantColoringAlgorithm{partition}(matrix_template, color)
+end
+
+function coloring(
+    A::AbstractMatrix,
+    ::ColoringProblem{:nonsymmetric,partition},
+    algo::ConstantColoringAlgorithm{partition};
+    decompression_eltype::Type=Float64,
+    symmetric_pattern::Bool=false,
+) where {partition}
+    @compat (; matrix_template, result) = algo
+    if size(A) != size(matrix_template)
+        throw(
+            DimensionMismatch(
+                "`ConstantColoringAlgorithm` expected matrix of size $(size(matrix_template)) but got matrix of size $(size(A))",
+            ),
+        )
+    else
+        return result
+    end
+end
+
+function ADTypes.column_coloring(
+    A::AbstractMatrix, algo::ConstantColoringAlgorithm{:column}
+)
+    problem = ColoringProblem{:nonsymmetric,:column}()
+    result = coloring(A, problem, algo)
+    return column_colors(result)
+end
+
+function ADTypes.row_coloring(A::AbstractMatrix, algo::ConstantColoringAlgorithm)
+    problem = ColoringProblem{:nonsymmetric,:row}()
+    result = coloring(A, problem, algo)
+    return row_colors(result)
+end

test/constant.jl

@@ -0,0 +1,37 @@
+using ADTypes: ADTypes
+using SparseMatrixColorings
+using Test
+
+matrix_template = ones(100, 200)
+
+@testset "Column coloring" begin
+    problem = ColoringProblem(; structure=:nonsymmetric, partition=:column)
+    color = rand(1:5, size(matrix_template, 2))
+    algo = ConstantColoringAlgorithm(matrix_template, color; partition=:column)
+    wrong_algo = ConstantColoringAlgorithm(matrix_template, color; partition=:row)
+    @test_throws DimensionMismatch coloring(transpose(matrix_template), problem, algo)
+    @test_throws MethodError coloring(matrix_template, problem, wrong_algo)
+    result = coloring(matrix_template, problem, algo)
+    @test column_colors(result) == color
+    @test ADTypes.column_coloring(matrix_template, algo) == color
+    @test_throws MethodError ADTypes.row_coloring(matrix_template, algo)
+end
+
+@testset "Row coloring" begin
+    problem = ColoringProblem(; structure=:nonsymmetric, partition=:row)
+    color = rand(1:5, size(matrix_template, 1))
+    algo = ConstantColoringAlgorithm(matrix_template, color; partition=:row)
+    @test_throws DimensionMismatch coloring(transpose(matrix_template), problem, algo)
+    result = coloring(matrix_template, problem, algo)
+    @test row_colors(result) == color
+    @test ADTypes.row_coloring(matrix_template, algo) == color
+    @test_throws MethodError ADTypes.column_coloring(matrix_template, algo)
+end
+
+@testset "Symmetric coloring" begin
+    wrong_problem = ColoringProblem(; structure=:symmetric, partition=:column)
+    color = rand(1:5, size(matrix_template, 2))
+    algo = ConstantColoringAlgorithm(matrix_template, color; partition=:column)
+    @test_throws MethodError coloring(matrix_template, wrong_problem, algo)
+    @test_throws MethodError ADTypes.symmetric_coloring(matrix_template, algo)
+end

test/runtests.jl

@@ -42,6 +42,9 @@ include("utils.jl")
         @testset "Constructors" begin
             include("constructors.jl")
         end
+        @testset "Constant coloring" begin
+            include("constant.jl")
+        end
     end
     @testset verbose = true "Correctness" begin
         @testset "Small instances" begin