TGCCN

Author: CarloLucibello

GraphNeuralNetworks/src/GraphNeuralNetworks.jl

@@ -55,8 +55,7 @@ export GNNRecurrence,
        GConvLSTM, GConvLSTMCell,
        DCGRU, DCGRUCell,
        EvolveGCNO, EvolveGCNOCell,
-       TGCN,
-       A3TGCN
+       TGCN, TGCNCell
 
 include("layers/pool.jl")
 export GlobalPool,

GraphNeuralNetworks/src/layers/temporalconv.jl

@@ -130,7 +130,7 @@ function (rnn::GNNRecurrence)(g, x)
     return rnn(g, x, initialstates(rnn))
 end
 
-function (rnn::GNNRecurrence)(g, x, state) where {T}
+function (rnn::GNNRecurrence)(g, x, state)
     return scan(rnn.cell, g, x, state)
 end
 
@@ -750,7 +750,60 @@ julia> size(y[end]) # (d_out, num_nodes[end])
 EvolveGCNO(args...; kws...) = GNNRecurrence(EvolveGCNOCell(args...; kws...))
 
 
+"""
+    TGCNCell(in => out; kws...)
+
+Recurrent graph convolutional cell from the paper
+[T-GCN: A Temporal Graph Convolutional
+Network for Traffic Prediction](https://arxiv.org/pdf/1811.05320).
+
+Uses two stacked [`GCNConv`](@ref) layers to model spatial dependencies,
+and a GRU mechanism to model temporal dependencies.
+
+`in` and `out` are the number of input and output node features, respectively.
+The keyword arguments are passed to the [`GCNConv`](@ref) constructor.
+
+# Forward 
+
+    cell(g::GNNGraph, x, [state])
+
+- `g`: The input graph.
+- `x`: The node features. It should be a matrix of size `in x num_nodes`.
+- `state`: The current state of the cell.
+    If not provided, it is generated by calling `Flux.initialstates(cell)`.
+    The state is a matrix of size `out x num_nodes`.
+
+Returns the updated node features and the updated state.
+
+# Examples
+
+```jldoctest
+julia> using GraphNeuralNetworks, Flux
+
+julia> num_nodes, num_edges = 5, 10;
+
+julia> d_in, d_out = 2, 3;
+
+julia> timesteps = 5;
+
+julia> g = rand_graph(num_nodes, num_edges);
 
+julia> x = [rand(Float32, d_in, num_nodes) for t in 1:timesteps];
+
+julia> cell = DCGRUCell(d_in => d_out, 2);
+
+julia> state = Flux.initialstates(cell);
+
+julia> y = state;
+
+julia> for xt in x
+           y, state = cell(g, xt, state)
+       end
+
+julia> size(y) # (d_out, num_nodes)
+(3, 5)
+```
+"""
 @concrete struct TGCNCell <: GNNLayer
     in::Int
     out::Int
@@ -795,4 +848,36 @@ end
 
 function Base.show(io::IO, cell::TGCNCell)
     print(io, "TGCNCell($(cell.in) => $(cell.out))")
-end
+end
+
+"""
+    TGCN(args...; kws...)
+
+Construct a recurrent layer corresponding to the [`TGCNCell`](@ref) cell.
+
+The arguments are passed to the [`TGCNCell`](@ref) constructor.
+See [`GNNRecurrence`](@ref) for more details.
+
+# Examples
+
+```jldoctest
+julia> num_nodes, num_edges = 5, 10;
+
+julia> d_in, d_out = 2, 3;
+
+julia> timesteps = 5;
+
+julia> g = rand_graph(num_nodes, num_edges);
+
+julia> x = rand(Float32, d_in, timesteps, num_nodes);
+
+julia> layer = TGCN(d_in => d_out)
+
+julia> y = layer(g, x);
+
+julia> size(y) # (d_out, timesteps, num_nodes)
+(3, 5, 5)
+```
+"""
+TGCN(args...; kws...) = GNNRecurrence(TGCNCell(args...; kws...))
+