update docs

Author: yuehhua

docs/bibliography.bib

@@ -183,3 +183,17 @@ @inproceedings{Hamilton2017
    title = {Inductive Representation Learning on Large Graphs},
    year = {2017},
 }
+
+@inproceedings{Satorras2021,
+   abstract = {This paper introduces a new model to learn graph neural networks equivariant to rotations, translations, reflections and permutations called E(n)-Equivariant Graph Neural Networks (EGNNs). In contrast with existing methods, our work does not require computationally expensive higher-order representations in intermediate layers while it still achieves competitive or better performance. In addition, whereas existing methods are limited to equivariance on 3 dimensional spaces, our model is easily scaled to higher-dimensional spaces. We demonstrate the effectiveness of our method on dynamical systems modelling, representation learning in graph autoencoders and predicting molecular properties.},
+   author = {Victor Garcia Satorras and Emiel Hoogeboom and Max Welling},
+   editor = {Marina Meila and Tong Zhang},
+   booktitle = {Proceedings of the 38th International Conference on Machine Learning},
+   month = {2},
+   pages = {9323-9332},
+   publisher = {PMLR},
+   title = {E(n) Equivariant Graph Neural Networks},
+   volume = {139},
+   url = {http://arxiv.org/abs/2102.09844},
+   year = {2021},
+}

docs/make.jl

@@ -42,8 +42,10 @@ makedocs(
              "Dynamic Graph Update" => "dynamicgraph.md",
              "Manual" => [
                "FeaturedGraph" => "manual/featuredgraph.md",
-               "Graph Convolutional Layers" => "manual/conv.md",
+               "Graph Convolutional Layers" => "manual/graph_conv.md",
                "Graph Pooling Layers" => "manual/pool.md",
+               "Group Convolutional Layers" => "manual/group_conv.md",
+               "Positional Encoding Layers" => "manual/positional.md",
                "Embeddings" => "manual/embedding.md",
                "Models" => "manual/models.md",
                "Linear Algebra" => "manual/linalg.md",

docs/src/manual/featuredgraph.md

@@ -9,6 +9,8 @@ GraphSignals.edge_feature
 GraphSignals.has_edge_feature
 GraphSignals.global_feature
 GraphSignals.has_global_feature
+GraphSignals.positional_feature
+GraphSignals.has_positional_feature
 GraphSignals.subgraph
 GraphSignals.ConcreteFeaturedGraph
 ```

docs/src/manual/conv.md renamed to docs/src/manual/graph_conv.md

@@ -0,0 +1,19 @@
+# Group Convolutional Layers
+
+## ``E(n)``-equivariant Convolutional Layer
+
+It employs message-passing scheme and can be defined by following functions:
+
+- message function (Eq. 3 from paper): ``m_{ij} = \phi_e(h_i^l, h_j^l, ||x_i^l - x_j^l||^2, a_{ij})``
+- aggregate (Eq. 5 from paper): ``m_i = \sum_j m_{ij}``
+- update function (Eq. 6 from paper): ``h_i^{l+1} = \phi_h(h_i^l, m_i)``
+
+where ``h_i^l`` and ``h_j^l`` denotes the node feature for node ``i`` and ``j``, respectively, in ``l``-th layer, as well as ``x_i^l`` and ``x_j^l`` denote the positional feature for node ``i`` and ``j``, respectively, in ``l``-th layer. ``a_{ij}`` is the edge feature for edge ``(i,j)``. ``\phi_e`` and ``\phi_h`` are neural network for edges and nodes.
+
+```@docs
+EEquivGraphConv
+```
+
+Reference: [Satorras2021](@cite)
+
+---

docs/src/manual/group_conv.md

@@ -0,0 +1,19 @@
+# Positional Encoding Layers
+
+## ``E(n)``-equivariant Positional Encoding Layer
+
+It employs message-passing scheme and can be defined by following functions:
+
+- message function: ``y_{ij}^l = (x_i^l - x_j^l)\phi_x(m_{ij})``
+- aggregate: ``y_i^l = \frac{1}{M} \sum_{j \in \mathcal{N}(i)} y_{ij}^l``
+- update function: ``x_i^{l+1} = x_i^l + y_i^l``
+
+where ``x_i^l`` and ``x_j^l`` denote the positional feature for node ``i`` and ``j``, respectively, in ``l``-th layer, ``\phi_x`` is the neural network for positional encoding and ``m_{ij}`` is the edge feature for edge ``(i,j)``. ``y_{ij}^l`` and ``y_i^l`` represent the encoded positional feature and aggregated positional feature, respectively, and ``M`` denotes number of neighbors of node ``i``.
+
+```@docs
+EEquivGraphPE
+```
+
+Reference: [Satorras2021](@cite)
+
+---

docs/src/manual/positional.md

@@ -9,34 +9,20 @@ E(n)-equivariant graph neural network layer.
 - `out_dim`: the output of the layer will have dimension `out_dim` + (dimension of input vector - `in_dim`).
 - `pos_dim::Int`: dimension of positional encoding.
 - `edge_dim::Int`: dimension of edge feature.
-- `init`: neural network initialization function, should be compatible with `Flux.Dense`.
+- `init`: neural network initialization function.
 
 # Examples
 
 ```jldoctest
-julia> in_dim, int_dim, out_dim = 3,6,5
-(3, 5, 5)
+julia> in_dim, out_dim, pos_dim = 3, 5, 2
+(3, 5, 2)
 
-julia> egnn = EEquivGraphConv(in_dim, int_dim, out_dim)
-EEquivGraphConv{Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}(Dense(8 => 5), Dense(5 => 1), Dense(8 => 5), 3, 5, 5)
-
-julia> m_len = 2*in_dim + 2
-8
-
-julia> nn_edge = Flux.Dense(m_len, int_dim)
-Dense(8 => 5)       # 45 parameters
-
-julia> nn_x = Flux.Dense(int_dim, 1)
-Dense(5 => 1)       # 6 parameters
-
-julia> nn_h = Flux.Dense(in_dim + int_dim, out_dim)
-Dense(8 => 5)       # 45 parameters
-
-julia> egnn = EEquivGraphConv(in_dim, nn_edge, nn_x, nn_h)
-EEquivGraphConv{Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}(Dense(8 => 5), Dense(5 => 1), Dense(8 => 5), 3, 5, 5)
+julia> egnn = EEquivGraphConv(in_dim=>out_dim, pos_dim, in_dim)
+EEquivGraphConv(ϕ_edge=Dense(10 => 5), ϕ_x=Dense(5 => 2), ϕ_h=Dense(8 => 5))
 ```
-"""
 
+See also [`WithGraph`](@ref) for training layer with static graph and [`EEquivGraphPE`](@ref) for positional encoding.
+"""
 struct EEquivGraphConv{X,E,H}
     pe::X
     nn_edge::E

src/layers/group_conv.jl

@@ -7,6 +7,30 @@ abstract type AbstractPE end
 
 positional_encode(l::AbstractPE, args...) = throw(ErrorException("positional_encode function for $l is not implemented."))
 
+"""
+    EEquivGraphPE(in_dim=>out_dim; init=glorot_uniform, bias=true)
+
+E(n)-equivariant positional encoding layer.
+
+# Arguments
+
+- `in_dim::Int`: dimension of input positional feature.
+- `out_dim::Int`:  dimension of output positional feature.
+- `init`: neural network initialization function.
+- `bias::Bool`: dimension of edge feature.
+
+# Examples
+
+```jldoctest
+julia> in_dim_edge, out_dim = 2, 5
+(2, 5)
+
+julia> l = EEquivGraphPE(in_dim_edge=>out_dim)
+EEquivGraphPE(2 => 5)
+```
+
+See also [`EEquivGraphConv`](@ref).
+"""
 struct EEquivGraphPE{X} <: MessagePassing
     nn::X
 end