test_basics.py 21.1 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
import unittest
9
10
11
12

D = 5
reduce_msg_shapes = set()

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

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

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

28
def generate_graph(grad=False):
29
    g = DGLGraph()
30
    g.add_nodes(10) # 10 nodes
31
    # create a graph where 0 is the source and 9 is the sink
Minjie Wang's avatar
Minjie Wang committed
32
    # 17 edges
33
34
35
36
37
    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)
38
    g = g.to(F.ctx())
39
40
41
42
43
44
    ncol = F.randn((10, D))
    ecol = F.randn((17, D))
    if grad:
        ncol = F.attach_grad(ncol)
        ecol = F.attach_grad(ecol)

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

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

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

VoVAllen's avatar
VoVAllen committed
131

132
133
def test_batch_setter_autograd():
    g = generate_graph(grad=True)
134
    h1 = g.ndata['h']
135
    # partial set
136
    v = F.tensor([1, 2, 8], g.idtype)
137
138
139
140
    hh = F.attach_grad(F.zeros((len(v), D)))
    with F.record_grad():
        g.nodes[v].data['h'] = hh
        h2 = g.ndata['h']
VoVAllen's avatar
VoVAllen committed
141
        F.backward(h2, F.ones((10, D)) * 2)
142
143
    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.]))
144

145
def _test_nx_conversion():
146
147
148
    # check conversion between networkx and DGLGraph

    def _check_nx_feature(nxg, nf, ef):
Minjie Wang's avatar
Minjie Wang committed
149
150
        # check node and edge feature of nxg
        # this is used to check to_networkx
151
152
153
154
155
156
157
        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]:
158
                    node_feat[k].append(F.unsqueeze(attr[k], 0))
159
            for k in node_feat:
160
161
                feat = F.cat(node_feat[k], 0)
                assert F.allclose(feat, nf[k])
162
163
164
165
166
167
168
169
        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:
170
                    edge_feat[k][eid] = F.unsqueeze(attr[k], 0)
171
            for k in edge_feat:
172
173
                feat = F.cat(edge_feat[k], 0)
                assert F.allclose(feat, ef[k])
174
175
176
        else:
            assert len(ef) == 0

177
178
179
180
181
    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))
182
    g = DGLGraph()
183
184
185
186
187
188
189
190
191
192
193
    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
194
    # convert to DGLGraph, nx graph has id in edge feature
195
    # use id feature to test non-tensor copy
196
    g = dgl.from_networkx(nxg, node_attrs=['n1'], edge_attrs=['e1', 'id'])
Minjie Wang's avatar
Minjie Wang committed
197
    # check graph size
198
199
    assert g.number_of_nodes() == 5
    assert g.number_of_edges() == 4
Minjie Wang's avatar
Minjie Wang committed
200
201
202
203
204
    # 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
205
    assert F.allclose(g.ndata['n1'], n1)
Minjie Wang's avatar
Minjie Wang committed
206
    # with id in nx edge feature, e1 should follow original order
207
    assert F.allclose(g.edata['e1'], e1)
208
    assert F.array_equal(F.astype(g.edata['id'], F.int64), F.copy_to(F.arange(0, 4), F.cpu()))
209

Minjie Wang's avatar
Minjie Wang committed
210
    # test conversion after modifying DGLGraph
211
    g.edata.pop('id') # pop id so we don't need to provide id when adding edges
212
213
    new_n = F.randn((2, 3))
    new_e = F.randn((3, 5))
214
215
216
    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})
217
218
    n1 = F.cat((n1, new_n), 0)
    e1 = F.cat((e1, new_e), 0)
219
220
221
222
223
224
    # 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
225
226
227
228
229
    # 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
230
    g = dgl.from_networkx(nxg, node_attrs=['n1'], edge_attrs=['e1'])
Minjie Wang's avatar
Minjie Wang committed
231
232
233
234
235
236
237
    # 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
238
    assert F.allclose(g.ndata['n1'], n1)
Minjie Wang's avatar
Minjie Wang committed
239
240
241
    # edge feature order follows nxg.edges()
    edge_feat = []
    for _, _, attr in nxg.edges(data=True):
242
243
244
        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
245

246
247
248
249
250
    # 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():
VoVAllen's avatar
VoVAllen committed
251
        nxg.nodes[u]['h'] = F.tensor([u])
252
253
254
    for u, v, d in nxg.edges(data=True):
        d['h'] = F.tensor([u, v])

255
    g = dgl.from_networkx(nxg, node_attrs=['h'], edge_attrs=['h'])
256
257
258
259
260
261
262
263
    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.]]))

