export aggregate_neighbors; improve docstrings

Author: CarloLucibello

docs/src/api/messagepassing.md

@@ -15,6 +15,7 @@ Pages   = ["messagepassing.md"]
 
 ```@docs
 apply_edges
+aggregate_neighbors
 propagate
 ```
 

src/GraphNeuralNetworks.jl

@@ -21,14 +21,21 @@ using .GNNGraphs: COO_T, ADJMAT_T, SPARSE_T,
 
 export
     # utils
-    reduce_nodes, reduce_edges, 
-    softmax_nodes, softmax_edges,
-    broadcast_nodes, broadcast_edges,
+    reduce_nodes, 
+    reduce_edges, 
+    softmax_nodes, 
+    softmax_edges,
+    broadcast_nodes, 
+    broadcast_edges,
     softmax_edge_neighbors,
 
     # msgpass
-    apply_edges, propagate,
-    copy_xj, copy_xi, xi_dot_xj,
+    apply_edges,
+    aggregate_neighbors, 
+    propagate,
+    copy_xj, 
+    copy_xi, 
+    xi_dot_xj,
 
     # layers/basic
     GNNLayer,

src/layers/conv.jl

@@ -824,10 +824,33 @@ function (l::AGNNConv)(g::GNNGraph, x::AbstractMatrix)
 end
 
 @doc raw"""
+    MEGNetConv(ϕe, ϕv; aggr=mean)
     MEGNetConv(in => out; aggr=mean)
 
 Convolution from [Graph Networks as a Universal Machine Learning Framework for Molecules and Crystals](https://arxiv.org/pdf/1812.05055.pdf)
-paper.
+paper. In the forward pass, takes as inputs node features `x` and edge features `e` and returns
+updated features `x̄, ē` according to 
+
+```math
+ē  = ϕe(vcat(xi, xj, e))
+x̄  = ϕv(vcat(x, \square_{j\in \mathcal{N}(i)} ē_{j\to i}))
+```
+`aggr` defines the aggregation to be performed.
+
+If the neural networks `ϕe` and  `ϕv` are not provided, they will be constructed from
+the `in` and `out` arguments instead as multi-layer perceptron with one hidden layer and `relu` 
+activations.
+````
+
+# Examples
+
+```julia
+g = rand_graph(10, 30)
+x = randn(3, 10)
+e = randn(3, 30)
+m = MEGNetConv(3 => 3)
+x̄, ē = m(g, x, e)
+```
 """
 struct MEGNetConv <: GNNLayer
     ϕe
@@ -837,6 +860,8 @@ end
 
 @functor MEGNetConv
 
+MEGNetConv(ϕe, ϕv; aggr=mean) = MEGNetConv(ϕe, ϕv, aggr)
+
 function MEGNetConv(ch::Pair{Int,Int}; aggr=mean)
     nin, nout = ch 
     ϕe = Chain(Dense(3nin, nout, relu),
@@ -845,7 +870,7 @@ function MEGNetConv(ch::Pair{Int,Int}; aggr=mean)
     ϕv = Chain(Dense(nin + nout, nout, relu),
                Dense(nout, nout))
 
-    MEGNetConv(ϕe, ϕv, aggr)
+    MEGNetConv(ϕe, ϕv; aggr)
 end
 
 function (l::MEGNetConv)(g::GNNGraph)

src/msgpass.jl

@@ -61,7 +61,7 @@ l = GNNConv(10 => 20)
 l(g, x)
 ```
 
-See also [`apply_edges`](@ref).
+See also [`apply_edges`](@ref) and [`aggregate_neighbors`](@ref).
 """
 function propagate end 
 
@@ -103,7 +103,7 @@ such tensors.
        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).
+See also [`propagate`](@ref) and [`aggregate_neighbors`](@ref).
 """
 function apply_edges end 
 
@@ -125,7 +125,19 @@ _gather(x::Nothing, i) = nothing
 
 
 ##  AGGREGATE NEIGHBORS
+@doc raw"""
+    aggregate_neighbors(g::GNNGraph, aggr, m)
+
+Given a graph `g`, edge features `m`, and an aggregation
+operator `aggr` (e.g `+, min, max, mean`), returns the new node
+features 
+```math
+\mathbf{x}_i = \square_{j \in \mathcal{N}(i)} \mathbf{m}_{j\to i}
+```
 
+Neighborhood aggregation is the second step of [`propagate`](@ref), 
+where it comes after [`apply_edges`](@ref).
+"""
 function aggregate_neighbors(g::GNNGraph, aggr, m)
     s, t = edge_index(g)
     return _scatter(aggr, m, t)

test/layers/conv.jl

@@ -186,7 +186,7 @@
     end
 
     @testset "MEGNetConv" begin
-        l = MEGNetConv(in_channel, edim => out_channel, tanh, aggr=+)
+        l = MEGNetConv(in_channel => out_channel, tanh, aggr=+)
         for g in test_graphs
             g = GNNGraph(g, edata=rand(T, in_channel, g.num_edges))
             test_layer(l, g, rtol=1e-5, outsize=(out_channel, g.num_nodes))