test_basics.py 26.2 KB
Newer Older
1
import backend as F
2
import dgl
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
3
4
import numpy as np
import scipy.sparse as ssp
5
import networkx as nx
6
from dgl import DGLGraph
7
from collections import defaultdict as ddict
8
9
10
11

D = 5
reduce_msg_shapes = set()

12
def message_func(edges):
13
14
    assert F.ndim(edges.src['h']) == 2
    assert F.shape(edges.src['h'])[1] == D
15
    return {'m' : edges.src['h']}
16

17
18
def reduce_func(nodes):
    msgs = nodes.mailbox['m']
Minjie Wang's avatar
Minjie Wang committed
19
    reduce_msg_shapes.add(tuple(msgs.shape))
20
21
22
    assert F.ndim(msgs) == 3
    assert F.shape(msgs)[2] == D
    return {'accum' : F.sum(msgs, 1)}
Minjie Wang's avatar
Minjie Wang committed
23

24
25
def apply_node_func(nodes):
    return {'h' : nodes.data['h'] + nodes.data['accum']}
Minjie Wang's avatar
Minjie Wang committed
26

27
def generate_graph(grad=False):
28
    g = DGLGraph()
29
    g.add_nodes(10) # 10 nodes
30
    # create a graph where 0 is the source and 9 is the sink
Minjie Wang's avatar
Minjie Wang committed
31
    # 17 edges
32
33
34
35
36
    for i in range(1, 9):
        g.add_edge(0, i)
        g.add_edge(i, 9)
    # add a back flow from 9 to 0
    g.add_edge(9, 0)
37
38
39
40
41
42
    ncol = F.randn((10, D))
    ecol = F.randn((17, D))
    if grad:
        ncol = F.attach_grad(ncol)
        ecol = F.attach_grad(ecol)

43
44
    g.ndata['h'] = ncol
    g.edata['w'] = ecol
45
46
    g.set_n_initializer(dgl.init.zero_initializer)
    g.set_e_initializer(dgl.init.zero_initializer)
47
48
49
50
    return g

def test_batch_setter_getter():
    def _pfc(x):
51
        return list(F.zerocopy_to_numpy(x)[:,0])
52
53
    g = generate_graph()
    # set all nodes
54
55
    g.ndata['h'] = F.zeros((10, D))
    assert F.allclose(g.ndata['h'], F.zeros((10, D)))
Minjie Wang's avatar
Minjie Wang committed
56
    # pop nodes
57
    old_len = len(g.ndata)
Minjie Wang's avatar
Minjie Wang committed
58
    assert _pfc(g.pop_n_repr('h')) == [0.] * 10
59
    assert len(g.ndata) == old_len - 1
60
    g.ndata['h'] = F.zeros((10, D))
61
    # set partial nodes
62
63
    u = F.tensor([1, 3, 5])
    g.nodes[u].data['h'] = F.ones((3, D))
64
    assert _pfc(g.ndata['h']) == [0., 1., 0., 1., 0., 1., 0., 0., 0., 0.]
65
    # get partial nodes
66
    u = F.tensor([1, 2, 3])
67
    assert _pfc(g.nodes[u].data['h']) == [1., 0., 1.]
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89

    '''
    s, d, eid
    0, 1, 0
    1, 9, 1
    0, 2, 2
    2, 9, 3
    0, 3, 4
    3, 9, 5
    0, 4, 6
    4, 9, 7
    0, 5, 8
    5, 9, 9
    0, 6, 10
    6, 9, 11
    0, 7, 12
    7, 9, 13
    0, 8, 14
    8, 9, 15
    9, 0, 16
    '''
    # set all edges
90
    g.edata['l'] = F.zeros((17, D))
91
    assert _pfc(g.edata['l']) == [0.] * 17
Minjie Wang's avatar
Minjie Wang committed
92
    # pop edges
93
    old_len = len(g.edata)
Minjie Wang's avatar
Minjie Wang committed
94
    assert _pfc(g.pop_e_repr('l')) == [0.] * 17
95
    assert len(g.edata) == old_len - 1
96
    g.edata['l'] = F.zeros((17, D))
Minjie Wang's avatar
Minjie Wang committed
97
    # set partial edges (many-many)
98
99
100
    u = F.tensor([0, 0, 2, 5, 9])
    v = F.tensor([1, 3, 9, 9, 0])
    g.edges[u, v].data['l'] = F.ones((5, D))
Minjie Wang's avatar
Minjie Wang committed
101
102
    truth = [0.] * 17
    truth[0] = truth[4] = truth[3] = truth[9] = truth[16] = 1.
103
    assert _pfc(g.edata['l']) == truth
Minjie Wang's avatar
Minjie Wang committed
104
    # set partial edges (many-one)
105
106
107
    u = F.tensor([3, 4, 6])
    v = F.tensor([9])
    g.edges[u, v].data['l'] = F.ones((3, D))
Minjie Wang's avatar
Minjie Wang committed
108
    truth[5] = truth[7] = truth[11] = 1.
