Make type-stable constructors part of the API (#80)

Author: gdalle

docs/src/api.md

@@ -5,12 +5,12 @@ CollapsedDocStrings = true
 CurrentModule = SparseMatrixColorings
 ```
 
+The docstrings on this page define the public API of the package.
+
 ```@docs
 SparseMatrixColorings
 ```
 
-The docstrings on this page define the public API of the package.
-
 ## Main function
 
 ```@docs
@@ -39,8 +39,6 @@ decompress!
 
 ## Orders
 
-These symbols are not exported but they are still part of the public API.
-
 ```@docs
 AbstractOrder
 NaturalOrder

docs/src/dev.md

@@ -43,7 +43,6 @@ SparseMatrixColorings.LinearSystemColoringResult
 ## Testing
 
 ```@docs
-SparseMatrixColorings.same_sparsity_pattern
 SparseMatrixColorings.directly_recoverable_columns
 SparseMatrixColorings.symmetrically_orthogonal_columns
 SparseMatrixColorings.structurally_orthogonal_columns
@@ -54,6 +53,7 @@ SparseMatrixColorings.structurally_orthogonal_columns
 ```@docs
 SparseMatrixColorings.respectful_similar
 SparseMatrixColorings.matrix_versions
+SparseMatrixColorings.same_pattern
 ```
 
 ## Examples

src/SparseMatrixColorings.jl

@@ -49,12 +49,11 @@ include("decompression.jl")
 include("check.jl")
 include("examples.jl")
 
+export NaturalOrder, RandomOrder, LargestFirst
 export ColoringProblem, GreedyColoringAlgorithm, AbstractColoringResult
 export coloring
 export column_colors, row_colors
 export column_groups, row_groups
 export compress, decompress, decompress!
 
-@compat public NaturalOrder, RandomOrder, LargestFirst
-
 end

src/decompression.jl

@@ -181,27 +181,16 @@ true
 - [`ColoringProblem`](@ref)
 - [`AbstractColoringResult`](@ref)
 """
-function decompress!(
-    A::AbstractMatrix{R},
-    B::AbstractMatrix{R},
-    result::AbstractColoringResult{structure,partition,decompression},
-) where {R<:Real,structure,partition,decompression}
-    # common checks
-    S = get_matrix(result)
-    structure == :symmetric && checksquare(A)
-    if !same_sparsity_pattern(A, S)
-        throw(DimensionMismatch("`A` and `S` must have the same sparsity pattern."))
-    end
-    return decompress_aux!(A, B, result)
-end
+function decompress! end
 
 ## NonSymmetricColoringResult
 
-function decompress_aux!(
+function decompress!(
     A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::NonSymmetricColoringResult{:column}
 ) where {R<:Real}
-    A .= zero(R)
     S = get_matrix(result)
+    check_same_pattern(A, S)
+    A .= zero(R)
     color = column_colors(result)
     rvS = rowvals(S)
     for j in axes(S, 2)
@@ -214,11 +203,12 @@ function decompress_aux!(
     return A
 end
 
-function decompress_aux!(
+function decompress!(
     A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::NonSymmetricColoringResult{:row}
 ) where {R<:Real}
-    A .= zero(R)
     S = get_matrix(result)
+    check_same_pattern(A, S)
+    A .= zero(R)
     color = row_colors(result)
     rvS = rowvals(S)
     for j in axes(S, 2)
@@ -231,9 +221,11 @@ function decompress_aux!(
     return A
 end
 
-function decompress_aux!(
+function decompress!(
     A::SparseMatrixCSC{R}, B::AbstractMatrix{R}, result::NonSymmetricColoringResult{:column}
 ) where {R<:Real}
+    S = get_matrix(result)
+    check_same_pattern(A, S)
     nzA = nonzeros(A)
     ind = result.compressed_indices
     for i in eachindex(nzA, ind)
@@ -242,9 +234,11 @@ function decompress_aux!(
     return A
 end
 
-function decompress_aux!(
+function decompress!(
     A::SparseMatrixCSC{R}, B::AbstractMatrix{R}, result::NonSymmetricColoringResult{:row}
 ) where {R<:Real}
+    S = get_matrix(result)
+    check_same_pattern(A, S)
     nzA = nonzeros(A)
     ind = result.compressed_indices
     for i in eachindex(nzA, ind)
@@ -255,11 +249,12 @@ end
 
 ## StarSetColoringResult
 
-function decompress_aux!(
+function decompress!(
     A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::StarSetColoringResult
 ) where {R<:Real}
-    A .= zero(R)
     S = get_matrix(result)
+    check_same_pattern(A, S)
+    A .= zero(R)
     color = column_colors(result)
     rvS = rowvals(S)
     for j in axes(S, 2)
@@ -272,9 +267,11 @@ function decompress_aux!(
     return A
 end
 
-function decompress_aux!(
+function decompress!(
     A::SparseMatrixCSC{R}, B::AbstractMatrix{R}, result::StarSetColoringResult
 ) where {R<:Real}
+    S = get_matrix(result)
+    check_same_pattern(A, S)
     nzA = nonzeros(A)
     ind = result.compressed_indices
     for i in eachindex(nzA, ind)
@@ -287,11 +284,12 @@ end
 
 # TODO: add method for A::SparseMatrixCSC
 
-function decompress_aux!(
+function decompress!(
     A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::TreeSetColoringResult
 ) where {R<:Real}
-    A .= zero(R)
     S = get_matrix(result)
+    check_same_pattern(A, S)
+    A .= zero(R)
     color = column_colors(result)
     @compat (; vertices_by_tree, reverse_bfs_orders, buffer) = result
 
@@ -326,10 +324,11 @@ end
 
 ## MatrixInverseColoringResult
 
-function decompress_aux!(
+function decompress!(
     A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::LinearSystemColoringResult
 ) where {R<:Real}
     S = get_matrix(result)
+    check_same_pattern(A, S)
     color = column_colors(result)
     @compat (; strict_upper_nonzero_inds, T_factorization, strict_upper_nonzeros_A) = result
 

src/interface.jl

@@ -1,39 +1,62 @@
+function check_valid_problem(structure::Symbol, partition::Symbol)
+    valid = (
+        (structure == :nonsymmetric && partition in (:column, :row)) ||
+        (structure == :symmetric && partition == :column)
+    )
+    if !valid
+        throw(
+            ArgumentError(
+                "The combination `($(repr(structure)), $(repr(partition)))` is not supported by `ColoringProblem`.",
+            ),
+        )
+    end
+end
+
+function check_valid_algorithm(decompression::Symbol)
+    valid = decompression in (:direct, :substitution)
+    if !valid
+        throw(
+            ArgumentError(
+                "The setting `decompression=$(repr(decompression))` is not supported by `GreedyColoringAlgorithm`.",
+            ),
+        )
+    end
+end
+
 """
     ColoringProblem{structure,partition}
 
 Selector type for the coloring problem to solve, enabling multiple dispatch.
 
 It is passed as an argument to the main function [`coloring`](@ref).
 
-# Constructor
+# Constructors
 
-    ColoringProblem(; structure::Symbol=:nonsymmetric, partition::Symbol=:column)
+    ColoringProblem{structure,partition}()
+    ColoringProblem(; structure=:nonsymmetric, partition=:column)
 
 - `structure::Symbol`: either `:nonsymmetric` or `:symmetric`
 - `partition::Symbol`: either `:column`, `:row` or `:bidirectional`
 
+!!! warning
+    The second constructor (based on keyword arguments) is type-unstable.
+
 #  Link to automatic differentiation
 
 Matrix coloring is often used in automatic differentiation, and here is the translation guide:
 
-| matrix   | mode                 | `structure`     | `partition`      |
-| -------- | -------------------- | --------------- | -----------------|
-| Jacobian | forward              | `:nonsymmetric` | `:column`        |
-| Jacobian | reverse              | `:nonsymmetric` | `:row`           |
-| Jacobian | forward + reverse    | `:nonsymmetric` | `:bidirectional` |
-| Hessian  | any                  | `:symmetric`    | `:column`        |
-
-!!! warning
-    With a `:symmetric` structure, you have to use a `:column` partition.
-
-!!! warning
-    At the moment, `:bidirectional` partitions are not implemented.
+| matrix   | mode    | `structure`     | `partition`      | implemented |
+| -------- | ------- | --------------- | ---------------- | ----------- |
+| Jacobian | forward | `:nonsymmetric` | `:column`        | yes         |
+| Jacobian | reverse | `:nonsymmetric` | `:row`           | yes         |
+| Jacobian | mixed   | `:nonsymmetric` | `:bidirectional` | no          |
+| Hessian  | -       | `:symmetric`    | `:column`        | yes         |
+| Hessian  | -       | `:symmetric`    | `:row`           | no          |
 """
 struct ColoringProblem{structure,partition} end
 
 function ColoringProblem(; structure::Symbol=:nonsymmetric, partition::Symbol=:column)
-    @assert structure in (:nonsymmetric, :symmetric)
-    @assert partition in (:column, :row, :bidirectional)
+    check_valid_problem(structure, partition)
     return ColoringProblem{structure,partition}()
 end
 
@@ -44,16 +67,17 @@ Greedy coloring algorithm for sparse matrices which colors columns or rows one a
 
 It is passed as an argument to the main function [`coloring`](@ref).
 
-# Constructor
+# Constructors
 
-    GreedyColoringAlgorithm(
-        order::AbstractOrder=NaturalOrder();
-        decompression::Symbol=:direct
-    )
+    GreedyColoringAlgorithm{decompression}(order=NaturalOrder())
+    GreedyColoringAlgorithm(order=NaturalOrder(); decompression=:direct)
 
 - `order::AbstractOrder`: the order in which the columns or rows are colored, which can impact the number of 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
+    The second constructor (based on keyword arguments) is type-unstable.
+
 # ADTypes coloring interface
 
 `GreedyColoringAlgorithm` is a subtype of [`ADTypes.AbstractColoringAlgorithm`](@extref ADTypes.AbstractColoringAlgorithm), which means the following methods are also applicable:
@@ -74,10 +98,17 @@ struct GreedyColoringAlgorithm{decompression,O<:AbstractOrder} <:
     order::O
 end
 
+function GreedyColoringAlgorithm{decompression}(
+    order::AbstractOrder=NaturalOrder()
+) where {decompression}
+    check_valid_algorithm(decompression)
+    return GreedyColoringAlgorithm{decompression,typeof(order)}(order)
+end
+
 function GreedyColoringAlgorithm(
     order::AbstractOrder=NaturalOrder(); decompression::Symbol=:direct
 )
-    @assert decompression in (:direct, :substitution)
+    check_valid_algorithm(decompression)
     return GreedyColoringAlgorithm{decompression,typeof(order)}(order)
 end
 
@@ -106,9 +137,9 @@ julia> S = sparse([
            0 1 1 0 0 0
        ]);
 
-julia> problem = ColoringProblem(structure=:nonsymmetric, partition=:column);
+julia> problem = ColoringProblem(; structure=:nonsymmetric, partition=:column);
 
-julia> algo = GreedyColoringAlgorithm();
+julia> algo = GreedyColoringAlgorithm(; decompression=:direct);
 
 julia> result = coloring(S, problem, algo);
 

src/matrices.jl

@@ -46,24 +46,23 @@ function respectful_similar(A::Adjoint, ::Type{T}) where {T}
 end
 
 """
-    same_sparsity_pattern(A::AbstractMatrix, B::AbstractMatrix)
+    same_pattern(A::AbstractMatrix, B::AbstractMatrix)
 
 Perform a partial equality check on the sparsity patterns of `A` and `B`:
 
 - if the return is `true`, they might have the same sparsity pattern but we're not sure
 - if the return is `false`, they definitely don't have the same sparsity pattern
 """
-function same_sparsity_pattern(A::AbstractMatrix, B::AbstractMatrix)
+function same_pattern(A::AbstractMatrix, B::AbstractMatrix)
     return size(A) == size(B)
 end
 
-function same_sparsity_pattern(A::SparseMatrixCSC, B::SparseMatrixCSC)
+function same_pattern(A::SparseMatrixCSC, B::SparseMatrixCSC)
     return size(A) == size(B) && nnz(A) == nnz(B)
 end
 
-function same_sparsity_pattern(
-    A::TransposeOrAdjoint{<:Any,<:SparseMatrixCSC},
-    B::TransposeOrAdjoint{<:Any,<:SparseMatrixCSC},
-)
-    return same_sparsity_pattern(parent(A), parent(B))
+function check_same_pattern(A::AbstractMatrix, S::AbstractMatrix)
+    if !same_pattern(A, S)
+        throw(DimensionMismatch("`A` and `S` must have the same sparsity pattern."))
+    end
 end

test/constructors.jl

@@ -0,0 +1,15 @@
+using SparseMatrixColorings
+using Test
+
+@test ColoringProblem{:nonsymmetric,:column}() == ColoringProblem()
+@test ColoringProblem{:symmetric,:column}() ==
+    ColoringProblem(; structure=:symmetric, partition=:column)
+
+@test_throws ArgumentError ColoringProblem(; structure=:weird, partition=:column)
+@test_throws ArgumentError ColoringProblem(; structure=:symmetric, partition=:row)
+
+@test GreedyColoringAlgorithm{:direct}() == GreedyColoringAlgorithm()
+@test GreedyColoringAlgorithm{:substitution}() ==
+    GreedyColoringAlgorithm(; decompression=:substitution)
+
+@test_throws ArgumentError GreedyColoringAlgorithm(decompression=:weird)

test/matrices.jl

@@ -1,5 +1,7 @@
+using LinearAlgebra
 using SparseArrays
-using SparseMatrixColorings: matrix_versions, respectful_similar, same_sparsity_pattern
+using SparseMatrixColorings:
+    check_same_pattern, matrix_versions, respectful_similar, same_pattern
 using StableRNGs
 using Test
 
@@ -30,3 +32,21 @@ same_view(::Adjoint, ::Adjoint) = true
         size(B) == size(A) && same_view(A, B)
     end
 end
+
+@testset "Sparsity pattern" begin
+    S = sparse([
+        0 1 1
+        0 1 0
+        1 1 0
+    ])
+
+    A1 = copy(S)
+    A2 = copy(S)
+    A2[1, 1] = 1
+
+    @test same_pattern(A1, S)
+    @test !same_pattern(A2, S)
+    @test same_pattern(Matrix(A2), S)
+
+    @test_throws DimensionMismatch check_same_pattern(A2, S)
+end