ADJMAT_T can store edge weights

Author: CarloLucibello

src/GNNGraphs/convert.jl

@@ -96,17 +96,18 @@ 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=Int; dir=:out, num_nodes=nothing)
+function to_dense(coo::COO_T, T::DataType=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
     A = fill!(similar(s, T, (n, n)), 0)
-    if isnothing(val)
-        A[s .+ n .* (t .- 1)] .= 1 # exploiting linear indexing
-    else    
-        A[s .+ n .* (t .- 1)] .= val # exploiting linear indexing
-    end
+    v = vec(A)
+    idxs = s .+ n .* (t .- 1) 
+    NNlib.scatter!(+, v, val, idxs)
+    # A[s .+ n .* (t .- 1)] .= val # exploiting linear indexing
     return A, n, length(s)
 end
 

src/GNNGraphs/query.jl

@@ -126,10 +126,15 @@ function Graphs.degree(g::GNNGraph{<:COO_T}, T=nothing; dir=:out, edge_weight=tr
 end
 
 function Graphs.degree(g::GNNGraph{<:ADJMAT_T}, T=Int; dir=:out, edge_weight=true)
-    @assert edge_weight === true
-    @assert dir ∈ (:in, :out)
+    @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)
-    return dir == :out ? vec(sum(A, dims=2)) : vec(sum(A, dims=1))
+    if (edge_weight === false) || (edge_weight === nothing)
+        A = map(>(0), A)
+    end
+    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)

test/GNNGraphs/query.jl

@@ -36,7 +36,6 @@
     @testset "degree" begin 
         s = [1, 1, 2, 3]
         t = [2, 2, 2, 4]
-        eweight = [0.1, 2.1, 1.2, 1]
         g = GNNGraph(s, t, graph_type=GRAPH_T)
 
         @test degree(g) == degree(g; dir=:out) == [2, 1, 1, 0] # default is outdegree
@@ -45,14 +44,14 @@
         @test eltype(degree(g, Float32)) == Float32
 
         # weighted degree
-        if GRAPH_T == :coo
+        # 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]
             @test degree(g, edge_weight=false) ==  [2, 1, 1, 0]
             @test degree(g, edge_weight=nothing) ==  [2, 1, 1, 0]
             @test degree(g, edge_weight=2*eweight) == [4.4, 2.4, 2.0, 0.0]
-        end
+        # end
 
         if TEST_GPU
             d = degree(g)

test/GNNGraphs/transform.jl

@@ -131,11 +131,13 @@
 
     @testset "negative_sample" begin
         if GRAPH_T == :coo
-            n, m = 10,30
+            n, m = 10, 30
             g = rand_graph(n, m, bidirected=true, graph_type=GRAPH_T)
 
             # check bidirected=is_bidirected(g) default
             gneg = negative_sample(g, num_neg_edges=20)
+            @test gneg.num_nodes == g.num_nodes
+            @test gneg.num_edges == 20
             @test is_bidirected(gneg)
             @test intersect(g, gneg).num_edges == 0
         end