support batched graphs

Author: CarloLucibello

src/GNNGraphs/convert.jl

@@ -69,7 +69,7 @@ function to_dense(A::ADJMAT_T, T=nothing; dir=:out, num_nodes=nothing)
     num_nodes = size(A, 1)
     @assert num_nodes == size(A, 2)
     # @assert all(x -> (x == 1) || (x == 0), A)
-    num_edges = round(Int, sum(A))
+    num_edges = numnonzeros(A)
     if dir == :in
         A = A'
     end
@@ -85,7 +85,7 @@ function to_dense(adj_list::ADJLIST_T, T=nothing; dir=:out, num_nodes=nothing)
     num_edges = sum(length.(adj_list))
     @assert num_nodes > 0
     T = T === nothing ? eltype(adj_list[1]) : T
-    A = similar(adj_list[1], T, (num_nodes, num_nodes))
+    A = fill!(similar(adj_list[1], T, (num_nodes, num_nodes)), 0)
     if dir == :out
         for (i, neigs) in enumerate(adj_list)
             A[i, neigs] .= 1
@@ -129,13 +129,13 @@ function to_sparse(A::ADJMAT_T, T=nothing; dir=:out, num_nodes=nothing)
     end
     if !(A isa AbstractSparseMatrix)
         A = sparse(A)
-    end 
+    end
     return A, num_nodes, num_edges
 end
 
 function to_sparse(adj_list::ADJLIST_T, T=nothing; dir=:out, num_nodes=nothing)
     coo, num_nodes, num_edges = to_coo(adj_list; dir, num_nodes)
-    return to_sparse(coo; dir, num_nodes)
+    return to_sparse(coo; num_nodes)
 end
 
 function to_sparse(coo::COO_T, T=nothing; dir=:out, num_nodes=nothing)

src/GNNGraphs/generate.jl

@@ -12,7 +12,7 @@ Use a `seed > 0` for reproducibility.
 
 Additional keyword arguments will be passed to the [`GNNGraph`](@ref) constructor.
 
-# Usage
+# Examples
 
 ```juliarepl
 julia> g = rand_graph(5, 4, bidirected=false)
@@ -46,15 +46,82 @@ function rand_graph(n::Integer, m::Integer; bidirected=true, seed=-1, kws...)
 end
 
 
-function knn_graph(points::AbstractMatrix, k::Int; self_loops=false, dir=:in, kws...)
+"""
+    knn_graph(points::AbstractMatrix, 
+              k::Int; 
+              graph_indicator = nothing,
+              self_loops = false, 
+              dir = :in, 
+              kws...)
+
+Create a `k`-nearest neighbor graph where each node is linked 
+to its `k` closest `points`.  
+
+# Arguments
+
+- `points`: A num_features × num_nodes matrix storing the Euclidean positions of the nodes.
+- `k`: The number of neighbors considered in the kNN algorithm.
+- `graph_indicator`: Either nothing or a vector containing the graph assigment of each node, 
+                     in which case the returned graph will be a batch of graphs. 
+- `self_loops`: If `true`, consider the node itself among its `k` nearest neighbors, in which
+                case the graph will contain self-loops. 
+- `dir`: The direction of the edges. If `dir=:in` edges go from the `k` 
+         neighbors to the central node. If `dir=:out` we have the opposite
+         direction.
+- `kws`: Further keyword arguments will be passed to the [`GNNGraph ](@ref) constructor.
+
+# Examples
+
+```juliarepl
+julia> n, k = 10, 3;
+
+julia> x = rand(3, n);
+
+julia> g = knn_graph(x, k)
+GNNGraph:
+    num_nodes = 10
+    num_edges = 30
+
+julia> graph_indicator = [1,1,1,1,1,2,2,2,2,2];
+
+julia> g = knn_graph(x, k; graph_indicator)
+GNNGraph:
+    num_nodes = 10
+    num_edges = 30
+    num_graphs = 2
+
+```
+"""
+function knn_graph(points::AbstractMatrix, k::Int; 
+        graph_indicator = nothing,
+        self_loops = false, 
+        dir = :in, 
+        kws...)
+
+    if graph_indicator !== nothing
+        d, n = size(points)
+        @assert graph_indicator isa AbstractVector{<:Integer}
+        @assert length(graph_indicator) == n
+        # All graphs in the batch must have at least k nodes. 
+        cm = StatsBase.countmap(graph_indicator)
+        @assert all(values(cm) .>= k)
+        
+        # Make sure that the distance between points in different graphs
+        # is always larger than any distance within the same graph.
+        points = points .- minimum(points)
+        points = points ./ maximum(points)
+        dummy_feature = 2d .* reshape(graph_indicator, 1, n)
+        points = vcat(points, dummy_feature)
+    end
+    
     kdtree = NearestNeighbors.KDTree(points)
-    sortres = false
     if !self_loops
         k += 1
     end
+    sortres = false
     idxs, dists = NearestNeighbors.knn(kdtree, points, k, sortres)
-    # return idxs
-    g = GNNGraph(idxs; dir, kws...)
+    
+    g = GNNGraph(idxs; dir, graph_indicator, kws...)
     if !self_loops
         g = remove_self_loops(g)
     end

src/GNNGraphs/gnngraph.jl

@@ -65,7 +65,7 @@ functionality from that library.
 - `edata`: Edge features. A named tuple of arrays whose last dimension has size `num_edges`.
 - `gdata`: Graph features. A named tuple of arrays whose last dimension has size `num_graphs`. 
 
-# Usage. 
+# Examples 
 
 ```julia
 using Flux, GraphNeuralNetworks
@@ -201,7 +201,7 @@ function Base.show(io::IO, g::GNNGraph)
     print(io, "GNNGraph:
     num_nodes = $(g.num_nodes)
     num_edges = $(g.num_edges)")
-    g.num_graphs > 1 && print("\nnum_graphs = $(g.num_graphs)")
+    g.num_graphs > 1 && print("\n    num_graphs = $(g.num_graphs)")
     if !isempty(g.ndata)
         print(io, "\n    ndata:")
         for k in keys(g.ndata)

src/GNNGraphs/transform.jl

@@ -55,6 +55,22 @@ function remove_self_loops(g::GNNGraph{<:COO_T})
             g.ndata, g.edata, g.gdata)
 end
 
+
+function remove_self_loops(g::GNNGraph{<:ADJMAT_T})
+    @assert g.edata === (;)
+    A = g.graph
+    A[diagind(A)] .= 0
+    if A isa AbstractSparseMatrix
+        dropzeros!(A)
+    end
+    num_edges = numnonzeros(A)
+    GNNGraph(A, 
+            g.num_nodes, num_edges, g.num_graphs, 
+            g.graph_indicator,
+            g.ndata, g.edata, g.gdata)
+end
+
+
 """
     remove_multi_edges(g::GNNGraph)
 
