Add GNNLux temporal convolutional layers docstrings and docs (#516)

* Add TGCN docs

* Add A3TGCN docs

* Add GConvGRU docs

* Add GConvLASTM docs

* Add DCGRU docs

* Add EvolveGNCO docs

* Change title

* Add temporal conv

Author: aurorarossi

GNNLux/docs/make.jl

@@ -23,7 +23,8 @@ makedocs(;
          pages = ["Home" => "index.md",
                    "API Reference" => [
                                     "Basic" => "api/basic.md",
-                                    "Convolutional layers" => "api/conv.md"]]
+                                    "Convolutional layers" => "api/conv.md",
+                                    "Temporal Convolutional layers" => "api/temporalconv.md",]]
          )
 
 

GNNLux/docs/src/api/basic.md

@@ -2,7 +2,8 @@
 CurrentModule = GNNLux
 ```
 
-## GNNLayer
+## Basic Layers
+
 ```@docs
 GNNLayer
 GNNChain

GNNLux/docs/src/api/temporalconv.md

@@ -0,0 +1,15 @@
+```@meta
+CurrentModule = GNNLux
+```
+
+# Temporal Graph-Convolutional Layers
+
+Convolutions for time-varying graphs (temporal graphs) such as the [`TemporalSnapshotsGNNGraph`](@ref).
+
+## Docs
+
+```@autodocs
+Modules = [GNNLux]
+Pages   = ["layers/temporalconv.jl"]
+Private = false
+```

GNNLux/src/layers/temporalconv.jl

@@ -56,8 +56,96 @@ function Base.show(io::IO, tgcn::TGCNCell)
     print(io, "TGCNCell($(tgcn.in_dims) => $(tgcn.out_dims))")
 end
 
+"""
+    TGCN(in => out; use_bias = true, init_weight = glorot_uniform, init_state = zeros32, init_bias = zeros32, add_self_loops = false, use_edge_weight = true) 
+
+Temporal Graph Convolutional Network (T-GCN) recurrent layer 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.
+- `use_bias`: Add learnable bias. Default `true`.
+- `init_weight`: Weights' initializer. Default `glorot_uniform`.
+- `init_state`: Initial state of the hidden stat of the GRU layer. Default `zeros32`.
+- `init_bias`: Bias initializer. Default `zeros32`.
+- `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`.
+
+
+
+# Examples
+
+```julia
+using GNNLux, Lux, Random
+
+# initialize random number generator
+rng = Random.default_rng()
+
+# create data
+g = rand_graph(rng, 5, 10)
+x = rand(rng, Float32, 2, 5)
+
+# create TGCN layer
+tgcn = TGCN(2 => 6)
+
+# setup layer
+ps, st = LuxCore.setup(rng, tgcn)
+
+# forward pass
+y, st = tgcn(g, x, ps, st)      # result size (6, 5)
+```
+"""
 TGCN(ch::Pair{Int, Int}; kwargs...) = GNNLux.StatefulRecurrentCell(TGCNCell(ch; kwargs...))
 
+"""
+    A3TGCN(in => out; use_bias = true, init_weight = glorot_uniform, init_state = zeros32, init_bias = zeros32, add_self_loops = false, use_edge_weight = true) 
+
+Attention Temporal Graph Convolutional Network (A3T-GCN) model from the paper [A3T-GCN: Attention Temporal Graph
+Convolutional Network for Traffic Forecasting](https://arxiv.org/pdf/2006.11583.pdf).
+
+Performs a TGCN layer, followed by a soft attention layer.
+
+# Arguments
+
+- `in`: Number of input features.
+- `out`: Number of output features.
+- `use_bias`: Add learnable bias. Default `true`.
+- `init_weight`: Weights' initializer. Default `glorot_uniform`.
+- `init_state`: Initial state of the hidden stat of the GRU layer. Default `zeros32`.
+- `init_bias`: Bias initializer. Default `zeros32`.
+- `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`.
+# Examples
+
+```julia
+using GNNLux, Lux, Random
+
+# initialize random number generator
+rng = Random.default_rng()
+
+# create data
+g = rand_graph(rng, 5, 10)
+x = rand(rng, Float32, 2, 5)
+
+# create A3TGCN layer
+l = A3TGCN(2 => 6)
+
+# setup layer
+ps, st = LuxCore.setup(rng, l)
+
+# forward pass
+y, st = l(g, x, ps, st)      # result size (6, 5)
+```
+"""
 @concrete struct A3TGCN <: GNNContainerLayer{(:tgcn, :dense1, :dense2)}
     in_dims::Int
     out_dims::Int
