test_transform.py 119 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
##
#   Copyright 2019-2021 Contributors
#
#   Licensed 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.
#

17
from scipy import sparse as spsp
18
19
import networkx as nx
import numpy as np
20
import os
21
22
import dgl
import dgl.function as fn
23
import dgl.partition
24
import backend as F
25
import unittest
26
import math
27
28
import pytest
from test_utils.graph_cases import get_cases
nv-dlasalle's avatar
nv-dlasalle committed
29
from test_utils import parametrize_idtype
30

31
32
from test_heterograph import create_test_heterograph3, create_test_heterograph4, \
    create_test_heterograph5
33

34
35
D = 5

36

37
# line graph related
38

39
def test_line_graph1():
40
    N = 5
41
    G = dgl.DGLGraph(nx.star_graph(N)).to(F.ctx())
42
    G.edata['h'] = F.randn((2 * N, D))
43
44
    L = G.line_graph(shared=True)
    assert L.number_of_nodes() == 2 * N
45
    assert F.allclose(L.ndata['h'], G.edata['h'])
46
    assert G.device == F.ctx()
47

48

nv-dlasalle's avatar
nv-dlasalle committed
49
@parametrize_idtype
50
def test_line_graph2(idtype):
51
    g = dgl.heterograph({
52
        ('user', 'follows', 'user'): ([0, 1, 1, 2, 2], [2, 0, 2, 0, 1])
53
    }, idtype=idtype)
54
    lg = dgl.line_graph(g)
55
56
57
58
59
60
61
62
    assert lg.number_of_nodes() == 5
    assert lg.number_of_edges() == 8
    row, col = lg.edges()
    assert np.array_equal(F.asnumpy(row),
                          np.array([0, 0, 1, 2, 2, 3, 4, 4]))
    assert np.array_equal(F.asnumpy(col),
                          np.array([3, 4, 0, 3, 4, 0, 1, 2]))

63
    lg = dgl.line_graph(g, backtracking=False)
64
65
66
67
68
69
70
    assert lg.number_of_nodes() == 5
    assert lg.number_of_edges() == 4
    row, col = lg.edges()
    assert np.array_equal(F.asnumpy(row),
                          np.array([0, 1, 2, 4]))
    assert np.array_equal(F.asnumpy(col),
                          np.array([4, 0, 3, 1]))
71
    g = dgl.heterograph({
72
        ('user', 'follows', 'user'): ([0, 1, 1, 2, 2], [2, 0, 2, 0, 1])
73
    }, idtype=idtype).formats('csr')
74
    lg = dgl.line_graph(g)
75
76
77
78
79
80
81
82
    assert lg.number_of_nodes() == 5
    assert lg.number_of_edges() == 8
    row, col = lg.edges()
    assert np.array_equal(F.asnumpy(row),
                          np.array([0, 0, 1, 2, 2, 3, 4, 4]))
    assert np.array_equal(F.asnumpy(col),
                          np.array([3, 4, 0, 3, 4, 0, 1, 2]))

83
    g = dgl.heterograph({
84
        ('user', 'follows', 'user'): ([0, 1, 1, 2, 2], [2, 0, 2, 0, 1])
85
    }, idtype=idtype).formats('csc')
86
    lg = dgl.line_graph(g)
87
88
89
90
91
92
93
94
95
96
97
    assert lg.number_of_nodes() == 5
    assert lg.number_of_edges() == 8
    row, col, eid = lg.edges('all')
    row = F.asnumpy(row)
    col = F.asnumpy(col)
    eid = F.asnumpy(eid).astype(int)
    order = np.argsort(eid)
    assert np.array_equal(row[order],
                          np.array([0, 0, 1, 2, 2, 3, 4, 4]))
    assert np.array_equal(col[order],
                          np.array([3, 4, 0, 3, 4, 0, 1, 2]))
98

99

100
101
102
103
104
105
def test_no_backtracking():
    N = 5
    G = dgl.DGLGraph(nx.star_graph(N))
    L = G.line_graph(backtracking=False)
    assert L.number_of_nodes() == 2 * N
    for i in range(1, N):
106
107
108
109
        e1 = G.edge_ids(0, i)
        e2 = G.edge_ids(i, 0)
        assert not L.has_edges_between(e1, e2)
        assert not L.has_edges_between(e2, e1)
110

111

112
# reverse graph related
nv-dlasalle's avatar
nv-dlasalle committed
113
@parametrize_idtype
114
def test_reverse(idtype):
115
    g = dgl.DGLGraph()
116
    g = g.astype(idtype).to(F.ctx())
117
118
119
    g.add_nodes(5)
    # The graph need not to be completely connected.
    g.add_edges([0, 1, 2], [1, 2, 1])
120
121
    g.ndata['h'] = F.tensor([[0.], [1.], [2.], [3.], [4.]])
    g.edata['h'] = F.tensor([[5.], [6.], [7.]])
122
123
124
125
126
127
    rg = g.reverse()

    assert g.is_multigraph == rg.is_multigraph

    assert g.number_of_nodes() == rg.number_of_nodes()
    assert g.number_of_edges() == rg.number_of_edges()
128
129
    assert F.allclose(F.astype(rg.has_edges_between(
        [1, 2, 1], [0, 1, 2]), F.float32), F.ones((3,)))
130
131
132
    assert g.edge_ids(0, 1) == rg.edge_ids(1, 0)
    assert g.edge_ids(1, 2) == rg.edge_ids(2, 1)
    assert g.edge_ids(2, 1) == rg.edge_ids(1, 2)
133

134
    # test dgl.reverse
135
136
137
138
    # test homogeneous graph
    g = dgl.graph((F.tensor([0, 1, 2]), F.tensor([1, 2, 0])))
    g.ndata['h'] = F.tensor([[0.], [1.], [2.]])
    g.edata['h'] = F.tensor([[3.], [4.], [5.]])
139
    g_r = dgl.reverse(g)
140
141
142
143
144
145
146
147
148
149
150
    assert g.number_of_nodes() == g_r.number_of_nodes()
    assert g.number_of_edges() == g_r.number_of_edges()
    u_g, v_g, eids_g = g.all_edges(form='all')
    u_rg, v_rg, eids_rg = g_r.all_edges(form='all')
    assert F.array_equal(u_g, v_rg)
    assert F.array_equal(v_g, u_rg)
    assert F.array_equal(eids_g, eids_rg)
    assert F.array_equal(g.ndata['h'], g_r.ndata['h'])
    assert len(g_r.edata) == 0

    # without share ndata
151
    g_r = dgl.reverse(g, copy_ndata=False)
152
153
154
155
156
157
    assert g.number_of_nodes() == g_r.number_of_nodes()
    assert g.number_of_edges() == g_r.number_of_edges()
    assert len(g_r.ndata) == 0
    assert len(g_r.edata) == 0

    # with share ndata and edata
158
    g_r = dgl.reverse(g, copy_ndata=True, copy_edata=True)
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
    assert g.number_of_nodes() == g_r.number_of_nodes()
    assert g.number_of_edges() == g_r.number_of_edges()
    assert F.array_equal(g.ndata['h'], g_r.ndata['h'])
    assert F.array_equal(g.edata['h'], g_r.edata['h'])

    # add new node feature to g_r
    g_r.ndata['hh'] = F.tensor([0, 1, 2])
    assert ('hh' in g.ndata) is False
    assert ('hh' in g_r.ndata) is True

    # add new edge feature to g_r
    g_r.edata['hh'] = F.tensor([0, 1, 2])
    assert ('hh' in g.edata) is False
    assert ('hh' in g_r.edata) is True

    # test heterogeneous graph
    g = dgl.heterograph({
176
177
178
179
        ('user', 'follows', 'user'): (
            [0, 1, 2, 4, 3, 1, 3], [1, 2, 3, 2, 0, 0, 1]),
        ('user', 'plays', 'game'): (
            [0, 0, 2, 3, 3, 4, 1], [1, 0, 1, 0, 1, 0, 0]),
180
181
        ('developer', 'develops', 'game'): ([0, 1, 1, 2], [0, 0, 1, 1])},
        idtype=idtype, device=F.ctx())
182
183
184
    g.nodes['user'].data['h'] = F.tensor([0, 1, 2, 3, 4])
    g.nodes['user'].data['hh'] = F.tensor([1, 1, 1, 1, 1])
    g.nodes['game'].data['h'] = F.tensor([0, 1])
185
    g.edges['follows'].data['h'] = F.tensor([0, 1, 2, 4, 3, 1, 3])
186
    g.edges['follows'].data['hh'] = F.tensor([1, 2, 3, 2, 0, 0, 1])
187
    g_r = dgl.reverse(g)
188
189
190
191
192
193
194
195
196

    for etype_g, etype_gr in zip(g.canonical_etypes, g_r.canonical_etypes):
        assert etype_g[0] == etype_gr[2]
        assert etype_g[1] == etype_gr[1]
        assert etype_g[2] == etype_gr[0]
        assert g.number_of_edges(etype_g) == g_r.number_of_edges(etype_gr)
    for ntype in g.ntypes:
        assert g.number_of_nodes(ntype) == g_r.number_of_nodes(ntype)
    assert F.array_equal(g.nodes['user'].data['h'], g_r.nodes['user'].data['h'])
197
198
    assert F.array_equal(g.nodes['user'].data['hh'],
                         g_r.nodes['user'].data['hh'])
199
200
    assert F.array_equal(g.nodes['game'].data['h'], g_r.nodes['game'].data['h'])
    assert len(g_r.edges['follows'].data) == 0
201
202
203
204
    u_g, v_g, eids_g = g.all_edges(form='all',
                                   etype=('user', 'follows', 'user'))
    u_rg, v_rg, eids_rg = g_r.all_edges(form='all',
                                        etype=('user', 'follows', 'user'))
205
206
207
208
    assert F.array_equal(u_g, v_rg)
    assert F.array_equal(v_g, u_rg)
    assert F.array_equal(eids_g, eids_rg)
    u_g, v_g, eids_g = g.all_edges(form='all', etype=('user', 'plays', 'game'))
209
210
    u_rg, v_rg, eids_rg = g_r.all_edges(form='all',
                                        etype=('game', 'plays', 'user'))
211
212
213
    assert F.array_equal(u_g, v_rg)
    assert F.array_equal(v_g, u_rg)
    assert F.array_equal(eids_g, eids_rg)
214
215
216
217
    u_g, v_g, eids_g = g.all_edges(form='all',
                                   etype=('developer', 'develops', 'game'))
    u_rg, v_rg, eids_rg = g_r.all_edges(form='all',
                                        etype=('game', 'develops', 'developer'))
218
219
220
221
222
    assert F.array_equal(u_g, v_rg)
    assert F.array_equal(v_g, u_rg)
    assert F.array_equal(eids_g, eids_rg)

    # withour share ndata
223
    g_r = dgl.reverse(g, copy_ndata=False)
224
225
226
227
228
229
230
231
232
233
    for etype_g, etype_gr in zip(g.canonical_etypes, g_r.canonical_etypes):
        assert etype_g[0] == etype_gr[2]
        assert etype_g[1] == etype_gr[1]
        assert etype_g[2] == etype_gr[0]
        assert g.number_of_edges(etype_g) == g_r.number_of_edges(etype_gr)
    for ntype in g.ntypes:
        assert g.number_of_nodes(ntype) == g_r.number_of_nodes(ntype)
    assert len(g_r.nodes['user'].data) == 0
    assert len(g_r.nodes['game'].data) == 0

234
    g_r = dgl.reverse(g, copy_ndata=True, copy_edata=True)
235
236
237
238
239
240
    print(g_r)
    for etype_g, etype_gr in zip(g.canonical_etypes, g_r.canonical_etypes):
        assert etype_g[0] == etype_gr[2]
        assert etype_g[1] == etype_gr[1]
        assert etype_g[2] == etype_gr[0]
        assert g.number_of_edges(etype_g) == g_r.number_of_edges(etype_gr)
241
242
243
244
    assert F.array_equal(g.edges['follows'].data['h'],
                         g_r.edges['follows'].data['h'])
    assert F.array_equal(g.edges['follows'].data['hh'],
                         g_r.edges['follows'].data['hh'])
245
246
247
248
249
250
251
252
253
254
255

    # add new node feature to g_r
    g_r.nodes['user'].data['hhh'] = F.tensor([0, 1, 2, 3, 4])
    assert ('hhh' in g.nodes['user'].data) is False
    assert ('hhh' in g_r.nodes['user'].data) is True

    # add new edge feature to g_r
    g_r.edges['follows'].data['hhh'] = F.tensor([1, 2, 3, 2, 0, 0, 1])
    assert ('hhh' in g.edges['follows'].data) is False
    assert ('hhh' in g_r.edges['follows'].data) is True

256

nv-dlasalle's avatar
nv-dlasalle committed
257
@parametrize_idtype
258
def test_reverse_shared_frames(idtype):
259
    g = dgl.DGLGraph()
260
    g = g.astype(idtype).to(F.ctx())
261
262
    g.add_nodes(3)
    g.add_edges([0, 1, 2], [1, 2, 1])
263
264
    g.ndata['h'] = F.tensor([[0.], [1.], [2.]])
    g.edata['h'] = F.tensor([[3.], [4.], [5.]])
265
266

    rg = g.reverse(share_ndata=True, share_edata=True)
267
268
269
    assert F.allclose(g.ndata['h'], rg.ndata['h'])
    assert F.allclose(g.edata['h'], rg.edata['h'])
    assert F.allclose(g.edges[[0, 2], [1, 1]].data['h'],
270
271
                      rg.edges[[1, 1], [0, 2]].data['h'])

272

273
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
274
def test_to_bidirected():
275
276
    # homogeneous graph
    elist = [(0, 0), (0, 1), (1, 0),
277
             (1, 1), (2, 1), (2, 2)]
278
    num_edges = 7
279
    g = dgl.graph(tuple(zip(*elist)))
280
281
282
283
284
285
286
287
288
289
    elist.append((1, 2))
    elist = set(elist)
    big = dgl.to_bidirected(g)
    assert big.number_of_edges() == num_edges
    src, dst = big.edges()
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == set(elist)

    # heterogeneous graph
    elist1 = [(0, 0), (0, 1), (1, 0),
290
              (1, 1), (2, 1), (2, 2)]
291
292
    elist2 = [(0, 0), (0, 1)]
    g = dgl.heterograph({
293
294
        ('user', 'wins', 'user'): tuple(zip(*elist1)),
        ('user', 'follows', 'user'): tuple(zip(*elist2))
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
    })
    g.nodes['user'].data['h'] = F.ones((3, 1))
    elist1.append((1, 2))
    elist1 = set(elist1)
    elist2.append((1, 0))
    elist2 = set(elist2)
    big = dgl.to_bidirected(g)
    assert big.number_of_edges('wins') == 7
    assert big.number_of_edges('follows') == 3
    src, dst = big.edges(etype='wins')
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == set(elist1)
    src, dst = big.edges(etype='follows')
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == set(elist2)

    big = dgl.to_bidirected(g, copy_ndata=True)
    assert F.array_equal(g.nodes['user'].data['h'], big.nodes['user'].data['h'])

314

315
def test_add_reverse_edges():
316
317
318
319
    # homogeneous graph
    g = dgl.graph((F.tensor([0, 1, 3, 1]), F.tensor([1, 2, 0, 2])))
    g.ndata['h'] = F.tensor([[0.], [1.], [2.], [1.]])
    g.edata['h'] = F.tensor([[3.], [4.], [5.], [6.]])
320
    bg = dgl.add_reverse_edges(g, copy_ndata=True, copy_edata=True)
321
322
323
324
325
    u, v = g.edges()
    ub, vb = bg.edges()
    assert F.array_equal(F.cat([u, v], dim=0), ub)
    assert F.array_equal(F.cat([v, u], dim=0), vb)
    assert F.array_equal(g.ndata['h'], bg.ndata['h'])
326
327
    assert F.array_equal(F.cat([g.edata['h'], g.edata['h']], dim=0),
                         bg.edata['h'])
328
329
330
331
332
333
    bg.ndata['hh'] = F.tensor([[0.], [1.], [2.], [1.]])
    assert ('hh' in g.ndata) is False
    bg.edata['hh'] = F.tensor([[0.], [1.], [2.], [1.], [0.], [1.], [2.], [1.]])
    assert ('hh' in g.edata) is False

    # donot share ndata and edata
334
    bg = dgl.add_reverse_edges(g, copy_ndata=False, copy_edata=False)
335
336
337
338
339
340
341
    ub, vb = bg.edges()
    assert F.array_equal(F.cat([u, v], dim=0), ub)
    assert F.array_equal(F.cat([v, u], dim=0), vb)
    assert ('h' in bg.ndata) is False
    assert ('h' in bg.edata) is False

    # zero edge graph
342
    g = dgl.graph(([], []))
343
344
    bg = dgl.add_reverse_edges(g, copy_ndata=True, copy_edata=True,
                               exclude_self=False)
345
346
347

    # heterogeneous graph
    g = dgl.heterograph({
348
349
        ('user', 'wins', 'user'): (
            F.tensor([0, 2, 0, 2, 2]), F.tensor([1, 1, 2, 1, 0])),
350
351
352
353
354
355
        ('user', 'plays', 'game'): (F.tensor([1, 2, 1]), F.tensor([2, 1, 1])),
        ('user', 'follows', 'user'): (F.tensor([1, 2, 1]), F.tensor([0, 0, 0]))
    })
    g.nodes['game'].data['hv'] = F.ones((3, 1))
    g.nodes['user'].data['hv'] = F.ones((3, 1))
    g.edges['wins'].data['h'] = F.tensor([0, 1, 2, 3, 4])
356
357
358
359
360
361
    bg = dgl.add_reverse_edges(g, copy_ndata=True, copy_edata=True,
                               ignore_bipartite=True)
    assert F.array_equal(g.nodes['game'].data['hv'],
                         bg.nodes['game'].data['hv'])
    assert F.array_equal(g.nodes['user'].data['hv'],
                         bg.nodes['user'].data['hv'])
362
363
364
365
    u, v = g.all_edges(order='eid', etype=('user', 'wins', 'user'))
    ub, vb = bg.all_edges(order='eid', etype=('user', 'wins', 'user'))
    assert F.array_equal(F.cat([u, v], dim=0), ub)
    assert F.array_equal(F.cat([v, u], dim=0), vb)
366
367
368
    assert F.array_equal(
        F.cat([g.edges['wins'].data['h'], g.edges['wins'].data['h']], dim=0),
        bg.edges['wins'].data['h'])
369
370
371
372
373
374
375
376
    u, v = g.all_edges(order='eid', etype=('user', 'follows', 'user'))
    ub, vb = bg.all_edges(order='eid', etype=('user', 'follows', 'user'))
    assert F.array_equal(F.cat([u, v], dim=0), ub)
    assert F.array_equal(F.cat([v, u], dim=0), vb)
    u, v = g.all_edges(order='eid', etype=('user', 'plays', 'game'))
    ub, vb = bg.all_edges(order='eid', etype=('user', 'plays', 'game'))
    assert F.array_equal(u, ub)
    assert F.array_equal(v, vb)
377
378
    assert set(bg.edges['plays'].data.keys()) == {dgl.EID}
    assert set(bg.edges['follows'].data.keys()) == {dgl.EID}
379
380

    # donot share ndata and edata
381
382
    bg = dgl.add_reverse_edges(g, copy_ndata=False, copy_edata=False,
                               ignore_bipartite=True)
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
    assert len(bg.edges['wins'].data) == 0
    assert len(bg.edges['plays'].data) == 0
    assert len(bg.edges['follows'].data) == 0
    assert len(bg.nodes['game'].data) == 0
    assert len(bg.nodes['user'].data) == 0
    u, v = g.all_edges(order='eid', etype=('user', 'wins', 'user'))
    ub, vb = bg.all_edges(order='eid', etype=('user', 'wins', 'user'))
    assert F.array_equal(F.cat([u, v], dim=0), ub)
    assert F.array_equal(F.cat([v, u], dim=0), vb)
    u, v = g.all_edges(order='eid', etype=('user', 'follows', 'user'))
    ub, vb = bg.all_edges(order='eid', etype=('user', 'follows', 'user'))
    assert F.array_equal(F.cat([u, v], dim=0), ub)
    assert F.array_equal(F.cat([v, u], dim=0), vb)
    u, v = g.all_edges(order='eid', etype=('user', 'plays', 'game'))
    ub, vb = bg.all_edges(order='eid', etype=('user', 'plays', 'game'))
    assert F.array_equal(u, ub)
    assert F.array_equal(v, vb)

