HeteroGraph implementation (#146)

* update based on Dicts

* suggestion

* add rand_heterograph

* normalize data

* simplify construction

* cleanup docs

* update show graph

* fix test

* renaming

* cl/hetero

* fix test

Author: CarloLucibello

Project.toml

@@ -47,10 +47,11 @@ julia = "1.6"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
+DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 InlineStrings = "842dd82b-1e85-43dc-bf29-5d0ee9dffc48"
 MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test", "Adapt", "InlineStrings", "Zygote", "FiniteDifferences", "ChainRulesTestUtils", "MLDatasets"]
+test = ["Test", "Adapt", "DataFrames", "InlineStrings", "Zygote", "FiniteDifferences", "ChainRulesTestUtils", "MLDatasets"]

docs/make.jl

@@ -24,6 +24,7 @@ makedocs(;
              "Message Passing" => "messagepassing.md",
              "Model Building" => "models.md",
              "Datasets" => "datasets.md",
+             "HeteroGraphs" => "gnnheterograph.md",
              "Tutorials" => tutorials,
              "API Reference" =>
                [

docs/src/gnngraph.md

@@ -47,8 +47,8 @@ julia> target = [2,3,1,3,1,2,4,3];
 
 julia> g = GNNGraph(source, target)
 GNNGraph:
-    num_nodes = 4
-    num_edges = 8
+  num_nodes: 4
+  num_edges: 8
 
 
 julia> @assert g.num_nodes == 4   # number of nodes
@@ -88,7 +88,7 @@ g = rand_graph(10,  60, ndata = rand(Float32, 32, 10))
 
 g.ndata.x  # `:x` is the default name for node features
 
-# For convinience, we can access the features through the shortcut
+# For convenience, we can access the features through the shortcut
 g.x 
 
 # You can have multiple feature arrays
@@ -131,8 +131,8 @@ julia> weight = [1.0, 0.5, 2.1, 2.3, 4, 4.1];
 
 julia> g = GNNGraph(source, target, weight)
 GNNGraph:
-    num_nodes = 3
-    num_edges = 6
+  num_nodes: 3
+  num_edges: 6
 
 julia> get_edge_weight(g)
 6-element Vector{Float64}:
@@ -221,10 +221,10 @@ using Flux: gpu
 g_gpu = g |> gpu
 ```
 
-## JuliaGraphs/Graphs.jl integration
+## Integraton with Graphs.jl integration
 
-Since `GNNGraph <: Graphs.AbstractGraph`, we can use any functionality from Graphs.jl. 
-Moreover, `GNNGraph`s can be constructed from `Graphs.Graph` and `Graphs.DiGraph`.
+Since `GNNGraph <: Graphs.AbstractGraph`, we can use any functionality from [Graphs.jl](https://github.com/JuliaGraphs/Graphs.jl) for querying and analyzing the graph structure. 
+Moreover, a `GNNGraph` can be easily constructed from a `Graphs.Graph` or a `Graphs.DiGraph`:
 
 ```julia
 julia> import Graphs
@@ -237,17 +237,17 @@ julia> gu = Graphs.erdos_renyi(10, 20)
 
 # Since GNNGraphs are undirected, the edges are doubled when converting 
 # to GNNGraph
-julia> GNNGraph(gu)  # Since GNNGraphs are 
+julia> GNNGraph(gu)
 GNNGraph:
-    num_nodes = 10
-    num_edges = 40
+  num_nodes: 10
+  num_edges: 40
 
 # A Graphs.jl directed graph
 julia> gd = Graphs.erdos_renyi(10, 20, is_directed=true)
 {10, 20} directed simple Int64 graph
 
 julia> GNNGraph(gd)
 GNNGraph:
-    num_nodes = 10
-    num_edges = 20
+  num_nodes: 10
+  num_edges: 20
 ```

docs/src/gnnheterograph.md

@@ -0,0 +1,25 @@
+# Heterogeneous Graphs
+
+!!! warning
+    Heterographs support is still experimental.
+    The interface could be subject to change in the future.
+
+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)`.
+
+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 
+the type [`GNNHeteroGraph`](@ref).
+
+
+```@docs
+GNNHeteroGraph
+rand_heterograph
+```
+
+

examples/node_classification_cora.jl

@@ -44,10 +44,9 @@ function train(; kws...)
     dataset = Cora()
     classes = dataset.metadata["classes"]
     g = mldataset2gnngraph(dataset) |> device
-    X = g.ndata.features
-    y = onehotbatch(g.ndata.targets |> cpu, classes) |> device # remove when https://github.com/FluxML/Flux.jl/pull/1959 tagged
-    (; train_mask, val_mask, test_mask) = g.ndata
-    ytrain = y[:,train_mask]
+    X = g.features
+    y = onehotbatch(g.targets |> cpu, classes) |> device # remove when https://github.com/FluxML/Flux.jl/pull/1959 tagged
+    ytrain = y[:, g.train_mask]
 
     nin, nhidden, nout = size(X,1), args.nhidden, length(classes)
 
@@ -63,8 +62,8 @@ function train(; kws...)
 
     ## LOGGING FUNCTION
     function report(epoch)
-        train = eval_loss_accuracy(X, y, train_mask, model, g)
-        test = eval_loss_accuracy(X, y, test_mask, model, g)        
+        train = eval_loss_accuracy(X, y, g.train_mask, model, g)
+        test = eval_loss_accuracy(X, y, g.test_mask, model, g)        
         println("Epoch: $epoch   Train: $(train)   Test: $(test)")
     end
 
@@ -73,7 +72,7 @@ function train(; kws...)
     for epoch in 1:args.epochs
         gs = Flux.gradient(ps) do
             ŷ = model(g, X)
-            logitcrossentropy(ŷ[:,train_mask], ytrain)
+            logitcrossentropy(ŷ[:, g.train_mask], ytrain)
         end
 
         Flux.Optimise.update!(opt, ps, gs)

src/GNNGraphs/GNNGraphs.jl

@@ -22,6 +22,9 @@ export GNNGraph,
        edge_features, 
        graph_features
 
+include("gnnheterograph.jl")
+export GNNHeteroGraph
+
 include("query.jl")
 export adjacency_list,
        edge_index,
@@ -60,6 +63,7 @@ export add_nodes,
 
 include("generate.jl")
 export rand_graph, 
+       rand_heterograph,    
        knn_graph,
        radius_graph
 
@@ -75,4 +79,6 @@ include("utils.jl")
 include("gatherscatter.jl")
 # _gather, _scatter
 
+
+
 end #module

src/GNNGraphs/convert.jl

@@ -1,13 +1,55 @@
 ### CONVERT_TO_COO REPRESENTATION ########
 
-function to_coo(coo::COO_T; dir=:out, num_nodes=nothing, weighted=true)
+function to_coo(data::EDict; num_nodes = nothing, kws...)
+    graph = EDict{Any}()
+    _num_nodes = NDict{Int}()
+    num_edges = EDict{Int}()
+    for k in keys(data)
+        d = data[k]
+        @assert d isa Tuple
+        if length(d) == 2
+            d = (d..., nothing)
+        end
+        if num_nodes !== nothing
+            n1 = get(num_nodes, k[1], nothing)
+            n2 = get(num_nodes, k[3], nothing)   
+        else
+            n1 = nothing
+            n2 = nothing
+        end
+        g, nnodes, nedges = to_coo(d; hetero=true, num_nodes=(n1,n2), kws...)
+        graph[k] = g
+        num_edges[k] = nedges
+        _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
+    return graph, _num_nodes, num_edges
+end
+
+function to_coo(coo::COO_T; dir=:out, num_nodes=nothing, weighted=true, hetero=false)
     s, t, val = coo   
-    num_nodes::Int = isnothing(num_nodes) ? max(maximum(s), maximum(t)) : num_nodes 
+    
+    if isnothing(num_nodes)
+        ns = maximum(s)
+        nt = maximum(t)
+        num_nodes = hetero ? (ns, nt) : max(ns, nt)
+    elseif num_nodes isa Integer
+        ns = num_nodes
+        nt = num_nodes
+    elseif num_nodes isa Tuple
+        ns = isnothing(num_nodes[1]) ? maximum(s) : num_nodes[1]
+        nt = isnothing(num_nodes[2]) ? maximum(t) : num_nodes[2]
+        num_nodes = (ns, nt)
+    else
+        error("Invalid num_nodes $num_nodes")
+    end
     @assert isnothing(val) || length(val) == length(s)
     @assert length(s) == length(t)
     if !isempty(s)
-        @assert min(minimum(s), minimum(t)) >= 1 
-        @assert max(maximum(s), maximum(t)) <= num_nodes 
+        @assert minimum(s) >= 1
+        @assert minimum(t) >= 1
+        @assert maximum(s) <= ns
+        @assert maximum(t) <= nt
     end
     num_edges = length(s)
     if !weighted

src/GNNGraphs/generate.jl

@@ -45,6 +45,47 @@ function rand_graph(n::Integer, m::Integer; bidirected=true, seed=-1, kws...)
     return GNNGraph(Graphs.erdos_renyi(n, m2; is_directed=!bidirected, seed); kws...)    
 end
 
+"""
+    rand_heterograph(n, m; seed=-1, 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.
+
+Use a `seed > 0` for reproducibility.
+
+Additional keyword arguments will be passed to the [`GNNHeteroGraph`](@ref) constructor.
+
+# Examples
+
+```juliarepl
+
+
+julia> g = rand_heterograph((:user => 10, :movie => 20),
+                            (:user, :rate, :movie) => 30)
+GNNHeteroGraph:
+  num_nodes: (:user => 10, :movie => 20)         
+  num_edges: ((:user, :rate, :movie) => 30,)
+```
+"""
+rand_heteropraph
+
+# for generic iterators of pairs
+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
+    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_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
 
 """
     knn_graph(points::AbstractMatrix, 

src/GNNGraphs/gnngraph.jl

@@ -230,27 +230,27 @@ function Base.show(io::IO, ::MIME"text/plain", g::GNNGraph)
         print(io, "GNNGraph($(g.num_nodes), $(g.num_edges))")
     else # if the following block is indented the printing is ruined
     print(io, "GNNGraph:
-    num_nodes = $(g.num_nodes)
-    num_edges = $(g.num_edges)")
-    g.num_graphs > 1 && print(io, "\n    num_graphs = $(g.num_graphs)")
-    if !isempty(g.ndata)
-        print(io, "\n    ndata:")
-        for k in keys(g.ndata)
-            print(io, "\n        $k => $(summary(g.ndata[k]))")
-        end
-    end
-    if !isempty(g.edata)
-        print(io, "\n    edata:")
-        for k in keys(g.edata)
-            print(io, "\n        $k => $(summary(g.edata[k]))")
-        end
-    end
-    if !isempty(g.gdata)
-        print(io, "\n    gdata:")
-        for k in keys(g.gdata)
-            print(io, "\n        $k => $(summary(g.gdata[k]))")
-        end
-    end
+  num_nodes: $(g.num_nodes)
+  num_edges: $(g.num_edges)")
+  g.num_graphs > 1 && print(io, "\n    num_graphs = $(g.num_graphs)")
+  if !isempty(g.ndata)
+      print(io, "\n  ndata:")
+      for k in keys(g.ndata)
+        print(io, "\n    $k = $(shortsummary(g.ndata[k]))")
+      end
+  end
+  if !isempty(g.edata)
+      print(io, "\n  edata:")
+      for k in keys(g.edata)
+        print(io, "\n    $k = $(shortsummary(g.edata[k]))")
+      end
+  end
+  if !isempty(g.gdata)
+      print(io, "\n  gdata:")
+      for k in keys(g.gdata)
+        print(io, "\n    $k = $(shortsummary(g.gdata[k]))")
+      end
+  end
     end #else
 end