@@ -94,7 +182,7 @@ function Base.show(io::IO, l::A3TGCN)
     print(io, "A3TGCN($(l.in_dims) => $(l.out_dims))")
 end
 
-@concrete struct GConvGRUCell <: GNNContainerLayer{(:conv_x_r, :conv_h_r, :conv_x_z, :conv_h_z, :conv_x_h, :conv_h_h)}
+ @concrete struct GConvGRUCell <: GNNContainerLayer{(:conv_x_r, :conv_h_r, :conv_x_z, :conv_h_z, :conv_x_h, :conv_h_h)}
     in_dims::Int
     out_dims::Int
     k::Int
@@ -147,6 +235,45 @@ end
 
 LuxCore.outputsize(l::GConvGRUCell) = (l.out_dims,)
 
+"""
+    GConvGRU(in => out, k; use_bias = true, init_weight = glorot_uniform, init_state = zeros32, init_bias = zeros32) 
+
+Graph Convolutional Gated Recurrent Unit (GConvGRU) recurrent layer from the paper [Structured Sequence Modeling with Graph Convolutional Recurrent Networks](https://arxiv.org/pdf/1612.07659).
+
+Performs a layer of ChebConv 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.
+- `k`: Chebyshev polynomial order.
+- `use_bias`: Add learnable bias. Default `true`.
+- `init_weight`: Weights' initializer. Default `glorot_uniform`.
+- `init_state`: Initial state of the hidden stat of the GRU layer. Default `zeros32`.
+- `init_bias`: Bias initializer. Default `zeros32`.
+
+# Examples
+
+```julia
+using GNNLux, Lux, Random
+
+# initialize random number generator
+rng = Random.default_rng()
+
+# create data
+g = rand_graph(rng, 5, 10)
+x = rand(rng, Float32, 2, 5)
+
+# create layer
+l = GConvGRU(2 => 5, 2)
+
+# setup layer
+ps, st = LuxCore.setup(rng, l)
+
+# forward pass
+y, st = l(g, x, ps, st)      # result size (5, 5)
+```
+"""
 GConvGRU(ch::Pair{Int, Int}, k::Int; kwargs...) = GNNLux.StatefulRecurrentCell(GConvGRUCell(ch, k; kwargs...))
 
 @concrete struct GConvLSTMCell <: GNNContainerLayer{(:conv_x_i, :conv_h_i, :dense_i, :conv_x_f, :conv_h_f, :dense_f, :conv_x_c, :conv_h_c, :dense_c, :conv_x_o, :conv_h_o, :dense_o)}
@@ -230,6 +357,45 @@ end
 
 LuxCore.outputsize(l::GConvLSTMCell) = (l.out_dims,)
 
+"""
+    GConvLSTM(in => out, k; use_bias = true, init_weight = glorot_uniform, init_state = zeros32, init_bias = zeros32) 
+
+Graph Convolutional Long Short-Term Memory (GConvLSTM) recurrent layer from the paper [Structured Sequence Modeling with Graph Convolutional Recurrent Networks](https://arxiv.org/pdf/1612.07659). 
+
+Performs a layer of ChebConv to model spatial dependencies, followed by a Long Short-Term Memory (LSTM) cell to model temporal dependencies.
+
+# Arguments
+
+- `in`: Number of input features.
+- `out`: Number of output features.
+- `k`: Chebyshev polynomial order.
+- `use_bias`: Add learnable bias. Default `true`.
+- `init_weight`: Weights' initializer. Default `glorot_uniform`.
+- `init_state`: Initial state of the hidden stat of the GRU layer. Default `zeros32`.
+- `init_bias`: Bias initializer. Default `zeros32`.
+
+# Examples
+
+```julia
+using GNNLux, Lux, Random
+
+# initialize random number generator
+rng = Random.default_rng()
+
+# create data
+g = rand_graph(rng, 5, 10)
+x = rand(rng, Float32, 2, 5)
+
+# create GConvLSTM layer
+l = GConvLSTM(2 => 5, 2)
+
+# setup layer
+ps, st = LuxCore.setup(rng, l)
+
+# forward pass
+y, st = l(g, x, ps, st)      # result size (5, 5)
+```
+"""
 GConvLSTM(ch::Pair{Int, Int}, k::Int; kwargs...) = GNNLux.StatefulRecurrentCell(GConvLSTMCell(ch, k; kwargs...))
 
 @concrete struct DCGRUCell <: GNNContainerLayer{(:dconv_u, :dconv_r, :dconv_c)}