401
402
403
404
405
406
407
408
    # test the case when some nodes have zero degree
    # homogeneous graph
    g = dgl.graph((F.tensor([0, 1, 3, 1]), F.tensor([1, 2, 0, 2])), num_nodes=6)
    g.ndata['h'] = F.tensor([[0.], [1.], [2.], [1.], [1.], [1.]])
    g.edata['h'] = F.tensor([[3.], [4.], [5.], [6.]])
    bg = dgl.add_reverse_edges(g, copy_ndata=True, copy_edata=True)
    assert g.number_of_nodes() == bg.number_of_nodes()
    assert F.array_equal(g.ndata['h'], bg.ndata['h'])
409
410
    assert F.array_equal(F.cat([g.edata['h'], g.edata['h']], dim=0),
                         bg.edata['h'])
411
412
413

    # heterogeneous graph
    g = dgl.heterograph({
414
415
        ('user', 'wins', 'user'): (
            F.tensor([0, 2, 0, 2, 2]), F.tensor([1, 1, 2, 1, 0])),
416
        ('user', 'plays', 'game'): (F.tensor([1, 2, 1]), F.tensor([2, 1, 1])),
417
418
        ('user', 'follows', 'user'): (
            F.tensor([1, 2, 1]), F.tensor([0, 0, 0]))},
419
420
421
422
423
424
425
        num_nodes_dict={
            'user': 5,
            'game': 3
        })
    g.nodes['game'].data['hv'] = F.ones((3, 1))
    g.nodes['user'].data['hv'] = F.ones((5, 1))
    g.edges['wins'].data['h'] = F.tensor([0, 1, 2, 3, 4])
426
427
    bg = dgl.add_reverse_edges(g, copy_ndata=True, copy_edata=True,
                               ignore_bipartite=True)
428
429
    assert g.number_of_nodes('user') == bg.number_of_nodes('user')
    assert g.number_of_nodes('game') == bg.number_of_nodes('game')
430
431
432
433
434
435
436
    assert F.array_equal(g.nodes['game'].data['hv'],
                         bg.nodes['game'].data['hv'])
    assert F.array_equal(g.nodes['user'].data['hv'],
                         bg.nodes['user'].data['hv'])
    assert F.array_equal(
        F.cat([g.edges['wins'].data['h'], g.edges['wins'].data['h']], dim=0),
        bg.edges['wins'].data['h'])
437

438
439
440
441
442
443
444
445
446
447
    # test exclude_self
    g = dgl.heterograph({
        ('A', 'r1', 'A'): (F.tensor([0, 0, 1, 1]), F.tensor([0, 1, 1, 2])),
        ('A', 'r2', 'A'): (F.tensor([0, 1]), F.tensor([1, 2]))
    })
    g.edges['r1'].data['h'] = F.tensor([0, 1, 2, 3])
    rg = dgl.add_reverse_edges(g, copy_edata=True, exclude_self=True)
    assert rg.num_edges('r1') == 6
    assert rg.num_edges('r2') == 4
    assert F.array_equal(rg.edges['r1'].data['h'], F.tensor([0, 1, 2, 3, 1, 3]))
448

449

450
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
451
452
453
454
455
456
457
458
459
460
def test_simple_graph():
    elist = [(0, 1), (0, 2), (1, 2), (0, 1)]
    g = dgl.DGLGraph(elist, readonly=True)
    assert g.is_multigraph
    sg = dgl.to_simple_graph(g)
    assert not sg.is_multigraph
    assert sg.number_of_edges() == 3
    src, dst = sg.edges()
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == set(elist)
461

462

463
464
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
def _test_bidirected_graph():
465
    def _test(in_readonly, out_readonly):
466
        elist = [(0, 0), (0, 1), (1, 0),
467
                 (1, 1), (2, 1), (2, 2)]
468
        num_edges = 7
469
470
471
        g = dgl.DGLGraph(elist, readonly=in_readonly)
        elist.append((1, 2))
        elist = set(elist)
472
        big = dgl.to_bidirected_stale(g, out_readonly)
473
        assert big.number_of_edges() == num_edges
474
475
476
477
478
479
480
481
482
        src, dst = big.edges()
        eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
        assert eset == set(elist)

    _test(True, True)
    _test(True, False)
    _test(False, True)
    _test(False, False)

483

484
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
485
486
487
488
def test_khop_graph():
    N = 20
    feat = F.randn((N, 5))

Mufei Li's avatar
Mufei Li committed
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
    def _test(g):
        for k in range(4):
            g_k = dgl.khop_graph(g, k)
            # use original graph to do message passing for k times.
            g.ndata['h'] = feat
            for _ in range(k):
                g.update_all(fn.copy_u('h', 'm'), fn.sum('m', 'h'))
            h_0 = g.ndata.pop('h')
            # use k-hop graph to do message passing for one time.
            g_k.ndata['h'] = feat
            g_k.update_all(fn.copy_u('h', 'm'), fn.sum('m', 'h'))
            h_1 = g_k.ndata.pop('h')
            assert F.allclose(h_0, h_1, rtol=1e-3, atol=1e-3)

    # Test for random undirected graphs
    g = dgl.DGLGraph(nx.erdos_renyi_graph(N, 0.3))
    _test(g)
    # Test for random directed graphs
    g = dgl.DGLGraph(nx.erdos_renyi_graph(N, 0.3, directed=True))
    _test(g)
509

510

511
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
512
513
514
515
516
def test_khop_adj():
    N = 20
    feat = F.randn((N, 5))
    g = dgl.DGLGraph(nx.erdos_renyi_graph(N, 0.3))
    for k in range(3):
517
        adj = F.tensor(F.swapaxes(dgl.khop_adj(g, k), 0, 1))
518
519
520
521
522
523
524
525
526
        # use original graph to do message passing for k times.
        g.ndata['h'] = feat
        for _ in range(k):
            g.update_all(fn.copy_u('h', 'm'), fn.sum('m', 'h'))
        h_0 = g.ndata.pop('h')
        # use k-hop adj to do message passing for one time.
        h_1 = F.matmul(adj, feat)
        assert F.allclose(h_0, h_1, rtol=1e-3, atol=1e-3)

527

528
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
529
530
531
532
533
534
535
def test_laplacian_lambda_max():
    N = 20
    eps = 1e-6
    # test DGLGraph
    g = dgl.DGLGraph(nx.erdos_renyi_graph(N, 0.3))
    l_max = dgl.laplacian_lambda_max(g)
    assert (l_max[0] < 2 + eps)
Zihao Ye's avatar
Zihao Ye committed
536
    # test batched DGLGraph
537
    '''
538
539
540
541
542
543
544
545
546
    N_arr = [20, 30, 10, 12]
    bg = dgl.batch([
        dgl.DGLGraph(nx.erdos_renyi_graph(N, 0.3))
        for N in N_arr
    ])
    l_max_arr = dgl.laplacian_lambda_max(bg)
    assert len(l_max_arr) == len(N_arr)
    for l_max in l_max_arr:
        assert l_max < 2 + eps
547
    '''
548

549

550
def create_large_graph(num_nodes, idtype=F.int64):
551
552
553
    row = np.random.choice(num_nodes, num_nodes * 10)
    col = np.random.choice(num_nodes, num_nodes * 10)
    spm = spsp.coo_matrix((np.ones(len(row)), (row, col)))
554
    spm.sum_duplicates()
555

556
    return dgl.from_scipy(spm, idtype=idtype)
557

558

559
# Disabled since everything will be on heterogeneous graphs
560
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
561
def test_partition_with_halo():
562
    g = create_large_graph(1000)
563
    node_part = np.random.choice(4, g.number_of_nodes())
564
565
    subgs, _, _ = dgl.transforms.partition_graph_with_halo(g, node_part, 2,
                                                           reshuffle=True)
566
567
568
    for part_id, subg in subgs.items():
        node_ids = np.nonzero(node_part == part_id)[0]
        lnode_ids = np.nonzero(F.asnumpy(subg.ndata['inner_node']))[0]
569
570
        orig_nids = F.asnumpy(subg.ndata['orig_id'])[lnode_ids]
        assert np.all(np.sort(orig_nids) == node_ids)
571
572
573
574
575
        assert np.all(F.asnumpy(subg.in_degrees(lnode_ids)) == F.asnumpy(
            g.in_degrees(orig_nids)))
        assert np.all(F.asnumpy(subg.out_degrees(lnode_ids)) == F.asnumpy(
            g.out_degrees(orig_nids)))

576

577
@unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
578
579
@unittest.skipIf(F._default_context_str == 'gpu',
                 reason="METIS doesn't support GPU")
nv-dlasalle's avatar
nv-dlasalle committed
580
@parametrize_idtype
581
def test_metis_partition(idtype):
Da Zheng's avatar
Da Zheng committed
582
    # TODO(zhengda) Metis fails to partition a small graph.
583
584
585
586
587
588
589
590
591
592
593
594
595
    g = create_large_graph(1000, idtype=idtype)
    if idtype == F.int64:
        check_metis_partition(g, 0)
        check_metis_partition(g, 1)
        check_metis_partition(g, 2)
        check_metis_partition_with_constraint(g)
    else:
        assert_fail = False
        try:
            check_metis_partition(g, 1)
        except:
            assert_fail = True
        assert assert_fail
596

597

598
599
def check_metis_partition_with_constraint(g):
    ntypes = np.zeros((g.number_of_nodes(),), dtype=np.int32)
600
601
602
603
    ntypes[0:int(g.number_of_nodes() / 4)] = 1
    ntypes[int(g.number_of_nodes() * 3 / 4):] = 2
    subgs = dgl.transforms.metis_partition(g, 4, extra_cached_hops=1,
                                           balance_ntypes=ntypes)
604
605
606
607
608
609
610
611
    if subgs is not None:
        for i in subgs:
            subg = subgs[i]
            parent_nids = F.asnumpy(subg.ndata[dgl.NID])
            sub_ntypes = ntypes[parent_nids]
            print('type0:', np.sum(sub_ntypes == 0))
            print('type1:', np.sum(sub_ntypes == 1))
            print('type2:', np.sum(sub_ntypes == 2))
612
    subgs = dgl.transforms.metis_partition(g, 4, extra_cached_hops=1,
613
614
                                           balance_ntypes=ntypes,
                                           balance_edges=True)
615
616
617
618
619
620
621
622
    if subgs is not None:
        for i in subgs:
            subg = subgs[i]
            parent_nids = F.asnumpy(subg.ndata[dgl.NID])
            sub_ntypes = ntypes[parent_nids]
            print('type0:', np.sum(sub_ntypes == 0))
            print('type1:', np.sum(sub_ntypes == 1))
            print('type2:', np.sum(sub_ntypes == 2))
Da Zheng's avatar
Da Zheng committed
623

624

Da Zheng's avatar
Da Zheng committed
625
def check_metis_partition(g, extra_hops):
626
    subgs = dgl.transforms.metis_partition(g, 4, extra_cached_hops=extra_hops)
627
628
629
630
    num_inner_nodes = 0
    num_inner_edges = 0
    if subgs is not None:
        for part_id, subg in subgs.items():
Da Zheng's avatar
Da Zheng committed
631
632
633
634
            lnode_ids = np.nonzero(F.asnumpy(subg.ndata['inner_node']))[0]
            ledge_ids = np.nonzero(F.asnumpy(subg.edata['inner_edge']))[0]
            num_inner_nodes += len(lnode_ids)
            num_inner_edges += len(ledge_ids)
635
636
            assert np.sum(F.asnumpy(subg.ndata['part_id']) == part_id) == len(
                lnode_ids)
637
638
639
        assert num_inner_nodes == g.number_of_nodes()
        print(g.number_of_edges() - num_inner_edges)

Da Zheng's avatar
Da Zheng committed
640
641
642
    if extra_hops == 0:
        return

643
    # partitions with node reshuffling
644
645
    subgs = dgl.transforms.metis_partition(g, 4, extra_cached_hops=extra_hops,
                                           reshuffle=True)
646
647
    num_inner_nodes = 0
    num_inner_edges = 0
Da Zheng's avatar
Da Zheng committed
648
    edge_cnts = np.zeros((g.number_of_edges(),))
649
650
651
652
653
654
    if subgs is not None:
        for part_id, subg in subgs.items():
            lnode_ids = np.nonzero(F.asnumpy(subg.ndata['inner_node']))[0]
            ledge_ids = np.nonzero(F.asnumpy(subg.edata['inner_edge']))[0]
            num_inner_nodes += len(lnode_ids)
            num_inner_edges += len(ledge_ids)
655
656
            assert np.sum(F.asnumpy(subg.ndata['part_id']) == part_id) == len(
                lnode_ids)
Da Zheng's avatar
Da Zheng committed
657
658
659
660
661
            nids = F.asnumpy(subg.ndata[dgl.NID])

            # ensure the local node Ids are contiguous.
            parent_ids = F.asnumpy(subg.ndata[dgl.NID])
            parent_ids = parent_ids[:len(lnode_ids)]
662
663
            assert np.all(
                parent_ids == np.arange(parent_ids[0], parent_ids[-1] + 1))
Da Zheng's avatar
Da Zheng committed
664
665
666
667
668
669
670
671
672
673
674
675
676
677

            # count the local edges.
            parent_ids = F.asnumpy(subg.edata[dgl.EID])[ledge_ids]
            edge_cnts[parent_ids] += 1

            orig_ids = subg.ndata['orig_id']
            inner_node = F.asnumpy(subg.ndata['inner_node'])
            for nid in range(subg.number_of_nodes()):
                neighs = subg.predecessors(nid)
                old_neighs1 = F.gather_row(orig_ids, neighs)
                old_nid = F.asnumpy(orig_ids[nid])
                old_neighs2 = g.predecessors(old_nid)
                # If this is an inner node, it should have the full neighborhood.
                if inner_node[nid]:
678
679
                    assert np.all(np.sort(F.asnumpy(old_neighs1)) == np.sort(
                        F.asnumpy(old_neighs2)))
Da Zheng's avatar
Da Zheng committed
680
681
682
        # Normally, local edges are only counted once.
        assert np.all(edge_cnts == 1)

683
684
685
        assert num_inner_nodes == g.number_of_nodes()
        print(g.number_of_edges() - num_inner_edges)

686
687
688

@unittest.skipIf(F._default_context_str == 'gpu',
                 reason="It doesn't support GPU")
Da Zheng's avatar
Da Zheng committed
689
def test_reorder_nodes():
690
    g = create_large_graph(1000)
Da Zheng's avatar
Da Zheng committed
691
692
    new_nids = np.random.permutation(g.number_of_nodes())
    # TODO(zhengda) we need to test both CSR and COO.
693
    new_g = dgl.partition.reorder_nodes(g, new_nids)
Da Zheng's avatar
Da Zheng committed
694
695
696
697
698
699
700
701
702
    new_in_deg = new_g.in_degrees()
    new_out_deg = new_g.out_degrees()
    in_deg = g.in_degrees()
    out_deg = g.out_degrees()
    new_in_deg1 = F.scatter_row(in_deg, F.tensor(new_nids), in_deg)
    new_out_deg1 = F.scatter_row(out_deg, F.tensor(new_nids), out_deg)
    assert np.all(F.asnumpy(new_in_deg == new_in_deg1))
    assert np.all(F.asnumpy(new_out_deg == new_out_deg1))
    orig_ids = F.asnumpy(new_g.ndata['orig_id'])
703
704
705
706
707
708
709
    for nid in range(g.number_of_nodes()):
        neighs = F.asnumpy(g.successors(nid))
        new_neighs1 = new_nids[neighs]
        new_nid = new_nids[nid]
        new_neighs2 = new_g.successors(new_nid)
        assert np.all(np.sort(new_neighs1) == np.sort(F.asnumpy(new_neighs2)))

Da Zheng's avatar
Da Zheng committed
710
711
712
713
714
715
716
717
718
719
720
721
722
    for nid in range(new_g.number_of_nodes()):
        neighs = F.asnumpy(new_g.successors(nid))
        old_neighs1 = orig_ids[neighs]
        old_nid = orig_ids[nid]
        old_neighs2 = g.successors(old_nid)
        assert np.all(np.sort(old_neighs1) == np.sort(F.asnumpy(old_neighs2)))

        neighs = F.asnumpy(new_g.predecessors(nid))
        old_neighs1 = orig_ids[neighs]
        old_nid = orig_ids[nid]
        old_neighs2 = g.predecessors(old_nid)
        assert np.all(np.sort(old_neighs1) == np.sort(F.asnumpy(old_neighs2)))

723

nv-dlasalle's avatar
nv-dlasalle committed
724
@parametrize_idtype
725
def test_compact(idtype):
726
    g1 = dgl.heterograph({
727
728
729
        ('user', 'follow', 'user'): ([1, 3], [3, 5]),
        ('user', 'plays', 'game'): ([2, 3, 2], [4, 4, 5]),
        ('game', 'wished-by', 'user'): ([6, 5], [7, 7])},
730
        {'user': 20, 'game': 10}, idtype=idtype, device=F.ctx())
731
732

    g2 = dgl.heterograph({
733
734
        ('game', 'clicked-by', 'user'): ([3], [1]),
        ('user', 'likes', 'user'): ([1, 8], [8, 9])},
735
        {'user': 20, 'game': 10}, idtype=idtype, device=F.ctx())
736

737
    g3 = dgl.heterograph({('user', '_E', 'user'): ((0, 1), (1, 2))},
738
                         {'user': 10}, idtype=idtype, device=F.ctx())
739
    g4 = dgl.heterograph({('user', '_E', 'user'): ((1, 3), (3, 5))},
740
                         {'user': 10}, idtype=idtype, device=F.ctx())
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760

    def _check(g, new_g, induced_nodes):
        assert g.ntypes == new_g.ntypes
        assert g.canonical_etypes == new_g.canonical_etypes

        for ntype in g.ntypes:
            assert -1 not in induced_nodes[ntype]

        for etype in g.canonical_etypes:
            g_src, g_dst = g.all_edges(order='eid', etype=etype)
            g_src = F.asnumpy(g_src)
            g_dst = F.asnumpy(g_dst)
            new_g_src, new_g_dst = new_g.all_edges(order='eid', etype=etype)
            new_g_src_mapped = induced_nodes[etype[0]][F.asnumpy(new_g_src)]
            new_g_dst_mapped = induced_nodes[etype[2]][F.asnumpy(new_g_dst)]
            assert (g_src == new_g_src_mapped).all()
            assert (g_dst == new_g_dst_mapped).all()

    # Test default
    new_g1 = dgl.compact_graphs(g1)
761
762
    induced_nodes = {ntype: new_g1.nodes[ntype].data[dgl.NID] for ntype in
                     new_g1.ntypes}
763
    induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()}
764
    assert new_g1.idtype == idtype
765
766
767
768
769
770
    assert set(induced_nodes['user']) == set([1, 3, 5, 2, 7])
    assert set(induced_nodes['game']) == set([4, 5, 6])
    _check(g1, new_g1, induced_nodes)

    # Test with always_preserve given a dict
    new_g1 = dgl.compact_graphs(
771
772
        g1, always_preserve={'game': F.tensor([4, 7], idtype)})
    assert new_g1.idtype == idtype
773
774
    induced_nodes = {ntype: new_g1.nodes[ntype].data[dgl.NID] for ntype in
                     new_g1.ntypes}
775
776
777
778
779
780
781
    induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()}
    assert set(induced_nodes['user']) == set([1, 3, 5, 2, 7])
    assert set(induced_nodes['game']) == set([4, 5, 6, 7])
    _check(g1, new_g1, induced_nodes)

    # Test with always_preserve given a tensor
    new_g3 = dgl.compact_graphs(
782
        g3, always_preserve=F.tensor([1, 7], idtype))
