wip on rand_split_edge

Author: CarloLucibello

examples/link_prediction_pubmed.jl

@@ -28,7 +28,7 @@ struct DotPredictor end
 
 function (::DotPredictor)(g, x)
     z = apply_edges((xi, xj, e) -> sum(xi .* xj, dims=1), g, xi=x, xj=x)
-    # z = apply_edges(xi_dot_xj, g, xi=x, xj=x) # Same with  buit-in methods
+    # z = apply_edges(xi_dot_xj, g, xi=x, xj=x) # Same with built-in method
     return vec(z)
 end
 
@@ -50,18 +50,23 @@ function train(; kws...)
     data = Cora.dataset()
     # data = PubMed.dataset()
     g = GNNGraph(data.adjacency_list)
+
+    # Print some info
     @info g
     @show is_bidirected(g)
     @show has_self_loops(g)
     @show has_multi_edges(g)
     @show mean(degree(g))
     isbidir = is_bidirected(g)  
 
+    # Move to device
     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)
+    #### TRAIN/TEST splits
+    # With bidirected graph, we make sure that an edge and its reverse
+    # are in the same split 
+    train_pos_g, test_pos_g = rand_edge_split(g, 0.9, bidirected=isbidir)
     test_neg_g = negative_sample(g, num_neg_edges=test_pos_g.num_edges, bidirected=isbidir)
 
     ### DEFINE MODEL #########