tests, docs, deprecations

Author: CarloLucibello

Project.toml

@@ -5,6 +5,7 @@ version = "0.1.2"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"

docs/src/api/messagepassing.md

@@ -14,8 +14,6 @@ Pages   = ["messagepassing.md"]
 ## Docs
 
 ```@docs
-compute_message
-update_node
-update_edge
+apply_edges
 propagate
 ```

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

@@ -3,12 +3,14 @@
 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.
+
+Among its features:
+
+* Integratation with the JuliaGraphs ecosystem.
+* Provides common graph convolutional layers.
+* Fast operations on batched graphs. 
+* Easy to define custom graph convolutional layers.
+* CUDA support.
 
 
 ## Package overview
@@ -50,17 +52,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);
@@ -86,7 +88,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,22 +1,19 @@
 # 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 
+The message passing is initiated by the [`propagate`](@ref) function
+and generally takes the form
 
 ```math
 \begin{aligned}
 \mathbf{m}_{j\to i} &= \phi(\mathbf{x}_i, \mathbf{x}_j, \mathbf{e}_{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)
+
+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. 
 
@@ -35,7 +32,6 @@ where ``c_{ij} = \sqrt{|N(i)||N(j)|}``. We will also add a bias and an activatio
 
 ```julia
 using Flux, LightGraphs, GraphNeuralNetworks
-import GraphNeuralNetworks: compute_message, update_node, propagate
 
 struct GCN{A<:AbstractMatrix, B, F} <: GNNLayer
     weight::A
@@ -52,12 +48,10 @@ function GCN(ch::Pair{Int,Int}, σ=identity)
     GCN(W, b, σ)
 end
 
-compute_message(l::GCN, xi, xj, eij) = l.weight * xj
-
 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 = propagate((xi, xj, e) -> l.weight * xj, g, +, xj=x)
     x = x .* c'
     return l.σ.(x .+ l.bias)
 end

docs/src/models.md

@@ -15,13 +15,12 @@ In the explicit modeling style, the model is created according to the following
 1. Define a new type for your model (`GNN` in the example below). Layers and submodels are fields.
 2. Apply `Flux.@functor` to the new type to make it Flux's compatible (parameters' collection, gpu movement, etc...)
 3. Optionally define a convenience constructor for your model.
-4. Define the forward pass by implementing the function call method for your type
+4. Define the forward pass by implementing the call method for your type.
 5. Instantiate the model. 
 
 Here is an example of this construction:
 ```julia
 using Flux, LightGraphs, GraphNeuralNetworks
-using Flux: @functor
 
 struct GNN                                # step 1
     conv1
@@ -31,7 +30,7 @@ struct GNN                                # step 1
     dense
 end
 
-@functor GNN                              # step 2
+Flux.@functor GNN                              # step 2
 
 function GNN(din::Int, d::Int, dout::Int) # step 3    
     GNN(GCNConv(din => d),
@@ -51,10 +50,12 @@ function (model::GNN)(g::GNNGraph, x)     # step 4
 end
 
 din, d, dout = 3, 4, 2 
-g = GNNGraph(random_regular_graph(10, 4))
-X = randn(Float32, din, 10)
 model = GNN(din, d, dout)                 # step 5
-y = model(g, X)
+
+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))
 ```
 
 ## Implicit modeling with GNNChains
@@ -81,8 +82,6 @@ model = GNNChain(GCNConv(din => d),
                  GCNConv(d => d, relu),
                  Dropout(0.5),
                  Dense(d, dout))
-
-y = model(g, X)  # output size: (dout, g.num_nodes)
 ```
 
 The `GNNChain` only propagates the graph and the node features. More complex scenarios, e.g. when also edge features are updated, have to be handled using the explicit definition of the forward pass. 

src/GraphNeuralNetworks.jl

@@ -31,7 +31,7 @@ export
     sprand, sparse, 
 
     # msgpass
-    update_node, update_edge, compute_message, propagate,
+    apply_edges, propagate,
 
     # layers/basic
     GNNLayer,

src/deprecations.jl

@@ -2,9 +2,30 @@
 
 @deprecate GINConv(nn; eps=0, aggr=+)  GINConv(nn, eps; aggr)
 
-# TO Deprecate
-# x, _ = propagate(l, g, l.aggr, x, e)
 
-# # TODO deprecate
-# propagate(l, g::GNNGraph, aggr, x, e=nothing) = propagate(l, g, aggr; x, e)
+# Deprecated in v0.2
+# TODO check if argument order is exact
+function compute_message end
+function update_node end
+function update_edge end
 
+compute_message(l, xi, xj, e) = compute_message(l, xi, xj)
+update_node(l, x, m̄) = m̄
+update_edge(l, e, m) = e
+
+function propagate(l::GNNLayer, g::GNNGraph, aggr, x, e=nothing)
+    @warn """
+          Passing a GNNLayer to propagate is deprecated, 
+          you should pass directly the message function.
+          The new signature is `propagate(f, g, aggr; [xi, xj, e])`.
+
+          Also the functions `compute_message`, `update_node`,
+          and `update_edge` have been deprecated. Please
+          refer to the documentation.
+          """
+    m = apply_edge((a...) -> compute_message(l, a...), g, x, x, e)
+    m̄ = aggregate_neighbors(g, aggr, m)
+    x = update_node(l, x, m̄)
+    e = update_edge(l, e, m)
+    return x, e
+end

src/gnngraph.jl

@@ -50,10 +50,10 @@ from the LightGraphs' graph library can be used on it.
     - `:sparse`. A sparse adjacency matrix representation.
     - `:dense`. A dense adjacency matrix representation.  
     Default `:coo`.
-- `dir`. The assumed edge direction when given adjacency matrix or adjacency list input data `g`. 
+- `dir`: The assumed edge direction when given adjacency matrix or adjacency list input data `g`. 
         Possible values are `:out` and `:in`. Default `:out`.
-- `num_nodes`. The number of nodes. If not specified, inferred from `g`. Default `nothing`.
-- `graph_indicator`. For batched graphs, a vector containing the graph assigment of each node. Default `nothing`.  
+- `num_nodes`: The number of nodes. If not specified, inferred from `g`. Default `nothing`.
+- `graph_indicator`: For batched graphs, a vector containing the graph assigment of each node. Default `nothing`.  
 - `ndata`: Node features. A named tuple of arrays whose last dimension has size num_nodes.
 - `edata`: Edge features. A named tuple of arrays whose whose last dimension has size num_edges.
 - `gdata`: Global features. A named tuple of arrays whose has size num_graphs.