783
784
    induced_nodes = {ntype: new_g3.nodes[ntype].data[dgl.NID] for ntype in
                     new_g3.ntypes}
785
    induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()}
786

787
    assert new_g3.idtype == idtype
788
789
790
791
792
    assert set(induced_nodes['user']) == set([0, 1, 2, 7])
    _check(g3, new_g3, induced_nodes)

    # Test multiple graphs
    new_g1, new_g2 = dgl.compact_graphs([g1, g2])
793
794
    induced_nodes = {ntype: new_g1.nodes[ntype].data[dgl.NID] for ntype in
                     new_g1.ntypes}
795
    induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()}
796
797
    assert new_g1.idtype == idtype
    assert new_g2.idtype == idtype
798
799
800
801
802
803
804
    assert set(induced_nodes['user']) == set([1, 3, 5, 2, 7, 8, 9])
    assert set(induced_nodes['game']) == set([3, 4, 5, 6])
    _check(g1, new_g1, induced_nodes)
    _check(g2, new_g2, induced_nodes)

    # Test multiple graphs with always_preserve given a dict
    new_g1, new_g2 = dgl.compact_graphs(
805
        [g1, g2], always_preserve={'game': F.tensor([4, 7], dtype=idtype)})
806
807
    induced_nodes = {ntype: new_g1.nodes[ntype].data[dgl.NID] for ntype in
                     new_g1.ntypes}
808
    induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()}
809
810
    assert new_g1.idtype == idtype
    assert new_g2.idtype == idtype
811
812
813
814
815
816
817
    assert set(induced_nodes['user']) == set([1, 3, 5, 2, 7, 8, 9])
    assert set(induced_nodes['game']) == set([3, 4, 5, 6, 7])
    _check(g1, new_g1, induced_nodes)
    _check(g2, new_g2, induced_nodes)

    # Test multiple graphs with always_preserve given a tensor
    new_g3, new_g4 = dgl.compact_graphs(
818
        [g3, g4], always_preserve=F.tensor([1, 7], dtype=idtype))
819
820
    induced_nodes = {ntype: new_g3.nodes[ntype].data[dgl.NID] for ntype in
                     new_g3.ntypes}
821
    induced_nodes = {k: F.asnumpy(v) for k, v in induced_nodes.items()}
822

823
824
825
    assert new_g3.idtype == idtype
    assert new_g4.idtype == idtype

826
827
828
829
    assert set(induced_nodes['user']) == set([0, 1, 2, 3, 5, 7])
    _check(g3, new_g3, induced_nodes)
    _check(g4, new_g4, induced_nodes)

830
831
832

@unittest.skipIf(F._default_context_str == 'gpu',
                 reason="GPU to simple not implemented")
nv-dlasalle's avatar
nv-dlasalle committed
833
@parametrize_idtype
834
def test_to_simple(idtype):
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
    # homogeneous graph
    g = dgl.graph((F.tensor([0, 1, 2, 1]), F.tensor([1, 2, 0, 2])))
    g.ndata['h'] = F.tensor([[0.], [1.], [2.]])
    g.edata['h'] = F.tensor([[3.], [4.], [5.], [6.]])
    sg, wb = dgl.to_simple(g, writeback_mapping=True)
    u, v = g.all_edges(form='uv', order='eid')
    u = F.asnumpy(u).tolist()
    v = F.asnumpy(v).tolist()
    uv = list(zip(u, v))
    eid_map = F.asnumpy(wb)

    su, sv = sg.all_edges(form='uv', order='eid')
    su = F.asnumpy(su).tolist()
    sv = F.asnumpy(sv).tolist()
    suv = list(zip(su, sv))
    sc = F.asnumpy(sg.edata['count'])
    assert set(uv) == set(suv)
    for i, e in enumerate(suv):
        assert sc[i] == sum(e == _e for _e in uv)
    for i, e in enumerate(uv):
        assert eid_map[i] == suv.index(e)
    # shared ndata
    assert F.array_equal(sg.ndata['h'], g.ndata['h'])
    assert 'h' not in sg.edata
    # new ndata to sg
    sg.ndata['hh'] = F.tensor([[0.], [1.], [2.]])
    assert 'hh' not in g.ndata

    sg = dgl.to_simple(g, writeback_mapping=False, copy_ndata=False)
    assert 'h' not in sg.ndata
    assert 'h' not in sg.edata

867
868
869
870
871
872
873
874
    # test coalesce edge feature
    sg = dgl.to_simple(g, copy_edata=True, aggregator='arbitrary')
    assert F.allclose(sg.edata['h'][1], F.tensor([4.]))
    sg = dgl.to_simple(g, copy_edata=True, aggregator='sum')
    assert F.allclose(sg.edata['h'][1], F.tensor([10.]))
    sg = dgl.to_simple(g, copy_edata=True, aggregator='mean')
    assert F.allclose(sg.edata['h'][1], F.tensor([5.]))

875
    # heterogeneous graph
876
    g = dgl.heterograph({
877
878
        ('user', 'follow', 'user'): ([0, 1, 2, 1, 1, 1],
                                     [1, 3, 2, 3, 4, 4]),
879
880
        ('user', 'plays', 'game'): (
            [3, 2, 1, 1, 3, 2, 2], [5, 3, 4, 4, 5, 3, 3])},
881
        idtype=idtype, device=F.ctx())
882
883
884
    g.nodes['user'].data['h'] = F.tensor([0, 1, 2, 3, 4])
    g.nodes['user'].data['hh'] = F.tensor([0, 1, 2, 3, 4])
    g.edges['follow'].data['h'] = F.tensor([0, 1, 2, 3, 4, 5])
885
886
    sg, wb = dgl.to_simple(g, return_counts='weights', writeback_mapping=True,
                           copy_edata=True)
887
    g.nodes['game'].data['h'] = F.tensor([0, 1, 2, 3, 4, 5])
888
889
890
891
892
893

    for etype in g.canonical_etypes:
        u, v = g.all_edges(form='uv', order='eid', etype=etype)
        u = F.asnumpy(u).tolist()
        v = F.asnumpy(v).tolist()
        uv = list(zip(u, v))
894
        eid_map = F.asnumpy(wb[etype])
895
896
897
898
899
900
901
902
903
904
905
906

        su, sv = sg.all_edges(form='uv', order='eid', etype=etype)
        su = F.asnumpy(su).tolist()
        sv = F.asnumpy(sv).tolist()
        suv = list(zip(su, sv))
        sw = F.asnumpy(sg.edges[etype].data['weights'])

        assert set(uv) == set(suv)
        for i, e in enumerate(suv):
            assert sw[i] == sum(e == _e for _e in uv)
        for i, e in enumerate(uv):
            assert eid_map[i] == suv.index(e)
907
908
    # shared ndata
    assert F.array_equal(sg.nodes['user'].data['h'], g.nodes['user'].data['h'])
909
910
    assert F.array_equal(sg.nodes['user'].data['hh'],
                         g.nodes['user'].data['hh'])
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
    assert 'h' not in sg.nodes['game'].data
    # new ndata to sg
    sg.nodes['user'].data['hhh'] = F.tensor([0, 1, 2, 3, 4])
    assert 'hhh' not in g.nodes['user'].data
    # share edata
    feat_idx = F.asnumpy(wb[('user', 'follow', 'user')])
    _, indices = np.unique(feat_idx, return_index=True)
    assert np.array_equal(F.asnumpy(sg.edges['follow'].data['h']),
                          F.asnumpy(g.edges['follow'].data['h'])[indices])

    sg = dgl.to_simple(g, writeback_mapping=False, copy_ndata=False)
    for ntype in g.ntypes:
        assert g.number_of_nodes(ntype) == sg.number_of_nodes(ntype)
    assert 'h' not in sg.nodes['user'].data
    assert 'hh' not in sg.nodes['user'].data
926

927
928
929
930
931
932
933
934
935
936
    # verify DGLGraph.edge_ids() after dgl.to_simple()
    # in case ids are not initialized in underlying coo2csr()
    u = F.tensor([0, 1, 2])
    v = F.tensor([1, 2, 3])
    eids = F.tensor([0, 1, 2])
    g = dgl.graph((u, v))
    assert F.array_equal(g.edge_ids(u, v), eids)
    sg = dgl.to_simple(g)
    assert F.array_equal(sg.edge_ids(u, v), eids)

937

nv-dlasalle's avatar
nv-dlasalle committed
938
@parametrize_idtype
939
def test_to_block(idtype):
940
    def check(g, bg, ntype, etype, dst_nodes, include_dst_in_src=True):
941
942
943
        if dst_nodes is not None:
            assert F.array_equal(bg.dstnodes[ntype].data[dgl.NID], dst_nodes)
        n_dst_nodes = bg.number_of_nodes('DST/' + ntype)
944
945
946
947
        if include_dst_in_src:
            assert F.array_equal(
                bg.srcnodes[ntype].data[dgl.NID][:n_dst_nodes],
                bg.dstnodes[ntype].data[dgl.NID])
948
949
950
951
952
953

        g = g[etype]
        bg = bg[etype]
        induced_src = bg.srcdata[dgl.NID]
        induced_dst = bg.dstdata[dgl.NID]
        induced_eid = bg.edata[dgl.EID]
954

955
956
957
958
959
960
961
962
963
964
965
        bg_src, bg_dst = bg.all_edges(order='eid')
        src_ans, dst_ans = g.all_edges(order='eid')

        induced_src_bg = F.gather_row(induced_src, bg_src)
        induced_dst_bg = F.gather_row(induced_dst, bg_dst)
        induced_src_ans = F.gather_row(src_ans, induced_eid)
        induced_dst_ans = F.gather_row(dst_ans, induced_eid)

        assert F.array_equal(induced_src_bg, induced_src_ans)
        assert F.array_equal(induced_dst_bg, induced_dst_ans)

966
    def checkall(g, bg, dst_nodes, include_dst_in_src=True):
967
968
        for etype in g.etypes:
            ntype = g.to_canonical_etype(etype)[2]
969
            if dst_nodes is not None and ntype in dst_nodes:
970
                check(g, bg, ntype, etype, dst_nodes[ntype], include_dst_in_src)
971
            else:
972
                check(g, bg, ntype, etype, None, include_dst_in_src)
973

974
    # homogeneous graph
975
976
    g = dgl.graph(
        (F.tensor([1, 2], dtype=idtype), F.tensor([2, 3], dtype=idtype)))
977
978
979
980
981
982
983
984
985
    dst_nodes = F.tensor([3, 2], dtype=idtype)
    bg = dgl.to_block(g, dst_nodes=dst_nodes)
    check(g, bg, '_N', '_E', dst_nodes)

    src_nodes = bg.srcnodes['_N'].data[dgl.NID]
    bg = dgl.to_block(g, dst_nodes=dst_nodes, src_nodes=src_nodes)
    check(g, bg, '_N', '_E', dst_nodes)

    # heterogeneous graph
986
    g = dgl.heterograph({
987
988
        ('A', 'AA', 'A'): ([0, 2, 1, 3], [1, 3, 2, 4]),
        ('A', 'AB', 'B'): ([0, 1, 3, 1], [1, 3, 5, 6]),
989
        ('B', 'BA', 'A'): ([2, 3], [3, 2])}, idtype=idtype, device=F.ctx())
990
991
992
993
994
    g.nodes['A'].data['x'] = F.randn((5, 10))
    g.nodes['B'].data['x'] = F.randn((7, 5))
    g.edges['AA'].data['x'] = F.randn((4, 3))
    g.edges['AB'].data['x'] = F.randn((4, 3))
    g.edges['BA'].data['x'] = F.randn((2, 3))
995
996
    g_a = g['AA']

997
998
999
1000
1001
    def check_features(g, bg):
        for ntype in bg.srctypes:
            for key in g.nodes[ntype].data:
                assert F.array_equal(
                    bg.srcnodes[ntype].data[key],
1002
1003
                    F.gather_row(g.nodes[ntype].data[key],
                                 bg.srcnodes[ntype].data[dgl.NID]))
1004
1005
1006
1007
        for ntype in bg.dsttypes:
            for key in g.nodes[ntype].data:
                assert F.array_equal(
                    bg.dstnodes[ntype].data[key],
1008
1009
                    F.gather_row(g.nodes[ntype].data[key],
                                 bg.dstnodes[ntype].data[dgl.NID]))
1010
1011
1012
1013
        for etype in bg.canonical_etypes:
            for key in g.edges[etype].data:
                assert F.array_equal(
                    bg.edges[etype].data[key],
1014
1015
                    F.gather_row(g.edges[etype].data[key],
                                 bg.edges[etype].data[dgl.EID]))
1016

1017
1018
    bg = dgl.to_block(g_a)
    check(g_a, bg, 'A', 'AA', None)
1019
    check_features(g_a, bg)
1020
1021
1022
1023
1024
    assert bg.number_of_src_nodes() == 5
    assert bg.number_of_dst_nodes() == 4

    bg = dgl.to_block(g_a, include_dst_in_src=False)
    check(g_a, bg, 'A', 'AA', None, False)
1025
    check_features(g_a, bg)
1026
1027
    assert bg.number_of_src_nodes() == 4
    assert bg.number_of_dst_nodes() == 4
1028

1029
    dst_nodes = F.tensor([4, 3, 2, 1], dtype=idtype)
1030
1031
    bg = dgl.to_block(g_a, dst_nodes)
    check(g_a, bg, 'A', 'AA', dst_nodes)
1032
    check_features(g_a, bg)
1033
1034
1035
1036

    g_ab = g['AB']

    bg = dgl.to_block(g_ab)
1037
    assert bg.idtype == idtype
1038
    assert bg.number_of_nodes('SRC/B') == 4
1039
1040
    assert F.array_equal(bg.srcnodes['B'].data[dgl.NID],
                         bg.dstnodes['B'].data[dgl.NID])
1041
    assert bg.number_of_nodes('DST/A') == 0
1042
    checkall(g_ab, bg, None)
1043
    check_features(g_ab, bg)
1044

1045
    dst_nodes = {'B': F.tensor([5, 6, 3, 1], dtype=idtype)}
1046
    bg = dgl.to_block(g, dst_nodes)
1047
    assert bg.number_of_nodes('SRC/B') == 4
1048
1049
    assert F.array_equal(bg.srcnodes['B'].data[dgl.NID],
                         bg.dstnodes['B'].data[dgl.NID])
1050
1051
    assert bg.number_of_nodes('DST/A') == 0
    checkall(g, bg, dst_nodes)
1052
    check_features(g, bg)
1053

1054
1055
    dst_nodes = {'A': F.tensor([4, 3, 2, 1], dtype=idtype),
                 'B': F.tensor([3, 5, 6, 1], dtype=idtype)}
1056
1057
    bg = dgl.to_block(g, dst_nodes=dst_nodes)
    checkall(g, bg, dst_nodes)
1058
    check_features(g, bg)
1059

1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
    # test specifying lhs_nodes with include_dst_in_src
    src_nodes = {}
    for ntype in dst_nodes.keys():
        # use the previous run to get the list of source nodes
        src_nodes[ntype] = bg.srcnodes[ntype].data[dgl.NID]
    bg = dgl.to_block(g, dst_nodes=dst_nodes, src_nodes=src_nodes)
    checkall(g, bg, dst_nodes)
    check_features(g, bg)

    # test without include_dst_in_src
1070
1071
    dst_nodes = {'A': F.tensor([4, 3, 2, 1], dtype=idtype),
                 'B': F.tensor([3, 5, 6, 1], dtype=idtype)}
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
    bg = dgl.to_block(g, dst_nodes=dst_nodes, include_dst_in_src=False)
    checkall(g, bg, dst_nodes, False)
    check_features(g, bg)

    # test specifying lhs_nodes without include_dst_in_src
    src_nodes = {}
    for ntype in dst_nodes.keys():
        # use the previous run to get the list of source nodes
        src_nodes[ntype] = bg.srcnodes[ntype].data[dgl.NID]
    bg = dgl.to_block(g, dst_nodes=dst_nodes, include_dst_in_src=False,
1082
                      src_nodes=src_nodes)
1083
1084
1085
1086
    checkall(g, bg, dst_nodes, False)
    check_features(g, bg)


1087
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
nv-dlasalle's avatar
nv-dlasalle committed
1088
@parametrize_idtype
1089
def test_remove_edges(idtype):
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
    def check(g1, etype, g, edges_removed):
        src, dst, eid = g.edges(etype=etype, form='all')
        src1, dst1 = g1.edges(etype=etype, order='eid')
        if etype is not None:
            eid1 = g1.edges[etype].data[dgl.EID]
        else:
            eid1 = g1.edata[dgl.EID]
        src1 = F.asnumpy(src1)
        dst1 = F.asnumpy(dst1)
        eid1 = F.asnumpy(eid1)
        src = F.asnumpy(src)
        dst = F.asnumpy(dst)
        eid = F.asnumpy(eid)
        sde_set = set(zip(src, dst, eid))

        for s, d, e in zip(src1, dst1, eid1):
            assert (s, d, e) in sde_set
        assert not np.isin(edges_removed, eid1).any()
1108
        assert g1.idtype == g.idtype
1109
1110
1111

    for fmt in ['coo', 'csr', 'csc']:
        for edges_to_remove in [[2], [2, 2], [3, 2], [1, 3, 1, 2]]:
1112
1113
            g = dgl.graph(([0, 2, 1, 3], [1, 3, 2, 4]), idtype=idtype).formats(
                fmt)
1114
            g1 = dgl.remove_edges(g, F.tensor(edges_to_remove, idtype))
1115
1116
            check(g1, None, g, edges_to_remove)

1117
            g = dgl.from_scipy(
1118
1119
                spsp.csr_matrix(([1, 1, 1, 1], ([0, 2, 1, 3], [1, 3, 2, 4])),
                                shape=(5, 5)),
1120
1121
                idtype=idtype).formats(fmt)
            g1 = dgl.remove_edges(g, F.tensor(edges_to_remove, idtype))
1122
1123
1124
            check(g1, None, g, edges_to_remove)

    g = dgl.heterograph({
1125
1126
1127
        ('A', 'AA', 'A'): ([0, 2, 1, 3], [1, 3, 2, 4]),
        ('A', 'AB', 'B'): ([0, 1, 3, 1], [1, 3, 5, 6]),
        ('B', 'BA', 'A'): ([2, 3], [3, 2])}, idtype=idtype)
1128
1129
1130
    g2 = dgl.remove_edges(g, {'AA': F.tensor([2], idtype),
                              'AB': F.tensor([3], idtype),
                              'BA': F.tensor([1], idtype)})
1131
1132
1133
    check(g2, 'AA', g, [2])
    check(g2, 'AB', g, [3])
    check(g2, 'BA', g, [1])
1134

1135
1136
1137
    g3 = dgl.remove_edges(g, {'AA': F.tensor([], idtype),
                              'AB': F.tensor([3], idtype),
                              'BA': F.tensor([1], idtype)})
1138
1139
1140
1141
    check(g3, 'AA', g, [])
    check(g3, 'AB', g, [3])
    check(g3, 'BA', g, [1])

1142
    g4 = dgl.remove_edges(g, {'AB': F.tensor([3, 1, 2, 0], idtype)})
1143
    check(g4, 'AA', g, [])
1144
    check(g4, 'AB', g, [3, 1, 2, 0])
1145
1146
    check(g4, 'BA', g, [])

1147

nv-dlasalle's avatar
nv-dlasalle committed
1148
@parametrize_idtype
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
def test_add_edges(idtype):
    # homogeneous graph
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    u = 0
    v = 1
    g = dgl.add_edges(g, u, v)
    assert g.device == F.ctx()
    assert g.number_of_nodes() == 3
    assert g.number_of_edges() == 3
    u = [0]
    v = [1]
    g = dgl.add_edges(g, u, v)
    assert g.device == F.ctx()
    assert g.number_of_nodes() == 3
    assert g.number_of_edges() == 4
    u = F.tensor(u, dtype=idtype)
    v = F.tensor(v, dtype=idtype)
    g = dgl.add_edges(g, u, v)
    assert g.device == F.ctx()
    assert g.number_of_nodes() == 3
    assert g.number_of_edges() == 5
    u, v = g.edges(form='uv', order='eid')
    assert F.array_equal(u, F.tensor([0, 1, 0, 0, 0], dtype=idtype))
