duplicate edge features if needed

Author: CarloLucibello

docs/src/gnngraph.md

@@ -45,9 +45,17 @@ g = GNNGraph(erdos_renyi(10,  30), ndata = rand(Float32, 32, 10))
 g = GNNGraph(erdos_renyi(10,  30), ndata = (; x=rand(Float32, 32, 10), y=rand(Float32, 10)))
 
 
-# Attach an array with edge features
+# Attach an array with edge features.
+# Since `GNNGraph`s are directed, the number of edges
+# will be double that of the original LightGraphs' undirected graph.
+g = GNNGraph(erdos_renyi(10,  30), edata = rand(Float32, 60))
+@assert g.num_edges == 60
+
+# If we pass only half of the edge features, they will be copied
+# on the reversed edges.
 g = GNNGraph(erdos_renyi(10,  30), edata = rand(Float32, 30))
 
+
 # Create a new graph from previous one, inheriting edge data
 # but replacing node data
 g′ = GNNGraph(g, ndata =(; z = ones(Float32, 16, 10)))

docs/src/index.md

@@ -38,7 +38,7 @@ julia> for _ in 1:1000
 julia> gbatch = Flux.batch(all_graphs)
 GNNGraph:
     num_nodes = 10000
-    num_edges = 20000
+    num_edges = 40000
     num_graphs = 1000
     ndata:
         x => (16, 10000)

src/gnngraph.jl