109
    assert _pfc(g.edata['l']) == truth
Minjie Wang's avatar
Minjie Wang committed
110
    # set partial edges (one-many)
111
112
113
    u = F.tensor([0])
    v = F.tensor([4, 5, 6])
    g.edges[u, v].data['l'] = F.ones((3, D))
Minjie Wang's avatar
Minjie Wang committed
114
    truth[6] = truth[8] = truth[10] = 1.
115
    assert _pfc(g.edata['l']) == truth
Minjie Wang's avatar
Minjie Wang committed
116
    # get partial edges (many-many)
117
118
    u = F.tensor([0, 6, 0])
    v = F.tensor([6, 9, 7])
119
    assert _pfc(g.edges[u, v].data['l']) == [1., 1., 0.]
Minjie Wang's avatar
Minjie Wang committed
120
    # get partial edges (many-one)
121
122
    u = F.tensor([5, 6, 7])
    v = F.tensor([9])
123
    assert _pfc(g.edges[u, v].data['l']) == [1., 1., 0.]
Minjie Wang's avatar
Minjie Wang committed
124
    # get partial edges (one-many)
125
126
    u = F.tensor([0])
    v = F.tensor([3, 4, 5])
127
    assert _pfc(g.edges[u, v].data['l']) == [1., 1., 1.]
128

129
130
def test_batch_setter_autograd():
    g = generate_graph(grad=True)
131
    h1 = g.ndata['h']
132
    # partial set
133
134
135
136
137
138
139
140
    v = F.tensor([1, 2, 8])
    hh = F.attach_grad(F.zeros((len(v), D)))
    with F.record_grad():
        g.nodes[v].data['h'] = hh
        h2 = g.ndata['h']
    F.backward(h2, F.ones((10, D)) * 2)
    assert F.array_equal(F.grad(h1)[:,0], F.tensor([2., 0., 0., 2., 2., 2., 2., 2., 0., 2.]))
    assert F.array_equal(F.grad(hh)[:,0], F.tensor([2., 2., 2.]))
141

142
143
144
145
def test_nx_conversion():
    # check conversion between networkx and DGLGraph

    def _check_nx_feature(nxg, nf, ef):
Minjie Wang's avatar
Minjie Wang committed
146
147
        # check node and edge feature of nxg
        # this is used to check to_networkx
148
149
150
151
152
153
154
        num_nodes = len(nxg)
        num_edges = nxg.size()
        if num_nodes > 0:
            node_feat = ddict(list)
            for nid, attr in nxg.nodes(data=True):
                assert len(attr) == len(nf)
                for k in nxg.nodes[nid]:
155
                    node_feat[k].append(F.unsqueeze(attr[k], 0))
156
            for k in node_feat:
157
158
                feat = F.cat(node_feat[k], 0)
                assert F.allclose(feat, nf[k])
159
160
161
162
163
164
165
166
        else:
            assert len(nf) == 0
        if num_edges > 0:
            edge_feat = ddict(lambda: [0] * num_edges)
            for u, v, attr in nxg.edges(data=True):
                assert len(attr) == len(ef) + 1 # extra id
                eid = attr['id']
                for k in ef:
167
                    edge_feat[k][eid] = F.unsqueeze(attr[k], 0)
168
            for k in edge_feat:
169
170
                feat = F.cat(edge_feat[k], 0)
                assert F.allclose(feat, ef[k])
171
172
173
        else:
            assert len(ef) == 0

174
175
176
177
178
    n1 = F.randn((5, 3))
    n2 = F.randn((5, 10))
    n3 = F.randn((5, 4))
    e1 = F.randn((4, 5))
    e2 = F.randn((4, 7))
179
180
181
182
183
184
185
186
187
188
189
190
    g = DGLGraph(multigraph=True)
    g.add_nodes(5)
    g.add_edges([0,1,3,4], [2,4,0,3])
    g.ndata.update({'n1': n1, 'n2': n2, 'n3': n3})
    g.edata.update({'e1': e1, 'e2': e2})

    # convert to networkx
    nxg = g.to_networkx(node_attrs=['n1', 'n3'], edge_attrs=['e1', 'e2'])
    assert len(nxg) == 5
    assert nxg.size() == 4
    _check_nx_feature(nxg, {'n1': n1, 'n3': n3}, {'e1': e1, 'e2': e2})

Minjie Wang's avatar
Minjie Wang committed
191
    # convert to DGLGraph, nx graph has id in edge feature
192
193
    # use id feature to test non-tensor copy
    g.from_networkx(nxg, node_attrs=['n1'], edge_attrs=['e1', 'id'])
Minjie Wang's avatar
Minjie Wang committed
194
    # check graph size
195
196
    assert g.number_of_nodes() == 5
    assert g.number_of_edges() == 4
Minjie Wang's avatar
Minjie Wang committed
197
198
199
200
201
    # check number of features
    # test with existing dglgraph (so existing features should be cleared)
    assert len(g.ndata) == 1
    assert len(g.edata) == 2
    # check feature values
