rng instead of seed for rand_graph

Author: CarloLucibello

.gitignore

@@ -9,4 +9,5 @@ Manifest.toml
 .vscode
 LocalPreferences.toml
 .DS_Store
-docs/src/democards/gridtheme.css
+docs/src/democards/gridtheme.css
+test.jl

GNNGraphs/src/abstracttypes.jl

@@ -1,5 +1,5 @@
 
-const COO_T = Tuple{T, T, V} where {T <: AbstractVector{<:Integer}, V}
+const COO_T = Tuple{T, T, V} where {T <: AbstractVector{<:Integer}, V <: AbstractVector}
 const ADJLIST_T = AbstractVector{T} where {T <: AbstractVector{<:Integer}}
 const ADJMAT_T = AbstractMatrix
 const SPARSE_T = AbstractSparseMatrix # subset of ADJMAT_T

GNNGraphs/src/generate.jl

@@ -36,15 +36,30 @@ GNNGraph:
 # Each edge has a reverse
 julia> edge_index(g)
 ([1, 3, 3, 4], [3, 4, 1, 3])
-
 ```
 """
-function rand_graph(n::Integer, m::Integer; bidirected = true, seed = -1, edge_weight = nothing, kws...)
+function rand_graph(n::Integer, m::Integer; seed=-1, kws...)
+    if seed != -1
+        Base.depwarn("Keyword argument `seed` is deprecated, pass an rng as first argument instead.", :rand_graph)
+        rng = MersenneTwister(seed)
+    else
+        rng = Random.default_rng()
+    end
+    return rand_graph(rng, n, m; kws...)
+end
+
+function rand_graph(rng::AbstractRNG, n::Integer, m::Integer; bidirected = true, edge_weight = nothing, kws...)
     if bidirected
-        @assert iseven(m) "Need even number of edges for bidirected graphs, given m=$m."
+        @assert iseven(m) lazy"Need even number of edges for bidirected graphs, given m=$m."
+        s, t, _ = _rand_edges(rng, n, m ÷ 2; directed=false, self_loops=false)
+        s, t = vcat(s, t), vcat(t, s)
+        if edge_weight !== nothing
+            edge_weight = vcat(edge_weight, edge_weight)
+        end
+    else
+        s, t, _ = _rand_edges(rng, n, m; directed=true, self_loops=false)
     end
-    m2 = bidirected ? m ÷ 2 : m
-    return GNNGraph(Graphs.erdos_renyi(n, m2; is_directed = !bidirected, seed); edge_weight, kws...)
+    return GNNGraph((s, t, edge_weight); kws...)
 end
 
 """
@@ -104,12 +119,6 @@ function _rand_bidirected_heterograph(rng, n::NDict, m::EDict; kws...)
     return GNNHeteroGraph(graphs; num_nodes = n, kws...)
 end
 
-function _rand_edges(rng, (n1, n2), m)
-    idx = StatsBase.sample(rng, 1:(n1 * n2), m, replace = false)
-    s, t = edge_decoding(idx, n1, n2)
-    val = nothing
-    return s, t, val
-end
 
 """
     rand_bipartite_heterograph(n1, n2, m; [bidirected, seed, node_t, edge_t, kws...])

GNNGraphs/src/utils.jl

@@ -205,17 +205,13 @@ end
 numnonzeros(a::AbstractSparseMatrix) = nnz(a)
 numnonzeros(a::AbstractMatrix) = count(!=(0), a)
 
-# 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
+## Map edges into a contiguous range of integers
+function edge_encoding(s, t, n; directed = true, self_loops = true)
+    if directed && self_loops
         maxid = n^2
-    else
-        # Undirected edges and self-loops allowed
+        idx = (s .- 1) .* n .+ t
+    elseif !directed && self_loops
         maxid = n * (n + 1) ÷ 2
-
         mask = s .> t
         snew = copy(s)
         tnew = copy(t)
@@ -228,18 +224,34 @@ function edge_encoding(s, t, n; directed = true)
         #     = ∑_{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)
+    elseif directed && !self_loops
+        @assert all(s .!= t)
+        maxid = n * (n - 1)
+        idx = (s .- 1) .* (n - 1) .+ t .- (t .> s)
+    elseif !directed && !self_loops
+        @assert all(s .!= t)
+        maxid = n * (n - 1) ÷ 2
+        mask = s .> t
+        snew = copy(s)
+        tnew = copy(t)
+        snew[mask] .= t[mask]
+        tnew[mask] .= s[mask]
+        s, t = snew, tnew
+
+        # idx(s,t) = ∑_{s',1<= s'<s} ∑_{t',s'< t' <=n} 1 + ∑_{t',s<t'<=t} 1
+        # idx(s,t) = ∑_{s',1<= s'<s} (n-s') + (t-s)
+        # idx(s,t) = (s-1)n - s*(s-1)/2 + (t-s)
+        idx = @. (s - 1) * n - s * (s - 1) ÷ 2 + (t - s)
     end
     return idx, maxid
 end
 
-# each edge is represented by a number in
-# 1:N^2
-function edge_decoding(idx, n; directed = true)
-    if directed
-        # g = remove_self_loops(g)
+# inverse of edge_encoding
+function edge_decoding(idx, n; directed = true, self_loops = true)
+    if directed && self_loops
         s = (idx .- 1) .÷ n .+ 1
         t = (idx .- 1) .% n .+ 1
-    else
+    elseif !directed && self_loops
         # We replace j=n in
         # idx = (i - 1)*(2*(n+1)-i)÷2 + (j - i + 1)
         # and obtain
@@ -252,19 +264,52 @@ function edge_decoding(idx, n; directed = true)
         s = @. ceil(Int, -sqrt((n + 1 / 2)^2 - 2 * idx) + n + 1 / 2)
         t = @. idx - (s - 1) * (2 * (n + 1) - s) ÷ 2 - 1 + s
         # t =  (idx .- 1) .% n .+ 1
+    elseif directed && !self_loops
+        s = (idx .- 1) .÷ (n - 1) .+ 1
+        t = (idx .- 1) .% (n - 1) .+ 1
+        t = t .+ (t .>= s)
+    elseif !directed && !self_loops
+        # Considering t = s + 1 in
+        # idx = @. (s - 1) * n - s * (s - 1) ÷ 2 + (t - s)
+        # and inverting for s we have
+        s = @. floor(Int, 1/2 + n - 1/2 * sqrt(9 - 4n + 4n^2 - 8*idx))
+        # now we can find t
+        t = @. idx - (s - 1) * n + s * (s - 1) ÷ 2 + s
     end
     return s, t
 end
 
-# each edge is represented by a number in
-# 1:n1*n2
+# for bipartite graphs
 function edge_decoding(idx, n1, n2)
     @assert all(1 .<= idx .<= n1 * n2)
     s = (idx .- 1) .÷ n2 .+ 1
     t = (idx .- 1) .% n2 .+ 1
     return s, t
 end
 
+function _rand_edges(rng, n::Int, m::Int; directed = true, self_loops = true)
+    idmax = if directed && self_loops
+                n^2
+            elseif !directed && self_loops
+                n * (n + 1) ÷ 2
+            elseif directed && !self_loops
+                n * (n - 1)
+            elseif !directed && !self_loops
+                n * (n - 1) ÷ 2
+            end
+    idx = StatsBase.sample(rng, 1:idmax, m, replace = false)
+    s, t = edge_decoding(idx, n; directed, self_loops)
+    val = nothing
+    return s, t, val
+end
+
+function _rand_edges(rng, (n1, n2), m)
+    idx = StatsBase.sample(rng, 1:(n1 * n2), m, replace = false)
+    s, t = edge_decoding(idx, n1, n2)
+    val = nothing
+    return s, t, val
+end
+
 binarize(x) = map(>(0), x)
 
 @non_differentiable binarize(x...)

GNNGraphs/test/utils.jl

@@ -47,6 +47,50 @@
     tnew[mask] .= s1[mask]
     @test sdec == snew
     @test tdec == tnew
+
+    @testset "directed=false, self_loops=false" begin
+        n = 5
+        edges = [(1,2), (3,1), (1,4), (1,5), (2,3), (2,4), (2,5), (3,4), (3,5), (4,5)]
+        s = [e[1] for e in edges]
+        t = [e[2] for e in edges]
+        g = GNNGraph(s, t)
+        idx, idxmax = GNNGraphs.edge_encoding(s, t, n, directed=false, self_loops=false)
+        @test idxmax == n * (n - 1) ÷ 2
+        @test idx == 1:idxmax
+        
+        snew, tnew = GNNGraphs.edge_decoding(idx, n, directed=false, self_loops=false)
+        @test snew == [1, 1, 1, 1, 2, 2, 2, 3, 3, 4]
+        @test tnew == [2, 3, 4, 5, 3, 4, 5, 4, 5, 5]
+    end
+
+    @testset "directed=false, self_loops=false" begin
+        n = 5
+        edges = [(1,2), (3,1), (1,4), (1,5), (2,3), (2,4), (2,5), (3,4), (3,5), (4,5)]
+        s = [e[1] for e in edges]
+        t = [e[2] for e in edges]
+
+        idx, idxmax = GNNGraphs.edge_encoding(s, t, n, directed=false, self_loops=false)
+        @test idxmax == n * (n - 1) ÷ 2
+        @test idx == 1:idxmax
+        
+        snew, tnew = GNNGraphs.edge_decoding(idx, n, directed=false, self_loops=false)
+        @test snew == [1, 1, 1, 1, 2, 2, 2, 3, 3, 4]
+        @test tnew == [2, 3, 4, 5, 3, 4, 5, 4, 5, 5]
+    end
+
+    @testset "directed=true, self_loops=false" begin
+        n = 5
+        edges = [(1,2), (3,1), (1,4), (1,5), (2,3), (2,4), (2,5), (3,4), (3,5), (4,5)]
+        s = [e[1] for e in edges]
+        t = [e[2] for e in edges]
+
+        idx, idxmax = GNNGraphs.edge_encoding(s, t, n, directed=true, self_loops=false)
+        @test idxmax == n^2 - n
+        @test idx == [1, 9, 3, 4, 6, 7, 8, 11, 12, 16]
+        snew, tnew = GNNGraphs.edge_decoding(idx, n, directed=true, self_loops=false)
+        @test snew == s
+        @test tnew == t
+    end
 end
 
 @testset "color_refinment" begin

test/runtests.jl

@@ -35,7 +35,9 @@ tests = [
 !CUDA.functional() && @warn("CUDA unavailable, not testing GPU support")
 
 # @testset "GraphNeuralNetworks: graph format $graph_type" for graph_type in (:coo, :dense, :sparse)
-for graph_type in (:coo, :dense, :sparse)
+# for graph_type in (:coo, :dense, :sparse)
+for graph_type in (:dense,)
+
     @info "Testing graph format :$graph_type"
     global GRAPH_T = graph_type
     global TEST_GPU = CUDA.functional() && (GRAPH_T != :sparse)