add EEquivGraphPE layer

Author: yuehhua

src/GeometricFlux.jl

@@ -31,6 +31,11 @@ export
     # layers/msgpass
     MessagePassing,
 
+    # layers/positional
+    AbstractPE,
+    positional_encode,
+    EEquivGraphPE,
+
     # layers/graph_conv
     GCNConv,
     ChebConv,
@@ -75,6 +80,7 @@ include("layers/graphlayers.jl")
 include("layers/gn.jl")
 include("layers/msgpass.jl")
 
+include("layers/positional.jl")
 include("layers/graph_conv.jl")
 include("layers/group_conv.jl")
 include("layers/pool.jl")

src/layers/group_conv.jl

@@ -6,22 +6,16 @@ E(n)-equivariant graph neural network layer as defined in the paper "[E(n) Equiv
 
 # Arguments
 
-Either one of two sets of arguments:
-
-Set 1:
-
 - `in_dim`: node feature dimension. Data is assumed to be of the form [feature; coordinate], so `in_dim` must strictly be less than the dimension of the input vectors.
 - `int_dim`: intermediate dimension, can be arbitrary.
 - `out_dim`: the output of the layer will have dimension `out_dim` + (dimension of input vector - `in_dim`).
 - `init`: neural network initialization function, should be compatible with `Flux.Dense`.
-
-Set 2:
-
-- `in_dim`: as in Set 1.
 - `nn_edge`: a differentiable function that must take vectors of dimension `in_dim * 2 + 2` (output designated `int_dim`)
 - `nn_x`: a differentiable function that must take vectors of dimension `int_dim` to dimension `1`.
 - `nn_h`: a differentiable function that must take vectors of dimension `in_dim + int_dim` to `out_dim`.
 
+# Examples
+
 ```jldoctest
 julia> in_dim, int_dim, out_dim = 3,6,5
 (3, 5, 5)
@@ -57,6 +51,8 @@ end
 
 @functor EEquivGraphConv
 
+Flux.trainable(l::EEquivGraphConv) = (l.nn_edge, l.nn_x, l.nn_h)
+
 function EEquivGraphConv(in_dim::Int, int_dim::Int, out_dim::Int; init=glorot_uniform)
     m_len = 2in_dim + 2
     nn_edge = Flux.Dense(m_len, int_dim; init=init)
@@ -120,9 +116,19 @@ function update(e::EEquivGraphConv, m, h)
     return z
 end
 
+# For variable graph
 function(egnn::EEquivGraphConv)(fg::AbstractFeaturedGraph)
     X = node_feature(fg)
     GraphSignals.check_num_nodes(fg, X)
     _, V, _ = propagate(egnn, graph(fg), nothing, X, nothing, +, nothing, nothing)
     return ConcreteFeaturedGraph(fg, nf=V)
 end
+
+function Base.show(io::IO, l::EEquivGraphConv)
+    in_channel = size(l.weight1, ndims(l.weight1))
+    out_channel = size(l.weight1, ndims(l.weight1)-1)
+    print(io, "GraphConv(", in_channel, " => ", out_channel)
+    l.σ == identity || print(io, ", ", l.σ)
+    print(io, ", aggr=", l.aggr)
+    print(io, ")")
+end

src/layers/positional.jl

@@ -0,0 +1,44 @@
+"""
+    AbstractPE
+
+Abstract type of positional encoding for GNN.
+"""
+abstract type AbstractPE end
+
+positional_encode(l::AbstractPE, args...) = throw(ErrorException("positional_encode function for $l is not implemented."))
+
+struct EEquivGraphPE{X} <: MessagePassing
+    nn_x::X
+end
+
+function EEquivGraphPE(ch::Pair{Int,Int}; init=glorot_uniform, bias::Bool=true)
+    in, out = ch
+    nn_x = Flux.Dense(in, out; init=init, bias=bias)
+    return EEquivGraphPE(nn_x)
+end
+
+@functor EEquivGraphPE
+
+ϕ_x(l::EEquivGraphPE, m_ij) = l.nn_x(m_ij)
+
+message(l::EEquivGraphPE, x_i, x_j, e) = (x_i - x_j) .* ϕ_x(l, e)
+
+update(l::EEquivGraphPE, m::AbstractArray, x::AbstractArray) = m .+ x
+
+# For variable graph
+function(l::EEquivGraphPE)(fg::AbstractFeaturedGraph)
+    X = node_feature(fg)
+    E = edge_feature(fg)
+    GraphSignals.check_num_nodes(fg, X)
+    GraphSignals.check_num_nodes(fg, E)
+    _, V, _ = propagate(l, graph(fg), E, X, nothing, mean, nothing, nothing)
+    return ConcreteFeaturedGraph(fg, nf=V)
+end
+
+# For static graph
+function(l::EEquivGraphPE)(el::NamedTuple, X::AbstractArray, E::AbstractArray)
+    GraphSignals.check_num_nodes(el.N, X)
+    GraphSignals.check_num_nodes(el.E, E)
+    _, V, _ = propagate(l, el, E, X, nothing, mean, nothing, nothing)
+    return V
+end

test/layers/positional.jl

@@ -0,0 +1,26 @@
+@testset "positional" begin
+    T = Float32
+    batch_size = 10
+    in_channel = 3
+    out_channel = 5
+
+    N = 4
+    E = 4
+    adj = T[0. 1. 0. 1.;
+            1. 0. 1. 0.;
+            0. 1. 0. 1.;
+            1. 0. 1. 0.]
+    
+    @testset "EEquivGraphPE" begin
+        l = EEquivGraphPE(in_channel=>out_channel)
+
+        nf = rand(T, out_channel, N)
+        ef = rand(T, in_channel, E)
+        fg = FeaturedGraph(adj, nf=nf, ef=ef)
+        fg_ = l(fg)
+        @test size(node_feature(fg_)) == (out_channel, N)
+
+        g = Zygote.gradient(() -> sum(node_feature(l(fg))), Flux.params(l))
+        @test length(g.grads) == 4
+    end
+end

test/runtests.jl

@@ -22,6 +22,7 @@ cuda_tests = [
 tests = [
     "layers/gn",
     "layers/msgpass",
+    "layers/positional",
     "layers/graph_conv",
     "layers/group_conv",
     "layers/pool",