Merge pull request #220 from emsal0/CGConv

Implement CGConv layer.

Author: yuehhua

docs/src/manual/conv.md

@@ -120,3 +120,17 @@ where ``f_{\Theta}`` denotes a neural network parametrized by ``\Theta``, *i.e.*
 GINConv
 ```
 Reference: [How Powerful are Graph Neural Networks?](https://arxiv.org/pdf/1810.00826.pdf)
+
+## Crystal Graph Convolutional Network
+
+```math
+\textbf{x}_i' = \textbf{x}_i + \sum_{j \in \mathcal{N}(i)} \sigma\left( \textbf{z}_{i,j} \textbf{W}_f + \textbf{b}_f \right) \odot \text{softplus}\left(\textbf{z}_{i,j} \textbf{W}_s + \textbf{b}_s \right)
+```
+
+where ``\textbf{z}_{i,j} = [\textbf{x}_i, \textbf{x}_j}, \textbf{e}_{i,j}]`` denotes the concatenation of node features, neighboring node features, and edge features. The operation ``\odot`` represents elementwise multiplication, and ``\sigma`` denotes the sigmoid function.
+
+```@docs
+CGConv
+```
+
+Reference: [Crystal Graph Convolutional Neural Networks for an Accurate and Interpretable Prediction of Material Properties](https://arxiv.org/pdf/1710.10324.pdf)

src/GeometricFlux.jl

@@ -29,6 +29,7 @@ export
     GatedGraphConv,
     EdgeConv,
     GINConv,
+    CGConv,
 
     # layer/pool
     GlobalPool,

src/layers/conv.jl

@@ -436,3 +436,77 @@ end
 
 (l::GINConv)(x::AbstractMatrix) = l(l.fg, x)
 (l::GINConv)(fg::FeaturedGraph) = FeaturedGraph(fg.graph, nf = l(fg, node_feature(fg)))
+
+
+"""
+    CGConv([fg,] (node_dim, edge_dim), out, init)
+
+Crystal Graph Convolutional network. Uses both node and edge features.
+
+# Arguments
+
+- `fg`: Optional [`FeaturedGraph`] argument(@ref)
+- `node_dim`: Dimensionality of the input node features. Also is necessarily the output dimensionality.
+- `edge_dim`: Dimensionality of the input edge features.
+- `out`: Dimensionality of the output features.
+- `init`: Initialization algorithm for each of the weight matrices
+- `bias`: Whether or not to learn an additive bias parameter.
+
+# Usage
+
+You can call `CGConv` in several different ways:
+                                    
+- Pass a FeaturedGraph: `CGConv(fg)`, returns `FeaturedGraph` 
+- Pass both node and edge features: `CGConv(X, E)` 
+- Pass one matrix, which can either be node features or edge features: `CGConv(M; edge)`:
+    `edge` is default false, meaning that `M` denotes node features.
+"""
+struct CGConv{V <: AbstractFeaturedGraph, T,
+              A <: AbstractMatrix{T}, B} <: MessagePassing
+    fg::V
+    Wf::A
+    Ws::A
+    bf::B
+    bs::B
+end
+
+@functor CGConv
+
+function CGConv(fg::AbstractFeaturedGraph, dims::NTuple{2,Int};
+                init=glorot_uniform, bias=true)
+    node_dim, edge_dim = dims
+    Wf = init(node_dim, 2*node_dim + edge_dim)
+    Ws = init(node_dim, 2*node_dim + edge_dim)
+    bf = Flux.create_bias(Wf, bias, node_dim)
+    bs = Flux.create_bias(Ws, bias, node_dim)
+    CGConv(fg, Wf, Ws, bf, bs)
+end
+
+function CGConv(dims::NTuple{2,Int}; init=glorot_uniform, bias=true)
+    CGConv(NullGraph(), dims; init=init, bias=bias)
+end
+
+message(c::CGConv,
+        x_i::AbstractVector, x_j::AbstractVector, e::AbstractVector) = begin
+    z = vcat(x_i, x_j, e)
+    σ.(c.Wf * z + c.bf) .* softplus.(c.Ws * z + c.bs)
+end
+update(c::CGConv, m::AbstractVector, x) = x + m
+
+function (c::CGConv)(fg::FeaturedGraph, X::AbstractMatrix, E::AbstractMatrix)
+    check_num_nodes(fg, X)
+    check_num_edges(fg, E)
+    _, Y = propagate(c, adjacency_list(fg), E, X, +)
+    Y
+end
+
+(l::CGConv)(fg::FeaturedGraph) = FeaturedGraph(fg, nf=l(fg, node_feature(fg),
+                                                        edge_feature(fg)),
+                                               ef=edge_feature(fg))
+(l::CGConv)(M::AbstractMatrix; as_edge=false) =
+    if as_edge
+        l(l.fg, node_feature(l.fg), M)
+    else
+        l(l.fg, M, edge_feature(l.fg))
+    end
+(l::CGConv)(X::AbstractMatrix, E::AbstractMatrix) = l(l.fg, X, E)

src/utils.jl

@@ -77,3 +77,7 @@ end
 function check_num_nodes(fg::FeaturedGraph, x::AbstractArray)
     @assert nv(fg) == size(x, ndims(x))    
 end
+
+function check_num_edges(fg::FeaturedGraph, e::AbstractArray)
+    @assert ne(fg) == size(e, ndims(e))
+end

test/layers/conv.jl

@@ -1,4 +1,5 @@
 in_channel = 3
+in_channel_edge = 1
 out_channel = 5
 N = 4
 T = Float32
@@ -379,4 +380,26 @@ fg_single_vertex = FeaturedGraph(adj_single_vertex)
             @test !in(:eps, Flux.trainable(gc))
         end
     end
+
+    @testset "CGConv" begin
+        fg = FeaturedGraph(adj)
+        X = rand(Float32, in_channel, N)
+        E = rand(Float32, in_channel_edge, ne(fg))
+        Xt = transpose(rand(Float32, N, in_channel))
+        @testset "layer with graph" begin
+            cgc = CGConv(FeaturedGraph(adj),
+                         (in_channel, in_channel_edge))
+            @test size(cgc.Wf) == (in_channel, 2 * in_channel + in_channel_edge)
+            @test size(cgc.Ws) == (in_channel, 2 * in_channel + in_channel_edge)
+            @test size(cgc.bf) == (in_channel,)
+            @test size(cgc.bs) == (in_channel,)
+
+            Y = cgc(X, E)
+            @test size(Y) == (in_channel, N)
+
+            Yg = cgc(FeaturedGraph(adj, nf=X, ef=E))
+            @test size(node_feature(Yg)) == (in_channel, N)
+            @test edge_feature(Yg) == E
+        end
+    end
 end