Coalesce function and state for coo graphs (#613)

* Add coalesce_graph function for COO graphs and is_coalesced flag

* Simplify coalesce_graph function

* Revert back to old  sort_edge_index function

* Rename remove_multi_edges to coalesce, update references, add deprecation

* Add is_coalesced function to check if graph is coalesced

* Add tests to verify is_coalesced for various graph transformations

Author: dferre97

GNNGraphs/src/GNNGraphs.jl

@@ -75,7 +75,7 @@ export add_nodes,
        rand_edge_split,
        remove_self_loops,
        remove_edges, 
-       remove_multi_edges,
+       coalesce,
        set_edge_weight,
        to_bidirected,
        to_unidirected,

GNNGraphs/src/deprecations.jl

@@ -11,3 +11,5 @@ function Base.getproperty(vds::Vector{DataStore}, s::Symbol)
         return [getdata(ds)[s] for ds in vds]
     end
 end
+
+@deprecate remove_multi_edges(g::GNNGraph{<:COO_T}; aggr = +) Base.coalesce(g; aggr = aggr)

GNNGraphs/src/gnngraph.jl

@@ -113,6 +113,20 @@ struct GNNGraph{T <: Union{COO_T, ADJMAT_T}} <: AbstractGNNGraph{T}
     ndata::DataStore
     edata::DataStore
     gdata::DataStore
+    is_coalesced::Bool # only for :coo, true if the graph is coalesced, i.e., indices ordered by row and no multi edges
+end
+
+# GNNGraph constructor setting the is_coalesced field to false
+function GNNGraph(graph::T,
+                  num_nodes::Int,
+                  num_edges::Int,
+                  num_graphs::Int,
+                  graph_indicator::Union{Nothing, AVecI},
+                  ndata::DataStore,
+                  edata::DataStore,
+                  gdata::DataStore) where {T <: Union{COO_T, ADJMAT_T}}
+    return GNNGraph{T}(graph, num_nodes, num_edges, num_graphs,
+                       graph_indicator, ndata, edata, gdata, false)
 end
 
 function GNNGraph(data::D;

GNNGraphs/src/query.jl

@@ -103,6 +103,18 @@ Graphs.ne(g::GNNGraph) = g.num_edges
 Graphs.has_vertex(g::GNNGraph, i::Int) = 1 <= i <= g.num_nodes
 Graphs.vertices(g::GNNGraph) = 1:(g.num_nodes)
 
+"""
+    is_coalesced(g::GNNGraph) -> Bool
+
+Check whether the given `GNNGraph` `g` is coalesced (see [`coalesce`](@ref)). Only meaningful for COO graphs.
+
+# Arguments
+- `g::GNNGraph`: The graph to check.
+
+# Returns
+- `Bool`: Whether the graph is coalesced. If the graph is not of type COO, this function will always return `false`.
+"""
+is_coalesced(g::GNNGraph) = g.is_coalesced
 
 """
     neighbors(g::GNNGraph, i::Integer; dir=:out)

GNNGraphs/src/transform.jl

@@ -44,7 +44,7 @@ end
 Return a graph constructed from `g` where self-loops (edges from a node to itself)
 are removed. 
 
-See also [`add_self_loops`](@ref) and [`remove_multi_edges`](@ref).
+See also [`add_self_loops`](@ref) and [`coalesce`](@ref).
 """
 function remove_self_loops(g::GNNGraph{<:COO_T})
     s, t = edge_index(g)
@@ -146,15 +146,14 @@ function remove_edges(g::GNNGraph{<:COO_T}, p = 0.5)
 end
 
 """
-    remove_multi_edges(g::GNNGraph; aggr=+)
+    coalesce(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`.
+Return a new GNNGraph where all multiple edges between the same pair of nodes are merged (using aggr for edge weights and features), and the edge indices are sorted lexicographically (by source, then target).
+This method is only applicable to graphs of type `:coo`.
 
-See also [`remove_self_loops`](@ref), [`has_multi_edges`](@ref), and [`to_bidirected`](@ref).
+`aggr` can take value `+`,`min`, `max` or `mean`.
 """
-function remove_multi_edges(g::GNNGraph{<:COO_T}; aggr = +)
+function Base.coalesce(g::GNNGraph{<:COO_T}; aggr = +)
     s, t = edge_index(g)
     w = get_edge_weight(g)
     edata = g.edata
@@ -181,7 +180,7 @@ function remove_multi_edges(g::GNNGraph{<:COO_T}; aggr = +)
     return GNNGraph((s, t, w),
              g.num_nodes, num_edges, g.num_graphs,
              g.graph_indicator,
-             g.ndata, edata, g.gdata)
+             g.ndata, edata, g.gdata, true)
 end
 
 """
@@ -441,7 +440,7 @@ end
     to_bidirected(g)
 
 Adds a reverse edge for each edge in the graph, then calls 
-[`remove_multi_edges`](@ref) with `mean` aggregation to simplify the graph. 
+[`coalesce`](@ref) with `mean` aggregation to simplify the graph. 
 
 See also [`is_bidirected`](@ref). 
 
@@ -505,7 +504,7 @@ function to_bidirected(g::GNNGraph{<:COO_T})
                  g.graph_indicator,
                  g.ndata, edata, g.gdata)
 
-    return remove_multi_edges(g; aggr = mean)
+    return coalesce(g; aggr = mean)
 end
 
 """
@@ -525,7 +524,7 @@ function to_unidirected(g::GNNGraph{<:COO_T})
                  g.graph_indicator,
                  g.ndata, g.edata, g.gdata)
 
