| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| |
| # pylint: skip-file |
| from __future__ import print_function |
| import numpy as np |
| import scipy as sp |
| import mxnet as mx |
| import random |
| import itertools |
| from numpy.testing import assert_allclose, assert_array_equal |
| from mxnet.test_utils import * |
| import unittest |
| |
| def check_uniform(out, num_hops, max_num_vertices): |
| sample_id = out[0] |
| sub_csr = out[1] |
| layer = out[2] |
| # check sample_id |
| assert (len(sample_id) == max_num_vertices+1) |
| num_vertices = sample_id[-1].asnumpy()[0] |
| # check sub_csr |
| sub_csr.check_format(full_check=True) |
| assert np.all((sub_csr.indptr[num_vertices:] == sub_csr.indptr[num_vertices]).asnumpy()) |
| # check layer |
| for data in layer[:num_vertices]: |
| assert(data <= num_hops) |
| |
| def check_non_uniform(out, num_hops, max_num_vertices): |
| sample_id = out[0] |
| sub_csr = out[1] |
| prob = out[2] |
| layer = out[3] |
| # check sample_id |
| assert (len(sample_id) == max_num_vertices+1) |
| num_vertices = sample_id[-1].asnumpy()[0] |
| # check sub_csr |
| sub_csr.check_format(full_check=True) |
| assert np.all((sub_csr.indptr[num_vertices:] == sub_csr.indptr[num_vertices]).asnumpy()) |
| # check prob |
| assert (len(prob) == max_num_vertices) |
| # check layer |
| for data in layer[:num_vertices]: |
| assert(data <= num_hops) |
| |
| def check_compact(csr, id_arr, num_nodes): |
| compact = mx.nd.contrib.dgl_graph_compact(csr, id_arr, graph_sizes=num_nodes, return_mapping=False) |
| assert compact.shape[0] == num_nodes |
| assert compact.shape[1] == num_nodes |
| assert mx.nd.sum(compact.indptr == csr.indptr[0:int(num_nodes + 1)]).asnumpy() == num_nodes + 1 |
| sub_indices = compact.indices.asnumpy() |
| indices = csr.indices.asnumpy() |
| id_arr = id_arr.asnumpy() |
| for i in range(len(sub_indices)): |
| sub_id = sub_indices[i] |
| assert id_arr[sub_id] == indices[i] |
| |
| def test_uniform_sample(): |
| shape = (5, 5) |
| data_np = np.array([1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20], dtype=np.int64) |
| indices_np = np.array([1,2,3,4,0,2,3,4,0,1,3,4,0,1,2,4,0,1,2,3], dtype=np.int64) |
| indptr_np = np.array([0,4,8,12,16,20], dtype=np.int64) |
| a = mx.nd.sparse.csr_matrix((data_np, indices_np, indptr_np), shape=shape) |
| |
| seed = mx.nd.array([0,1,2,3,4], dtype=np.int64) |
| out = mx.nd.contrib.dgl_csr_neighbor_uniform_sample(a, seed, num_args=2, num_hops=1, num_neighbor=2, max_num_vertices=5) |
| assert (len(out) == 3) |
| check_uniform(out, num_hops=1, max_num_vertices=5) |
| num_nodes = out[0][-1].asnumpy() |
| assert num_nodes > 0 |
| assert num_nodes < len(out[0]) |
| check_compact(out[1], out[0], num_nodes) |
| |
| seed = mx.nd.array([0], dtype=np.int64) |
| out = mx.nd.contrib.dgl_csr_neighbor_uniform_sample(a, seed, num_args=2, num_hops=1, num_neighbor=1, max_num_vertices=4) |
| assert (len(out) == 3) |
| check_uniform(out, num_hops=1, max_num_vertices=4) |
| num_nodes = out[0][-1].asnumpy() |
| assert num_nodes > 0 |
| assert num_nodes < len(out[0]) |
| check_compact(out[1], out[0], num_nodes) |
| |
| seed = mx.nd.array([0], dtype=np.int64) |
| out = mx.nd.contrib.dgl_csr_neighbor_uniform_sample(a, seed, num_args=2, num_hops=2, num_neighbor=1, max_num_vertices=3) |
| assert (len(out) == 3) |
| check_uniform(out, num_hops=2, max_num_vertices=3) |
| num_nodes = out[0][-1].asnumpy() |
| assert num_nodes > 0 |
| assert num_nodes < len(out[0]) |
| check_compact(out[1], out[0], num_nodes) |
| |
| seed = mx.nd.array([0,2,4], dtype=np.int64) |
| out = mx.nd.contrib.dgl_csr_neighbor_uniform_sample(a, seed, num_args=2, num_hops=1, num_neighbor=2, max_num_vertices=5) |
| assert (len(out) == 3) |
| check_uniform(out, num_hops=1, max_num_vertices=5) |
| num_nodes = out[0][-1].asnumpy() |
| assert num_nodes > 0 |
| assert num_nodes < len(out[0]) |
| check_compact(out[1], out[0], num_nodes) |
| |
| seed = mx.nd.array([0,4], dtype=np.int64) |
| out = mx.nd.contrib.dgl_csr_neighbor_uniform_sample(a, seed, num_args=2, num_hops=1, num_neighbor=2, max_num_vertices=5) |
| assert (len(out) == 3) |
| check_uniform(out, num_hops=1, max_num_vertices=5) |
| num_nodes = out[0][-1].asnumpy() |
| assert num_nodes > 0 |
| assert num_nodes < len(out[0]) |
| check_compact(out[1], out[0], num_nodes) |
| |
| seed = mx.nd.array([0,4], dtype=np.int64) |
| out = mx.nd.contrib.dgl_csr_neighbor_uniform_sample(a, seed, num_args=2, num_hops=2, num_neighbor=2, max_num_vertices=5) |
| assert (len(out) == 3) |
| check_uniform(out, num_hops=2, max_num_vertices=5) |
| num_nodes = out[0][-1].asnumpy() |
| assert num_nodes > 0 |
| assert num_nodes < len(out[0]) |
| check_compact(out[1], out[0], num_nodes) |
| |
| seed = mx.nd.array([0,4], dtype=np.int64) |
| out = mx.nd.contrib.dgl_csr_neighbor_uniform_sample(a, seed, num_args=2, num_hops=1, num_neighbor=2, max_num_vertices=5) |
| assert (len(out) == 3) |
| check_uniform(out, num_hops=1, max_num_vertices=5) |
| num_nodes = out[0][-1].asnumpy() |
| assert num_nodes > 0 |
| assert num_nodes < len(out[0]) |
| check_compact(out[1], out[0], num_nodes) |
| |
| def test_non_uniform_sample(): |
| shape = (5, 5) |
| prob = mx.nd.array([0.9, 0.8, 0.2, 0.4, 0.1], dtype=np.float32) |
| data_np = np.array([1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20], dtype=np.int64) |
| indices_np = np.array([1,2,3,4,0,2,3,4,0,1,3,4,0,1,2,4,0,1,2,3], dtype=np.int64) |
| indptr_np = np.array([0,4,8,12,16,20], dtype=np.int64) |
| a = mx.nd.sparse.csr_matrix((data_np, indices_np, indptr_np), shape=shape) |
| |
| seed = mx.nd.array([0,1,2,3,4], dtype=np.int64) |
| out = mx.nd.contrib.dgl_csr_neighbor_non_uniform_sample(a, prob, seed, num_args=3, num_hops=1, num_neighbor=2, max_num_vertices=5) |
| assert (len(out) == 4) |
| check_non_uniform(out, num_hops=1, max_num_vertices=5) |
| |
| seed = mx.nd.array([0], dtype=np.int64) |
| out = mx.nd.contrib.dgl_csr_neighbor_non_uniform_sample(a, prob, seed, num_args=3, num_hops=1, num_neighbor=1, max_num_vertices=4) |
| assert (len(out) == 4) |
| check_non_uniform(out, num_hops=1, max_num_vertices=4) |
| |
| seed = mx.nd.array([0], dtype=np.int64) |
| out = mx.nd.contrib.dgl_csr_neighbor_non_uniform_sample(a, prob, seed, num_args=3, num_hops=2, num_neighbor=1, max_num_vertices=4) |
| assert (len(out) == 4) |
| check_non_uniform(out, num_hops=2, max_num_vertices=4) |
| |
| seed = mx.nd.array([0,2,4], dtype=np.int64) |
| out = mx.nd.contrib.dgl_csr_neighbor_non_uniform_sample(a, prob, seed, num_args=3, num_hops=1, num_neighbor=2, max_num_vertices=5) |
| assert (len(out) == 4) |
| check_non_uniform(out, num_hops=1, max_num_vertices=5) |
| |
| seed = mx.nd.array([0,4], dtype=np.int64) |
| out = mx.nd.contrib.dgl_csr_neighbor_non_uniform_sample(a, prob, seed, num_args=3, num_hops=1, num_neighbor=2, max_num_vertices=5) |
| assert (len(out) == 4) |
| check_non_uniform(out, num_hops=1, max_num_vertices=5) |
| |
| seed = mx.nd.array([0,4], dtype=np.int64) |
| out = mx.nd.contrib.dgl_csr_neighbor_non_uniform_sample(a, prob, seed, num_args=3, num_hops=2, num_neighbor=2, max_num_vertices=5) |
| assert (len(out) == 4) |
| check_non_uniform(out, num_hops=2, max_num_vertices=5) |
| |
| seed = mx.nd.array([0,4], dtype=np.int64) |
| out = mx.nd.contrib.dgl_csr_neighbor_non_uniform_sample(a, prob, seed, num_args=3, num_hops=1, num_neighbor=2, max_num_vertices=5) |
| assert (len(out) == 4) |
| check_non_uniform(out, num_hops=1, max_num_vertices=5) |
| |
| def test_edge_id(): |
| shape = rand_shape_2d() |
| data = rand_ndarray(shape, stype='csr', density=0.4) |
| ground_truth = np.zeros(shape, dtype=np.float32) |
| ground_truth -= 1.0 |
| indptr_np = data.indptr.asnumpy() |
| data_np = data.data.asnumpy() |
| indices_np = data.indices.asnumpy() |
| for i in range(shape[0]): |
| for j in range(indptr_np[i], indptr_np[i+1]): |
| idx = indices_np[j] |
| ground_truth[i, idx] = data_np[j] |
| |
| np_u = np.random.randint(0, shape[0], size=(5, )) |
| np_v = np.random.randint(0, shape[1], size=(5, )) |
| mx_u = mx.nd.array(np_u) |
| mx_v = mx.nd.array(np_v) |
| assert_almost_equal(mx.nd.contrib.edge_id(data, mx_u, mx_v).asnumpy(), |
| ground_truth[np_u, np_v], rtol=1e-5, atol=1e-6) |
| |
| def generate_graph(n): |
| arr = sp.sparse.random(n, n, density=0.2, format='coo') |
| arr.data = np.arange(0, len(arr.row), dtype=np.float32) |
| return arr.tocsr(), mx.nd.sparse.csr_matrix(arr.tocsr()).astype(np.int64) |
| |
| def test_subgraph(): |
| sp_g, g = generate_graph(100) |
| vertices = np.unique(np.random.randint(0, 100, size=(20))) |
| subgs = mx.nd.contrib.dgl_subgraph(g, mx.nd.array(vertices, dtype=np.int64), |
| return_mapping=True) |
| subgs[0].check_format() |
| subgs[1].check_format() |
| assert_array_equal(subgs[0].indptr, subgs[1].indptr) |
| assert_array_equal(subgs[0].indices, subgs[1].indices) |
| sp_subg = subgs[1].asscipy() |
| for i in range(len(subgs[0].indptr) - 1): |
| subv1 = i |
| v1 = vertices[subv1] |
| row_start = int(subgs[0].indptr[subv1].asnumpy()[0]) |
| row_end = int(subgs[0].indptr[subv1 + 1].asnumpy()[0]) |
| if row_start >= len(subgs[0].indices): |
| remain = subgs[0].indptr[subv1:].asnumpy() |
| assert np.sum(remain == row_start) == len(remain) |
| break |
| row = subgs[0].indices[row_start:row_end] |
| for _, subv2 in enumerate(row.asnumpy()): |
| v2 = vertices[subv2] |
| assert sp_g[v1, v2] == sp_subg[subv1, subv2] |
| |
| def test_adjacency(): |
| sp_g, g = generate_graph(100) |
| adj = mx.nd.contrib.dgl_adjacency(g) |
| assert adj.dtype == np.float32 |
| assert adj.shape == g.shape |
| assert_array_equal(adj.indptr, g.indptr) |
| assert_array_equal(adj.indices, g.indices) |
| assert_array_equal(adj.data, mx.nd.ones(shape=g.indices.shape)) |
| |
