support edge_weights in degree

Author: CarloLucibello

.github/workflows/CI.yml

@@ -6,9 +6,8 @@ on:
   push:
     branches:
       - master
-    tags: '*'
 jobs:
-  test:
+  test:https://piazza.com/unibocconi.it/fall2020/csdb/resources
     name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }}
     runs-on: ${{ matrix.os }}
     strategy:
@@ -45,5 +44,4 @@ jobs:
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v1
         with:
-          # token: ${{ secrets.CODECOV_TOKEN }}
           file: lcov.info

src/GNNGraphs/query.jl

@@ -13,9 +13,9 @@ edge_index(g::GNNGraph{<:COO_T}) = g.graph[1:2]
 
 edge_index(g::GNNGraph{<:ADJMAT_T}) = to_coo(g.graph, num_nodes=g.num_nodes)[1][1:2]
 
-edge_weight(g::GNNGraph{<:COO_T}) = g.graph[3]
+get_edge_weight(g::GNNGraph{<:COO_T}) = g.graph[3]
 
-edge_weight(g::GNNGraph{<:ADJMAT_T}) = to_coo(g.graph, num_nodes=g.num_nodes)[1][3]
+get_edge_weight(g::GNNGraph{<:ADJMAT_T}) = to_coo(g.graph, num_nodes=g.num_nodes)[1][3]
 
 Graphs.edges(g::GNNGraph) = zip(edge_index(g)...)
 