1172
1173
1174
1175
1176
1177
1178
1179
1180
    assert F.array_equal(v, F.tensor([1, 2, 1, 1, 1], dtype=idtype))
    g = dgl.add_edges(g, [], [])
    g = dgl.add_edges(g, 0, [])
    g = dgl.add_edges(g, [], 0)
    assert g.device == F.ctx()
    assert g.number_of_nodes() == 3
    assert g.number_of_edges() == 5
    u, v = g.edges(form='uv', order='eid')
    assert F.array_equal(u, F.tensor([0, 1, 0, 0, 0], dtype=idtype))
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
    assert F.array_equal(v, F.tensor([1, 2, 1, 1, 1], dtype=idtype))

    # node id larger than current max node id
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    u = F.tensor([0, 1], dtype=idtype)
    v = F.tensor([2, 3], dtype=idtype)
    g = dgl.add_edges(g, u, v)
    assert g.number_of_nodes() == 4
    assert g.number_of_edges() == 4
    u, v = g.edges(form='uv', order='eid')
    assert F.array_equal(u, F.tensor([0, 1, 0, 1], dtype=idtype))
    assert F.array_equal(v, F.tensor([1, 2, 2, 3], dtype=idtype))

    # has data
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    g.ndata['h'] = F.copy_to(F.tensor([1, 1, 1], dtype=idtype), ctx=F.ctx())
    g.edata['h'] = F.copy_to(F.tensor([1, 1], dtype=idtype), ctx=F.ctx())
    u = F.tensor([0, 1], dtype=idtype)
    v = F.tensor([2, 3], dtype=idtype)
1200
1201
    e_feat = {'h': F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx()),
              'hh': F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
    g = dgl.add_edges(g, u, v, e_feat)
    assert g.number_of_nodes() == 4
    assert g.number_of_edges() == 4
    u, v = g.edges(form='uv', order='eid')
    assert F.array_equal(u, F.tensor([0, 1, 0, 1], dtype=idtype))
    assert F.array_equal(v, F.tensor([1, 2, 2, 3], dtype=idtype))
    assert F.array_equal(g.ndata['h'], F.tensor([1, 1, 1, 0], dtype=idtype))
    assert F.array_equal(g.edata['h'], F.tensor([1, 1, 2, 2], dtype=idtype))
    assert F.array_equal(g.edata['hh'], F.tensor([0, 0, 2, 2], dtype=idtype))

    # zero data graph
1213
    g = dgl.graph(([], []), num_nodes=0, idtype=idtype, device=F.ctx())
1214
1215
    u = F.tensor([0, 1], dtype=idtype)
    v = F.tensor([2, 2], dtype=idtype)
1216
1217
    e_feat = {'h': F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx()),
              'hh': F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
    g = dgl.add_edges(g, u, v, e_feat)
    assert g.number_of_nodes() == 3
    assert g.number_of_edges() == 2
    u, v = g.edges(form='uv', order='eid')
    assert F.array_equal(u, F.tensor([0, 1], dtype=idtype))
    assert F.array_equal(v, F.tensor([2, 2], dtype=idtype))
    assert F.array_equal(g.edata['h'], F.tensor([2, 2], dtype=idtype))
    assert F.array_equal(g.edata['hh'], F.tensor([2, 2], dtype=idtype))

    # bipartite graph
1228
    g = dgl.heterograph(
1229
1230
        {('user', 'plays', 'game'): ([0, 1], [1, 2])}, idtype=idtype,
        device=F.ctx())
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
    u = 0
    v = 1
    g = dgl.add_edges(g, u, v)
    assert g.device == F.ctx()
    assert g.number_of_nodes('user') == 2
    assert g.number_of_nodes('game') == 3
    assert g.number_of_edges() == 3
    u = [0]
    v = [1]
    g = dgl.add_edges(g, u, v)
    assert g.device == F.ctx()
    assert g.number_of_nodes('user') == 2
    assert g.number_of_nodes('game') == 3
    assert g.number_of_edges() == 4
    u = F.tensor(u, dtype=idtype)
    v = F.tensor(v, dtype=idtype)
    g = dgl.add_edges(g, u, v)
    assert g.device == F.ctx()
    assert g.number_of_nodes('user') == 2
    assert g.number_of_nodes('game') == 3
    assert g.number_of_edges() == 5
    u, v = g.edges(form='uv')
    assert F.array_equal(u, F.tensor([0, 1, 0, 0, 0], dtype=idtype))
    assert F.array_equal(v, F.tensor([1, 2, 1, 1, 1], dtype=idtype))

    # node id larger than current max node id
1257
    g = dgl.heterograph(
1258
1259
        {('user', 'plays', 'game'): ([0, 1], [1, 2])}, idtype=idtype,
        device=F.ctx())
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
    u = F.tensor([0, 2], dtype=idtype)
    v = F.tensor([2, 3], dtype=idtype)
    g = dgl.add_edges(g, u, v)
    assert g.device == F.ctx()
    assert g.number_of_nodes('user') == 3
    assert g.number_of_nodes('game') == 4
    assert g.number_of_edges() == 4
    u, v = g.edges(form='uv', order='eid')
    assert F.array_equal(u, F.tensor([0, 1, 0, 2], dtype=idtype))
    assert F.array_equal(v, F.tensor([1, 2, 2, 3], dtype=idtype))

    # has data
1272
    g = dgl.heterograph(
1273
1274
1275
1276
1277
1278
        {('user', 'plays', 'game'): ([0, 1], [1, 2])}, idtype=idtype,
        device=F.ctx())
    g.nodes['user'].data['h'] = F.copy_to(F.tensor([1, 1], dtype=idtype),
                                          ctx=F.ctx())
    g.nodes['game'].data['h'] = F.copy_to(F.tensor([2, 2, 2], dtype=idtype),
                                          ctx=F.ctx())
1279
1280
1281
    g.edata['h'] = F.copy_to(F.tensor([1, 1], dtype=idtype), ctx=F.ctx())
    u = F.tensor([0, 2], dtype=idtype)
    v = F.tensor([2, 3], dtype=idtype)
1282
1283
    e_feat = {'h': F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx()),
              'hh': F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}
1284
1285
1286
1287
1288
1289
1290
    g = dgl.add_edges(g, u, v, e_feat)
    assert g.number_of_nodes('user') == 3
    assert g.number_of_nodes('game') == 4
    assert g.number_of_edges() == 4
    u, v = g.edges(form='uv', order='eid')
    assert F.array_equal(u, F.tensor([0, 1, 0, 2], dtype=idtype))
    assert F.array_equal(v, F.tensor([1, 2, 2, 3], dtype=idtype))
1291
1292
1293
1294
    assert F.array_equal(g.nodes['user'].data['h'],
                         F.tensor([1, 1, 0], dtype=idtype))
    assert F.array_equal(g.nodes['game'].data['h'],
                         F.tensor([2, 2, 2, 0], dtype=idtype))
1295
1296
1297
1298
    assert F.array_equal(g.edata['h'], F.tensor([1, 1, 2, 2], dtype=idtype))
    assert F.array_equal(g.edata['hh'], F.tensor([0, 0, 2, 2], dtype=idtype))

    # heterogeneous graph
1299
    g = create_test_heterograph3(idtype)
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
    u = F.tensor([0, 2], dtype=idtype)
    v = F.tensor([2, 3], dtype=idtype)
    g = dgl.add_edges(g, u, v, etype='plays')
    assert g.number_of_nodes('user') == 3
    assert g.number_of_nodes('game') == 4
    assert g.number_of_nodes('developer') == 2
    assert g.number_of_edges('plays') == 6
    assert g.number_of_edges('develops') == 2
    u, v = g.edges(form='uv', order='eid', etype='plays')
    assert F.array_equal(u, F.tensor([0, 1, 1, 2, 0, 2], dtype=idtype))
    assert F.array_equal(v, F.tensor([0, 0, 1, 1, 2, 3], dtype=idtype))
1311
1312
1313
1314
1315
1316
    assert F.array_equal(g.nodes['user'].data['h'],
                         F.tensor([1, 1, 1], dtype=idtype))
    assert F.array_equal(g.nodes['game'].data['h'],
                         F.tensor([2, 2, 0, 0], dtype=idtype))
    assert F.array_equal(g.edges['plays'].data['h'],
                         F.tensor([1, 1, 1, 1, 0, 0], dtype=idtype))
1317
1318
1319
1320
1321

    # add with feature
    e_feat = {'h': F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}
    u = F.tensor([0, 2], dtype=idtype)
    v = F.tensor([2, 3], dtype=idtype)
1322
1323
    g.nodes['game'].data['h'] = F.copy_to(F.tensor([2, 2, 1, 1], dtype=idtype),
                                          ctx=F.ctx())
1324
1325
1326
1327
1328
1329
1330
1331
1332
    g = dgl.add_edges(g, u, v, data=e_feat, etype='develops')
    assert g.number_of_nodes('user') == 3
    assert g.number_of_nodes('game') == 4
    assert g.number_of_nodes('developer') == 3
    assert g.number_of_edges('plays') == 6
    assert g.number_of_edges('develops') == 4
    u, v = g.edges(form='uv', order='eid', etype='develops')
    assert F.array_equal(u, F.tensor([0, 1, 0, 2], dtype=idtype))
    assert F.array_equal(v, F.tensor([0, 1, 2, 3], dtype=idtype))
1333
1334
1335
1336
1337
1338
1339
    assert F.array_equal(g.nodes['developer'].data['h'],
                         F.tensor([3, 3, 0], dtype=idtype))
    assert F.array_equal(g.nodes['game'].data['h'],
                         F.tensor([2, 2, 1, 1], dtype=idtype))
    assert F.array_equal(g.edges['develops'].data['h'],
                         F.tensor([0, 0, 2, 2], dtype=idtype))

1340

nv-dlasalle's avatar
nv-dlasalle committed
1341
@parametrize_idtype
1342
1343
1344
def test_add_nodes(idtype):
    # homogeneous Graphs
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
1345
    g.ndata['h'] = F.copy_to(F.tensor([1, 1, 1], dtype=idtype), ctx=F.ctx())
1346
1347
1348
1349
1350
1351
    new_g = dgl.add_nodes(g, 1)
    assert g.number_of_nodes() == 3
    assert new_g.number_of_nodes() == 4
    assert F.array_equal(new_g.ndata['h'], F.tensor([1, 1, 1, 0], dtype=idtype))

    # zero node graph
1352
    g = dgl.graph(([], []), num_nodes=3, idtype=idtype, device=F.ctx())
1353
1354
1355
    g.ndata['h'] = F.copy_to(F.tensor([1, 1, 1], dtype=idtype), ctx=F.ctx())
    g = dgl.add_nodes(g, 1, data={
        'h': F.copy_to(F.tensor([2], dtype=idtype), ctx=F.ctx())})
1356
1357
1358
1359
    assert g.number_of_nodes() == 4
    assert F.array_equal(g.ndata['h'], F.tensor([1, 1, 1, 2], dtype=idtype))

    # bipartite graph
1360
    g = dgl.heterograph(
1361
1362
1363
1364
1365
        {('user', 'plays', 'game'): ([0, 1], [1, 2])}, idtype=idtype,
        device=F.ctx())
    g = dgl.add_nodes(g, 2, data={
        'h': F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())},
                      ntype='user')
1366
1367
    assert g.number_of_nodes('user') == 4
    assert g.number_of_nodes('game') == 3
1368
1369
    assert F.array_equal(g.nodes['user'].data['h'],
                         F.tensor([0, 0, 2, 2], dtype=idtype))
1370
1371
1372
1373
1374
    g = dgl.add_nodes(g, 2, ntype='game')
    assert g.number_of_nodes('user') == 4
    assert g.number_of_nodes('game') == 5

    # heterogeneous graph
1375
    g = create_test_heterograph3(idtype)
1376
    g = dgl.add_nodes(g, 1, ntype='user')
1377
1378
1379
    g = dgl.add_nodes(g, 2, data={
        'h': F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())},
                      ntype='game')
1380
1381
1382
    assert g.number_of_nodes('user') == 4
    assert g.number_of_nodes('game') == 4
    assert g.number_of_nodes('developer') == 2
1383
1384
1385
1386
1387
    assert F.array_equal(g.nodes['user'].data['h'],
                         F.tensor([1, 1, 1, 0], dtype=idtype))
    assert F.array_equal(g.nodes['game'].data['h'],
                         F.tensor([2, 2, 2, 2], dtype=idtype))

1388

