Flux.batch for GNNHeteroGraph (#309)

* batchhetero

* fix all

Author: CarloLucibello

src/GNNGraphs/transform.jl

@@ -472,6 +472,60 @@ function Flux.batch(g::GNNGraph)
     throw(ArgumentError("Cannot batch a `GNNGraph` (containing $(g.num_graphs) graphs). Pass a vector of `GNNGraph`s instead."))
 end
 
+
+function Flux.batch(gs::AbstractVector{<:GNNHeteroGraph{T}}) where {T <: COO_T}
+    @assert length(gs) > 0
+    ntypes = union([g.ntypes for g in gs]...)
+    etypes = union([g.etypes for g in gs]...)
+    # TODO remove these constraints
+    @assert ntypes == gs[1].ntypes
+    @assert etypes == gs[1].etypes
+    
+    v_num_nodes = Dict(node_t => [get(g.num_nodes,node_t,0) for g in gs] for node_t in ntypes)
+    num_nodes = Dict(node_t => sum(v_num_nodes[node_t]) for node_t in ntypes)
+    num_edges = Dict(edge_t => sum(g.num_edges[edge_t] for g in gs) for edge_t in etypes)
+    edge_indices = Dict(edge_t => [edge_index(g, edge_t) for g in gs] for edge_t in etypes)
+    nodesum = Dict(node_t => cumsum([0; v_num_nodes[node_t]])[1:(end - 1)] for node_t in ntypes)
+    graphs = []
+    for edge_t in etypes
+        src_t, _, dst_t = edge_t
+        # @show edge_t edge_indices[edge_t] first(edge_indices[edge_t])
+        # for ei in edge_indices[edge_t]
+        #     @show ei[1]
+        # end 
+        # # [ei[1] for (ii, ei) in enumerate(edge_indices[edge_t])]
+        s = cat_features([ei[1] .+ nodesum[src_t][ii] for (ii, ei) in enumerate(edge_indices[edge_t])])
+        t = cat_features([ei[2] .+ nodesum[dst_t][ii] for (ii, ei) in enumerate(edge_indices[edge_t])])
+        w = cat_features([get_edge_weight(g, edge_t) for g in gs])
+        push!(graphs, edge_t => (s, t, w))
+    end
+    graph = Dict(graphs...)
+
+    #TODO relax this restriction
+    @assert all(g -> g.num_graphs == 1, gs) 
+
+    s = edge_index(gs[1], etypes[1])[1] # grab any source vector
+
+    function materialize_graph_indicator(g, node_t)     
+        ones_like(s, g.num_nodes[node_t])
+    end
+    v_gi = Dict(node_t => [materialize_graph_indicator(g, node_t) for g in gs] for node_t in ntypes)
+    v_num_graphs = [g.num_graphs for g in gs]
+    graphsum = cumsum([0; v_num_graphs])[1:(end - 1)]
+    v_gi = Dict(node_t => [ng .+ gi for (ng, gi) in zip(graphsum, v_gi[node_t])] for node_t in ntypes)
+    graph_indicator = Dict(node_t => cat_features(v_gi[node_t]) for node_t in ntypes)
+
+    GNNHeteroGraph(graph,
+             num_nodes,
+             num_edges,
+             sum(v_num_graphs),
+             graph_indicator,
+             cat_features([g.ndata for g in gs]),
+             cat_features([g.edata for g in gs]),
+             cat_features([g.gdata for g in gs]),
+             ntypes, etypes)
+end
+
 """
     unbatch(g::GNNGraph)
 

src/GNNGraphs/utils.jl

@@ -83,9 +83,14 @@ function cat_features(x1::Dict{Symbol, T}, x2::Dict{Symbol, T}) where {T}
     sort(collect(keys(x1))) == sort(collect(keys(x2))) ||
         @error "cannot concatenate feature data with different keys"
 
-    return Dict{Symbol, T}(k => cat_features(x1[k], x2[k]) for k in keys(x1))
+    return Dict{Symbol, T}([k => cat_features(x1[k], x2[k]) for k in keys(x1)]...)
 end
 
+function cat_features(x::Dict)
+    return Dict([k => cat_features(v) for (k, v) in pairs(x)]...)
+end
+
+
 function cat_features(xs::AbstractVector{<:AbstractArray{T, N}}) where {T <: Number, N}
     cat(xs...; dims = N)
 end
@@ -104,17 +109,30 @@ function cat_features(xs::AbstractVector{<:NamedTuple})
     NamedTuple(k => cat_features([x[k] for x in xs]) for k in syms)
 end
 
-function cat_features(xs::AbstractVector{Dict{Symbol, T}}) where {T}
-    symbols = [sort(collect(keys(x))) for x in xs]
-    all(y -> y == symbols[1], symbols) ||
+# function cat_features(xs::AbstractVector{Dict{Symbol, T}}) where {T}
+#     symbols = [sort(collect(keys(x))) for x in xs]
+#     all(y -> y == symbols[1], symbols) ||
+#         @error "cannot concatenate feature data with different keys"
+#     length(xs) == 1 && return xs[1]
+
+#     # concatenate 
+#     syms = symbols[1]
+#     return Dict{Symbol, T}([k => cat_features([x[k] for x in xs]) for k in syms]...)
+# end
+
+function cat_features(xs::AbstractVector{<:Dict})
+    _allkeys = [sort(collect(keys(x))) for x in xs]
+    _keys = _allkeys[1]
+    all(y -> y == _keys, _allkeys) ||
         @error "cannot concatenate feature data with different keys"
     length(xs) == 1 && return xs[1]
 
     # concatenate 
-    syms = symbols[1]
-    return Dict{Symbol, T}(k => cat_features([x[k] for x in xs]) for k in syms)
+    return Dict([k => cat_features([x[k] for x in xs]) for k in _keys]...)
 end
 
+
+
 # Turns generic type into named tuple
 normalize_graphdata(data::Nothing; n, kws...) = DataStore(n)
 

test/GNNGraphs/transform.jl

@@ -292,3 +292,27 @@ end
                      0.5 1.0 0.5
                      0.0 0.0 0.0]
 end
+
+@testset "batch heterograph" begin
+    gs = [rand_bipartite_heterograph((10, 15), 20) for _ in 1:5]
+    g = Flux.batch(gs)
+    @test g.num_nodes[:A] == 50
+    @test g.num_nodes[:B] == 75
+    @test g.num_edges[(:A,:to,:B)] == 100
+    @test g.num_edges[(:B,:to,:A)] == 100
+    @test g.num_graphs == 5
+    @test g.graph_indicator == Dict(:A => vcat([fill(i, 10) for i in 1:5]...),
+                                    :B => vcat([fill(i, 15) for i in 1:5]...))
+
+    for gi in gs
+        gi.ndata[:A].x = ones(2, 10)
+        gi.ndata[:A].y = zeros(10)
+        gi.edata[(:A,:to,:B)].e = fill(2, 20)
+        gi.gdata.u = 7
+    end
+    g = Flux.batch(gs)
+    @test g.ndata[:A].x == ones(2, 50)
+    @test g.ndata[:A].y == zeros(50)
+    @test g.edata[(:A,:to,:B)].e == fill(2, 100)
+    @test g.gdata.u == fill(7, 5)
+end