support self loops

Author: CarloLucibello

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/query.jl

@@ -113,9 +113,9 @@ function Graphs.adjacency_matrix(g::GNNGraph{<:ADJMAT_T}, T::DataType=nodetype(g
 end
 
 function _get_edge_weight(g, edge_weight)
-    if edge_weight === true
+    if edge_weight === true || edge_weight === nothing 
         ew = get_edge_weight(g)
-    elseif (edge_weight === false) || (edge_weight === nothing)
+    elseif edge_weight === false
         ew = nothing 
     elseif edge_weight isa AbstractVector
         ew = edge_weight 
@@ -147,7 +147,7 @@ function Graphs.degree(g::GNNGraph{<:COO_T}, T=nothing; dir=:out, edge_weight=tr
     s, t = edge_index(g)
 
     edge_weight = _get_edge_weight(g, edge_weight)
-    edge_weight = isnothing(edge_weight) ? eltype(s)(1) : 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)
@@ -161,19 +161,20 @@ function Graphs.degree(g::GNNGraph{<:COO_T}, T=nothing; dir=:out, edge_weight=tr
 end
 
 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) || (edge_weight === nothing)) && !(Nt <: Integer) 
+        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) || (edge_weight === nothing)
+    if edge_weight === false
         A = map(>(0), A)
     end
     A = eltype(A) != T ? T.(A) : 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 get_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)
@@ -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

src/layers/conv.jl

@@ -1,5 +1,5 @@
 @doc raw"""
-    GCNConv(in => out, σ=identity; bias=true, init=glorot_uniform, add_self_loops=true, edge_weight=true)
+    GCNConv(in => out, σ=identity; [bias, init, add_self_loops, use_edge_weight])
 
 Graph convolutional layer from paper [Semi-supervised Classification with Graph Convolutional Networks](https://arxiv.org/abs/1609.02907).
 
@@ -9,7 +9,7 @@ Performs the operation
 ```
 where ``a_{ij} = 1 / \sqrt{|N(i)||N(j)|}`` is a normalization factor computed from the node degrees. 
 
-If the input graph has weighted edges and `edge_weight=true`, than ``c_{ij}`` will be computed as
+If the input graph has weighted edges and `use_edge_weight=true`, than ``a_{ij}`` will be computed as
 ```math
 a_{ij} = \frac{e_{j\to i}}{\sqrt{\sum_{j \in N(i)}  e_{j\to i}} \sqrt{\sum_{i \in N(j)}  e_{i\to j}}}
 ```
@@ -25,15 +25,40 @@ and optionally an edge weight vector.
 - `bias`: Add learnable bias. Default `true`.
 - `init`: Weights' initializer. Default `glorot_uniform`.
 - `add_self_loops`: Add self loops to the graph before performing the convolution. Default `false`.
-- `edge_weight`. If `true`, consider the edge weights in the input graph (if available).
-                 Not compatible with `add_self_loops=true` at the moment. Default `true`.
+- `use_edge_weight`. If `true`, consider the edge weights in the input graph (if available).
+                 If `add_self_loops=true` the new weights will be set to 1. Default `false`.
+
+# Examples
+
+```julia
+# create data
+s = [1,1,2,3]
+t = [2,3,1,1]
+g = GNNGraph(s, t)
+x = randn(3, g.num_nodes)
+
+# create layer
+l = GCNConv(3 => 5) 
+
+# forward pass
+y = l(g, x)       # size:  5 × num_nodes
+
+# convolution with edge weights
+w = [1.1, 0.1, 2.3, 0.5]
+y = l(g, x, w)
+
+# Edge weights can also be embedded in the graph.
+g = GNNGraph(s, t, w)
+l = GCNConv(3 => 5, use_edge_weight=true) 
+y = l(g, x) # same as l(g, x, w) 
+```
 """
 struct GCNConv{A<:AbstractMatrix, B, F} <: GNNLayer
     weight::A
     bias::B
     σ::F
     add_self_loops::Bool
-    edge_weight::Bool
+    use_edge_weight::Bool
 end
 
 @functor GCNConv