nv-dlasalle's avatar
nv-dlasalle committed
1389
@parametrize_idtype
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
def test_remove_edges(idtype):
    # homogeneous Graphs
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    e = 0
    g = dgl.remove_edges(g, e)
    assert g.number_of_edges() == 1
    u, v = g.edges(form='uv', order='eid')
    assert F.array_equal(u, F.tensor([1], dtype=idtype))
    assert F.array_equal(v, F.tensor([2], dtype=idtype))
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    e = [0]
    g = dgl.remove_edges(g, e)
    assert g.number_of_edges() == 1
    u, v = g.edges(form='uv', order='eid')
    assert F.array_equal(u, F.tensor([1], dtype=idtype))
    assert F.array_equal(v, F.tensor([2], dtype=idtype))
    e = F.tensor([0], dtype=idtype)
    g = dgl.remove_edges(g, e)
    assert g.number_of_edges() == 0

    # has node data
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    g.ndata['h'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
    g = dgl.remove_edges(g, 1)
    assert g.number_of_edges() == 1
    assert F.array_equal(g.ndata['h'], F.tensor([1, 2, 3], dtype=idtype))

    # has edge data
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    g.edata['h'] = F.copy_to(F.tensor([1, 2], dtype=idtype), ctx=F.ctx())
    g = dgl.remove_edges(g, 0)
    assert g.number_of_edges() == 1
    assert F.array_equal(g.edata['h'], F.tensor([2], dtype=idtype))

    # invalid eid
    assert_fail = False
    try:
        g = dgl.remove_edges(g, 1)
    except:
        assert_fail = True
    assert assert_fail

    # bipartite graph
1433
    g = dgl.heterograph(
1434
1435
        {('user', 'plays', 'game'): ([0, 1], [1, 2])}, idtype=idtype,
        device=F.ctx())
1436
1437
1438
1439
1440
1441
    e = 0
    g = dgl.remove_edges(g, e)
    assert g.number_of_edges() == 1
    u, v = g.edges(form='uv', order='eid')
    assert F.array_equal(u, F.tensor([1], dtype=idtype))
    assert F.array_equal(v, F.tensor([2], dtype=idtype))
1442
    g = dgl.heterograph(
1443
1444
        {('user', 'plays', 'game'): ([0, 1], [1, 2])}, idtype=idtype,
        device=F.ctx())
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
    e = [0]
    g = dgl.remove_edges(g, e)
    assert g.number_of_edges() == 1
    u, v = g.edges(form='uv', order='eid')
    assert F.array_equal(u, F.tensor([1], dtype=idtype))
    assert F.array_equal(v, F.tensor([2], dtype=idtype))
    e = F.tensor([0], dtype=idtype)
    g = dgl.remove_edges(g, e)
    assert g.number_of_edges() == 0

    # has data
1456
    g = dgl.heterograph(
1457
1458
1459
1460
1461
1462
        {('user', 'plays', 'game'): ([0, 1], [1, 2])}, idtype=idtype,
        device=F.ctx())
    g.nodes['user'].data['h'] = F.copy_to(F.tensor([1, 1], dtype=idtype),
                                          ctx=F.ctx())
    g.nodes['game'].data['h'] = F.copy_to(F.tensor([2, 2, 2], dtype=idtype),
                                          ctx=F.ctx())
1463
1464
1465
    g.edata['h'] = F.copy_to(F.tensor([1, 2], dtype=idtype), ctx=F.ctx())
    g = dgl.remove_edges(g, 1)
    assert g.number_of_edges() == 1
1466
1467
1468
1469
    assert F.array_equal(g.nodes['user'].data['h'],
                         F.tensor([1, 1], dtype=idtype))
    assert F.array_equal(g.nodes['game'].data['h'],
                         F.tensor([2, 2, 2], dtype=idtype))
1470
1471
1472
    assert F.array_equal(g.edata['h'], F.tensor([1], dtype=idtype))

    # heterogeneous graph
1473
    g = create_test_heterograph3(idtype)
1474
1475
    g.edges['plays'].data['h'] = F.copy_to(F.tensor([1, 2, 3, 4], dtype=idtype),
                                           ctx=F.ctx())
1476
1477
1478
1479
1480
    g = dgl.remove_edges(g, 1, etype='plays')
    assert g.number_of_edges('plays') == 3
    u, v = g.edges(form='uv', order='eid', etype='plays')
    assert F.array_equal(u, F.tensor([0, 1, 2], dtype=idtype))
    assert F.array_equal(v, F.tensor([0, 1, 1], dtype=idtype))
1481
1482
    assert F.array_equal(g.edges['plays'].data['h'],
                         F.tensor([1, 3, 4], dtype=idtype))
1483
1484
1485
    # remove all edges of 'develops'
    g = dgl.remove_edges(g, [0, 1], etype='develops')
    assert g.number_of_edges('develops') == 0
1486
1487
1488
1489
1490
1491
    assert F.array_equal(g.nodes['user'].data['h'],
                         F.tensor([1, 1, 1], dtype=idtype))
    assert F.array_equal(g.nodes['game'].data['h'],
                         F.tensor([2, 2], dtype=idtype))
    assert F.array_equal(g.nodes['developer'].data['h'],
                         F.tensor([3, 3], dtype=idtype))
1492

1493
1494
1495
1496
1497
1498
1499
1500
1501
    # batched graph
    ctx = F.ctx()
    g1 = dgl.graph(([0, 1], [1, 2]), num_nodes=5, idtype=idtype, device=ctx)
    g2 = dgl.graph(([], []), idtype=idtype, device=ctx)
    g3 = dgl.graph(([2, 3, 4], [3, 2, 1]), idtype=idtype, device=ctx)
    bg = dgl.batch([g1, g2, g3])
    bg_r = dgl.remove_edges(bg, 2)
    assert bg.batch_size == bg_r.batch_size
    assert F.array_equal(bg.batch_num_nodes(), bg_r.batch_num_nodes())
1502
1503
    assert F.array_equal(bg_r.batch_num_edges(),
                         F.tensor([2, 0, 2], dtype=F.int64))
1504
1505
1506
1507

    bg_r = dgl.remove_edges(bg, [0, 2])
    assert bg.batch_size == bg_r.batch_size
    assert F.array_equal(bg.batch_num_nodes(), bg_r.batch_num_nodes())
1508
1509
    assert F.array_equal(bg_r.batch_num_edges(),
                         F.tensor([1, 0, 2], dtype=F.int64))
1510
1511
1512
1513

    bg_r = dgl.remove_edges(bg, F.tensor([0, 2], dtype=idtype))
    assert bg.batch_size == bg_r.batch_size
    assert F.array_equal(bg.batch_num_nodes(), bg_r.batch_num_nodes())
1514
1515
    assert F.array_equal(bg_r.batch_num_edges(),
                         F.tensor([1, 0, 2], dtype=F.int64))
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535

    # batched heterogeneous graph
    g1 = dgl.heterograph({
        ('user', 'follows', 'user'): ([0, 1], [1, 2]),
        ('user', 'plays', 'game'): ([1, 3], [0, 1])
    }, num_nodes_dict={'user': 4, 'game': 3}, idtype=idtype, device=ctx)
    g2 = dgl.heterograph({
        ('user', 'follows', 'user'): ([0, 2], [3, 4]),
        ('user', 'plays', 'game'): ([], [])
    }, num_nodes_dict={'user': 6, 'game': 2}, idtype=idtype, device=ctx)
    g3 = dgl.heterograph({
        ('user', 'follows', 'user'): ([], []),
        ('user', 'plays', 'game'): ([1, 2], [1, 2])
    }, idtype=idtype, device=ctx)
    bg = dgl.batch([g1, g2, g3])
    bg_r = dgl.remove_edges(bg, 1, etype='follows')
    assert bg.batch_size == bg_r.batch_size
    ntypes = bg.ntypes
    for nty in ntypes:
        assert F.array_equal(bg.batch_num_nodes(nty), bg_r.batch_num_nodes(nty))
1536
1537
1538
1539
    assert F.array_equal(bg_r.batch_num_edges('follows'),
                         F.tensor([1, 2, 0], dtype=F.int64))
    assert F.array_equal(bg_r.batch_num_edges('plays'),
                         bg.batch_num_edges('plays'))
1540
1541
1542
1543
1544

    bg_r = dgl.remove_edges(bg, 2, etype='plays')
    assert bg.batch_size == bg_r.batch_size
    for nty in ntypes:
        assert F.array_equal(bg.batch_num_nodes(nty), bg_r.batch_num_nodes(nty))
1545
1546
1547
1548
    assert F.array_equal(bg.batch_num_edges('follows'),
                         bg_r.batch_num_edges('follows'))
    assert F.array_equal(bg_r.batch_num_edges('plays'),
                         F.tensor([2, 0, 1], dtype=F.int64))
1549
1550
1551
1552
1553

    bg_r = dgl.remove_edges(bg, [0, 1, 3], etype='follows')
    assert bg.batch_size == bg_r.batch_size
    for nty in ntypes:
        assert F.array_equal(bg.batch_num_nodes(nty), bg_r.batch_num_nodes(nty))
1554
1555
1556
1557
    assert F.array_equal(bg_r.batch_num_edges('follows'),
                         F.tensor([0, 1, 0], dtype=F.int64))
    assert F.array_equal(bg.batch_num_edges('plays'),
                         bg_r.batch_num_edges('plays'))
1558
1559
1560
1561
1562

    bg_r = dgl.remove_edges(bg, [1, 2], etype='plays')
    assert bg.batch_size == bg_r.batch_size
    for nty in ntypes:
        assert F.array_equal(bg.batch_num_nodes(nty), bg_r.batch_num_nodes(nty))
1563
1564
1565
1566
    assert F.array_equal(bg.batch_num_edges('follows'),
                         bg_r.batch_num_edges('follows'))
    assert F.array_equal(bg_r.batch_num_edges('plays'),
                         F.tensor([1, 0, 1], dtype=F.int64))
1567

1568
1569
    bg_r = dgl.remove_edges(bg, F.tensor([0, 1, 3], dtype=idtype),
                            etype='follows')
1570
1571
1572
    assert bg.batch_size == bg_r.batch_size
    for nty in ntypes:
        assert F.array_equal(bg.batch_num_nodes(nty), bg_r.batch_num_nodes(nty))
1573
1574
1575
1576
    assert F.array_equal(bg_r.batch_num_edges('follows'),
                         F.tensor([0, 1, 0], dtype=F.int64))
    assert F.array_equal(bg.batch_num_edges('plays'),
                         bg_r.batch_num_edges('plays'))
1577
1578
1579
1580
1581

    bg_r = dgl.remove_edges(bg, F.tensor([1, 2], dtype=idtype), etype='plays')
    assert bg.batch_size == bg_r.batch_size
    for nty in ntypes:
        assert F.array_equal(bg.batch_num_nodes(nty), bg_r.batch_num_nodes(nty))
1582
1583
1584
1585
1586
    assert F.array_equal(bg.batch_num_edges('follows'),
                         bg_r.batch_num_edges('follows'))
    assert F.array_equal(bg_r.batch_num_edges('plays'),
                         F.tensor([1, 0, 1], dtype=F.int64))

1587

nv-dlasalle's avatar
nv-dlasalle committed
1588
@parametrize_idtype
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
def test_remove_nodes(idtype):
    # homogeneous Graphs
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    n = 0
    g = dgl.remove_nodes(g, n)
    assert g.number_of_nodes() == 2
    assert g.number_of_edges() == 1
    u, v = g.edges(form='uv', order='eid')
    assert F.array_equal(u, F.tensor([0], dtype=idtype))
    assert F.array_equal(v, F.tensor([1], dtype=idtype))
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    n = [1]
    g = dgl.remove_nodes(g, n)
    assert g.number_of_nodes() == 2
    assert g.number_of_edges() == 0
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    n = F.tensor([2], dtype=idtype)
    g = dgl.remove_nodes(g, n)
    assert g.number_of_nodes() == 2
    assert g.number_of_edges() == 1
    u, v = g.edges(form='uv', order='eid')
    assert F.array_equal(u, F.tensor([0], dtype=idtype))
    assert F.array_equal(v, F.tensor([1], dtype=idtype))

    # invalid nid
    assert_fail = False
    try:
        g.remove_nodes(3)
    except:
        assert_fail = True
    assert assert_fail

    # has node and edge data
    g = dgl.graph(([0, 0, 2], [0, 1, 2]), idtype=idtype, device=F.ctx())
    g.ndata['hv'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
    g.edata['he'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
    g = dgl.remove_nodes(g, F.tensor([0], dtype=idtype))
    assert g.number_of_nodes() == 2
    assert g.number_of_edges() == 1
    u, v = g.edges(form='uv', order='eid')
    assert F.array_equal(u, F.tensor([1], dtype=idtype))
    assert F.array_equal(v, F.tensor([1], dtype=idtype))
    assert F.array_equal(g.ndata['hv'], F.tensor([2, 3], dtype=idtype))
    assert F.array_equal(g.edata['he'], F.tensor([3], dtype=idtype))

    # node id larger than current max node id
1635
    g = dgl.heterograph(
1636
1637
        {('user', 'plays', 'game'): ([0, 1], [1, 2])}, idtype=idtype,
        device=F.ctx())
1638
1639
1640
1641
1642
1643
1644
1645
    n = 0
    g = dgl.remove_nodes(g, n, ntype='user')
    assert g.number_of_nodes('user') == 1
    assert g.number_of_nodes('game') == 3
    assert g.number_of_edges() == 1
    u, v = g.edges(form='uv', order='eid')
    assert F.array_equal(u, F.tensor([0], dtype=idtype))
    assert F.array_equal(v, F.tensor([2], dtype=idtype))
1646
    g = dgl.heterograph(
1647
1648
        {('user', 'plays', 'game'): ([0, 1], [1, 2])}, idtype=idtype,
        device=F.ctx())
1649
1650
1651
1652
1653
1654
1655
1656
    n = [1]
    g = dgl.remove_nodes(g, n, ntype='user')
    assert g.number_of_nodes('user') == 1
    assert g.number_of_nodes('game') == 3
    assert g.number_of_edges() == 1
    u, v = g.edges(form='uv', order='eid')
    assert F.array_equal(u, F.tensor([0], dtype=idtype))
    assert F.array_equal(v, F.tensor([1], dtype=idtype))
1657
    g = dgl.heterograph(
1658
1659
        {('user', 'plays', 'game'): ([0, 1], [1, 2])}, idtype=idtype,
        device=F.ctx())
1660
1661
1662
1663
1664
1665
1666
    n = F.tensor([0], dtype=idtype)
    g = dgl.remove_nodes(g, n, ntype='game')
    assert g.number_of_nodes('user') == 2
    assert g.number_of_nodes('game') == 2
    assert g.number_of_edges() == 2
    u, v = g.edges(form='uv', order='eid')
    assert F.array_equal(u, F.tensor([0, 1], dtype=idtype))
1667
    assert F.array_equal(v, F.tensor([0, 1], dtype=idtype))
1668
1669

    # heterogeneous graph
1670
    g = create_test_heterograph3(idtype)
1671
1672
    g.edges['plays'].data['h'] = F.copy_to(F.tensor([1, 2, 3, 4], dtype=idtype),
                                           ctx=F.ctx())
1673
1674
1675
1676
1677
1678
    g = dgl.remove_nodes(g, 0, ntype='game')
    assert g.number_of_nodes('user') == 3
    assert g.number_of_nodes('game') == 1
    assert g.number_of_nodes('developer') == 2
    assert g.number_of_edges('plays') == 2
    assert g.number_of_edges('develops') == 1
1679
1680
    assert F.array_equal(g.nodes['user'].data['h'],
                         F.tensor([1, 1, 1], dtype=idtype))
1681
    assert F.array_equal(g.nodes['game'].data['h'], F.tensor([2], dtype=idtype))
1682
1683
    assert F.array_equal(g.nodes['developer'].data['h'],
                         F.tensor([3, 3], dtype=idtype))
1684
1685
1686
    u, v = g.edges(form='uv', order='eid', etype='plays')
    assert F.array_equal(u, F.tensor([1, 2], dtype=idtype))
    assert F.array_equal(v, F.tensor([0, 0], dtype=idtype))
1687
1688
    assert F.array_equal(g.edges['plays'].data['h'],
                         F.tensor([3, 4], dtype=idtype))
1689
1690
1691
1692
    u, v = g.edges(form='uv', order='eid', etype='develops')
    assert F.array_equal(u, F.tensor([1], dtype=idtype))
    assert F.array_equal(v, F.tensor([0], dtype=idtype))

1693
1694
1695
1696
1697
1698
1699
1700
    # batched graph
    ctx = F.ctx()
    g1 = dgl.graph(([0, 1], [1, 2]), num_nodes=5, idtype=idtype, device=ctx)
    g2 = dgl.graph(([], []), idtype=idtype, device=ctx)
    g3 = dgl.graph(([2, 3, 4], [3, 2, 1]), idtype=idtype, device=ctx)
    bg = dgl.batch([g1, g2, g3])
    bg_r = dgl.remove_nodes(bg, 1)
    assert bg_r.batch_size == bg.batch_size
1701
1702
1703
1704
    assert F.array_equal(bg_r.batch_num_nodes(),
                         F.tensor([4, 0, 5], dtype=F.int64))
    assert F.array_equal(bg_r.batch_num_edges(),
                         F.tensor([0, 0, 3], dtype=F.int64))
1705
1706
1707

    bg_r = dgl.remove_nodes(bg, [1, 7])
    assert bg_r.batch_size == bg.batch_size
1708
1709
1710
1711
    assert F.array_equal(bg_r.batch_num_nodes(),
                         F.tensor([4, 0, 4], dtype=F.int64))
    assert F.array_equal(bg_r.batch_num_edges(),
                         F.tensor([0, 0, 1], dtype=F.int64))
1712
1713
1714

    bg_r = dgl.remove_nodes(bg, F.tensor([1, 7], dtype=idtype))
    assert bg_r.batch_size == bg.batch_size
1715
1716
1717
1718
    assert F.array_equal(bg_r.batch_num_nodes(),
                         F.tensor([4, 0, 4], dtype=F.int64))
    assert F.array_equal(bg_r.batch_num_edges(),
                         F.tensor([0, 0, 1], dtype=F.int64))
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735

    # batched heterogeneous graph
    g1 = dgl.heterograph({
        ('user', 'follows', 'user'): ([0, 1], [1, 2]),
        ('user', 'plays', 'game'): ([1, 3], [0, 1])
    }, num_nodes_dict={'user': 4, 'game': 3}, idtype=idtype, device=ctx)
    g2 = dgl.heterograph({
        ('user', 'follows', 'user'): ([0, 2], [3, 4]),
        ('user', 'plays', 'game'): ([], [])
    }, num_nodes_dict={'user': 6, 'game': 2}, idtype=idtype, device=ctx)
    g3 = dgl.heterograph({
        ('user', 'follows', 'user'): ([], []),
        ('user', 'plays', 'game'): ([1, 2], [1, 2])
    }, idtype=idtype, device=ctx)
    bg = dgl.batch([g1, g2, g3])
    bg_r = dgl.remove_nodes(bg, 1, ntype='user')
    assert bg_r.batch_size == bg.batch_size
1736
1737
1738
1739
1740
1741
1742
1743
    assert F.array_equal(bg_r.batch_num_nodes('user'),
                         F.tensor([3, 6, 3], dtype=F.int64))
    assert F.array_equal(bg.batch_num_nodes('game'),
                         bg_r.batch_num_nodes('game'))
    assert F.array_equal(bg_r.batch_num_edges('follows'),
                         F.tensor([0, 2, 0], dtype=F.int64))
    assert F.array_equal(bg_r.batch_num_edges('plays'),
                         F.tensor([1, 0, 2], dtype=F.int64))
1744
1745
1746

    bg_r = dgl.remove_nodes(bg, 6, ntype='game')
    assert bg_r.batch_size == bg.batch_size
1747
1748
1749
1750
1751
1752
1753
1754
    assert F.array_equal(bg.batch_num_nodes('user'),
                         bg_r.batch_num_nodes('user'))
    assert F.array_equal(bg_r.batch_num_nodes('game'),
                         F.tensor([3, 2, 2], dtype=F.int64))
    assert F.array_equal(bg.batch_num_edges('follows'),
                         bg_r.batch_num_edges('follows'))
    assert F.array_equal(bg_r.batch_num_edges('plays'),
                         F.tensor([2, 0, 1], dtype=F.int64))
1755
1756
1757

    bg_r = dgl.remove_nodes(bg, [1, 5, 6, 11], ntype='user')
    assert bg_r.batch_size == bg.batch_size
1758
1759
1760
1761
1762
1763
1764
1765
    assert F.array_equal(bg_r.batch_num_nodes('user'),
                         F.tensor([3, 4, 2], dtype=F.int64))
    assert F.array_equal(bg.batch_num_nodes('game'),
                         bg_r.batch_num_nodes('game'))
    assert F.array_equal(bg_r.batch_num_edges('follows'),
                         F.tensor([0, 1, 0], dtype=F.int64))
    assert F.array_equal(bg_r.batch_num_edges('plays'),
                         F.tensor([1, 0, 1], dtype=F.int64))
1766
1767
1768

    bg_r = dgl.remove_nodes(bg, [0, 3, 4, 7], ntype='game')
    assert bg_r.batch_size == bg.batch_size
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
    assert F.array_equal(bg.batch_num_nodes('user'),
                         bg_r.batch_num_nodes('user'))
    assert F.array_equal(bg_r.batch_num_nodes('game'),
                         F.tensor([2, 0, 2], dtype=F.int64))
    assert F.array_equal(bg.batch_num_edges('follows'),
                         bg_r.batch_num_edges('follows'))
    assert F.array_equal(bg_r.batch_num_edges('plays'),
                         F.tensor([1, 0, 1], dtype=F.int64))

    bg_r = dgl.remove_nodes(bg, F.tensor([1, 5, 6, 11], dtype=idtype),
                            ntype='user')
1780
    assert bg_r.batch_size == bg.batch_size
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
    assert F.array_equal(bg_r.batch_num_nodes('user'),
                         F.tensor([3, 4, 2], dtype=F.int64))
    assert F.array_equal(bg.batch_num_nodes('game'),
                         bg_r.batch_num_nodes('game'))
    assert F.array_equal(bg_r.batch_num_edges('follows'),
                         F.tensor([0, 1, 0], dtype=F.int64))
    assert F.array_equal(bg_r.batch_num_edges('plays'),
                         F.tensor([1, 0, 1], dtype=F.int64))

    bg_r = dgl.remove_nodes(bg, F.tensor([0, 3, 4, 7], dtype=idtype),
                            ntype='game')
1792
    assert bg_r.batch_size == bg.batch_size
1793
1794
1795
1796
1797
1798
1799
1800
1801
    assert F.array_equal(bg.batch_num_nodes('user'),
                         bg_r.batch_num_nodes('user'))
    assert F.array_equal(bg_r.batch_num_nodes('game'),
                         F.tensor([2, 0, 2], dtype=F.int64))
    assert F.array_equal(bg.batch_num_edges('follows'),
                         bg_r.batch_num_edges('follows'))
    assert F.array_equal(bg_r.batch_num_edges('plays'),
                         F.tensor([1, 0, 1], dtype=F.int64))

1802

nv-dlasalle's avatar
nv-dlasalle committed
1803
@parametrize_idtype
1804
1805
def test_add_selfloop(idtype):
    # homogeneous graph
1806
1807

    # test for fill_data is float
1808
1809
    g = dgl.graph(([0, 0, 2], [2, 1, 0]), idtype=idtype, device=F.ctx())
    g.edata['he'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
1810
1811
    g.edata['he1'] = F.copy_to(F.tensor([[0., 1.], [2., 3.], [4., 5.]]),
                               ctx=F.ctx())
1812
1813
1814
1815
1816
1817
1818
    g.ndata['hn'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
    g = dgl.add_self_loop(g)
    assert g.number_of_nodes() == 3
    assert g.number_of_edges() == 6
    u, v = g.edges(form='uv', order='eid')
    assert F.array_equal(u, F.tensor([0, 0, 2, 0, 1, 2], dtype=idtype))
    assert F.array_equal(v, F.tensor([2, 1, 0, 0, 1, 2], dtype=idtype))
1819
1820
    assert F.array_equal(g.edata['he'],
                         F.tensor([1, 2, 3, 1, 1, 1], dtype=idtype))
1821
    assert F.array_equal(g.edata['he1'], F.tensor([[0., 1.], [2., 3.], [4., 5.],
1822
1823
                                                   [1., 1.], [1., 1.],
                                                   [1., 1.]]))
1824
1825
1826
1827

    # test for fill_data is int
    g = dgl.graph(([0, 0, 2], [2, 1, 0]), idtype=idtype, device=F.ctx())
    g.edata['he'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
1828
1829
    g.edata['he1'] = F.copy_to(F.tensor([[0, 1], [2, 3], [4, 5]], dtype=idtype),
                               ctx=F.ctx())
1830
1831
1832
1833
1834
1835
1836
    g.ndata['hn'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
    g = dgl.add_self_loop(g, fill_data=1)
    assert g.number_of_nodes() == 3
    assert g.number_of_edges() == 6
    u, v = g.edges(form='uv', order='eid')
    assert F.array_equal(u, F.tensor([0, 0, 2, 0, 1, 2], dtype=idtype))
    assert F.array_equal(v, F.tensor([2, 1, 0, 0, 1, 2], dtype=idtype))
1837
1838
    assert F.array_equal(g.edata['he'],
                         F.tensor([1, 2, 3, 1, 1, 1], dtype=idtype))
1839
    assert F.array_equal(g.edata['he1'], F.tensor([[0, 1], [2, 3], [4, 5],
1840
1841
                                                   [1, 1], [1, 1], [1, 1]],
                                                  dtype=idtype))
1842
1843
1844
1845

    # test for fill_data is str
    g = dgl.graph(([0, 0, 2], [2, 1, 0]), idtype=idtype, device=F.ctx())
    g.edata['he'] = F.copy_to(F.tensor([1., 2., 3.]), ctx=F.ctx())
1846
1847
    g.edata['he1'] = F.copy_to(F.tensor([[0., 1.], [2., 3.], [4., 5.]]),
                               ctx=F.ctx())
1848
1849
1850
1851
1852
1853
1854
1855
1856
    g.ndata['hn'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
    g = dgl.add_self_loop(g, fill_data='sum')
    assert g.number_of_nodes() == 3
    assert g.number_of_edges() == 6
    u, v = g.edges(form='uv', order='eid')
    assert F.array_equal(u, F.tensor([0, 0, 2, 0, 1, 2], dtype=idtype))
    assert F.array_equal(v, F.tensor([2, 1, 0, 0, 1, 2], dtype=idtype))
    assert F.array_equal(g.edata['he'], F.tensor([1., 2., 3., 3., 2., 1.]))
    assert F.array_equal(g.edata['he1'], F.tensor([[0., 1.], [2., 3.], [4., 5.],
1857
1858
                                                   [4., 5.], [2., 3.],
                                                   [0., 1.]]))
1859
1860

    # bipartite graph
1861
    g = dgl.heterograph(
1862
1863
        {('user', 'plays', 'game'): ([0, 1, 2], [1, 2, 2])}, idtype=idtype,
        device=F.ctx())
1864
1865
1866
1867
1868
1869
1870
1871
    # nothing will happend
    raise_error = False
    try:
        g = dgl.add_self_loop(g)
    except:
        raise_error = True
    assert raise_error

1872
    # test for fill_data is float
1873
    g = create_test_heterograph5(idtype)
1874
1875
    g.edges['follows'].data['h1'] = F.copy_to(F.tensor([[0., 1.], [1., 2.]]),
                                              ctx=F.ctx())
1876
1877
1878
1879
1880
1881
1882
1883
    g = dgl.add_self_loop(g, etype='follows')
    assert g.number_of_nodes('user') == 3
    assert g.number_of_nodes('game') == 2
    assert g.number_of_edges('follows') == 5
    assert g.number_of_edges('plays') == 2
    u, v = g.edges(form='uv', order='eid', etype='follows')
    assert F.array_equal(u, F.tensor([1, 2, 0, 1, 2], dtype=idtype))
    assert F.array_equal(v, F.tensor([0, 1, 0, 1, 2], dtype=idtype))
1884
1885
1886
1887
1888
1889
1890
    assert F.array_equal(g.edges['follows'].data['h'],
                         F.tensor([1, 2, 1, 1, 1], dtype=idtype))
    assert F.array_equal(g.edges['follows'].data['h1'],
                         F.tensor([[0., 1.], [1., 2.], [1., 1.],
                                   [1., 1.], [1., 1.]]))
    assert F.array_equal(g.edges['plays'].data['h'],
                         F.tensor([1, 2], dtype=idtype))
1891
1892
1893

    # test for fill_data is int
    g = create_test_heterograph5(idtype)
1894
1895
    g.edges['follows'].data['h1'] = F.copy_to(
        F.tensor([[0, 1], [1, 2]], dtype=idtype), ctx=F.ctx())
1896
1897
1898
1899
1900
1901
1902
1903
    g = dgl.add_self_loop(g, fill_data=1, etype='follows')
    assert g.number_of_nodes('user') == 3
    assert g.number_of_nodes('game') == 2
    assert g.number_of_edges('follows') == 5
    assert g.number_of_edges('plays') == 2
    u, v = g.edges(form='uv', order='eid', etype='follows')
    assert F.array_equal(u, F.tensor([1, 2, 0, 1, 2], dtype=idtype))
    assert F.array_equal(v, F.tensor([0, 1, 0, 1, 2], dtype=idtype))
1904
1905
1906
1907
1908
1909
1910
    assert F.array_equal(g.edges['follows'].data['h'],
                         F.tensor([1, 2, 1, 1, 1], dtype=idtype))
    assert F.array_equal(g.edges['follows'].data['h1'],
                         F.tensor([[0, 1], [1, 2], [1, 1],
                                   [1, 1], [1, 1]], dtype=idtype))
    assert F.array_equal(g.edges['plays'].data['h'],
                         F.tensor([1, 2], dtype=idtype))
1911

1912
1913
1914
1915
1916
1917
1918
    # test for fill_data is str
    g = dgl.heterograph({
        ('user', 'follows', 'user'): (F.tensor([1, 2], dtype=idtype),
                                      F.tensor([0, 1], dtype=idtype)),
        ('user', 'plays', 'game'): (F.tensor([0, 1], dtype=idtype),
                                    F.tensor([0, 1], dtype=idtype))},
        idtype=idtype, device=F.ctx())
1919
1920
1921
1922
    g.nodes['user'].data['h'] = F.copy_to(F.tensor([1, 1, 1], dtype=idtype),
                                          ctx=F.ctx())
    g.nodes['game'].data['h'] = F.copy_to(F.tensor([2, 2], dtype=idtype),
                                          ctx=F.ctx())
1923
    g.edges['follows'].data['h'] = F.copy_to(F.tensor([1., 2.]), ctx=F.ctx())
1924
1925
    g.edges['follows'].data['h1'] = F.copy_to(F.tensor([[0., 1.], [1., 2.]]),
                                              ctx=F.ctx())
1926
1927
1928
1929
1930
1931
1932
1933
1934
    g.edges['plays'].data['h'] = F.copy_to(F.tensor([1., 2.]), ctx=F.ctx())
    g = dgl.add_self_loop(g, fill_data='mean', etype='follows')
    assert g.number_of_nodes('user') == 3
    assert g.number_of_nodes('game') == 2
    assert g.number_of_edges('follows') == 5
    assert g.number_of_edges('plays') == 2
    u, v = g.edges(form='uv', order='eid', etype='follows')
    assert F.array_equal(u, F.tensor([1, 2, 0, 1, 2], dtype=idtype))
    assert F.array_equal(v, F.tensor([0, 1, 0, 1, 2], dtype=idtype))
1935
1936
1937
1938
1939
    assert F.array_equal(g.edges['follows'].data['h'],
                         F.tensor([1., 2., 1., 2., 0.]))
    assert F.array_equal(g.edges['follows'].data['h1'],
                         F.tensor([[0., 1.], [1., 2.], [0., 1.],
                                   [1., 2.], [0., 0.]]))
1940
1941
    assert F.array_equal(g.edges['plays'].data['h'], F.tensor([1., 2.]))

1942
1943
1944
1945
1946
1947
1948
    raise_error = False
    try:
        g = dgl.add_self_loop(g, etype='plays')
    except:
        raise_error = True
    assert raise_error

1949

nv-dlasalle's avatar
nv-dlasalle committed
1950
@parametrize_idtype
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
def test_remove_selfloop(idtype):
    # homogeneous graph
    g = dgl.graph(([0, 0, 0, 1], [1, 0, 0, 2]), idtype=idtype, device=F.ctx())
    g.edata['he'] = F.copy_to(F.tensor([1, 2, 3, 4], dtype=idtype), ctx=F.ctx())
    g = dgl.remove_self_loop(g)
    assert g.number_of_nodes() == 3
    assert g.number_of_edges() == 2
    assert F.array_equal(g.edata['he'], F.tensor([1, 4], dtype=idtype))

    # bipartite graph
1961
    g = dgl.heterograph(
1962
1963
        {('user', 'plays', 'game'): ([0, 1, 2], [1, 2, 2])}, idtype=idtype,
        device=F.ctx())
1964
1965
1966
1967
1968
1969
1970
1971
    # nothing will happend
    raise_error = False
    try:
        g = dgl.remove_self_loop(g, etype='plays')
    except:
        raise_error = True
    assert raise_error

1972
    g = create_test_heterograph4(idtype)
1973
1974
1975
1976
1977
1978
1979
1980
    g = dgl.remove_self_loop(g, etype='follows')
    assert g.number_of_nodes('user') == 3
    assert g.number_of_nodes('game') == 2
    assert g.number_of_edges('follows') == 2
    assert g.number_of_edges('plays') == 2
    u, v = g.edges(form='uv', order='eid', etype='follows')
    assert F.array_equal(u, F.tensor([1, 2], dtype=idtype))
    assert F.array_equal(v, F.tensor([0, 1], dtype=idtype))
1981
1982
1983
1984
    assert F.array_equal(g.edges['follows'].data['h'],
                         F.tensor([2, 4], dtype=idtype))
    assert F.array_equal(g.edges['plays'].data['h'],
                         F.tensor([1, 2], dtype=idtype))
1985
1986
1987
1988
1989
1990
1991

    raise_error = False
    try:
        g = dgl.remove_self_loop(g, etype='plays')
    except:
        raise_error = True
    assert raise_error
1992

1993
    # batch information
1994
1995
    g = dgl.graph(([0, 0, 0, 1, 3, 3, 4], [1, 0, 0, 2, 3, 4, 4]), idtype=idtype,
                  device=F.ctx())
1996
1997
1998
1999
2000
2001
2002
2003
    g.set_batch_num_nodes(F.tensor([3, 2], dtype=F.int64))
    g.set_batch_num_edges(F.tensor([4, 3], dtype=F.int64))
    g = dgl.remove_self_loop(g)
    assert g.number_of_nodes() == 5
    assert g.number_of_edges() == 3
    assert F.array_equal(g.batch_num_nodes(), F.tensor([3, 2], dtype=F.int64))
    assert F.array_equal(g.batch_num_edges(), F.tensor([2, 1], dtype=F.int64))

2004

nv-dlasalle's avatar
nv-dlasalle committed
2005
@parametrize_idtype
2006
def test_reorder_graph(idtype):
2007
2008
2009
2010
2011
    g = dgl.graph(([0, 1, 2, 3, 4], [2, 2, 3, 2, 3]),
                  idtype=idtype, device=F.ctx())
    g.ndata['h'] = F.copy_to(F.randn((g.num_nodes(), 3)), ctx=F.ctx())
    g.edata['w'] = F.copy_to(F.randn((g.num_edges(), 2)), ctx=F.ctx())

2012
    # call with default: node_permute_algo=None, edge_permute_algo='src'
2013
    rg = dgl.reorder_graph(g)
2014
2015
2016
2017
2018
2019
    assert dgl.EID in rg.edata.keys()
    src = F.asnumpy(rg.edges()[0])
    assert np.array_equal(src, np.sort(src))

    # call with 'rcmk' node_permute_algo
    rg = dgl.reorder_graph(g, node_permute_algo='rcmk')
2020
2021
2022
2023
2024
2025
2026
2027
2028
2029
2030
2031
2032
2033
2034
2035
2036
    assert dgl.NID in rg.ndata.keys()
    assert dgl.EID in rg.edata.keys()
    src = F.asnumpy(rg.edges()[0])
    assert np.array_equal(src, np.sort(src))

    # call with 'dst' edge_permute_algo
    rg = dgl.reorder_graph(g, edge_permute_algo='dst')
    dst = F.asnumpy(rg.edges()[1])
    assert np.array_equal(dst, np.sort(dst))

    # call with unknown edge_permute_algo
    raise_error = False
    try:
        dgl.reorder_graph(g, edge_permute_algo='none')
    except:
        raise_error = True
    assert raise_error
2037
2038

    # reorder back to original according to stored ids
2039
    rg = dgl.reorder_graph(g, node_permute_algo='rcmk')
2040
2041
    rg2 = dgl.reorder_graph(rg, 'custom', permute_config={
        'nodes_perm': np.argsort(F.asnumpy(rg.ndata[dgl.NID]))})
2042
2043
2044
2045
    assert F.array_equal(g.ndata['h'], rg2.ndata['h'])
    assert F.array_equal(g.edata['w'], rg2.edata['w'])

    # do not store ids
2046
    rg = dgl.reorder_graph(g, store_ids=False)
2047
2048
2049
2050
2051
2052
2053
2054
2055
2056
2057
2058
2059
    assert not dgl.NID in rg.ndata.keys()
    assert not dgl.EID in rg.edata.keys()

    # metis does not work on windows.
    if os.name == 'nt':
        pass
    else:
        # metis_partition may fail for small graph.
        mg = create_large_graph(1000).to(F.ctx())

        # call with metis strategy, but k is not specified
        raise_error = False
        try:
2060
            dgl.reorder_graph(mg, node_permute_algo='metis')
2061
2062
2063
2064
2065
2066
2067
        except:
            raise_error = True
        assert raise_error

        # call with metis strategy, k is specified
        raise_error = False
        try:
2068
            dgl.reorder_graph(mg,
2069
2070
                              node_permute_algo='metis',
                              permute_config={'k': 2})
2071
2072
2073
2074
2075
2076
2077
2078
        except:
            raise_error = True
        assert not raise_error

    # call with qualified nodes_perm specified
    nodes_perm = np.random.permutation(g.num_nodes())
    raise_error = False
    try:
2079
2080
        dgl.reorder_graph(g, node_permute_algo='custom', permute_config={
            'nodes_perm': nodes_perm})
2081
2082
2083
2084
2085
2086
2087
    except:
        raise_error = True
    assert not raise_error

    # call with unqualified nodes_perm specified
    raise_error = False
    try:
2088
        dgl.reorder_graph(g, node_permute_algo='custom', permute_config={
2089
            'nodes_perm': nodes_perm[:g.num_nodes() - 1]})
2090
2091
2092
2093
2094
2095
2096
    except:
        raise_error = True
    assert raise_error

    # call with unsupported strategy
    raise_error = False
    try:
2097
        dgl.reorder_graph(g, node_permute_algo='cmk')
2098
2099
2100
2101
2102
2103
2104
2105
2106
    except:
        raise_error = True
    assert raise_error

    # heterograph: not supported
    raise_error = False
    try:
        hg = dgl.heterogrpah({('user', 'follow', 'user'): (
            [0, 1], [1, 2])}, idtype=idtype, device=F.ctx())
2107
        dgl.reorder_graph(hg)
2108
2109
2110
2111
    except:
        raise_error = True
    assert raise_error

2112
2113
    # TODO: shall we fix them?
    # add 'csc' format if needed
2114
2115
2116
2117
    # fg = g.formats('csr')
    # assert 'csc' not in sum(fg.formats().values(), [])
    # rfg = dgl.reorder_graph(fg)
    # assert 'csc' in sum(rfg.formats().values(), [])
2118

2119
2120
2121

@unittest.skipIf(dgl.backend.backend_name == "tensorflow",
                 reason="TF doesn't support a slicing operation")
nv-dlasalle's avatar
nv-dlasalle committed
2122
@parametrize_idtype
Mufei Li's avatar
Mufei Li committed
2123
2124
2125
2126
2127
2128
2129
2130
2131
2132
2133
2134
2135
2136
def test_norm_by_dst(idtype):
    # Case1: A homogeneous graph
    g = dgl.graph(([0, 1, 1], [1, 1, 2]), idtype=idtype, device=F.ctx())
    eweight = dgl.norm_by_dst(g)
    assert F.allclose(eweight, F.tensor([0.5, 0.5, 1.0]))

    # Case2: A heterogeneous graph
    g = dgl.heterograph({
        ('user', 'follows', 'user'): ([0, 1], [1, 2]),
        ('user', 'plays', 'game'): ([0, 1, 1], [1, 1, 2])
    }, idtype=idtype, device=F.ctx())
    eweight = dgl.norm_by_dst(g, etype=('user', 'plays', 'game'))
    assert F.allclose(eweight, F.tensor([0.5, 0.5, 1.0]))

2137

nv-dlasalle's avatar
nv-dlasalle committed
2138
@parametrize_idtype
2139
2140
2141
2142
2143
2144
2145
2146
2147
2148
2149
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162
2163
2164
2165
2166
2167
2168
2169
2170
2171
2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
2182
2183
2184
2185
2186
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
2198
2199
2200
2201
2202
2203
2204
2205
2206
2207
2208
2209
2210
2211
2212
2213
2214
2215
def test_module_add_self_loop(idtype):
    g = dgl.graph(([1, 1], [1, 2]), idtype=idtype, device=F.ctx())
    g.ndata['h'] = F.randn((g.num_nodes(), 2))
    g.edata['w'] = F.randn((g.num_edges(), 3))

    # Case1: add self-loops with the default setting
    transform = dgl.AddSelfLoop()
    new_g = transform(g)
    assert new_g.device == g.device
    assert new_g.idtype == g.idtype
    assert new_g.num_nodes() == g.num_nodes()
    assert new_g.num_edges() == 4
    src, dst = new_g.edges()
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(0, 0), (1, 1), (1, 2), (2, 2)}
    assert 'h' in new_g.ndata
    assert 'w' in new_g.edata

    # Case2: Remove self-loops first to avoid duplicate ones
    transform = dgl.AddSelfLoop(allow_duplicate=True)
    new_g = transform(g)
    assert new_g.device == g.device
    assert new_g.idtype == g.idtype
    assert new_g.num_nodes() == g.num_nodes()
    assert new_g.num_edges() == 5
    src, dst = new_g.edges()
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(0, 0), (1, 1), (1, 2), (2, 2)}
    assert 'h' in new_g.ndata
    assert 'w' in new_g.edata

    # Create a heterogeneous graph
    g = dgl.heterograph({
        ('user', 'plays', 'game'): ([0], [1]),
        ('user', 'follows', 'user'): ([1], [3])
    }, idtype=idtype, device=F.ctx())
    g.nodes['user'].data['h1'] = F.randn((4, 2))
    g.edges['plays'].data['w1'] = F.randn((1, 3))
    g.nodes['game'].data['h2'] = F.randn((2, 4))
    g.edges['follows'].data['w2'] = F.randn((1, 5))

    # Case3: add self-loops for a heterogeneous graph
    new_g = transform(g)
    assert new_g.device == g.device
    assert new_g.idtype == g.idtype
    assert new_g.ntypes == g.ntypes
    assert new_g.canonical_etypes == g.canonical_etypes
    for nty in new_g.ntypes:
        assert new_g.num_nodes(nty) == g.num_nodes(nty)
    assert new_g.num_edges('plays') == 1
    assert new_g.num_edges('follows') == 5
    assert 'h1' in new_g.nodes['user'].data
    assert 'h2' in new_g.nodes['game'].data
    assert 'w1' in new_g.edges['plays'].data
    assert 'w2' in new_g.edges['follows'].data

    # Case4: add self-etypes for a heterogeneous graph
    transform = dgl.AddSelfLoop(new_etypes=True)
    new_g = transform(g)
    assert new_g.device == g.device
    assert new_g.idtype == g.idtype
    assert new_g.ntypes == g.ntypes
    assert set(new_g.canonical_etypes) == {
        ('user', 'plays', 'game'), ('user', 'follows', 'user'),
        ('user', 'self', 'user'), ('game', 'self', 'game')
    }
    for nty in new_g.ntypes:
        assert new_g.num_nodes(nty) == g.num_nodes(nty)
    assert new_g.num_edges('plays') == 1
    assert new_g.num_edges('follows') == 5
    assert new_g.num_edges(('user', 'self', 'user')) == 4
    assert new_g.num_edges(('game', 'self', 'game')) == 2
    assert 'h1' in new_g.nodes['user'].data
    assert 'h2' in new_g.nodes['game'].data
    assert 'w1' in new_g.edges['plays'].data
    assert 'w2' in new_g.edges['follows'].data

2216

nv-dlasalle's avatar
nv-dlasalle committed
2217
@parametrize_idtype
2218
2219
2220
2221
2222
2223
2224
2225
2226
2227
2228
2229
2230
2231
2232
2233
2234
2235
2236
2237
2238
2239
2240
2241
2242
2243
2244
2245
2246
2247
2248
2249
2250
2251
2252
2253
2254
2255
2256
2257
2258
2259
def test_module_remove_self_loop(idtype):
    transform = dgl.RemoveSelfLoop()

    # Case1: homogeneous graph
    g = dgl.graph(([1, 1], [1, 2]), idtype=idtype, device=F.ctx())
    g.ndata['h'] = F.randn((g.num_nodes(), 2))
    g.edata['w'] = F.randn((g.num_edges(), 3))
    new_g = transform(g)
    assert new_g.device == g.device
    assert new_g.idtype == g.idtype
    assert new_g.num_nodes() == g.num_nodes()
    assert new_g.num_edges() == 1
    src, dst = new_g.edges()
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(1, 2)}
    assert 'h' in new_g.ndata
    assert 'w' in new_g.edata

    # Case2: heterogeneous graph
    g = dgl.heterograph({
        ('user', 'plays', 'game'): ([0, 1], [1, 1]),
        ('user', 'follows', 'user'): ([1, 2], [2, 2])
    }, idtype=idtype, device=F.ctx())
    g.nodes['user'].data['h1'] = F.randn((3, 2))
    g.edges['plays'].data['w1'] = F.randn((2, 3))
    g.nodes['game'].data['h2'] = F.randn((2, 4))
    g.edges['follows'].data['w2'] = F.randn((2, 5))

    new_g = transform(g)
    assert new_g.device == g.device
    assert new_g.idtype == g.idtype
    assert new_g.ntypes == g.ntypes
    assert new_g.canonical_etypes == g.canonical_etypes
    for nty in new_g.ntypes:
        assert new_g.num_nodes(nty) == g.num_nodes(nty)
    assert new_g.num_edges('plays') == 2
    assert new_g.num_edges('follows') == 1
    assert 'h1' in new_g.nodes['user'].data
    assert 'h2' in new_g.nodes['game'].data
    assert 'w1' in new_g.edges['plays'].data
    assert 'w2' in new_g.edges['follows'].data

2260

nv-dlasalle's avatar
nv-dlasalle committed
2261
@parametrize_idtype
2262
2263
2264
2265
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276
2277
def test_module_add_reverse(idtype):
    transform = dgl.AddReverse()

    # Case1: Add reverse edges for a homogeneous graph
    g = dgl.graph(([0], [1]), idtype=idtype, device=F.ctx())
    g.ndata['h'] = F.randn((g.num_nodes(), 3))
    g.edata['w'] = F.randn((g.num_edges(), 2))
    new_g = transform(g)
    assert new_g.device == g.device
    assert new_g.idtype == g.idtype
    assert g.num_nodes() == new_g.num_nodes()
    src, dst = new_g.edges()
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(0, 1), (1, 0)}
    assert F.allclose(g.ndata['h'], new_g.ndata['h'])
    assert F.allclose(g.edata['w'], F.narrow_row(new_g.edata['w'], 0, 1))
2278
2279
    assert F.allclose(F.narrow_row(new_g.edata['w'], 1, 2),
                      F.zeros((1, 2), F.float32, F.ctx()))
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303

    # Case2: Add reverse edges for a homogeneous graph and copy edata
    transform = dgl.AddReverse(copy_edata=True)
    new_g = transform(g)
    assert new_g.device == g.device
    assert new_g.idtype == g.idtype
    assert g.num_nodes() == new_g.num_nodes()
    src, dst = new_g.edges()
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(0, 1), (1, 0)}
    assert F.allclose(g.ndata['h'], new_g.ndata['h'])
    assert F.allclose(g.edata['w'], F.narrow_row(new_g.edata['w'], 0, 1))
    assert F.allclose(g.edata['w'], F.narrow_row(new_g.edata['w'], 1, 2))

    # Case3: Add reverse edges for a heterogeneous graph
    g = dgl.heterograph({
        ('user', 'plays', 'game'): ([0, 1], [1, 1]),
        ('user', 'follows', 'user'): ([1, 2], [2, 2])
    }, device=F.ctx())
    new_g = transform(g)
    assert new_g.device == g.device
    assert new_g.idtype == g.idtype
    assert g.ntypes == new_g.ntypes
    assert set(new_g.canonical_etypes) == {
2304
2305
        ('user', 'plays', 'game'), ('user', 'follows', 'user'),
        ('game', 'rev_plays', 'user')}
2306
2307
2308
2309
2310
2311
2312
2313
2314
2315
2316
2317
2318
2319
2320
2321
2322
2323
2324
2325
2326
2327
2328
2329
2330
2331
2332
2333
2334
2335
2336
2337
2338
2339
2340
2341
2342
2343
2344
2345
2346
2347
2348
    for nty in g.ntypes:
        assert g.num_nodes(nty) == new_g.num_nodes(nty)

    src, dst = new_g.edges(etype='plays')
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(0, 1), (1, 1)}

    src, dst = new_g.edges(etype='follows')
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(1, 2), (2, 2), (2, 1)}

    src, dst = new_g.edges(etype='rev_plays')
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(1, 1), (1, 0)}

    # Case4: Enforce reverse edge types for symmetric canonical edge types
    transform = dgl.AddReverse(sym_new_etype=True)
    new_g = transform(g)
    assert new_g.device == g.device
    assert new_g.idtype == g.idtype
    assert g.ntypes == new_g.ntypes
    assert set(new_g.canonical_etypes) == {
        ('user', 'plays', 'game'), ('user', 'follows', 'user'),
        ('game', 'rev_plays', 'user'), ('user', 'rev_follows', 'user')}
    for nty in g.ntypes:
        assert g.num_nodes(nty) == new_g.num_nodes(nty)

    src, dst = new_g.edges(etype='plays')
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(0, 1), (1, 1)}

    src, dst = new_g.edges(etype='follows')
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(1, 2), (2, 2)}

    src, dst = new_g.edges(etype='rev_plays')
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(1, 1), (1, 0)}

    src, dst = new_g.edges(etype='rev_follows')
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(2, 1), (2, 2)}

