support for aggregating edge feature in remove_multi_edges (#115)

* work on remove multi edges

* improve remove_self_loops

Author: CarloLucibello

src/GNNGraphs/GNNGraphs.jl

@@ -67,5 +67,7 @@ include("operators.jl")
 include("convert.jl")
 include("utils.jl")
 
-    
+include("gatherscatter.jl")
+# _gather, _scatter
+
 end #module

src/GNNGraphs/gatherscatter.jl

@@ -0,0 +1,9 @@
+_gather(x::NamedTuple, i) = map(x -> _gather(x, i), x)
+_gather(x::Tuple, i) = map(x -> _gather(x, i), x)
+_gather(x::AbstractArray, i) = NNlib.gather(x, i)
+_gather(x::Nothing, i) = nothing
+
+_scatter(aggr, m::NamedTuple, t; dstsize=nothing) = map(m -> _scatter(aggr, m, t; dstsize), m)
+_scatter(aggr, m::Tuple, t; dstsize=nothing) = map(m -> _scatter(aggr, m, t; dstsize), m)
+_scatter(aggr, m::AbstractArray, t; dstsize=nothing) = NNlib.scatter(aggr, m, t; dstsize)
+_scatter(aggr, m::Nothing, t; dstsize=nothing) = nothing

src/GNNGraphs/transform.jl

@@ -41,18 +41,19 @@ end
 
 function remove_self_loops(g::GNNGraph{<:COO_T})
     s, t = edge_index(g)
-    # TODO remove these constraints
-    @assert g.edata === (;)
-    @assert get_edge_weight(g) === nothing
+    w = get_edge_weight(g)
+    edata = g.edata
 
     mask_old_loops = s .!= t
     s = s[mask_old_loops]
     t = t[mask_old_loops]
-
-    GNNGraph((s, t, nothing), 
+    edata = getobs(edata, mask_old_loops)
+    w = isnothing(w) ? nothing : getobs(w, mask_old_loops)
+    
+    GNNGraph((s, t, w), 
             g.num_nodes, length(s), g.num_graphs, 
             g.graph_indicator,
-            g.ndata, g.edata, g.gdata)
+            g.ndata, edata, g.gdata)
 end
 
 
@@ -72,25 +73,43 @@ end
 
 
 """
-    remove_multi_edges(g::GNNGraph)
+    remove_multi_edges(g::GNNGraph; aggr=+)
 
 Remove multiple edges (also called parallel edges or repeated edges) from graph `g`.
+Possible edge features are aggregated according to `aggr`, that can take value 
+`+`,`min`, `max` or `mean`.
+
+See also [`remove_self_loops`](@ref).
 """
-function remove_multi_edges(g::GNNGraph{<:COO_T})
+function remove_multi_edges(g::GNNGraph{<:COO_T}; aggr=+)
     s, t = edge_index(g)
-    # TODO remove these constraints
-    @assert g.num_graphs == 1
-    @assert g.edata === (;)
-    @assert get_edge_weight(g) === nothing
+    w = get_edge_weight(g)
+    edata = g.edata
+    num_edges = g.num_edges
 
     idxs, idxmax = edge_encoding(s, t, g.num_nodes)
-    union!(idxs)
-    s, t = edge_decoding(idxs, g.num_nodes)
+    
+    perm = sortperm(idxs)
+    idxs = idxs[perm]
+    s, t = s[perm], t[perm]
+    edata = getobs(edata, perm)
+    w = isnothing(w) ? nothing : getobs(w, perm)
+    idxs = [-1; idxs]
+    mask = idxs[2:end] .> idxs[1:end-1]
+    if !all(mask)
+        s, t = s[mask], t[mask]
+        idxs = similar(s, num_edges)
+        idxs .= 1:num_edges
+        idxs .= idxs .- cumsum(.!mask)
+        num_edges = length(s)
+        w = _scatter(aggr, w, idxs)
+        edata = _scatter(aggr, edata, idxs)
+    end
 
-    GNNGraph((s, t, nothing), 
-            g.num_nodes, length(s), g.num_graphs, 
+    GNNGraph((s, t, w),
+            g.num_nodes, num_edges, g.num_graphs,
             g.graph_indicator,
-            g.ndata, g.edata, g.gdata)
+            g.ndata, edata, g.gdata)
 end
 
 """

src/msgpass.jl

@@ -112,18 +112,12 @@ apply_edges(l, g::GNNGraph; xi=nothing, xj=nothing, e=nothing) =
 
 function apply_edges(f, g::GNNGraph, xi, xj, e)
     s, t = edge_index(g)
-    xi = _gather(xi, t)   # size: (D, num_nodes) -> (D, num_edges)
-    xj = _gather(xj, s)
+    xi = GNNGraphs._gather(xi, t)   # size: (D, num_nodes) -> (D, num_edges)
+    xj = GNNGraphs._gather(xj, s)
     m = f(xi, xj, e)
     return m
 end
 
-_gather(x::NamedTuple, i) = map(x -> _gather(x, i), x)
-_gather(x::Tuple, i) = map(x -> _gather(x, i), x)
-_gather(x::AbstractArray, i) = NNlib.gather(x, i)
-_gather(x::Nothing, i) = nothing
-
-
 ##  AGGREGATE NEIGHBORS
 @doc raw"""
     aggregate_neighbors(g::GNNGraph, aggr, m)
@@ -140,14 +134,9 @@ where it comes after [`apply_edges`](@ref).
 """
 function aggregate_neighbors(g::GNNGraph, aggr, m)
     s, t = edge_index(g)
-    return _scatter(aggr, m, t)
+    return GNNGraphs._scatter(aggr, m, t)
 end
 
-_scatter(aggr, m::NamedTuple, t) = map(m -> _scatter(aggr, m, t), m)
-_scatter(aggr, m::Tuple, t) = map(m -> _scatter(aggr, m, t), m)
-_scatter(aggr, m::AbstractArray, t) = NNlib.scatter(aggr, m, t)
-
-
 
 ### MESSAGE FUNCTIONS ###
 """

test/GNNGraphs/transform.jl

@@ -107,14 +107,25 @@
     end
 
     @testset "remove_self_loops" begin
-        if GRAPH_T == :coo
+        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;])
             @test g1.num_edges == g.num_edges + 5
