fixes

Author: CarloLucibello

src/featuredgraph.jl

@@ -298,29 +298,32 @@ function LightGraphs.laplacian_matrix(fg::FeaturedGraph, T::DataType=Int; dir::S
 end
 
 """
-    normalized_laplacian(fg, T=Float32; selfloop=false, dir=:out)
+    normalized_laplacian(fg, T=Float32; add_self_loops=false, dir=:out)
 
 Normalized Laplacian matrix of graph `g`.
 
 # Arguments
 
 - `fg`: A `FeaturedGraph`.
 - `T`: result element type.
-- `selfloop`: adding self loop while calculating the matrix.
+- `add_self_loops`: add self-loops while calculating the matrix.
 - `dir`: the edge directionality considered (:out, :in, :both).
 """
-function normalized_laplacian(fg::FeaturedGraph, T::DataType=Float32; selfloop::Bool=false, dir::Symbol=:out)
+function normalized_laplacian(fg::FeaturedGraph, T::DataType=Float32; 
+                        add_self_loops::Bool=false, dir::Symbol=:out)
+    Ã = normalized_adjacency(fg, T; dir, add_self_loops)
+    return I - Ã
+end
+
+function normalized_adjacency(fg::FeaturedGraph, T::DataType=Float32; 
+                        add_self_loops::Bool=false, dir::Symbol=:out)
     A = adjacency_matrix(fg, T; dir=dir)
-    sz = size(A)
-    @assert sz[1] == sz[2]
-    if selfloop
-        A += I - Diagonal(A)
-    else
-        A -= Diagonal(A) 
+    if add_self_loops
+        A += I
     end
     degs = vec(sum(A; dims=2))
     inv_sqrtD = Diagonal(inv.(sqrt.(degs)))
-    return I - inv_sqrtD * A * inv_sqrtD
+    return inv_sqrtD * A * inv_sqrtD
 end
 
 @doc raw"""
@@ -350,18 +353,23 @@ _eigmax(A) = KrylovKit.eigsolve(Symmetric(A), 1, :LR)[1][1] # also eigs(A, x0, n
 # https://discourse.julialang.org/t/cuda-eigenvalues-of-a-sparse-matrix/46851/5
 
 """
-    add_self_loops(fg::FeaturedGraph)
+    add_self_loops(fg::FeaturedGraph; add_to_existing=true)
 
 Return a featured graph with the same features as `fg`
 but also adding edges connecting the nodes to themselves.
+
+If `add_to_existing=true`, nodes with already existing
+self-loops will obtain a second self-loop.
 """
-function add_self_loops(fg::FeaturedGraph{<:COO_T})
+function add_self_loops(fg::FeaturedGraph{<:COO_T}; add_to_existing=true)
     s, t = edge_index(fg)
     @assert edge_feature(fg) === nothing
     @assert edge_weight(fg) === nothing
-    mask_old_loops = s .!= t
-    s = s[mask_old_loops]
-    t = t[mask_old_loops]
+    if !add_to_existing
+        mask_old_loops = s .!= t
+        s = s[mask_old_loops]
+        t = t[mask_old_loops]
+    end
     n = fg.num_nodes
     nodes = convert(typeof(s), [1:n;])
     s = [s; nodes]
@@ -371,11 +379,16 @@ function add_self_loops(fg::FeaturedGraph{<:COO_T})
         node_feature(fg), edge_feature(fg), global_feature(fg))
 end
 
-function add_self_loops(fg::FeaturedGraph{<:ADJMAT_T})
+function add_self_loops(fg::FeaturedGraph{<:ADJMAT_T}; add_to_existing=true)
     A = graph(fg)
     @assert edge_feature(fg) === nothing
-    nold = sum(Diagonal(A)) |> Int
-    A = A - Diagonal(A) + I
+    if add_to_existing
+        nold = 0
+        A += I
+    else
+        nold = sum(Diagonal(A)) |> Int
+        A += I - Diagonal(A)
+    end
     num_edges =  fg.num_edges - nold + fg.num_nodes
     FeaturedGraph(A, fg.num_nodes, num_edges,
         node_feature(fg), edge_feature(fg), global_feature(fg))
@@ -396,6 +409,7 @@ function remove_self_loops(fg::FeaturedGraph{<:COO_T})
 end
 
 @non_differentiable normalized_laplacian(x...)
+@non_differentiable normalized_adjacency(x...)
 @non_differentiable scaled_laplacian(x...)
 @non_differentiable adjacency_matrix(x...)
 @non_differentiable adjacency_list(x...)

src/layers/conv.jl

@@ -35,15 +35,15 @@ end
 ## therefore fallback to message passing framework on gpu for the time being
 
 function (l::GCNConv)(fg::FeaturedGraph, x::AbstractMatrix)
-    L̃ = normalized_laplacian(fg, eltype(x); selfloop=true)
-    l.σ.(l.weight * x * L̃ .+ l.bias)
+    Ã = normalized_adjacency(fg, eltype(x); dir=:out, add_self_loops=true)
+    l.σ.(l.weight * x * Ã .+ l.bias)
 end
 
 message(l::GCNConv, xi, xj) = xj
 update(l::GCNConv, m, x) = m
 
 function (l::GCNConv)(fg::FeaturedGraph, x::CuMatrix)
-    fg = add_self_loops(fg)
+    fg = add_self_loops(fg; add_to_existing=true)
     T = eltype(l.weight)
     # cout = sqrt.(degree(fg, dir=:out))
     cin = 1 ./ reshape(sqrt.(T.(degree(fg, dir=:in))), 1, :)