Fixes and more tests in remove_nodes function (#424)

* need to fix

* bounds check

* added more tests

* Update src/GNNGraphs/transform.jl

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

* Update transform.jl (re added previous change)

---------

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

Author: rbSparky

src/GNNGraphs/transform.jl

@@ -219,9 +219,16 @@ function remove_nodes(g::GNNGraph{<:COO_T}, nodes_to_remove::AbstractVector)
     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)
+    
+    function find_edges_to_remove(nodes, nodes_to_remove)
+        return findall(node_id -> begin
+            idx = searchsortedlast(nodes_to_remove, node_id)
+            idx >= 1 && idx <= length(nodes_to_remove) && nodes_to_remove[idx] == node_id
+        end, nodes)
+    end
+    
+    edges_to_remove_s = find_edges_to_remove(s, nodes_to_remove)
+    edges_to_remove_t = find_edges_to_remove(t, nodes_to_remove)
     edges_to_remove = union(edges_to_remove_s, edges_to_remove_t)
 
     mask_edges_to_keep = trues(length(s))

test/GNNGraphs/transform.jl

@@ -150,6 +150,7 @@ end
 end
 
 @testset "remove_nodes" begin if GRAPH_T == :coo
+    #single node
     s = [1, 1, 2, 3]
     t = [2, 3, 4, 5]
     eweights = [0.1, 0.2, 0.3, 0.4]
@@ -179,6 +180,35 @@ end
     @test eweights_new == eweightstest
     @test ndata_new == ndatatest
     @test edata_new == edatatest
+
+    # multiple nodes
+    s = [1, 5, 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,4])
+    snew = [3,2]
+    tnew = [2,3]
+    eweights_new = [0.2,0.4]
+    ndata_new = [2.0,3.0,5.0]
+    edata_new = ['b','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 == 3
+    @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