Propagate arbitrary integer types as indices (#222)

* Propagate arbitrary integer types as indices

* Fixes

* Single type

* Coverage

* Fix

* Add tests

* Update src/coloring.jl

Co-authored-by: Alexis Montoison <35051714+amontoison@users.noreply.github.com>

* Fix type

* Update src/graph.jl

Co-authored-by: Alexis Montoison <35051714+amontoison@users.noreply.github.com>

* Fix

* Coverage

---------

Co-authored-by: Alexis Montoison <35051714+amontoison@users.noreply.github.com>

Author: gdalle

src/adtypes.jl

@@ -5,7 +5,7 @@ function coloring(
     kwargs...,
 )
     bg = BipartiteGraph(A)
-    color = convert(Vector{Int}, ADTypes.column_coloring(A, algo))
+    color = convert(Vector{eltype(bg)}, ADTypes.column_coloring(A, algo))
     return ColumnColoringResult(A, bg, color)
 end
 
@@ -16,6 +16,6 @@ function coloring(
     kwargs...,
 )
     bg = BipartiteGraph(A)
-    color = convert(Vector{Int}, ADTypes.row_coloring(A, algo))
+    color = convert(Vector{eltype(bg)}, ADTypes.row_coloring(A, algo))
     return RowColoringResult(A, bg, color)
 end

src/check.jl

@@ -262,8 +262,8 @@ function directly_recoverable_columns(
 end
 
 """
-    valid_dynamic_order(g::AdjacencyGraph, π::AbstractVector{Int}, order::DynamicDegreeBasedOrder)
-    valid_dynamic_order(bg::AdjacencyGraph, ::Val{side}, π::AbstractVector{Int}, order::DynamicDegreeBasedOrder)
+    valid_dynamic_order(g::AdjacencyGraph, π::AbstractVector{<:Integer}, order::DynamicDegreeBasedOrder)
+    valid_dynamic_order(bg::AdjacencyGraph, ::Val{side}, π::AbstractVector{<:Integer}, order::DynamicDegreeBasedOrder)
 
 Check that a permutation `π` corresponds to a valid application of a [`DynamicDegreeBasedOrder`](@ref).
 
@@ -273,7 +273,9 @@ This is done by checking, for each ordered vertex, that its back- or forward-deg
     This function is not coded with efficiency in mind, it is designed for small-scale tests.
 """
 function valid_dynamic_order(
-    g::AdjacencyGraph, π::AbstractVector{Int}, ::DynamicDegreeBasedOrder{degtype,direction}
+    g::AdjacencyGraph,
+    π::AbstractVector{<:Integer},
+    ::DynamicDegreeBasedOrder{degtype,direction},
 ) where {degtype,direction}
     length(π) != nb_vertices(g) && return false
     length(unique(π)) != nb_vertices(g) && return false
@@ -300,7 +302,7 @@ end
 function valid_dynamic_order(
     g::BipartiteGraph,
     ::Val{side},
-    π::AbstractVector{Int},
+    π::AbstractVector{<:Integer},
     ::DynamicDegreeBasedOrder{degtype,direction},
 ) where {side,degtype,direction}
     length(π) != nb_vertices(g, Val(side)) && return false

src/coloring.jl

@@ -17,18 +17,18 @@ The vertices are colored in a greedy fashion, following the `order` supplied.
 > [_What Color Is Your Jacobian? Graph Coloring for Computing Derivatives_](https://epubs.siam.org/doi/10.1137/S0036144504444711), Gebremedhin et al. (2005), Algorithm 3.2
 """
 function partial_distance2_coloring(
-    bg::BipartiteGraph, ::Val{side}, order::AbstractOrder
-) where {side}
-    color = Vector{Int}(undef, nb_vertices(bg, Val(side)))
-    forbidden_colors = Vector{Int}(undef, nb_vertices(bg, Val(side)))
+    bg::BipartiteGraph{T}, ::Val{side}, order::AbstractOrder
+) where {T,side}
+    color = Vector{T}(undef, nb_vertices(bg, Val(side)))
+    forbidden_colors = Vector{T}(undef, nb_vertices(bg, Val(side)))
     vertices_in_order = vertices(bg, Val(side), order)
     partial_distance2_coloring!(color, forbidden_colors, bg, Val(side), vertices_in_order)
     return color
 end
 
 function partial_distance2_coloring!(
-    color::Vector{Int},
-    forbidden_colors::Vector{Int},
+    color::AbstractVector{<:Integer},
+    forbidden_colors::AbstractVector{<:Integer},
     bg::BipartiteGraph,
     ::Val{side},
     vertices_in_order::AbstractVector{<:Integer},
@@ -76,16 +76,18 @@ 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 4.1
 """
-function star_coloring(g::AdjacencyGraph, order::AbstractOrder, postprocessing::Bool)
+function star_coloring(
+    g::AdjacencyGraph{T}, order::AbstractOrder, postprocessing::Bool
+) where {T<:Integer}
     # Initialize data structures
     nv = nb_vertices(g)
     ne = nb_edges(g)
-    color = zeros(Int, nv)
-    forbidden_colors = zeros(Int, nv)
-    first_neighbor = fill((0, 0, 0), nv)  # at first no neighbors have been encountered
-    treated = zeros(Int, nv)
-    star = Vector{Int}(undef, ne)
-    hub = Int[]  # one hub for each star, including the trivial ones
+    color = zeros(T, nv)
+    forbidden_colors = zeros(T, nv)
+    first_neighbor = fill((zero(T), zero(T), zero(T)), nv)  # at first no neighbors have been encountered
+    treated = zeros(T, nv)
+    star = Vector{T}(undef, ne)
+    hub = T[]  # one hub for each star, including the trivial ones
     vertices_in_order = vertices(g, order)
 
     for v in vertices_in_order
@@ -196,11 +198,11 @@ Encode a set of 2-colored stars resulting from the [`star_coloring`](@ref) algor
 
 $TYPEDFIELDS
 """
-struct StarSet
+struct StarSet{T}
     "a mapping from edges (pair of vertices) to their star index"
-    star::Vector{Int}
+    star::Vector{T}
     "a mapping from star indices to their hub (undefined hubs for single-edge stars are the negative value of one of the vertices, picked arbitrarily)"
-    hub::Vector{Int}
+    hub::Vector{T}
 end
 
 """
@@ -226,15 +228,17 @@ 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, order::AbstractOrder, postprocessing::Bool)
+function acyclic_coloring(
+    g::AdjacencyGraph{T}, order::AbstractOrder, postprocessing::Bool
+) where {T<:Integer}
     # Initialize data structures
     nv = nb_vertices(g)
     ne = nb_edges(g)
-    color = zeros(Int, nv)
-    forbidden_colors = zeros(Int, nv)
-    first_neighbor = fill((0, 0, 0), nv)  # at first no neighbors have been encountered
-    first_visit_to_tree = fill((0, 0), ne)
-    forest = Forest{Int}(ne)
+    color = zeros(T, nv)
+    forbidden_colors = zeros(T, nv)
+    first_neighbor = fill((zero(T), zero(T), zero(T)), nv)  # at first no neighbors have been encountered
+    first_visit_to_tree = fill((zero(T), zero(T)), ne)
+    forest = Forest{T}(ne)
     vertices_in_order = vertices(g, order)
 
     for v in vertices_in_order
@@ -367,23 +371,23 @@ Encode a set of 2-colored trees resulting from the [`acyclic_coloring`](@ref) al
 
 $TYPEDFIELDS
 """
-struct TreeSet
-    reverse_bfs_orders::Vector{Vector{Tuple{Int,Int}}}
+struct TreeSet{T}
+    reverse_bfs_orders::Vector{Vector{Tuple{T,T}}}
     is_star::Vector{Bool}
 end
 
-function TreeSet(g::AdjacencyGraph, forest::Forest{Int})
+function TreeSet(g::AdjacencyGraph{T}, forest::Forest{T}) where {T}
     S = pattern(g)
     edge_to_index = edge_indices(g)
     nv = nb_vertices(g)
     nt = forest.num_trees
 
     # dictionary that maps a tree's root to the index of the tree
-    roots = Dict{Int,Int}()
+    roots = Dict{T,T}()
     sizehint!(roots, nt)
 
     # vector of dictionaries where each dictionary stores the neighbors of each vertex in a tree
-    trees = [Dict{Int,Vector{Int}}() for i in 1:nt]
+    trees = [Dict{T,Vector{T}}() for i in 1:nt]
 
     # current number of roots found
     nr = 0
@@ -423,10 +427,10 @@ function TreeSet(g::AdjacencyGraph, forest::Forest{Int})
     end
 
     # degrees is a vector of integers that stores the degree of each vertex in a tree
-    degrees = Vector{Int}(undef, nv)
+    degrees = Vector{T}(undef, nv)
 
     # reverse breadth first (BFS) traversal order for each tree in the forest
-    reverse_bfs_orders = [Tuple{Int,Int}[] for i in 1:nt]
+    reverse_bfs_orders = [Tuple{T,T}[] for i in 1:nt]
 
     # nvmax is the number of vertices of the biggest tree in the forest
     nvmax = 0
@@ -436,7 +440,7 @@ function TreeSet(g::AdjacencyGraph, forest::Forest{Int})
     end
 
     # Create a queue with a fixed size nvmax
-    queue = Vector{Int}(undef, nvmax)
+    queue = Vector{T}(undef, nvmax)
 
     # Specify if each tree in the forest is a star,
     # meaning that one vertex is directly connected to all other vertices in the tree
@@ -519,7 +523,7 @@ function postprocess!(
     color::AbstractVector{<:Integer},
     star_or_tree_set::Union{StarSet,TreeSet},
     g::AdjacencyGraph,
-    offsets::Vector{Int},
+    offsets::AbstractVector{<:Integer},
 )
     S = pattern(g)
     edge_to_index = edge_indices(g)

src/constant.jl

@@ -14,7 +14,7 @@ Indeed, for symmetric coloring problems, we need more than just the vector of co
 
 - `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`).
+- `color::Vector{<:Integer}`: vector of integer colors, one for each row or column (depending on `partition`).
 
 !!! warning
     The second constructor (based on keyword arguments) is type-unstable.
@@ -65,33 +65,36 @@ julia> column_colors(result)
 - [`ADTypes.row_coloring`](@extref ADTypes.row_coloring)
 """
 struct ConstantColoringAlgorithm{
-    partition,M<:AbstractMatrix,R<:AbstractColoringResult{:nonsymmetric,partition,:direct}
+    partition,
+    M<:AbstractMatrix,
+    T<:Integer,
+    R<:AbstractColoringResult{:nonsymmetric,partition,:direct},
 } <: ADTypes.AbstractColoringAlgorithm
     matrix_template::M
-    color::Vector{Int}
+    color::Vector{T}
     result::R
 end
 
 function ConstantColoringAlgorithm{:column}(
-    matrix_template::AbstractMatrix, color::Vector{Int}
+    matrix_template::AbstractMatrix, color::Vector{<:Integer}
 )
     bg = BipartiteGraph(matrix_template)
     result = ColumnColoringResult(matrix_template, bg, color)
-    M, R = typeof(matrix_template), typeof(result)
-    return ConstantColoringAlgorithm{:column,M,R}(matrix_template, color, result)
+    T, M, R = eltype(bg), typeof(matrix_template), typeof(result)
+    return ConstantColoringAlgorithm{:column,M,T,R}(matrix_template, color, result)
 end
 
 function ConstantColoringAlgorithm{:row}(
-    matrix_template::AbstractMatrix, color::Vector{Int}
+    matrix_template::AbstractMatrix, color::Vector{<:Integer}
 )
     bg = BipartiteGraph(matrix_template)
     result = RowColoringResult(matrix_template, bg, color)
-    M, R = typeof(matrix_template), typeof(result)
-    return ConstantColoringAlgorithm{:row,M,R}(matrix_template, color, result)
+    T, M, R = eltype(bg), typeof(matrix_template), typeof(result)
+    return ConstantColoringAlgorithm{:row,M,T,R}(matrix_template, color, result)
 end
 
 function ConstantColoringAlgorithm(
-    matrix_template::AbstractMatrix, color::Vector{Int}; partition=:column
+    matrix_template::AbstractMatrix, color::Vector{<:Integer}; partition::Symbol=:column
 )
     return ConstantColoringAlgorithm{partition}(matrix_template, color)
 end

src/decompression.jl

@@ -677,12 +677,12 @@ function decompress!(
     result::LinearSystemColoringResult,
     uplo::Symbol=:F,
 )
-    (; color, strict_upper_nonzero_inds, T_factorization, strict_upper_nonzeros_A) = result
+    (; color, strict_upper_nonzero_inds, M_factorization, strict_upper_nonzeros_A) = result
     S = result.ag.S
     uplo == :F && check_same_pattern(A, S)
 
     # TODO: for some reason I cannot use ldiv! with a sparse QR
-    strict_upper_nonzeros_A = T_factorization \ vec(B)
+    strict_upper_nonzeros_A = M_factorization \ vec(B)
     fill!(A, zero(eltype(A)))
     for i in axes(A, 1)
         if !iszero(S[i, i])

src/graph.jl

@@ -23,8 +23,9 @@ end
 
 SparsityPatternCSC(A::SparseMatrixCSC) = SparsityPatternCSC(A.m, A.n, A.colptr, A.rowval)
 
+Base.eltype(::SparsityPatternCSC{T}) where {T} = T
 Base.size(S::SparsityPatternCSC) = (S.m, S.n)
-Base.size(S::SparsityPatternCSC, d) = d::Integer <= 2 ? size(S)[d] : 1
+Base.size(S::SparsityPatternCSC, d::Integer) = d::Integer <= 2 ? size(S)[d] : 1
 Base.axes(S::SparsityPatternCSC, d::Integer) = Base.OneTo(size(S, d))
 
 SparseArrays.nnz(S::SparsityPatternCSC) = length(S.rowval)
@@ -222,11 +223,13 @@ The adjacency graph of a symmetric matrix `A ∈ ℝ^{n × n}` is `G(A) = (V, E)
 
 > [_What Color Is Your Jacobian? SparsityPatternCSC Coloring for Computing Derivatives_](https://epubs.siam.org/doi/10.1137/S0036144504444711), Gebremedhin et al. (2005)
 """
-struct AdjacencyGraph{T,has_diagonal}
+struct AdjacencyGraph{T<:Integer,has_diagonal}
     S::SparsityPatternCSC{T}
     edge_to_index::Vector{T}
 end
 
+Base.eltype(::AdjacencyGraph{T}) where {T} = T
+
 function AdjacencyGraph(
     S::SparsityPatternCSC{T},
     edge_to_index::Vector{T}=build_edge_to_index(S);
@@ -298,7 +301,7 @@ function has_neighbor(g::AdjacencyGraph, v::Integer, u::Integer)
     return false
 end
 
-function degree_in_subset(g::AdjacencyGraph, v::Integer, subset::AbstractVector{Int})
+function degree_in_subset(g::AdjacencyGraph, v::Integer, subset::AbstractVector{<:Integer})
     d = 0
     for u in subset
         if has_neighbor(g, v, u)
@@ -338,11 +341,13 @@ When `symmetric_pattern` is `true`, this construction is more efficient.
 
 > [_What Color Is Your Jacobian? SparsityPatternCSC Coloring for Computing Derivatives_](https://epubs.siam.org/doi/10.1137/S0036144504444711), Gebremedhin et al. (2005)
 """
-struct BipartiteGraph{T}
+struct BipartiteGraph{T<:Integer}
     S1::SparsityPatternCSC{T}
     S2::SparsityPatternCSC{T}
 end
 
+Base.eltype(::BipartiteGraph{T}) where {T} = T
+
 function BipartiteGraph(A::AbstractMatrix; symmetric_pattern::Bool=false)
     return BipartiteGraph(SparseMatrixCSC(A); symmetric_pattern)
 end
@@ -425,7 +430,7 @@ function has_neighbor_dist2(
 end
 
 function degree_dist2_in_subset(
-    bg::BipartiteGraph, ::Val{side}, v::Integer, subset::AbstractVector{Int}
+    bg::BipartiteGraph, ::Val{side}, v::Integer, subset::AbstractVector{<:Integer}
 ) where {side}
     d = 0
     for u in subset

src/interface.jl

@@ -306,7 +306,9 @@ function _coloring(
         symmetric_result = StarSetColoringResult(A_and_Aᵀ, ag, color, star_set)
         return BicoloringResult(A, ag, symmetric_result, R)
     else
-        row_color, column_color, _ = remap_colors(color, maximum(color), size(A)...)
+        row_color, column_color, _ = remap_colors(
+            eltype(ag), color, maximum(color), size(A)...
+        )
         return row_color, column_color
     end
 end
@@ -326,7 +328,9 @@ function _coloring(
         symmetric_result = TreeSetColoringResult(A_and_Aᵀ, ag, color, tree_set, R)
         return BicoloringResult(A, ag, symmetric_result, R)
     else
-        row_color, column_color, _ = remap_colors(color, maximum(color), size(A)...)
+        row_color, column_color, _ = remap_colors(
+            eltype(ag), color, maximum(color), size(A)...
+        )
         return row_color, column_color
     end
 end