Temporal Graph-Convolutional Layers

DCGRUCell(in => out, k; use_bias = true, init_weight = glorot_uniform, init_bias = zeros32)

Diffusion Convolutional Recurrent Neural Network (DCGRU) cell from the paper Diffusion Convolutional Recurrent Neural Network: Data-driven Traffic Forecasting.

Uses a DConv layer to model spatial dependencies, in combination with a Gated Recurrent Unit (GRU) cell to model temporal dependencies.

Arguments

• in => out: A pair where in is the number of input node features and out the number of output node features.
• k: Diffusion step for the DConv.
• use_bias: Add learnable bias. Default true.
• init_weight: Convolution weights' initializer. Default glorot_uniform.
• init_bias: Bias initializer. Default zeros32.

Forward

cell(g, x, ps, st)
cell(g, (x, h), ps, st)

Performs one recurrence step and returns (h, h), st, where h is the updated hidden state of size out x num_nodes. If the carry h is not provided, it is initialized to zeros.

EvolveGCNOCell(in => out; use_bias = true, init_weight = glorot_uniform, init_bias = zeros32)

Evolving Graph Convolutional Network cell of type "-O" from the paper EvolveGCN: Evolving Graph Convolutional Networks for Dynamic Graphs.

Uses a GCNConv layer whose weight matrix is evolved across the temporal sequence by an LSTMCell. It can work with time-varying graphs and node features.

Arguments

• in => out: A pair where in is the number of input node features and out the number of output node features.
• use_bias: Add learnable bias for the convolution and the LSTM cell. Default true.
• init_weight: Weights' initializer. Default glorot_uniform.
• init_bias: Bias initializer. Default zeros32.

Forward

cell(g, x, ps, st)
cell(g, (x, state), ps, st)

Performs one recurrence step and returns (y, state), st, where y is the convolution output of size out x num_nodes and state is the updated (weight, lstm_carry) carry. If the carry is not provided, the convolution weight is initialized from the parameters and the LSTM carry to zeros.

GConvGRUCell(in => out, k; use_bias = true, init_weight = glorot_uniform, init_bias = zeros32)

Graph Convolutional Gated Recurrent Unit (GConvGRU) recurrent cell from the paper Structured Sequence Modeling with Graph Convolutional Recurrent Networks.

Uses ChebConv to model spatial dependencies, followed by a Gated Recurrent Unit (GRU) cell to model temporal dependencies.

Arguments

• in => out: A pair where in is the number of input node features and out the number of output node features.
• k: Chebyshev polynomial order.
• use_bias: Add learnable bias. Default true.
• init_weight: Weights' initializer. Default glorot_uniform.
• init_bias: Bias initializer. Default zeros32.

Forward

cell(g, x, ps, st)
cell(g, (x, h), ps, st)

Performs one recurrence step and returns (h, h), st, where h is the updated hidden state of size out x num_nodes. If the carry h is not provided, it is initialized to zeros.

GConvLSTMCell(in => out, k; use_bias = true, init_weight = glorot_uniform, init_bias = zeros32)

Graph Convolutional Long Short-Term Memory (GConvLSTM) recurrent cell from the paper Structured Sequence Modeling with Graph Convolutional Recurrent Networks.

Uses ChebConv to model spatial dependencies, followed by a Long Short-Term Memory (LSTM) cell with peephole connections to model temporal dependencies.

Arguments

• in => out: A pair where in is the number of input node features and out the number of output node features.
• k: Chebyshev polynomial order.
• use_bias: Add learnable bias. Default true.
• init_weight: Weights' initializer. Default glorot_uniform.
• init_bias: Bias initializer. Default zeros32.

Forward

cell(g, x, ps, st)
cell(g, (x, (h, c)), ps, st)

Performs one recurrence step and returns (h, (h, c)), st, where h is the updated hidden state and c the updated cell state, both of size out x num_nodes. If the carry (h, c) is not provided, it is initialized to zeros.

GNNRecurrence(cell; return_sequence = true)

Recurrent layer analogous to Lux.Recurrence that wraps a graph recurrent cell and applies it over an entire temporal sequence of node features at once.

The cell has to follow the recurrent-cell interface (out, carry), st = cell(g, (x, carry), ps, st), with the convenience method (out, carry), st = cell(g, x, ps, st) initializing the carry to zeros.

The layer constructors TGCN, GConvGRU, GConvLSTM, DCGRU and EvolveGCNO all return a GNNRecurrence wrapping the corresponding cell.

Arguments

• cell: A graph recurrent cell (e.g. TGCNCell).
• return_sequence: If true the whole sequence of outputs is returned, otherwise only the last output. Default true.

Forward

layer(g, x, ps, st)

• g: The input GNNGraph or TemporalSnapshotsGNNGraph.
  • If a GNNGraph, the same graph is used at every timestep.
  • If a TemporalSnapshotsGNNGraph, a different graph (snapshot) is used at each timestep. Not all cells support this.
• x: The time-varying node features.
  • If g is a GNNGraph, an array of size in x timesteps x num_nodes.
  • If g is a TemporalSnapshotsGNNGraph, a vector of length timesteps whose t-th element has size in x num_nodes_t.

Returns the updated node features and state:

• If return_sequence == true and g is a GNNGraph, the output is an array of size out x timesteps x num_nodes; if g is a TemporalSnapshotsGNNGraph, it is a vector of length timesteps.
• If return_sequence == false, only the last timestep's output is returned.

Examples

using GNNLux, Lux, Random

rng = Random.default_rng()
num_nodes, num_edges = 5, 10
d_in, d_out, timesteps = 2, 3, 5

g = rand_graph(rng, num_nodes, num_edges)
x = rand(rng, Float32, d_in, timesteps, num_nodes)

cell = GConvLSTMCell(d_in => d_out, 2)
layer = GNNRecurrence(cell)
ps, st = LuxCore.setup(rng, layer)

y, st = layer(g, x, ps, st)   # size(y) == (d_out, timesteps, num_nodes)

TGCNCell(in => out; use_bias = true, init_weight = glorot_uniform, init_bias = zeros32,
         add_self_loops = true, use_edge_weight = false, act = relu)

Recurrent graph convolutional cell from the paper T-GCN: A Temporal Graph Convolutional Network for Traffic Prediction.

Uses two stacked GCNConv layers to model spatial dependencies and a GRU mechanism to model temporal dependencies.

Arguments

• in => out: A pair where in is the number of input node features and out the number of output node features.
• use_bias: Add learnable bias. Default true.
• init_weight: Convolution weights' initializer. Default glorot_uniform.
• init_bias: Bias initializer. Default zeros32.
• add_self_loops: Add self loops to the graph before the convolution. Default true.
• use_edge_weight: If true, consider the edge weights in the input graph (if available). Default false.
• act: Activation function of the first GCNConv layer. Default relu.

Forward

cell(g, x, ps, st)
cell(g, (x, h), ps, st)

Performs one recurrence step and returns (h, h), st, where h is the updated hidden state of size out x num_nodes. If the carry h is not provided, it is initialized to zeros.

Examples

using GNNLux, Lux, Random

rng = Random.default_rng()
g = rand_graph(rng, 5, 10)
x = rand(rng, Float32, 2, 5)

cell = TGCNCell(2 => 6)
ps, st = LuxCore.setup(rng, cell)
(y, h), st = cell(g, x, ps, st)   # size(y) == (6, 5)

DCGRU(in => out, k; kws...)

Construct a GNNRecurrence layer from a DCGRUCell. The arguments are passed to the DCGRUCell constructor.

Examples

using GNNLux, Lux, Random

rng = Random.default_rng()
g = rand_graph(rng, 5, 10)
x = rand(rng, Float32, 2, 5, 5)  # (in, timesteps, num_nodes)

layer = DCGRU(2 => 5, 2)
ps, st = LuxCore.setup(rng, layer)
y, st = layer(g, x, ps, st)   # size(y) == (5, 5, 5)

EvolveGCNO(in => out; kws...)

Construct a GNNRecurrence layer from an EvolveGCNOCell. It can process an entire temporal sequence of graphs and node features at once. The arguments are passed to the EvolveGCNOCell constructor.

Examples

using GNNLux, Lux, Random

rng = Random.default_rng()
tg = TemporalSnapshotsGNNGraph([rand_graph(rng, 10, 20), rand_graph(rng, 10, 14), rand_graph(rng, 10, 22)])
x = [rand(rng, Float32, 4, 10) for _ in 1:tg.num_snapshots]

layer = EvolveGCNO(4 => 5)
ps, st = LuxCore.setup(rng, layer)
y, st = layer(tg, x, ps, st)   # length(y) == 3, size(y[1]) == (5, 10)

GConvGRU(in => out, k; kws...)

Construct a GNNRecurrence layer from a GConvGRUCell. The arguments are passed to the GConvGRUCell constructor.

Examples

using GNNLux, Lux, Random

rng = Random.default_rng()
g = rand_graph(rng, 5, 10)
x = rand(rng, Float32, 2, 5, 5)  # (in, timesteps, num_nodes)

layer = GConvGRU(2 => 5, 2)
ps, st = LuxCore.setup(rng, layer)
y, st = layer(g, x, ps, st)   # size(y) == (5, 5, 5)

GConvLSTM(in => out, k; kws...)

Construct a GNNRecurrence layer from a GConvLSTMCell. The arguments are passed to the GConvLSTMCell constructor.

Examples

using GNNLux, Lux, Random

rng = Random.default_rng()
g = rand_graph(rng, 5, 10)
x = rand(rng, Float32, 2, 5, 5)  # (in, timesteps, num_nodes)

layer = GConvLSTM(2 => 5, 2)
ps, st = LuxCore.setup(rng, layer)
y, st = layer(g, x, ps, st)   # size(y) == (5, 5, 5)

TGCN(in => out; kws...)

Construct a GNNRecurrence layer from a TGCNCell. The arguments are passed to the TGCNCell constructor.

Examples

using GNNLux, Lux, Random

rng = Random.default_rng()
g = rand_graph(rng, 5, 10)
x = rand(rng, Float32, 2, 5, 5)  # (in, timesteps, num_nodes)

layer = TGCN(2 => 6)
ps, st = LuxCore.setup(rng, layer)
y, st = layer(g, x, ps, st)   # size(y) == (6, 5, 5)