implement knn_graph

Author: CarloLucibello

Project.toml

@@ -17,6 +17,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 NNlibCUDA = "a00861dc-f156-4864-bf3c-e6376f28a68d"
+NearestNeighbors = "b8a86587-4115-5ab1-83bc-aa920d37bbce"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
@@ -34,6 +35,7 @@ Graphs = "1.4"
 KrylovKit = "0.5"
 LearnBase = "0.4, 0.5"
 MacroTools = "0.5"
+NearestNeighbors = "0.4"
 NNlib = "0.7"
 NNlibCUDA = "0.1"
 Reexport = "1"

src/GNNGraphs/GNNGraphs.jl

@@ -7,6 +7,7 @@ import Graphs
 using Graphs: AbstractGraph, outneighbors, inneighbors, adjacency_matrix, degree, has_self_loops, is_directed
 import Flux
 using Flux: batch
+import NearestNeighbors
 import NNlib
 import LearnBase
 import StatsBase
@@ -53,7 +54,8 @@ export add_nodes,
        blockdiag
 
 include("generate.jl")
-export rand_graph
+export rand_graph, 
+       knn_graph
 
 include("operators.jl")
 # Base.intersect

src/GNNGraphs/generate.jl

@@ -44,3 +44,19 @@ function rand_graph(n::Integer, m::Integer; bidirected=true, seed=-1, kws...)
     m2 = bidirected ? m÷2 : m
     return GNNGraph(Graphs.erdos_renyi(n, m2; is_directed=!bidirected, seed); kws...)    
 end
+
+
+function knn_graph(points::AbstractMatrix, k::Int; self_loops=false, dir=:in, kws...)
+    kdtree = NearestNeighbors.KDTree(points)
+    sortres = false
+    if !self_loops
+        k += 1
+    end
+    idxs, dists = NearestNeighbors.knn(kdtree, points, k, sortres)
+    # return idxs
+    g = GNNGraph(idxs; dir, kws...)
+    if !self_loops
+        g = remove_self_loops(g)
+    end
+    return g
+end

src/GNNGraphs/gnngraph.jl

@@ -198,26 +198,26 @@ function GNNGraph(g::GNNGraph; ndata=g.ndata, edata=g.edata, gdata=g.gdata, grap
 end
 
 function Base.show(io::IO, g::GNNGraph)
-    println(io, "GNNGraph:
+    print(io, "GNNGraph:
     num_nodes = $(g.num_nodes)
     num_edges = $(g.num_edges)")
-    g.num_graphs > 1 && println("num_graphs = $(g.num_graphs)")
+    g.num_graphs > 1 && print("\nnum_graphs = $(g.num_graphs)")
     if !isempty(g.ndata)
-        println(io, "    ndata:")
+        print(io, "\n    ndata:")
         for k in keys(g.ndata)
-            println(io, "        $k => $(size(g.ndata[k]))")
+            print(io, "\n        $k => $(size(g.ndata[k]))")
         end
     end
     if !isempty(g.edata)
-        println(io, "    edata:")
+        print(io, "\n    edata:")
         for k in keys(g.edata)
-            println(io, "        $k => $(size(g.edata[k]))")
+            print(io, "\n        $k => $(size(g.edata[k]))")
         end
     end
     if !isempty(g.gdata)
-        println(io, "    gdata:")
+        print(io, "\n    gdata:")
         for k in keys(g.gdata)
-            println(io, "        $k => $(size(g.gdata[k]))")
+            print(io, "\n        $k => $(size(g.gdata[k]))")
         end
     end
 end

test/GNNGraphs/generate.jl

@@ -24,4 +24,21 @@
         g2 = rand_graph(n, m, bidirected=false, seed=17, graph_type=GRAPH_T)
         @test edge_index(g2) == edge_index(g)
     end
+
+    @testset "knn_graph" begin
+        n = 10
+        k = 3
+        x = rand(3, n)
+        g = knn_graph(x, k; graph_type=GRAPH_T)
+        @test g.num_nodes == 10
+        @test g.num_edges == n*k
+        @test degree(g, dir=:in) == fill(k, n)        
+        @test has_self_loops(g) == false 
+
+        g = knn_graph(x, k; dir=:out, self_loops=true, graph_type=GRAPH_T)
+        @test g.num_nodes == 10
+        @test g.num_edges == n*k
+        @test degree(g, dir=:out) == fill(k, n)        
+        @test has_self_loops(g) == true 
+    end
 end