Implemented khop_adj (#239)

* Add khop_adj function

* Add tests khop_adj function

* Add export khop_adj

* Align with adjacency_matrix signature

* Add more tests

* Cleaner inputs

Co-authored-by: Carlo Lucibello <[email protected]>

Co-authored-by: Carlo Lucibello <[email protected]>

Author: aurorarossi

src/GNNGraphs/GNNGraphs.jl

@@ -41,7 +41,8 @@ export adjacency_list,
        has_self_loops,
        has_isolated_nodes,
        inneighbors,
-       outneighbors 
+       outneighbors,
+       khop_adj 
 
 include("transform.jl")
 export add_nodes,

src/GNNGraphs/query.jl

@@ -404,6 +404,17 @@ function has_multi_edges(g::GNNGraph)
     length(union(idxs)) < length(idxs)
 end
 
+"""
+    khop_adj(g::GNNGraph,k::Int,T::DataType=eltype(g); dir=:out, weighted=true)
+
+Return ``A^k`` where ``A`` is the adjacency matrix of the graph 'g'.
+
+"""
+function khop_adj(g::GNNGraph,k::Int, T::DataType=eltype(g); dir=:out, weighted=true)
+    return (adjacency_matrix(g, T; dir, weighted))^k
+end
+
+
 @non_differentiable edge_index(x...)
 @non_differentiable adjacency_list(x...)
 @non_differentiable graph_indicator(x...)

test/GNNGraphs/query.jl

@@ -164,5 +164,18 @@
 
             @test gw == [1,1,1]
         end
+
+        @testset "khop_adj" begin
+            s = [1,2,3]
+            t = [2,3,1]
+            w = [0.1,0.1,0.2]
+            g = GNNGraph(s, t, w)
+            @test khop_adj(g,2)== adjacency_matrix(g)*adjacency_matrix(g)
+            @test khop_adj(g,2,Int8;weighted=false) == sparse([0 0 1;1 0 0;0 1 0])
+            @test khop_adj(g,2,Int8;dir=in,weighted=false) == sparse([0 0 1;1 0 0;0 1 0]')
+            @test khop_adj(g,1) == adjacency_matrix(g)
+            @test eltype(khop_adj(g,4)) == Float64
+            @test eltype(khop_adj(g,10,Float32)) == Float32    
+        end
     end
 end