Avoid allocation when creating adjacency graph (#104)

Author: gdalle

src/SparseMatrixColorings.jl

@@ -9,9 +9,10 @@ $EXPORTS
 """
 module SparseMatrixColorings
 
-using DataStructures: DisjointSets, find_root!, root_union!, num_groups
 using ADTypes: ADTypes
+using Base.Iterators: Iterators
 using Compat: @compat, stack
+using DataStructures: DisjointSets, find_root!, root_union!, num_groups
 using DocStringExtensions: README, EXPORTS, SIGNATURES, TYPEDEF, TYPEDFIELDS
 using LinearAlgebra:
     Adjoint,

src/coloring.jl

@@ -74,7 +74,7 @@ The vertices are colored in a greedy fashion, following the `order` supplied.
 
 > [_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::Graph, order::AbstractOrder)
+function star_coloring(g::Graph{false}, order::AbstractOrder)
     # Initialize data structures
     nvertices = length(g)
     color = zeros(Int, nvertices)
@@ -267,7 +267,7 @@ The vertices are colored in a greedy fashion, following the `order` supplied.
 
 > [_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::Graph, order::AbstractOrder)
+function acyclic_coloring(g::Graph{false}, order::AbstractOrder)
     # Initialize data structures
     nvertices = length(g)
     nedges = nnz(g) ÷ 2  # symmetric sparse matrix with empty diagonal

src/graph.jl

@@ -1,28 +1,70 @@
 ## Standard graph
 
 """
-    Graph{T}
+    Graph{loops,T}
 
 Undirected graph structure stored in Compressed Sparse Column (CSC) format.
+Note that the underyling CSC may not be square, so the term "graph" is slightly abusive.
+
+The type parameter `loops` must be set to:
+- `true` if coefficients `(i, i)` present in the CSC are counted as edges in the graph (e.g. for each half of a bipartite graph)
+- `false` otherwise (e.g. for an adjacency graph)
 
 # Fields
 
 - `colptr::Vector{T}`: same as for `SparseMatrixCSC`
 - `rowval::Vector{T}`: same as for `SparseMatrixCSC`
 """
-struct Graph{T<:Integer}
+struct Graph{loops,T<:Integer}
     colptr::Vector{T}
     rowval::Vector{T}
 end
 
-Graph(S::SparseMatrixCSC) = Graph(S.colptr, S.rowval)
+function Graph{loops}(S::SparseMatrixCSC{Tv,Ti}) where {loops,Tv,Ti}
+    return Graph{loops,Ti}(S.colptr, S.rowval)
+end
 
 Base.length(g::Graph) = length(g.colptr) - 1
-SparseArrays.nnz(g::Graph) = length(g.rowval)
+
+SparseArrays.nnz(g::Graph{true}) = length(g.rowval)
+
+function SparseArrays.nnz(g::Graph{false})
+    n = 0
+    for j in 1:(length(g.colptr) - 1)
+        for k in g.colptr[j]:(g.colptr[j + 1] - 1)
+            i = g.rowval[k]
+            if i != j
+                n += 1
+            end
+        end
+    end
+    return n
+end
 
 vertices(g::Graph) = 1:length(g)
-neighbors(g::Graph, v::Integer) = view(g.rowval, g.colptr[v]:(g.colptr[v + 1] - 1))
-degree(g::Graph, v::Integer) = length(g.colptr[v]:(g.colptr[v + 1] - 1))
+
+function neighbors(g::Graph{true}, v::Integer)
+    return view(g.rowval, g.colptr[v]:(g.colptr[v + 1] - 1))
+end
+
+function neighbors(g::Graph{false}, v::Integer)
+    neighbors_with_loops = view(g.rowval, g.colptr[v]:(g.colptr[v + 1] - 1))
+    return Iterators.filter(!=(v), neighbors_with_loops)  # TODO: optimize
+end
+
+function degree(g::Graph{true}, v::Integer)
+    return length(g.colptr[v]:(g.colptr[v + 1] - 1))
+end
+
+function degree(g::Graph{false}, v::Integer)
+    d = length(g.colptr[v]:(g.colptr[v + 1] - 1))
+    for k in g.colptr[v]:(g.colptr[v + 1] - 1)
+        if g.rowval[k] == v
+            d -= 1
+        end
+    end
+    return d
+end
 
 maximum_degree(g::Graph) = maximum(Base.Fix1(degree, g), vertices(g))
 minimum_degree(g::Graph) = minimum(Base.Fix1(degree, g), vertices(g))
@@ -42,8 +84,8 @@ A bipartite graph has two "sides", which we number `1` and `2`.
 - `g2::Graph{T}`: contains the neighbors for vertices on side `2`
 """
 struct BipartiteGraph{T<:Integer}
-    g1::Graph{T}
-    g2::Graph{T}
+    g1::Graph{true,T}
+    g2::Graph{true,T}
 end
 
 Base.length(bg::BipartiteGraph, ::Val{1}) = length(bg.g1)
@@ -92,7 +134,7 @@ The adjacency graph of a symmetrix matric `A ∈ ℝ^{n × n}` is `G(A) = (V, E)
 
 > [_What Color Is Your Jacobian? Graph Coloring for Computing Derivatives_](https://epubs.siam.org/doi/10.1137/S0036144504444711), Gebremedhin et al. (2005)
 """
-adjacency_graph(A::SparseMatrixCSC) = Graph(A - Diagonal(A))
+adjacency_graph(A::SparseMatrixCSC) = Graph{false}(A)
 
 """
     bipartite_graph(A::SparseMatrixCSC; symmetric_pattern::Bool)
@@ -112,12 +154,12 @@ When `symmetric_pattern` is `true`, this construction is more efficient.
 > [_What Color Is Your Jacobian? Graph Coloring for Computing Derivatives_](https://epubs.siam.org/doi/10.1137/S0036144504444711), Gebremedhin et al. (2005)
 """
 function bipartite_graph(A::SparseMatrixCSC; symmetric_pattern::Bool=false)
-    g2 = Graph(A)  # columns to rows
+    g2 = Graph{true}(A)  # columns to rows
     if symmetric_pattern
         checksquare(A)  # proxy for checking full symmetry
         g1 = g2
     else
-        g1 = Graph(sparse(transpose(A)))  # rows to columns
+        g1 = Graph{true}(sparse(transpose(A)))  # rows to columns
     end
     return BipartiteGraph(g1, g2)
 end

test/graph.jl

@@ -1,21 +1,40 @@
 using LinearAlgebra
 using SparseArrays
-using SparseMatrixColorings: Graph, adjacency_graph, bipartite_graph, neighbors
+using SparseMatrixColorings: Graph, adjacency_graph, bipartite_graph, degree, neighbors
 using Test
 
 ## Standard graph
 
 @testset "Graph" begin
-    g = Graph(sparse([
+    g = Graph{true}(sparse([
         1 0 1
         1 1 0
         0 0 0
     ]))
 
     @test length(g) == 3
+    @test nnz(g) == 4
     @test neighbors(g, 1) == [1, 2]
     @test neighbors(g, 2) == [2]
     @test neighbors(g, 3) == [1]
+    @test degree(g, 1) == 2
+    @test degree(g, 2) == 1
+    @test degree(g, 3) == 1
+
+    g = Graph{false}(sparse([
+        1 0 1
+        1 1 0
+        0 0 0
+    ]))
+
+    @test length(g) == 3
+    @test nnz(g) == 2
+    @test collect(neighbors(g, 1)) == [2]
+    @test collect(neighbors(g, 2)) == Int[]
+    @test collect(neighbors(g, 3)) == [1]
+    @test degree(g, 1) == 1
+    @test degree(g, 2) == 0
+    @test degree(g, 3) == 1
 end;
 
 ## Bipartite graph (fig 3.1 of "What color is your Jacobian?")
@@ -76,12 +95,12 @@ end;
     B = transpose(A) * A
     g = adjacency_graph(B - Diagonal(B))
     @test length(g) == 8
-    @test neighbors(g, 1) == [6, 7, 8]
-    @test neighbors(g, 2) == [5, 7, 8]
-    @test neighbors(g, 3) == [5, 6, 8]
-    @test neighbors(g, 4) == [5, 6, 7]
-    @test neighbors(g, 5) == [2, 3, 4, 6, 7, 8]
-    @test neighbors(g, 6) == [1, 3, 4, 5, 7, 8]
-    @test neighbors(g, 7) == [1, 2, 4, 5, 6, 8]
-    @test neighbors(g, 8) == [1, 2, 3, 5, 6, 7]
+    @test collect(neighbors(g, 1)) == [6, 7, 8]
+    @test collect(neighbors(g, 2)) == [5, 7, 8]
+    @test collect(neighbors(g, 3)) == [5, 6, 8]
+    @test collect(neighbors(g, 4)) == [5, 6, 7]
+    @test collect(neighbors(g, 5)) == [2, 3, 4, 6, 7, 8]
+    @test collect(neighbors(g, 6)) == [1, 3, 4, 5, 7, 8]
+    @test collect(neighbors(g, 7)) == [1, 2, 4, 5, 6, 8]
+    @test collect(neighbors(g, 8)) == [1, 2, 3, 5, 6, 7]
 end