2349
2350
2351

@unittest.skipIf(F._default_context_str == 'gpu',
                 reason="GPU not supported for to_simple")
nv-dlasalle's avatar
nv-dlasalle committed
2352
@parametrize_idtype
2353
2354
2355
2356
2357
2358
2359
2360
2361
2362
2363
2364
2365
2366
2367
2368
2369
2370
2371
2372
2373
2374
2375
2376
2377
2378
2379
2380
2381
2382
2383
2384
2385
2386
2387
2388
2389
2390
def test_module_to_simple(idtype):
    transform = dgl.ToSimple()
    g = dgl.graph(([0, 1, 1], [1, 2, 2]), idtype=idtype, device=F.ctx())
    g.ndata['h'] = F.randn((g.num_nodes(), 2))
    g.edata['w'] = F.tensor([[0.1], [0.2], [0.3]])
    sg = transform(g)
    assert sg.device == g.device
    assert sg.idtype == g.idtype
    assert sg.num_nodes() == g.num_nodes()
    assert sg.num_edges() == 2
    src, dst = sg.edges()
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(0, 1), (1, 2)}
    assert F.allclose(sg.edata['count'], F.tensor([1, 2]))
    assert F.allclose(sg.ndata['h'], g.ndata['h'])

    g = dgl.heterograph({
        ('user', 'follows', 'user'): ([0, 1, 1], [1, 2, 2]),
        ('user', 'plays', 'game'): ([0, 1, 0], [1, 1, 1])
    })
    sg = transform(g)
    assert sg.device == g.device
    assert sg.idtype == g.idtype
    assert sg.ntypes == g.ntypes
    assert sg.canonical_etypes == g.canonical_etypes
    for nty in sg.ntypes:
        assert sg.num_nodes(nty) == g.num_nodes(nty)
    for ety in sg.canonical_etypes:
        assert sg.num_edges(ety) == 2

    src, dst = sg.edges(etype='follows')
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(0, 1), (1, 2)}

    src, dst = sg.edges(etype='plays')
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(0, 1), (1, 1)}