@@ -183,7 +199,7 @@ containing the total number of original nodes and edges.
 Equivalent to [`SparseArrays.blockdiag`](@ref).
 See also [`Flux.unbatch`](@ref).
 
-# Usage
+# Examples
 
 ```juliarepl
 julia> g1 = rand_graph(4, 6, ndata=ones(8, 4))
@@ -231,7 +247,7 @@ an array of the individual graphs batched together in `g`.
 
 See also [`Flux.batch`](@ref) and [`getgraph`](@ref).
 
-# Usage
+# Examples
 
 ```juliarepl
 julia> gbatched = Flux.batch([rand_graph(5, 6), rand_graph(10, 8), rand_graph(4,2)])

src/GNNGraphs/utils.jl

@@ -71,6 +71,8 @@ ones_like(x::AbstractArray, T=eltype(x), sz=size(x)) = fill!(similar(x, T, sz),
 ones_like(x::SparseMatrixCSC, T=eltype(x), sz=size(x)) = ones(T, sz)
 ones_like(x::CUMAT_T, T=eltype(x), sz=size(x)) = CUDA.ones(T, sz)
 
+numnonzeros(a::AbstractSparseMatrix) = nnz(a)
+numnonzeros(a::AbstractMatrix) = count(!=(0), a)
 
 # each edge is represented by a number in
 # 1:N^2

src/layers/conv.jl

@@ -713,7 +713,7 @@ as the input size.
 - `init`: Weights' initializer.
 - `residual`: Add a residual connection.
 
-# Usage 
+# Examples 
 
 ```julia
 x = rand(Float32, 2, g.num_nodes)

src/msgpass.jl

@@ -31,7 +31,7 @@ providing as input `f` a closure.
       with the same batch size.
 - `aggr`: Neighborhood aggregation operator. Use `+`, `mean`, `max`, or `min`. 
 
-# Usage Examples
+# Examples
 
 ```julia
 using GraphNeuralNetworks, Flux

test/GNNGraphs/generate.jl

@@ -26,8 +26,7 @@
     end
 
     @testset "knn_graph" begin
-        n = 10
-        k = 3
+        n, k = 10, 3
         x = rand(3, n)
         g = knn_graph(x, k; graph_type=GRAPH_T)
         @test g.num_nodes == 10
@@ -40,5 +39,15 @@
         @test g.num_edges == n*k
         @test degree(g, dir=:out) == fill(k, n)        
         @test has_self_loops(g) == true 
+
+        graph_indicator = [1,1,1,1,1,2,2,2,2,2]
+        g = knn_graph(x, k; graph_indicator, graph_type=GRAPH_T)
+        @test g.num_graphs == 2
+        s, t = edge_index(g)
+        ne = n*k÷2
+        @test all(1 .<= s[1:ne] .<= 5)
+        @test all(1 .<= t[1:ne] .<= 5)     
+        @test all(6 .<= s[ne+1:end] .<= 10)
+        @test all(6 .<= t[ne+1:end] .<= 10)     
     end
 end