ndata, edata, gdata

Author: CarloLucibello

Project.toml

@@ -9,11 +9,13 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
+LearnBase = "7f8f8fb0-2700-5f03-b4bd-41f8cfc144b6"
 LightGraphs = "093fc24a-ae57-5d10-9952-331d41423f4d"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 NNlibCUDA = "a00861dc-f156-4864-bf3c-e6376f28a68d"
+PrettyPrinting = "54e16d92-306c-5ea0-a30b-337be88ac337"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

perf/perf.jl

@@ -11,7 +11,7 @@ function run_single_benchmark(N, c, D, CONV; gtype=:lg)
     data = erdos_renyi(N, c / (N-1), seed=17)
     X = randn(Float32, D, N)
 
-    g = GNNGraph(data; nf=X, graph_type=gtype)
+    g = GNNGraph(data; ndata=X, graph_type=gtype)
     g_gpu = g |> gpu    
 
     m = CONV(D => D)

src/GraphNeuralNetworks.jl

@@ -1,8 +1,5 @@
 module GraphNeuralNetworks
 
-using Core: apply_type
-using NNlib: similar
-using LinearAlgebra: similar, fill!
 using Statistics: mean
 using LinearAlgebra
 using SparseArrays
@@ -12,17 +9,17 @@ using CUDA
 using Flux
 using Flux: glorot_uniform, leakyrelu, GRUCell, @functor
 using MacroTools: @forward
+using LearnBase: getobs
 using NNlib, NNlibCUDA
 using ChainRulesCore
 import LightGraphs
-using LightGraphs: AbstractGraph, outneighbors, inneighbors, is_directed, ne, nv, 
-                  adjacency_matrix, degree
+using LightGraphs: AbstractGraph, outneighbors, inneighbors, adjacency_matrix, degree
 
 export
     # gnngraph
     GNNGraph,
     edge_index,
-    node_feature, edge_feature, global_feature,
+    node_features, edge_features, global_features,
     adjacency_list, normalized_laplacian, scaled_laplacian,
     add_self_loops, remove_self_loops,
     subgraph,
@@ -52,7 +49,6 @@ export
     topk_index
 
 
-
 
 include("gnngraph.jl")
 include("graph_conversions.jl")

src/gnngraph.jl

@@ -11,8 +11,8 @@ const ADJMAT_T = AbstractMatrix
 const SPARSE_T = AbstractSparseMatrix # subset of ADJMAT_T
 
 """
-    GNNGraph(data; [graph_type, nf, ef, gf, num_nodes, graph_indicator, dir])
-    GNNGraph(g::GNNGraph; [nf, ef, gf])
+    GNNGraph(data; [graph_type, ndata, edata, gdata, num_nodes, graph_indicator, dir])
+    GNNGraph(g::GNNGraph; [ndata, edata, gdata])
 
 A type representing a graph structure and storing also arrays 
 that contain features associated to nodes, edges, and the whole graph. 
@@ -50,10 +50,10 @@ from the LightGraphs' graph library can be used on it.
 - `dir`. The assumed edge direction when given adjacency matrix or adjacency list input data `g`. 
         Possible values are `:out` and `:in`. Default `:out`.
 - `num_nodes`. The number of nodes. If not specified, inferred from `g`. Default `nothing`.
-- `graph_indicator`. For batched graphs, a vector containeing the graph assigment of each node. Default `nothing`.  
-- `nf`: Node features. Either nothing, or an array whose last dimension has size num_nodes. Default `nothing`.
-- `ef`: Edge features. Either nothing, or an array whose last dimension has size num_edges. Default `nothing`.
-- `gf`: Global features. Default `nothing`. 
+- `graph_indicator`. For batched graphs, a vector containing the graph assigment of each node. Default `nothing`.  
+- `ndata`: Node features. A named tuple of arrays whose last dimension has size num_nodes.
+- `edata`: Edge features. A named tuple of arrays whose whose last dimension has size num_edges.
+- `gdata`: Global features. A named tuple of arrays whose has size num_graphs. 
 
 # Usage. 
 
@@ -77,7 +77,7 @@ g = GNNGraph(s, t)
 g = GNNGraph(erdos_renyi(100, 20))
 
 # Copy graph while also adding node features
-g = GNNGraph(g, nf=rand(100, 5))
+g = GNNGraph(g, ndata = (x = rand(100, g.num_nodes),))
 
 # Send to gpu
 g = g |> gpu
@@ -86,38 +86,28 @@ g = g |> gpu
 # Both source and target are vectors of length num_edges
 source, target = edge_index(g)
 ```
-
-See also [`graph`](@ref), [`edge_index`](@ref), [`node_feature`](@ref), [`edge_feature`](@ref), and [`global_feature`](@ref) 
 """
 struct GNNGraph{T<:Union{COO_T,ADJMAT_T}}
     graph::T
     num_nodes::Int
     num_edges::Int
     num_graphs::Int
     graph_indicator