202
    assert F.allclose(g.ndata['n1'], n1)
Minjie Wang's avatar
Minjie Wang committed
203
    # with id in nx edge feature, e1 should follow original order
204
    assert F.allclose(g.edata['e1'], e1)
205
    assert F.array_equal(g.get_e_repr()['id'], F.copy_to(F.arange(0, 4), F.cpu()))
206

Minjie Wang's avatar
Minjie Wang committed
207
208
    # test conversion after modifying DGLGraph
    g.pop_e_repr('id') # pop id so we don't need to provide id when adding edges
209
210
    new_n = F.randn((2, 3))
    new_e = F.randn((3, 5))
211
212
213
    g.add_nodes(2, data={'n1': new_n})
    # add three edges, one is a multi-edge
    g.add_edges([3, 6, 0], [4, 5, 2], data={'e1': new_e})
214
215
    n1 = F.cat((n1, new_n), 0)
    e1 = F.cat((e1, new_e), 0)
216
217
218
219
220
221
    # convert to networkx again
    nxg = g.to_networkx(node_attrs=['n1'], edge_attrs=['e1'])
    assert len(nxg) == 7
    assert nxg.size() == 7
    _check_nx_feature(nxg, {'n1': n1}, {'e1': e1})

Minjie Wang's avatar
Minjie Wang committed
222
223
224
225
226
227
228
229
230
231
232
233
234
235
    # now test convert from networkx without id in edge feature
    # first pop id in edge feature
    for _, _, attr in nxg.edges(data=True):
        attr.pop('id')
    # test with a new graph
    g = DGLGraph(multigraph=True)
    g.from_networkx(nxg, node_attrs=['n1'], edge_attrs=['e1'])
    # check graph size
    assert g.number_of_nodes() == 7
    assert g.number_of_edges() == 7
    # check number of features
    assert len(g.ndata) == 1
    assert len(g.edata) == 1
    # check feature values
236
    assert F.allclose(g.ndata['n1'], n1)
Minjie Wang's avatar
Minjie Wang committed
237
238
239
    # edge feature order follows nxg.edges()
    edge_feat = []
    for _, _, attr in nxg.edges(data=True):
240
241
242
        edge_feat.append(F.unsqueeze(attr['e1'], 0))
    edge_feat = F.cat(edge_feat, 0)
    assert F.allclose(g.edata['e1'], edge_feat)
Minjie Wang's avatar
Minjie Wang committed
243

244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
    # Test converting from a networkx graph whose nodes are
    # not labeled with consecutive-integers.
    nxg = nx.cycle_graph(5)
    nxg.remove_nodes_from([0, 4])
    for u in nxg.nodes():
        nxg.node[u]['h'] = F.tensor([u])
    for u, v, d in nxg.edges(data=True):
        d['h'] = F.tensor([u, v])

    g = dgl.DGLGraph()
    g.from_networkx(nxg, node_attrs=['h'], edge_attrs=['h'])
    assert g.number_of_nodes() == 3
    assert g.number_of_edges() == 4
    assert g.has_edge_between(0, 1)
    assert g.has_edge_between(1, 2)
    assert F.allclose(g.ndata['h'], F.tensor([[1.], [2.], [3.]]))
    assert F.allclose(g.edata['h'], F.tensor([[1., 2.], [1., 2.],
                                              [2., 3.], [2., 3.]]))

263
264
def test_batch_send():
    g = generate_graph()
265
    def _fmsg(edges):
266
        assert tuple(F.shape(edges.src['h'])) == (5, D)
267
        return {'m' : edges.src['h']}
Minjie Wang's avatar
Minjie Wang committed
268
    g.register_message_func(_fmsg)
269
    # many-many send
270
271
    u = F.tensor([0, 0, 0, 0, 0])
    v = F.tensor([1, 2, 3, 4, 5])
272
    g.send((u, v))
273
    # one-many send
274
275
    u = F.tensor([0])
    v = F.tensor([1, 2, 3, 4, 5])
276
    g.send((u, v))
277
    # many-one send
278
279
    u = F.tensor([1, 2, 3, 4, 5])
    v = F.tensor([9])
280
    g.send((u, v))
281

282
def test_batch_recv():
Minjie Wang's avatar
Minjie Wang committed
283
    # basic recv test
284
    g = generate_graph()
Minjie Wang's avatar
Minjie Wang committed
285
286
287
    g.register_message_func(message_func)
    g.register_reduce_func(reduce_func)
    g.register_apply_node_func(apply_node_func)
288
289
    u = F.tensor([0, 0, 0, 4, 5, 6])
    v = F.tensor([1, 2, 3, 9, 9, 9])
Minjie Wang's avatar
Minjie Wang committed
290
    reduce_msg_shapes.clear()
291
    g.send((u, v))
292
    g.recv(F.unique(v))
