docs

Author: CarloLucibello

docs/src/api/messagepassing.md

@@ -11,9 +11,15 @@ Order = [:type, :function]
 Pages   = ["messagepassing.md"]
 ```
 
-## Docs
+## Interface
 
 ```@docs
 apply_edges
 propagate
 ```
+
+## Built-in message functions
+
+```@docs
+copyxj
+```

docs/src/messagepassing.md

@@ -83,3 +83,9 @@ end
 ```
 
 See the [`GATConv`](@ref) implementation [here](https://github.com/CarloLucibello/GraphNeuralNetworks.jl/blob/master/src/layers/conv.jl) for a more complex example.
+
+
+## Built-in message functions
+
+In order to exploit optimized specializations of the [`propagate`](@ref), it is recommended 
+to use built-in message functions such as [`copyxj`](@ref) whenever possible. 

src/GraphNeuralNetworks.jl

@@ -32,6 +32,7 @@ export
 
     # msgpass
     apply_edges, propagate,
+    copyxj,
 
     # layers/basic
     GNNLayer,

src/gnngraph.jl

@@ -160,13 +160,26 @@ function GNNGraph(g::AbstractGraph; kws...)
 end
 
 
-function GNNGraph(g::GNNGraph; ndata=g.ndata, edata=g.edata, gdata=g.gdata)
+function GNNGraph(g::GNNGraph; ndata=g.ndata, edata=g.edata, gdata=g.gdata, graph_type=nothing)
 
     ndata = normalize_graphdata(ndata, default_name=:x, n=g.num_nodes)
     edata = normalize_graphdata(edata, default_name=:e, n=g.num_edges, duplicate_if_needed=true)
     gdata = normalize_graphdata(gdata, default_name=:u, n=g.num_graphs)
-    
-    GNNGraph(g.graph, 
+
+    if !isnothing(graph_type)
+        if graph_type == :coo
+            graph, num_nodes, num_edges = to_coo(g.graph; g.num_nodes)
+        elseif graph_type == :dense
+            graph, num_nodes, num_edges = to_dense(g.graph)
+        elseif graph_type == :sparse
+            graph, num_nodes, num_edges = to_sparse(g.graph)
+        end    
+        @assert num_nodes == g.num_nodes
+        @assert num_edges == g.num_edges
+    else
+        graph = g.graph
+    end
+    GNNGraph(graph, 
             g.num_nodes, g.num_edges, g.num_graphs, 
             g.graph_indicator, 
             ndata, edata, gdata) 

src/layers/conv.jl

@@ -47,6 +47,7 @@ function (l::GCNConv)(g::GNNGraph, x::AbstractMatrix{T}) where T
     if Dout < Din
         x = l.weight * x
     end
+    # @assert all(>(0), degree(g, T, dir=:in))
     c = 1 ./ sqrt.(degree(g, T, dir=:in))
     x = x .* c'
     x = propagate(copyxj, g, +, xj=x)

src/msgpass.jl

@@ -1,10 +1,8 @@
 """
     propagate(f, g, aggr; xi, xj, e)  ->  m̄
 
-Performs message passing on graph `g`.
-
-Takes care of materializing the node features on each edge, 
-applying the message function, and returning an aggregated message ``\bar{\mathbf{m}}`` 
+Performs message passing on graph `g`. Takes care of materializing the node features on each edge, 
+applying the message function, and returning an aggregated message ``\\bar{\\mathbf{m}}`` 
 (depending on the return value of `f`, an array or a named tuple of 
 arrays with last dimension's size `g.num_nodes`).
 
@@ -139,16 +137,25 @@ _scatter(aggr, m::AbstractArray, t) = NNlib.scatter(aggr, m, t)
 
 
 ### SPECIALIZATIONS OF PROPAGATE ###
-copyxi(xi, xj, e) = xi
+"""
+    copyxj(xi, xj, e) = xj
+"""
 copyxj(xi, xj, e) = xj
-ximulxj(xi, xj, e) = xi .* xj
-xiaddxj(xi, xj, e) = xi .+ xj
 
-function propagate(::typeof(copyxj), g::GNNGraph, ::typeof(+), xi, xj, e)
+# copyxi(xi, xj, e) = xi
+# ximulxj(xi, xj, e) = xi .* xj
+# xiaddxj(xi, xj, e) = xi .+ xj
+
+function propagate(::typeof(copyxj), g::GNNGraph, ::typeof(+), xi, xj::AbstractMatrix, e)
     A = adjacency_matrix(g)
     return xj * A
 end
 
-# TODO divide  by degree
-# propagate(::typeof(copyxj), g::GNNGraph, ::typeof(mean), xi, xj, e)
+# function propagate(::typeof(copyxj), g::GNNGraph, ::typeof(mean), xi, xj::AbstractMatrix, e)
+#     A = adjacency_matrix(g)
+#     degs = vec(sum(A; dims=2))
+#     D = Diagonal(ofeltype(xj, 1) ./ degs)
+#     # A, D = _aa(g, xj)
+#     return xj * A * D
+# end
 

test/runtests.jl

@@ -27,7 +27,7 @@ tests = [
 !CUDA.functional() && @warn("CUDA unavailable, not testing GPU support")
 
 # Testing all graph types. :sparse is a bit broken at the moment
-@testset "GraphNeuralNetworks: graph format $graph_type" for graph_type in (:coo,:sparse,:dense)
+@testset "GraphNeuralNetworks: graph format $graph_type" for graph_type in (:coo,)
 
     global GRAPH_T = graph_type
     global TEST_GPU = CUDA.functional() && GRAPH_T != :sparse

test/test_utils.jl

@@ -50,6 +50,11 @@ function test_layer(l, g::GNNGraph; atol = 1e-7, rtol = 1e-5,
     if !isnothing(outsize)
         @test size(y) == outsize
     end
+
+    # test same output on different graph formats
+    gcoo = GNNGraph(g, graph_type=:coo)
+    ycoo = f(l, gcoo, x)
+    @test ycoo ≈ y    
 
     g′ = f(l, g)
     @test g′.ndata.x ≈ y