heterograph guide (#310)

Author: CarloLucibello

docs/src/heterograph.md

@@ -3,12 +3,95 @@
 Heterogeneus graphs (also called heterographs), are graphs where each node has a type,
 that we denote with symbols such as `:user` and `:movie`,
 and edges also represent different relations identified
-by a triple of symbols, `(source_nodes, edge_type, target_nodes)`, as in `(:user, :rate, :movie)`.
+by a triplet of symbols, `(source_node_type, edge_type, target_node_type)`, as in `(:user, :rate, :movie)`.
 
 Different node/edge types can store different group of features
 and this makes heterographs a very flexible modeling tools 
-and data containers.
-
-In GraphNeuralNetworks.jl heterographs are implemented in 
+and data containers. In GraphNeuralNetworks.jl heterographs are implemented in 
 the type [`GNNHeteroGraph`](@ref).
 
+
+## Creating a Heterograph
+
+A heterograph can be created by passing pairs of edge types and data to the constructor.
+```julia-repl
+julia> g = GNNHeteroGraph((:user, :rate, :movie) => ([1,1,2,3], [7,13,5,7]))
+GNNHeteroGraph:
+  num_nodes: Dict(:movie => 13, :user => 3)
+  num_edges: Dict((:user, :rate, :movie) => 4)
+```
+New relations, possibly with new node types, can be added with the function [`add_edges`](@ref).
+```julia-repl
+julia> g = add_edges(g, (:user, :like, :actor) => ([1,2,3,3,3], [3,5,1,9,4]))
+GNNHeteroGraph:
+  num_nodes: Dict(:actor => 9, :movie => 13, :user => 3)
+  num_edges: Dict((:user, :like, :actor) => 5, (:user, :rate, :movie) => 4)
+```
+See [`rand_heterograph`](@ref), [`rand_bipartite_heterograph`](@ref)
+for generating random heterographs. 
+
+```julia-repl
+julia> g = rand_bipartite_heterograph((10, 15), 20)
+GNNHeteroGraph:
+  num_nodes: Dict(:A => 10, :B => 15)
+  num_edges: Dict((:A, :to, :B) => 20, (:B, :to, :A) => 20)
+```
+
+## Basic Queries
+
+Basic queries are similar to those for homogeneous graphs:
+```julia-repl
+julia> g = GNNHeteroGraph((:user, :rate, :movie) => ([1,1,2,3], [7,13,5,7]))
+GNNHeteroGraph:
+  num_nodes: Dict(:movie => 13, :user => 3)
+  num_edges: Dict((:user, :rate, :movie) => 4)
+
+julia> g.num_nodes
+Dict{Symbol, Int64} with 2 entries:
+  :user  => 3
+  :movie => 13
+
+julia> g.num_edges
+Dict{Tuple{Symbol, Symbol, Symbol}, Int64} with 1 entry:
+  (:user, :rate, :movie) => 4
+
+# source and target node for a given relation
+julia> edge_index(g, (:user, :rate, :movie))
+([1, 1, 2, 3], [7, 13, 5, 7])
+
+# node types
+julia> g.ntypes
+2-element Vector{Symbol}:
+ :user
+ :movie
+
+# edge types
+julia> g.etypes
+1-element Vector{Tuple{Symbol, Symbol, Symbol}}:
+ (:user, :rate, :movie)
+```
+
+## Data Features
+
+Node, edge, and graph features can be added at constuction time or later using:
+```julia-repl
+# equivalent to g.ndata[:user][:x] = ...
+julia> g[:user].x = rand(Float32, 64, 3);
+
+julia> g[:movie].z = rand(Float32, 64, 13);
+
+# equivalent to g.edata[(:user, :rate, :movie)][:e] = ...
+julia> g[:user, :rate, :movie].e = rand(Float32, 64, 4);
+
+julia> g
+GNNHeteroGraph:
+  num_nodes: Dict(:movie => 13, :user => 3)
+  num_edges: Dict((:user, :rate, :movie) => 4)
+  ndata:
+        :movie  =>  DataStore(z = [64×13 Matrix{Float32}])
+        :user  =>  DataStore(x = [64×3 Matrix{Float32}])
+  edata:
+        (:user, :rate, :movie)  =>  DataStore(e = [64×4 Matrix{Float32}])
+```
+
+

src/GNNGraphs/gnnheterograph.jl

@@ -145,9 +145,9 @@ function GNNHeteroGraph(data::EDict;
 end
 
 function show_sorted_dict(io::IO, d::Dict, compact::Bool)
-    if compact
+    # if compact
         print(io, "Dict")
-    end
+    # end
     print(io, "(")
     if !isempty(d)
         _keys = sort!(collect(keys(d)))
@@ -156,6 +156,9 @@ function show_sorted_dict(io::IO, d::Dict, compact::Bool)
         end
         print(io, "$(_str(_keys[end])) => $(d[_keys[end]])")
     end
+    # if length(d) == 1
+    #     print(io, ",")
+    # end
     print(io, ")")
 end
 
@@ -180,15 +183,17 @@ function Base.show(io::IO, ::MIME"text/plain", g::GNNHeteroGraph)
         print(io, "\n  num_edges: ")
         show_sorted_dict(io, g.num_edges, false)
         g.num_graphs > 1 && print(io, "\n  num_graphs: $(g.num_graphs)")
-        if !isempty(g.ndata)
+        if !isempty(g.ndata) && !all(isempty, values(g.ndata))
             print(io, "\n  ndata:")
             for k in sort(collect(keys(g.ndata)))
+                isempty(g.ndata[k]) && continue    
                 print(io, "\n\t", _str(k), "  =>  $(shortsummary(g.ndata[k]))")
             end
         end
-        if !isempty(g.edata)
+        if !isempty(g.edata) && !all(isempty, values(g.edata))
             print(io, "\n  edata:")
             for k in sort(collect(keys(g.edata)))
+                isempty(g.edata[k]) && continue
                 print(io, "\n\t$k  =>  $(shortsummary(g.edata[k]))")
             end
         end
@@ -271,20 +276,12 @@ end
 
 @non_differentiable _ntypes_from_edges(::Any...)
 
-
 function Base.getindex(g::GNNHeteroGraph, node_t::NType)
-    if !haskey(g.ndata, node_t) && node_t in g.ntypes
-        g.ndata[node_t] = DataStore(g.num_nodes[node_t])
-    end
     return g.ndata[node_t]
 end
 
-Base.setindex!(g::GNNHeteroGraph, node_t::NType, x) = g.ndata[node_t] = x
+Base.getindex(g::GNNHeteroGraph, n1_t::Symbol, rel::Symbol, n2_t::Symbol) = g[(n1_t, rel, n2_t)]
 
 function Base.getindex(g::GNNHeteroGraph, edge_t::EType)
-    if !haskey(g.edata, node_t) && edge_t in g.etypes
-        g.ndata[node_t] = DataStore(g.num_edges[edge_t])
-    end
     return g.edata[edge_t]
 end
-Base.setindex!(g::GNNHeteroGraph, edge_t::EType, x) = g.edata[edge_t] = x

src/GNNGraphs/transform.jl

@@ -120,7 +120,8 @@ end
     add_edges(g::GNNGraph, s::AbstractVector, t::AbstractVector; [edata])
 
 Add to graph `g` the edges with source nodes `s` and target nodes `t`.
-Optionally, pass the features  `edata` for the new edges. 
+Optionally, pass the features  `edata` for the new edges.
+Returns a new graph sharing part of the underlying data with `g`.
 """
 function add_edges(g::GNNGraph{<:COO_T},
                    snew::AbstractVector{<:Integer},
@@ -143,14 +144,33 @@ function add_edges(g::GNNGraph{<:COO_T},
              g.ndata, edata, g.gdata)
 end
 
+
+"""
+    add_edges(g::GNNHeteroGraph, edge_t, s, t; [edata, num_nodes])
+    add_edges(g::GNNHeteroGraph, edge_t => (s, t); [edata, num_nodes])
+
+Add to heterograph `g` the releation of type `edge_t` with source node vector `s` and target node vector `t`.
+Optionally, pass the features  `edata` for the new edges.
+`edge_t` is a triplet of symbols `(srctype, etype, dsttype)`. 
+
+If the edge type is not already present in the graph, it is added. If it involves new node types, they are added to the graph as well.
+In this case, a dictionary or named tuple of `num_nodes` can be passed to specify the number of nodes of the new types,
+otherwise the number of nodes is inferred from the maximum node id in `s` and `t`.
+"""
+add_edges(g::GNNHeteroGraph{<:COO_T}, data::Pair{EType, <:Tuple}; kws...) = add_edges(g, data.first, data.second...; kws...)
+
 function add_edges(g::GNNHeteroGraph{<:COO_T},
                    edge_t::EType,
                    snew::AbstractVector{<:Integer},
                    tnew::AbstractVector{<:Integer};
-                   edata = nothing)
+                   edata = nothing,
+                   num_nodes = Dict{Symbol,Int}())
     @assert length(snew) == length(tnew)
+    is_existing_rel = haskey(g.graph, edge_t)
     # TODO remove this constraint
-    @assert get_edge_weight(g, edge_t) === nothing
+    if is_existing_rel
+        @assert get_edge_weight(g, edge_t) === nothing
+    end
 
     edata = normalize_graphdata(edata, default_name = :e, n = length(snew))
     g_edata = g.edata |> copy
@@ -164,23 +184,41 @@ function add_edges(g::GNNHeteroGraph{<:COO_T},
 
     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)
+    ntypes = g.ntypes |> copy
+    _num_nodes = g.num_nodes |> copy
+    ndata = g.ndata |> copy
+    if !is_existing_rel
+        for (node_t, st) in [(edge_t[1], snew), (edge_t[3], tnew)]
+            if node_t ∉ ntypes
+                push!(ntypes, node_t)
+                if haskey(num_nodes, node_t)
+                    _num_nodes[node_t] == num_nodes[node_t]
+                else
+                    _num_nodes[node_t] = maximum(st)
+                end
+                ndata[node_t] = DataStore(_num_nodes[node_t])
+            end
+        end
+        push!(etypes, edge_t)
     else
         s, t = edge_index(g, edge_t)
         snew = [s; snew]
         tnew = [t; tnew]
     end
+    @assert maximum(snew) <= _num_nodes[edge_t[1]]
+    @assert maximum(tnew) <= _num_nodes[edge_t[3]]
+    @assert minimum(snew) >= 1
+    @assert minimum(tnew) >= 1
+
     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,
+            _num_nodes, num_edges, g.num_graphs,
              g.graph_indicator,
              g.ndata, g_edata, g.gdata,
-             g.ntypes, etypes)
+             ntypes, etypes)
 end
 
 

test/GNNGraphs/gnnheterograph.jl

@@ -34,8 +34,20 @@ end
     @test size(hg.ndata[:B].y) == (4, 20)
     @test size(hg.edata[(:A, :rel1, :B)].e) == (5, 30)
     @test hg.gdata == DataStore(u = 1)
+
+end
+
+@testset "indexing syntax" begin
+    g = GNNHeteroGraph((:user, :rate, :movie) => ([1,1,2,3], [7,13,5,7]))
+    g[:movie].z = rand(Float32, 64, 13);
+    g[:user, :rate, :movie].e = rand(Float32, 64, 4);
+    g[:user].x = rand(Float32, 64, 3);
+    @test size(g.ndata[:user].x) == (64, 3)
+    @test size(g.ndata[:movie].z) == (64, 13)
+    @test size(g.edata[(:user, :rate, :movie)].e) == (64, 4) 
 end
 
+
 @testset "simplified constructor" begin
     hg = rand_heterograph((:A => 10, :B => 20),
                             ((:A, :rel1, :B) => 30, (:B, :rel2, :A) => 10),

test/GNNGraphs/transform.jl

@@ -121,6 +121,23 @@ end
             @test has_edge(hg, (:B,:to,:A), 1, 2)
             @test !has_edge(hg, (:B,:to,:A), 2, 1)
             @test !has_edge(hg, (:B,:to,:A), 2, 2)
+
+            @testset "new nodes" begin
+                hg = rand_bipartite_heterograph((2, 2), 3)
+                hg = add_edges(hg, (:C,:rel,:B) => ([1, 3], [1,2]))
+                @test hg.num_nodes == Dict(:A => 2, :B => 2, :C => 3)
+                @test hg.num_edges == Dict((:A,:to,:B) => 3, (:B,:to,:A) => 3, (:C,:rel,:B) => 2)
+                s, t = edge_index(hg, (:C,:rel,:B))
+                @test s == [1, 3]
+                @test t == [1, 2]
+
+                hg = add_edges(hg, (:D,:rel,:F) => ([1, 3], [1,2]))
+                @test hg.num_nodes == Dict(:A => 2, :B => 2, :C => 3, :D => 3, :F => 2)
+                @test hg.num_edges == Dict((:A,:to,:B) => 3, (:B,:to,:A) => 3, (:C,:rel,:B) => 2, (:D,:rel,:F) => 2)
+                s, t = edge_index(hg, (:D,:rel,:F))
+                @test s == [1, 3]
+                @test t == [1, 2]
+            end
         end
     end 
 end