Minjie Wang's avatar
Minjie Wang committed
293
294
295
    assert(reduce_msg_shapes == {(1, 3, D), (3, 1, D)})
    reduce_msg_shapes.clear()

296
297
298
299
300
301
302
def test_apply_nodes():
    def _upd(nodes):
        return {'h' : nodes.data['h'] * 2}
    g = generate_graph()
    g.register_apply_node_func(_upd)
    old = g.ndata['h']
    g.apply_nodes()
303
304
    assert F.allclose(old * 2, g.ndata['h'])
    u = F.tensor([0, 3, 4, 6])
305
    g.apply_nodes(lambda nodes : {'h' : nodes.data['h'] * 0.}, u)
306
    assert F.allclose(F.gather_row(g.ndata['h'], u), F.zeros((4, D)))
307

308
def test_apply_edges():
309
310
311
    def _upd(edges):
        return {'w' : edges.data['w'] * 2}
    g = generate_graph()
312
    g.register_apply_edge_func(_upd)
313
    old = g.edata['w']
314
    g.apply_edges()
315
316
317
    assert F.allclose(old * 2, g.edata['w'])
    u = F.tensor([0, 0, 0, 4, 5, 6])
    v = F.tensor([1, 2, 3, 9, 9, 9])
318
    g.apply_edges(lambda edges : {'w' : edges.data['w'] * 0.}, (u, v))
319
    eid = F.tensor(g.edge_ids(u, v))
320
    assert F.allclose(F.gather_row(g.edata['w'], eid), F.zeros((6, D)))
321

322
323
def test_update_routines():
    g = generate_graph()
Minjie Wang's avatar
Minjie Wang committed
324
325
326
    g.register_message_func(message_func)
    g.register_reduce_func(reduce_func)
    g.register_apply_node_func(apply_node_func)
327

328
    # send_and_recv
329
    reduce_msg_shapes.clear()
330
331
    u = [0, 0, 0, 4, 5, 6]
    v = [1, 2, 3, 9, 9, 9]
332
    g.send_and_recv((u, v))
333
334
    assert(reduce_msg_shapes == {(1, 3, D), (3, 1, D)})
    reduce_msg_shapes.clear()
335
336
337
    try:
        g.send_and_recv([u, v])
        assert False
338
    except:
339
        pass
340

341
    # pull
342
    v = F.tensor([1, 2, 3, 9])
343
    reduce_msg_shapes.clear()
344
    g.pull(v)
345
346
347
    assert(reduce_msg_shapes == {(1, 8, D), (3, 1, D)})
    reduce_msg_shapes.clear()

348
    # push
349
    v = F.tensor([0, 1, 2, 3])
350
    reduce_msg_shapes.clear()
351
    g.push(v)
352
353
354
355
356
357
358
359
360
    assert(reduce_msg_shapes == {(1, 3, D), (8, 1, D)})
    reduce_msg_shapes.clear()

    # update_all
    reduce_msg_shapes.clear()
    g.update_all()
    assert(reduce_msg_shapes == {(1, 8, D), (9, 1, D)})
    reduce_msg_shapes.clear()

361
362
363
364
365
366
367
368
def test_recv_0deg():
    # test recv with 0deg nodes;
    g = DGLGraph()
    g.add_nodes(2)
    g.add_edge(0, 1)
    def _message(edges):
        return {'m' : edges.src['h']}
    def _reduce(nodes):
369
        return {'h' : nodes.data['h'] + F.sum(nodes.mailbox['m'], 1)}
370
371
372
    def _apply(nodes):
        return {'h' : nodes.data['h'] * 2}
    def _init2(shape, dtype, ctx, ids):
373
        return 2 + F.zeros(shape, dtype, ctx)
374
375
376
377
378
    g.register_message_func(_message)
    g.register_reduce_func(_reduce)
    g.register_apply_node_func(_apply)
    g.set_n_initializer(_init2, 'h')
    # test#1: recv both 0deg and non-0deg nodes
379
    old = F.randn((2, 5))
380
381
382
383
384
    g.ndata['h'] = old
    g.send((0, 1))
    g.recv([0, 1])
    new = g.ndata.pop('h')
    # 0deg check: initialized with the func and got applied
385
    assert F.allclose(new[0], F.full_1d(5, 4, F.float32))
386
    # non-0deg check
387
    assert F.allclose(new[1], F.sum(old, 0) * 2)
388
389

    # test#2: recv only 0deg node is equal to apply
390
    old = F.randn((2, 5))
391
392
393
394
395
    g.ndata['h'] = old
    g.send((0, 1))
    g.recv(0)
    new = g.ndata.pop('h')
    # 0deg check: equal to apply_nodes
396
    assert F.allclose(new[0], 2 * old[0])
397
    # non-0deg check: untouched
398
    assert F.allclose(new[1], old[1])
399
400
401
402
403
404
405
406
407

