rand_bipartite

Author: CarloLucibello

GNNGraphs/src/generate.jl

@@ -48,7 +48,9 @@ function rand_graph(n::Integer, m::Integer; seed=-1, kws...)
     return rand_graph(rng, n, m; kws...)
 end
 
-function rand_graph(rng::AbstractRNG, n::Integer, m::Integer; bidirected = true, edge_weight = nothing, kws...)
+function rand_graph(rng::AbstractRNG, n::Integer, m::Integer; 
+            bidirected::Bool = true, 
+            edge_weight::Union{AbstractVector, Nothing} = nothing, kws...)
     if bidirected
         @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)
@@ -63,13 +65,13 @@ function rand_graph(rng::AbstractRNG, n::Integer, m::Integer; bidirected = true,
 end
 
 """
-    rand_heterograph(n, m; seed=-1, bidirected=false, kws...)
+    rand_heterograph([rng,] n, m; bidirected=false, kws...)
 
-Construct an [`GNNHeteroGraph`](@ref) with number of nodes and edges 
+Construct an [`GNNHeteroGraph`](@ref) with random edges and with number of nodes and edges 
 specified by `n` and `m` respectively. `n` and `m` can be any iterable of pairs
 specifing node/edge types and their numbers.
 
-Use a `seed > 0` for reproducibility.
+Pass a random number generator as a first argument to make the generation reproducible.
 
 Setting `bidirected=true` will generate a bidirected graph, i.e. each edge will have a reverse edge.
 Therefore, for each edge type `(:A, :rel, :B)` a corresponding reverse edge type `(:B, :rel, :A)`
@@ -92,16 +94,25 @@ function rand_heterograph end
 # for generic iterators of pairs
 rand_heterograph(n, m; kws...) = rand_heterograph(Dict(n), Dict(m); kws...)
 
-function rand_heterograph(n::NDict, m::EDict; bidirected = false, seed = -1, kws...)
-    rng = seed > 0 ? MersenneTwister(seed) : Random.GLOBAL_RNG
+function  rand_heterograph(n::NDict, m::EDict; seed=-1, kws...)
+    if seed != -1
+        Base.depwarn("Keyword argument `seed` is deprecated, pass an rng as first argument instead.", :rand_heterograph)
+        rng = MersenneTwister(seed)
+    else
+        rng = Random.default_rng()
+    end
+    return rand_heterograph(rng, n, m; kws...)
+end
+
+function rand_heterograph(rng::AbstractRNG, n::NDict, m::EDict; bidirected::Bool = false, kws...)
     if bidirected
         return _rand_bidirected_heterograph(rng, n, m; kws...)
     end
     graphs = Dict(k => _rand_edges(rng, (n[k[1]], n[k[3]]), m[k]) for k in keys(m))
     return GNNHeteroGraph(graphs; num_nodes = n, kws...)
 end
 
-function _rand_bidirected_heterograph(rng, n::NDict, m::EDict; kws...)
+function _rand_bidirected_heterograph(rng::AbstractRNG, n::NDict, m::EDict; kws...)
     for k in keys(m)
         if reverse(k) ∈ keys(m)
             @assert m[k] == m[reverse(k)] "Number of edges must be the same in reverse edge types for bidirected graphs."
@@ -121,35 +132,58 @@ end
 
 
 """
-    rand_bipartite_heterograph(n1, n2, m; [bidirected, seed, node_t, edge_t, kws...])
-    rand_bipartite_heterograph((n1, n2), m; ...)
-    rand_bipartite_heterograph((n1, n2), (m1, m2); ...)
+    rand_bipartite_heterograph([rng,] 
+                               (n1, n2), (m12, m21); 
+                               bidirected = true, 
+                               node_t = (:A, :B), 
+                               edge_t = :to, 
+                               kws...)
 
-Construct an [`GNNHeteroGraph`](@ref) with number of nodes and edges
-specified by `n1`, `n2` and `m1` and `m2` respectively.
+Construct an [`GNNHeteroGraph`](@ref) with random edges representing a bipartite graph.
+The graph will have two types of nodes, and edges will only connect nodes of different types.
 
-See [`rand_heterograph`](@ref) for a more general version.
+The first argument is a tuple `(n1, n2)` specifying the number of nodes of each type.
+The second argument is a tuple `(m12, m21)` specifying the number of edges connecting nodes of type `1` to nodes of type `2` 
+and vice versa.
 
-# Keyword arguments
+The type of nodes and edges can be specified with the `node_t` and `edge_t` keyword arguments,
+which default to `(:A, :B)` and `:to` respectively.
 
-- `bidirected`: whether to generate a bidirected graph. Default is `true`.
-- `seed`: random seed. Default is `-1` (no seed).
-- `node_t`: node types. If `bipartite=true`, this should be a tuple of two node types, otherwise it should be a single node type.
-- `edge_t`: edge types. If `bipartite=true`, this should be a tuple of two edge types, otherwise it should be a single edge type.
-"""
-function rand_bipartite_heterograph end
+If `bidirected=true` (default), the reverse edge of each edge will be present. In this case
+`m12 == m21` is required.
 
-rand_bipartite_heterograph(n1::Int, n2::Int, m::Int; kws...) = rand_bipartite_heterograph((n1, n2), (m, m); kws...)
+A random number generator can be passed as the first argument to make the generation reproducible.
 
-rand_bipartite_heterograph((n1, n2)::NTuple{2,Int}, m::Int; kws...) = rand_bipartite_heterograph((n1, n2), (m, m); kws...)
+Additional keyword arguments will be passed to the [`GNNHeteroGraph`](@ref) constructor.
 
-function rand_bipartite_heterograph((n1, n2)::NTuple{2,Int}, (m1, m2)::NTuple{2,Int}; bidirected=true, 
-                        node_t = (:A, :B), edge_t = :to, kws...)
-    if edge_t isa Symbol
-        edge_t = (edge_t, edge_t)
+See [`rand_heterograph`](@ref) for a more general version.
+
+# Examples
+
+```julia-repl
+julia> g = rand_bipartite_heterograph((10, 15), 20)
+GNNHeteroGraph:
+  num_nodes: (:A => 10, :B => 15)
+  num_edges: ((:A, :to, :B) => 20, (:B, :to, :A) => 20)
+
+julia> g = rand_bipartite_heterograph((10, 15), (20, 0), node_t=(:user, :item), edge_t=:-, bidirected=false)
+GNNHeteroGraph:
+  num_nodes: Dict(:item => 15, :user => 10)
+  num_edges: Dict((:item, :-, :user) => 0, (:user, :-, :item) => 20)
+```
+"""
+rand_bipartite_heterograph(n, m; kws...) = rand_bipartite_heterograph(Random.default_rng(), n, m; kws...)
+
+function rand_bipartite_heterograph(rng::AbstractRNG, (n1, n2)::NTuple{2,Int}, m; bidirected=true, 
+                        node_t = (:A, :B), edge_t::Symbol = :to, kws...)
+    if m isa Integer
+        m12 = m21 = m
+    else
+        m12, m21 = m
     end
-    return rand_heterograph(Dict(node_t[1] => n1, node_t[2] => n2), 
-                            Dict((node_t[1], edge_t[1], node_t[2]) => m1, (node_t[2], edge_t[2], node_t[1]) => m2); 
+
+    return rand_heterograph(rng, Dict(node_t[1] => n1, node_t[2] => n2), 
+                            Dict((node_t[1], edge_t, node_t[2]) => m12, (node_t[2], edge_t, node_t[1]) => m21); 
                             bidirected, kws...)
 end
 

GNNGraphs/test/generate.jl

@@ -81,7 +81,7 @@ end
 end
 
 @testset "rand_bipartite_heterograph" begin
-    g = rand_bipartite_heterograph(10, 15, 20)
+    g = rand_bipartite_heterograph((10, 15), (20, 20))
     @test g.num_nodes == Dict(:A => 10, :B => 15)
     @test g.num_edges == Dict((:A, :to, :B) => 20, (:B, :to, :A) => 20)
     sA, tB = edge_index(g, (:A, :to, :B))

GNNGraphs/test/gnnheterograph.jl

@@ -123,7 +123,7 @@ end
 
 @testset "get/set node features" begin
     d, n = 3, 5
-    g = rand_bipartite_heterograph(n, 2*n, 15)
+    g = rand_bipartite_heterograph((n, 2*n), 15)
     g[:A].x = rand(Float32, d, n)
     g[:B].y = rand(Float32, d, 2*n)
 
@@ -133,7 +133,7 @@ end
 
 @testset "add_edges" begin
     d, n = 3, 5
-    g = rand_bipartite_heterograph(n, 2 * n, 15)
+    g = rand_bipartite_heterograph((n, 2 * n), 15)
     s, t = [1, 2, 3], [3, 2, 1]
     ## Keep the same ntypes - construct with args
     g1 = add_edges(g, (:A, :rel1, :B), s, t)

GNNLux/test/layers/basic_tests.jl

@@ -1,6 +1,6 @@
 @testitem "layers/basic" setup=[SharedTestSetup] begin
     rng = StableRNG(17)
-    g = rand_graph(10, 40, seed=17)
+    g = rand_graph(rng, 10, 40)
     x = randn(rng, Float32, 3, 10)        
 
     @testset "GNNLayer" begin

GNNLux/test/layers/conv_tests.jl

@@ -1,6 +1,6 @@
 @testitem "layers/conv" setup=[SharedTestSetup] begin
     rng = StableRNG(1234)
-    g = rand_graph(10, 40, seed=1234)
+    g = rand_graph(rng, 10, 40)
     in_dims = 3
     out_dims = 5
     x = randn(rng, Float32, in_dims, 10)

test/layers/heteroconv.jl

@@ -1,6 +1,6 @@
 @testset "HeteroGraphConv" begin
     d, n = 3, 5
-    g = rand_bipartite_heterograph(n, 2*n, 15)
+    g = rand_bipartite_heterograph((n, 2*n), 15)
     hg = rand_bipartite_heterograph((2,3), 6)
 
     model = HeteroGraphConv([(:A,:to,:B) => GraphConv(d => d),