extend reduce_nodes and graph_indicator for heterographs (#339)

* graph indicator and reduce nodes

* more tests

Author: CarloLucibello

docs/src/heterograph.md

@@ -2,7 +2,7 @@
 
 Heterogeneous graphs (also called heterographs), are graphs where each node has a type,
 that we denote with symbols such as `:user` and `:movie`.
-Also edges have a type, such as `:rate` or `:like`, and they can connect nodes of different types. We call a triplet `(source_node_type, edge_type, target_node_type)` the type of a *relation*, e.g. `(:user, :rate, :movie)`.
+Releations such as `:rate` or `:like` can connect nodes of different types. We call a triplet `(source_node_type, relation_type, target_node_type)` the type of a edge, e.g. `(:user, :rate, :movie)`.
 
 Different node/edge types can store different groups of features
 and this makes heterographs a very flexible modeling tools 
@@ -12,7 +12,7 @@ the type [`GNNHeteroGraph`](@ref).
 
 ## Creating a Heterograph
 
-A heterograph can be created by passing pairs of relation type and data to the constructor.
+A heterograph can be created by passing pairs `edge_type => data` to the constructor.
 ```jldoctest
 julia> g = GNNHeteroGraph((:user, :like, :actor) => ([1,2,2,3], [1,3,2,9]),
                           (:user, :rate, :movie) => ([1,1,2,3], [7,13,5,7]))

src/GNNGraphs/query.jl

@@ -13,11 +13,11 @@ 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)`.
+for each edges in `g` of type `edge_t = (src_t, rel_t, trg_t)`.
 
 If `edge_t` is not provided, it will error if `g` has more than one edge type.
 """
@@ -454,12 +454,13 @@ _rand_dense_vector(A::AbstractMatrix{T}) where {T} = randn(float(T), size(A, 1))
 # https://discourse.julialang.org/t/cuda-eigenvalues-of-a-sparse-matrix/46851/5
 
 """
-    graph_indicator(g)
+    graph_indicator(g::GNNGraph; edges=false)
 
 Return a vector containing the graph membership
 (an integer from `1` to `g.num_graphs`) of each node in the graph.
+If `edges=true`, return the graph membership of each edge instead.
 """
-function graph_indicator(g; edges = false)
+function graph_indicator(g::GNNGraph; edges = false)
     if isnothing(g.graph_indicator)
         gi = ones_like(edge_index(g)[1], Int, g.num_nodes)
     else
@@ -473,6 +474,29 @@ function graph_indicator(g; edges = false)
     end
 end
 
+"""
+    graph_indicator(g::GNNHeteroGraph, [node_t])
+
+Return a Dict of vectors containing the graph membership
+(an integer from `1` to `g.num_graphs`) of each node in the graph for each node type.
+If `node_t` is provided, return the graph membership of each node of type `node_t` instead.
+
+See also [`batch`](@ref).
+"""
+function graph_indicator(g::GNNHeteroGraph)
+    return g.graph_indicator
+end
+
+function graph_indicator(g::GNNHeteroGraph, node_t::Symbol)
+    @assert node_t ∈ g.ntypes
+    if isnothing(g.graph_indicator)
+        gi = ones_like(edge_index(g, first(g.etypes))[1], Int, g.num_nodes[node_t])
+    else
+        gi = g.graph_indicator[node_t]
+    end
+    return gi
+end
+
 function node_features(g::GNNGraph)
     if isempty(g.ndata)
         return nothing

src/GNNGraphs/transform.jl

@@ -193,15 +193,16 @@ function add_edges(g::GNNGraph{<:COO_T}, data::COO_T; edata = nothing)
 end
 
 """
-    add_edges(g::GNNHeteroGraph, rel_t, s, t; [edata, num_nodes])
-    add_edges(g::GNNHeteroGraph, rel_t => (s, t); [edata, num_nodes])
-    add_edges(g::GNNHeteroGraph, rel_t => (s, t, w); [edata, num_nodes])
+    add_edges(g::GNNHeteroGraph, edge_t, s, t; [edata, num_nodes])
+    add_edges(g::GNNHeteroGraph, edge_t => (s, t); [edata, num_nodes])
+    add_edges(g::GNNHeteroGraph, edge_t => (s, t, w); [edata, num_nodes])
 
-Add to heterograph `g` the relation of type `rel_t` with source node vector `s` and target node vector `t`.
+Add to heterograph `g` edges of type `edge_t` with source node vector `s` and target node vector `t`.
 Optionally, pass the  edge weights `w` or the features  `edata` for the new edges.
-`rel_t` is a triplet of symbols `(src_t, edge_t, dst_t)`. 
+`edge_t` is a triplet of symbols `(src_t, rel_t, dst_t)`. 
 
-If the relation is not already present in the graph, it is added. If it involves new node types, they are added to the graph as well.
+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`.
 """

src/utils.jl

@@ -6,13 +6,27 @@ ofeltype(x, y) = convert(float(eltype(x)), y)
 For a batched graph `g`, return the graph-wise aggregation of the node
 features `x`. The aggregation operator `aggr` can be `+`, `mean`, `max`, or `min`.
 The returned array will have last dimension `g.num_graphs`.
+
+See also: [`reduce_edges`](@ref).
 """
 function reduce_nodes(aggr, g::GNNGraph, x)
     @assert size(x)[end] == g.num_nodes
     indexes = graph_indicator(g)
     return NNlib.scatter(aggr, x, indexes)
 end
 
+"""
+    reduce_nodes(aggr, indicator::AbstractVector, x)
+
+Return the graph-wise aggregation of the node features `x` given the
+graph indicator `indicator`. The aggregation operator `aggr` can be `+`, `mean`, `max`, or `min`.
+
+See also [`graph_indicator`](@ref).
+"""
+function reduce_nodes(aggr, indicator::AbstractVector, x)
+    return NNlib.scatter(aggr, x, indicator)
+end
+
 """
     reduce_edges(aggr, g, e)
 

test/GNNGraphs/query.jl

@@ -232,3 +232,19 @@ end
         @test eltype(khop_adj(g, 10, Float32)) == Float32
     end
 end
+
+if GRAPH_T == :coo
+    @testset "HeteroGraph" begin
+        @testset "graph_indicator" begin
+            gs = [rand_heterograph(Dict(:user => 10, :movie => 20, :actor => 30), 
+                                Dict((:user,:like,:movie) => 10,
+                                     (:actor,:rate,:movie)=>20)) for _ in 1:3]
+            g = MLUtils.batch(gs)
+            @test graph_indicator(g) == Dict(:user => [repeat([1], 10); repeat([2], 10); repeat([3], 10)],
+                                      :movie => [repeat([1], 20); repeat([2], 20); repeat([3], 20)],
+                                        :actor => [repeat([1], 30); repeat([2], 30); repeat([3], 30)])
+            @test graph_indicator(g, :movie) == [repeat([1], 20); repeat([2], 20); repeat([3], 20)]
+        end 
+    end
+end       
+

test/utils.jl

@@ -10,6 +10,9 @@ e = g.edata.e
     r = reduce_nodes(mean, g, x)
     @test size(r) == (Dx, g.num_graphs)
     @test r[:, 2] ≈ mean(getgraph(g, 2).ndata.x, dims = 2)
+
+    r2 = reduce_nodes(mean, graph_indicator(g), x)
+    @test r2 == r
 end
 
 @testset "reduce_edges" begin