@@ -95,21 +95,33 @@ function Graphs.adjacency_matrix(g::GNNGraph{<:ADJMAT_T}, T::DataType=eltype(g.g
     return dir == :out ? A : A'
 end
 
-function Graphs.degree(g::GNNGraph{<:COO_T}, T=nothing; dir=:out)
+function Graphs.degree(g::GNNGraph{<:COO_T}, T=nothing; dir=:out, edge_weight=true)
     s, t = edge_index(g)
-    T = isnothing(T) ? eltype(s) : T
+
+    if edge_weight === true
+        edge_weight = get_edge_weight(g)
+    elseif (edge_weight === false) || (edge_weight === nothing)
+        edge_weight = nothing 
+    elseif edge_weight isa AbstractVector
+        edge_weight = edge_weight 
+    else 
+        error("Invalid edge_weight argument.")
+    end
+    edge_weight = isnothing(edge_weight) ? eltype(s)(1) : edge_weight
+
+    T = isnothing(T) ? eltype(edge_weight) : T
     degs = fill!(similar(s, T, g.num_nodes), 0)
-    src = 1
     if dir ∈ [:out, :both]
-        NNlib.scatter!(+, degs, src, s)
+        NNlib.scatter!(+, degs, edge_weight, s)
     end
     if dir ∈ [:in, :both]
-        NNlib.scatter!(+, degs, src, t)
+        NNlib.scatter!(+, degs, edge_weight, t)
     end
     return degs 
 end
 
-function Graphs.degree(g::GNNGraph{<:ADJMAT_T}, T=Int; dir=:out)
+function Graphs.degree(g::GNNGraph{<:ADJMAT_T}, T=Int; dir=:out, edge_weight=true)
+    @assert edge_weight === true
     @assert dir ∈ (:in, :out)
     A = adjacency_matrix(g, T)
     return dir == :out ? vec(sum(A, dims=2)) : vec(sum(A, dims=1))

src/GNNGraphs/transform.jl

@@ -11,7 +11,7 @@ self-loops will obtain a second self-loop.
 function add_self_loops(g::GNNGraph{<:COO_T})
     s, t = edge_index(g)
     @assert g.edata === (;)
-    @assert edge_weight(g) === nothing
+    @assert get_edge_weight(g) === nothing
     n = g.num_nodes
     nodes = convert(typeof(s), [1:n;])
     s = [s; nodes]
@@ -39,7 +39,7 @@ function remove_self_loops(g::GNNGraph{<:COO_T})
     s, t = edge_index(g)
     # TODO remove these constraints
     @assert g.edata === (;)
-    @assert edge_weight(g) === nothing
+    @assert get_edge_weight(g) === nothing
 
     mask_old_loops = s .!= t
     s = s[mask_old_loops]
@@ -61,7 +61,7 @@ function remove_multi_edges(g::GNNGraph{<:COO_T})
     # TODO remove these constraints
     @assert g.num_graphs == 1
     @assert g.edata === (;)
-    @assert edge_weight(g) === nothing
+    @assert get_edge_weight(g) === nothing
 
     idxs, idxmax = edge_encoding(s, t, g.num_nodes)
     union!(idxs)
@@ -85,7 +85,7 @@ function add_edges(g::GNNGraph{<:COO_T},
 
     @assert length(snew) == length(tnew)
     # TODO remove this constraint
-    @assert edge_weight(g) === nothing
+    @assert get_edge_weight(g) === nothing
 
     edata = normalize_graphdata(edata, default_name=:e, n=length(snew))
     edata = cat_features(g.edata, edata)
@@ -126,7 +126,7 @@ function SparseArrays.blockdiag(g1::GNNGraph, g2::GNNGraph)
         s2, t2 = edge_index(g2)
         s = vcat(s1, nv1 .+ s2)
         t = vcat(t1, nv1 .+ t2)
-        w = cat_features(edge_weight(g1), edge_weight(g2))
+        w = cat_features(get_edge_weight(g1), get_edge_weight(g2))
         graph = (s, t, w)
         ind1 = isnothing(g1.graph_indicator) ? ones_like(s1, Int, nv1) : g1.graph_indicator 
         ind2 = isnothing(g2.graph_indicator) ? ones_like(s2, Int, nv2) : g2.graph_indicator     
@@ -288,7 +288,7 @@ function getgraph(g::GNNGraph, i::AbstractVector{Int}; nmap=false)
     graph_indicator = [graphmap[i] for i in g.graph_indicator[node_mask]]
 
     s, t = edge_index(g)
-    w = edge_weight(g)
+    w = get_edge_weight(g)
     edge_mask = s .∈ Ref(nodes) 
 
     if g.graph isa COO_T 

test/GNNGraphs/gnngraph.jl

@@ -70,18 +70,6 @@
             adjacency_matrix(g; dir=:in) == adj_mat
         end 
 
-        @testset "degree" begin
-            @test degree(g, dir=:out) == vec(sum(adj_mat, dims=2))
-            @test degree(g, dir=:in) == vec(sum(adj_mat, dims=1))
-
-            if TEST_GPU
-                d = degree(g)
-                d_gpu = degree(g_gpu)
-                @test d_gpu isa CuVector
-                @test Array(d_gpu) == d
-            end
-        end
-
         if TEST_GPU
             @testset "functor" begin                
                 s_cpu, t_cpu = edge_index(g)
@@ -134,11 +122,6 @@
         @test adjacency_list(g, dir=:out) ==  adj_list_out
         @test adjacency_matrix(g, dir=:in) ==  adj_mat_in
         @test adjacency_list(g, dir=:in) ==  adj_list_in
-
-        @testset "degree" begin
-            @test degree(g, dir=:out) == vec(sum(adj_mat_out, dims=2))
-            @test degree(g, dir=:in) == vec(sum(adj_mat_out, dims=1))
-        end
     end
 
     @testset "Graphs constructor" begin

test/GNNGraphs/query.jl

@@ -32,4 +32,33 @@
         g = GNNGraph(s, t, graph_type=GRAPH_T)
         @test !has_self_loops(g)
     end
+
+    @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
+        @test degree(g; dir=:in) == [0, 3, 0, 1]
+        @test degree(g; dir=:both) == [2, 4, 1, 1]
+        @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]
+            @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
+
+        if TEST_GPU
+            d = degree(g)
+            d_gpu = degree(g_gpu)
+            @test d_gpu isa CuVector
+            @test Array(d_gpu) == d
+        end
+    end
 end