def test_recv_0deg_newfld():
    # test recv with 0deg nodes; the reducer also creates a new field
    g = DGLGraph()
    g.add_nodes(2)
    g.add_edge(0, 1)
    def _message(edges):
        return {'m' : edges.src['h']}
    def _reduce(nodes):
408
        return {'h1' : nodes.data['h'] + F.sum(nodes.mailbox['m'], 1)}
409
410
411
    def _apply(nodes):
        return {'h1' : nodes.data['h1'] * 2}
    def _init2(shape, dtype, ctx, ids):
412
        return 2 + F.zeros(shape, dtype=dtype, ctx=ctx)
413
414
415
416
    g.register_message_func(_message)
    g.register_reduce_func(_reduce)
    g.register_apply_node_func(_apply)
    # test#1: recv both 0deg and non-0deg nodes
417
    old = F.randn((2, 5))
418
419
420
421
422
423
    g.set_n_initializer(_init2, 'h1')
    g.ndata['h'] = old
    g.send((0, 1))
    g.recv([0, 1])
    new = g.ndata.pop('h1')
    # 0deg check: initialized with the func and got applied
424
    assert F.allclose(new[0], F.full_1d(5, 4, dtype=F.float32))
425
    # non-0deg check
426
    assert F.allclose(new[1], F.sum(old, 0) * 2)
427
428

    # test#2: recv only 0deg node
429
    old = F.randn((2, 5))
430
    g.ndata['h'] = old
431
    g.ndata['h1'] = F.full((2, 5), -1, F.int64)  # this is necessary
432
433
434
435
    g.send((0, 1))
    g.recv(0)
    new = g.ndata.pop('h1')
    # 0deg check: fallback to apply
436
    assert F.allclose(new[0], F.full_1d(5, -2, F.int64))
437
    # non-0deg check: not changed
438
    assert F.allclose(new[1], F.full_1d(5, -1, F.int64))
439
440
441

def test_update_all_0deg():
    # test#1
442
    g = DGLGraph()
Minjie Wang's avatar
Minjie Wang committed
443
    g.add_nodes(5)
444
445
446
447
    g.add_edge(1, 0)
    g.add_edge(2, 0)
    g.add_edge(3, 0)
    g.add_edge(4, 0)
448
449
450
    def _message(edges):
        return {'m' : edges.src['h']}
    def _reduce(nodes):
451
        return {'h' : nodes.data['h'] + F.sum(nodes.mailbox['m'], 1)}
452
453
    def _apply(nodes):
        return {'h' : nodes.data['h'] * 2}
454
    def _init2(shape, dtype, ctx, ids):
455
        return 2 + F.zeros(shape, dtype, ctx)
456
    g.set_n_initializer(_init2, 'h')
457
    old_repr = F.randn((5, 5))
458
    g.ndata['h'] = old_repr
459
    g.update_all(_message, _reduce, _apply)
460
    new_repr = g.ndata['h']
461
462
    # the first row of the new_repr should be the sum of all the node
    # features; while the 0-deg nodes should be initialized by the
463
    # initializer and applied with UDF.
464
465
    assert F.allclose(new_repr[1:], 2*(2+F.zeros((4,5))))
    assert F.allclose(new_repr[0], 2 * F.sum(old_repr, 0))
466

467
468
469
470
471
472
473
474
    # test#2: graph with no edge
    g = DGLGraph()
    g.add_nodes(5)
    g.set_n_initializer(_init2, 'h')
    g.ndata['h'] = old_repr
    g.update_all(_message, _reduce, _apply)
    new_repr = g.ndata['h']
    # should fallback to apply
475
    assert F.allclose(new_repr, 2*old_repr)
476

477
def test_pull_0deg():
478
    g = DGLGraph()
Minjie Wang's avatar
Minjie Wang committed
479
    g.add_nodes(2)
480
    g.add_edge(0, 1)
481
482
483
    def _message(edges):
        return {'m' : edges.src['h']}
    def _reduce(nodes):
484
        return {'h' : nodes.data['h'] + F.sum(nodes.mailbox['m'], 1)}
485
486
487
    def _apply(nodes):
        return {'h' : nodes.data['h'] * 2}
    def _init2(shape, dtype, ctx, ids):
488
        return 2 + F.zeros(shape, dtype, ctx)
489
490
491
492
493
    g.register_message_func(_message)
    g.register_reduce_func(_reduce)
    g.register_apply_node_func(_apply)
    g.set_n_initializer(_init2, 'h')
    # test#1: pull both 0deg and non-0deg nodes
494
    old = F.randn((2, 5))
495
496
497
498
    g.ndata['h'] = old
    g.pull([0, 1])
    new = g.ndata.pop('h')
    # 0deg check: initialized with the func and got applied
499
    assert F.allclose(new[0], F.full_1d(5, 4, dtype=F.float32))
500
    # non-0deg check
501
    assert F.allclose(new[1], F.sum(old, 0) * 2)
502
503

    # test#2: pull only 0deg node
504
    old = F.randn((2, 5))