@@ -124,20 +124,20 @@ function GNNGraph(data;
     @assert dir ∈ [:in, :out]
 
     if graph_type == :coo
-        g, num_nodes, num_edges = to_coo(data; num_nodes, dir)
+        graph, num_nodes, num_edges = to_coo(data; num_nodes, dir)
     elseif graph_type == :dense
-        g, num_nodes, num_edges = to_dense(data; dir)
+        graph, num_nodes, num_edges = to_dense(data; dir)
     elseif graph_type == :sparse
-        g, num_nodes, num_edges = to_sparse(data; dir)
+        graph, num_nodes, num_edges = to_sparse(data; dir)
     end
 
     num_graphs = !isnothing(graph_indicator) ? maximum(graph_indicator) : 1
 
-    ndata = normalize_graphdata(ndata, :x)
-    edata = normalize_graphdata(edata, :e)
-    gdata = normalize_graphdata(gdata, :u)
+    ndata = normalize_graphdata(ndata, default_name=:x, n=num_nodes)
+    edata = normalize_graphdata(edata, default_name=:e, n=num_edges, duplicate_if_needed=true)
+    gdata = normalize_graphdata(gdata, default_name=:u, n=num_graphs)
 
-    GNNGraph(g, 
+    GNNGraph(graph, 
             num_nodes, num_edges, num_graphs, 
             graph_indicator,
             ndata, edata, gdata)
@@ -149,21 +149,22 @@ GNNGraph((s, t)::NTuple{2}; kws...) = GNNGraph((s, t, nothing); kws...)
 
 # GNNGraph(g::AbstractGraph; kws...) = GNNGraph(adjacency_matrix(g, dir=:out); kws...)
 
-function GNNGraph(g::AbstractGraph; kws...)
+function GNNGraph(g::AbstractGraph; edata=(;), kws...)
     s = LightGraphs.src.(LightGraphs.edges(g))
     t = LightGraphs.dst.(LightGraphs.edges(g))
     if !LightGraphs.is_directed(g) 
         # add reverse edges since GNNGraph are directed
-        s, t = [s; t], [t; s]
+        s, t = [s; t], [t; s]    
     end
-    GNNGraph((s, t); num_nodes = LightGraphs.nv(g), kws...)
+    GNNGraph((s, t); edata, num_nodes=LightGraphs.nv(g), kws...)
 end
 
+
 function GNNGraph(g::GNNGraph; ndata=g.ndata, edata=g.edata, gdata=g.gdata)
 
-    ndata = normalize_graphdata(ndata, :x)
-    edata = normalize_graphdata(edata, :e)
-    gdata = normalize_graphdata(gdata, :u)
+    ndata = normalize_graphdata(ndata, default_name=:x, n=g.num_nodes)
+    edata = normalize_graphdata(edata, default_name=:e, n=g.num_edges, duplicate_if_needed=true)
+    gdata = normalize_graphdata(gdata, default_name=:u, n=g.num_graphs)
 
     GNNGraph(g.graph, 
             g.num_nodes, g.num_edges, g.num_graphs, 

src/layers/conv.jl

@@ -1,13 +1,13 @@
 @doc raw"""
-    GCNConv(in => out, σ=identity; bias=true, init=glorot_uniform)
+    GCNConv(in => out, σ=identity; bias=true, init=glorot_uniform, add_self_loops=true)
 
 Graph convolutional layer from paper [Semi-supervised Classification with Graph Convolutional Networks](https://arxiv.org/abs/1609.02907).
 
 Performs the operation
 ```math
-\mathbf{x}'_i = \sum_{j\in \{i\} \cup N(i)} \frac{1}{c_{ij}} W \mathbf{x}_j
+\mathbf{x}'_i = \sum_{j\in N(i)} \frac{1}{c_{ij}} W \mathbf{x}_j
 ```
-where ``c_{ij} = \sqrt{(1+|N(i)|)(1+|N(j)|)}``.
+where ``c_{ij} = \sqrt{|N(i)||N(j)|}``.
 
 The input to the layer is a node feature array `X` 
 of size `(num_features, num_nodes)`.
@@ -19,37 +19,42 @@ of size `(num_features, num_nodes)`.
 - `σ`: Activation function.
 - `bias`: Add learnable bias.
 - `init`: Weights' initializer.
+- `add_self_loops`: Add self loops to the graph before performing the convolution.
 """
 struct GCNConv{A<:AbstractMatrix, B, F} <: GNNLayer
     weight::A
     bias::B
     σ::F
+    add_self_loops::Bool
 end
 
 @functor GCNConv
 
 function GCNConv(ch::Pair{Int,Int}, σ=identity;
-                 init=glorot_uniform, bias::Bool=true)
+                 init=glorot_uniform, bias::Bool=true,
+                 add_self_loops=true)
     in, out = ch
     W = init(out, in)
     b = bias ? Flux.create_bias(W, true, out) : false
-    GCNConv(W, b, σ)
+    GCNConv(W, b, σ, add_self_loops)
 end
 
 ## Matrix operations are more performant, 
-## but cannot compute the normalized laplacian of sparse cuda matrices yet,
+## but cannot compute the normalized adjacency of sparse cuda matrices yet,
 ## therefore fallback to message passing framework on gpu for the time being
 
 function (l::GCNConv)(g::GNNGraph, x::AbstractMatrix{T}) where T
-    Ã = normalized_adjacency(g, T; dir=:out, add_self_loops=true)
+    Ã = normalized_adjacency(g, T; dir=:out, l.add_self_loops)
     l.σ.(l.weight * x * Ã .+ l.bias)
 end
 
 compute_message(l::GCNConv, xi, xj, eij) = xj
 update_node(l::GCNConv, m, x) = m
 
 function (l::GCNConv)(g::GNNGraph, x::CuMatrix{T}) where T
-    g = add_self_loops(g)
+    if l.add_self_loops
+        g = add_self_loops(g)
+    end
     c = 1 ./ sqrt.(degree(g, T, dir=:in))
     x = x .* c'
     x, _ = propagate(l, g, +, x)

src/utils.jl

@@ -24,6 +24,28 @@ function cat_features(x1::NamedTuple, x2::NamedTuple)
 end
 
 # Turns generic type into named tuple
-normalize_graphdata(data::NamedTuple, s::Symbol) = data
-normalize_graphdata(data::Nothing, s::Symbol) = NamedTuple()
-normalize_graphdata(data, s::Symbol) = NamedTuple{(s,)}((data,)) 
+normalize_graphdata(data::Nothing; kws...) = NamedTuple()
+
+normalize_graphdata(data; default_name::Symbol, kws...) = 
+    normalize_graphdata(NamedTuple{(default_name,)}((data,)); default_name, kws...) 
+
+function normalize_graphdata(data::NamedTuple; default_name=:z, n, duplicate_if_needed=false)
+    sz = map(x -> x isa AbstractArray ? size(x)[end] : 0, data)
+    
+    if duplicate_if_needed # used to copy edge features on reverse edges    
+        @assert all(s -> s == 0 ||  s == n || s == n÷2, sz)
+
+        function replace(k, v)
+            if v isa AbstractArray && size(v)[end] == n÷2
+                v = cat(v, v, dims=ndims(v))
+            end
+            k => v
+        end
+
+        data = NamedTuple(replace(k,v) for (k,v) in pairs(data))
+    else
+        @assert all(s -> s == 0 ||  s == n, sz)
+    end
+    return data
+end
+

test/gnngraph.jl

@@ -182,6 +182,21 @@
         @test g.ndata.x2 ≈ 2X
         @test g.edata.e2 ≈ 2E
         @test g.gdata.u2 ≈ 2U
+
+        # Dimension checks
+        @test_throws AssertionError GNNGraph(erdos_renyi(10,  30), edata=rand(29), graph_type=GRAPH_T)
+        @test_throws AssertionError GNNGraph(erdos_renyi(10,  30), edata=rand(2, 29), graph_type=GRAPH_T)
+        @test_throws AssertionError GNNGraph(erdos_renyi(10,  30), edata=(; x=rand(30), y=rand(29)), graph_type=GRAPH_T)
+
+        # Copy features on reverse edge
+        e = rand(30)
+        g = GNNGraph(erdos_renyi(10,  30), edata=e, graph_type=GRAPH_T)
+        @test g.edata.e == [e; e]
+
+
+        # Attach non array data
+        g = GNNGraph(erdos_renyi(10,  30), edata="ciao", graph_type=GRAPH_T)
+        @test g.edata.e == "ciao"
     end 
 
     @testset "LearnBase and DataLoader compat" begin