Added remove_nodes function (#420)

* remove nodes

* Update transform.jl

* remove nodes

* Update src/GNNGraphs/transform.jl

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

* fix

* made efficient using searchsortedlast

* more tests

* remove ndata getobs

* Update test/GNNGraphs/transform.jl

---------

Author: rbSparky

src/GNNGraphs/GNNGraphs.jl

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

src/GNNGraphs/transform.jl

@@ -188,6 +188,66 @@ function remove_multi_edges(g::GNNGraph{<:COO_T}; aggr = +)
              g.ndata, edata, g.gdata)
 end
 
+"""
+    remove_nodes(g::GNNGraph, nodes_to_remove::AbstractVector)
+
+Remove specified nodes, and their associated edges, from a GNNGraph. This operation reindexes the remaining nodes to maintain a continuous sequence of node indices, starting from 1. Similarly, edges are reindexed to account for the removal of edges connected to the removed nodes.
+
+# Arguments
+- `g`: The input graph from which nodes (and their edges) will be removed.
+- `nodes_to_remove`: Vector of node indices to be removed.
+
+# Returns
+A new GNNGraph with the specified nodes and all edges associated with these nodes removed. 
+
+# Example
+```julia
+using GraphNeuralNetworks
+
+g = GNNGraph([1, 1, 2, 2, 3], [2, 3, 1, 3, 1])
+
+# Remove nodes with indices 2 and 3, for example
+g_new = remove_nodes(g, [2, 3])
+
+# g_new now does not contain nodes 2 and 3, and any edges that were connected to these nodes.
+println(g_new)
+```
+"""
+function remove_nodes(g::GNNGraph{<:COO_T}, nodes_to_remove::AbstractVector)
+    nodes_to_remove = sort(union(nodes_to_remove))
+    s, t = edge_index(g)
+    w = get_edge_weight(g)
+    edata = g.edata
+    ndata = g.ndata
+
+    edges_to_remove_s = findall(x -> nodes_to_remove[searchsortedlast(nodes_to_remove, x)] == x, s)
+    edges_to_remove_t = findall(x -> nodes_to_remove[searchsortedlast(nodes_to_remove, x)] == x, t)
+    edges_to_remove = union(edges_to_remove_s, edges_to_remove_t)
+
+    mask_edges_to_keep = trues(length(s))
+    mask_edges_to_keep[edges_to_remove] .= false
+    s = s[mask_edges_to_keep]
+    t = t[mask_edges_to_keep]
+
+    w = isnothing(w) ? nothing : getobs(w, mask_edges_to_keep)
+
+    for node in sort(nodes_to_remove, rev=true) 
+        s[s .> node] .-= 1
+        t[t .> node] .-= 1
+    end
+
+    nodes_to_keep = setdiff(1:g.num_nodes, nodes_to_remove)
+    ndata = getobs(ndata, nodes_to_keep)
+    edata = getobs(edata, mask_edges_to_keep)
+
+    num_nodes = g.num_nodes - length(nodes_to_remove)
+    
+    return GNNGraph((s, t, w),
+             num_nodes, length(s), g.num_graphs,
+             g.graph_indicator,
+             ndata, edata, g.gdata)
+end
+
 """
     add_edges(g::GNNGraph, s::AbstractVector, t::AbstractVector; [edata])
     add_edges(g::GNNGraph, (s, t); [edata])

test/GNNGraphs/transform.jl

@@ -149,6 +149,37 @@ end
     end 
 end
 
+@testset "remove_nodes" begin if GRAPH_T == :coo
+    s = [1, 1, 2, 3]
+    t = [2, 3, 4, 5]
+    eweights = [0.1, 0.2, 0.3, 0.4]
+    ndata = [1.0, 2.0, 3.0, 4.0, 5.0]
+    edata = ['a', 'b', 'c', 'd']
+
+    g = GNNGraph(s, t, eweights, ndata = ndata, edata = edata, graph_type = GRAPH_T)
+
+    gnew = remove_nodes(g, [1])
+
+    snew = [1, 2]
+    tnew = [3, 4]
+    eweights_new = [0.3, 0.4]
+    ndata_new = [2.0, 3.0, 4.0, 5.0]
+    edata_new = ['c', 'd']
+
+    stest, ttest = edge_index(gnew)
+    eweightstest = get_edge_weight(gnew)
+    ndatatest = gnew.ndata.x
+    edatatest = gnew.edata.e
+
+
+    @test gnew.num_edges == 2
+    @test gnew.num_nodes == 4
+    @test snew == stest
+    @test tnew == ttest
+    @test eweights_new == eweightstest
+    @test ndata_new == ndatatest
+    @test edata_new == edatatest
+end end
 
 @testset "add_nodes" begin if GRAPH_T == :coo
     g = rand_graph(6, 4, ndata = rand(2, 6), graph_type = GRAPH_T)