add to_bidirected (#116)

* add to_bidirected(g)

* docs

Author: CarloLucibello

src/GNNGraphs/GNNGraphs.jl

@@ -13,7 +13,7 @@ import LearnBase
 import StatsBase
 import KrylovKit
 using ChainRulesCore
-using LinearAlgebra, Random
+using LinearAlgebra, Random, Statistics
 
 include("gnngraph.jl")
 export GNNGraph, 
@@ -48,6 +48,7 @@ export add_nodes,
        remove_self_loops,
        remove_multi_edges,
        set_edge_weight,
+       to_bidirected,
        # from Flux
        batch,
        unbatch,

src/GNNGraphs/transform.jl

@@ -79,7 +79,7 @@ Remove multiple edges (also called parallel edges or repeated edges) from graph
 Possible edge features are aggregated according to `aggr`, that can take value 
 `+`,`min`, `max` or `mean`.
 
-See also [`remove_self_loops`](@ref).
+See also [`remove_self_loops`](@ref), [`has_multi_edges`](@ref), and [`to_bidirected`](@ref).
 """
 function remove_multi_edges(g::GNNGraph{<:COO_T}; aggr=+)
     s, t = edge_index(g)
@@ -119,14 +119,14 @@ Add to graph `g` the edges with source nodes `s` and target nodes `t`.
 Optionally, pass the features  `edata` for the new edges. 
 """
 function add_edges(g::GNNGraph{<:COO_T}, 
-        snew::AbstractVector{<:Integer}, 
-        tnew::AbstractVector{<:Integer};
-        edata=nothing)
+                snew::AbstractVector{<:Integer}, 
+                tnew::AbstractVector{<:Integer};
+                edata=nothing)
 
     @assert length(snew) == length(tnew)
     # TODO remove this constraint
     @assert get_edge_weight(g) === nothing
-    
+
     edata = normalize_graphdata(edata, default_name=:e, n=length(snew))
     edata = cat_features(g.edata, edata)
 
@@ -157,6 +157,77 @@ end
 #     return true
 # 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. 
+
+See also [`is_bidirected`](@ref). 
+
+# Examples
+
+```juliarepl
+julia> s, t = [1, 2, 3, 3, 4], [2, 3, 4, 4, 4];
+
+julia> w = [1.0, 2.0, 3.0, 4.0, 5.0];
+
+julia> e = [10.0, 20.0, 30.0, 40.0, 50.0];
+
+julia> g = GNNGraph(s, t, w, edata = e)
+GNNGraph:
+    num_nodes = 4
+    num_edges = 5
+    edata:
+        e => (5,)
+
+julia> g2 = to_bidirected(g)
+GNNGraph:
+    num_nodes = 4
+    num_edges = 7
+    edata:
+        e => (7,)
+
+julia> edge_index(g2)
+([1, 2, 2, 3, 3, 4, 4], [2, 1, 3, 2, 4, 3, 4])
+
+julia> get_edge_weight(g2)
+7-element Vector{Float64}:
+ 1.0
+ 1.0
+ 2.0
+ 2.0
+ 3.5
+ 3.5
+ 5.0
+
+julia> g2.edata.e
+7-element Vector{Float64}:
+ 10.0
+ 10.0
+ 20.0
+ 20.0
+ 35.0
+ 35.0
+ 50.0
+```
+"""
+function to_bidirected(g::GNNGraph{<:COO_T})
+    s, t = edge_index(g)
+    w = get_edge_weight(g)
+    snew = [s; t]
+    tnew = [t; s] 
+    w = cat_features(w, w)
+    edata = cat_features(g.edata, g.edata)   
+
+    g = GNNGraph((snew, tnew, w),
+                g.num_nodes, length(snew), g.num_graphs,
+                g.graph_indicator,
+                g.ndata, edata, g.gdata)
+
+    return remove_multi_edges(g; aggr=mean)
+end
+
 
 """
     add_nodes(g::GNNGraph, n; [ndata])

test/GNNGraphs/transform.jl

@@ -209,4 +209,27 @@
         gw2 = set_edge_weight(g2, w)
         @test get_edge_weight(gw2) == w
     end
+
+    @testset "to_bidirected" begin
+        if GRAPH_T == :coo
+            s, t = [1, 2, 3, 3, 4], [2, 3, 4, 4, 4] 
+            w = [1.0, 2.0, 3.0, 4.0, 5.0]
+            e = [10.0, 20.0, 30.0, 40.0, 50.0]
+            g = GNNGraph(s, t, w, edata = e)
+
+            g2 = to_bidirected(g)
+            @test g2.num_nodes == g.num_nodes
+            @test g2.num_edges == 7
+            s2, t2 = edge_index(g2)         
+            @test is_bidirected(g2)
+            @test !has_multi_edges(g2)
+
+            s2, t2 = edge_index(g2)
+            w2 = get_edge_weight(g2)
+            @test s2 == [1, 2, 2, 3, 3, 4, 4]
+            @test t2 == [2, 1, 3, 2, 4, 3, 4]
+            @test w2 == [1, 1, 2, 2, 3.5, 3.5, 5]
+            @test g2.edata.e == [10.0, 10.0, 20.0, 20.0, 35.0, 35.0, 50.0]
+        end
+    end
 end