264
265
266
267
268
def test_apply_nodes():
    def _upd(nodes):
        return {'h' : nodes.data['h'] * 2}
    g = generate_graph()
    old = g.ndata['h']
269
    g.apply_nodes(_upd)
270
    assert F.allclose(old * 2, g.ndata['h'])
271
    u = F.tensor([0, 3, 4, 6], g.idtype)
272
    g.apply_nodes(lambda nodes : {'h' : nodes.data['h'] * 0.}, u)
273
    assert F.allclose(F.gather_row(g.ndata['h'], u), F.zeros((4, D)))
274

275
def test_apply_edges():
276
277
278
279
    def _upd(edges):
        return {'w' : edges.data['w'] * 2}
    g = generate_graph()
    old = g.edata['w']
280
    g.apply_edges(_upd)
281
    assert F.allclose(old * 2, g.edata['w'])
282
283
    u = F.tensor([0, 0, 0, 4, 5, 6], g.idtype)
    v = F.tensor([1, 2, 3, 9, 9, 9], g.idtype)
284
    g.apply_edges(lambda edges : {'w' : edges.data['w'] * 0.}, (u, v))
285
    eid = F.tensor(g.edge_ids(u, v))
286
    assert F.allclose(F.gather_row(g.edata['w'], eid), F.zeros((6, D)))
287

288
289
290
def test_update_routines():
    g = generate_graph()

291
    # send_and_recv
292
    reduce_msg_shapes.clear()
293
294
    u = [0, 0, 0, 4, 5, 6]
    v = [1, 2, 3, 9, 9, 9]
295
    g.send_and_recv((u, v), message_func, reduce_func, apply_node_func)
296
297
    assert(reduce_msg_shapes == {(1, 3, D), (3, 1, D)})
    reduce_msg_shapes.clear()
298
299
300
    try:
        g.send_and_recv([u, v])
        assert False
301
    except:
302
        pass
303

304
    # pull
305
    v = F.tensor([1, 2, 3, 9], g.idtype)
306
    reduce_msg_shapes.clear()
307
    g.pull(v, message_func, reduce_func, apply_node_func)
308
309
310
    assert(reduce_msg_shapes == {(1, 8, D), (3, 1, D)})
    reduce_msg_shapes.clear()

311
    # push
312
    v = F.tensor([0, 1, 2, 3], g.idtype)
313
    reduce_msg_shapes.clear()
314
    g.push(v, message_func, reduce_func, apply_node_func)
315
316
317
318
319
    assert(reduce_msg_shapes == {(1, 3, D), (8, 1, D)})
    reduce_msg_shapes.clear()

    # update_all
    reduce_msg_shapes.clear()
320
    g.update_all(message_func, reduce_func, apply_node_func)
321
322
323
    assert(reduce_msg_shapes == {(1, 8, D), (9, 1, D)})
    reduce_msg_shapes.clear()

324
325
def test_update_all_0deg():
    # test#1
326
    g = DGLGraph()
327
    g = g.to(F.ctx())
Minjie Wang's avatar
Minjie Wang committed
328
    g.add_nodes(5)
329
330
331
332
    g.add_edge(1, 0)
    g.add_edge(2, 0)
    g.add_edge(3, 0)
    g.add_edge(4, 0)
333
334
335
    def _message(edges):
        return {'m' : edges.src['h']}
    def _reduce(nodes):
336
        return {'h' : nodes.data['h'] + F.sum(nodes.mailbox['m'], 1)}
337
338
    def _apply(nodes):
        return {'h' : nodes.data['h'] * 2}
339
    def _init2(shape, dtype, ctx, ids):
340
        return 2 + F.zeros(shape, dtype, ctx)
341
    g.set_n_initializer(_init2, 'h')
342
    old_repr = F.randn((5, 5))
343
    g.ndata['h'] = old_repr
344
    g.update_all(_message, _reduce, _apply)
345
    new_repr = g.ndata['h']
346
347
    # 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
348
    # initializer and applied with UDF.
349
350
    assert F.allclose(new_repr[1:], 2*(2+F.zeros((4,5))))
    assert F.allclose(new_repr[0], 2 * F.sum(old_repr, 0))
351

352
353
    # test#2: graph with no edge
    g = DGLGraph()
354
    g = g.to(F.ctx())
355
356
357
358
359
360
    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
361
    assert F.allclose(new_repr, 2*old_repr)
362

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

    # test#2: pull only 0deg node
388
    old = F.randn((2, 5))
389
    g.ndata['h'] = old
