undirected edge encoding

Author: CarloLucibello

src/GNNGraphs/transform.jl

@@ -393,36 +393,6 @@ function rand_edge_split(g::GNNGraph, frac)
 end
 
 
-# each edge is represented by a number in
-# 1:N^2
-function edge_encoding(s, t, n; directed=true)
-    if directed
-        # directed edges and self-loops allowed
-        idx = (s .- 1) .* n .+ t
-        maxid = n^2
-    else 
-        # undirected edges and self-loops allowed
-        maxid = n * (n - 1) ÷ 2
-        mask = s .<= t
-        s1, t1 = s[mask], t[mask]
-        t2, s2 = s[.!mask], t[.!mask]
-        s, t = [s1; s2], [t1; t2] 
-        offset1 = (n .* 0:n-1) .- cumsum(0:n-1)
-        offset2 = 0:n-1
-        idx = offset1[s] .+ (t .- offset2)
-    end
-    return idx, maxid
-end
-
-# each edge is represented by a number in
-# 1:N^2
-function edge_decoding(idx, n)
-    # g = remove_self_loops(g)
-    s =  (idx .- 1) .÷ n .+ 1
-    t =  (idx .- 1) .% n .+ 1
-    return s, t
-end
-
 # """
 # Transform vector of cartesian indexes into a tuple of vectors containing integers.
 # """

src/GNNGraphs/utils.jl

@@ -70,3 +70,47 @@ end
 ones_like(x::AbstractArray, T=eltype(x), sz=size(x)) = fill!(similar(x, T, sz), 1)
 ones_like(x::SparseMatrixCSC, T=eltype(x), sz=size(x)) = ones(T, sz)
 ones_like(x::CUMAT_T, T=eltype(x), sz=size(x)) = CUDA.ones(T, sz)
+
+
+# each edge is represented by a number in
+# 1:N^2
+function edge_encoding(s, t, n; directed=true)
+    if directed
+        # directed edges and self-loops allowed
+        idx = (s .- 1) .* n .+ t
+        maxid = n^2
+    else 
+        # Undirected edges and self-loops allowed
+        # In this encoding, each edge has 2 possible encodings (also the self-loops).
+        # We return the canonical one given by the upper triangular adj matrix
+        maxid = n * (n + 1) ÷ 2
+        mask = s .> t
+        # s1, t1 = s[mask], t[mask]
+        # t2, s2 = s[.!mask], t[.!mask]
+        snew = copy(s)
+        tnew = copy(t)
+        snew[mask] .= t[mask]
+        tnew[mask] .= s[mask]
+        s, t = snew, tnew
+        
+        # idx = ∑_{i',i'<i} ∑_{j',j'>=i'}^n 1 + ∑_{j',i<=j'<=j} 1 
+        #     = ∑_{i',i'<i} ∑_{j',j'>=i'}^n 1 + j - i + 1
+        #     = ∑_{i',i'<i} (n - i' + 1) + (j - i + 1)
+        #     = (i - 1)*(2*(n+1)-i)÷2 + (j - i + 1)
+        idx = @. (s-1)*(2*(n+1)-s)÷2 + (t-s+1)
+    end
+    return idx, maxid
+end
+
+# each edge is represented by a number in
+# 1:N^2
+function edge_decoding(idx, n; directed=true)
+    # g = remove_self_loops(g)
+    s =  (idx .- 1) .÷ n .+ 1
+    t =  (idx .- 1) .% n .+ 1
+    return s, t
+end
+
+@non_differentiable edge_encoding(x...)
+@non_differentiable edge_decoding(x...)
+

test/GNNGraphs/utils.jl

@@ -0,0 +1,23 @@
+@testset "Utils" begin
+    @testset "edge encoding/decoding" begin
+        # not is_bidirected
+        n = 5
+        s = [1,1,2,3,3,4,5]
+        t = [1,3,1,1,2,5,5]
+        
+        # directed=true
+        idx, maxid = GNNGraphs.edge_encoding(s, t, n) 
+        @test maxid == n^2
+        @test idx == [1, 3, 6, 11, 12, 20, 25]
+
+        sdec, tdec = GNNGraphs.edge_decoding(idx, n) 
+        @test sdec == s
+        @test tdec == t
+
+
+        # directed=false
+        idx, maxid = GNNGraphs.edge_encoding(s, t, n, directed=false)
+        @test maxid == n * (n+1)÷2
+        @test idx == [1, 3, 2, 3, 7, 14, 15]
+    end
+end