505
506
507
508
    g.ndata['h'] = old
    g.pull(0)
    new = g.ndata.pop('h')
    # 0deg check: fallback to apply
509
    assert F.allclose(new[0], 2*old[0])
510
    # non-0deg check: not touched
511
    assert F.allclose(new[1], old[1])
512

513
514
515
516
517
518
519
520
def test_send_multigraph():
    g = DGLGraph(multigraph=True)
    g.add_nodes(3)
    g.add_edge(0, 1)
    g.add_edge(0, 1)
    g.add_edge(0, 1)
    g.add_edge(2, 1)

521
522
523
524
525
    def _message_a(edges):
        return {'a': edges.data['a']}
    def _message_b(edges):
        return {'a': edges.data['a'] * 3}
    def _reduce(nodes):
526
        return {'a': F.max(nodes.mailbox['a'], 1)}
527
528

    def answer(*args):
529
        return F.max(F.stack(args, 0), 0)
530
531

    # send by eid
532
533
    old_repr = F.randn((4, 5))
    g.ndata['a'] = F.zeros((3, 5))
534
535
536
537
    g.edata['a'] = old_repr
    g.send([0, 2], message_func=_message_a)
    g.recv(1, _reduce)
    new_repr = g.ndata['a']
538
    assert F.allclose(new_repr[1], answer(old_repr[0], old_repr[2]))
539

540
    g.ndata['a'] = F.zeros((3, 5))
541
542
543
544
    g.edata['a'] = old_repr
    g.send([0, 2, 3], message_func=_message_a)
    g.recv(1, _reduce)
    new_repr = g.ndata['a']
545
    assert F.allclose(new_repr[1], answer(old_repr[0], old_repr[2], old_repr[3]))
546
547

    # send on multigraph
548
    g.ndata['a'] = F.zeros((3, 5))
549
550
551
552
    g.edata['a'] = old_repr
    g.send(([0, 2], [1, 1]), _message_a)
    g.recv(1, _reduce)
    new_repr = g.ndata['a']
553
    assert F.allclose(new_repr[1], F.max(old_repr, 0))
554
555

    # consecutive send and send_on
556
    g.ndata['a'] = F.zeros((3, 5))
557
558
559
560
561
    g.edata['a'] = old_repr
    g.send((2, 1), _message_a)
    g.send([0, 1], message_func=_message_b)
    g.recv(1, _reduce)
    new_repr = g.ndata['a']
562
    assert F.allclose(new_repr[1], answer(old_repr[0] * 3, old_repr[1] * 3, old_repr[3]))
563
564

    # consecutive send_on
565
    g.ndata['a'] = F.zeros((3, 5))
566
567
568
569
570
    g.edata['a'] = old_repr
    g.send(0, message_func=_message_a)
    g.send(1, message_func=_message_b)
    g.recv(1, _reduce)
    new_repr = g.ndata['a']
571
    assert F.allclose(new_repr[1], answer(old_repr[0], old_repr[1] * 3))
572
573

    # send_and_recv_on
574
    g.ndata['a'] = F.zeros((3, 5))
575
576
577
    g.edata['a'] = old_repr
    g.send_and_recv([0, 2, 3], message_func=_message_a, reduce_func=_reduce)
    new_repr = g.ndata['a']
578
579
    assert F.allclose(new_repr[1], answer(old_repr[0], old_repr[2], old_repr[3]))
    assert F.allclose(new_repr[[0, 2]], F.zeros((2, 5)))
580

581
582
583
584
585
586
587
588
def test_dynamic_addition():
    N = 3
    D = 1

    g = DGLGraph()

    # Test node addition
    g.add_nodes(N)
589
590
    g.ndata.update({'h1': F.randn((N, D)),
                    'h2': F.randn((N, D))})
591
    g.add_nodes(3)
592
    assert g.ndata['h1'].shape[0] == g.ndata['h2'].shape[0] == N + 3
593
594
595
596

    # Test edge addition
    g.add_edge(0, 1)
    g.add_edge(1, 0)
597
598
    g.edata.update({'h1': F.randn((2, D)),
                    'h2': F.randn((2, D))})
599
    assert g.edata['h1'].shape[0] == g.edata['h2'].shape[0] == 2
600
601

    g.add_edges([0, 2], [2, 0])
602
    g.edata['h1'] = F.randn((4, D))
603
    assert g.edata['h1'].shape[0] == g.edata['h2'].shape[0] == 4
604
605

    g.add_edge(1, 2)
606
    g.edges[4].data['h1'] = F.randn((1, D))
607
    assert g.edata['h1'].shape[0] == g.edata['h2'].shape[0] == 5
608

609
    # test add edge with part of the features
610
    g.add_edge(2, 1, {'h1': F.randn((1, D))})
611
612
    assert len(g.edata['h1']) == len(g.edata['h2'])

613

Haibin Lin's avatar
Haibin Lin committed
614
615
616
617
618
def test_repr():
    G = dgl.DGLGraph()
    G.add_nodes(10)
    G.add_edge(0, 1)
    repr_string = G.__repr__()
