Add TGCNCell layer (#314)

* Add `TGCNCell` layer

* Export

* Add show function

* Add docstring

* Use `Flux.GRUv3Cell`

* Improve

Author: aurorarossi

src/GraphNeuralNetworks.jl

@@ -70,6 +70,9 @@ export
 # layers/heteroconv
       HeteroGraphConv,
 
+# layers/temporalconv
+      TGCNCell,
+
 # layers/pool
       GlobalPool,
       GlobalAttentionPool,
@@ -84,6 +87,7 @@ include("utils.jl")
 include("layers/basic.jl")
 include("layers/conv.jl")
 include("layers/heteroconv.jl")
+include("layers/temporalconv.jl")
 include("layers/pool.jl")
 include("msgpass.jl")
 include("mldatasets.jl")

src/layers/temporalconv.jl

@@ -0,0 +1,52 @@
+"""
+    TGCNCell(in => out; [bias, init, add_self_loops, use_edge_weight])
+
+Temporal Graph Convolutional Network (T-GCN) cell from the paper [T-GCN: A Temporal Graph Convolutional Network for Traffic Prediction](https://arxiv.org/pdf/1811.05320.pdf).
+
+Performs a layer of GCNConv to model spatial dependencies, followed by a Gated Recurrent Unit (GRU) cell to model temporal dependencies.
+
+# Arguments
+
+- `in`: Number of input features.
+- `out`: Number of output features.
+- `bias`: Add learnable bias. Default `true`.
+- `init`: Weights' initializer. Default `glorot_uniform`.
+- `add_self_loops`: Add self loops to the graph before performing the convolution. Default `false`.
+- `use_edge_weight`: If `true`, consider the edge weights in the input graph (if available).
+                     If `add_self_loops=true` the new weights will be set to 1. 
+                     This option is ignored if the `edge_weight` is explicitly provided in the forward pass.
+                     Default `false`.
+"""
+struct TGCNCell <: GNNLayer
+    conv::GCNConv
+    gru::Flux.GRUv3Cell
+    state0
+    in::Int
+    out::Int
+end
+
+Flux.@functor TGCNCell
+
+function TGCNCell(ch::Pair{Int, Int};
+                  bias::Bool = true,
+                  init = Flux.glorot_uniform,
+                  init_state = Flux.zeros32,
+                  add_self_loops = false,
+                  use_edge_weight = true)
+    in, out = ch
+    conv = GCNConv(in => out, sigmoid; init, bias, add_self_loops,
+                   use_edge_weight)
+    gru = Flux.GRUv3Cell(out, out)
+    state0 = init_state(out,1)
+    return TGCNCell(conv, gru, state0, in,out)
+end
+
+function (tgcn::TGCNCell)(h, g::GNNGraph, x::AbstractArray)
+    x̃ = tgcn.conv(g, x)
+    h, x̃ = tgcn.gru(h, x̃)
+    return h, x̃
+end
+
+function Base.show(io::IO, tgcn::TGCNCell)
+    print(io, "TGCNCell($(tgcn.in) => $(tgcn.out))")
+end

test/layers/temporalconv.jl

@@ -0,0 +1,16 @@
+in_channel = 3
+out_channel = 5
+N = 4
+T = Float32
+
+g1 = GNNGraph(rand_graph(N,8),
+                ndata = rand(T, in_channel, N),
+                graph_type = :sparse)
+
+@testset "TGCNCell" begin
+    tgcn = TGCNCell(in_channel => out_channel)
+    h, x̃ = tgcn(tgcn.state0, g1, g1.ndata.x)
+    @test size(h) == (out_channel, N)
+    @test size(x̃) == (out_channel, N)
+    @test h == x̃
+end

test/runtests.jl

@@ -41,6 +41,7 @@ tests = [
     "layers/basic",
     "layers/conv",
     "layers/heteroconv",
+    "layers/temporalconv",
     "layers/pool",
     "mldatasets",
     "examples/node_classification_cora",