better negative sampling

Author: CarloLucibello

docs/src/index.md

@@ -1,17 +1,17 @@
 # GraphNeuralNetworks
 
-This is the documentation page for the [GraphNeuralNetworks.jl](https://github.com/CarloLucibello/GraphNeuralNetworks.jl) library.
-
-A graph neural network library for Julia based on the deep learning framework [Flux.jl](https://github.com/FluxML/Flux.jl). GNN.jl is largely inspired by python's libraries [PyTorch Geometric](https://pytorch-geometric.readthedocs.io/en/latest/) and [Deep Graph Library](https://docs.dgl.ai/),
-and by julia's [GeometricFlux](https://fluxml.ai/GeometricFlux.jl/stable/).
+This is the documentation page for [GraphNeuralNetworks.jl](https://github.com/CarloLucibello/GraphNeuralNetworks.jl), a graph neural network library written in Julia and based on the deep learning framework [Flux.jl](https://github.com/FluxML/Flux.jl).
+GNN.jl is largely inspired by [PyTorch Geometric](https://pytorch-geometric.readthedocs.io/en/latest/),[Deep Graph Library](https://docs.dgl.ai/),
+and [GeometricFlux.jl](https://fluxml.ai/GeometricFlux.jl/stable/).
 
 Among its features:
 
-* Integratation with the JuliaGraphs ecosystem.
+* Integration with the JuliaGraphs ecosystem.
 * Implementation of common graph convolutional layers.
 * Fast operations on batched graphs. 
 * Easy to define custom layers.
 * CUDA support.
+* Examples of node-wise, edge-wise, and graph-wise machine learning tasks. 
 
 
 ## Package overview

src/GNNGraphs/transform.jl

@@ -328,15 +328,60 @@ end
 """
     negative_sample(g::GNNGraph; num_neg_edges=g.num_edges)
 
-Return a graph containing random negative edges (i.e. non-edges) from graph `g`.
+Return a graph containing random negative edges (i.e. non-edges) from graph `g` as edges.
 """
-function negative_sample(g::GNNGraph; num_neg_edges=g.num_edges)
-    adj = adjacency_matrix(g)
-    adj_neg = 1 .- adj - I
-    neg_s, neg_t = ci2t(findall(adj_neg .> 0), 2)
-    neg_eids = randperm(length(neg_s))[1:num_neg_edges]
-    neg_s, neg_t = neg_s[neg_eids], neg_t[neg_eids]
-    return GNNGraph(neg_s, neg_t, num_nodes=g.num_nodes)
+function negative_sample(g::GNNGraph; 
+        max_trials=3, 
+        num_neg_edges=g.num_edges)
+
+    @assert g.num_graphs == 1
+    # Consider self-loops as positive edges
+    # Construct new graph dropping features
+    g = add_self_loops(GNNGraph(edge_index(g))) 
+    
+    s, t = edge_index(g)
+    n = g.num_nodes
+    if s isa CuArray
+        # Convert to gpu since set operations and sampling are not supported by CUDA.jl
+        device = Flux.gpu 
+        s, t = Flux.cpu(s), Flux.cpu(t) 
+    else 
+        device = Flux.cpu
+    end
+    idx_pos, maxid = edge_encoding(s, t, n)
+    
+    pneg = 1 - g.num_edges / maxid # prob of selecting negative edge 
+    # pneg * sample_prob * maxid == num_neg_edges  
+    sample_prob = min(1, num_neg_edges / (pneg * maxid) * 1.1)
+    idx_neg = Int[]
+    for _ in 1:max_trials
+        rnd = randsubseq(1:maxid, sample_prob)
+        setdiff!(rnd, idx_pos)
+        union!(idx_neg, rnd)
+        if length(idx_neg) >= num_neg_edges
+            idx_neg = idx_neg[1:num_neg_edges]
+            break
+        end
+    end
+    s_neg, t_neg = edge_decoding(idx_neg, n)
+    return GNNGraph(s_neg, t_neg, num_nodes=n) |> device
+end
+
+# each edge is represented by a number in
+# 1:N^2
+function edge_encoding(s, t, n)
+    idx = (s .- 1) .* n .+ t
+    maxid = n^2 
+    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
 
 # """