diff --git a/ext/GraphNeuralNetworksCUDAExt/GNNGraphs/query.jl b/ext/GraphNeuralNetworksCUDAExt/GNNGraphs/query.jl new file mode 100644 index 000000000..0e74f725e --- /dev/null +++ b/ext/GraphNeuralNetworksCUDAExt/GNNGraphs/query.jl @@ -0,0 +1,2 @@ + +GNNGraphs._rand_dense_vector(A::CUMAT_T) = CUDA.randn(size(A, 1)) diff --git a/ext/GraphNeuralNetworksCUDAExt/GNNGraphs/transform.jl b/ext/GraphNeuralNetworksCUDAExt/GNNGraphs/transform.jl new file mode 100644 index 000000000..d2ee417fc --- /dev/null +++ b/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) diff --git a/ext/GraphNeuralNetworksCUDAExt/GNNGraphs/utils.jl b/ext/GraphNeuralNetworksCUDAExt/GNNGraphs/utils.jl new file mode 100644 index 000000000..c3d78e9c1 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/ext/GraphNeuralNetworksCUDAExt/GraphNeuralNetworksCUDAExt.jl b/ext/GraphNeuralNetworksCUDAExt/GraphNeuralNetworksCUDAExt.jl index 538083043..5e0132890 100644 --- a/ext/GraphNeuralNetworksCUDAExt/GraphNeuralNetworksCUDAExt.jl +++ b/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 diff --git a/ext/GraphNeuralNetworksSimpleWeightedGraphsExt/GraphNeuralNetworksSimpleWeightedGraphsExt.jl b/ext/GraphNeuralNetworksSimpleWeightedGraphsExt/GraphNeuralNetworksSimpleWeightedGraphsExt.jl new file mode 100644 index 000000000..aabc13443 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/GNNGraphs/GNNGraphs.jl b/src/GNNGraphs/GNNGraphs.jl index 2757c9f92..bea3385f9 100644 --- a/src/GNNGraphs/GNNGraphs.jl +++ b/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 diff --git a/src/GNNGraphs/transform.jl b/src/GNNGraphs/transform.jl index a732b9e92..ee6b2fc1f 100644 --- a/src/GNNGraphs/transform.jl +++ b/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 diff --git a/test/GNNGraphs/chainrules.jl b/test/GNNGraphs/chainrules.jl new file mode 100644 index 000000000..f0df6b6ca --- /dev/null +++ b/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 diff --git a/test/GNNGraphs/convert.jl b/test/GNNGraphs/convert.jl new file mode 100644 index 000000000..898a8d771 --- /dev/null +++ b/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 diff --git a/test/GNNGraphs/datastore.jl b/test/GNNGraphs/datastore.jl new file mode 100644 index 000000000..1c8cfdc1c --- /dev/null +++ b/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 diff --git a/test/GNNGraphs/generate.jl b/test/GNNGraphs/generate.jl new file mode 100644 index 000000000..d9f281fb2 --- /dev/null +++ b/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 diff --git a/test/GNNGraphs/gnngraph.jl b/test/GNNGraphs/gnngraph.jl new file mode 100644 index 000000000..f1c952cb1 --- /dev/null +++ b/test/GNNGraphs/gnngraph.jl @@ -0,0 +1,380 @@ +@testset "Constructor: adjacency matrix" begin + A = sprand(10, 10, 0.5) + sA, tA, vA = findnz(A) + + g = GNNGraph(A, graph_type = GRAPH_T) + s, t = edge_index(g) + v = get_edge_weight(g) + @test s == sA + @test t == tA + @test v == vA + + g = GNNGraph(Matrix(A), graph_type = GRAPH_T) + s, t = edge_index(g) + v = get_edge_weight(g) + @test s == sA + @test t == tA + @test v == vA + + g = GNNGraph([0 0 0 + 0 0 1 + 0 1 0], graph_type = GRAPH_T) + @test g.num_nodes == 3 + @test g.num_edges == 2 + + g = GNNGraph([0 1 0 + 1 0 0 + 0 0 0], graph_type = GRAPH_T) + @test g.num_nodes == 3 + @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 + + g = GNNGraph(graph_type = GRAPH_T) + @test g.num_nodes == 0 +end + +@testset "symmetric graph" begin + s = [1, 1, 2, 2, 3, 3, 4, 4] + t = [2, 4, 1, 3, 2, 4, 1, 3] + adj_mat = [0 1 0 1 + 1 0 1 0 + 0 1 0 1 + 1 0 1 0] + adj_list_out = [[2, 4], [1, 3], [2, 4], [1, 3]] + adj_list_in = [[2, 4], [1, 3], [2, 4], [1, 3]] + + # core functionality + g = GNNGraph(s, t; graph_type = GRAPH_T) + if TEST_GPU + g_gpu = g |> gpu + end + + @test g.num_edges == 8 + @test g.num_nodes == 4 + @test nv(g) == g.num_nodes + @test ne(g) == g.num_edges + @test Tuple.(collect(edges(g))) |> sort == collect(zip(s, t)) |> sort + @test sort(outneighbors(g, 1)) == [2, 4] + @test sort(inneighbors(g, 1)) == [2, 4] + @test is_directed(g) == true + s1, t1 = sort_edge_index(edge_index(g)) + @test s1 == s + @test t1 == t + @test vertices(g) == 1:(g.num_nodes) + + @test sort.(adjacency_list(g; dir = :in)) == adj_list_in + @test sort.(adjacency_list(g; dir = :out)) == adj_list_out + + @testset "adjacency_matrix" begin + @test adjacency_matrix(g) == adj_mat + @test adjacency_matrix(g; dir = :in) == adj_mat + @test adjacency_matrix(g; dir = :out) == adj_mat + + if TEST_GPU + # See https://github.com/JuliaGPU/CUDA.jl/pull/1093 + mat_gpu = adjacency_matrix(g_gpu) + @test mat_gpu isa ACUMatrix{Int} + @test Array(mat_gpu) == adj_mat + end + end + + @testset "normalized_laplacian" begin + mat = normalized_laplacian(g) + if TEST_GPU + mat_gpu = normalized_laplacian(g_gpu) + @test mat_gpu isa ACUMatrix{Float32} + @test Array(mat_gpu) == mat + end + end + + @testset "scaled_laplacian" begin if TEST_GPU + @test_broken begin + mat = scaled_laplacian(g) + mat_gpu = scaled_laplacian(g_gpu) + @test mat_gpu isa ACUMatrix{Float32} + @test Array(mat_gpu) == mat + end + end end + + @testset "constructors" begin + adjacency_matrix(g; dir = :out) == adj_mat + adjacency_matrix(g; dir = :in) == adj_mat + end + + if TEST_GPU + @testset "functor" begin + s_cpu, t_cpu = edge_index(g) + s_gpu, t_gpu = edge_index(g_gpu) + @test s_gpu isa CuVector{Int} + @test Array(s_gpu) == s_cpu + @test t_gpu isa CuVector{Int} + @test Array(t_gpu) == t_cpu + end + end +end + +@testset "asymmetric graph" begin + s = [1, 2, 3, 4] + t = [2, 3, 4, 1] + adj_mat_out = [0 1 0 0 + 0 0 1 0 + 0 0 0 1 + 1 0 0 0] + adj_list_out = [[2], [3], [4], [1]] + + adj_mat_in = [0 0 0 1 + 1 0 0 0 + 0 1 0 0 + 0 0 1 0] + adj_list_in = [[4], [1], [2], [3]] + + # core functionality + g = GNNGraph(s, t; graph_type = GRAPH_T) + if TEST_GPU + g_gpu = g |> gpu + end + + @test g.num_edges == 4 + @test g.num_nodes == 4 + @test length(edges(g)) == 4 + @test sort(outneighbors(g, 1)) == [2] + @test sort(inneighbors(g, 1)) == [4] + @test is_directed(g) == true + @test is_directed(typeof(g)) == true + s1, t1 = sort_edge_index(edge_index(g)) + @test s1 == s + @test t1 == t + + # adjacency + @test adjacency_matrix(g) == adj_mat_out + @test adjacency_list(g) == adj_list_out + @test adjacency_matrix(g, dir = :out) == adj_mat_out + @test adjacency_list(g, dir = :out) == adj_list_out + @test adjacency_matrix(g, dir = :in) == adj_mat_in + @test adjacency_list(g, dir = :in) == adj_list_in +end + +@testset "zero" begin + g = rand_graph(4, 6, graph_type = GRAPH_T) + G = typeof(g) + @test zero(G) == G(0) +end + +@testset "Graphs.jl constructor" begin + lg = random_regular_graph(10, 4) + @test !Graphs.is_directed(lg) + g = GNNGraph(lg) + @test g.num_edges == 2 * ne(lg) # g in undirected + @test Graphs.is_directed(g) + for e in Graphs.edges(lg) + i, j = src(e), dst(e) + @test has_edge(g, i, j) + @test has_edge(g, j, i) + end + + @testset "SimpleGraph{Int32}" begin + g = GNNGraph(SimpleGraph{Int32}(6), graph_type = GRAPH_T) + @test g.num_nodes == 6 + end +end + +@testset "Features" begin + g = GNNGraph(sprand(10, 10, 0.3), graph_type = GRAPH_T) + + # default names + X = rand(10, g.num_nodes) + E = rand(10, g.num_edges) + U = rand(10, g.num_graphs) + + g = GNNGraph(g, ndata = X, edata = E, gdata = U) + @test g.ndata.x === X + @test g.edata.e === E + @test g.gdata.u === U + @test g.x === g.ndata.x + @test g.e === g.edata.e + @test g.u === g.gdata.u + + # Check no args + g = GNNGraph(g) + @test g.ndata.x === X + @test g.edata.e === E + @test g.gdata.u === U + + # multiple features names + g = GNNGraph(g, ndata = (x2 = 2X, g.ndata...), edata = (e2 = 2E, g.edata...), + gdata = (u2 = 2U, g.gdata...)) + @test g.ndata.x === X + @test g.edata.e === E + @test g.gdata.u === U + @test g.ndata.x2 ≈ 2X + @test g.edata.e2 ≈ 2E + @test g.gdata.u2 ≈ 2U + @test g.x === g.ndata.x + @test g.e === g.edata.e + @test g.u === g.gdata.u + @test g.x2 === g.ndata.x2 + @test g.e2 === g.edata.e2 + @test g.u2 === g.gdata.u2 + + # Dimension checks + @test_throws AssertionError GNNGraph(erdos_renyi(10, 30), edata = rand(29), + graph_type = GRAPH_T) + @test_throws AssertionError GNNGraph(erdos_renyi(10, 30), edata = rand(2, 29), + graph_type = GRAPH_T) + @test_throws AssertionError GNNGraph(erdos_renyi(10, 30), + edata = (; x = rand(30), y = rand(29)), + graph_type = GRAPH_T) + + # Copy features on reverse edge + e = rand(30) + g = GNNGraph(erdos_renyi(10, 30), edata = e, graph_type = GRAPH_T) + @test g.edata.e == [e; e] + + # non-array global + g = rand_graph(10, 30, gdata = "ciao", graph_type = GRAPH_T) + @test g.gdata.u == "ciao" + + # vectors stays vectors + g = rand_graph(10, 30, ndata = rand(10), + edata = rand(30), + gdata = (u = rand(2), z = rand(1), q = 1), + graph_type = GRAPH_T) + @test size(g.ndata.x) == (10,) + @test size(g.edata.e) == (30,) + @test size(g.gdata.u) == (2, 1) + @test size(g.gdata.z) == (1,) + @test g.gdata.q === 1 + + # Error for non-array ndata + @test_throws AssertionError rand_graph(10, 30, ndata = "ciao", graph_type = GRAPH_T) + @test_throws AssertionError rand_graph(10, 30, ndata = 1, graph_type = GRAPH_T) + + # Error for Ambiguous getproperty + g = rand_graph(10, 20, ndata = rand(2, 10), edata = (; x = rand(3, 20)), + graph_type = GRAPH_T) + @test size(g.ndata.x) == (2, 10) + @test size(g.edata.x) == (3, 20) + @test_throws ArgumentError g.x +end + +@testset "MLUtils and DataLoader compat" begin + n, m, num_graphs = 10, 30, 50 + X = rand(10, n) + E = rand(10, m) + U = rand(10, 1) + data = [rand_graph(n, m, ndata = X, edata = E, gdata = U, graph_type = GRAPH_T) + for _ in 1:num_graphs] + g = Flux.batch(data) + + @testset "batch then pass to dataloader" begin + @test MLUtils.getobs(g, 3) == getgraph(g, 3) + @test MLUtils.getobs(g, 3:5) == getgraph(g, 3:5) + @test MLUtils.numobs(g) == g.num_graphs + + d = Flux.DataLoader(g, batchsize = 2, shuffle = false) + @test first(d) == getgraph(g, 1:2) + end + + @testset "pass to dataloader and no automatic collation" begin + @test MLUtils.getobs(data, 3) == data[3] + @test MLUtils.getobs(data, 3:5) isa Vector{<:GNNGraph} + @test MLUtils.getobs(data, 3:5) == [data[3], data[4], data[5]] + @test MLUtils.numobs(data) == g.num_graphs + + d = Flux.DataLoader(data, batchsize = 2, shuffle = false) + @test first(d) == [data[1], data[2]] + end +end + +@testset "Graphs.jl integration" begin + g = GNNGraph(erdos_renyi(10, 20), graph_type = GRAPH_T) + @test g isa Graphs.AbstractGraph +end + +@testset "==" begin + g1 = rand_graph(5, 6, ndata = rand(5), edata = rand(6), graph_type = GRAPH_T) + @test g1 == g1 + @test g1 == deepcopy(g1) + @test g1 !== deepcopy(g1) + + g2 = GNNGraph(g1, graph_type = GRAPH_T) + @test g1 == g2 + @test g1 === g2 # this is true since GNNGraph is immutable + + g2 = GNNGraph(g1, ndata = rand(5), graph_type = GRAPH_T) + @test g1 != g2 + @test g1 !== g2 + + g2 = GNNGraph(g1, edata = rand(6), graph_type = GRAPH_T) + @test g1 != g2 + @test g1 !== g2 +end + +@testset "hash" begin + g1 = rand_graph(5, 6, ndata = rand(5), edata = rand(6), graph_type = GRAPH_T) + @test hash(g1) == hash(g1) + @test hash(g1) == hash(deepcopy(g1)) + @test hash(g1) == hash(GNNGraph(g1, ndata = g1.ndata, graph_type = GRAPH_T)) + @test hash(g1) == hash(GNNGraph(g1, ndata = g1.ndata, graph_type = GRAPH_T)) + @test hash(g1) != hash(GNNGraph(g1, ndata = rand(5), graph_type = GRAPH_T)) + @test hash(g1) != hash(GNNGraph(g1, edata = rand(6), graph_type = GRAPH_T)) +end + +@testset "copy" begin + g1 = rand_graph(10, 4, ndata = rand(2, 10), graph_type = GRAPH_T) + g2 = copy(g1) + @test g1 === g2 # shallow copies are identical for immutable objects + + g2 = copy(g1, deep = true) + @test g1 == g2 + @test g1 !== g2 +end + +## Cannot test this because DataStore is not an ordered collection +## Uncomment when/if it will be based on OrderedDict +# @testset "show" begin +# @test sprint(show, rand_graph(10, 20)) == "GNNGraph(10, 20) with no data" +# @test sprint(show, rand_graph(10, 20, ndata=rand(5, 10))) == "GNNGraph(10, 20) with x: 5×10 data" +# @test sprint(show, rand_graph(10, 20, ndata=(a=rand(5, 10), b=rand(3, 10)), edata=rand(2, 20), gdata=(q=rand(1, 1), p=rand(3, 1)))) == "GNNGraph(10, 20) with (a: 5×10, b: 3×10), e: 2×20, (q: 1×1, p: 3×1) data" +# @test sprint(show, rand_graph(10, 20, ndata=(a=rand(5, 10),))) == "GNNGraph(10, 20) with a: 5×10 data" +# @test sprint(show, rand_graph(10, 20, ndata=rand(5, 10), edata=rand(2, 20))) == "GNNGraph(10, 20) with x: 5×10, e: 2×20 data" +# @test sprint(show, rand_graph(10, 20, ndata=rand(5, 10), gdata=rand(1, 1))) == "GNNGraph(10, 20) with x: 5×10, u: 1×1 data" +# @test sprint(show, rand_graph(10, 20, ndata=rand(5, 10), edata=(e=rand(2, 20), f=rand(2, 20), h=rand(3, 20)), gdata=rand(1, 1))) == "GNNGraph(10, 20) with x: 5×10, (e: 2×20, f: 2×20, h: 3×20), u: 1×1 data" +# @test sprint(show, rand_graph(10, 20, ndata=(a=rand(5, 10), b=rand(3, 10)), edata=rand(2, 20))) == "GNNGraph(10, 20) with (a: 5×10, b: 3×10), e: 2×20 data" +# @test sprint(show, rand_graph(10, 20, ndata=(a=rand(5,5, 10), b=rand(3,2, 10)), edata=rand(2, 20))) == "GNNGraph(10, 20) with (a: 5×5×10, b: 3×2×10), e: 2×20 data" +# end + +# @testset "show plain/text compact true" begin +# @test sprint(show, MIME("text/plain"), rand_graph(10, 20); context=:compact => true) == "GNNGraph(10, 20) with no data" +# @test sprint(show, MIME("text/plain"), rand_graph(10, 20, ndata=rand(5, 10)); context=:compact => true) == "GNNGraph(10, 20) with x: 5×10 data" +# @test sprint(show, MIME("text/plain"), rand_graph(10, 20, ndata=(a=rand(5, 10), b=rand(3, 10)), edata=rand(2, 20), gdata=(q=rand(1, 1), p=rand(3, 1))); context=:compact => true) == "GNNGraph(10, 20) with (a: 5×10, b: 3×10), e: 2×20, (q: 1×1, p: 3×1) data" +# @test sprint(show, MIME("text/plain"), rand_graph(10, 20, ndata=(a=rand(5, 10),)); context=:compact => true) == "GNNGraph(10, 20) with a: 5×10 data" +# @test sprint(show, MIME("text/plain"), rand_graph(10, 20, ndata=rand(5, 10), edata=rand(2, 20)); context=:compact => true) == "GNNGraph(10, 20) with x: 5×10, e: 2×20 data" +# @test sprint(show, MIME("text/plain"), rand_graph(10, 20, ndata=rand(5, 10), gdata=rand(1, 1)); context=:compact => true) == "GNNGraph(10, 20) with x: 5×10, u: 1×1 data" +# @test sprint(show, MIME("text/plain"), rand_graph(10, 20, ndata=rand(5, 10), edata=(e=rand(2, 20), f=rand(2, 20), h=rand(3, 20)), gdata=rand(1, 1)); context=:compact => true) == "GNNGraph(10, 20) with x: 5×10, (e: 2×20, f: 2×20, h: 3×20), u: 1×1 data" +# @test sprint(show, MIME("text/plain"), rand_graph(10, 20, ndata=(a=rand(5, 10), b=rand(3, 10)), edata=rand(2, 20)); context=:compact => true) == "GNNGraph(10, 20) with (a: 5×10, b: 3×10), e: 2×20 data" +# @test sprint(show, MIME("text/plain"), rand_graph(10, 20, ndata=(a=rand(5,5, 10), b=rand(3,2, 10)), edata=rand(2, 20)); context=:compact => true) == "GNNGraph(10, 20) with (a: 5×5×10, b: 3×2×10), e: 2×20 data" +# end + +# @testset "show plain/text compact false" begin +# @test sprint(show, MIME("text/plain"), rand_graph(10, 20); context=:compact => false) == "GNNGraph:\n num_nodes: 10\n num_edges: 20" +# @test sprint(show, MIME("text/plain"), rand_graph(10, 20, ndata=rand(5, 10)); context=:compact => false) == "GNNGraph:\n num_nodes: 10\n num_edges: 20\n ndata:\n\tx = 5×10 Matrix{Float64}" +# @test sprint(show, MIME("text/plain"), rand_graph(10, 20, ndata=(a=rand(5, 10), b=rand(3, 10)), edata=rand(2, 20), gdata=(q=rand(1, 1), p=rand(3, 1))); context=:compact => false) == "GNNGraph:\n num_nodes: 10\n num_edges: 20\n ndata:\n\ta = 5×10 Matrix{Float64}\n\tb = 3×10 Matrix{Float64}\n edata:\n\te = 2×20 Matrix{Float64}\n gdata:\n\tq = 1×1 Matrix{Float64}\n\tp = 3×1 Matrix{Float64}" +# @test sprint(show, MIME("text/plain"), rand_graph(10, 20, ndata=(a=rand(5, 10),)); context=:compact => false) == "GNNGraph:\n num_nodes: 10\n num_edges: 20\n ndata:\n\ta = 5×10 Matrix{Float64}" +# @test sprint(show, MIME("text/plain"), rand_graph(10, 20, ndata=rand(5, 10), edata=rand(2, 20)); context=:compact => false) == "GNNGraph:\n num_nodes: 10\n num_edges: 20\n ndata:\n\tx = 5×10 Matrix{Float64}\n edata:\n\te = 2×20 Matrix{Float64}" +# @test sprint(show, MIME("text/plain"), rand_graph(10, 20, ndata=rand(5, 10), gdata=rand(1, 1)); context=:compact => false) == "GNNGraph:\n num_nodes: 10\n num_edges: 20\n ndata:\n\tx = 5×10 Matrix{Float64}\n gdata:\n\tu = 1×1 Matrix{Float64}" +# @test sprint(show, MIME("text/plain"), rand_graph(10, 20, ndata=rand(5, 10), edata=(e=rand(2, 20), f=rand(2, 20), h=rand(3, 20)), gdata=rand(1, 1)); context=:compact => false) == "GNNGraph:\n num_nodes: 10\n num_edges: 20\n ndata:\n\tx = 5×10 Matrix{Float64}\n edata:\n\te = 2×20 Matrix{Float64}\n\tf = 2×20 Matrix{Float64}\n\th = 3×20 Matrix{Float64}\n gdata:\n\tu = 1×1 Matrix{Float64}" +# @test sprint(show, MIME("text/plain"), rand_graph(10, 20, ndata=(a=rand(5, 10), b=rand(3, 10)), edata=rand(2, 20)); context=:compact => false) == "GNNGraph:\n num_nodes: 10\n num_edges: 20\n ndata:\n\ta = 5×10 Matrix{Float64}\n\tb = 3×10 Matrix{Float64}\n edata:\n\te = 2×20 Matrix{Float64}" +# @test sprint(show, MIME("text/plain"), rand_graph(10, 20, ndata=(a=rand(5, 5, 10), b=rand(3, 2, 10)), edata=rand(2, 20)); context=:compact => false) == "GNNGraph:\n num_nodes: 10\n num_edges: 20\n ndata:\n\ta = 5×5×10 Array{Float64, 3}\n\tb = 3×2×10 Array{Float64, 3}\n edata:\n\te = 2×20 Matrix{Float64}" +# end diff --git a/test/GNNGraphs/gnnheterograph.jl b/test/GNNGraphs/gnnheterograph.jl new file mode 100644 index 000000000..e17159dc0 --- /dev/null +++ b/test/GNNGraphs/gnnheterograph.jl @@ -0,0 +1,209 @@ + + +@testset "Empty constructor" begin + g = GNNHeteroGraph() + @test isempty(g.num_nodes) + g = add_edges(g, (:user, :like, :actor) => ([1,2,3,3,3], [3,5,1,9,4])) + @test g.num_nodes[:user] == 3 + @test g.num_nodes[:actor] == 9 + @test g.num_edges[(:user, :like, :actor)] == 5 +end + +@testset "Constructor from pairs" begin + hg = GNNHeteroGraph((:A, :e1, :B) => ([1,2,3,4], [3,2,1,5])) + @test hg.num_nodes == Dict(:A => 4, :B => 5) + @test hg.num_edges == Dict((:A, :e1, :B) => 4) + + hg = GNNHeteroGraph((:A, :e1, :B) => ([1,2,3], [3,2,1]), + (:A, :e2, :C) => ([1,2,3], [4,5,6])) + @test hg.num_nodes == Dict(:A => 3, :B => 3, :C => 6) + @test hg.num_edges == Dict((:A, :e1, :B) => 3, (:A, :e2, :C) => 3) +end + +@testset "Generation" begin + hg = rand_heterograph(Dict(:A => 10, :B => 20), + Dict((:A, :rel1, :B) => 30, (:B, :rel2, :A) => 10)) + + @test hg.num_nodes == Dict(:A => 10, :B => 20) + @test hg.num_edges == Dict((:A, :rel1, :B) => 30, (:B, :rel2, :A) => 10) + @test hg.graph_indicator === nothing + @test hg.num_graphs == 1 + @test hg.ndata isa Dict{Symbol, DataStore} + @test hg.edata isa Dict{Tuple{Symbol, Symbol, Symbol}, DataStore} + @test isempty(hg.gdata) + @test sort(hg.ntypes) == [:A, :B] + @test sort(hg.etypes) == [(:A, :rel1, :B), (:B, :rel2, :A)] + +end + +@testset "features" begin + hg = rand_heterograph(Dict(:A => 10, :B => 20), + Dict((:A, :rel1, :B) => 30, (:B, :rel2, :A) => 10), + ndata = Dict(:A => rand(2, 10), + :B => (x = rand(3, 20), y = rand(4, 20))), + edata = Dict((:A, :rel1, :B) => rand(5, 30)), + gdata = 1) + + @test size(hg.ndata[:A].x) == (2, 10) + @test size(hg.ndata[:B].x) == (3, 20) + @test size(hg.ndata[:B].y) == (4, 20) + @test size(hg.edata[(:A, :rel1, :B)].e) == (5, 30) + @test hg.gdata == DataStore(u = 1) + +end + +@testset "indexing syntax" begin + g = GNNHeteroGraph((:user, :rate, :movie) => ([1,1,2,3], [7,13,5,7])) + g[:movie].z = rand(Float32, 64, 13); + g[:user, :rate, :movie].e = rand(Float32, 64, 4); + g[:user].x = rand(Float32, 64, 3); + @test size(g.ndata[:user].x) == (64, 3) + @test size(g.ndata[:movie].z) == (64, 13) + @test size(g.edata[(:user, :rate, :movie)].e) == (64, 4) +end + + +@testset "simplified constructor" begin + hg = rand_heterograph((:A => 10, :B => 20), + ((:A, :rel1, :B) => 30, (:B, :rel2, :A) => 10), + ndata = (:A => rand(2, 10), + :B => (x = rand(3, 20), y = rand(4, 20))), + edata = (:A, :rel1, :B) => rand(5, 30), + gdata = 1) + + @test hg.num_nodes == Dict(:A => 10, :B => 20) + @test hg.num_edges == Dict((:A, :rel1, :B) => 30, (:B, :rel2, :A) => 10) + @test hg.graph_indicator === nothing + @test hg.num_graphs == 1 + @test size(hg.ndata[:A].x) == (2, 10) + @test size(hg.ndata[:B].x) == (3, 20) + @test size(hg.ndata[:B].y) == (4, 20) + @test size(hg.edata[(:A, :rel1, :B)].e) == (5, 30) + @test hg.gdata == DataStore(u = 1) + + nA, nB = 10, 20 + edges1 = rand(1:nA, 20), rand(1:nB, 20) + edges2 = rand(1:nB, 30), rand(1:nA, 30) + hg = GNNHeteroGraph(((:A, :rel1, :B) => edges1, (:B, :rel2, :A) => edges2)) + @test hg.num_edges == Dict((:A, :rel1, :B) => 20, (:B, :rel2, :A) => 30) + + nA, nB = 10, 20 + edges1 = rand(1:nA, 20), rand(1:nB, 20) + edges2 = rand(1:nB, 30), rand(1:nA, 30) + hg = GNNHeteroGraph(((:A, :rel1, :B) => edges1, (:B, :rel2, :A) => edges2); + num_nodes = (:A => nA, :B => nB)) + @test hg.num_nodes == Dict(:A => 10, :B => 20) + @test hg.num_edges == Dict((:A, :rel1, :B) => 20, (:B, :rel2, :A) => 30) +end + +@testset "num_edge_types / num_node_types" begin + hg = rand_heterograph((:A => 10, :B => 20), + ((:A, :rel1, :B) => 30, (:B, :rel2, :A) => 10), + ndata = (:A => rand(2, 10), + :B => (x = rand(3, 20), y = rand(4, 20))), + edata = (:A, :rel1, :B) => rand(5, 30), + gdata = 1) + @test num_edge_types(hg) == 2 + @test num_node_types(hg) == 2 + + g = rand_graph(10, 20) + @test num_edge_types(g) == 1 + @test num_node_types(g) == 1 +end + +@testset "numobs" begin + hg = rand_heterograph((:A => 10, :B => 20), + ((:A, :rel1, :B) => 30, (:B, :rel2, :A) => 10), + ndata = (:A => rand(2, 10), + :B => (x = rand(3, 20), y = rand(4, 20))), + edata = (:A, :rel1, :B) => rand(5, 30), + gdata = 1) + @test MLUtils.numobs(hg) == 1 +end + +@testset "get/set node features" begin + d, n = 3, 5 + g = rand_bipartite_heterograph(n, 2*n, 15) + g[:A].x = rand(Float32, d, n) + g[:B].y = rand(Float32, d, 2*n) + + @test size(g[:A].x) == (d, n) + @test size(g[:B].y) == (d, 2*n) +end + +@testset "add_edges" begin + d, n = 3, 5 + g = rand_bipartite_heterograph(n, 2 * n, 15) + s, t = [1, 2, 3], [3, 2, 1] + ## Keep the same ntypes - construct with args + g1 = add_edges(g, (:A, :rel1, :B), s, t) + @test num_node_types(g1) == 2 + @test num_edge_types(g1) == 3 + for i in eachindex(s, t) + @test has_edge(g1, (:A, :rel1, :B), s[i], t[i]) + end + # no change to num_nodes + @test g1.num_nodes[:A] == n + @test g1.num_nodes[:B] == 2n + + ## Keep the same ntypes - construct with a pair + g2 = add_edges(g, (:A, :rel1, :B) => (s, t)) + @test num_node_types(g2) == 2 + @test num_edge_types(g2) == 3 + for i in eachindex(s, t) + @test has_edge(g2, (:A, :rel1, :B), s[i], t[i]) + end + # no change to num_nodes + @test g2.num_nodes[:A] == n + @test g2.num_nodes[:B] == 2n + + ## New ntype with num_nodes (applies only to the new ntype) and edata + edata = rand(Float32, d, length(s)) + g3 = add_edges(g, + (:A, :rel1, :C) => (s, t); + num_nodes = Dict(:A => 1, :B => 1, :C => 10), + edata) + @test num_node_types(g3) == 3 + @test num_edge_types(g3) == 3 + for i in eachindex(s, t) + @test has_edge(g3, (:A, :rel1, :C), s[i], t[i]) + end + # added edata + @test g3.edata[(:A, :rel1, :C)].e == edata + # no change to existing num_nodes + @test g3.num_nodes[:A] == n + @test g3.num_nodes[:B] == 2n + # new num_nodes added as per kwarg + @test g3.num_nodes[:C] == 10 +end + +@testset "add self loops" begin + g1 = GNNHeteroGraph((:A, :to, :B) => ([1,2,3,4], [3,2,1,5])) + g2 = add_self_loops(g1, (:A, :to, :B)) + @test g2.num_edges[(:A, :to, :B)] === g1.num_edges[(:A, :to, :B)] + g1 = GNNHeteroGraph((:A, :to, :A) => ([1,2,3,4], [3,2,1,5])) + g2 = add_self_loops(g1, (:A, :to, :A)) + @test g2.num_edges[(:A, :to, :A)] === g1.num_edges[(:A, :to, :A)] + g1.num_nodes[(:A)] +end + +## Cannot test this because DataStore is not an ordered collection +## Uncomment when/if it will be based on OrderedDict +# @testset "show" begin +# num_nodes = Dict(:A => 10, :B => 20); +# edges1 = rand(1:num_nodes[:A], 20), rand(1:num_nodes[:B], 20) +# edges2 = rand(1:num_nodes[:B], 30), rand(1:num_nodes[:A], 30) +# eindex = ((:A, :rel1, :B) => edges1, (:B, :rel2, :A) => edges2) +# ndata = Dict(:A => (x = rand(2, num_nodes[:A]), y = rand(3, num_nodes[:A])),:B => rand(10, num_nodes[:B])) +# edata= Dict((:A, :rel1, :B) => (x = rand(2, 20), y = rand(3, 20)),(:B, :rel2, :A) => rand(10, 30)) +# hg1 = GraphNeuralNetworks.GNNHeteroGraph(eindex; num_nodes) +# hg2 = GraphNeuralNetworks.GNNHeteroGraph(eindex; num_nodes, ndata,edata) +# hg3 = GraphNeuralNetworks.GNNHeteroGraph(eindex; num_nodes, ndata) +# @test sprint(show, hg1) == "GNNHeteroGraph(Dict(:A => 10, :B => 20), Dict((:A, :rel1, :B) => 20, (:B, :rel2, :A) => 30))" +# @test sprint(show, hg2) == sprint(show, hg1) +# @test sprint(show, MIME("text/plain"), hg1; context=:compact => true) == "GNNHeteroGraph(Dict(:A => 10, :B => 20), Dict((:A, :rel1, :B) => 20, (:B, :rel2, :A) => 30))" +# @test sprint(show, MIME("text/plain"), hg2; context=:compact => true) == sprint(show, MIME("text/plain"), hg1;context=:compact => true) +# @test sprint(show, MIME("text/plain"), hg1; context=:compact => false) == "GNNHeteroGraph:\n num_nodes: (:A => 10, :B => 20)\n num_edges: ((:A, :rel1, :B) => 20, (:B, :rel2, :A) => 30)" +# @test sprint(show, MIME("text/plain"), hg2; context=:compact => false) == "GNNHeteroGraph:\n num_nodes: (:A => 10, :B => 20)\n num_edges: ((:A, :rel1, :B) => 20, (:B, :rel2, :A) => 30)\n ndata:\n\t:A => (x = 2×10 Matrix{Float64}, y = 3×10 Matrix{Float64})\n\t:B => x = 10×20 Matrix{Float64}\n edata:\n\t(:A, :rel1, :B) => (x = 2×20 Matrix{Float64}, y = 3×20 Matrix{Float64})\n\t(:B, :rel2, :A) => e = 10×30 Matrix{Float64}" +# @test sprint(show, MIME("text/plain"), hg3; context=:compact => false) =="GNNHeteroGraph:\n num_nodes: (:A => 10, :B => 20)\n num_edges: ((:A, :rel1, :B) => 20, (:B, :rel2, :A) => 30)\n ndata:\n\t:A => (x = 2×10 Matrix{Float64}, y = 3×10 Matrix{Float64})\n\t:B => x = 10×20 Matrix{Float64}" +# @test sprint(show, MIME("text/plain"), hg2; context=:compact => false) != sprint(show, MIME("text/plain"), hg3; context=:compact => false) +# end diff --git a/test/GNNGraphs/operators.jl b/test/GNNGraphs/operators.jl new file mode 100644 index 000000000..9ba65ae91 --- /dev/null +++ b/test/GNNGraphs/operators.jl @@ -0,0 +1,4 @@ +@testset "intersect" begin + g = rand_graph(10, 20, graph_type = GRAPH_T) + @test intersect(g, g).num_edges == 20 +end diff --git a/test/GNNGraphs/query.jl b/test/GNNGraphs/query.jl new file mode 100644 index 000000000..b0f03a262 --- /dev/null +++ b/test/GNNGraphs/query.jl @@ -0,0 +1,257 @@ +@testset "is_bidirected" begin + g = rand_graph(10, 20, bidirected = true, graph_type = GRAPH_T) + @test is_bidirected(g) + + g = rand_graph(10, 20, bidirected = false, graph_type = GRAPH_T) + @test !is_bidirected(g) +end + +@testset "has_multi_edges" begin if GRAPH_T == :coo + s = [1, 1, 2, 3] + t = [2, 2, 2, 4] + g = GNNGraph(s, t, graph_type = GRAPH_T) + @test has_multi_edges(g) + + s = [1, 2, 2, 3] + t = [2, 1, 2, 4] + g = GNNGraph(s, t, graph_type = GRAPH_T) + @test !has_multi_edges(g) +end end + +@testset "edges" begin + g = rand_graph(4, 10, graph_type = GRAPH_T) + @test edgetype(g) <: Graphs.Edge + for e in edges(g) + @test e isa Graphs.Edge + end +end + +@testset "has_isolated_nodes" begin + s = [1, 2, 3] + t = [2, 3, 2] + g = GNNGraph(s, t, graph_type = GRAPH_T) + @test has_isolated_nodes(g) == false + @test has_isolated_nodes(g, dir = :in) == true +end + +@testset "has_self_loops" begin + s = [1, 1, 2, 3] + t = [2, 2, 2, 4] + g = GNNGraph(s, t, graph_type = GRAPH_T) + @test has_self_loops(g) + + s = [1, 1, 2, 3] + t = [2, 2, 3, 4] + g = GNNGraph(s, t, graph_type = GRAPH_T) + @test !has_self_loops(g) +end + +@testset "degree" begin + @testset "unweighted" begin + s = [1, 1, 2, 3] + t = [2, 2, 2, 4] + g = GNNGraph(s, t, graph_type = GRAPH_T) + + @test degree(g) isa Vector{Int} + @test degree(g) == degree(g; dir = :out) == [2, 1, 1, 0] # default is outdegree + @test degree(g; dir = :in) == [0, 3, 0, 1] + @test degree(g; dir = :both) == [2, 4, 1, 1] + @test eltype(degree(g, Float32)) == Float32 + + if TEST_GPU + g_gpu = g |> gpu + d = degree(g) + d_gpu = degree(g_gpu) + @test d_gpu isa CuVector{Int} + @test Array(d_gpu) == d + end + end + + @testset "weighted" begin + # weighted degree + s = [1, 1, 2, 3] + t = [2, 2, 2, 4] + eweight = Float32[0.1, 2.1, 1.2, 1] + g = GNNGraph((s, t, eweight), graph_type = GRAPH_T) + @test degree(g) ≈ [2.2, 1.2, 1.0, 0.0] + d = degree(g, edge_weight = false) + if GRAPH_T == :coo + @test d == [2, 1, 1, 0] + else + # Adjacency matrix representation cannot disambiguate multiple edges + # and edge weights + @test d == [1, 1, 1, 0] + end + @test eltype(d) <: Integer + @test degree(g, edge_weight = 2 * eweight) ≈ [4.4, 2.4, 2.0, 0.0] broken = (GRAPH_T != :coo) + + if TEST_GPU + g_gpu = g |> gpu + d = degree(g) + d_gpu = degree(g_gpu) + @test d_gpu isa CuVector{Float32} + @test Array(d_gpu) ≈ d + end + @testset "gradient" begin + gw = gradient(eweight) do w + g = GNNGraph((s, t, w), graph_type = GRAPH_T) + sum(degree(g, edge_weight = false)) + end[1] + + @test gw === nothing + + gw = gradient(eweight) do w + g = GNNGraph((s, t, w), graph_type = GRAPH_T) + sum(degree(g, edge_weight = true)) + end[1] + + @test gw isa AbstractVector{Float32} + @test gw isa Vector{Float32} broken = (GRAPH_T == :sparse) + @test gw ≈ ones(Float32, length(gw)) + + gw = gradient(eweight) do w + g = GNNGraph((s, t, w), graph_type = GRAPH_T) + sum(degree(g, dir=:both, edge_weight=true)) + end[1] + + @test gw isa AbstractVector{Float32} + @test gw isa Vector{Float32} broken = (GRAPH_T == :sparse) + @test gw ≈ 2 * ones(Float32, length(gw)) + + grad = gradient(g) do g + sum(degree(g, edge_weight=false)) + end[1] + @test grad === nothing + + grad = gradient(g) do g + sum(degree(g, edge_weight=true)) + end[1] + + if GRAPH_T == :coo + @test grad.graph[3] isa Vector{Float32} + @test grad.graph[3] ≈ ones(Float32, length(gw)) + else + if GRAPH_T == :sparse + @test grad.graph isa AbstractSparseMatrix{Float32} + end + @test grad.graph isa AbstractMatrix{Float32} + + @test grad.graph ≈ [0.0 1.0 0.0 0.0 + 0.0 1.0 0.0 0.0 + 0.0 0.0 0.0 1.0 + 0.0 0.0 0.0 0.0] + end + + @testset "directed, degree dir=$dir" for dir in [:in, :out, :both] + g = rand_graph(10, 30, bidirected=false) + w = rand(Float32, 30) + s, t = edge_index(g) + + grad = gradient(w) do w + g = GNNGraph((s, t, w), graph_type = GRAPH_T) + sum(tanh.(degree(g; dir, edge_weight=true))) + end[1] + + ngrad = ngradient(w) do w + g = GNNGraph((s, t, w), graph_type = GRAPH_T) + sum(tanh.(degree(g; dir, edge_weight=true))) + end[1] + + @test grad ≈ ngrad + end + + @testset "heterognn, degree" begin + g = GNNHeteroGraph((:A, :to, :B) => ([1,1,2,3], [7,13,5,7])) + @test degree(g, (:A, :to, :B), dir = :out) == [2, 1, 1] + @test degree(g, (:A, :to, :B), dir = :in) == [0, 0, 0, 0, 1, 0, 2, 0, 0, 0, 0, 0, 1] + @test degree(g, (:A, :to, :B)) == [2, 1, 1] + end + end + end +end + +@testset "laplacian_matrix" begin + g = rand_graph(10, 30, graph_type = GRAPH_T) + A = adjacency_matrix(g) + D = Diagonal(vec(sum(A, dims = 2))) + L = laplacian_matrix(g) + @test eltype(L) == eltype(g) + @test L ≈ D - A +end + +@testset "laplacian_lambda_max" begin + s = [1, 2, 3, 4, 5, 1, 2, 3, 4, 5] + t = [2, 3, 4, 5, 1, 5, 1, 2, 3, 4] + g = GNNGraph(s, t) + @test laplacian_lambda_max(g) ≈ Float32(1.809017) + data1 = [g for i in 1:5] + gall1 = Flux.batch(data1) + @test laplacian_lambda_max(gall1) ≈ [Float32(1.809017) for i in 1:5] + data2 = [rand_graph(10, 20) for i in 1:3] + gall2 = Flux.batch(data2) + @test length(laplacian_lambda_max(gall2, add_self_loops=true)) == 3 +end + +@testset "adjacency_matrix" begin + a = sprand(5, 5, 0.5) + abin = map(x -> x > 0 ? 1 : 0, a) + + g = GNNGraph(a, graph_type = GRAPH_T) + A = adjacency_matrix(g, Float32) + @test A ≈ a + @test eltype(A) == Float32 + + Abin = adjacency_matrix(g, Float32, weighted = false) + @test Abin ≈ abin + @test eltype(Abin) == Float32 + + @testset "gradient" begin + s = [1, 2, 3] + t = [2, 3, 1] + w = [0.1, 0.1, 0.2] + gw = gradient(w) do w + g = GNNGraph(s, t, w, graph_type = GRAPH_T) + A = adjacency_matrix(g, weighted = false) + sum(A) + end[1] + @test gw === nothing + + gw = gradient(w) do w + g = GNNGraph(s, t, w, graph_type = GRAPH_T) + A = adjacency_matrix(g, weighted = true) + sum(A) + end[1] + + @test gw == [1, 1, 1] + end + + @testset "khop_adj" begin + s = [1, 2, 3] + t = [2, 3, 1] + w = [0.1, 0.1, 0.2] + g = GNNGraph(s, t, w) + @test khop_adj(g, 2) == adjacency_matrix(g) * adjacency_matrix(g) + @test khop_adj(g, 2, Int8; weighted = false) == sparse([0 0 1; 1 0 0; 0 1 0]) + @test khop_adj(g, 2, Int8; dir = in, weighted = false) == + sparse([0 0 1; 1 0 0; 0 1 0]') + @test khop_adj(g, 1) == adjacency_matrix(g) + @test eltype(khop_adj(g, 4)) == Float64 + @test eltype(khop_adj(g, 10, Float32)) == Float32 + end +end + +if GRAPH_T == :coo + @testset "HeteroGraph" begin + @testset "graph_indicator" begin + gs = [rand_heterograph(Dict(:user => 10, :movie => 20, :actor => 30), + Dict((:user,:like,:movie) => 10, + (:actor,:rate,:movie)=>20)) for _ in 1:3] + g = MLUtils.batch(gs) + @test graph_indicator(g) == Dict(:user => [repeat([1], 10); repeat([2], 10); repeat([3], 10)], + :movie => [repeat([1], 20); repeat([2], 20); repeat([3], 20)], + :actor => [repeat([1], 30); repeat([2], 30); repeat([3], 30)]) + @test graph_indicator(g, :movie) == [repeat([1], 20); repeat([2], 20); repeat([3], 20)] + end + end +end + diff --git a/test/GNNGraphs/sampling.jl b/test/GNNGraphs/sampling.jl new file mode 100644 index 000000000..5dfb63ab2 --- /dev/null +++ b/test/GNNGraphs/sampling.jl @@ -0,0 +1,46 @@ +@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 diff --git a/test/GNNGraphs/temporalsnapshotsgnngraph.jl b/test/GNNGraphs/temporalsnapshotsgnngraph.jl new file mode 100644 index 000000000..90ddeafbf --- /dev/null +++ b/test/GNNGraphs/temporalsnapshotsgnngraph.jl @@ -0,0 +1,117 @@ +@testset "Constructor array TemporalSnapshotsGNNGraph" begin + snapshots = [rand_graph(10, 20) for i in 1:5] + tsg = TemporalSnapshotsGNNGraph(snapshots) + @test tsg.num_nodes == [10 for i in 1:5] + @test tsg.num_edges == [20 for i in 1:5] + wrsnapshots = [rand_graph(10,20), rand_graph(12,22)] + @test_throws AssertionError TemporalSnapshotsGNNGraph(wrsnapshots) +end + +@testset "==" begin + snapshots = [rand_graph(10, 20) for i in 1:5] + tsg1 = TemporalSnapshotsGNNGraph(snapshots) + tsg2 = TemporalSnapshotsGNNGraph(snapshots) + @test tsg1 == tsg2 + tsg3 = TemporalSnapshotsGNNGraph(snapshots[1:3]) + @test tsg1 != tsg3 + @test tsg1 !== tsg3 +end + +@testset "getindex" begin + snapshots = [rand_graph(10, 20) for i in 1:5] + tsg = TemporalSnapshotsGNNGraph(snapshots) + @test tsg[3] == snapshots[3] + @test tsg[[1,2]] == TemporalSnapshotsGNNGraph([10,10], [20,20], 2, snapshots[1:2], tsg.tgdata) +end + +@testset "getproperty" begin + x = rand(10) + snapshots = [rand_graph(10, 20, ndata = x) for i in 1:5] + tsg = TemporalSnapshotsGNNGraph(snapshots) + @test tsg.tgdata == DataStore() + @test tsg.x == tsg.ndata.x == [x for i in 1:5] + @test_throws KeyError tsg.ndata.w + @test_throws ArgumentError tsg.w +end + +@testset "add/remove_snapshot" begin + snapshots = [rand_graph(10, 20) for i in 1:5] + tsg = TemporalSnapshotsGNNGraph(snapshots) + g = rand_graph(10, 20) + tsg = add_snapshot(tsg, 3, g) + @test tsg.num_nodes == [10 for i in 1:6] + @test tsg.num_edges == [20 for i in 1:6] + @test tsg.snapshots[3] == g + tsg = remove_snapshot(tsg, 3) + @test tsg.num_nodes == [10 for i in 1:5] + @test tsg.num_edges == [20 for i in 1:5] + @test tsg.snapshots == snapshots +end + +@testset "add/remove_snapshot" begin + snapshots = [rand_graph(10, 20) for i in 1:5] + tsg = TemporalSnapshotsGNNGraph(snapshots) + g = rand_graph(10, 20) + tsg2 = add_snapshot(tsg, 3, g) + @test tsg2.num_nodes == [10 for i in 1:6] + @test tsg2.num_edges == [20 for i in 1:6] + @test tsg2.snapshots[3] == g + @test tsg2.num_snapshots == 6 + @test tsg.num_nodes == [10 for i in 1:5] + @test tsg.num_edges == [20 for i in 1:5] + @test tsg.snapshots[2] === tsg2.snapshots[2] + @test tsg.snapshots[3] === tsg2.snapshots[4] + @test length(tsg.snapshots) == 5 + @test tsg.num_snapshots == 5 + + tsg21 = add_snapshot(tsg2, 7, g) + @test tsg21.num_snapshots == 7 + + tsg3 = remove_snapshot(tsg, 3) + @test tsg3.num_nodes == [10 for i in 1:4] + @test tsg3.num_edges == [20 for i in 1:4] + @test tsg3.snapshots == snapshots[[1,2,4,5]] +end + + +# @testset "add/remove_snapshot!" begin +# snapshots = [rand_graph(10, 20) for i in 1:5] +# tsg = TemporalSnapshotsGNNGraph(snapshots) +# g = rand_graph(10, 20) +# tsg2 = add_snapshot!(tsg, 3, g) +# @test tsg2.num_nodes == [10 for i in 1:6] +# @test tsg2.num_edges == [20 for i in 1:6] +# @test tsg2.snapshots[3] == g +# @test tsg2.num_snapshots == 6 +# @test tsg2 === tsg + +# tsg3 = remove_snapshot!(tsg, 3) +# @test tsg3.num_nodes == [10 for i in 1:4] +# @test tsg3.num_edges == [20 for i in 1:4] +# @test length(tsg3.snapshots) === 4 +# @test tsg3 === tsg +# end + +@testset "show" begin + snapshots = [rand_graph(10, 20) for i in 1:5] + tsg = TemporalSnapshotsGNNGraph(snapshots) + @test sprint(show,tsg) == "TemporalSnapshotsGNNGraph(5) with no data" + @test sprint(show, MIME("text/plain"), tsg; context=:compact => true) == "TemporalSnapshotsGNNGraph(5) with no data" + @test sprint(show, MIME("text/plain"), tsg; context=:compact => false) == "TemporalSnapshotsGNNGraph:\n num_nodes: [10, 10, 10, 10, 10]\n num_edges: [20, 20, 20, 20, 20]\n num_snapshots: 5" + tsg.tgdata.x=rand(4) + @test sprint(show,tsg) == "TemporalSnapshotsGNNGraph(5) with x: 4-element data" +end + +if TEST_GPU + @testset "gpu" begin + snapshots = [rand_graph(10, 20; ndata = rand(5,10)) for i in 1:5] + tsg = TemporalSnapshotsGNNGraph(snapshots) + tsg.tgdata.x = rand(5) + tsg = Flux.gpu(tsg) + @test tsg.snapshots[1].ndata.x isa CuArray + @test tsg.snapshots[end].ndata.x isa CuArray + @test tsg.tgdata.x isa CuArray + @test tsg.num_nodes isa CuArray + @test tsg.num_edges isa CuArray + end +end diff --git a/test/GNNGraphs/transform.jl b/test/GNNGraphs/transform.jl new file mode 100644 index 000000000..70570d155 --- /dev/null +++ b/test/GNNGraphs/transform.jl @@ -0,0 +1,626 @@ +@testset "add self-loops" begin + A = [1 1 0 0 + 0 0 1 0 + 0 0 0 1 + 1 0 0 0] + A2 = [2 1 0 0 + 0 1 1 0 + 0 0 1 1 + 1 0 0 1] + + g = GNNGraph(A; graph_type = GRAPH_T) + fg2 = add_self_loops(g) + @test adjacency_matrix(g) == A + @test g.num_edges == sum(A) + @test adjacency_matrix(fg2) == A2 + @test fg2.num_edges == sum(A2) +end + +@testset "batch" begin + g1 = GNNGraph(random_regular_graph(10, 2), ndata = rand(16, 10), + graph_type = GRAPH_T) + g2 = GNNGraph(random_regular_graph(4, 2), ndata = rand(16, 4), graph_type = GRAPH_T) + g3 = GNNGraph(random_regular_graph(7, 2), ndata = rand(16, 7), graph_type = GRAPH_T) + + g12 = Flux.batch([g1, g2]) + g12b = blockdiag(g1, g2) + @test g12 == g12b + + g123 = Flux.batch([g1, g2, g3]) + @test g123.graph_indicator == [fill(1, 10); fill(2, 4); fill(3, 7)] + + # Allow wider eltype + g123 = Flux.batch(GNNGraph[g1, g2, g3]) + @test g123.graph_indicator == [fill(1, 10); fill(2, 4); fill(3, 7)] + + + s, t = edge_index(g123) + @test s == [edge_index(g1)[1]; 10 .+ edge_index(g2)[1]; 14 .+ edge_index(g3)[1]] + @test t == [edge_index(g1)[2]; 10 .+ edge_index(g2)[2]; 14 .+ edge_index(g3)[2]] + @test node_features(g123)[:, 11:14] ≈ node_features(g2) + + # scalar graph features + g1 = GNNGraph(g1, gdata = rand()) + g2 = GNNGraph(g2, gdata = rand()) + g3 = GNNGraph(g3, gdata = rand()) + g123 = Flux.batch([g1, g2, g3]) + @test g123.gdata.u == [g1.gdata.u, g2.gdata.u, g3.gdata.u] + + # Batch of batches + g123123 = Flux.batch([g123, g123]) + @test g123123.graph_indicator == + [fill(1, 10); fill(2, 4); fill(3, 7); fill(4, 10); fill(5, 4); fill(6, 7)] + @test g123123.num_graphs == 6 +end + +@testset "unbatch" begin + g1 = rand_graph(10, 20, graph_type = GRAPH_T) + g2 = rand_graph(5, 10, graph_type = GRAPH_T) + g12 = Flux.batch([g1, g2]) + gs = Flux.unbatch([g1, g2]) + @test length(gs) == 2 + @test gs[1].num_nodes == 10 + @test gs[1].num_edges == 20 + @test gs[1].num_graphs == 1 + @test gs[2].num_nodes == 5 + @test gs[2].num_edges == 10 + @test gs[2].num_graphs == 1 +end + +@testset "batch/unbatch roundtrip" begin + n = 20 + c = 3 + ngraphs = 10 + gs = [rand_graph(n, c * n, ndata = rand(2, n), edata = rand(3, c * n), + graph_type = GRAPH_T) + for _ in 1:ngraphs] + gall = Flux.batch(gs) + gs2 = Flux.unbatch(gall) + @test gs2[1] == gs[1] + @test gs2[end] == gs[end] +end + +@testset "getgraph" begin + g1 = GNNGraph(random_regular_graph(10, 2), ndata = rand(16, 10), + graph_type = GRAPH_T) + g2 = GNNGraph(random_regular_graph(4, 2), ndata = rand(16, 4), graph_type = GRAPH_T) + g3 = GNNGraph(random_regular_graph(7, 2), ndata = rand(16, 7), graph_type = GRAPH_T) + g = Flux.batch([g1, g2, g3]) + + g2b, nodemap = getgraph(g, 2, nmap = true) + s, t = edge_index(g2b) + @test s == edge_index(g2)[1] + @test t == edge_index(g2)[2] + @test node_features(g2b) ≈ node_features(g2) + + g2c = getgraph(g, 2) + @test g2c isa GNNGraph{typeof(g.graph)} + + g1b, nodemap = getgraph(g1, 1, nmap = true) + @test g1b === g1 + @test nodemap == 1:(g1.num_nodes) +end + +@testset "remove_edges" begin + if GRAPH_T == :coo + s = [1, 1, 2, 3] + t = [2, 3, 4, 5] + w = [0.1, 0.2, 0.3, 0.4] + edata = ['a', 'b', 'c', 'd'] + g = GNNGraph(s, t, w, edata = edata, graph_type = GRAPH_T) + + # single edge removal + gnew = remove_edges(g, [1]) + new_s, new_t = edge_index(gnew) + @test gnew.num_edges == 3 + @test new_s == s[2:end] + @test new_t == t[2:end] + + # multiple edge removal + gnew = remove_edges(g, [1,2,4]) + new_s, new_t = edge_index(gnew) + new_w = get_edge_weight(gnew) + new_edata = gnew.edata.e + @test gnew.num_edges == 1 + @test new_s == [2] + @test new_t == [4] + @test new_w == [0.3] + @test new_edata == ['c'] + end +end + +@testset "add_edges" begin + if GRAPH_T == :coo + s = [1, 1, 2, 3] + t = [2, 3, 4, 5] + g = GNNGraph(s, t, graph_type = GRAPH_T) + snew = [1] + tnew = [4] + gnew = add_edges(g, snew, tnew) + @test gnew.num_edges == 5 + @test sort(inneighbors(gnew, 4)) == [1, 2] + + gnew2 = add_edges(g, (snew, tnew)) + @test gnew2 == gnew + @test get_edge_weight(gnew2) === nothing + + g = GNNGraph(s, t, edata = (e1 = rand(2, 4), e2 = rand(3, 4)), graph_type = GRAPH_T) + # @test_throws ErrorException add_edges(g, snew, tnew) + gnew = add_edges(g, snew, tnew, edata = (e1 = ones(2, 1), e2 = zeros(3, 1))) + @test all(gnew.edata.e1[:, 5] .== 1) + @test all(gnew.edata.e2[:, 5] .== 0) + + @testset "adding new nodes" begin + g = GNNGraph() + g = add_edges(g, ([1,3], [2, 1])) + @test g.num_nodes == 3 + @test g.num_edges == 2 + @test sort(inneighbors(g, 1)) == [3] + @test sort(outneighbors(g, 1)) == [2] + end + @testset "also add weights" begin + s = [1, 1, 2, 3] + t = [2, 3, 4, 5] + w = [1.0, 2.0, 3.0, 4.0] + snew = [1] + tnew = [4] + wnew = [5.] + + g = GNNGraph((s, t), graph_type = GRAPH_T) + gnew = add_edges(g, (snew, tnew, wnew)) + @test get_edge_weight(gnew) == [ones(length(s)); wnew] + + g = GNNGraph((s, t, w), graph_type = GRAPH_T) + gnew = add_edges(g, (snew, tnew, wnew)) + @test get_edge_weight(gnew) == [w; wnew] + end + end +end + +@testset "perturb_edges" begin if GRAPH_T == :coo + s, t = [1, 2, 3, 4, 5], [2, 3, 4, 5, 1] + g = GNNGraph((s, t)) + rng = MersenneTwister(42) + g_per = perturb_edges(g, 0.5, rng=rng) + @test g_per.num_edges == 8 +end end + +@testset "remove_nodes" begin if GRAPH_T == :coo + #single node + s = [1, 1, 2, 3] + t = [2, 3, 4, 5] + eweights = [0.1, 0.2, 0.3, 0.4] + ndata = [1.0, 2.0, 3.0, 4.0, 5.0] + edata = ['a', 'b', 'c', 'd'] + + g = GNNGraph(s, t, eweights, ndata = ndata, edata = edata, graph_type = GRAPH_T) + + gnew = remove_nodes(g, [1]) + + snew = [1, 2] + tnew = [3, 4] + eweights_new = [0.3, 0.4] + ndata_new = [2.0, 3.0, 4.0, 5.0] + edata_new = ['c', 'd'] + + stest, ttest = edge_index(gnew) + eweightstest = get_edge_weight(gnew) + ndatatest = gnew.ndata.x + edatatest = gnew.edata.e + + + @test gnew.num_edges == 2 + @test gnew.num_nodes == 4 + @test snew == stest + @test tnew == ttest + @test eweights_new == eweightstest + @test ndata_new == ndatatest + @test edata_new == edatatest + + # multiple nodes + s = [1, 5, 2, 3] + t = [2, 3, 4, 5] + eweights = [0.1, 0.2, 0.3, 0.4] + ndata = [1.0, 2.0, 3.0, 4.0, 5.0] + edata = ['a', 'b', 'c', 'd'] + + g = GNNGraph(s, t, eweights, ndata = ndata, edata = edata, graph_type = GRAPH_T) + + gnew = remove_nodes(g, [1,4]) + snew = [3,2] + tnew = [2,3] + eweights_new = [0.2,0.4] + ndata_new = [2.0,3.0,5.0] + edata_new = ['b','d'] + + stest, ttest = edge_index(gnew) + eweightstest = get_edge_weight(gnew) + ndatatest = gnew.ndata.x + edatatest = gnew.edata.e + + @test gnew.num_edges == 2 + @test gnew.num_nodes == 3 + @test snew == stest + @test tnew == ttest + @test eweights_new == eweightstest + @test ndata_new == ndatatest + @test edata_new == edatatest +end end + +@testset "drop_nodes" begin + if GRAPH_T == :coo + Random.seed!(42) + s = [1, 1, 2, 3] + t = [2, 3, 4, 5] + g = GNNGraph(s, t, graph_type = GRAPH_T) + + gnew = drop_nodes(g, Float32(0.5)) + @test gnew.num_nodes == 3 + + gnew = drop_nodes(g, Float32(1.0)) + @test gnew.num_nodes == 0 + + gnew = drop_nodes(g, Float32(0.0)) + @test gnew.num_nodes == 5 + end +end + +@testset "add_nodes" begin if GRAPH_T == :coo + g = rand_graph(6, 4, ndata = rand(2, 6), graph_type = GRAPH_T) + gnew = add_nodes(g, 5, ndata = ones(2, 5)) + @test gnew.num_nodes == g.num_nodes + 5 + @test gnew.num_edges == g.num_edges + @test gnew.num_graphs == g.num_graphs + @test all(gnew.ndata.x[:, 7:11] .== 1) +end end + +@testset "remove_self_loops" begin if GRAPH_T == :coo # add_edges and set_edge_weight only implemented for coo + g = rand_graph(10, 20, graph_type = GRAPH_T) + g1 = add_edges(g, [1:5;], [1:5;]) + @test g1.num_edges == g.num_edges + 5 + g2 = remove_self_loops(g1) + @test g2.num_edges == g.num_edges + @test sort_edge_index(edge_index(g2)) == sort_edge_index(edge_index(g)) + + # with edge features and weights + g1 = GNNGraph(g1, edata = (e1 = ones(3, g1.num_edges), e2 = 2 * ones(g1.num_edges))) + g1 = set_edge_weight(g1, 3 * ones(g1.num_edges)) + g2 = remove_self_loops(g1) + @test g2.num_edges == g.num_edges + @test sort_edge_index(edge_index(g2)) == sort_edge_index(edge_index(g)) + @test size(get_edge_weight(g2)) == (g2.num_edges,) + @test size(g2.edata.e1) == (3, g2.num_edges) + @test size(g2.edata.e2) == (g2.num_edges,) +end end + +@testset "remove_multi_edges" begin if GRAPH_T == :coo + g = rand_graph(10, 20, graph_type = GRAPH_T) + s, t = edge_index(g) + g1 = add_edges(g, s[1:5], t[1:5]) + @test g1.num_edges == g.num_edges + 5 + g2 = remove_multi_edges(g1, aggr = +) + @test g2.num_edges == g.num_edges + @test sort_edge_index(edge_index(g2)) == sort_edge_index(edge_index(g)) + + # Default aggregation is + + g1 = GNNGraph(g1, edata = (e1 = ones(3, g1.num_edges), e2 = 2 * ones(g1.num_edges))) + g1 = set_edge_weight(g1, 3 * ones(g1.num_edges)) + g2 = remove_multi_edges(g1) + @test g2.num_edges == g.num_edges + @test sort_edge_index(edge_index(g2)) == sort_edge_index(edge_index(g)) + @test count(g2.edata.e1[:, i] == 2 * ones(3) for i in 1:(g2.num_edges)) == 5 + @test count(g2.edata.e2[i] == 4 for i in 1:(g2.num_edges)) == 5 + w2 = get_edge_weight(g2) + @test count(w2[i] == 6 for i in 1:(g2.num_edges)) == 5 +end end + +@testset "negative_sample" begin if GRAPH_T == :coo + n, m = 10, 30 + g = rand_graph(n, m, bidirected = true, graph_type = GRAPH_T) + + # check bidirected=is_bidirected(g) default + gneg = negative_sample(g, num_neg_edges = 20) + @test gneg.num_nodes == g.num_nodes + @test gneg.num_edges == 20 + @test is_bidirected(gneg) + @test intersect(g, gneg).num_edges == 0 +end end + +@testset "rand_edge_split" begin if GRAPH_T == :coo + n, m = 100, 300 + + g = rand_graph(n, m, bidirected = true, graph_type = GRAPH_T) + # check bidirected=is_bidirected(g) default + g1, g2 = rand_edge_split(g, 0.9) + @test is_bidirected(g1) + @test is_bidirected(g2) + @test intersect(g1, g2).num_edges == 0 + @test g1.num_edges + g2.num_edges == g.num_edges + @test g2.num_edges < 50 + + g = rand_graph(n, m, bidirected = false, graph_type = GRAPH_T) + # check bidirected=is_bidirected(g) default + g1, g2 = rand_edge_split(g, 0.9) + @test !is_bidirected(g1) + @test !is_bidirected(g2) + @test intersect(g1, g2).num_edges == 0 + @test g1.num_edges + g2.num_edges == g.num_edges + @test g2.num_edges < 50 + + g1, g2 = rand_edge_split(g, 0.9, bidirected = false) + @test !is_bidirected(g1) + @test !is_bidirected(g2) + @test intersect(g1, g2).num_edges == 0 + @test g1.num_edges + g2.num_edges == g.num_edges + @test g2.num_edges < 50 +end end + +@testset "set_edge_weight" begin + g = rand_graph(10, 20, graph_type = GRAPH_T) + w = rand(20) + + gw = set_edge_weight(g, w) + @test get_edge_weight(gw) == w + + # now from weighted graph + s, t = edge_index(g) + g2 = GNNGraph(s, t, rand(20), graph_type = GRAPH_T) + gw2 = set_edge_weight(g2, w) + @test get_edge_weight(gw2) == w +end + +@testset "to_bidirected" begin if GRAPH_T == :coo + s, t = [1, 2, 3, 3, 4], [2, 3, 4, 4, 4] + w = [1.0, 2.0, 3.0, 4.0, 5.0] + e = [10.0, 20.0, 30.0, 40.0, 50.0] + g = GNNGraph(s, t, w, edata = e) + + g2 = to_bidirected(g) + @test g2.num_nodes == g.num_nodes + @test g2.num_edges == 7 + @test is_bidirected(g2) + @test !has_multi_edges(g2) + + s2, t2 = edge_index(g2) + w2 = get_edge_weight(g2) + @test s2 == [1, 2, 2, 3, 3, 4, 4] + @test t2 == [2, 1, 3, 2, 4, 3, 4] + @test w2 == [1, 1, 2, 2, 3.5, 3.5, 5] + @test g2.edata.e == [10.0, 10.0, 20.0, 20.0, 35.0, 35.0, 50.0] +end end + +@testset "to_unidirected" begin if GRAPH_T == :coo + s = [1, 2, 3, 4, 4] + t = [2, 3, 4, 3, 4] + w = [1.0, 2.0, 3.0, 4.0, 5.0] + e = [10.0, 20.0, 30.0, 40.0, 50.0] + g = GNNGraph(s, t, w, edata = e) + + g2 = to_unidirected(g) + @test g2.num_nodes == g.num_nodes + @test g2.num_edges == 4 + @test !has_multi_edges(g2) + + s2, t2 = edge_index(g2) + w2 = get_edge_weight(g2) + @test s2 == [1, 2, 3, 4] + @test t2 == [2, 3, 4, 4] + @test w2 == [1, 2, 3.5, 5] + @test g2.edata.e == [10.0, 20.0, 35.0, 50.0] +end end + +@testset "Graphs.Graph from GNNGraph" begin + g = rand_graph(10, 20, graph_type = GRAPH_T) + + G = Graphs.Graph(g) + @test nv(G) == g.num_nodes + @test ne(G) == g.num_edges ÷ 2 + + DG = Graphs.DiGraph(g) + @test nv(DG) == g.num_nodes + @test ne(DG) == g.num_edges +end + +@testset "random_walk_pe" begin + s = [1, 2, 2, 3] + t = [2, 1, 3, 2] + ndata = [-1, 0, 1] + g = GNNGraph(s, t, graph_type = GRAPH_T, ndata = ndata) + output = random_walk_pe(g, 3) + @test output == [0.0 0.0 0.0 + 0.5 1.0 0.5 + 0.0 0.0 0.0] +end + +@testset "HeteroGraphs" begin + @testset "batch" begin + gs = [rand_bipartite_heterograph((10, 15), 20) for _ in 1:5] + g = Flux.batch(gs) + @test g.num_nodes[:A] == 50 + @test g.num_nodes[:B] == 75 + @test g.num_edges[(:A,:to,:B)] == 100 + @test g.num_edges[(:B,:to,:A)] == 100 + @test g.num_graphs == 5 + @test g.graph_indicator == Dict(:A => vcat([fill(i, 10) for i in 1:5]...), + :B => vcat([fill(i, 15) for i in 1:5]...)) + + for gi in gs + gi.ndata[:A].x = ones(2, 10) + gi.ndata[:A].y = zeros(10) + gi.edata[(:A,:to,:B)].e = fill(2, 20) + gi.gdata.u = 7 + end + g = Flux.batch(gs) + @test g.ndata[:A].x == ones(2, 50) + @test g.ndata[:A].y == zeros(50) + @test g.edata[(:A,:to,:B)].e == fill(2, 100) + @test g.gdata.u == fill(7, 5) + + # Allow for wider eltype + g = Flux.batch(GNNHeteroGraph[g for g in gs]) + @test g.ndata[:A].x == ones(2, 50) + @test g.ndata[:A].y == zeros(50) + @test g.edata[(:A,:to,:B)].e == fill(2, 100) + @test g.gdata.u == fill(7, 5) + end + + @testset "batch non-similar edge types" begin + gs = [rand_heterograph((:A =>10, :B => 14), ((:A, :to1, :A) => 5, (:A, :to1, :B) => 20)), + rand_heterograph((:A => 10, :B => 15), ((:A, :to1, :B) => 5, (:B, :to2, :B) => 16)), + rand_heterograph((:B => 15, :C => 5), ((:C, :to1, :B) => 5, (:B, :to2, :C) => 21)), + rand_heterograph((:A => 10, :B => 10, :C => 10), ((:A, :to1, :C) => 5, (:A, :to1, :B) => 5)), + rand_heterograph((:C => 20), ((:C, :to3, :C) => 10)) + ] + g = Flux.batch(gs) + + @test g.num_nodes[:A] == 10 + 10 + 10 + @test g.num_nodes[:B] == 14 + 15 + 15 + 10 + @test g.num_nodes[:C] == 5 + 10 + 20 + @test g.num_edges[(:A,:to1,:A)] == 5 + @test g.num_edges[(:A,:to1,:B)] == 20 + 5 + 5 + @test g.num_edges[(:A,:to1,:C)] == 5 + + @test g.num_edges[(:B,:to2,:B)] == 16 + @test g.num_edges[(:B,:to2,:C)] == 21 + + @test g.num_edges[(:C,:to1,:B)] == 5 + @test g.num_edges[(:C,:to3,:C)] == 10 + @test length(keys(g.num_edges)) == 7 + @test g.num_graphs == 5 + + function ndata_if_key(g, key, subkey, value) + if haskey(g.ndata, key) + g.ndata[key][subkey] = reduce(hcat, fill(value, g.num_nodes[key])) + end + end + + function edata_if_key(g, key, subkey, value) + if haskey(g.edata, key) + g.edata[key][subkey] = reduce(hcat, fill(value, g.num_edges[key])) + end + end + + for gi in gs + ndata_if_key(gi, :A, :x, [0]) + ndata_if_key(gi, :A, :y, ones(2)) + ndata_if_key(gi, :B, :x, ones(3)) + ndata_if_key(gi, :C, :y, zeros(4)) + edata_if_key(gi, (:A,:to1,:B), :x, [0]) + gi.gdata.u = 7 + end + + g = Flux.batch(gs) + + @test g.ndata[:A].x == reduce(hcat, fill(0, 10 + 10 + 10)) + @test g.ndata[:A].y == ones(2, 10 + 10 + 10) + @test g.ndata[:B].x == ones(3, 14 + 15 + 15 + 10) + @test g.ndata[:C].y == zeros(4, 5 + 10 + 20) + + @test g.edata[(:A,:to1,:B)].x == reduce(hcat, fill(0, 20 + 5 + 5)) + + @test g.gdata.u == fill(7, 5) + + # Allow for wider eltype + g = Flux.batch(GNNHeteroGraph[g for g in gs]) + @test g.ndata[:A].x == reduce(hcat, fill(0, 10 + 10 + 10)) + @test g.ndata[:A].y == ones(2, 10 + 10 + 10) + @test g.ndata[:B].x == ones(3, 14 + 15 + 15 + 10) + @test g.ndata[:C].y == zeros(4, 5 + 10 + 20) + + @test g.edata[(:A,:to1,:B)].x == reduce(hcat, fill(0, 20 + 5 + 5)) + + @test g.gdata.u == fill(7, 5) + end + + @testset "add_edges" begin + hg = rand_bipartite_heterograph((2, 2), (4, 0), bidirected=false) + hg = add_edges(hg, (:B,:to,:A), [1, 1], [1,2]) + @test hg.num_edges == Dict((:A,:to,:B) => 4, (:B,:to,:A) => 2) + @test has_edge(hg, (:B,:to,:A), 1, 1) + @test has_edge(hg, (:B,:to,:A), 1, 2) + @test !has_edge(hg, (:B,:to,:A), 2, 1) + @test !has_edge(hg, (:B,:to,:A), 2, 2) + + @testset "new nodes" begin + hg = rand_bipartite_heterograph((2, 2), 3) + hg = add_edges(hg, (:C,:rel,:B) => ([1, 3], [1,2])) + @test hg.num_nodes == Dict(:A => 2, :B => 2, :C => 3) + @test hg.num_edges == Dict((:A,:to,:B) => 3, (:B,:to,:A) => 3, (:C,:rel,:B) => 2) + s, t = edge_index(hg, (:C,:rel,:B)) + @test s == [1, 3] + @test t == [1, 2] + + hg = add_edges(hg, (:D,:rel,:F) => ([1, 3], [1,2])) + @test hg.num_nodes == Dict(:A => 2, :B => 2, :C => 3, :D => 3, :F => 2) + @test hg.num_edges == Dict((:A,:to,:B) => 3, (:B,:to,:A) => 3, (:C,:rel,:B) => 2, (:D,:rel,:F) => 2) + s, t = edge_index(hg, (:D,:rel,:F)) + @test s == [1, 3] + @test t == [1, 2] + end + + @testset "also add weights" begin + hg = GNNHeteroGraph((:user, :rate, :movie) => ([1,1,2,3], [7,13,5,7], [0.1, 0.2, 0.3, 0.4])) + hgnew = add_edges(hg, (:user, :like, :actor) => ([1, 2], [3, 4], [0.5, 0.6])) + @test hgnew.num_nodes[:user] == 3 + @test hgnew.num_nodes[:movie] == 13 + @test hgnew.num_nodes[:actor] == 4 + @test hgnew.num_edges == Dict((:user, :rate, :movie) => 4, (:user, :like, :actor) => 2) + @test get_edge_weight(hgnew, (:user, :rate, :movie)) == [0.1, 0.2, 0.3, 0.4] + @test get_edge_weight(hgnew, (:user, :like, :actor)) == [0.5, 0.6] + + hgnew2 = add_edges(hgnew, (:user, :like, :actor) => ([6, 7], [8, 10], [0.7, 0.8])) + @test hgnew2.num_nodes[:user] == 7 + @test hgnew2.num_nodes[:movie] == 13 + @test hgnew2.num_nodes[:actor] == 10 + @test hgnew2.num_edges == Dict((:user, :rate, :movie) => 4, (:user, :like, :actor) => 4) + @test get_edge_weight(hgnew2, (:user, :rate, :movie)) == [0.1, 0.2, 0.3, 0.4] + @test get_edge_weight(hgnew2, (:user, :like, :actor)) == [0.5, 0.6, 0.7, 0.8] + end + end + + @testset "add self-loops heterographs" begin + g = rand_heterograph((:A =>10, :B => 14), ((:A, :to1, :A) => 5, (:A, :to1, :B) => 20)) + # Case in which haskey(g.graph, edge_t) passes + g = add_self_loops(g, (:A, :to1, :A)) + + @test g.num_edges[(:A, :to1, :A)] == 5 + 10 + @test g.num_edges[(:A, :to1, :B)] == 20 + # This test should not use length(keys(g.num_edges)) since that may be undefined behavior + @test sum(1 for k in keys(g.num_edges) if g.num_edges[k] != 0) == 2 + + # Case in which haskey(g.graph, edge_t) fails + g = add_self_loops(g, (:A, :to3, :A)) + + @test g.num_edges[(:A, :to1, :A)] == 5 + 10 + @test g.num_edges[(:A, :to1, :B)] == 20 + @test g.num_edges[(:A, :to3, :A)] == 10 + @test sum(1 for k in keys(g.num_edges) if g.num_edges[k] != 0) == 3 + + # Case with edge weights + g = GNNHeteroGraph(Dict((:A, :to1, :A) => ([1, 2, 3], [3, 2, 1], [2, 2, 2]), (:A, :to2, :B) => ([1, 4, 5], [1, 2, 3]))) + n = g.num_nodes[:A] + g = add_self_loops(g, (:A, :to1, :A)) + + @test g.graph[(:A, :to1, :A)][3] == vcat([2, 2, 2], fill(1, n)) + end +end + +@testset "ppr_diffusion" begin + if GRAPH_T == :coo + s = [1, 1, 2, 3] + t = [2, 3, 4, 5] + eweights = [0.1, 0.2, 0.3, 0.4] + + g = GNNGraph(s, t, eweights) + + g_new = ppr_diffusion(g) + w_new = get_edge_weight(g_new) + + check_ew = Float32[0.012749999 + 0.025499998 + 0.038249996 + 0.050999995] + + @test w_new ≈ check_ew + end +end \ No newline at end of file diff --git a/test/GNNGraphs/utils.jl b/test/GNNGraphs/utils.jl new file mode 100644 index 000000000..db65b6357 --- /dev/null +++ b/test/GNNGraphs/utils.jl @@ -0,0 +1,62 @@ +@testset "edge encoding/decoding" begin + # not is_bidirected + n = 5 + s = [1, 1, 2, 3, 3, 4, 5] + t = [1, 3, 1, 1, 2, 5, 5] + + # directed=true + idx, maxid = GNNGraphs.edge_encoding(s, t, n) + @test maxid == n^2 + @test idx == [1, 3, 6, 11, 12, 20, 25] + + sdec, tdec = GNNGraphs.edge_decoding(idx, n) + @test sdec == s + @test tdec == t + + n1, m1 = 10, 30 + g = rand_graph(n1, m1) + s1, t1 = edge_index(g) + idx, maxid = GNNGraphs.edge_encoding(s1, t1, n1) + sdec, tdec = GNNGraphs.edge_decoding(idx, n1) + @test sdec == s1 + @test tdec == t1 + + # directed=false + idx, maxid = GNNGraphs.edge_encoding(s, t, n, directed = false) + @test maxid == n * (n + 1) ÷ 2 + @test idx == [1, 3, 2, 3, 7, 14, 15] + + mask = s .> t + snew = copy(s) + tnew = copy(t) + snew[mask] .= t[mask] + tnew[mask] .= s[mask] + sdec, tdec = GNNGraphs.edge_decoding(idx, n, directed = false) + @test sdec == snew + @test tdec == tnew + + n1, m1 = 6, 8 + g = rand_graph(n1, m1) + s1, t1 = edge_index(g) + idx, maxid = GNNGraphs.edge_encoding(s1, t1, n1, directed = false) + sdec, tdec = GNNGraphs.edge_decoding(idx, n1, directed = false) + mask = s1 .> t1 + snew = copy(s1) + tnew = copy(t1) + snew[mask] .= t1[mask] + tnew[mask] .= s1[mask] + @test sdec == snew + @test tdec == tnew +end + +@testset "color_refinment" begin + g = rand_graph(10, 20, seed=17, graph_type = GRAPH_T) + x0 = ones(Int, 10) + x, ncolors, niters = color_refinement(g, x0) + @test ncolors == 8 + @test niters == 2 + @test x == [4, 5, 6, 7, 8, 5, 8, 9, 10, 11] + + x2, _, _ = color_refinement(g) + @test x2 == x +end \ No newline at end of file diff --git a/test/runtests.jl b/test/runtests.jl index cc7a93d0f..150ca8c53 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -26,6 +26,17 @@ ENV["DATADEPS_ALWAYS_ACCEPT"] = true # for MLDatasets include("test_utils.jl") tests = [ + "GNNGraphs/chainrules", + "GNNGraphs/datastore", + "GNNGraphs/gnngraph", + "GNNGraphs/convert", + "GNNGraphs/transform", + "GNNGraphs/operators", + "GNNGraphs/generate", + "GNNGraphs/query", + "GNNGraphs/sampling", + "GNNGraphs/gnnheterograph", + "GNNGraphs/temporalsnapshotsgnngraph", "utils", "msgpass", "layers/basic",