-    return remove_multi_edges(g; aggr = mean)
+    return coalesce(g; aggr = mean)
 end
 
 function Graphs.SimpleGraph(g::GNNGraph)

GNNGraphs/test/transform.jl

@@ -345,22 +345,22 @@ end
     end
 end
 
-@testitem "remove_multi_edges" setup=[GraphsTestModule] begin
+@testitem "coalesce" setup=[GraphsTestModule] begin
     using .GraphsTestModule
     for GRAPH_T in GRAPH_TYPES
         if GRAPH_T == :coo
             g = rand_graph(10, 20, graph_type = GRAPH_T)
             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(g1, aggr = +)
+            g2 = coalesce(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)
+            g2 = coalesce(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
@@ -714,3 +714,51 @@ end
         end
     end
 end
+
+@testitem "graph transform ops set is_coalesced=false" setup=[GraphsTestModule] begin
+    using .GraphsTestModule
+    g = rand_graph(5, 10, graph_type=:coo)
+    g = coalesce(g)  # ensure the graph is coalesced to start with
+
+    # add_self_loops
+    g1 = add_self_loops(g)
+    @test g1.is_coalesced == false
+
+    # remove_self_loops
+    g2 = add_self_loops(g)  # ensure there are self-loops to remove
+    g2 = remove_self_loops(g2)
+    @test g2.is_coalesced == false
+
+    # remove_edges
+    g3 = remove_edges(g, [1])
+    @test g3.is_coalesced == false
+
+    # add_edges
+    g4 = add_edges(g, [1], [2])
+    @test g4.is_coalesced == false
+
+    # perturb_edges
+    g5 = perturb_edges(g, 0.5)
+    @test g5.is_coalesced == false
+
+    # remove_nodes
+    g6 = remove_nodes(g, [1])
+    @test g6.is_coalesced == false
+
+    # add_nodes
+    g7 = add_nodes(g, 2)
+    @test g7.is_coalesced == false
+
+    # rand_edge_split returns two graphs
+    g8a, g8b = rand_edge_split(g, 0.5)
+    @test g8a.is_coalesced == false
+    @test g8b.is_coalesced == false
+
+    # negative_sample
+    g9 = negative_sample(g, num_neg_edges=3)
+    @test g9.is_coalesced == false
+
+    # ppr_diffusion
+    g11 = ppr_diffusion(g)
+    @test g11.is_coalesced == false
+end