fixes to degree

Author: CarloLucibello

src/GNNGraphs/convert.jl

@@ -63,8 +63,9 @@ end
 
 to_dense(A::AbstractSparseMatrix, x...; kws...) = to_dense(collect(A), x...; kws...)
 
-function to_dense(A::ADJMAT_T, T::DataType=eltype(A); dir=:out, num_nodes=nothing)
+function to_dense(A::ADJMAT_T, T=nothing; dir=:out, num_nodes=nothing)
     @assert dir ∈ [:out, :in]
+    T = T === nothing ? eltype(A) : T
     num_nodes = size(A, 1)
     @assert num_nodes == size(A, 2)
     # @assert all(x -> (x == 1) || (x == 0), A)
@@ -78,11 +79,12 @@ function to_dense(A::ADJMAT_T, T::DataType=eltype(A); dir=:out, num_nodes=nothin
     return A, num_nodes, num_edges
 end
 
-function to_dense(adj_list::ADJLIST_T, T::DataType=Int; dir=:out, num_nodes=nothing)
+function to_dense(adj_list::ADJLIST_T, T=nothing; dir=:out, num_nodes=nothing)
     @assert dir ∈ [:out, :in]
     num_nodes = length(adj_list)
     num_edges = sum(length.(adj_list))
     @assert num_nodes > 0
+    T = T === nothing ? eltype(adj_list[1]) : T
     A = similar(adj_list[1], T, (num_nodes, num_nodes))
     if dir == :out
         for (i, neigs) in enumerate(adj_list)
@@ -96,13 +98,13 @@ function to_dense(adj_list::ADJLIST_T, T::DataType=Int; dir=:out, num_nodes=noth
     A, num_nodes, num_edges
 end
 
-function to_dense(coo::COO_T, T::DataType=Nothing; dir=:out, num_nodes=nothing)
+function to_dense(coo::COO_T, T=nothing; dir=:out, num_nodes=nothing)
     # `dir` will be ignored since the input `coo` is always in source -> target format.
     # The output will always be a adjmat in :out format (e.g. A[i,j] denotes from i to j)
     s, t, val = coo
     n = isnothing(num_nodes) ? max(maximum(s), maximum(t)) : num_nodes
     val = isnothing(val) ? eltype(s)(1) : val
-    T = T == Nothing ? eltype(val) : T
+    T = T === nothing ? eltype(val) : T
     A = fill!(similar(s, T, (n, n)), 0)
     v = vec(A)
     idxs = s .+ n .* (t .- 1) 
@@ -113,10 +115,11 @@ end
 
 ### SPARSE #############
 
-function to_sparse(A::ADJMAT_T, T::DataType=eltype(A); dir=:out, num_nodes=nothing)
+function to_sparse(A::ADJMAT_T, T=nothing; dir=:out, num_nodes=nothing)
     @assert dir ∈ [:out, :in]
     num_nodes = size(A, 1)
     @assert num_nodes == size(A, 2)
+    T = T === nothing ? eltype(A) : T
     num_edges = A isa AbstractSparseMatrix ? nnz(A) : count(!=(0), A)
     if dir == :in
         A = A'
@@ -130,13 +133,14 @@ function to_sparse(A::ADJMAT_T, T::DataType=eltype(A); dir=:out, num_nodes=nothi
     return A, num_nodes, num_edges
 end
 
-function to_sparse(adj_list::ADJLIST_T, T::DataType=Int; dir=:out, num_nodes=nothing)
+function to_sparse(adj_list::ADJLIST_T, T=nothing; dir=:out, num_nodes=nothing)
     coo, num_nodes, num_edges = to_coo(adj_list; dir, num_nodes)
     return to_sparse(coo; dir, num_nodes)
 end
 
-function to_sparse(coo::COO_T, T::DataType=Int; dir=:out, num_nodes=nothing)
+function to_sparse(coo::COO_T, T=nothing; dir=:out, num_nodes=nothing)
     s, t, eweight  = coo
+    T = T === nothing ? eltype(s) : T
     eweight = isnothing(eweight) ? fill!(similar(s, T), 1) : eweight
     num_nodes = isnothing(num_nodes) ? max(maximum(s), maximum(t)) : num_nodes 
     A = sparse(s, t, eweight, num_nodes, num_nodes)

src/GNNGraphs/query.jl

@@ -20,6 +20,23 @@ get_edge_weight(g::GNNGraph{<:ADJMAT_T}) = to_coo(g.graph, num_nodes=g.num_nodes
 Graphs.edges(g::GNNGraph) = zip(edge_index(g)...)
 
 Graphs.edgetype(g::GNNGraph) = Tuple{Int, Int}
+nodetype(g::GNNGraph) = Base.eltype(g)
+
+"""
+    nodetype(g::GNNGraph)
+
+Type of nodes in `g`,
+an integer type like `Int`, `Int32`, `Uint16`, ....
+"""
+function nodetype(g::GNNGraph{<:COO_T}, T=nothing)
+    s, t = edge_index(g)
+    return eltype(s)
+end
+
+function nodetype(g::GNNGraph{<:ADJMAT_T}, T=nothing)
+    T !== nothing && return T
+    return eltype(g.graph)
+end
 
 function Graphs.has_edge(g::GNNGraph{<:COO_T}, i::Integer, j::Integer)
     s, t = edge_index(g)
@@ -77,7 +94,7 @@ function adjacency_list(g::GNNGraph; dir=:out)
     return [fneighs(g, i) for i in 1:g.num_nodes]
 end
 
-function Graphs.adjacency_matrix(g::GNNGraph{<:COO_T}, T::DataType=Int; dir=:out)
+function Graphs.adjacency_matrix(g::GNNGraph{<:COO_T}, T::DataType=nodetype(g); dir=:out)
     if g.graph[1] isa CuVector
         # TODO revisit after https://github.com/JuliaGPU/CUDA.jl/pull/1152
         A, n, m = to_dense(g.graph, T, num_nodes=g.num_nodes)
@@ -88,7 +105,7 @@ function Graphs.adjacency_matrix(g::GNNGraph{<:COO_T}, T::DataType=Int; dir=:out
     return dir == :out ? A : A'
 end
 
-function Graphs.adjacency_matrix(g::GNNGraph{<:ADJMAT_T}, T::DataType=eltype(g.graph); dir=:out)
+function Graphs.adjacency_matrix(g::GNNGraph{<:ADJMAT_T}, T::DataType=nodetype(g); dir=:out)
     @assert dir ∈ [:in, :out]
     A = g.graph
     A = T != eltype(A) ? T.(A) : A
@@ -125,27 +142,37 @@ function Graphs.degree(g::GNNGraph{<:COO_T}, T=nothing; dir=:out, edge_weight=tr
     return degs 
 end
 
-function Graphs.degree(g::GNNGraph{<:ADJMAT_T}, T=Int; dir=:out, edge_weight=true)
+function Graphs.degree(g::GNNGraph{<:ADJMAT_T}, T=nothing; dir=:out, edge_weight=true)
     @assert !(edge_weight isa AbstractArray) "passing the edge weights is not support by adjacency matrix representations" 
     @assert dir ∈ (:in, :out, :both)
-    A = adjacency_matrix(g, T)
+    if T === nothing
+        Nt = nodetype(g)
+        if ((edge_weight === false) || (edge_weight === nothing)) && !(Nt <: Integer) 
+            T = Nt == Float32 ? Int32 : 
+                Nt == Float16 ? Int16 : Int
+        else
+            T = Nt
+        end
+    end
+    A = adjacency_matrix(g)
     if (edge_weight === false) || (edge_weight === nothing)
-        A = map(>(0), A)
+        A = map(x -> x > 0 ? T(1) : T(0), A)
     end
+    A = eltype(A) != T ? T.(A) : A
     return dir == :out ? vec(sum(A, dims=2)) : 
            dir == :in  ? vec(sum(A, dims=1)) :
                   vec(sum(A, dims=1)) .+ vec(sum(A, dims=2)) 
 end
 
-function Graphs.laplacian_matrix(g::GNNGraph, T::DataType=Int; dir::Symbol=:out)
+function Graphs.laplacian_matrix(g::GNNGraph, T=nothing; dir::Symbol=:out)
     A = adjacency_matrix(g, T; dir=dir)
     D = Diagonal(vec(sum(A; dims=2)))
     return D - A
 end
 
 
 """
-    normalized_laplacian(g, T=Float32; add_self_loops=false, dir=:out)
+    normalized_laplacian(g, T=nothing; add_self_loops=false, dir=:out)
 
 Normalized Laplacian matrix of graph `g`.
 
@@ -156,13 +183,13 @@ Normalized Laplacian matrix of graph `g`.
 - `add_self_loops`: add self-loops while calculating the matrix.
 - `dir`: the edge directionality considered (:out, :in, :both).
 """
-function normalized_laplacian(g::GNNGraph, T::DataType=Float32; 
+function normalized_laplacian(g::GNNGraph, T=nodetype(g); 
                         add_self_loops::Bool=false, dir::Symbol=:out)
     Ã = normalized_adjacency(g, T; dir, add_self_loops)
     return I - Ã
 end
 
-function normalized_adjacency(g::GNNGraph, T::DataType=Float32; 
+function normalized_adjacency(g::GNNGraph, T=nodetype(g); 
                         add_self_loops::Bool=false, dir::Symbol=:out)
     A = adjacency_matrix(g, T; dir=dir)
     if add_self_loops
@@ -174,7 +201,7 @@ function normalized_adjacency(g::GNNGraph, T::DataType=Float32;
 end
 
 @doc raw"""
-    scaled_laplacian(g, T=Float32; dir=:out)
+    scaled_laplacian(g, T=nothing; dir=:out)
 
 Scaled Laplacian matrix of graph `g`,
 defined as ``\hat{L} = \frac{2}{\lambda_{max}} L - I`` where ``L`` is the normalized Laplacian matrix.
@@ -185,7 +212,7 @@ defined as ``\hat{L} = \frac{2}{\lambda_{max}} L - I`` where ``L`` is the normal
 - `T`: result element type.
 - `dir`: the edge directionality considered (:out, :in, :both).
 """
-function scaled_laplacian(g::GNNGraph, T::DataType=Float32; dir=:out)
+function scaled_laplacian(g::GNNGraph, T=nothing; dir=:out)
     L = normalized_laplacian(g, T)
     @assert issymmetric(L) "scaled_laplacian only works with symmetric matrices"
     λmax = _eigmax(L)

test/GNNGraphs/query.jl

@@ -44,14 +44,21 @@
         @test eltype(degree(g, Float32)) == Float32
 
         # weighted degree
-        # if GRAPH_T == :coo
-            eweight = [0.1, 2.1, 1.2, 1]
-            g = GNNGraph((s, t, eweight), graph_type=GRAPH_T)
-            @test degree(g) ==  [2.2, 1.2, 1.0, 0.0]
+        eweight = [0.1, 2.1, 1.2, 1]
+        g = GNNGraph((s, t, eweight), graph_type=GRAPH_T)
+        @test degree(g) ==  [2.2, 1.2, 1.0, 0.0]
+        if GRAPH_T == :coo
             @test degree(g, edge_weight=false) ==  [2, 1, 1, 0]
             @test degree(g, edge_weight=nothing) ==  [2, 1, 1, 0]
+        else
+            # Adjacency matrix representation cannot disambiguate multiple edges
+            # and edge weights
+            @test degree(g, edge_weight=false) ==  [1, 1, 1, 0]
+            @test degree(g, edge_weight=nothing) ==  [1, 1, 1, 0]
+        end
+        if GRAPH_T == :coo
             @test degree(g, edge_weight=2*eweight) == [4.4, 2.4, 2.0, 0.0]
-        # end
+        end
 
         if TEST_GPU
             d = degree(g)