Added node_shuffle

src/GNNGraphs/GNNGraphs.jl

@@ -78,6 +78,7 @@ export add_nodes,
        to_unidirected,
        random_walk_pe,
        remove_nodes,
+       node_shuffle,
 # from Flux
        batch,
        unbatch,

src/GNNGraphs/transform.jl

@@ -248,6 +248,28 @@ function remove_nodes(g::GNNGraph{<:COO_T}, nodes_to_remove::AbstractVector)
              ndata, edata, g.gdata)
 end
 
+"""
+    node_shuffle(g::GNNGraph{<:COO_T})
+
+Shuffle the node features in the given GNNGraph as shown in paper [On Node Features for Graph Neural Networks](https://arxiv.org/ftp/arxiv/papers/1911/1911.08795.pdf)
+
+# Arguments
+- `g::GNNGraph`: The input graph represented as a GNNGraph.
+
+# Returns
+- `g::GNNGraph`: The GNNGraph with shuffled node features.
+"""
+function node_shuffle(g::GNNGraph{<:COO_T})
+    features = g.ndata.x
+    num_nodes = size(features, 1) 
+
+    shuffled_indices = randperm(num_nodes)
+    shuffled_features = features[shuffled_indices, :]
+
+    g.ndata.x = shuffled_features
+    return g  
+end
+
 """
     add_edges(g::GNNGraph, s::AbstractVector, t::AbstractVector; [edata])
     add_edges(g::GNNGraph, (s, t); [edata])

test/GNNGraphs/transform.jl

@@ -190,6 +190,14 @@ end end
     @test all(gnew.ndata.x[:, 7:11] .== 1)
 end end
 
+@testset "node_shuffle" begin if GRAPH_T == :coo
+    g = rand_graph(5, 4, ndata = rand(2, 5), graph_type = GRAPH_T)
+    g_new = node_shuffle(g)
+
+    @test g_new.num_edges == 4
+    @test g_new.num_nodes == 5
+end end
+
 @testset "remove_self_loops" begin if GRAPH_T == :coo # add_edges and set_edge_weight only implemented for coo
     g = rand_graph(10, 20, graph_type = GRAPH_T)
     g1 = add_edges(g, [1:5;], [1:5;])