Add random_walk_pe

src/GNNGraphs/GNNGraphs.jl

@@ -62,6 +62,7 @@ export add_nodes,
        set_edge_weight,
        to_bidirected,
        to_unidirected,
+       add_RandomWalkPE!,
 # from Flux
        batch,
        unbatch,

src/GNNGraphs/transform.jl

@@ -694,6 +694,25 @@ function rand_edge_split(g::GNNGraph, frac; bidirected = is_bidirected(g))
     return g1, g2
 end
 
+function add_RandomWalkPE!(g::GNNGraph, walk_length::Int, symbol::Symbol=:p)
+    matrix=zeros(walk_length,g.num_nodes)
+    adj = adjacency_matrix(g)
+    if adj isa CuArray
+        matrix = CuArray(matrix)
+    end
+    deg = degree(g)
+    deg_inv = inv.(deg)
+    deg_inv[isinf.(deg_inv)] .= 0
+    RW = adj * Diagonal(deg_inv)
+    out = RW
+    matrix[1,:] .= diag(RW)
+    for i in 2:walk_length
+        out = out * RW
+        matrix[i,:] .= diag(out)
+    end
+    setproperty!(g.ndata, symbol, matrix)
+end
+
 # """
 # Transform vector of cartesian indexes into a tuple of vectors containing integers.
 # """

test/GNNGraphs/transform.jl

@@ -1,12 +1,12 @@
 @testset "add self-loops" begin
     A = [1 1 0 0
-            0 0 1 0
-            0 0 0 1
-            1 0 0 0]
+         0 0 1 0
+         0 0 0 1
+         1 0 0 0]
     A2 = [2 1 0 0
-            0 1 1 0
-            0 0 1 1
-            1 0 0 1]
+          0 1 1 0
+          0 0 1 1
+          1 0 0 1]
 
     g = GNNGraph(A; graph_type = GRAPH_T)
     fg2 = add_self_loops(g)
@@ -18,7 +18,7 @@ end
 
 @testset "batch" begin
     g1 = GNNGraph(random_regular_graph(10, 2), ndata = rand(16, 10),
-                    graph_type = GRAPH_T)
+                  graph_type = GRAPH_T)
     g2 = GNNGraph(random_regular_graph(4, 2), ndata = rand(16, 4), graph_type = GRAPH_T)
     g3 = GNNGraph(random_regular_graph(7, 2), ndata = rand(16, 7), graph_type = GRAPH_T)
 
@@ -44,7 +44,7 @@ end
     # Batch of batches
     g123123 = Flux.batch([g123, g123])
     @test g123123.graph_indicator ==
-            [fill(1, 10); fill(2, 4); fill(3, 7); fill(4, 10); fill(5, 4); fill(6, 7)]
+          [fill(1, 10); fill(2, 4); fill(3, 7); fill(4, 10); fill(5, 4); fill(6, 7)]
     @test g123123.num_graphs == 6
 end
 
@@ -67,8 +67,8 @@ end
     c = 3
     ngraphs = 10
     gs = [rand_graph(n, c * n, ndata = rand(2, n), edata = rand(3, c * n),
-                        graph_type = GRAPH_T)
-            for _ in 1:ngraphs]
+                     graph_type = GRAPH_T)
+          for _ in 1:ngraphs]
     gall = Flux.batch(gs)
     gs2 = Flux.unbatch(gall)
     @test gs2[1] == gs[1]
@@ -77,7 +77,7 @@ end
 
 @testset "getgraph" begin
     g1 = GNNGraph(random_regular_graph(10, 2), ndata = rand(16, 10),
-                    graph_type = GRAPH_T)
+                  graph_type = GRAPH_T)
     g2 = GNNGraph(random_regular_graph(4, 2), ndata = rand(16, 4), graph_type = GRAPH_T)
     g3 = GNNGraph(random_regular_graph(7, 2), ndata = rand(16, 7), graph_type = GRAPH_T)
     g = Flux.batch([g1, g2, g3])
@@ -268,3 +268,9 @@ end end
     @test nv(DG) == g.num_nodes
     @test ne(DG) == g.num_edges
 end
+
+@testset "RandomWalkPE" begin
+    g = rand_graph(10, 20, graph_type = GRAPH_T)
+    add_RandomWalkPE!(g, 3)
+    @test size(g.ndata.p) == (3, 10)
+end