2391

nv-dlasalle's avatar
nv-dlasalle committed
2392
@parametrize_idtype
2393
2394
2395
2396
2397
2398
2399
2400
2401
2402
2403
2404
2405
2406
2407
2408
2409
2410
2411
2412
2413
2414
def test_module_line_graph(idtype):
    transform = dgl.LineGraph()
    g = dgl.graph(([0, 1, 1], [1, 0, 2]), idtype=idtype, device=F.ctx())
    g.ndata['h'] = F.tensor([[0.], [1.], [2.]])
    g.edata['w'] = F.tensor([[0.], [0.1], [0.2]])
    new_g = transform(g)
    assert new_g.device == g.device
    assert new_g.idtype == g.idtype
    assert new_g.num_nodes() == g.num_edges()
    src, dst = new_g.edges()
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(0, 1), (0, 2), (1, 0)}

    transform = dgl.LineGraph(backtracking=False)
    new_g = transform(g)
    assert new_g.device == g.device
    assert new_g.idtype == g.idtype
    assert new_g.num_nodes() == g.num_edges()
    src, dst = new_g.edges()
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(0, 2)}

2415

nv-dlasalle's avatar
nv-dlasalle committed
2416
@parametrize_idtype
2417
2418
2419
2420
2421
2422
2423
2424
2425
2426
2427
2428
2429
def test_module_khop_graph(idtype):
    transform = dgl.KHopGraph(2)
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    g.ndata['h'] = F.randn((g.num_nodes(), 2))
    new_g = transform(g)
    assert new_g.device == g.device
    assert new_g.idtype == g.idtype
    assert new_g.num_nodes() == g.num_nodes()
    assert F.allclose(g.ndata['h'], new_g.ndata['h'])
    src, dst = new_g.edges()
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(0, 2)}

2430

nv-dlasalle's avatar
nv-dlasalle committed
2431
@parametrize_idtype
2432
2433
2434
2435
2436
2437
2438
2439
2440
2441
2442
def test_module_add_metapaths(idtype):
    g = dgl.heterograph({
        ('person', 'author', 'paper'): ([0, 0, 1], [1, 2, 2]),
        ('paper', 'accepted', 'venue'): ([1], [0]),
        ('paper', 'rejected', 'venue'): ([2], [1])
    }, idtype=idtype, device=F.ctx())
    g.nodes['venue'].data['h'] = F.randn((g.num_nodes('venue'), 2))
    g.edges['author'].data['h'] = F.randn((g.num_edges('author'), 3))

    # Case1: keep_orig_edges is True
    metapaths = {
2443
2444
2445
2446
        'accepted': [('person', 'author', 'paper'),
                     ('paper', 'accepted', 'venue')],
        'rejected': [('person', 'author', 'paper'),
                     ('paper', 'rejected', 'venue')]
2447
2448
2449
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460
2461
    }
    transform = dgl.AddMetaPaths(metapaths)
    new_g = transform(g)
    assert new_g.device == g.device
    assert new_g.idtype == g.idtype
    assert new_g.ntypes == g.ntypes
    assert set(new_g.canonical_etypes) == {
        ('person', 'author', 'paper'), ('paper', 'accepted', 'venue'),
        ('paper', 'rejected', 'venue'), ('person', 'accepted', 'venue'),
        ('person', 'rejected', 'venue')
    }
    for nty in new_g.ntypes:
        assert new_g.num_nodes(nty) == g.num_nodes(nty)
    for ety in g.canonical_etypes:
        assert new_g.num_edges(ety) == g.num_edges(ety)
2462
2463
2464
2465
    assert F.allclose(g.nodes['venue'].data['h'],
                      new_g.nodes['venue'].data['h'])
    assert F.allclose(g.edges['author'].data['h'],
                      new_g.edges['author'].data['h'])
2466
2467
2468
2469
2470
2471
2472
2473
2474
2475
2476
2477
2478
2479
2480
2481
2482
2483

    src, dst = new_g.edges(etype=('person', 'accepted', 'venue'))
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(0, 0)}

    src, dst = new_g.edges(etype=('person', 'rejected', 'venue'))
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(0, 1), (1, 1)}

    # Case2: keep_orig_edges is False
    transform = dgl.AddMetaPaths(metapaths, keep_orig_edges=False)
    new_g = transform(g)
    assert new_g.device == g.device
    assert new_g.idtype == g.idtype
    assert new_g.ntypes == g.ntypes
    assert len(new_g.canonical_etypes) == 2
    for nty in new_g.ntypes:
        assert new_g.num_nodes(nty) == g.num_nodes(nty)
2484
2485
    assert F.allclose(g.nodes['venue'].data['h'],
                      new_g.nodes['venue'].data['h'])
2486
2487
2488
2489
2490
2491
2492
2493
2494

    src, dst = new_g.edges(etype=('person', 'accepted', 'venue'))
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(0, 0)}

    src, dst = new_g.edges(etype=('person', 'rejected', 'venue'))
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(0, 1), (1, 1)}

2495

nv-dlasalle's avatar
nv-dlasalle committed
2496
@parametrize_idtype
2497
2498
2499
2500
2501
2502
2503
2504
2505
2506
2507
2508
def test_module_compose(idtype):
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    transform = dgl.Compose([dgl.AddReverse(), dgl.AddSelfLoop()])
    new_g = transform(g)
    assert new_g.device == g.device
    assert new_g.idtype == g.idtype
    assert new_g.num_edges() == 7

    src, dst = new_g.edges()
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(0, 1), (1, 2), (1, 0), (2, 1), (0, 0), (1, 1), (2, 2)}

2509

nv-dlasalle's avatar
nv-dlasalle committed
2510
@parametrize_idtype
2511
2512
2513
2514
2515
2516
2517
2518
2519
2520
2521
def test_module_gcnnorm(idtype):
    g = dgl.heterograph({
        ('A', 'r1', 'A'): ([0, 1, 2], [0, 0, 1]),
        ('A', 'r2', 'B'): ([0, 0], [1, 1]),
        ('B', 'r3', 'B'): ([0, 1, 2], [0, 0, 1])
    }, idtype=idtype, device=F.ctx())
    g.edges['r3'].data['w'] = F.tensor([0.1, 0.2, 0.3])
    transform = dgl.GCNNorm()
    new_g = transform(g)
    assert 'w' not in new_g.edges[('A', 'r2', 'B')].data
    assert F.allclose(new_g.edges[('A', 'r1', 'A')].data['w'],
2522
2523
2524
                      F.tensor([1. / 2, 1. / math.sqrt(2), 0.]))
    assert F.allclose(new_g.edges[('B', 'r3', 'B')].data['w'],
                      F.tensor([1. / 3, 2. / 3, 0.]))
2525

2526
2527
2528

@unittest.skipIf(dgl.backend.backend_name != 'pytorch',
                 reason='Only support PyTorch for now')
nv-dlasalle's avatar
nv-dlasalle committed
2529
@parametrize_idtype
2530
def test_module_ppr(idtype):
2531
2532
    g = dgl.graph(([0, 1, 2, 3, 4], [2, 3, 4, 5, 3]), idtype=idtype,
                  device=F.ctx())
2533
2534
2535
2536
2537
2538
2539
2540
2541
    g.ndata['h'] = F.randn((6, 2))
    transform = dgl.PPR(avg_degree=2)
    new_g = transform(g)
    assert new_g.idtype == g.idtype
    assert new_g.device == g.device
    assert new_g.num_nodes() == g.num_nodes()
    src, dst = new_g.edges()
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(0, 0), (0, 2), (0, 4), (1, 1), (1, 3), (1, 5), (2, 2),
2542
2543
                    (2, 3), (2, 4), (3, 3), (3, 5), (4, 3), (4, 4), (4, 5),
                    (5, 5)}
2544
2545
2546
2547
2548
2549
2550
2551
2552
2553
2554
    assert F.allclose(g.ndata['h'], new_g.ndata['h'])
    assert 'w' in new_g.edata

    # Prior edge weights
    g.edata['w'] = F.tensor([0.1, 0.2, 0.3, 0.4, 0.5])
    new_g = transform(g)
    src, dst = new_g.edges()
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(0, 0), (1, 1), (1, 3), (2, 2), (2, 3), (2, 4),
                    (3, 3), (3, 5), (4, 3), (4, 4), (4, 5), (5, 5)}

2555
2556
2557

@unittest.skipIf(dgl.backend.backend_name != 'pytorch',
                 reason='Only support PyTorch for now')
nv-dlasalle's avatar
nv-dlasalle committed
2558
@parametrize_idtype
2559
2560
def test_module_heat_kernel(idtype):
    # Case1: directed graph
2561
2562
    g = dgl.graph(([0, 1, 2, 3, 4], [2, 3, 4, 5, 3]), idtype=idtype,
                  device=F.ctx())
2563
2564
2565
2566
2567
2568
2569
2570
2571
2572
2573
2574
2575
2576
2577
2578
2579
    g.ndata['h'] = F.randn((6, 2))
    transform = dgl.HeatKernel(avg_degree=1)
    new_g = transform(g)
    assert new_g.idtype == g.idtype
    assert new_g.device == g.device
    assert new_g.num_nodes() == g.num_nodes()
    assert F.allclose(g.ndata['h'], new_g.ndata['h'])
    assert 'w' in new_g.edata

    # Case2: weighted undirected graph
    g = dgl.graph(([0, 1, 2, 3], [1, 0, 3, 2]), idtype=idtype, device=F.ctx())
    g.edata['w'] = F.tensor([0.1, 0.2, 0.3, 0.4])
    new_g = transform(g)
    src, dst = new_g.edges()
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(0, 0), (1, 1), (2, 2), (3, 3)}

2580
2581
2582

@unittest.skipIf(dgl.backend.backend_name != 'pytorch',
                 reason='Only support PyTorch for now')
nv-dlasalle's avatar
nv-dlasalle committed
2583
@parametrize_idtype
2584
2585
def test_module_gdc(idtype):
    transform = dgl.GDC([0.1, 0.2, 0.1], avg_degree=1)
2586
2587
    g = dgl.graph(([0, 1, 2, 3, 4], [2, 3, 4, 5, 3]), idtype=idtype,
                  device=F.ctx())
2588
2589
2590
2591
2592
2593
2594
    g.ndata['h'] = F.randn((6, 2))
    new_g = transform(g)
    assert new_g.idtype == g.idtype
    assert new_g.device == g.device
    assert new_g.num_nodes() == g.num_nodes()
    src, dst = new_g.edges()
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
2595
2596
    assert eset == {(0, 0), (0, 2), (0, 4), (1, 1), (1, 3), (1, 5), (2, 2),
                    (2, 3),
2597
2598
2599
2600
2601
2602
2603
2604
2605
2606
2607
                    (2, 4), (3, 3), (3, 5), (4, 3), (4, 4), (4, 5), (5, 5)}
    assert F.allclose(g.ndata['h'], new_g.ndata['h'])
    assert 'w' in new_g.edata

    # Prior edge weights
    g.edata['w'] = F.tensor([0.1, 0.2, 0.3, 0.4, 0.5])
    new_g = transform(g)
    src, dst = new_g.edges()
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(0, 0), (1, 1), (2, 2), (3, 3), (4, 3), (4, 4), (5, 5)}

2608
2609
2610

@unittest.skipIf(dgl.backend.backend_name == "tensorflow",
                 reason="TF doesn't support a slicing operation")
nv-dlasalle's avatar
nv-dlasalle committed
2611
@parametrize_idtype
2612
2613
2614
2615
2616
def test_module_node_shuffle(idtype):
    transform = dgl.NodeShuffle()
    g = dgl.heterograph({
        ('A', 'r', 'B'): ([0, 1], [1, 2]),
    }, idtype=idtype, device=F.ctx())
2617
2618
    g.nodes['B'].data['h'] = F.randn((g.num_nodes('B'), 2))
    old_nfeat = g.nodes['B'].data['h']
2619
    new_g = transform(g)
2620
2621
    new_nfeat = g.nodes['B'].data['h']
    assert F.allclose(old_nfeat, new_nfeat)
2622

2623
2624
2625

@unittest.skipIf(dgl.backend.backend_name != 'pytorch',
                 reason='Only support PyTorch for now')
nv-dlasalle's avatar
nv-dlasalle committed
2626
@parametrize_idtype
2627
2628
2629
2630
2631
def test_module_drop_node(idtype):
    transform = dgl.DropNode()
    g = dgl.heterograph({
        ('A', 'r', 'B'): ([0, 1], [1, 2]),
    }, idtype=idtype, device=F.ctx())
2632
    num_nodes_old = g.num_nodes()
2633
2634
2635
2636
2637
    new_g = transform(g)
    assert new_g.idtype == g.idtype
    assert new_g.device == g.device
    assert new_g.ntypes == g.ntypes
    assert new_g.canonical_etypes == g.canonical_etypes
2638
2639
2640
    num_nodes_new = g.num_nodes()
    # Ensure that the original graph is not corrupted
    assert num_nodes_old == num_nodes_new
2641

2642
2643
2644

@unittest.skipIf(dgl.backend.backend_name != 'pytorch',
                 reason='Only support PyTorch for now')
nv-dlasalle's avatar
nv-dlasalle committed
2645
@parametrize_idtype
2646
2647
2648
2649
2650
2651
def test_module_drop_edge(idtype):
    transform = dgl.DropEdge()
    g = dgl.heterograph({
        ('A', 'r1', 'B'): ([0, 1], [1, 2]),
        ('C', 'r2', 'C'): ([3, 4, 5], [6, 7, 8])
    }, idtype=idtype, device=F.ctx())
2652
    num_edges_old = g.num_edges()
2653
2654
2655
2656
2657
    new_g = transform(g)
    assert new_g.idtype == g.idtype
    assert new_g.device == g.device
    assert new_g.ntypes == g.ntypes
    assert new_g.canonical_etypes == g.canonical_etypes
