[heterograph] add has_edge, get_index, add_edges (#298)

* [heterograph] has_edge and get_index

* add_edges

* cleanup

Author: CarloLucibello

docs/make.jl

@@ -24,7 +24,7 @@ makedocs(;
              "Message Passing" => "messagepassing.md",
              "Model Building" => "models.md",
              "Datasets" => "datasets.md",
-            #  "Tutorials" => tutorials,
+             "Tutorials" => tutorials,
              "API Reference" => [
                  "GNNGraph" => "api/gnngraph.md",
                  "Basic Layers" => "api/basic.md",

src/GNNGraphs/GNNGraphs.jl

@@ -81,6 +81,7 @@ export add_nodes,
 include("generate.jl")
 export rand_graph,
        rand_heterograph,
+       rand_bipartite_heterograph,
        knn_graph,
        radius_graph
 

src/GNNGraphs/convert.jl

@@ -1,7 +1,7 @@
 ### CONVERT_TO_COO REPRESENTATION ########
 
 function to_coo(data::EDict; num_nodes = nothing, kws...)
-    graph = EDict{Any}()
+    graph = EDict{COO_T}()
     _num_nodes = NDict{Int}()
     num_edges = EDict{Int}()
     for k in keys(data)
@@ -23,6 +23,7 @@ function to_coo(data::EDict; num_nodes = nothing, kws...)
         _num_nodes[k[1]] = max(get(_num_nodes, k[1], 0), nnodes[1])
         _num_nodes[k[3]] = max(get(_num_nodes, k[3], 0), nnodes[2])
     end
+    graph = Dict(k => v for (k, v) in pairs(graph)) # try to restrict the key/value types
     return graph, _num_nodes, num_edges
 end
 

src/GNNGraphs/generate.jl

@@ -48,21 +48,23 @@ function rand_graph(n::Integer, m::Integer; bidirected = true, seed = -1, edge_w
 end
 
 """
-    rand_heterograph(n, m; seed=-1, kws...)
+    rand_heterograph(n, m; seed=-1, bidirected=false, kws...)
 
 Construct an [`GNNHeteroGraph`](@ref) with number of nodes and edges 
-specified by `n` and `m` respectively.
-`n` and `m` can be any iterable of pairs.
+specified by `n` and `m` respectively. `n` and `m` can be any iterable of pairs
+specifing node/edge types and their numbers.
 
 Use a `seed > 0` for reproducibility.
 
+Setting `bidirected=true` will generate a bidirected graph, i.e. each edge will have a reverse edge.
+Therefore, for each edge type `(:A, :rel, :B)` a corresponding reverse edge type `(:B, :rel, :A)`
+will be generated.
+
 Additional keyword arguments will be passed to the [`GNNHeteroGraph`](@ref) constructor.
 
 # Examples
 
 ```julia-repl
-
-
 julia> g = rand_heterograph((:user => 10, :movie => 20),
                             (:user, :rate, :movie) => 30)
 GNNHeteroGraph:
@@ -76,19 +78,72 @@ function rand_heterograph end
 rand_heterograph(n, m; kws...) = rand_heterograph(Dict(n), Dict(m); kws...)
 
 function rand_heterograph(n::NDict, m::EDict; bidirected = false, seed = -1, kws...)
-    @assert !bidirected "Bidirected graphs not supported yet."
     rng = seed > 0 ? MersenneTwister(seed) : Random.GLOBAL_RNG
+    if bidirected
+        return _rand_bidirected_heterograph(rng, n, m; kws...)
+    end
     graphs = Dict(k => _rand_edges(rng, (n[k[1]], n[k[3]]), m[k]) for k in keys(m))
     return GNNHeteroGraph(graphs; num_nodes = n, kws...)
 end
 
+function _rand_bidirected_heterograph(rng, n::NDict, m::EDict; kws...)
+    for k in keys(m)
+        if reverse(k) ∈ keys(m)
+            @assert m[k] == m[reverse(k)] "Number of edges must be the same in reverse edge types for bidirected graphs."
+        else
+            m[reverse(k)] = m[k]
+        end
+    end
+    graphs = Dict{EType, Tuple{Vector{Int}, Vector{Int}, Nothing}}()
+    for k in keys(m)
+        reverse(k) ∈ keys(graphs) && continue
+        s, t, val =  _rand_edges(rng, (n[k[1]], n[k[3]]), m[k])
+        graphs[k] = s, t, val
+        graphs[reverse(k)] = t, s, val
+    end
+    return GNNHeteroGraph(graphs; num_nodes = n, kws...)
+end
+
 function _rand_edges(rng, (n1, n2), m)
     idx = StatsBase.sample(rng, 1:(n1 * n2), m, replace = false)
     s, t = edge_decoding(idx, n1, n2)
     val = nothing
     return s, t, val
 end
 
+"""
+    rand_bipartite_heterograph(n1, n2, m; [bidirected, seed, node_t, edge_t, kws...])
+    rand_bipartite_heterograph((n1, n2), m; ...)
+    rand_bipartite_heterograph((n1, n2), (m1, m2); ...)
+
+Construct an [`GNNHeteroGraph`](@ref) with number of nodes and edges
+specified by `n1`, `n2` and `m1` and `m2` respectively.
+
+See [`rand_heterograph`](@ref) for a more general version.
+
+# Keyword arguments
+
+- `bidirected`: whether to generate a bidirected graph. Default is `true`.
+- `seed`: random seed. Default is `-1` (no seed).
+- `node_t`: node types. If `bipartite=true`, this should be a tuple of two node types, otherwise it should be a single node type.
+- `edge_t`: edge types. If `bipartite=true`, this should be a tuple of two edge types, otherwise it should be a single edge type.
+"""
+function rand_bipartite_heterograph end
+
+rand_bipartite_heterograph(n1::Int, n2::Int, m::Int; kws...) = rand_bipartite_heterograph((n1, n2), (m, m); kws...)
+
+rand_bipartite_heterograph((n1, n2)::NTuple{2,Int}, m::Int; kws...) = rand_bipartite_heterograph((n1, n2), (m, m); kws...)
+
+function rand_bipartite_heterograph((n1, n2)::NTuple{2,Int}, (m1, m2)::NTuple{2,Int}; bidirected=true, 
+                        node_t = (:A, :B), edge_t = :to, kws...)
+    if edge_t isa Symbol
+        edge_t = (edge_t, edge_t)
+    end
+    return rand_heterograph(Dict(node_t[1] => n1, node_t[2] => n2), 
+                            Dict((node_t[1], edge_t[1], node_t[2]) => m1, (node_t[2], edge_t[2], node_t[1]) => m2); 
+                            bidirected, kws...)
+end
+
 """
     knn_graph(points::AbstractMatrix, 
               k::Int; 

src/GNNGraphs/gnnheterograph.jl

@@ -1,6 +1,8 @@
 
-const EDict{T} = Dict{Tuple{Symbol, Symbol, Symbol}, T}
-const NDict{T} = Dict{Symbol, T}
+const EType = Tuple{Symbol, Symbol, Symbol} 
+const NType = Symbol
+const EDict{T} = Dict{EType, T}
+const NDict{T} = Dict{NType, T}
 
 """
     GNNHeteroGraph(data; [ndata, edata, gdata, num_nodes])
@@ -78,30 +80,35 @@ julia> hg.ndata[:A].x
 
 See also [`GNNGraph`](@ref) for a homogeneous graph type and [`rand_heterograph`](@ref) for a function to generate random heterographs.
 """
-struct GNNHeteroGraph
-    graph::EDict
+struct GNNHeteroGraph{T <: Union{COO_T, ADJMAT_T}}
+    graph::EDict{T}
     num_nodes::NDict{Int}
     num_edges::EDict{Int}
     num_graphs::Int
     graph_indicator::Union{Nothing, NDict}
     ndata::NDict{DataStore}
     edata::EDict{DataStore}
     gdata::DataStore
-    ntypes::Vector{Symbol}
-    etypes::Vector{Symbol}
+    ntypes::Vector{NType}
+    etypes::Vector{EType}
 end
 
 @functor GNNHeteroGraph
 
 GNNHeteroGraph(data; kws...) = GNNHeteroGraph(Dict(data); kws...)
 
+function GNNHeteroGraph(data::Dict; kws...)
+    all(k -> k isa EType, keys(data)) || throw(ArgumentError("Keys of data must be tuples of the form (source_type, edge_type, target_type)"))
+    return GNNHeteroGraph(Dict(k => v for (k, v) in pairs(data)); kws...)
+end
+
 function GNNHeteroGraph(data::EDict;
                         num_nodes = nothing,
                         graph_indicator = nothing,
                         graph_type = :coo,
                         dir = :out,
-                        ndata = NDict{NamedTuple}(),
-                        edata = EDict{NamedTuple}(),
+                        ndata = NDict{DataStore}(),
+                        edata = EDict{DataStore}(),
                         gdata = (;))
     @assert graph_type ∈ [:coo, :dense, :sparse] "Invalid graph_type $graph_type requested"
     @assert dir ∈ [:in, :out]
@@ -112,7 +119,7 @@ function GNNHeteroGraph(data::EDict;
     end
 
     ntypes = union([[k[1] for k in keys(data)]; [k[3] for k in keys(data)]])
-    etypes = [k[2] for k in keys(data)]
+    etypes = collect(keys(data))
 
     if graph_type == :coo
         graph, num_nodes, num_edges = to_coo(data; num_nodes, dir)

src/GNNGraphs/query.jl

@@ -13,10 +13,19 @@ edge_index(g::GNNGraph{<:COO_T}) = g.graph[1:2]
 
 edge_index(g::GNNGraph{<:ADJMAT_T}) = to_coo(g.graph, num_nodes = g.num_nodes)[1][1:2]
 
+""""
+    edge_index(g::GNNHeteroGraph, edge_t)
+
+Return a tuple containing two vectors, respectively storing the source and target nodes for each edges in `g` of type `edge_t = (:node1_t, :rel, :node2_t)`.
+"""
+edge_index(g::GNNHeteroGraph{<:COO_T}, edge_t::EType) = g.graph[edge_t][1:2]
+
 get_edge_weight(g::GNNGraph{<:COO_T}) = g.graph[3]
 
 get_edge_weight(g::GNNGraph{<:ADJMAT_T}) = to_coo(g.graph, num_nodes = g.num_nodes)[1][3]
 
+get_edge_weight(g::GNNHeteroGraph{<:COO_T}, edge_t::EType) = g.graph[edge_t][3]
+
 Graphs.edges(g::GNNGraph) = Graphs.Edge.(edge_index(g)...)
 
 Graphs.edgetype(g::GNNGraph) = Graphs.Edge{eltype(g)}
@@ -42,6 +51,31 @@ end
 
 Graphs.has_edge(g::GNNGraph{<:ADJMAT_T}, i::Integer, j::Integer) = g.graph[i, j] != 0
 
+"""
+    has_edge(g::GNNHeteroGraph, edge_t, i, j)
+
+Return `true` if there is an edge of type `edge_t` from node `i` to node `j` in `g`.
+
+# Examples
+
+```julia-repl
+julia> g = rand_bipartite_heterograph((2, 2), (4, 0), bidirected=false)
+GNNHeteroGraph:
+  num_nodes: (:A => 2, :B => 2)
+  num_edges: ((:A, :to, :B) => 4, (:B, :to, :A) => 0)
+
+julia> has_edge(g, (:A,:to,:B), 1, 1)
+true
+
+julia> has_edge(g, (:B,:to,:A), 1, 1)
+false
+```
+"""
+function Graphs.has_edge(g::GNNHeteroGraph, edge_t::EType, i::Integer, j::Integer)
+    s, t = edge_index(g, edge_t)
+    return any((s .== i) .& (t .== j))
+end
+
 graph_type_symbol(::GNNGraph{<:COO_T}) = :coo
 graph_type_symbol(::GNNGraph{<:SPARSE_T}) = :sparse
 graph_type_symbol(::GNNGraph{<:ADJMAT_T}) = :dense

src/GNNGraphs/transform.jl

@@ -21,7 +21,7 @@ function add_self_loops(g::GNNGraph{<:COO_T})
         ew = [ew; fill!(similar(ew, n), 1)]
     end
 
-    GNNGraph((s, t, ew),
+    return GNNGraph((s, t, ew),
              g.num_nodes, length(s), g.num_graphs,
              g.graph_indicator,
              g.ndata, g.edata, g.gdata)
@@ -32,7 +32,7 @@ function add_self_loops(g::GNNGraph{<:ADJMAT_T})
     @assert isempty(g.edata)
     num_edges = g.num_edges + g.num_nodes
     A = A + I
-    GNNGraph(A,
+    return GNNGraph(A,
              g.num_nodes, num_edges, g.num_graphs,
              g.graph_indicator,
              g.ndata, g.edata, g.gdata)
@@ -71,7 +71,7 @@ function remove_self_loops(g::GNNGraph{<:ADJMAT_T})
         dropzeros!(A)
     end
     num_edges = numnonzeros(A)
-    GNNGraph(A,
+    return GNNGraph(A,
              g.num_nodes, num_edges, g.num_graphs,
              g.graph_indicator,
              g.ndata, g.edata, g.gdata)
@@ -110,7 +110,7 @@ function remove_multi_edges(g::GNNGraph{<:COO_T}; aggr = +)
         edata = _scatter(aggr, edata, idxs, num_edges)
     end
 
-    GNNGraph((s, t, w),
+    return GNNGraph((s, t, w),
              g.num_nodes, num_edges, g.num_graphs,
              g.graph_indicator,
              g.ndata, edata, g.gdata)
@@ -137,13 +137,55 @@ function add_edges(g::GNNGraph{<:COO_T},
     s = [s; snew]
     t = [t; tnew]
 
-    GNNGraph((s, t, nothing),
+    return GNNGraph((s, t, nothing),
              g.num_nodes, length(s), g.num_graphs,
              g.graph_indicator,
              g.ndata, edata, g.gdata)
 end
 
-### TODO Cannot implement this since GNNGraph is immutable (cannot change num_edges)
+function add_edges(g::GNNHeteroGraph{<:COO_T},
+                   edge_t::EType,
+                   snew::AbstractVector{<:Integer},
+                   tnew::AbstractVector{<:Integer};
+                   edata = nothing)
+    @assert length(snew) == length(tnew)
+    # TODO remove this constraint
+    @assert get_edge_weight(g, edge_t) === nothing
+
+    edata = normalize_graphdata(edata, default_name = :e, n = length(snew))
+    g_edata = g.edata |> copy
+    if !isempty(g.edata)
+        if haskey(g_edata, edge_t)
+            g_edata[edge_t] = cat_features(g.edata[edge_t], edata)
+        else
+            g_edata[edge_t] = edata
+        end
+    end
+
+    graph = g.graph |> copy
+    etypes = g.etypes |> copy
+    if !haskey(graph, edge_t)
+        @assert edge_t[1] ∈ g.ntypes && edge_t[3] ∈ g.ntypes
+        push!(g.etypes, edge_t)
+    else
+        s, t = edge_index(g, edge_t)
+        snew = [s; snew]
+        tnew = [t; tnew]
+    end
+    graph[edge_t] = (snew, tnew, nothing)
+    num_edges = g.num_edges |> copy
+    num_edges[edge_t] = length(graph[edge_t][1])
+
+    return GNNHeteroGraph(graph,
+             g.num_nodes, num_edges, g.num_graphs,
+             g.graph_indicator,
+             g.ndata, g_edata, g.gdata,
+             g.ntypes, etypes)
+end
+
+
+
+### TODO Cannot implement this since GNNGraph is immutable (cannot change num_edges). make it mutable
 # function Graphs.add_edge!(g::GNNGraph{<:COO_T}, snew::T, tnew::T; edata=nothing) where T<:Union{Integer, AbstractVector}
 #     s, t = edge_index(g)
 #     @assert length(snew) == length(tnew)

test/GNNGraphs/generate.jl

@@ -75,3 +75,20 @@ end
     s, t = edge_index(g)
     @test (s .> 5) == (t .> 5)
 end
+
+@testset "rand_bipartite_heterograph" begin
+    g = rand_bipartite_heterograph(10, 15, 20)
+    @test g.num_nodes == Dict(:A => 10, :B => 15)
+    @test g.num_edges == Dict((:A, :to, :B) => 20, (:B, :to, :A) => 20)
+    sA, tB = edge_index(g, (:A, :to, :B))
+    for (s, t) in zip(sA, tB)
+        @test 1 <= s <= 10
+        @test 1 <= t <= 15
+        @test has_edge(g, (:A,:to,:B), s, t)
+        @test has_edge(g, (:B,:to,:A), t, s)
+    end
+
+    g = rand_bipartite_heterograph((2, 2), (4, 0), bidirected=false)
+    @test has_edge(g, (:A,:to,:B), 1, 1)
+    @test !has_edge(g, (:B,:to,:A), 1, 1)
+end

test/GNNGraphs/gnnheterograph.jl

@@ -18,7 +18,7 @@
     @test hg.edata == Dict()
     @test isempty(hg.gdata)
     @test sort(hg.ntypes) == [:A, :B]
-    @test sort(hg.etypes) == [:rel1, :rel2]
+    @test sort(hg.etypes) == [(:A, :rel1, :B), (:B, :rel2, :A)]
 end
 
 @testset "features" begin