docs

Author: CarloLucibello

docs/src/messagepassing.md

@@ -1,7 +1,6 @@
 # Message Passing
 
-The message passing is initiated by the [`propagate`](@ref) function
-and generally takes the form
+A generic message passing on graph takes the form
 
 ```math
 \begin{aligned}
@@ -13,22 +12,41 @@ and generally takes the form
 ```
 
 where we refer to ``\phi`` as to the message function, 
-and to ``\gamma_x`` and ``\gamma_e`` as the node update and edge update function
-respectively. The generic aggregation ``\square`` usually is given by a summation
-``\sum``, a max or a mean operation. 
+and to ``\gamma_x`` and ``\gamma_e`` as to the node update and edge update function
+respectively. The aggregation ``\square`` is over the neighborhood ``N(i)`` of node ``i``, 
+and it is usually set to summation ``\sum``, a max or a mean operation. 
 
-The message propagation mechanism internally relies on the [`NNlib.gather`](@ref) 
+In GNN.jl, the function [`propagate`](@ref) takes care of materializing the
+node features on each edge, applying the message function, performing the
+aggregation, and returning ``\bar{\mathbf{m}}``. 
+It is then left to the user to perform further node and edge updates,
+manypulating arrays of size ``D_{node} \times num_nodes`` and   
+``D_{edge} \times num_edges``.
+
+As part of the [`propagate`](@ref) pipeline, we have the function
+[`apply_edges`](@ref). It can be independently used to materialize 
+node features on edges and perform edge-related computation without
+the following neighborhood aggregation one finds in `propagate`.
+
+The whole propagation mechanism internally relies on the [`NNlib.gather`](@ref) 
 and [`NNlib.scatter`](@ref) methods.
 
-## An example: implementing the GCNConv
 
-Let's (re-)implement the [`GCNConv`](@ref) layer use the message passing framework.
+## Examples
+
+### Basic use propagate and apply_edges 
+
+
+
+### Implementing a custom Graph Convolutional Layer
+
+Let's implement a simple graph convolutional layer using the message passing framework.
 The convolution reads 
 
 ```math
-\mathbf{x}'_i = \sum_{j \in N(i)} \frac{1}{c_{ij}} W \mathbf{x}_j
+\mathbf{x}'_i = W \cdot \sum_{j \in N(i)}  \mathbf{x}_j
 ```
-where ``c_{ij} = \sqrt{|N(i)||N(j)|}``. We will also add a bias and an activation function.
+We will also add a bias and an activation function.
 
 ```julia
 using Flux, LightGraphs, GraphNeuralNetworks
@@ -49,11 +67,18 @@ function GCN(ch::Pair{Int,Int}, σ=identity)
 end
 
 function (l::GCN)(g::GNNGraph, x::AbstractMatrix{T}) where T
-    c = 1 ./ sqrt.(degree(g, T, dir=:in))
-    x = x .* c'
-    x = propagate((xi, xj, e) -> l.weight * xj, g, +, xj=x)
-    x = x .* c'
-    return l.σ.(x .+ l.bias)
+    @assert size(x, 2) == g.num_nodes
+
+    # Computes messages from source/neighbour nodes (j) to target/root nodes (i).
+    # The message function will have to handle matrices of size (*, num_edges).
+    # In this simple case we just let the neighbor features go through.
+    message(xi, xj, e) = xj 
+
+    # The + operator gives the sum aggregation.
+    # `mean`, `max`, `min`, and `*` are other possibilities.
+    x = propagate(message, g, +, xj=x) 
+
+    return l.σ.(l.weight * x .+ l.bias)
 end
 ```
 

docs/src/models.md

@@ -54,6 +54,7 @@ model = GNN(din, d, dout)                 # step 5
 
 g = GNNGraph(random_regular_graph(10, 4))
 X = randn(Float32, din, 10) 
