negative_sample

Author: CarloLucibello

examples/link_prediction_pubmed.jl

@@ -55,13 +55,14 @@ function train(; kws...)
     @show has_self_loops(g)
     @show has_multi_edges(g)
     @show mean(degree(g))
+    isbidir = is_bidirected(g)  
 
     g = g |> device
     X = data.node_features |> device
 
     #### SPLIT INTO NEGATIVE AND POSITIVE SAMPLES
     train_pos_g, test_pos_g = rand_edge_split(g, 0.9)
-    test_neg_g = negative_sample(g, num_neg_edges=test_pos_g.num_edges)
+    test_neg_g = negative_sample(g, num_neg_edges=test_pos_g.num_edges, bidirected=isbidir)
 
     ### DEFINE MODEL #########
     nin, nhidden = size(X,1), args.nhidden
@@ -82,7 +83,7 @@ function train(; kws...)
         h = model(X)
         if neg_g === nothing
             # We sample a negative graph at each training step
-            neg_g = negative_sample(pos_g)
+            neg_g = negative_sample(pos_g, bidirected=isbidir)
         end
         pos_score = pred(pos_g, h)
         neg_score = pred(neg_g, h)

src/GNNGraphs/transform.jl

@@ -321,13 +321,21 @@ function getgraph(g::GNNGraph, i::AbstractVector{Int}; nmap=false)
 end
 
 """
-    negative_sample(g::GNNGraph; num_neg_edges=g.num_edges)
+    negative_sample(g::GNNGraph; 
+                    num_neg_edges = g.num_edges, 
+                    bidirected = is_bidirected(g))
 
 Return a graph containing random negative edges (i.e. non-edges) from graph `g` as edges.
+
+Is `bidirected=true`, the output graph will be bidirected and there will be no
+leakage from the origin graph. 
+
+See also [`is_bidirected`](@ref).
 """
 function negative_sample(g::GNNGraph; 
         max_trials=3, 
-        num_neg_edges=g.num_edges)
+        num_neg_edges=g.num_edges, 
+        bidirected = is_bidirected(g))
 
     @assert g.num_graphs == 1
     # Consider self-loops as positive edges
@@ -344,8 +352,12 @@ function negative_sample(g::GNNGraph;
         device = Flux.cpu
     end
     idx_pos, maxid = edge_encoding(s, t, n)
-    
-    pneg = 1 - g.num_edges / maxid # prob of selecting negative edge 
+    if bidirected
+        num_neg_edges = num_neg_edges ÷ 2
+        pneg = 1 - g.num_edges / 2maxid # prob of selecting negative edge 
+    else 
+        pneg = 1 - g.num_edges / 2maxid # prob of selecting negative edge 
+    end    
     # pneg * sample_prob * maxid == num_neg_edges  
     sample_prob = min(1, num_neg_edges / (pneg * maxid) * 1.1)
     idx_neg = Int[]
@@ -359,6 +371,9 @@ function negative_sample(g::GNNGraph;
         end
     end
     s_neg, t_neg = edge_decoding(idx_neg, n)
+    if bidirected
+        s_neg, t_neg = [s_neg; t_neg], [t_neg; s_neg] 
+    end
     return GNNGraph(s_neg, t_neg, num_nodes=n) |> device
 end
 
@@ -372,6 +387,7 @@ while `g2` wil contain the rest.
 Useful for train/test splits in link prediction tasks.
 """
 function rand_edge_split(g::GNNGraph, frac)
+    # TODO add bidirected version
     s, t = edge_index(g)
     eids = randperm(g.num_edges)
     size1 = round(Int, g.num_edges * frac)

test/GNNGraphs/utils.jl

@@ -27,5 +27,7 @@
         sdec, tdec = GNNGraphs.edge_decoding(idx, n, directed=false) 
         @test sdec == snew
         @test tdec == tnew
+
+        g = rand_graph(10, 30, bidirected=true)
     end
 end