diff --git a/src/coloring.jl b/src/coloring.jl index 61940043..13f9ddf6 100644 --- a/src/coloring.jl +++ b/src/coloring.jl @@ -82,19 +82,19 @@ function star_coloring(g::AdjacencyGraph, order::AbstractOrder; postprocessing:: ne = nb_edges(g) color = zeros(Int, nv) forbidden_colors = zeros(Int, nv) - first_neighbor = fill((0, 0), nv) # at first no neighbors have been encountered + first_neighbor = fill((0, 0, -1), nv) # at first no neighbors have been encountered treated = zeros(Int, nv) - star = Dict{Tuple{Int,Int},Int}() - sizehint!(star, ne) + star = Vector{Int}(undef, ne) hub = Int[] # one hub for each star, including the trivial ones nb_spokes = Int[] # number of spokes for each star vertices_in_order = vertices(g, order) for v in vertices_in_order - for w in neighbors(g, v) + for (w, e_vw) in neighbors_and_edgeinds(g, v) + w == v && continue iszero(color[w]) && continue forbidden_colors[color[w]] = v - (p, q) = first_neighbor[color[w]] + (p, q, _) = first_neighbor[color[w]] if p == v # Case 1 if treated[q] != v # forbid colors of neighbors of q @@ -103,11 +103,11 @@ function star_coloring(g::AdjacencyGraph, order::AbstractOrder; postprocessing:: # forbid colors of neighbors of w _treat!(treated, forbidden_colors, g, v, w, color) else - first_neighbor[color[w]] = (v, w) - for x in neighbors(g, w) + first_neighbor[color[w]] = (v, w, e_vw) + for (x, e_wx) in neighbors_and_edgeinds(g, w) + x == w && continue (x == v || iszero(color[x])) && continue - wx = _sort(w, x) - if x == hub[star[wx]] # potential Case 2 (which is always false for trivial stars with two vertices, since the associated hub is negative) + if x == hub[star[e_wx]] # potential Case 2 (which is always false for trivial stars with two vertices, since the associated hub is negative) forbidden_colors[color[x]] = v end end @@ -121,7 +121,7 @@ function star_coloring(g::AdjacencyGraph, order::AbstractOrder; postprocessing:: end _update_stars!(star, hub, nb_spokes, g, v, color, first_neighbor) end - star_set = StarSet(star, hub, nb_spokes) + star_set = StarSet(g, star, hub, nb_spokes) if postprocessing # Reuse the vector forbidden_colors to compute offsets during post-processing postprocess!(color, star_set, g, forbidden_colors) @@ -138,32 +138,48 @@ Encode a set of 2-colored stars resulting from the [`star_coloring`](@ref) algor $TYPEDFIELDS """ -struct StarSet +struct StarSet{V<:AbstractVector{Int}} "a mapping from edges (pair of vertices) to their star index" - star::Dict{Tuple{Int,Int},Int} + star::Vector{Int} "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} "a mapping from star indices to the vector of their spokes" - spokes::Vector{Vector{Int}} + spokes::Vector{V} end -function StarSet(star::Dict{Tuple{Int,Int},Int}, hub::Vector{Int}, nb_spokes::Vector{Int}) +function StarSet( + g::AdjacencyGraph, star::Vector{Int}, hub::Vector{Int}, nb_spokes::Vector{Int} +) + (; S, edgeindex) = g # Create a list of spokes for each star, preallocating their sizes based on nb_spokes - spokes = [Vector{Int}(undef, ns) for ns in nb_spokes] + spokes_storage = Vector{Int}(undef, nb_edges(g)) + spokes = Vector{typeof(view(spokes_storage, 1:0))}(undef, length(nb_spokes)) + cum_nb_spokes = 0 + for (s, ns) in enumerate(nb_spokes) + spokes[s] = view(spokes_storage, (cum_nb_spokes + 1):(cum_nb_spokes + ns)) + cum_nb_spokes += ns + end # Reuse nb_spokes as counters to track the current index while filling the spokes fill!(nb_spokes, 0) - for ((i, j), s) in pairs(star) - h = abs(hub[s]) - nb_spokes[s] += 1 - index = nb_spokes[s] - - # Assign the non-hub vertex (spoke) to the correct position in spokes - if i == h - spokes[s][index] = j - elseif j == h - spokes[s][index] = i + rvS = rowvals(S) + for j in axes(S, 2) + for k in nzrange(S, j) + i = rvS[k] + i >= j && continue + e = edgeindex[k] + s = star[e] + h = abs(hub[s]) + nb_spokes[s] += 1 + index = nb_spokes[s] + + # Assign the non-hub vertex (spoke) to the correct position in spokes + if i == h + spokes[s][index] = j + elseif j == h + spokes[s][index] = i + end end end return StarSet(star, hub, spokes) @@ -182,6 +198,7 @@ function _treat!( color::AbstractVector{<:Integer}, ) for x in neighbors(g, w) + x == w && continue iszero(color[x]) && continue forbidden_colors[color[x]] = v end @@ -191,7 +208,7 @@ end function _update_stars!( # modified - star::Dict{<:Tuple,<:Integer}, + star::Vector{Int}, hub::AbstractVector{<:Integer}, nb_spokes::AbstractVector{<:Integer}, # not modified @@ -200,33 +217,32 @@ function _update_stars!( color::AbstractVector{<:Integer}, first_neighbor::AbstractVector{<:Tuple}, ) - for w in neighbors(g, v) + for (w, e_vw) in neighbors_and_edgeinds(g, v) + w == v && continue iszero(color[w]) && continue - vw = _sort(v, w) x_exists = false - for x in neighbors(g, w) + for (x, e_wx) in neighbors_and_edgeinds(g, w) + x == w && continue if x != v && color[x] == color[v] # vw, wx ∈ E - wx = _sort(w, x) - star_wx = star[wx] + star_wx = star[e_wx] hub[star_wx] = w # this may already be true nb_spokes[star_wx] += 1 - star[vw] = star_wx + star[e_vw] = star_wx x_exists = true break end end if !x_exists - (p, q) = first_neighbor[color[w]] + (p, q, e_pq) = first_neighbor[color[w]] if p == v && q != w # vw, vq ∈ E and color[w] = color[q] - vq = _sort(v, q) - star_vq = star[vq] + star_vq = star[e_pq] hub[star_vq] = v # this may already be true nb_spokes[star_vq] += 1 - star[vw] = star_vq + star[e_vw] = star_vq else # vw forms a new star push!(hub, -max(v, w)) # star is trivial (composed only of two vertices) so we set the hub to a negative value, but it allows us to choose one of the two vertices push!(nb_spokes, 1) - star[vw] = length(hub) + star[e_vw] = length(hub) end end end @@ -249,7 +265,7 @@ This function corresponds to algorithm `DirectRecover2` in the paper. > [_Efficient Computation of Sparse Hessians Using Coloring and Automatic Differentiation_](https://pubsonline.informs.org/doi/abs/10.1287/ijoc.1080.0286), Gebremedhin et al. (2009), Figure 3 """ function symmetric_coefficient( - i::Integer, j::Integer, color::AbstractVector{<:Integer}, star_set::StarSet + i::Integer, j::Integer, e::Integer, color::AbstractVector{<:Integer}, star_set::StarSet ) (; star, hub) = star_set if i == j @@ -260,7 +276,7 @@ function symmetric_coefficient( # star only contains one triangle i, j = j, i end - star_id = star[i, j] + star_id = star[e] h = abs(hub[star_id]) if h == j # i is the spoke @@ -307,12 +323,15 @@ function acyclic_coloring(g::AdjacencyGraph, order::AbstractOrder; postprocessin for v in vertices_in_order for w in neighbors(g, v) + w == v && continue iszero(color[w]) && continue forbidden_colors[color[w]] = v end for w in neighbors(g, v) + w == v && continue iszero(color[w]) && continue for x in neighbors(g, w) + x == w && continue iszero(color[x]) && continue if forbidden_colors[color[x]] != v _prevent_cycle!( @@ -328,12 +347,15 @@ function acyclic_coloring(g::AdjacencyGraph, order::AbstractOrder; postprocessin end end for w in neighbors(g, v) # grow two-colored stars around the vertex v + w == v && continue iszero(color[w]) && continue _grow_star!(v, w, color, first_neighbor, forest) end for w in neighbors(g, v) + w == v && continue iszero(color[w]) && continue for x in neighbors(g, w) + x == w && continue (x == v || iszero(color[x])) && continue if color[x] == color[v] _merge_trees!(v, w, x, forest) # merge trees T₁ ∋ vw and T₂ ∋ wx if T₁ != T₂ diff --git a/src/graph.jl b/src/graph.jl index 2cd9c7fd..6ff09c95 100644 --- a/src/graph.jl +++ b/src/graph.jl @@ -182,6 +182,7 @@ The adjacency graph of a symmetric matrix `A ∈ ℝ^{n × n}` is `G(A) = (V, E) # Fields - `S::SparsityPatternCSC{T}`: Underlying sparsity pattern, whose diagonal is empty whenever `has_diagonal` is `false` +- `edgeindex::Vector{T}`: a vector mapping each nonzero of `S` to a unique edge index (ignoring diagonal and accounting for symmetry, so that `(i, j)` and `(j, i)` get the same index) # References @@ -189,10 +190,36 @@ The adjacency graph of a symmetric matrix `A ∈ ℝ^{n × n}` is `G(A) = (V, E) """ struct AdjacencyGraph{T,has_diagonal} S::SparsityPatternCSC{T} + edgeindex::Vector{T} +end + +function build_edgeindex(S::SparsityPatternCSC{T}) where {T} + offsets = zeros(T, size(S, 1)) + edgeindex = Vector{T}(undef, nnz(S)) + counter = zero(T) + rvS = rowvals(S) + for j in axes(S, 2) + for k in nzrange(S, j) + i = rvS[k] + if i > j + counter += 1 + # index lower triangle + edgeindex[k] = counter + # index upper triangle + k2 = S.colptr[i] + offsets[i] + edgeindex[k2] = counter + offsets[i] += 1 + elseif i == j + # this should never be used, make sure it errors + edgeindex[k] = -1 + end + end + end + return edgeindex end function AdjacencyGraph(S::SparsityPatternCSC{T}; has_diagonal::Bool=true) where {T} - return AdjacencyGraph{T,has_diagonal}(S) + return AdjacencyGraph{T,has_diagonal}(S, build_edgeindex(S)) end function AdjacencyGraph(A::SparseMatrixCSC; has_diagonal::Bool=true) @@ -204,21 +231,23 @@ function AdjacencyGraph(A::AbstractMatrix; has_diagonal::Bool=true) end pattern(g::AdjacencyGraph) = g.S +edgeindex(g::AdjacencyGraph) = g.edgeindex nb_vertices(g::AdjacencyGraph) = pattern(g).n vertices(g::AdjacencyGraph) = 1:nb_vertices(g) has_diagonal(::AdjacencyGraph{T,hl}) where {T,hl} = hl -function neighbors(g::AdjacencyGraph{T,true}, v::Integer) where {T} +function neighbors(g::AdjacencyGraph{T}, v::Integer) where {T} S = pattern(g) neighbors_v = view(rowvals(S), nzrange(S, v)) - return Iterators.filter(!=(v), neighbors_v) # TODO: optimize + return neighbors_v # includes diagonal end -function neighbors(g::AdjacencyGraph{T,false}, v::Integer) where {T} +function neighbors_and_edgeinds(g::AdjacencyGraph{T}, v::Integer) where {T} S = pattern(g) neighbors_v = view(rowvals(S), nzrange(S, v)) - return neighbors_v + einds = view(edgeindex(g), nzrange(S, v)) + return zip(neighbors_v, einds) # includes diagonal end function degree(g::AdjacencyGraph, v::Integer) @@ -235,6 +264,7 @@ function nb_edges(g::AdjacencyGraph) ne = 0 for v in vertices(g) for u in neighbors(g, v) + u == v && continue ne += 1 end end @@ -246,6 +276,7 @@ minimum_degree(g::AdjacencyGraph) = minimum(Base.Fix1(degree, g), vertices(g)) function has_neighbor(g::AdjacencyGraph, v::Integer, u::Integer) for w in neighbors(g, v) + w == v && continue if w == u return true end diff --git a/src/order.jl b/src/order.jl index 319279fb..7eedaef7 100644 --- a/src/order.jl +++ b/src/order.jl @@ -209,6 +209,7 @@ function vertices( u = pop_next_candidate!(db; direction) direction == :low2high ? push!(π, u) : pushfirst!(π, u) for v in neighbors(g, u) + v == u && continue already_ordered(db, v) && continue update_bucket!(db, v; degtype, direction) end diff --git a/src/result.jl b/src/result.jl index 8bbefddf..9b829369 100644 --- a/src/result.jl +++ b/src/result.jl @@ -228,20 +228,20 @@ $TYPEDFIELDS - [`AbstractColoringResult`](@ref) """ -struct StarSetColoringResult{M<:AbstractMatrix,G<:AdjacencyGraph,V} <: +struct StarSetColoringResult{M<:AbstractMatrix,G<:AdjacencyGraph,V1,V2} <: AbstractColoringResult{:symmetric,:column,:direct} A::M ag::G color::Vector{Int} - group::V - star_set::StarSet + group::V1 + star_set::StarSet{V2} compressed_indices::Vector{Int} end function StarSetColoringResult( A::AbstractMatrix, ag::AdjacencyGraph, color::Vector{Int}, star_set::StarSet ) - S = ag.S + (; S, edgeindex) = ag group = group_by_color(color) n = size(S, 1) rv = rowvals(S) @@ -249,7 +249,8 @@ function StarSetColoringResult( for j in axes(S, 2) for k in nzrange(S, j) i = rv[k] - l, c = symmetric_coefficient(i, j, color, star_set) + e = edgeindex[k] + l, c = symmetric_coefficient(i, j, e, color, star_set) # A[i, j] = B[l, c] compressed_indices[k] = (c - 1) * n + l end diff --git a/test/order.jl b/test/order.jl index 347e4b28..efd030ed 100644 --- a/test/order.jl +++ b/test/order.jl @@ -35,7 +35,7 @@ rng = StableRNG(63) end; @testset "RandomOrder" begin - A = sprand(rng, Bool, 10, 10, 0.5) + A = sparse(Symmetric(sprand(rng, Bool, 10, 10, 0.5))) ag = AdjacencyGraph(A) @test sort(vertices(ag, RandomOrder(rng))) == 1:10 @test sort(vertices(ag, RandomOrder())) == 1:10 @@ -63,7 +63,7 @@ end; @testset "LargestFirst" begin A = sparse([ 0 1 0 - 1 0 0 + 1 0 1 0 1 0 ]) ag = AdjacencyGraph(A)