-    nf
-    ef
-    gf
-    ## possible future property stores
-    # ndata::Dict{String, Any} # https://github.com/FluxML/Zygote.jl/issues/717        
-    # edata::Dict{String, Any}
-    # gdata::Dict{String, Any}
+    ndata::NamedTuple
+    edata::NamedTuple
+    gdata::NamedTuple
 end
 
 @functor GNNGraph
 
 function GNNGraph(data; 
                         num_nodes = nothing,
-                        num_graphs = 1,
                         graph_indicator = nothing, 
                         graph_type = :coo,
                         dir = :out,
-                        nf = nothing, 
-                        ef = nothing, 
-                        gf = nothing,
-                        # ndata = Dict{String, Any}(), 
-                        # edata = Dict{String, Any}(),
-                        # gdata = Dict{String, Any}()
+                        ndata = (;), 
+                        edata = (;), 
+                        gdata = (;),
                         )
 
     @assert graph_type ∈ [:coo, :dense, :sparse] "Invalid graph_type $graph_type requested"
@@ -133,18 +123,20 @@ function GNNGraph(data;
 
     num_graphs = !isnothing(graph_indicator) ? maximum(graph_indicator) : 1
 
-    ## Possible future implementation of feature maps. 
-    ## Currently this doesn't play well with zygote due to 
-    ## https://github.com/FluxML/Zygote.jl/issues/717    
-    # ndata["x"] = nf
-    # edata["e"] = ef
-    # gdata["g"] = gf
+    ndata = normalize_graphdata(ndata, :X)
+    edata = normalize_graphdata(edata, :E)
+    gdata = normalize_graphdata(gdata, :U)
 
-    GNNGraph(g, num_nodes, num_edges, 
-            num_graphs, graph_indicator,
-            nf, ef, gf)
+    GNNGraph(g, 
+            num_nodes, num_edges, num_graphs, 
+            graph_indicator,
+            ndata, edata, gdata)
 end
 
+normalize_graphdata(data::NamedTuple, s) = data
+normalize_graphdata(data::Nothing, s) = NamedTuple()
+normalize_graphdata(data, s) = NamedTuple{(s,)}((data,)) 
+
 # COO convenience constructors
 GNNGraph(s::AbstractVector, t::AbstractVector, v = nothing; kws...) = GNNGraph((s, t, v); kws...)
 GNNGraph((s, t)::NTuple{2}; kws...) = GNNGraph((s, t, nothing); kws...)
@@ -154,14 +146,19 @@ GNNGraph((s, t)::NTuple{2}; kws...) = GNNGraph((s, t, nothing); kws...)
 function GNNGraph(g::AbstractGraph; kws...)
     s = LightGraphs.src.(LightGraphs.edges(g))
     t = LightGraphs.dst.(LightGraphs.edges(g))
-    GNNGraph((s, t); num_nodes = nv(g), kws...)
+    GNNGraph((s, t); num_nodes = LightGraphs.nv(g), kws...)
 end
 
-function GNNGraph(g::GNNGraph; 
-                nf=node_feature(g), ef=edge_feature(g), gf=global_feature(g))
-                # ndata=copy(g.ndata), edata=copy(g.edata), gdata=copy(g.gdata), # copy keeps the refs to old data 
+function GNNGraph(g::GNNGraph; ndata=g.ndata, edata=g.edata, gdata=g.gdata)
+
+    ndata = normalize_graphdata(ndata, :X)
+    edata = normalize_graphdata(edata, :E)
+    gdata = normalize_graphdata(gdata, :U)
 
-    GNNGraph(g.graph, g.num_nodes, g.num_edges, g.num_graphs, g.graph_indicator, nf, ef, gf) #   ndata, edata, gdata, 
+    GNNGraph(g.graph, 
+            g.num_nodes, g.num_edges, g.num_graphs, 
+            g.graph_indicator, 
+            ndata, edata, gdata) 
 end
 
 
@@ -266,44 +263,6 @@ function LightGraphs.degree(g::GNNGraph{<:ADJMAT_T}, T=Int; dir=:out)
     return dir == :out ? vec(sum(A, dims=2)) : vec(sum(A, dims=1))
 end
 
-# node_feature(g::GNNGraph) = g.ndata["x"]
-# edge_feature(g::GNNGraph) = g.edata["e"]
-# global_feature(g::GNNGraph) = g.gdata["g"]
-
-
-"""
-    node_feature(g::GNNGraph)
-
-Return the node features of `g`.
-"""
-node_feature(g::GNNGraph) = g.nf
-
-"""
-    edge_feature(g::GNNGraph)
-
-Return the edge features of `g`.
-"""
-edge_feature(g::GNNGraph) = g.ef
-
-"""
-    global_feature(g::GNNGraph)
-
-Return the global features of `g`.
-"""
-global_feature(g::GNNGraph) = g.gf
-
-# function Base.getproperty(g::GNNGraph, sym::Symbol)
-#     if sym === :nf
-#         return g.ndata["x"]
-#     elseif sym === :ef
-#         return g.edata["e"]
-#     elseif sym === :gf
-#         return g.gdata["g"]
-#     else # fallback to getfield
-#         return getfield(g, sym)
-#     end
-# end
-
 function LightGraphs.laplacian_matrix(g::GNNGraph, T::DataType=Int; dir::Symbol=:out)
     A = adjacency_matrix(g, T; dir=dir)
     D = Diagonal(vec(sum(A; dims=2)))
@@ -376,41 +335,44 @@ self-loops will obtain a second self-loop.
 """
 function add_self_loops(g::GNNGraph{<:COO_T})
     s, t = edge_index(g)
-    @assert edge_feature(g) === nothing
+    @assert g.edata === (;)
     @assert edge_weight(g) === nothing
     n = g.num_nodes
     nodes = convert(typeof(s), [1:n;])
     s = [s; nodes]
     t = [t; nodes]
 
-    GNNGraph((s, t, nothing), g.num_nodes, length(s),
-        g.num_graphs, g.graph_indicator,
-        node_feature(g), edge_feature(g), global_feature(g))
+    GNNGraph((s, t, nothing), 
+        g.num_nodes, length(s), g.num_graphs, 
+        g.graph_indicator,
+        g.ndata, g.edata, g.gdata)
 end
 
-function add_self_loops(g::GNNGraph{<:ADJMAT_T}; add_to_existing=true)
-    A = graph(g)
-    @assert edge_feature(g) === nothing
+function add_self_loops(g::GNNGraph{<:ADJMAT_T})
+    A = adjaceny_matrix(g)
+    @assert g.edata === (;)
     A += I
     num_edges =  g.num_edges + g.num_nodes
-    GNNGraph(A, g.num_nodes, num_edges,
-        g.num_graphs, g.graph_indicator,
-        node_feature(g), edge_feature(g), global_feature(g))
+    GNNGraph(A, 
+            g.num_nodes, num_edges, g.num_graphs, 
+            g.graph_indicator,
+            g.ndata, g.edata, g.gdata)
 end
 
 function remove_self_loops(g::GNNGraph{<:COO_T})
     s, t = edge_index(g)
     # TODO remove these constraints
-    @assert edge_feature(g) === nothing
+    @assert g.edata === (;)
     @assert edge_weight(g) === nothing
 
     mask_old_loops = s .!= t
     s = s[mask_old_loops]
     t = t[mask_old_loops]
 
-    GNNGraph((s, t, nothing), g.num_nodes, length(s), 
-        g.num_graphs, g.graph_indicator,
-        node_feature(g), edge_feature(g), global_feature(g))
+    GNNGraph((s, t, nothing), 
+            g.num_nodes, length(s), g.num_graphs, 
+            g.graph_indicator,
+            g.ndata, g.edata, g.gdata)
 end
 
 function _catgraphs(g1::GNNGraph{<:COO_T}, g2::GNNGraph{<:COO_T})
@@ -425,14 +387,12 @@ function _catgraphs(g1::GNNGraph{<:COO_T}, g2::GNNGraph{<:COO_T})
     ind2 = isnothing(g2.graph_indicator) ? fill!(similar(s2, Int, nv2), 1) : g2.graph_indicator 
     graph_indicator = vcat(ind1, g1.num_graphs .+ ind2)
 
-    GNNGraph(
-        (s, t, w),
-        nv1 + nv2, g1.num_edges + g2.num_edges, 
-        g1.num_graphs + g2.num_graphs, graph_indicator,
-        cat_features(node_feature(g1), node_feature(g2)),
-        cat_features(edge_feature(g1), edge_feature(g2)),
-        cat_features(global_feature(g1), global_feature(g2)),
-    )
+    GNNGraph((s, t, w),
+            nv1 + nv2, g1.num_edges + g2.num_edges, g1.num_graphs + g2.num_graphs, 
+            graph_indicator,
+            cat_features(g1.ndata, g2.ndata),
+            cat_features(g1.edata, g2.edata),
+            cat_features(g1.gdata, g2.gdata))
 end
 
 ### Cat public interfaces #############
@@ -490,9 +450,9 @@ function subgraph(g::GNNGraph{<:COO_T}, i::AbstractVector)
     s = [nodemap[i] for i in s[edge_mask]]
     t = [nodemap[i] for i in t[edge_mask]]
     w = isnothing(w) ? nothing : w[edge_mask]
-    nf = isnothing(g.nf) ? nothing : g.nf[:,node_mask]
-    ef = isnothing(g.ef) ? nothing : g.ef[:,edge_mask]
-    gf = isnothing(g.gf) ? nothing : g.gf[:,i]
+    ndata = getobs(g.ndata, node_mask)
+    edata = getobs(g.ndata, edge_mask)
+    gdata = getobs(g.gdata, i)
 
     num_nodes = length(graph_indicator)
     num_edges = length(s)
@@ -501,10 +461,43 @@ function subgraph(g::GNNGraph{<:COO_T}, i::AbstractVector)
     gnew = GNNGraph((s,t,w), 
                 num_nodes, num_edges, num_graphs,
                 graph_indicator,
-                nf, ef, gf)
+                ndata, edata, gdata)
     return gnew, nodes
 end
 
+### TO DEPRECATE ?? ###
+function node_features(g::GNNGraph)
+    if isempty(g.ndata)
+        return nothing
+    elseif length(g.ndata) > 1
+        @error "Multiple feature arrays, access directly with g.ndata.X"
+    else
+        return g.ndata[1]
+    end
+end
+
+function edge_features(g::GNNGraph)
+    if isempty(g.edata)
+        return nothing
+    elseif length(g.edata) > 1
+        @error "Multiple feature arrays, access directly with g.edata.E"
+    else
+        return g.edata[1]
+    end
+end
+
+function global_features(g::GNNGraph)
+    if isempty(g.gdata)
+        return nothing
+    elseif length(g.gdata) > 1
+        @error "Multiple feature arrays, access directly with g.gdata.U"
+    else
+        return g.gdata[1]
+    end
+end
+#########
+
+
 @non_differentiable normalized_laplacian(x...)
 @non_differentiable normalized_adjacency(x...)
 @non_differentiable scaled_laplacian(x...)

src/layers/basic.jl

@@ -7,6 +7,9 @@ See also [`GNNChain`](@ref).
 """
 abstract type GNNLayer end
 
+#TODO extend to store also edge and global features
+(l::GNNLayer)(g::GNNGraph) = GNNGraph(g, ndata=l(g, node_features(g)))
+
 """
     GNNChain(layers...)
     GNNChain(name = layer, ...)
@@ -39,7 +42,7 @@ julia> m(g, x)
  -0.0134364  -0.0120716  -0.0172505
 ```
 """
-struct GNNChain{T}
+struct GNNChain{T} <: GNNLayer
   layers::T
 
   GNNChain(xs...) = new{typeof(xs)}(xs)