add sample_neighbors (#93)

* add GNNGraph(num_nodes)

* sample_neighbors

* fix dense constructor

Author: CarloLucibello

docs/src/api/gnngraph.md

@@ -67,3 +67,11 @@ Private = false
 ```@docs
 Graphs.intersect
 ```
+
+## Sampling 
+
+```@autodocs
+Modules = [GraphNeuralNetworks.GNNGraphs]
+Pages   = ["sampling.jl"]
+Private = false
+```

src/GNNGraphs/GNNGraphs.jl

@@ -57,6 +57,9 @@ include("generate.jl")
 export rand_graph, 
        knn_graph
 
+include("sampling.jl")
+export sample_neighbors
+
 include("operators.jl")
 # Base.intersect
 

src/GNNGraphs/gnngraph.jl

@@ -117,25 +117,25 @@ end
 
 @functor GNNGraph
 
-function GNNGraph(data; 
+function GNNGraph(data::D; 
                         num_nodes = nothing,
                         graph_indicator = nothing, 
                         graph_type = :coo,
                         dir = :out,
                         ndata = (;), 
                         edata = (;), 
                         gdata = (;),
-                        )
+                        ) where D <: Union{COO_T, ADJMAT_T, ADJLIST_T}
 
     @assert graph_type ∈ [:coo, :dense, :sparse] "Invalid graph_type $graph_type requested"
     @assert dir ∈ [:in, :out]
 
     if graph_type == :coo
         graph, num_nodes, num_edges = to_coo(data; num_nodes, dir)
     elseif graph_type == :dense
-        graph, num_nodes, num_edges = to_dense(data; dir)
+        graph, num_nodes, num_edges = to_dense(data; num_nodes, dir)
     elseif graph_type == :sparse
-        graph, num_nodes, num_edges = to_sparse(data; dir)
+        graph, num_nodes, num_edges = to_sparse(data; num_nodes, dir)
     end
 
     num_graphs = !isnothing(graph_indicator) ? maximum(graph_indicator) : 1
@@ -150,9 +150,9 @@ function GNNGraph(data;
             ndata, edata, gdata)
 end
 
-function GNNGraph(n::T; graph_type=:coo, kws...) where {T<:Integer}
+function (::Type{<:GNNGraph})(num_nodes::T; kws...) where {T<:Integer}
     s, t = T[], T[] 
-    return GNNGraph(s, t; graph_type, num_nodes=n, kws...)
+    return GNNGraph(s, t; num_nodes, kws...)
 end
 
 # COO convenience constructors
@@ -182,10 +182,10 @@ function GNNGraph(g::GNNGraph; ndata=g.ndata, edata=g.edata, gdata=g.gdata, grap
         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)
+            graph, num_nodes, num_edges = to_dense(g.graph; g.num_nodes)
         elseif graph_type == :sparse
-            graph, num_nodes, num_edges = to_sparse(g.graph)
-        end    
+            graph, num_nodes, num_edges = to_sparse(g.graph; g.num_nodes)
+        end
         @assert num_nodes == g.num_nodes
         @assert num_edges == g.num_edges
     else

src/GNNGraphs/query.jl

@@ -51,6 +51,15 @@ Graphs.ne(g::GNNGraph) = g.num_edges
 Graphs.has_vertex(g::GNNGraph, i::Int) = 1 <= i <= g.num_nodes
 Graphs.vertices(g::GNNGraph) = 1:g.num_nodes
 
+function Graphs.neighbors(g::GNNGraph, i; dir=:out)
+    @assert dir ∈ (:in, :out)
+    if dir == :out
+        outneighbors(g, i)
+    else
+        inneighbors(g, i)
+    end
+end
+
 function Graphs.outneighbors(g::GNNGraph{<:COO_T}, i::Integer)
     s, t = edge_index(g)
     return t[s .== i]
@@ -76,6 +85,7 @@ Graphs.is_directed(::Type{<:GNNGraph}) = true
 
 """
     adjacency_list(g; dir=:out)
+    adjacency_list(g, nodes; dir=:out)
 
 Return the adjacency list representation (a vector of vectors)
 of the graph `g`.
@@ -84,13 +94,44 @@ Calling `a` the adjacency list, if `dir=:out` than
 `a[i]` will contain the neighbors of node `i` through
 outgoing edges. If `dir=:in`, it will contain neighbors from
 incoming edges instead.
+
+If `nodes` is given, return the neighborhood of the nodes in `nodes` only.
 """
-function adjacency_list(g::GNNGraph; dir=:out)
+function adjacency_list(g::GNNGraph, nodes; dir=:out, with_eid=false)
     @assert dir ∈ [:out, :in]
-    fneighs = dir == :out ? outneighbors : inneighbors
-    return [fneighs(g, i) for i in 1:g.num_nodes]
+    s, t = edge_index(g)
+    if dir == :in
+        s, t = t, s
+    end
+    T = eltype(s)
+    idict = 0
+    dmap = Dict(n => (idict += 1) for n in nodes)
+    adjlist = [T[] for _=1:length(dmap)]
+    eidlist = [T[] for _=1:length(dmap)]
+    for (eid, (i, j)) in enumerate(zip(s, t))
+        inew = get(dmap, i, 0)
+        inew == 0 && continue
+        push!(adjlist[inew], j)
+        push!(eidlist[inew], eid)
+    end
+    if with_eid
+        return adjlist, eidlist
+    else
+        return adjlist
+    end
 end
 
+# function adjacency_list(g::GNNGraph, nodes; dir=:out)
+#     @assert dir ∈ [:out, :in]
+#     fneighs = dir == :out ? outneighbors : inneighbors
+#     return [fneighs(g, i) for i in nodes]
+# end
+
+
+
+adjacency_list(g::GNNGraph; dir=:out) = adjacency_list(g, 1:g.num_nodes; dir)
+
+
 function Graphs.adjacency_matrix(g::GNNGraph{<:COO_T}, T::DataType=eltype(g); dir=:out)
     if g.graph[1] isa CuVector
         # TODO revisit after https://github.com/JuliaGPU/CUDA.jl/pull/1152
@@ -140,7 +181,8 @@ Return a vector containing the degrees of the nodes in `g`.
                 In alternative, you can also pass a vector of weights to be used
                 instead of the graph's own weights.
 """
-function Graphs.degree(g::GNNGraph{<:COO_T}, T=nothing; dir=:out, edge_weight=true)
+function Graphs.degree(g::GNNGraph{<:COO_T}, T::TT=nothing; dir=:out, edge_weight=true) where 
+        TT<:Union{Nothing, Type{<:Number}}
     s, t = edge_index(g)
 
     edge_weight = _get_edge_weight(g, edge_weight)
@@ -157,7 +199,11 @@ function Graphs.degree(g::GNNGraph{<:COO_T}, T=nothing; dir=:out, edge_weight=tr
     return degs 
 end
 
-function Graphs.degree(g::GNNGraph{<:ADJMAT_T}, T=nothing; dir=:out, edge_weight=true)
+# TODO:: Make efficient
+Graphs.degree(g::GNNGraph, i::Union{Int, AbstractVector}; dir=:out) = degree(g; dir)[i]
+
+function Graphs.degree(g::GNNGraph{<:ADJMAT_T}, T::TT=nothing; dir=:out, edge_weight=true) where TT
+        TT<:Union{Nothing, Type{<:Number}}
     # edge_weight=true or edge_weight=nothing act the same here
     @assert !(edge_weight isa AbstractArray) "passing the edge weights is not support by adjacency matrix representations" 
     @assert dir ∈ (:in, :out, :both)

src/GNNGraphs/sampling.jl

@@ -0,0 +1,117 @@
+"""
+    sample_neighbors(g, nodes, K=-1; dir=:in, replace=false, dropnodes=false)
+
+Sample neighboring edges of the given nodes and return the induced subgraph.
+For each node, a number of inbound (or outbound when `dir = :out``) edges will be randomly chosen. 
+If `dropnodes=false`, the graph returned will then contain all the nodes in the original graph, 
+but only the sampled edges.
+
+The returned graph will contain an edge feature `EID` corresponding to the id of the edge
+in the original graph. If `dropnodes=true`, it will also contain a node feature `NID` with
+the node ids in the original graph.
+
+# Arguments
+
+- `g`. The graph.
+- `nodes`. A list of node IDs to sample neighbors from.
+- `K`. The maximum number of edges to be sampled for each node.
+       If -1, all the neighboring edges will be selected.
+- `dir`. Determines whether to sample inbound (`:in`) or outbound (``:out`) edges (Default `:in`).
+- `replace`. If `true`, sample with replacement.
+- `dropnodes`. If `true`, the resulting subgraph will contain only the nodes involved in the sampled edges.
+     
+# Examples
+
+```julia
+julia> g = rand_graph(20, 100)
+GNNGraph:
+    num_nodes = 20
+    num_edges = 100
+
+julia> sample_neighbors(g, 2:3)
+GNNGraph:
+    num_nodes = 20
+    num_edges = 9
+    edata:
+        EID => (9,)
+
+julia> sg = sample_neighbors(g, 2:3, dropnodes=true)
+GNNGraph:
+    num_nodes = 10
+    num_edges = 9
+    ndata:
+        NID => (10,)
+    edata:
+        EID => (9,)
+
+julia> sg.ndata.NID
+10-element Vector{Int64}:
+  2
+  3
+ 17
+ 14
+ 18
+ 15
+ 16
+ 20
+  7
+ 10
+
+julia> sample_neighbors(g, 2:3, 5, replace=true)
+GNNGraph:
+    num_nodes = 20
+    num_edges = 10
+    edata:
+        EID => (10,)
+```
+"""
+function sample_neighbors(g::GNNGraph{<:COO_T}, nodes, K=-1; 
+        dir=:in, replace=false, dropnodes=false)
+    @assert dir ∈ (:in, :out)
+    _, eidlist = adjacency_list(g, nodes; dir, with_eid=true)
+    for i in 1:length(eidlist)
+        if replace 
+            k = K > 0 ? K : length(eidlist[i])
+        else
+            k = K > 0 ? min(length(eidlist[i]), K) : length(eidlist[i])
+        end
+        eidlist[i] = StatsBase.sample(eidlist[i], k; replace)
+    end
+    eids = reduce(vcat, eidlist)
+    s, t = edge_index(g)
+    w = get_edge_weight(g)
+    s = s[eids]
+    t = t[eids]
+    w = isnothing(w) ? nothing : w[eids]
+
+    edata = getobs(g.edata, eids)
+    edata = (edata..., EID = eids)
+
+    num_edges = length(eids)
+
+    if !dropnodes
+        graph = (s, t, w)
+    
+        gnew = GNNGraph(graph, 
+                    g.num_nodes, num_edges, g.num_graphs,
+                    g.graph_indicator,
+                    g.ndata, edata, g.gdata)
+    else    
+        nodes_other = dir == :in ? setdiff(s, nodes) : setdiff(t, nodes)
+        nodes_all = [nodes; nodes_other]
+        nodemap = Dict(n => i for (i, n) in enumerate(nodes_all))
+        s = [nodemap[s] for s in s]
+        t = [nodemap[t] for t in t]
+        graph = (s, t, w)
+        graph_indicator = g.graph_indicator !== nothing ? g.graph_indicator[nodes_all] : nothing
+        num_nodes = length(nodes_all)
+        ndata = getobs(g.ndata, nodes_all)
+        ndata = (ndata..., NID = nodes_all)        
+
+        gnew = GNNGraph(graph, 
+                    num_nodes, num_edges, g.num_graphs,
+                    graph_indicator,
+                    ndata, edata, g.gdata)
+    end
+    return gnew
+end

src/GNNGraphs/transform.jl

@@ -120,6 +120,23 @@ function add_edges(g::GNNGraph{<:COO_T},
             g.ndata, edata, g.gdata)
 end
 
+### TODO Cannot implement this since GNNGraph is immutable (cannot change num_edges)
+# function Graphs.add_edge!(g::GNNGraph{<:COO_T}, snew::T, tnew::T; edata=nothing) where T<:Union{Integer, AbstractVector}
+#     s, t = edge_index(g)
+#     @assert length(snew) == length(tnew)
+#     # TODO remove this constraint
+#     @assert get_edge_weight(g) === nothing
+    
+#     edata = normalize_graphdata(edata, default_name=:e, n=length(snew))
+#     edata = cat_features(g.edata, edata)
+    
+#     s, t = edge_index(g)
+#     append!(s, snew)
+#     append!(t, tnew)
+#     g.num_edges += length(snew)
+#     return true
+# end
+
 
 """
     add_nodes(g::GNNGraph, n; [ndata])

test/GNNGraphs/gnngraph.jl

@@ -31,6 +31,18 @@
         @test g.num_edges == 2
     end
 
+    @testset "Constructor: integer" begin
+        g = GNNGraph(10, graph_type=GRAPH_T)
+        @test g.num_nodes == 10
+        @test g.num_edges == 0
+
+        g2 = rand_graph(10, 30, graph_type=GRAPH_T)
+        G = typeof(g2)
+        g = G(10)
+        @test g.num_nodes == 10
+        @test g.num_edges == 0
+    end
+
     @testset "symmetric graph" begin
         s = [1, 1, 2, 2, 3, 3, 4, 4]
         t = [2, 4, 1, 3, 2, 4, 1, 3]

test/GNNGraphs/sampling.jl

@@ -0,0 +1,50 @@
+@testset "sampling.jl" begin
+    if GRAPH_T == :coo
+        @testset "sample_neighbors" begin
+            # replace = false
+            dir = :in
+            nodes = 2:3
+            g = rand_graph(10, 40, bidirected=false, graph_type=GRAPH_T)
+            sg = sample_neighbors(g, nodes; dir)
+            @test sg.num_nodes == 10
+            @test sg.num_edges == sum(degree(g, i; dir) for i in nodes)
+            @test size(sg.edata.EID) == (sg.num_edges,)
+            @test length(union(sg.edata.EID)) == length(sg.edata.EID)
+            adjlist = adjacency_list(g; dir)
+            s, t = edge_index(sg)
+            @test all(t .∈ Ref(nodes))
+            for i in nodes
+                @test sort(neighbors(sg, i; dir)) == sort(neighbors(g, i; dir))
+            end
+
+            # replace = true
+            dir = :out
+            nodes = 2:3
+            K = 2
+            g = rand_graph(10, 40, bidirected=false, graph_type=GRAPH_T)
+            sg = sample_neighbors(g, nodes, K; dir, replace=true)
+            @test sg.num_nodes == 10
+            @test sg.num_edges == sum(K for i in nodes)
+            @test size(sg.edata.EID) == (sg.num_edges,)
+            adjlist = adjacency_list(g; dir)
+            s, t = edge_index(sg)
+            @test all(s .∈ Ref(nodes))
+            for i in nodes
+                @test issubset(neighbors(sg, i; dir), adjlist[i])
+            end
+
+            # dropnodes = true
+            dir = :in
+            nodes = 2:3
+            g = rand_graph(10, 40, bidirected=false, graph_type=GRAPH_T)
+            g = GNNGraph(g, ndata=(x1=rand(10),), edata=(e1=rand(40),))
+            sg = sample_neighbors(g, nodes; dir, dropnodes=true)
+            @test sg.num_edges == sum(degree(g, i; dir) for i in nodes)
+            @test size(sg.edata.EID) == (sg.num_edges,)
+            @test size(sg.ndata.NID) == (sg.num_nodes,)
+            @test sg.edata.e1 == g.edata.e1[sg.edata.EID]
+            @test sg.ndata.x1 == g.ndata.x1[sg.ndata.NID]
+            @test length(union(sg.ndata.NID)) == length(sg.ndata.NID)
+        end
+    end
+end