Propagate arbitrary integer types as indices

Author: gdalle

src/check.jl

@@ -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,15 +76,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 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)
+    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 = Int[]  # one hub for each star, including the trivial ones
     vertices_in_order = vertices(g, 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) 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,33 @@ 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,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(color), 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(color), 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/graph.jl

@@ -24,7 +24,7 @@ end
 SparsityPatternCSC(A::SparseMatrixCSC) = SparsityPatternCSC(A.m, A.n, A.colptr, A.rowval)
 
 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,7 +222,7 @@ 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
@@ -298,7 +298,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,7 +338,7 @@ 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
@@ -425,7 +425,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/order.jl

@@ -84,11 +84,11 @@ function vertices(g::AdjacencyGraph, ::LargestFirst)
     return sort(vertices(g); by=criterion, rev=true)
 end
 
-function vertices(bg::BipartiteGraph, ::Val{side}, ::LargestFirst) where {side}
+function vertices(bg::BipartiteGraph{T}, ::Val{side}, ::LargestFirst) where {T,side}
     other_side = 3 - side
     n = nb_vertices(bg, Val(side))
     visited = falses(n)  # necessary for distance-2 neighborhoods
-    degrees_dist2 = zeros(Int, n)
+    degrees_dist2 = zeros(T, n)
     for v in vertices(bg, Val(side))
         fill!(visited, false)
         for u in neighbors(bg, Val(side), v)
@@ -129,27 +129,27 @@ Instance of [`AbstractOrder`](@ref) which sorts vertices using a dynamically com
 """
 struct DynamicDegreeBasedOrder{degtype,direction} <: AbstractOrder end
 
-struct DegreeBuckets
-    degrees::Vector{Int}
-    bucket_storage::Vector{Int}
-    bucket_low::Vector{Int}
-    bucket_high::Vector{Int}
-    positions::Vector{Int}
+struct DegreeBuckets{T}
+    degrees::Vector{T}
+    bucket_storage::Vector{T}
+    bucket_low::Vector{T}
+    bucket_high::Vector{T}
+    positions::Vector{T}
 end
 
-function DegreeBuckets(degrees::Vector{Int}, dmax)
+function DegreeBuckets(::Type{T}, degrees::Vector{<:Integer}, dmax::Integer) where {T}
     # number of vertices per degree class
-    deg_count = zeros(Int, dmax + 1)
+    deg_count = zeros(T, dmax + 1)
     for d in degrees
         deg_count[d + 1] += 1
     end
     # bucket limits
     bucket_high = cumsum(deg_count)
-    bucket_low = vcat(0, @view(bucket_high[1:(end - 1)]))
+    bucket_low = vcat(zero(T), @view(bucket_high[1:(end - 1)]))
     bucket_low .+= 1
     # assign each vertex to the correct position inside its degree class
-    bucket_storage = similar(degrees, Int)
-    positions = similar(degrees, Int)
+    bucket_storage = similar(degrees, T)
+    positions = similar(degrees, T)
     for v in eachindex(positions, degrees)
         d = degrees[v]
         positions[v] = bucket_high[d + 1] - deg_count[d + 1] + 1
@@ -168,9 +168,9 @@ function degree_increasing(; degtype, direction)
     return increasing
 end
 
-function mark_ordered!(db::DegreeBuckets, v::Integer)
+function mark_ordered!(db::DegreeBuckets{T}, v::Integer) where {T}
     db.degrees[v] = -1
-    db.positions[v] = typemin(Int)
+    db.positions[v] = typemin(T)
     return nothing
 end
 
@@ -248,15 +248,15 @@ function update_bucket!(db::DegreeBuckets, v::Integer; degtype, direction)
 end
 
 function vertices(
-    g::AdjacencyGraph, ::DynamicDegreeBasedOrder{degtype,direction}
-) where {degtype,direction}
+    g::AdjacencyGraph{T}, ::DynamicDegreeBasedOrder{degtype,direction}
+) where {T<:Integer,degtype,direction}
     if degree_increasing(; degtype, direction)
-        degrees = zeros(Int, nb_vertices(g))
+        degrees = zeros(T, nb_vertices(g))
     else
         degrees = [degree(g, v) for v in vertices(g)]
     end
-    db = DegreeBuckets(degrees, maximum_degree(g))
-    π = Int[]
+    db = DegreeBuckets(T, degrees, maximum_degree(g))
+    π = T[]
     sizehint!(π, nb_vertices(g))
     for _ in 1:nb_vertices(g)
         u = pop_next_candidate!(db; direction)
@@ -271,12 +271,12 @@ function vertices(
 end
 
 function vertices(
-    g::BipartiteGraph, ::Val{side}, ::DynamicDegreeBasedOrder{degtype,direction}
-) where {side,degtype,direction}
+    g::BipartiteGraph{T}, ::Val{side}, ::DynamicDegreeBasedOrder{degtype,direction}
+) where {T<:Integer,side,degtype,direction}
     other_side = 3 - side
     # compute dist-2 degrees in an optimized way
     n = nb_vertices(g, Val(side))
-    degrees_dist2 = zeros(Int, n)
+    degrees_dist2 = zeros(T, n)
     dist2_neighbor = falses(n)
     for v in vertices(g, Val(side))
         fill!(dist2_neighbor, false)
@@ -288,13 +288,13 @@ function vertices(
         degrees_dist2[v] = sum(dist2_neighbor)
     end
     if degree_increasing(; degtype, direction)
-        degrees = zeros(Int, n)
+        degrees = zeros(T, n)
     else
         degrees = degrees_dist2
     end
     maxd2 = maximum(degrees_dist2)
-    db = DegreeBuckets(degrees, maxd2)
-    π = Int[]
+    db = DegreeBuckets(T, degrees, maxd2)
+    π = T[]
     sizehint!(π, n)
     visited = falses(n)
     for _ in 1:nb_vertices(g, Val(side))