temporarily reintegrate GNNGraphs tests (#449)

* add back GNNGraph tests

* reintegrate extensions

* don't use DeviceUtils

* import Flux

ext/GraphNeuralNetworksCUDAExt/GNNGraphs/query.jl

@@ -0,0 +1,2 @@
+
+GNNGraphs._rand_dense_vector(A::CUMAT_T) = CUDA.randn(size(A, 1))

ext/GraphNeuralNetworksCUDAExt/GNNGraphs/transform.jl

@@ -0,0 +1,2 @@
+
+GNNGraphs.dense_zeros_like(a::CUMAT_T, T::Type, sz = size(a)) = CUDA.zeros(T, sz)

ext/GraphNeuralNetworksCUDAExt/GNNGraphs/utils.jl

@@ -0,0 +1,8 @@
+
+GNNGraphs.iscuarray(x::AnyCuArray) = true
+
+
+function sort_edge_index(u::AnyCuArray, v::AnyCuArray)
+    #TODO proper cuda friendly implementation
+    sort_edge_index(u |> Flux.cpu, v |> Flux.cpu) |> Flux.gpu
+end

ext/GraphNeuralNetworksCUDAExt/GraphNeuralNetworksCUDAExt.jl

@@ -9,6 +9,9 @@ import GraphNeuralNetworks: propagate
 
 const CUMAT_T = Union{CUDA.AnyCuMatrix, CUDA.CUSPARSE.CuSparseMatrix}
 
+include("GNNGraphs/query.jl")
+include("GNNGraphs/transform.jl")
+include("GNNGraphs/utils.jl")
 include("msgpass.jl")
 
 end #module

ext/GraphNeuralNetworksSimpleWeightedGraphsExt/GraphNeuralNetworksSimpleWeightedGraphsExt.jl

@@ -0,0 +1,12 @@
+module GraphNeuralNetworksSimpleWeightedGraphsExt
+
+using GraphNeuralNetworks
+using Graphs
+using SimpleWeightedGraphs
+
+function GraphNeuralNetworks.GNNGraph(g::T; kws...) where 
+                 {T <: Union{SimpleWeightedGraph, SimpleWeightedDiGraph}}
+   return GNNGraph(g.weights, kws...)
+end
+
+end #module

src/GNNGraphs/GNNGraphs.jl

@@ -7,6 +7,8 @@ using Graphs: AbstractGraph, outneighbors, inneighbors, adjacency_matrix, degree
               has_self_loops, is_directed
 import NearestNeighbors
 import NNlib
+import Flux
+using Flux: batch
 import StatsBase
 import KrylovKit
 using ChainRulesCore

src/GNNGraphs/transform.jl

@@ -1126,10 +1126,13 @@ function negative_sample(g::GNNGraph;
 
     s, t = edge_index(g)
     n = g.num_nodes
-    device = get_device(s)
-    cdevice = cpu_device()
-    # Convert to gpu since set operations and sampling are not supported by CUDA.jl
-    s, t = cdevice(s), cdevice(t)
+    if iscuarray(s)
+        # Convert to gpu since set operations and sampling are not supported by CUDA.jl
+        device = Flux.gpu
+        s, t = Flux.cpu(s), Flux.cpu(t)
+    else
+        device = Flux.cpu
+    end
     idx_pos, maxid = edge_encoding(s, t, n)
     if bidirected
         num_neg_edges = num_neg_edges ÷ 2
@@ -1156,6 +1159,7 @@ function negative_sample(g::GNNGraph;
     return GNNGraph(s_neg, t_neg, num_nodes = n) |> device
 end
 
+
 """
     rand_edge_split(g::GNNGraph, frac; bidirected=is_bidirected(g)) -> g1, g2
 

test/GNNGraphs/chainrules.jl

@@ -0,0 +1,24 @@
+@testset "dict constructor" begin
+    grad = gradient(1.) do x
+        d = Dict([:x => x, :y => 5]...)    
+        return sum(d[:x].^2)
+    end[1]
+
+    @test grad == 2
+
+    ## BROKEN Constructors
+    # grad = gradient(1.) do x
+    #     d = Dict([(:x => x), (:y => 5)])    
+    #     return sum(d[:x].^2)
+    # end[1]
+
+    # @test grad == 2
+
+
+    # grad = gradient(1.) do x
+    #     d = Dict([(:x => x), (:y => 5)])    
+    #     return sum(d[:x].^2)
+    # end[1]
+
+    # @test grad == 2
+end

test/GNNGraphs/convert.jl

@@ -0,0 +1,20 @@
+if TEST_GPU
+    @testset "to_coo(dense) on gpu" begin
+        get_st(A) = GNNGraphs.to_coo(A)[1][1:2]
+        get_val(A) = GNNGraphs.to_coo(A)[1][3]
+
+        A = cu([0 2 2; 2.0 0 2; 2 2 0])
+
+        y = get_val(A)
+        @test y isa CuVector{Float32}
+        @test Array(y) ≈ [2, 2, 2, 2, 2, 2]
+
+        s, t = get_st(A)
+        @test s isa CuVector{<:Integer}
+        @test t isa CuVector{<:Integer}
+        @test Array(s) == [2, 3, 1, 3, 1, 2]
+        @test Array(t) == [1, 1, 2, 2, 3, 3]
+
+        @test gradient(A -> sum(get_val(A)), A)[1] isa CuMatrix{Float32}
+    end
+end

test/GNNGraphs/datastore.jl

@@ -0,0 +1,101 @@
+
+@testset "constructor" begin
+    @test_throws AssertionError DataStore(10, (:x => rand(10), :y => rand(2, 4)))
+
+    @testset "keyword args" begin
+        ds = DataStore(10, x = rand(10), y = rand(2, 10))
+        @test size(ds.x) == (10,)
+        @test size(ds.y) == (2, 10)
+
+        ds = DataStore(x = rand(10), y = rand(2, 10))
+        @test size(ds.x) == (10,)
+        @test size(ds.y) == (2, 10)
+    end
+end
+
+@testset "getproperty / setproperty!" begin
+    x = rand(10)
+    ds = DataStore(10, (:x => x, :y => rand(2, 10)))
+    @test ds.x == ds[:x] == x
+    @test_throws DimensionMismatch ds.z=rand(12)
+    ds.z = [1:10;]
+    @test ds.z == [1:10;]
+    vec = [DataStore(10, (:x => x,)), DataStore(10, (:x => x, :y => rand(2, 10)))]
+    @test vec.x == [x, x]
+    @test_throws KeyError vec.z 
+    @test vec._n == [10, 10]
+    @test vec._data == [Dict(:x => x), Dict(:x => x, :y => vec[2].y)]
+end
+
+@testset "setindex!" begin
+    ds = DataStore(10)
+    x = rand(10)
+    @test (ds[:x] = x) == x # Tests setindex!
+    @test ds.x == ds[:x] == x
+end
+
+@testset "map" begin
+    ds = DataStore(10, (:x => rand(10), :y => rand(2, 10)))
+    ds2 = map(x -> x .+ 1, ds)
+    @test ds2.x == ds.x .+ 1
+    @test ds2.y == ds.y .+ 1
+
+    @test_throws AssertionError ds2=map(x -> [x; x], ds)
+end
+
+@testset "getdata / getn" begin
+    ds = DataStore(10, (:x => rand(10), :y => rand(2, 10)))
+    @test getdata(ds) == getfield(ds, :_data)
+    @test_throws KeyError ds.data
+    @test getn(ds) == getfield(ds, :_n)
+    @test_throws KeyError ds.n
+end
+
+@testset "cat empty" begin
+    ds1 = DataStore(2, (:x => rand(2)))
+    ds2 = DataStore(1, (:x => rand(1)))
+    dsempty = DataStore(0, (:x => rand(0)))
+
+    ds = GNNGraphs.cat_features(ds1, ds2)
+    @test getn(ds) == 3
+    ds = GNNGraphs.cat_features(ds1, dsempty)
+    @test getn(ds) == 2
+
+    # issue #280
+    g = GNNGraph([1], [2])
+    h = add_edges(g, Int[], Int[])  # adds no edges
+    @test getn(g.edata) == 1
+    @test getn(h.edata) == 1
+end
+
+
+@testset "gradient" begin
+    ds = DataStore(10, (:x => rand(10), :y => rand(2, 10)))
+
+    f1(ds) = sum(ds.x)
+    grad = gradient(f1, ds)[1]
+    @test grad._data[:x] ≈ ngradient(f1, ds)[1][:x]
+
+    g = rand_graph(5, 2)
+    x = rand(2, 5)
+    grad = gradient(x -> sum(exp, GNNGraph(g, ndata = x).ndata.x), x)[1]
+    @test grad == exp.(x)
+end
+
+@testset "functor" begin
+    ds = DataStore(10, (:x => zeros(10), :y => ones(2, 10)))
+    p, re = Functors.functor(ds)
+    @test p[1] === getn(ds)
+    @test p[2] === getdata(ds)
+    @test ds == re(p)
+
+    ds2 = Functors.fmap(ds) do x
+        if x isa AbstractArray
+            x .+ 1
+        else
+            x
+        end
+    end
+    @test ds isa DataStore
+    @test ds2.x == ds.x .+ 1
+end

test/GNNGraphs/generate.jl

@@ -0,0 +1,122 @@
+@testset "rand_graph" begin
+    n, m = 10, 20
+    m2 = m ÷ 2
+    x = rand(3, n)
+    e = rand(4, m2)
+
+    g = rand_graph(n, m, ndata = x, edata = e, graph_type = GRAPH_T)
+    @test g.num_nodes == n
+    @test g.num_edges == m
+    @test g.ndata.x === x
+    if GRAPH_T == :coo
+        s, t = edge_index(g)
+        @test s[1:m2] == t[(m2 + 1):end]
+        @test t[1:m2] == s[(m2 + 1):end]
+        @test g.edata.e[:, 1:m2] == e
+        @test g.edata.e[:, (m2 + 1):end] == e
+    end
+
+    g = rand_graph(n, m, bidirected = false, seed = 17, graph_type = GRAPH_T)
+    @test g.num_nodes == n
+    @test g.num_edges == m
+
+    g2 = rand_graph(n, m, bidirected = false, seed = 17, graph_type = GRAPH_T)
+    @test edge_index(g2) == edge_index(g)
+
+    ew = rand(m2)
+    g = rand_graph(n, m, bidirected = true, seed = 17, graph_type = GRAPH_T, edge_weight = ew)
+    @test get_edge_weight(g) == [ew; ew] broken=(GRAPH_T != :coo)
+
+    ew = rand(m)
+    g = rand_graph(n, m, bidirected = false, seed = 17, graph_type = GRAPH_T, edge_weight = ew)
+    @test get_edge_weight(g) == ew broken=(GRAPH_T != :coo)
+end
+
+@testset "knn_graph" begin
+    n, k = 10, 3
+    x = rand(3, n)
+    g = knn_graph(x, k; graph_type = GRAPH_T)
+    @test g.num_nodes == 10
+    @test g.num_edges == n * k
+    @test degree(g, dir = :in) == fill(k, n)
+    @test has_self_loops(g) == false
+
+    g = knn_graph(x, k; dir = :out, self_loops = true, graph_type = GRAPH_T)
+    @test g.num_nodes == 10
+    @test g.num_edges == n * k
+    @test degree(g, dir = :out) == fill(k, n)
+    @test has_self_loops(g) == true
+
+    graph_indicator = [1, 1, 1, 1, 1, 2, 2, 2, 2, 2]
+    g = knn_graph(x, k; graph_indicator, graph_type = GRAPH_T)
+    @test g.num_graphs == 2
+    s, t = edge_index(g)
+    ne = n * k ÷ 2
+    @test all(1 .<= s[1:ne] .<= 5)
+    @test all(1 .<= t[1:ne] .<= 5)
+    @test all(6 .<= s[(ne + 1):end] .<= 10)
+    @test all(6 .<= t[(ne + 1):end] .<= 10)
+end
+
+@testset "radius_graph" begin
+    n, r = 10, 0.5
+    x = rand(3, n)
+    g = radius_graph(x, r; graph_type = GRAPH_T)
+    @test g.num_nodes == 10
+    @test has_self_loops(g) == false
+
+    g = radius_graph(x, r; dir = :out, self_loops = true, graph_type = GRAPH_T)
+    @test g.num_nodes == 10
+    @test has_self_loops(g) == true
+
+    graph_indicator = [1, 1, 1, 1, 1, 2, 2, 2, 2, 2]
+    g = radius_graph(x, r; graph_indicator, graph_type = GRAPH_T)
+    @test g.num_graphs == 2
+    s, t = edge_index(g)
+    @test (s .> 5) == (t .> 5)
+end
+
+@testset "rand_bipartite_heterograph" begin
+    g = rand_bipartite_heterograph(10, 15, 20)
+    @test g.num_nodes == Dict(:A => 10, :B => 15)
+    @test g.num_edges == Dict((:A, :to, :B) => 20, (:B, :to, :A) => 20)
+    sA, tB = edge_index(g, (:A, :to, :B))
+    for (s, t) in zip(sA, tB)
+        @test 1 <= s <= 10
+        @test 1 <= t <= 15
+        @test has_edge(g, (:A,:to,:B), s, t)
+        @test has_edge(g, (:B,:to,:A), t, s)
+    end
+
+    g = rand_bipartite_heterograph((2, 2), (4, 0), bidirected=false)
+    @test has_edge(g, (:A,:to,:B), 1, 1)
+    @test !has_edge(g, (:B,:to,:A), 1, 1)
+end
+
+@testset "rand_temporal_radius_graph" begin
+    number_nodes = 30
+    number_snapshots = 5
+    r = 0.1
+    speed = 0.1
+    tg = rand_temporal_radius_graph(number_nodes, number_snapshots, speed, r)
+    @test tg.num_nodes == [number_nodes for i in 1:number_snapshots]
+    @test tg.num_snapshots == number_snapshots
+    r2 = 0.95
+    tg2 = rand_temporal_radius_graph(number_nodes, number_snapshots, speed, r2)
+    @test mean(mean(degree.(tg.snapshots)))<=mean(mean(degree.(tg2.snapshots)))
+end
+
+@testset "rand_temporal_hyperbolic_graph" begin
+    @test GraphNeuralNetworks.GNNGraphs._hyperbolic_distance([1.0,1.0],[1.0,1.0];ζ=1)==0
+    @test GraphNeuralNetworks.GNNGraphs._hyperbolic_distance([0.23,0.11],[0.98,0.55];ζ=1)==GraphNeuralNetworks.GNNGraphs._hyperbolic_distance([0.98,0.55],[0.23,0.11];ζ=1)
+    number_nodes = 30
+    number_snapshots = 5
+    α, R, speed, ζ = 1, 1, 0.1, 1
+
+    tg = rand_temporal_hyperbolic_graph(number_nodes, number_snapshots; α, R, speed, ζ)
+    @test tg.num_nodes == [number_nodes for i in 1:number_snapshots]
+    @test tg.num_snapshots == number_snapshots
+    R = 10
+    tg1 = rand_temporal_hyperbolic_graph(number_nodes, number_snapshots; α, R, speed, ζ)
+    @test mean(mean(degree.(tg1.snapshots)))<=mean(mean(degree.(tg.snapshots)))
+end