Minjie Wang's avatar
Minjie Wang committed
619
    print(repr_string)
620
    G.ndata['x'] = F.zeros((10, 5))
Haibin Lin's avatar
Haibin Lin committed
621
    G.add_edges([0, 1], 2)
622
    G.edata['y'] = F.zeros((3, 4))
Haibin Lin's avatar
Haibin Lin committed
623
    repr_string = G.__repr__()
Minjie Wang's avatar
Minjie Wang committed
624
    print(repr_string)
Haibin Lin's avatar
Haibin Lin committed
625

626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647

def test_group_apply_edges():
    def edge_udf(edges):
        h = F.sum(edges.data['feat'] * (edges.src['h'] + edges.dst['h']), dim=2)
        normalized_feat = F.softmax(h, dim=1)
        return {"norm_feat": normalized_feat}

    g = DGLGraph()
    g.add_nodes(10)
    g.add_edges(0, [1, 2, 3, 4, 5, 6, 7, 8])
    g.add_edges(1, [2, 3, 4, 6, 7, 8])
    g.add_edges(2, [2, 3, 4, 5, 6, 7, 8])

    g.ndata['h'] = F.randn((g.number_of_nodes(), D))
    g.edata['feat'] = F.randn((g.number_of_edges(), D))

    def _test(group_by):
        g.group_apply_edges(group_by=group_by, func=edge_udf)
        if group_by == 'src':
            u, v, eid = g.out_edges(1, form='all')
        else:
            u, v, eid = g.in_edges(5, form='all')
648
649
        out_feat = g.edges[eid].data['norm_feat']
        result = (g.nodes[u].data['h'] + g.nodes[v].data['h']) * g.edges[eid].data['feat']
650
651
652
653
654
655
656
657
658
        result = F.softmax(F.sum(result, dim=1), dim=0)
        assert F.allclose(out_feat, result)

    # test group by source nodes
    _test('src')

    # test group by destination nodes
    _test('dst')

Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680

# GitHub issue #1036
def test_group_apply_edges2():
    m = ssp.random(10, 10, 0.2)
    g = DGLGraph(m, readonly=True)
    g.ndata['deg'] = g.in_degrees()
    g.ndata['id'] = F.arange(0, g.number_of_nodes())
    g.edata['id'] = F.arange(0, g.number_of_edges())

    def apply(edges):
        w = edges.data['id']
        n_nodes, deg = w.shape

        dst = edges.dst['id'][:, 0]
        eid1 = F.asnumpy(g.in_edges(dst, 'eid')).reshape(n_nodes, deg).sort(1)
        eid2 = F.asnumpy(edges.data['id']).sort(1)
        assert np.array_equal(eid1, eid2)
        return {'id2': w}

    g.group_apply_edges('dst', apply, inplace=True)


681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
def test_local_var():
    g = DGLGraph(nx.path_graph(5))
    g.ndata['h'] = F.zeros((g.number_of_nodes(), 3))
    g.edata['w'] = F.zeros((g.number_of_edges(), 4))
    # test override
    def foo(g):
        g = g.local_var()
        g.ndata['h'] = F.ones((g.number_of_nodes(), 3))
        g.edata['w'] = F.ones((g.number_of_edges(), 4))
    foo(g)
    assert F.allclose(g.ndata['h'], F.zeros((g.number_of_nodes(), 3)))
    assert F.allclose(g.edata['w'], F.zeros((g.number_of_edges(), 4)))
    # test out-place update
    def foo(g):
        g = g.local_var()
        g.nodes[[2, 3]].data['h'] = F.ones((2, 3))
        g.edges[[2, 3]].data['w'] = F.ones((2, 4))
    foo(g)
    assert F.allclose(g.ndata['h'], F.zeros((g.number_of_nodes(), 3)))
    assert F.allclose(g.edata['w'], F.zeros((g.number_of_edges(), 4)))
    # test out-place update 2
    def foo(g):
        g = g.local_var()
        g.apply_nodes(lambda nodes: {'h' : nodes.data['h'] + 10}, [2, 3])
        g.apply_edges(lambda edges: {'w' : edges.data['w'] + 10}, [2, 3])
    foo(g)
    assert F.allclose(g.ndata['h'], F.zeros((g.number_of_nodes(), 3)))
    assert F.allclose(g.edata['w'], F.zeros((g.number_of_edges(), 4)))
    # test auto-pop
    def foo(g):
        g = g.local_var()
        g.ndata['hh'] = F.ones((g.number_of_nodes(), 3))
        g.edata['ww'] = F.ones((g.number_of_edges(), 4))
    foo(g)
    assert 'hh' not in g.ndata
    assert 'ww' not in g.edata