2658
2659
2660
    num_edges_new = g.num_edges()
    # Ensure that the original graph is not corrupted
    assert num_edges_old == num_edges_new
2661

2662

nv-dlasalle's avatar
nv-dlasalle committed
2663
@parametrize_idtype
2664
2665
2666
2667
2668
2669
def test_module_add_edge(idtype):
    transform = dgl.AddEdge()
    g = dgl.heterograph({
        ('A', 'r1', 'B'): ([0, 1, 2, 3, 4], [1, 2, 3, 4, 5]),
        ('C', 'r2', 'C'): ([0, 1, 2, 3, 4], [1, 2, 3, 4, 5])
    }, idtype=idtype, device=F.ctx())
2670
    num_edges_old = g.num_edges()
2671
2672
2673
2674
2675
2676
2677
    new_g = transform(g)
    assert new_g.num_edges(('A', 'r1', 'B')) == 6
    assert new_g.num_edges(('C', 'r2', 'C')) == 6
    assert new_g.idtype == g.idtype
    assert new_g.device == g.device
    assert new_g.ntypes == g.ntypes
    assert new_g.canonical_etypes == g.canonical_etypes
2678
2679
2680
    num_edges_new = g.num_edges()
    # Ensure that the original graph is not corrupted
    assert num_edges_old == num_edges_new
2681

2682

nv-dlasalle's avatar
nv-dlasalle committed
2683
@parametrize_idtype
2684
2685
2686
2687
def test_module_random_walk_pe(idtype):
    transform = dgl.RandomWalkPE(2, 'rwpe')
    g = dgl.graph(([0, 1, 1], [1, 1, 0]), idtype=idtype, device=F.ctx())
    new_g = transform(g)
2688
    tgt = F.copy_to(F.tensor([[0., 0.5], [0.5, 0.75]]), g.device)
2689
2690
    assert F.allclose(new_g.ndata['rwpe'], tgt)

2691

nv-dlasalle's avatar
nv-dlasalle committed
2692
@parametrize_idtype
2693
def test_module_laplacian_pe(idtype):
2694
2695
2696
2697
2698
    g = dgl.graph(([2, 1, 0, 3, 1, 1], [3, 1, 1, 2, 1, 0]), idtype=idtype,
                  device=F.ctx())
    tgt_eigval = F.copy_to(
        F.repeat(F.tensor([[1.1534e-17, 1.3333e+00, 2., np.nan, np.nan]]),
                 g.num_nodes(), dim=0), g.device)
2699
    tgt_pe = F.copy_to(F.tensor([[0.5, 0.86602539, 0., 0., 0.],
2700
2701
2702
                                 [0.86602539, 0.5, 0., 0., 0.],
                                 [0., 0., 0.70710677, 0., 0.],
                                 [0., 0., 0.70710677, 0., 0.]]), g.device)
2703
2704
2705

    # without padding (k<n)
    transform = dgl.LaplacianPE(2, feat_name='lappe')
2706
2707
2708
    new_g = transform(g)
    # tensorflow has no abs() api
    if dgl.backend.backend_name == 'tensorflow':
2709
        assert F.allclose(new_g.ndata['lappe'].__abs__(), tgt_pe[:, :2])
2710
2711
    # pytorch & mxnet
    else:
2712
        assert F.allclose(new_g.ndata['lappe'].abs(), tgt_pe[:, :2])
2713
2714
2715
2716
2717
2718
2719
2720
2721
2722
2723
2724

    # with padding (k>=n)
    transform = dgl.LaplacianPE(5, feat_name='lappe', padding=True)
    new_g = transform(g)
    # tensorflow has no abs() api
    if dgl.backend.backend_name == 'tensorflow':
        assert F.allclose(new_g.ndata['lappe'].__abs__(), tgt_pe)
    # pytorch & mxnet
    else:
        assert F.allclose(new_g.ndata['lappe'].abs(), tgt_pe)

    # with eigenvalues
2725
2726
    transform = dgl.LaplacianPE(5, feat_name='lappe', eigval_name='eigval',
                                padding=True)
2727
2728
2729
    new_g = transform(g)
    # tensorflow has no abs() api
    if dgl.backend.backend_name == 'tensorflow':
2730
        assert F.allclose(new_g.ndata['eigval'][:, :3], tgt_eigval[:, :3])
2731
2732
2733
        assert F.allclose(new_g.ndata['lappe'].__abs__(), tgt_pe)
    # pytorch & mxnet
    else:
2734
        assert F.allclose(new_g.ndata['eigval'][:, :3], tgt_eigval[:, :3])
2735
        assert F.allclose(new_g.ndata['lappe'].abs(), tgt_pe)
2736

2737
2738
2739

@unittest.skipIf(dgl.backend.backend_name != 'pytorch',
                 reason='Only support PyTorch for now')
2740
2741
2742
@pytest.mark.parametrize('g', get_cases(['has_scalar_e_feature']))
def test_module_sign(g):
    import torch
2743

Mufei Li's avatar
Mufei Li committed
2744
    atol = 1e-06
2745
2746
2747
2748
2749
2750
2751
2752
2753
2754
2755
2756
2757
2758

    ctx = F.ctx()
    g = g.to(ctx)
    adj = g.adj(transpose=True, scipy_fmt='coo').todense()
    adj = torch.tensor(adj).float().to(ctx)

    weight_adj = g.adj(transpose=True, scipy_fmt='coo').astype(float).todense()
    weight_adj = torch.tensor(weight_adj).float().to(ctx)
    src, dst = g.edges()
    src, dst = src.long(), dst.long()
    weight_adj[dst, src] = g.edata['scalar_w']

    # raw
    transform = dgl.SIGNDiffusion(k=1, in_feat_name='h', diffuse_op='raw')
2759
    g = transform(g)
Mufei Li's avatar
Mufei Li committed
2760
    target = torch.matmul(adj, g.ndata['h'])
Mufei Li's avatar
Mufei Li committed
2761
    assert torch.allclose(g.ndata['out_feat_1'], target, atol=atol)
2762

2763
2764
    transform = dgl.SIGNDiffusion(k=1, in_feat_name='h',
                                  eweight_name='scalar_w', diffuse_op='raw')
2765
    g = transform(g)
Mufei Li's avatar
Mufei Li committed
2766
    target = torch.matmul(weight_adj, g.ndata['h'])
Mufei Li's avatar
Mufei Li committed
2767
    assert torch.allclose(g.ndata['out_feat_1'], target, atol=atol)
2768
2769
2770
2771

    # rw
    adj_rw = torch.matmul(torch.diag(1 / adj.sum(dim=1)), adj)
    transform = dgl.SIGNDiffusion(k=1, in_feat_name='h', diffuse_op='rw')
2772
    g = transform(g)
Mufei Li's avatar
Mufei Li committed
2773
    target = torch.matmul(adj_rw, g.ndata['h'])
Mufei Li's avatar
Mufei Li committed
2774
    assert torch.allclose(g.ndata['out_feat_1'], target, atol=atol)
2775

2776
2777
2778
2779
    weight_adj_rw = torch.matmul(torch.diag(1 / weight_adj.sum(dim=1)),
                                 weight_adj)
    transform = dgl.SIGNDiffusion(k=1, in_feat_name='h',
                                  eweight_name='scalar_w', diffuse_op='rw')
2780
    g = transform(g)
Mufei Li's avatar
Mufei Li committed
2781
    target = torch.matmul(weight_adj_rw, g.ndata['h'])
Mufei Li's avatar
Mufei Li committed
2782
    assert torch.allclose(g.ndata['out_feat_1'], target, atol=atol)
2783
2784
2785
2786

    # gcn
    raw_eweight = g.edata['scalar_w']
    gcn_norm = dgl.GCNNorm()
2787
    g = gcn_norm(g)
2788
2789
2790
    adj_gcn = adj.clone()
    adj_gcn[dst, src] = g.edata.pop('w')
    transform = dgl.SIGNDiffusion(k=1, in_feat_name='h', diffuse_op='gcn')
2791
2792
    g = transform(g)
    target = torch.matmul(adj_gcn, g.ndata['h'])
Mufei Li's avatar
Mufei Li committed
2793
    assert torch.allclose(g.ndata['out_feat_1'], target, atol=atol)
2794
2795

    gcn_norm = dgl.GCNNorm('scalar_w')
2796
    g = gcn_norm(g)
2797
2798
2799
2800
2801
    weight_adj_gcn = weight_adj.clone()
    weight_adj_gcn[dst, src] = g.edata['scalar_w']
    g.edata['scalar_w'] = raw_eweight
    transform = dgl.SIGNDiffusion(k=1, in_feat_name='h',
                                  eweight_name='scalar_w', diffuse_op='gcn')
2802
2803
    g = transform(g)
    target = torch.matmul(weight_adj_gcn, g.ndata['h'])
Mufei Li's avatar
Mufei Li committed
2804
    assert torch.allclose(g.ndata['out_feat_1'], target, atol=atol)
2805
2806
2807

    # ppr
    alpha = 0.2
2808
2809
    transform = dgl.SIGNDiffusion(k=1, in_feat_name='h', diffuse_op='ppr',
                                  alpha=alpha)
2810
    g = transform(g)
2811
2812
    target = (1 - alpha) * torch.matmul(adj_gcn, g.ndata['h']) + alpha * \
             g.ndata['h']
Mufei Li's avatar
Mufei Li committed
2813
    assert torch.allclose(g.ndata['out_feat_1'], target, atol=atol)
2814

2815
2816
    transform = dgl.SIGNDiffusion(k=1, in_feat_name='h',
                                  eweight_name='scalar_w',
2817
                                  diffuse_op='ppr', alpha=alpha)
2818
    g = transform(g)
2819
2820
    target = (1 - alpha) * torch.matmul(weight_adj_gcn, g.ndata['h']) + alpha * \
             g.ndata['h']
Mufei Li's avatar
Mufei Li committed
2821
    assert torch.allclose(g.ndata['out_feat_1'], target, atol=atol)
2822

2823
2824
2825

@unittest.skipIf(dgl.backend.backend_name != 'pytorch',
                 reason='Only support PyTorch for now')
nv-dlasalle's avatar
nv-dlasalle committed
2826
@parametrize_idtype
2827
2828
def test_module_row_feat_normalizer(idtype):
    # Case1: Normalize features of a homogeneous graph.
2829
    transform = dgl.RowFeatNormalizer(subtract_min=True,
2830
2831
                                      node_feat_names=['h'],
                                      edge_feat_names=['w'])
2832
2833
2834
2835
2836
2837
2838
2839
2840
2841
    g = dgl.rand_graph(5, 5, idtype=idtype, device=F.ctx())
    g.ndata['h'] = F.randn((g.num_nodes(), 128))
    g.edata['w'] = F.randn((g.num_edges(), 128))
    g = transform(g)
    assert g.ndata['h'].shape == (g.num_nodes(), 128)
    assert g.edata['w'].shape == (g.num_edges(), 128)
    assert F.allclose(g.ndata['h'].sum(1), F.tensor([1.0, 1.0, 1.0, 1.0, 1.0]))
    assert F.allclose(g.edata['w'].sum(1), F.tensor([1.0, 1.0, 1.0, 1.0, 1.0]))

    # Case2: Normalize features of a heterogeneous graph.
2842
    transform = dgl.RowFeatNormalizer(subtract_min=True,
2843
2844
                                      node_feat_names=['h', 'h2'],
                                      edge_feat_names=['w'])
2845
2846
2847
2848
2849
2850
    g = dgl.heterograph({
        ('user', 'follows', 'user'): (F.tensor([1, 2]), F.tensor([3, 4])),
        ('player', 'plays', 'game'): (F.tensor([2, 2]), F.tensor([1, 1]))
    }, idtype=idtype, device=F.ctx())
    g.ndata['h'] = {'game': F.randn((2, 128)), 'player': F.randn((3, 128))}
    g.ndata['h2'] = {'user': F.randn((5, 128))}
2851
2852
    g.edata['w'] = {('user', 'follows', 'user'): F.randn((2, 128)),
                    ('player', 'plays', 'game'): F.randn((2, 128))}
2853
2854
2855
2856
2857
2858
2859
2860
    g = transform(g)
    assert g.ndata['h']['game'].shape == (2, 128)
    assert g.ndata['h']['player'].shape == (3, 128)
    assert g.ndata['h2']['user'].shape == (5, 128)
    assert g.edata['w'][('user', 'follows', 'user')].shape == (2, 128)
    assert g.edata['w'][('player', 'plays', 'game')].shape == (2, 128)
    assert F.allclose(g.ndata['h']['game'].sum(1), F.tensor([1.0, 1.0]))
    assert F.allclose(g.ndata['h']['player'].sum(1), F.tensor([1.0, 1.0, 1.0]))
2861
2862
2863
2864
2865
2866
2867
    assert F.allclose(g.ndata['h2']['user'].sum(1),
                      F.tensor([1.0, 1.0, 1.0, 1.0, 1.0]))
    assert F.allclose(g.edata['w'][('user', 'follows', 'user')].sum(1),
                      F.tensor([1.0, 1.0]))
    assert F.allclose(g.edata['w'][('player', 'plays', 'game')].sum(1),
                      F.tensor([1.0, 1.0]))

2868

2869
2870
@unittest.skipIf(dgl.backend.backend_name != 'pytorch',
                 reason='Only support PyTorch for now')
nv-dlasalle's avatar
nv-dlasalle committed
2871
@parametrize_idtype
2872
2873
def test_module_feat_mask(idtype):
    # Case1: Mask node and edge feature tensors of a homogeneous graph.
2874
    transform = dgl.FeatMask(node_feat_names=['h'], edge_feat_names=['w'])
2875
2876
2877
2878
2879
2880
2881
2882
2883
2884
2885
2886
2887
2888
2889
2890
2891
2892
2893
2894
2895
2896
2897
2898
2899
    g = dgl.rand_graph(5, 20, idtype=idtype, device=F.ctx())
    g.ndata['h'] = F.ones((g.num_nodes(), 10))
    g.edata['w'] = F.ones((g.num_edges(), 20))
    g = transform(g)
    assert g.device == g.device
    assert g.idtype == g.idtype
    assert g.ndata['h'].shape == (g.num_nodes(), 10)
    assert g.edata['w'].shape == (g.num_edges(), 20)

    # Case2: Mask node and edge feature tensors of a heterogeneous graph.
    g = dgl.heterograph({
        ('user', 'follows', 'user'): (F.tensor([1, 2]), F.tensor([3, 4])),
        ('player', 'plays', 'game'): (F.tensor([2, 2]), F.tensor([1, 1]))
    }, idtype=idtype, device=F.ctx())
    g.ndata['h'] = {'game': F.randn((2, 5)), 'player': F.randn((3, 5))}
    g.edata['w'] = {('user', 'follows', 'user'): F.randn((2, 5)),
                    ('player', 'plays', 'game'): F.randn((2, 5))}
    g = transform(g)
    assert g.device == g.device
    assert g.idtype == g.idtype
    assert g.ndata['h']['game'].shape == (2, 5)
    assert g.ndata['h']['player'].shape == (3, 5)
    assert g.edata['w'][('user', 'follows', 'user')].shape == (2, 5)
    assert g.edata['w'][('player', 'plays', 'game')].shape == (2, 5)

2900

2901
2902
2903
2904
2905
2906
2907
2908
2909
2910
2911
2912
2913
2914
2915
2916
2917
2918
2919
2920
2921
2922
2923
2924
2925
2926
2927
2928
2929
2930
2931
2932
@parametrize_idtype
def test_shortest_dist(idtype):
    g = dgl.graph(([0, 1, 1, 2], [2, 0, 3, 3]), idtype=idtype, device=F.ctx())

    # case 1: directed single source
    dist = dgl.shortest_dist(g, root=0)
    tgt = F.copy_to(F.tensor([0, -1, 1, 2], dtype=F.int64), g.device)
    assert F.array_equal(dist, tgt)

    # case 2: undirected all pairs
    dist, paths = dgl.shortest_dist(g, root=None, return_paths=True)
    tgt_dist = F.copy_to(
        F.tensor([
            [0, -1, 1, 2],
            [1, 0, 2, 1],
            [-1, -1, 0, 1],
            [-1, -1, -1, 0]
        ], dtype=F.int64),
        g.device
    )
    tgt_paths = F.copy_to(
        F.tensor([
            [[-1, -1], [-1, -1], [0, -1], [0, 3]],
            [[1, -1], [-1, -1], [1, 0], [2, -1]],
            [[-1, -1], [-1, -1], [-1, -1], [3, -1]],
            [[-1, -1], [-1, -1], [-1, -1], [-1, -1]]
        ], dtype=F.int64),
        g.device
    )
    assert F.array_equal(dist, tgt_dist)
    assert F.array_equal(paths, tgt_paths)

2933

2934
2935
2936
2937
2938
2939
2940
2941
2942
2943
2944
2945
2946
2947
2948
2949
2950
2951
2952
2953
2954
2955
2956
2957
2958
2959
2960
2961
@parametrize_idtype
def test_module_to_levi(idtype):
    transform = dgl.ToLevi()
    g = dgl.graph(([0, 1, 2, 3], [1, 2, 3, 0]), idtype=idtype, device=F.ctx())
    g.ndata['h'] = F.randn((g.num_nodes(), 2))
    g.edata['w'] = F.randn((g.num_edges(), 2))
    lg = transform(g)
    assert lg.device == g.device
    assert lg.idtype == g.idtype
    assert lg.ntypes == ['edge', 'node']
    assert lg.canonical_etypes == [('edge', 'e2n', 'node'),
                                   ('node', 'n2e', 'edge')]
    assert lg.num_nodes('node') == g.num_nodes()
    assert lg.num_nodes('edge') == g.num_edges()
    assert lg.num_edges('n2e') == g.num_edges()
    assert lg.num_edges('e2n') == g.num_edges()

    src, dst = lg.edges(etype='n2e')
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(0, 0), (1, 1), (2, 2), (3, 3)}

    src, dst = lg.edges(etype='e2n')
    eset = set(zip(list(F.asnumpy(src)), list(F.asnumpy(dst))))
    assert eset == {(0, 1), (1, 2), (2, 3), (3, 0)}

    assert F.allclose(lg.nodes['node'].data['h'], g.ndata['h'])
    assert F.allclose(lg.nodes['edge'].data['w'], g.edata['w'])

2962
2963
2964
2965
2966
2967
2968
2969
2970
2971
2972
2973
2974
2975
2976
2977
2978
2979
2980
2981
2982
2983
2984
2985
2986
2987
2988
2989
2990
2991
2992
2993
2994
2995
2996
2997
2998
2999

@parametrize_idtype
def test_module_svd_pe(idtype):
    g = dgl.graph(
        (
            [0, 0, 1, 1, 2, 2, 2, 2, 3, 3, 4, 4],
            [2, 3, 0, 2, 0, 2, 3, 4, 3, 4, 0, 1]
        ),
        idtype=idtype,
        device=F.ctx(),
    )
    # without padding
    tgt_pe = F.copy_to(
        F.tensor(
            [
                [0.6669, 0.3068, 0.7979, 0.8477],
                [0.6311, 0.6101, 0.1248, 0.5137],
                [1.1993, 0.0665, 0.9183, 0.1455],
                [0.5682, 0.6766, 0.8952, 0.6449],
                [0.3393, 0.8363, 0.6500, 0.4564],
            ]
        ),
        g.device,
    )
    transform_1 = dgl.SVDPE(k=2, feat_name="svd_pe")
    g1 = transform_1(g)
    if dgl.backend.backend_name == "tensorflow":
        assert F.allclose(g1.ndata["svd_pe"].__abs__(), tgt_pe)
    else:
        assert F.allclose(g1.ndata["svd_pe"].abs(), tgt_pe)

    # with padding
    transform_2 = dgl.SVDPE(k=6, feat_name="svd_pe", padding=True)
    g2 = transform_2(g)
    assert F.shape(g2.ndata["svd_pe"]) == (5, 12)


if __name__ == "__main__":
3000
    test_partition_with_halo()
3001
    test_module_heat_kernel(F.int32)