test_traversal.py 4.01 KB
Newer Older
GaiYu0's avatar
GaiYu0 committed
1
2
3
import random
import sys
import time
4
import unittest
GaiYu0's avatar
GaiYu0 committed
5
6
7
8
9

import dgl
import networkx as nx
import numpy as np
import scipy.sparse as sp
10
import backend as F
GaiYu0's avatar
GaiYu0 committed
11

Gan Quan's avatar
Gan Quan committed
12
import itertools
nv-dlasalle's avatar
nv-dlasalle committed
13
from test_utils import parametrize_idtype
Gan Quan's avatar
Gan Quan committed
14

GaiYu0's avatar
GaiYu0 committed
15
16
np.random.seed(42)

Gan Quan's avatar
Gan Quan committed
17
def toset(x):
VoVAllen's avatar
VoVAllen committed
18
    # F.zerocopy_to_numpy may return a int
19
    return set(F.zerocopy_to_numpy(x).tolist())
Gan Quan's avatar
Gan Quan committed
20

nv-dlasalle's avatar
nv-dlasalle committed
21
@parametrize_idtype
22
def test_bfs(idtype, n=100):
Minjie Wang's avatar
Minjie Wang committed
23
24
25
26
27
28
29
30
31
    def _bfs_nx(g_nx, src):
        edges = nx.bfs_edges(g_nx, src)
        layers_nx = [set([src])]
        edges_nx = []
        frontier = set()
        edge_frontier = set()
        for u, v in edges:
            if u in layers_nx[-1]:
                frontier.add(v)
32
                edge_frontier.add(g.edge_ids(int(u), int(v)))
Minjie Wang's avatar
Minjie Wang committed
33
34
35
36
            else:
                layers_nx.append(frontier)
                edges_nx.append(edge_frontier)
                frontier = set([v])
37
                edge_frontier = set([g.edge_ids(u, v)])
38
39
40
41
        # avoids empty successors
        if len(frontier) > 0 and len(edge_frontier) > 0:
            layers_nx.append(frontier)
            edges_nx.append(edge_frontier)
Minjie Wang's avatar
Minjie Wang committed
42
        return layers_nx, edges_nx
GaiYu0's avatar
GaiYu0 committed
43

44
    a = sp.random(n, n, 3 / n, data_rvs=lambda n: np.ones(n))
45
    g = dgl.from_scipy(a).astype(idtype)
46

Minjie Wang's avatar
Minjie Wang committed
47
    g_nx = g.to_networkx()
GaiYu0's avatar
GaiYu0 committed
48
    src = random.choice(range(n))
Minjie Wang's avatar
Minjie Wang committed
49
    layers_nx, _ = _bfs_nx(g_nx, src)
50
    layers_dgl = dgl.bfs_nodes_generator(g, src)
GaiYu0's avatar
GaiYu0 committed
51
52
53
    assert len(layers_dgl) == len(layers_nx)
    assert all(toset(x) == y for x, y in zip(layers_dgl, layers_nx))

Minjie Wang's avatar
Minjie Wang committed
54
    g_nx = nx.random_tree(n, seed=42)
55
    g = dgl.from_networkx(g_nx).astype(idtype)
Minjie Wang's avatar
Minjie Wang committed
56
57
    src = 0
    _, edges_nx = _bfs_nx(g_nx, src)
Gan Quan's avatar
Gan Quan committed
58
59
60
61
    edges_dgl = dgl.bfs_edges_generator(g, src)
    assert len(edges_dgl) == len(edges_nx)
    assert all(toset(x) == y for x, y in zip(edges_dgl, edges_nx))

nv-dlasalle's avatar
nv-dlasalle committed
62
@parametrize_idtype
63
def test_topological_nodes(idtype, n=100):
64
    a = sp.random(n, n, 3 / n, data_rvs=lambda n: np.ones(n))
GaiYu0's avatar
GaiYu0 committed
65
    b = sp.tril(a, -1).tocoo()
66
    g = dgl.from_scipy(b).astype(idtype)
GaiYu0's avatar
GaiYu0 committed
67
68
69

    layers_dgl = dgl.topological_nodes_generator(g)

70
    adjmat = g.adjacency_matrix(transpose=True)
GaiYu0's avatar
GaiYu0 committed
71
72
    def tensor_topo_traverse():
        n = g.number_of_nodes()
73
        mask = F.copy_to(F.ones((n, 1)), F.cpu())
74
75
76
        degree = F.spmm(adjmat, mask)
        while F.reduce_sum(mask) != 0.:
            v = F.astype((degree == 0.), F.float32)
GaiYu0's avatar
GaiYu0 committed
77
78
            v = v * mask
            mask = mask - v
79
            frontier = F.copy_to(F.nonzero_1d(F.squeeze(v, 1)), F.cpu())
GaiYu0's avatar
GaiYu0 committed
80
            yield frontier
81
            degree -= F.spmm(adjmat, v)
GaiYu0's avatar
GaiYu0 committed
82
83
84
85
86
87
88

    layers_spmv = list(tensor_topo_traverse())

    assert len(layers_dgl) == len(layers_spmv)
    assert all(toset(x) == toset(y) for x, y in zip(layers_dgl, layers_spmv))

DFS_LABEL_NAMES = ['forward', 'reverse', 'nontree']
nv-dlasalle's avatar
nv-dlasalle committed
89
@parametrize_idtype
90
91
def test_dfs_labeled_edges(idtype, example=False):
    dgl_g = dgl.DGLGraph().astype(idtype)
Gan Quan's avatar
Gan Quan committed
92
93
    dgl_g.add_nodes(6)
    dgl_g.add_edges([0, 1, 0, 3, 3], [1, 2, 2, 4, 5])
GaiYu0's avatar
GaiYu0 committed
94
    dgl_edges, dgl_labels = dgl.dfs_labeled_edges_generator(
Gan Quan's avatar
Gan Quan committed
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
            dgl_g, [0, 3], has_reverse_edge=True, has_nontree_edge=True)
    dgl_edges = [toset(t) for t in dgl_edges]
    dgl_labels = [toset(t) for t in dgl_labels]
    g1_solutions = [
            # edges           labels
            [[0, 1, 1, 0, 2], [0, 0, 1, 1, 2]],
            [[2, 2, 0, 1, 0], [0, 1, 0, 2, 1]],
    ]
    g2_solutions = [
            # edges        labels
            [[3, 3, 4, 4], [0, 1, 0, 1]],
            [[4, 4, 3, 3], [0, 1, 0, 1]],
    ]

    def combine_frontiers(sol):
        es, ls = zip(*sol)
        es = [set(i for i in t if i is not None)
              for t in itertools.zip_longest(*es)]
        ls = [set(i for i in t if i is not None)
              for t in itertools.zip_longest(*ls)]
        return es, ls

    for sol_set in itertools.product(g1_solutions, g2_solutions):
        es, ls = combine_frontiers(sol_set)
        if es == dgl_edges and ls == dgl_labels:
            break
    else:
        assert False

GaiYu0's avatar
GaiYu0 committed
124
if __name__ == '__main__':
125
126
127
    test_bfs(idtype='int32')
    test_topological_nodes(idtype='int32')
    test_dfs_labeled_edges(idtype='int32')