Mufei Li's avatar
Mufei Li committed
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
    # test initializer1
    g = DGLGraph()
    g.add_nodes(2)
    g.add_edges([0, 1], [1, 1])
    g.set_n_initializer(dgl.init.zero_initializer)
    def foo(g):
        g = g.local_var()
        g.nodes[0].data['h'] = F.ones((1, 1))
        assert F.allclose(g.ndata['h'], F.tensor([[1.], [0.]]))
    foo(g)
    # test initializer2
    def foo_e_initializer(shape, dtype, ctx, id_range):
        return F.ones(shape)
    g.set_e_initializer(foo_e_initializer, field='h')
    def foo(g):
        g = g.local_var()
        g.edges[0, 1].data['h'] = F.ones((1, 1))
        assert F.allclose(g.edata['h'], F.ones((2, 1)))
        g.edges[0, 1].data['w'] = F.ones((1, 1))
        assert F.allclose(g.edata['w'], F.tensor([[1.], [0.]]))
    foo(g)

740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
def test_local_scope():
    g = DGLGraph(nx.path_graph(5))
    g.ndata['h'] = F.zeros((g.number_of_nodes(), 3))
    g.edata['w'] = F.zeros((g.number_of_edges(), 4))
    # test override
    def foo(g):
        with g.local_scope():
            g.ndata['h'] = F.ones((g.number_of_nodes(), 3))
            g.edata['w'] = F.ones((g.number_of_edges(), 4))
    foo(g)
    assert F.allclose(g.ndata['h'], F.zeros((g.number_of_nodes(), 3)))
    assert F.allclose(g.edata['w'], F.zeros((g.number_of_edges(), 4)))
    # test out-place update
    def foo(g):
        with g.local_scope():
            g.nodes[[2, 3]].data['h'] = F.ones((2, 3))
            g.edges[[2, 3]].data['w'] = F.ones((2, 4))
    foo(g)
    assert F.allclose(g.ndata['h'], F.zeros((g.number_of_nodes(), 3)))
    assert F.allclose(g.edata['w'], F.zeros((g.number_of_edges(), 4)))
    # test out-place update 2
    def foo(g):
        with g.local_scope():
            g.apply_nodes(lambda nodes: {'h' : nodes.data['h'] + 10}, [2, 3])
            g.apply_edges(lambda edges: {'w' : edges.data['w'] + 10}, [2, 3])
    foo(g)
    assert F.allclose(g.ndata['h'], F.zeros((g.number_of_nodes(), 3)))
    assert F.allclose(g.edata['w'], F.zeros((g.number_of_edges(), 4)))
    # test auto-pop
    def foo(g):
        with g.local_scope():
            g.ndata['hh'] = F.ones((g.number_of_nodes(), 3))
            g.edata['ww'] = F.ones((g.number_of_edges(), 4))
    foo(g)
    assert 'hh' not in g.ndata
    assert 'ww' not in g.edata

    # test nested scope
    def foo(g):
        with g.local_scope():
            g.ndata['hh'] = F.ones((g.number_of_nodes(), 3))
            g.edata['ww'] = F.ones((g.number_of_edges(), 4))
            with g.local_scope():
                g.ndata['hhh'] = F.ones((g.number_of_nodes(), 3))
                g.edata['www'] = F.ones((g.number_of_edges(), 4))
            assert 'hhh' not in g.ndata
            assert 'www' not in g.edata
    foo(g)
    assert 'hh' not in g.ndata
    assert 'ww' not in g.edata
790

Mufei Li's avatar
Mufei Li committed
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
    # test initializer1
    g = DGLGraph()
    g.add_nodes(2)
    g.add_edges([0, 1], [1, 1])
    g.set_n_initializer(dgl.init.zero_initializer)
    def foo(g):
        with g.local_scope():
            g.nodes[0].data['h'] = F.ones((1, 1))
            assert F.allclose(g.ndata['h'], F.tensor([[1.], [0.]]))
    foo(g)
    # test initializer2
    def foo_e_initializer(shape, dtype, ctx, id_range):
        return F.ones(shape)
    g.set_e_initializer(foo_e_initializer, field='h')
    def foo(g):
        with g.local_scope():
            g.edges[0, 1].data['h'] = F.ones((1, 1))
            assert F.allclose(g.edata['h'], F.ones((2, 1)))
            g.edges[0, 1].data['w'] = F.ones((1, 1))
            assert F.allclose(g.edata['w'], F.tensor([[1.], [0.]]))
    foo(g)

813
if __name__ == '__main__':
814
    test_nx_conversion()
815
    test_batch_setter_getter()
816
    test_batch_setter_autograd()
817
    test_batch_send()
818
    test_batch_recv()
819
    test_apply_nodes()
820
    test_apply_edges()
821
    test_update_routines()
822
823
824
    test_recv_0deg()
    test_recv_0deg_newfld()
    test_update_all_0deg()
825
    test_pull_0deg()
826
    test_send_multigraph()
827
    test_dynamic_addition()
Haibin Lin's avatar
Haibin Lin committed
828
    test_repr()
829
    test_group_apply_edges()
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
830
    test_group_apply_edges2()
831
832
    test_local_var()
    test_local_scope()