390
    g.pull(0, _message, _reduce, _apply)
391
392
    new = g.ndata.pop('h')
    # 0deg check: fallback to apply
393
    assert F.allclose(new[0], 2*old[0])
394
    # non-0deg check: not touched
395
    assert F.allclose(new[1], old[1])
396

397
398
399
400
401
def test_dynamic_addition():
    N = 3
    D = 1

    g = DGLGraph()
402
    g = g.to(F.ctx())
403
404
405

    # Test node addition
    g.add_nodes(N)
406
407
    g.ndata.update({'h1': F.randn((N, D)),
                    'h2': F.randn((N, D))})
408
    g.add_nodes(3)
409
    assert g.ndata['h1'].shape[0] == g.ndata['h2'].shape[0] == N + 3
410
411
412
413

    # Test edge addition
    g.add_edge(0, 1)
    g.add_edge(1, 0)
414
415
    g.edata.update({'h1': F.randn((2, D)),
                    'h2': F.randn((2, D))})
416
    assert g.edata['h1'].shape[0] == g.edata['h2'].shape[0] == 2
417
418

    g.add_edges([0, 2], [2, 0])
419
    g.edata['h1'] = F.randn((4, D))
420
    assert g.edata['h1'].shape[0] == g.edata['h2'].shape[0] == 4
421
422

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

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

430

Haibin Lin's avatar
Haibin Lin committed
431
def test_repr():
432
433
434
435
436
    g = dgl.DGLGraph()
    g = g.to(F.ctx())
    g.add_nodes(10)
    g.add_edge(0, 1)
    repr_string = g.__repr__()
Minjie Wang's avatar
Minjie Wang committed
437
    print(repr_string)
438
439
440
441
    g.ndata['x'] = F.zeros((10, 5))
    g.add_edges([0, 1], 2)
    g.edata['y'] = F.zeros((3, 4))
    repr_string = g.__repr__()
Minjie Wang's avatar
Minjie Wang committed
442
    print(repr_string)
Haibin Lin's avatar
Haibin Lin committed
443

444
445
446
447
448
449
450
451

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()
452
    g = g.to(F.ctx())
453
454
455
456
457
458
459
460
461
462
463
464
465
466
    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')
467
468
        out_feat = g.edges[eid].data['norm_feat']
        result = (g.nodes[u].data['h'] + g.nodes[v].data['h']) * g.edges[eid].data['feat']
469
470
471
472
473
474
475
476
477
        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
478
479

# GitHub issue #1036
480
@unittest.skipIf(dgl.backend.backend_name == "tensorflow", reason="TF doesn't support inplace update")
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
481
482
483
def test_group_apply_edges2():
    m = ssp.random(10, 10, 0.2)
    g = DGLGraph(m, readonly=True)
484
    g = g.to(F.ctx())
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
485
    g.ndata['deg'] = g.in_degrees()
486
487
    g.ndata['id'] = F.arange(0, g.number_of_nodes(), g.idtype)
    g.edata['id'] = F.arange(0, g.number_of_edges(), g.idtype)
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
488
489
490
491
492
493

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

        dst = edges.dst['id'][:, 0]
494
        # TODO: tmp hack
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
495
496
497
498
499
500
501
502
        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)


503
504
def test_local_var():
    g = DGLGraph(nx.path_graph(5))
505
    g = g.to(F.ctx())
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
    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
541
542
    # test initializer1
    g = DGLGraph()
543
    g = g.to(F.ctx())
Mufei Li's avatar
Mufei Li committed
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
    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)

564
565
def test_local_scope():
    g = DGLGraph(nx.path_graph(5))
566
    g = g.to(F.ctx())
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
    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
615

Mufei Li's avatar
Mufei Li committed
616
617
    # test initializer1
    g = DGLGraph()
618
    g = g.to(F.ctx())
Mufei Li's avatar
Mufei Li committed
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
    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)

639
if __name__ == '__main__':
640
    #test_nx_conversion()
641
    test_batch_setter_getter()
642
    test_batch_setter_autograd()
643
    test_batch_send()
644
    test_batch_recv()
645
    test_apply_nodes()
646
    test_apply_edges()
647
    test_update_routines()
648
649
650
    test_recv_0deg()
    test_recv_0deg_newfld()
    test_update_all_0deg()
651
    test_pull_0deg()
652
    test_send_multigraph()
653
    test_dynamic_addition()
Haibin Lin's avatar
Haibin Lin committed
654
    test_repr()
655
    test_group_apply_edges()
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
656
    test_group_apply_edges2()
657
658
    test_local_var()
    test_local_scope()