-            g2 = remove_self_loops(g)
+            g2 = remove_self_loops(g1)
             @test g2.num_edges == g.num_edges
             @test sort_edge_index(edge_index(g2)) == sort_edge_index(edge_index(g))
-        end 
+
+            # with edge features and weights
+            g1 = GNNGraph(g1, edata=(e1=ones(3,g1.num_edges), e2=2*ones(g1.num_edges)))
+            g1 = set_edge_weight(g1, 3*ones(g1.num_edges))
+            g2 = remove_self_loops(g1)
+            @test g2.num_edges == g.num_edges
+            @test sort_edge_index(edge_index(g2)) == sort_edge_index(edge_index(g))
+            @test size(get_edge_weight(g2)) == (g2.num_edges,) 
+            @test size(g2.edata.e1) == (3, g2.num_edges) 
+            @test size(g2.edata.e2) == (g2.num_edges,) 
+
+        end
     end
 
     @testset "remove_multi_edges" begin
@@ -123,9 +134,20 @@
             s, t = edge_index(g)
             g1 = add_edges(g, s[1:5], t[1:5])
             @test g1.num_edges == g.num_edges + 5
-            g2 = remove_multi_edges(g)
+            g2 = remove_multi_edges(g1, aggr=+)
+            @test g2.num_edges == g.num_edges
+            @test sort_edge_index(edge_index(g2)) == sort_edge_index(edge_index(g))
+
+            # Default aggregation is +
+            g1 = GNNGraph(g1, edata=(e1=ones(3,g1.num_edges), e2=2*ones(g1.num_edges)))
+            g1 = set_edge_weight(g1, 3*ones(g1.num_edges))
+            g2 = remove_multi_edges(g1)
             @test g2.num_edges == g.num_edges
             @test sort_edge_index(edge_index(g2)) == sort_edge_index(edge_index(g))
+            @test count(g2.edata.e1[:,i] == 2*ones(3) for i in 1:g2.num_edges) == 5
+            @test count(g2.edata.e2[i] == 4 for i in 1:g2.num_edges) == 5
+            w2 = get_edge_weight(g2)
+            @test count(w2[i] == 6 for i in 1:g2.num_edges) == 5
         end
     end