@@ -272,8 +438,86 @@ end
 
 LuxCore.outputsize(l::DCGRUCell) = (l.out_dims,)
 
+"""
+    DCGRU(in => out, k; use_bias = true, init_weight = glorot_uniform, init_state = zeros32, init_bias = zeros32) 
+
+Diffusion Convolutional Recurrent Neural Network (DCGRU) layer from the paper [Diffusion Convolutional Recurrent Neural
+Network: Data-driven Traffic Forecasting](https://arxiv.org/pdf/1707.01926).
+
+Performs a Diffusion Convolutional layer 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.
+- `k`: Diffusion step.
+- `use_bias`: Add learnable bias. Default `true`.
+- `init_weight`: Weights' initializer. Default `glorot_uniform`.
+- `init_state`: Initial state of the hidden stat of the GRU layer. Default `zeros32`.
+- `init_bias`: Bias initializer. Default `zeros32`.
+
+# Examples
+
+```julia
+using GNNLux, Lux, Random
+
+# initialize random number generator
+rng = Random.default_rng()
+
+# create data
+g = rand_graph(rng, 5, 10)
+x = rand(rng, Float32, 2, 5)
+
+# create layer
+l = DCGRU(2 => 5, 2)
+
+# setup layer
+ps, st = LuxCore.setup(rng, l)
+
+# forward pass
+y, st = l(g, x, ps, st)      # result size (5, 5)
+```
+"""
 DCGRU(ch::Pair{Int, Int}, k::Int; kwargs...) = GNNLux.StatefulRecurrentCell(DCGRUCell(ch, k; kwargs...))
 
+"""
+    EvolveGCNO(ch; use_bias = true, init_weight = glorot_uniform, init_state = zeros32, init_bias = zeros32)
+
+Evolving Graph Convolutional Network (EvolveGCNO) layer from the paper [EvolveGCN: Evolving Graph Convolutional Networks for Dynamic Graphs](https://arxiv.org/pdf/1902.10191).
+
+Perfoms a Graph Convolutional layer with parameters derived from a Long Short-Term Memory (LSTM) layer across the snapshots of the temporal graph.
+
+
+# Arguments
+
+- `in`: Number of input features.
+- `out`: Number of output features.
+- `use_bias`: Add learnable bias. Default `true`.
+- `init_weight`: Weights' initializer. Default `glorot_uniform`.
+- `init_state`: Initial state of the hidden stat of the GRU layer. Default `zeros32`.
+- `init_bias`: Bias initializer. Default `zeros32`.
+
+# Examples
+
+```julia
+using GNNLux, Lux, Random
+
+# initialize random number generator
+rng = Random.default_rng()
+
+# create data
+tg = TemporalSnapshotsGNNGraph([rand_graph(rng, 10, 20; ndata = rand(rng, 4, 10)), rand_graph(rng, 10, 14; ndata = rand(rng, 4, 10)), rand_graph(rng, 10, 22; ndata = rand(rng, 4, 10))])
+
+# create layer
+l = EvolveGCNO(4 => 5)
+
+# setup layer
+ps, st = LuxCore.setup(rng, l)
+
+# forward pass
+y, st = l(tg, tg.ndata.x , ps, st)      # result size 3, size y[1] (5, 10)
+```
+"""
 @concrete struct EvolveGCNO <: GNNLayer
     in_dims::Int
     out_dims::Int