Merge pull request #76 from CarloLucibello/cl/eweight

Support edge weights in GCNConv

Author: CarloLucibello

.github/workflows/CI.yml

@@ -6,7 +6,6 @@ on:
   push:
     branches:
       - master
-    tags: '*'
 jobs:
   test:
     name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }}
@@ -45,5 +44,4 @@ jobs:
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v1
         with:
-          # token: ${{ secrets.CODECOV_TOKEN }}
           file: lcov.info

Project.toml

@@ -22,7 +22,6 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
-TestEnv = "1e6cf692-eddd-4d53-88a5-2d735e33781b"
 
 [compat]
 Adapt = "3"
@@ -39,7 +38,6 @@ NNlib = "0.7"
 NNlibCUDA = "0.1"
 Reexport = "1"
 StatsBase = "0.32, 0.33"
-TestEnv = "1"
 julia = "1.6"
 
 [extras]

docs/src/api/messagepassing.md

@@ -24,4 +24,5 @@ propagate
 copy_xi
 copy_xj
 xi_dot_xj
+e_mul_xj
 ```

docs/src/gnngraph.md

@@ -114,6 +114,25 @@ g.ndata.z
 g.edata.e
 ```
 
+## Edge weights
+
+It is common to denote scalar edge features as edge weights. The `GNNGraph` has specific support
+for edge weights: they can be stored as part of internal representions of the graph (COO or adjacency matrix). Some graph convolutional layers, most notably the [`GCNConv`](@ref), can use the edge weights to perform weighted sums over the nodes' neighborhoods.
+
+```julia
+julia> source = [1, 1, 2, 2, 3, 3];
+
+julia> target = [2, 3, 1, 3, 1, 2];
+
+julia> weight = [1.0, 0.5, 2.1, 2.3, 4, 4.1];
+
+julia> g = GNNGraph(source, target, weight)
+GNNGraph:
+    num_nodes = 3
+    num_edges = 6
+    
+```
+
 ## Batches and Subgraphs
 
 Multiple `GNNGraph`s can be batched togheter into a single graph

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,26 +98,28 @@ 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=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
 
 ### 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'
@@ -129,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

@@ -13,13 +13,30 @@ 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)...)
 
 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,34 +105,85 @@ 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
     return dir == :out ? A : A'
 end
 
-function Graphs.degree(g::GNNGraph{<:COO_T}, T=nothing; dir=:out)
+function _get_edge_weight(g, edge_weight)
+    if edge_weight === true || edge_weight === nothing 
+        ew = get_edge_weight(g)
+    elseif edge_weight === false
+        ew = nothing 
+    elseif edge_weight isa AbstractVector
+        ew = edge_weight 
+    else
+        error("Invalid edge_weight argument.")
+    end
+    return ew
+end
+
+"""
+    degree(g::GNNGraph, T=nothing; dir=:out, edge_weight=true)
+
+Return a vector containing the degrees of the nodes in `g`.
+
+# Arguments
+- `g`: A graph.
+- `T`: Element type of the returned vector. If `nothing`, is
+       chosen based on the graph type and will be an integer
+       if `edge_weight=false`.
+- `dir`: For `dir=:out` the degree of a node is counted based on the outgoing edges.
+         For `dir=:in`, the ingoing edges are used. If `dir=:both` we have the sum of the two.
+- `edge_weight`: If `true` and the graph contains weighted edges, the degree will 
+                be weighted. Set to `false` instead to just count the number of
+                outgoing/ingoing edges.
+                In alternative, you can also pass a vector of weights to be used
+                instead of the graph's own weights.
+"""
+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
+
+    edge_weight = _get_edge_weight(g, edge_weight)
+    edge_weight = edge_weight === nothing ? 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)
-    @assert dir ∈ (:in, :out)
-    A = adjacency_matrix(g, T)
-    return dir == :out ? vec(sum(A, dims=2)) : vec(sum(A, dims=1))
+function Graphs.degree(g::GNNGraph{<:ADJMAT_T}, T=nothing; dir=:out, edge_weight=true)
+    # edge_weight=true or edge_weight=nothing act the same here
+    @assert !(edge_weight isa AbstractArray) "passing the edge weights is not support by adjacency matrix representations" 
+    @assert dir ∈ (:in, :out, :both)
+    if T === nothing
+        Nt = nodetype(g)
+        if edge_weight === false && !(Nt <: Integer) 
+            T = Nt == Float32 ? Int32 : 
+                Nt == Float16 ? Int16 : Int
+        else
+            T = Nt
+        end
+    end
+    A = adjacency_matrix(g)
+    if edge_weight === false
+        A = map(>(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::DataType=nodetype(g); dir::Symbol=:out)
     A = adjacency_matrix(g, T; dir=dir)
     D = Diagonal(vec(sum(A; dims=2)))
     return D - A

src/GNNGraphs/transform.jl

@@ -5,19 +5,23 @@
 Return a graph with the same features as `g`
 but also adding edges connecting the nodes to themselves.
 
-Nodes with already existing
-self-loops will obtain a second self-loop.
+Nodes with already existing self-loops will obtain a second self-loop.
+
+If the graphs has edge weights, the new edges will have weight 1.
 """
 function add_self_loops(g::GNNGraph{<:COO_T})
     s, t = edge_index(g)
     @assert g.edata === (;)
-    @assert edge_weight(g) === nothing
+    ew = get_edge_weight(g)
     n = g.num_nodes
     nodes = convert(typeof(s), [1:n;])
     s = [s; nodes]
     t = [t; nodes]
+    if ew !== nothing
+        ew = [ew; fill!(similar(ew, n), 1)]
+    end
 
-    GNNGraph((s, t, nothing), 
+    GNNGraph((s, t, ew), 
         g.num_nodes, length(s), g.num_graphs, 
         g.graph_indicator,
         g.ndata, g.edata, g.gdata)
@@ -39,7 +43,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 +65,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 +89,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 +130,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 +292,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 
@@ -340,7 +344,7 @@ function negative_sample(g::GNNGraph;
     @assert g.num_graphs == 1
     # Consider self-loops as positive edges
     # Construct new graph dropping features
-    g = add_self_loops(GNNGraph(edge_index(g))) 
+    g = add_self_loops(GNNGraph(edge_index(g), num_nodes=g.num_nodes)) 
 
     s, t = edge_index(g)
     n = g.num_nodes