Merge pull request #47 from CarloLucibello/cl/redesign

redesign message passing mechanism

Project.toml

@@ -1,7 +1,7 @@
 name = "GraphNeuralNetworks"
 uuid = "cffab07f-9bc2-4db1-8861-388f63bf7694"
 authors = ["Carlo Lucibello and contributors"]
-version = "0.1.2"
+version = "0.2.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -25,7 +25,7 @@ Adapt = "3"
 CUDA = "3.3"
 ChainRulesCore = "1"
 DataStructures = "0.18"
-Flux = "0.12"
+Flux = "0.12.7"
 KrylovKit = "0.5"
 LearnBase = "0.4, 0.5"
 LightGraphs = "1.3"

README.md

@@ -5,13 +5,13 @@
 ![](https://github.com/CarloLucibello/GraphNeuralNetworks.jl/actions/workflows/ci.yml/badge.svg)
 [![codecov](https://codecov.io/gh/CarloLucibello/GraphNeuralNetworks.jl/branch/master/graph/badge.svg)](https://codecov.io/gh/CarloLucibello/GraphNeuralNetworks.jl)
 
-A graph neural network library for Julia based on the deep learning framework [Flux.jl](https://github.com/FluxML/Flux.jl).
-Its most relevant features are:
-* Provides CUDA support.
-* It's integrated with the JuliaGraphs ecosystem.
-* Implements many common graph convolutional layers.
-* Performs fast operations on batched graphs. 
-* Makes it easy to define custom graph convolutional layers.
+A graph neural network library for Julia based on the deep learning framework [Flux.jl](https://github.com/FluxML/Flux.jl). Its features include:
+
+* Integratation with the JuliaGraphs ecosystem.
+* Implementation of common graph convolutional layers.
+* Fast operations on batched graphs. 
+* Easy to define custom layers.
+* CUDA support.
 
 ## Installation
 
@@ -28,4 +28,4 @@ Usage examples can be found in the [examples](https://github.com/CarloLucibello/
 
 ## Acknowledgements
 
-A big thank you goes to @yuehhua for creating [GeometricFlux.jl](https://github.com/FluxML/GeometricFlux.jl) of which GraphNeuralNetworks.jl is a radical redesign. 
+A big thanks goes to @yuehhua for creating [GeometricFlux.jl](https://github.com/FluxML/GeometricFlux.jl) of which GraphNeuralNetworks.jl is a radical redesign. 

docs/src/api/messagepassing.md

@@ -11,11 +11,15 @@ Order = [:type, :function]
 Pages   = ["messagepassing.md"]
 ```
 
-## Docs
+## Interface
 
 ```@docs
-compute_message
-update_node
-update_edge
+apply_edges
 propagate
 ```
+
+## Built-in message functions
+
+```@docs
+copyxj
+```

docs/src/gnngraph.md

@@ -1,6 +1,6 @@
 # Graphs
 
-The fundamental graph type in GraphNeuralNetworks.jl is the [`GNNGraph`](@ref), 
+The fundamental graph type in GraphNeuralNetworks.jl is the [`GNNGraph`](@ref).
 A GNNGraph `g` is a directed graph with nodes labeled from 1 to `g.num_nodes`.
 The underlying implementation allows for efficient application of graph neural network
 operators, gpu movement, and storage of node/edge/graph related feature arrays.
@@ -32,25 +32,50 @@ g = GNNGraph(source, target)
 
 See also the related methods [`adjacency_matrix`](@ref), [`edge_index`](@ref), and [`adjacency_list`](@ref).
 
+## Basic Queries
+
+```julia
+source = [1,1,2,2,3,3,3,4]
+target = [2,3,1,3,1,2,4,3]
+g = GNNGraph(source, target)
+
+@assert g.num_nodes == 4   # number of nodes
+@assert g.num_edges == 8   # number of edges
+@assert g.num_graphs == 1  # number of subgraphs (a GNNGraph can batch many graphs together)
+is_directed(g)      # a GGNGraph is always directed
+```
 
 ## Data Features
 
+One or more arrays can be associated to nodes, edges, and (sub)graphs of a `GNNGraph`.
+They will be stored in the fields `g.ndata`, `g.edata`, and `g.gdata` respectivaly.
+The data fields are `NamedTuple`s. The array they contain must have last dimension
+equal to `num_nodes` (in `ndata`), `num_edges` (in `edata`), or `num_graphs` (in `gdata`).
+
 ```julia
 # Create a graph with a single feature array `x` associated to nodes
 g = GNNGraph(erdos_renyi(10,  30), ndata = (; x = rand(Float32, 32, 10)))
-# Equivalent definition
+
+g.ndata.x  # access the features
+
+# Equivalent definition passing directly the array
 g = GNNGraph(erdos_renyi(10,  30), ndata = rand(Float32, 32, 10))
 
+g.ndata.x  # `:x` is the default name for node features
+
 # You can have multiple feature arrays
 g = GNNGraph(erdos_renyi(10,  30), ndata = (; x=rand(Float32, 32, 10), y=rand(Float32, 10)))
 
+g.ndata.y, g.ndata.x
 
 # Attach an array with edge features.
 # Since `GNNGraph`s are directed, the number of edges
 # will be double that of the original LightGraphs' undirected graph.
 g = GNNGraph(erdos_renyi(10,  30), edata = rand(Float32, 60))
 @assert g.num_edges == 60
 
+g.edata.e
+
 # If we pass only half of the edge features, they will be copied
 # on the reversed edges.
 g = GNNGraph(erdos_renyi(10,  30), edata = rand(Float32, 30))
@@ -59,28 +84,30 @@ g = GNNGraph(erdos_renyi(10,  30), edata = rand(Float32, 30))
 # Create a new graph from previous one, inheriting edge data
 # but replacing node data
 g′ = GNNGraph(g, ndata =(; z = ones(Float32, 16, 10)))
-```
-
-
-## Graph Manipulation
 
-```julia
-g′ = add_self_loops(g)
-
-g′ = remove_self_loops(g)
+g.ndata.z
+g.edata.e
 ```
 
 ## Batches and Subgraphs
 
+Multiple `GNNGraph`s can be batched togheter into a single graph
+containing the total number of the original nodes 
+and where the original graphs are disjoint subgraphs.
+
 ```julia
 using Flux
 
 gall = Flux.batch([GNNGraph(erdos_renyi(10, 30), ndata=rand(Float32,3,10)) for _ in 1:160])
 
-g23 = getgraph(gall, 2:3)
+@assert gall.num_graphs == 160 
+@assert gall.num_nodes == 1600   # 10 nodes x 160 graphs
+@assert gall.num_edges == 9600  # 30 undirected edges x 2 directions x 160 graphs
+
+g23, _ = getgraph(gall, 2:3)
 @assert g23.num_graphs == 2
-@assert g23.num_nodes == 20
-@assert g23.num_edges == 120 # 30 undirected edges x 2 graphs
+@assert g23.num_nodes == 20   # 10 nodes x 160 graphs
+@assert g23.num_edges == 120  # 30 undirected edges x 2 directions x 2 graphs x
 
 
 # DataLoader compatibility
@@ -92,6 +119,17 @@ for g in train_loader
     @assert size(g.ndata.x) = (3, 160)    
     .....
 end
+
+# Access the nodes' graph memberships through 
+gall.graph_indicator
+```
+
+## Graph Manipulation
+
+```julia
+g′ = add_self_loops(g)
+
+g′ = remove_self_loops(g)
 ```
 
 ## JuliaGraphs ecosystem integration

docs/src/index.md

@@ -2,26 +2,29 @@
 
 This is the documentation page for the [GraphNeuralNetworks.jl](https://github.com/CarloLucibello/GraphNeuralNetworks.jl) library.
 
-A graph neural network library for Julia based on the deep learning framework [Flux.jl](https://github.com/FluxML/Flux.jl).
-Its most relevant features are:
-* Provides CUDA support.
-* It's integrated with the JuliaGraphs ecosystem.
-* Implements many common graph convolutional layers.
-* Performs fast operations on batched graphs. 
-* Makes it easy to define custom graph convolutional layers.
+A graph neural network library for Julia based on the deep learning framework [Flux.jl](https://github.com/FluxML/Flux.jl). GNN.jl is largely inspired by python's libraries [PyTorch Geometric](https://pytorch-geometric.readthedocs.io/en/latest/) and [Deep Graph Library](https://docs.dgl.ai/),
+and by julia's [GeometricFlux](https://fluxml.ai/GeometricFlux.jl/stable/).
+
+Among its features:
+
+* Integratation with the JuliaGraphs ecosystem.
+* Implementation of common graph convolutional layers.
+* Fast operations on batched graphs. 
+* Easy to define custom layers.
+* CUDA support.
 
 
 ## Package overview
 
-Let's give a brief overview of the package solving a  
-graph regression problem on fake data. 
+Let's give a brief overview of the package by solving a  
+graph regression problem with synthetic data. 
 
 Usage examples on real datasets can be found in the [examples](https://github.com/CarloLucibello/GraphNeuralNetworks.jl/tree/master/examples) folder. 
 
 ### Data preparation
 
 First, we create our dataset consisting in multiple random graphs and associated data features. 
-that we batch together into a unique graph.
+Then we batch the graphs together into a unique graph.
 
 ```julia
 julia> using GraphNeuralNetworks, LightGraphs, Flux, CUDA, Statistics
@@ -50,17 +53,17 @@ GNNGraph:
 
 ### Model building 
 
-We concisely define our model using as a [`GNNChain`](@ref) containing 2 graph convolutaional 
+We concisely define our model as a [`GNNChain`](@ref) containing 2 graph convolutaional 
 layers. If CUDA is available, our model will live on the gpu.
 
 ```julia
 julia> device = CUDA.functional() ? Flux.gpu : Flux.cpu;
 
 julia> model = GNNChain(GCNConv(16 => 64),
-                        BatchNorm(64),
-                        x -> relu.(x),
+                        BatchNorm(64),     # Apply batch normalization on node features (nodes dimension is batch dimension)
+                        x -> relu.(x),     
                         GCNConv(64 => 64, relu),
-                        GlobalPool(mean),
+                        GlobalPool(mean),  # aggregate node-wise features into graph-wise features
                         Dense(64, 1)) |> device;
 
 julia> ps = Flux.params(model);
@@ -75,8 +78,8 @@ Flux's DataLoader iterates over mini-batches of graphs
 (batched together into a `GNNGraph` object). 
 
 ```julia
-gtrain, _ = getgraph(gbatch, 1:800)
-gtest, _ = getgraph(gbatch, 801:gbatch.num_graphs)
+gtrain = getgraph(gbatch, 1:800)
+gtest = getgraph(gbatch, 801:gbatch.num_graphs)
 train_loader = Flux.Data.DataLoader(gtrain, batchsize=32, shuffle=true)
 test_loader = Flux.Data.DataLoader(gtest, batchsize=32, shuffle=false)
 
@@ -86,7 +89,7 @@ loss(loader) = mean(loss(g |> device) for g in loader)
 
 for epoch in 1:100
     for g in train_loader
-        g = g |> gpu
+        g = g |> device
         grad = gradient(() -> loss(g), ps)
         Flux.Optimise.update!(opt, ps, grad)
     end

docs/src/messagepassing.md

@@ -1,38 +1,55 @@
 # Message Passing
 
-The message passing is initiated by [`propagate`](@ref)
-and can be customized for a specific layer by overloading the methods
-[`compute_message`](@ref), [`update_node`](@ref), and [`update_edge`](@ref).
-
-The message passing corresponds to the following operations 
+A generic message passing on graph takes the form
 
 ```math
 \begin{aligned}
 \mathbf{m}_{j\to i} &= \phi(\mathbf{x}_i, \mathbf{x}_j, \mathbf{e}_{j\to i}) \\
-\mathbf{x}_{i}' &= \gamma_x(\mathbf{x}_{i}, \square_{j\in N(i)}  \mathbf{m}_{j\to i})\\
+\bar{\mathbf{m}}_{i} &= \square_{j\in N(i)}  \mathbf{m}_{j\to i} \\
+\mathbf{x}_{i}' &= \gamma_x(\mathbf{x}_{i}, \bar{\mathbf{m}}_{i})\\
 \mathbf{e}_{j\to i}^\prime &=  \gamma_e(\mathbf{e}_{j \to i},\mathbf{m}_{j \to i})
 \end{aligned}
 ```
-where ``\phi`` is expressed by the [`compute_message`](@ref) function, 
-``\gamma_x`` and ``\gamma_e`` by [`update_node`](@ref) and [`update_edge`](@ref)
-respectively.
 
-The message propagation mechanism internally relies on the [`NNlib.gather`](@ref) 
+where we refer to ``\phi`` as to the message function, 
+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. 
+
+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
-import GraphNeuralNetworks: compute_message, update_node, propagate
 
 struct GCN{A<:AbstractMatrix, B, F} <: GNNLayer
     weight::A
@@ -49,16 +66,26 @@ function GCN(ch::Pair{Int,Int}, σ=identity)
     GCN(W, b, σ)
 end
 
-compute_message(l::GCN, xi, xj, eij) = l.weight * xj
-update_node(l::GCN, m, x) = m
-
 function (l::GCN)(g::GNNGraph, x::AbstractMatrix{T}) where T
-    c = 1 ./ sqrt.(degree(g, T, dir=:in))
-    x = x .* c'
-    x, _ = propagate(l, g, +, 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
 ```
 
 See the [`GATConv`](@ref) implementation [here](https://github.com/CarloLucibello/GraphNeuralNetworks.jl/blob/master/src/layers/conv.jl) for a more complex example.
+
+
+## Built-in message functions
+
+In order to exploit optimized specializations of the [`propagate`](@ref), it is recommended 
+to use built-in message functions such as [`copyxj`](@ref) whenever possible.