+
 y = model(g, X)  # output size: (dout, g.num_nodes)
 gs = gradient(() -> sum(model(g, X)), Flux.params(model))
 ```

src/msgpass.jl

@@ -1,26 +1,39 @@
 """
     propagate(f, g, aggr; xi, xj, e)  ->  m̄
 
-Performs the message passing scheme on graph `g`.
-Returns the aggregated node features `m̄` computed 
+Performs message passing on graph `g`.
 
-The computational steps are the following:
+Takes care of materializing the node features on each edge, 
+applying the message function, and returning an aggregated message ``\bar{\mathbf{m}}`` 
+(depending on the return value of `f`, an array or a named tuple of 
+arrays with last dimension's size `g.num_nodes`).
+
+It can be decomposed in two steps:
 
 ```julia
 m = apply_edges(f, g, xi, xj, e)
 m̄ = aggregate_neighbors(g, aggr, m)
 ```
 
-GNN layers typically call propagate in their forward pass.
+GNN layers typically call `propagate` in their forward pass,
+providing as input `f` a closure.  
 
 # Arguments
 
+- `g`: A `GNNGraph`.
+- `xi`: An array or a named tuple containing arrays whose last dimension's size 
+        is `g.num_nodes`. It will be appropriately materialized on the
+        target node of each edge (see also [`edge_index`](@ref)).
+- `xj`: As `xj`, but to be materialized on edges' sources. 
+- `e`: An array or a named tuple containing arrays whose last dimension's size is `g.num_edges`.
 - `f`: A generic function that will be passed over to [`apply_edges`](@ref). 
-      Takes as inputs `xi`, `xj`, and `e`
-       (target nodes' features, source nodes' features, and edge features
-       respetively) and returns new edge features `m`.
+      Has to take as inputs the edge-materialized `xi`, `xj`, and `e` 
+      (arrays or named tuples of arrays whose last dimension' size is the size of 
+      a batch of edges). Its output has to be an array or a named tuple of arrays
+      with the same batch size.
+- `aggr`: Neighborhood aggregation operator. Use `+`, `mean`, `max`, or `min`. 
 
-# Usage example
+# Usage Examples
 
 ```julia
 using GraphNeuralNetworks, Flux
@@ -68,29 +81,28 @@ end
 """
     apply_edges(f, xi, xj, e)
 
-Message function for the message-passing scheme
-started by [`propagate`](@ref).
 Returns the message from node `j` to node `i` .
 In the message-passing scheme, the incoming messages 
 from the neighborhood of `i` will later be aggregated
 in order to update the features of node `i`.
 
 The function operates on batches of edges, therefore
 `xi`, `xj`, and `e` are tensors whose last dimension
-is the batch size, or can be tuple/namedtuples of 
-such tensors, according to the input to propagate.
-
-By default, the function returns `xj`.
-Custom layer should specialize this method with the desired behavior.
-
+is the batch size, or can be named tuples of 
+such tensors.
+    
 # Arguments
 
-- `f`: A function that takes as inputs `xi`, `xj`, and `e`
-    (target nodes' features, source nodes' features, and edge features
-    respetively) and returns new edge features `m`.
-- `xi`: Features of the central node `i`.
-- `xj`: Features of the neighbor `j` of node `i`.
-- `eij`: Features of edge `(i,j)`.
+- `g`: A `GNNGraph`.
+- `xi`: An array or a named tuple containing arrays whose last dimension's size 
+        is `g.num_nodes`. It will be appropriately materialized on the
+        target node of each edge (see also [`edge_index`](@ref)).
+- `xj`: As `xj`, but to be materialized on edges' sources. 
+- `e`: An array or a named tuple containing arrays whose last dimension's size is `g.num_edges`.
+- `f`: A function that takes as inputs the edge-materialized `xi`, `xj`, and `e`.
+       These are arrays (or named tuples of arrays) whose last dimension' size is the size of
+       a batch of edges. The output of `f` has to be an array (or a named tuple of arrays)
+       with the same batch size. 
 
 See also [`propagate`](@ref).
 """