@@ -42,46 +67,56 @@ function GCNConv(ch::Pair{Int,Int}, σ=identity;
                  init=glorot_uniform, 
                  bias::Bool=true,
                  add_self_loops=true,
-                 edge_weight=false)
+                 use_edge_weight=false)
     in, out = ch
     W = init(out, in)
     b = bias ? Flux.create_bias(W, true, out) : false
-    GCNConv(W, b, σ, add_self_loops, edge_weight)
+    GCNConv(W, b, σ, add_self_loops, use_edge_weight)
 end
 
-function (l::GCNConv)(g::GNNGraph, x::AbstractMatrix)
-    # Extract edge_weight from g if available and l.edge_weight == false,
+function (l::GCNConv)(g::GNNGraph{<:COO_T}, x::AbstractMatrix)
+    # Extract edge_weight from g if available and l.edge_weight == true,
     # otherwise return nothing.
-    edge_weight = GNNGraphs._get_edge_weight(g, l.edge_weight) # vector or nothing
+    edge_weight = GNNGraphs._get_edge_weight(g, l.use_edge_weight) # vector or nothing
+    return l(g, x, edge_weight)
+end
+
+function (l::GCNConv)(g::GNNGraph{<:ADJMAT_T}, x::AbstractMatrix)
+    edge_weight = nothing
     return l(g, x, edge_weight)
 end
 
 function (l::GCNConv)(g::GNNGraph, x::AbstractMatrix{T}, edge_weight::EW) where 
     {T, EW<:Union{Nothing,AbstractVector}}
 
     if l.add_self_loops
-        @assert edge_weight === nothing
         g = add_self_loops(g)
+        if edge_weight !== nothing
+            edge_weight = [edge_weight; fill!(similar(edge_weight, g.num_nodes), 1)]
+            @assert length(edge_weight) == g.num_edges
+        end
     end
     Dout, Din = size(l.weight)
     if Dout < Din
+        # multiply before convolution if it is more convenient, otherwise multiply after
         x = l.weight * x
     end
-    # @assert all(>(0), degree(g, T, dir=:in))
-    c = 1 ./ sqrt.(degree(g, T; dir=:in, edge_weight))
+    d = degree(g, T; dir=:in, edge_weight)
+    c = 1 ./ sqrt.(d)
     x = x .* c'
     if edge_weight === nothing
         x = propagate(copy_xj, g, +, xj=x)
     else        
         x = propagate(e_mul_xj, g, +, xj=x, e=edge_weight)
     end
-        x = x .* c'
+    x = x .* c'
     if Dout >= Din
         x = l.weight * x
     end
     return l.σ.(x .+ l.bias)
 end
 
+
 function Base.show(io::IO, l::GCNConv)
     out, in = size(l.weight)
     print(io, "GCNConv($in => $out")

test/GNNGraphs/query.jl

@@ -60,15 +60,14 @@
             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=nothing) ==  degree(g)
             d = degree(g, edge_weight=false)
             if GRAPH_T == :coo
                 @test d ==  [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 d ==  [1, 1, 1, 0]
-                @test degree(g, edge_weight=nothing) ==  [1, 1, 1, 0]
             end
             @test eltype(d) <: Integer
             if GRAPH_T == :coo

test/layers/conv.jl

@@ -37,6 +37,23 @@
 
         l = GCNConv(in_channel => out_channel, add_self_loops=false)
         test_layer(l, g1, rtol=1e-5, outsize=(out_channel, g1.num_nodes))
+
+        @testset "edge weights" begin
+            s = [2,3,1,3,1,2]
+            t = [1,1,2,2,3,3]
+            w = [1,2,3,4,5,6]
+            g = GNNGraph((s, t, w), graph_type=GRAPH_T)
+            x = ones(1, g.num_nodes)
+            l = GCNConv(1 => 1, add_self_loops=false, use_edge_weight=true)
+            l.weight .= 1
+            d = degree(g, dir=:in)
+            y = l(g, x)
+            @test y[1,1] ≈ w[1] / √(d[1]*d[2]) + w[2] / √(d[1]*d[3]) 
+            @test y[1,2] ≈ w[3] / √(d[2]*d[1]) + w[4] / √(d[2]*d[3])
+            if GRAPH_T == :coo
+                @test y ≈ l(g, x, w)                
+            end
+        end
     end
 
     @testset "ChebConv" begin