test_heterograph.py 120 KB
Newer Older
1
2
3
4
5
6
7
import itertools
import multiprocessing as mp
import unittest
from collections import Counter

import backend as F

8
9
import dgl
import dgl.function as fn
10
import networkx as nx
11
import numpy as np
12
import pytest
13
import scipy.sparse as ssp
14
from dgl import DGLError
15
from scipy.sparse import rand
16
17
18
19
20
21
from utils import (
    assert_is_identical_hetero,
    check_graph_equal,
    get_cases,
    parametrize_idtype,
)
22

23

24
def create_test_heterograph(idtype):
25
    # test heterograph from the docstring, plus a user -- wishes -- game relation
Minjie Wang's avatar
Minjie Wang committed
26
27
28
29
30
31
    # 3 users, 2 games, 2 developers
    # metagraph:
    #    ('user', 'follows', 'user'),
    #    ('user', 'plays', 'game'),
    #    ('user', 'wishes', 'game'),
    #    ('developer', 'develops', 'game')])
32

33
34
35
36
37
38
39
40
41
42
    g = dgl.heterograph(
        {
            ("user", "follows", "user"): ([0, 1], [1, 2]),
            ("user", "plays", "game"): ([0, 1, 2, 1], [0, 0, 1, 1]),
            ("user", "wishes", "game"): ([0, 2], [1, 0]),
            ("developer", "develops", "game"): ([0, 1], [0, 1]),
        },
        idtype=idtype,
        device=F.ctx(),
    )
43
44
    assert g.idtype == idtype
    assert g.device == F.ctx()
45
46
    return g

47

48
def create_test_heterograph1(idtype):
Minjie Wang's avatar
Minjie Wang committed
49
    edges = []
50
51
52
53
54
    edges.extend([(0, 1), (1, 2)])  # follows
    edges.extend([(0, 3), (1, 3), (2, 4), (1, 4)])  # plays
    edges.extend([(0, 4), (2, 3)])  # wishes
    edges.extend([(5, 3), (6, 4)])  # develops
    edges = tuple(zip(*edges))
Minjie Wang's avatar
Minjie Wang committed
55
56
    ntypes = F.tensor([0, 0, 0, 1, 1, 2, 2])
    etypes = F.tensor([0, 0, 1, 1, 1, 1, 2, 2, 3, 3])
57
    g0 = dgl.graph(edges, idtype=idtype, device=F.ctx())
Minjie Wang's avatar
Minjie Wang committed
58
59
    g0.ndata[dgl.NTYPE] = ntypes
    g0.edata[dgl.ETYPE] = etypes
60
61
62
63
64
65
    return dgl.to_heterogeneous(
        g0,
        ["user", "game", "developer"],
        ["follows", "plays", "wishes", "develops"],
    )

Minjie Wang's avatar
Minjie Wang committed
66

67
def create_test_heterograph2(idtype):
68
69
70
71
72
73
74
75
76
77
78
79
80
    g = dgl.heterograph(
        {
            ("user", "follows", "user"): ([0, 1], [1, 2]),
            ("user", "plays", "game"): ([0, 1, 2, 1], [0, 0, 1, 1]),
            ("user", "wishes", "game"): ("csr", ([0, 1, 1, 2], [1, 0], [])),
            ("developer", "develops", "game"): (
                "csc",
                ([0, 1, 2], [0, 1], [0, 1]),
            ),
        },
        idtype=idtype,
        device=F.ctx(),
    )
81
82
    assert g.idtype == idtype
    assert g.device == F.ctx()
83
84
    return g

85

86
def create_test_heterograph3(idtype):
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
    g = dgl.heterograph(
        {
            ("user", "plays", "game"): (
                F.tensor([0, 1, 1, 2], dtype=idtype),
                F.tensor([0, 0, 1, 1], dtype=idtype),
            ),
            ("developer", "develops", "game"): (
                F.tensor([0, 1], dtype=idtype),
                F.tensor([0, 1], dtype=idtype),
            ),
        },
        idtype=idtype,
        device=F.ctx(),
    )

    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()
    )
    g.nodes["developer"].data["h"] = F.copy_to(
        F.tensor([3, 3], dtype=idtype), ctx=F.ctx()
    )
    g.edges["plays"].data["h"] = F.copy_to(
        F.tensor([1, 1, 1, 1], dtype=idtype), ctx=F.ctx()
    )
114
115
    return g

116

117
def create_test_heterograph4(idtype):
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
    g = dgl.heterograph(
        {
            ("user", "follows", "user"): (
                F.tensor([0, 1, 1, 2, 2, 2], dtype=idtype),
                F.tensor([0, 0, 1, 1, 2, 2], dtype=idtype),
            ),
            ("user", "plays", "game"): (
                F.tensor([0, 1], dtype=idtype),
                F.tensor([0, 1], dtype=idtype),
            ),
        },
        idtype=idtype,
        device=F.ctx(),
    )
    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()
    )
    g.edges["follows"].data["h"] = F.copy_to(
        F.tensor([1, 2, 3, 4, 5, 6], dtype=idtype), ctx=F.ctx()
    )
    g.edges["plays"].data["h"] = F.copy_to(
        F.tensor([1, 2], dtype=idtype), ctx=F.ctx()
    )
144
145
    return g

146

147
def create_test_heterograph5(idtype):
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
    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(),
    )
    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()
    )
    g.edges["follows"].data["h"] = F.copy_to(
        F.tensor([1, 2], dtype=idtype), ctx=F.ctx()
    )
    g.edges["plays"].data["h"] = F.copy_to(
        F.tensor([1, 2], dtype=idtype), ctx=F.ctx()
    )
174
175
    return g

176

Minjie Wang's avatar
Minjie Wang committed
177
178
179
def get_redfn(name):
    return getattr(F, name)

180

nv-dlasalle's avatar
nv-dlasalle committed
181
@parametrize_idtype
182
183
184
185
186
def test_create(idtype):
    device = F.ctx()
    g0 = create_test_heterograph(idtype)
    g1 = create_test_heterograph1(idtype)
    g2 = create_test_heterograph2(idtype)
187
    assert set(g0.ntypes) == set(g1.ntypes) == set(g2.ntypes)
188
189
190
191
192
    assert (
        set(g0.canonical_etypes)
        == set(g1.canonical_etypes)
        == set(g2.canonical_etypes)
    )
Minjie Wang's avatar
Minjie Wang committed
193

194
195
196
197
198
    # Create a bipartite graph from a SciPy matrix
    src_ids = np.array([2, 3, 4])
    dst_ids = np.array([1, 2, 3])
    eweight = np.array([0.2, 0.3, 0.5])
    sp_mat = ssp.coo_matrix((eweight, (src_ids, dst_ids)))
199
200
201
202
203
204
205
206
    g = dgl.bipartite_from_scipy(
        sp_mat,
        utype="user",
        etype="plays",
        vtype="game",
        idtype=idtype,
        device=device,
    )
207
208
209
210
211
212
213
214
    assert g.idtype == idtype
    assert g.device == device
    assert g.num_src_nodes() == 5
    assert g.num_dst_nodes() == 4
    assert g.num_edges() == 3
    src, dst = g.edges()
    assert F.allclose(src, F.tensor([2, 3, 4], dtype=idtype))
    assert F.allclose(dst, F.tensor([1, 2, 3], dtype=idtype))
215
216
217
218
219
220
221
222
223
224
    g = dgl.bipartite_from_scipy(
        sp_mat,
        utype="_U",
        etype="_E",
        vtype="_V",
        eweight_name="w",
        idtype=idtype,
        device=device,
    )
    assert F.allclose(g.edata["w"], F.tensor(eweight))
225
226
227

    # Create a bipartite graph from a NetworkX graph
    nx_g = nx.DiGraph()
228
229
230
    nx_g.add_nodes_from(
        [1, 3], bipartite=0, feat1=np.zeros((2)), feat2=np.ones((2))
    )
231
232
233
    nx_g.add_nodes_from([2, 4, 5], bipartite=1, feat3=np.zeros((3)))
    nx_g.add_edge(1, 4, weight=np.ones((1)), eid=np.array([1]))
    nx_g.add_edge(3, 5, weight=np.ones((1)), eid=np.array([0]))
234
235
236
237
238
239
240
241
    g = dgl.bipartite_from_networkx(
        nx_g,
        utype="user",
        etype="plays",
        vtype="game",
        idtype=idtype,
        device=device,
    )
242
243
    assert g.idtype == idtype
    assert g.device == device
244
245
246
247
248
249
    assert g.num_src_nodes() == 2
    assert g.num_dst_nodes() == 3
    assert g.num_edges() == 2
    src, dst = g.edges()
    assert F.allclose(src, F.tensor([0, 1], dtype=idtype))
    assert F.allclose(dst, F.tensor([1, 2], dtype=idtype))
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
    g = dgl.bipartite_from_networkx(
        nx_g,
        utype="_U",
        etype="_E",
        vtype="V",
        u_attrs=["feat1", "feat2"],
        e_attrs=["weight"],
        v_attrs=["feat3"],
    )
    assert F.allclose(g.srcdata["feat1"], F.tensor(np.zeros((2, 2))))
    assert F.allclose(g.srcdata["feat2"], F.tensor(np.ones((2, 2))))
    assert F.allclose(g.dstdata["feat3"], F.tensor(np.zeros((3, 3))))
    assert F.allclose(g.edata["weight"], F.tensor(np.ones((2, 1))))
    g = dgl.bipartite_from_networkx(
        nx_g,
        utype="_U",
        etype="_E",
        vtype="V",
        edge_id_attr_name="eid",
        idtype=idtype,
        device=device,
    )
272
273
274
    src, dst = g.edges()
    assert F.allclose(src, F.tensor([1, 0], dtype=idtype))
    assert F.allclose(dst, F.tensor([2, 1], dtype=idtype))
Minjie Wang's avatar
Minjie Wang committed
275
276

    # create from scipy
277
    spmat = ssp.coo_matrix(([1, 1, 1], ([0, 0, 1], [2, 3, 2])), shape=(4, 4))
278
279
280
    g = dgl.from_scipy(spmat, idtype=idtype, device=device)
    assert g.num_nodes() == 4
    assert g.num_edges() == 3
281
282
    assert g.idtype == idtype
    assert g.device == device
Minjie Wang's avatar
Minjie Wang committed
283

284
    # test inferring number of nodes for heterograph
285
286
287
288
289
290
291
292
293
294
295
296
    g = dgl.heterograph(
        {
            ("l0", "e0", "l1"): ([0, 0], [1, 2]),
            ("l0", "e1", "l2"): ([2], [2]),
            ("l2", "e2", "l2"): ([1, 3], [1, 3]),
        },
        idtype=idtype,
        device=device,
    )
    assert g.num_nodes("l0") == 3
    assert g.num_nodes("l1") == 3
    assert g.num_nodes("l2") == 4
297
298
    assert g.idtype == idtype
    assert g.device == device
299

300
301
    # test if validate flag works
    # homo graph
302
    with pytest.raises(DGLError):
303
304
        g = dgl.graph(
            ([0, 0, 0, 1, 1, 2], [0, 1, 2, 0, 1, 2]),
305
            num_nodes=2,
306
307
            idtype=idtype,
            device=device,
308
        )
309

310
311
    # bipartite graph
    def _test_validate_bipartite(card):
312
        with pytest.raises(DGLError):
313
314
315
316
317
318
            g = dgl.heterograph(
                {("_U", "_E", "_V"): ([0, 0, 1, 1, 2], [1, 1, 2, 2, 3])},
                {"_U": card[0], "_V": card[1]},
                idtype=idtype,
                device=device,
            )
319
320
321
322

    _test_validate_bipartite((3, 3))
    _test_validate_bipartite((2, 4))

323
324
325
    # test from_scipy
    num_nodes = 10
    density = 0.25
326
    for fmt in ["csr", "coo", "csc"]:
327
        adj = rand(num_nodes, num_nodes, density=density, format=fmt)
328
        g = dgl.from_scipy(adj, eweight_name="w", idtype=idtype)
329
330
        assert g.idtype == idtype
        assert g.device == F.cpu()
331
332
333
334
        assert F.array_equal(
            g.edata["w"], F.copy_to(F.tensor(adj.data), F.cpu())
        )

335

336
337
338
339
340
341
342
343
def test_create2():
    mat = ssp.random(20, 30, 0.1)

    # coo
    mat = mat.tocoo()
    row = F.tensor(mat.row, dtype=F.int64)
    col = F.tensor(mat.col, dtype=F.int64)
    g = dgl.heterograph(
344
345
346
        {("A", "AB", "B"): ("coo", (row, col))},
        num_nodes_dict={"A": 20, "B": 30},
    )
347
348
349
350
351
352
353

    # csr
    mat = mat.tocsr()
    indptr = F.tensor(mat.indptr, dtype=F.int64)
    indices = F.tensor(mat.indices, dtype=F.int64)
    data = F.tensor([], dtype=F.int64)
    g = dgl.heterograph(
354
355
356
        {("A", "AB", "B"): ("csr", (indptr, indices, data))},
        num_nodes_dict={"A": 20, "B": 30},
    )
357
358
359
360
361
362
363

    # csc
    mat = mat.tocsc()
    indptr = F.tensor(mat.indptr, dtype=F.int64)
    indices = F.tensor(mat.indices, dtype=F.int64)
    data = F.tensor([], dtype=F.int64)
    g = dgl.heterograph(
364
365
366
367
        {("A", "AB", "B"): ("csc", (indptr, indices, data))},
        num_nodes_dict={"A": 20, "B": 30},
    )

368

nv-dlasalle's avatar
nv-dlasalle committed
369
@parametrize_idtype
370
371
def test_query(idtype):
    g = create_test_heterograph(idtype)
372

373
    ntypes = ["user", "game", "developer"]
Minjie Wang's avatar
Minjie Wang committed
374
    canonical_etypes = [
375
376
377
378
379
380
        ("user", "follows", "user"),
        ("user", "plays", "game"),
        ("user", "wishes", "game"),
        ("developer", "develops", "game"),
    ]
    etypes = ["follows", "plays", "wishes", "develops"]
381
382

    # node & edge types
Minjie Wang's avatar
Minjie Wang committed
383
384
385
    assert set(ntypes) == set(g.ntypes)
    assert set(etypes) == set(g.etypes)
    assert set(canonical_etypes) == set(g.canonical_etypes)
386
387

    # metagraph
388
    mg = g.metagraph()
Minjie Wang's avatar
Minjie Wang committed
389
    assert set(g.ntypes) == set(mg.nodes)
390
    etype_triplets = [(u, v, e) for u, v, e in mg.edges(keys=True)]
391
392
393
394
395
396
397
398
    assert set(
        [
            ("user", "user", "follows"),
            ("user", "game", "plays"),
            ("user", "game", "wishes"),
            ("developer", "game", "develops"),
        ]
    ) == set(etype_triplets)
Minjie Wang's avatar
Minjie Wang committed
399
400
    for i in range(len(etypes)):
        assert g.to_canonical_etype(etypes[i]) == canonical_etypes[i]
401

402
403
    def _test(g):
        # number of nodes
404
        assert [g.num_nodes(ntype) for ntype in ntypes] == [3, 2, 2]
405

406
        # number of edges
407
        assert [g.num_edges(etype) for etype in etypes] == [2, 4, 2, 2]
408

409
        # has_nodes
410
        for ntype in ntypes:
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
411
            n = g.num_nodes(ntype)
412
            for i in range(n):
413
414
                assert g.has_nodes(i, ntype)
            assert not g.has_nodes(n, ntype)
415
            assert np.array_equal(
416
417
                F.asnumpy(g.has_nodes([0, n], ntype)).astype("int32"), [1, 0]
            )
Minjie Wang's avatar
Minjie Wang committed
418

419
        assert not g.is_multigraph
Minjie Wang's avatar
Minjie Wang committed
420
421
422
423

        for etype in etypes:
            srcs, dsts = edges[etype]
            for src, dst in zip(srcs, dsts):
424
                assert g.has_edges_between(src, dst, etype)
Minjie Wang's avatar
Minjie Wang committed
425
426
427
428
            assert F.asnumpy(g.has_edges_between(srcs, dsts, etype)).all()

            srcs, dsts = negative_edges[etype]
            for src, dst in zip(srcs, dsts):
429
                assert not g.has_edges_between(src, dst, etype)
Minjie Wang's avatar
Minjie Wang committed
430
431
432
433
434
435
436
            assert not F.asnumpy(g.has_edges_between(srcs, dsts, etype)).any()

            srcs, dsts = edges[etype]
            n_edges = len(srcs)

            # predecessors & in_edges & in_degree
            pred = [s for s, d in zip(srcs, dsts) if d == 0]
437
438
439
            assert set(F.asnumpy(g.predecessors(0, etype)).tolist()) == set(
                pred
            )
Minjie Wang's avatar
Minjie Wang committed
440
441
442
            u, v = g.in_edges([0], etype=etype)
            assert F.asnumpy(v).tolist() == [0] * len(pred)
            assert set(F.asnumpy(u).tolist()) == set(pred)
443
            assert g.in_degrees(0, etype) == len(pred)
Minjie Wang's avatar
Minjie Wang committed
444
445
446
447
448
449
450

            # successors & out_edges & out_degree
            succ = [d for s, d in zip(srcs, dsts) if s == 0]
            assert set(F.asnumpy(g.successors(0, etype)).tolist()) == set(succ)
            u, v = g.out_edges([0], etype=etype)
            assert F.asnumpy(u).tolist() == [0] * len(succ)
            assert set(F.asnumpy(v).tolist()) == set(succ)
451
            assert g.out_degrees(0, etype) == len(succ)
Minjie Wang's avatar
Minjie Wang committed
452

453
            # edge_ids
Minjie Wang's avatar
Minjie Wang committed
454
            for i, (src, dst) in enumerate(zip(srcs, dsts)):
455
456
457
                assert g.edge_ids(src, dst, etype=etype) == i
                _, _, eid = g.edge_ids(src, dst, etype=etype, return_uv=True)
                assert eid == i
458
459
460
            assert F.asnumpy(
                g.edge_ids(srcs, dsts, etype=etype)
            ).tolist() == list(range(n_edges))
461
            u, v, e = g.edge_ids(srcs, dsts, etype=etype, return_uv=True)
462
463
464
            u, v, e = F.asnumpy(u), F.asnumpy(v), F.asnumpy(e)
            assert u[e].tolist() == srcs
            assert v[e].tolist() == dsts
465

Minjie Wang's avatar
Minjie Wang committed
466
            # find_edges
467
468
469
470
471
            for eid in [
                list(range(n_edges)),
                np.arange(n_edges),
                F.astype(F.arange(0, n_edges), g.idtype),
            ]:
472
                u, v = g.find_edges(eid, etype)
473
474
                assert F.asnumpy(u).tolist() == srcs
                assert F.asnumpy(v).tolist() == dsts
Minjie Wang's avatar
Minjie Wang committed
475
476

            # all_edges.
477
478
            for order in ["eid"]:
                u, v, e = g.edges("all", order, etype)
Minjie Wang's avatar
Minjie Wang committed
479
480
481
482
483
484
485
486
487
488
                assert F.asnumpy(u).tolist() == srcs
                assert F.asnumpy(v).tolist() == dsts
                assert F.asnumpy(e).tolist() == list(range(n_edges))

            # in_degrees & out_degrees
            in_degrees = F.asnumpy(g.in_degrees(etype=etype))
            out_degrees = F.asnumpy(g.out_degrees(etype=etype))
            src_count = Counter(srcs)
            dst_count = Counter(dsts)
            utype, _, vtype = g.to_canonical_etype(etype)
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
489
            for i in range(g.num_nodes(utype)):
Minjie Wang's avatar
Minjie Wang committed
490
                assert out_degrees[i] == src_count[i]
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
491
            for i in range(g.num_nodes(vtype)):
Minjie Wang's avatar
Minjie Wang committed
492
493
494
                assert in_degrees[i] == dst_count[i]

    edges = {
495
496
497
498
        "follows": ([0, 1], [1, 2]),
        "plays": ([0, 1, 2, 1], [0, 0, 1, 1]),
        "wishes": ([0, 2], [1, 0]),
        "develops": ([0, 1], [0, 1]),
Minjie Wang's avatar
Minjie Wang committed
499
500
501
    }
    # edges that does not exist in the graph
    negative_edges = {
502
503
504
505
        "follows": ([0, 1], [0, 1]),
        "plays": ([0, 2], [1, 0]),
        "wishes": ([0, 1], [0, 1]),
        "develops": ([0, 1], [1, 0]),
Minjie Wang's avatar
Minjie Wang committed
506
    }
507
    g = create_test_heterograph(idtype)
Minjie Wang's avatar
Minjie Wang committed
508
    _test(g)
509
    g = create_test_heterograph1(idtype)
510
    _test(g)
511
    if F._default_context_str != "gpu":
512
        # XXX: CUDA COO operators have not been live yet.
513
        g = create_test_heterograph2(idtype)
514
        _test(g)
Minjie Wang's avatar
Minjie Wang committed
515
516
517

    etypes = canonical_etypes
    edges = {
518
519
520
521
        ("user", "follows", "user"): ([0, 1], [1, 2]),
        ("user", "plays", "game"): ([0, 1, 2, 1], [0, 0, 1, 1]),
        ("user", "wishes", "game"): ([0, 2], [1, 0]),
        ("developer", "develops", "game"): ([0, 1], [0, 1]),
Minjie Wang's avatar
Minjie Wang committed
522
523
524
    }
    # edges that does not exist in the graph
    negative_edges = {
525
526
527
528
529
        ("user", "follows", "user"): ([0, 1], [0, 1]),
        ("user", "plays", "game"): ([0, 2], [1, 0]),
        ("user", "wishes", "game"): ([0, 1], [0, 1]),
        ("developer", "develops", "game"): ([0, 1], [1, 0]),
    }
530
    g = create_test_heterograph(idtype)
Minjie Wang's avatar
Minjie Wang committed
531
    _test(g)
532
    g = create_test_heterograph1(idtype)
533
    _test(g)
534
    if F._default_context_str != "gpu":
535
        # XXX: CUDA COO operators have not been live yet.
536
        g = create_test_heterograph2(idtype)
537
        _test(g)
Minjie Wang's avatar
Minjie Wang committed
538
539
540
541

    # test repr
    print(g)

542

nv-dlasalle's avatar
nv-dlasalle committed
543
@parametrize_idtype
544
545
546
547
548
549
550
551
552
553
554
555
def test_empty_query(idtype):
    g = dgl.graph(([1, 2, 3], [0, 4, 5]), idtype=idtype, device=F.ctx())
    g.add_nodes(0)
    g.add_edges([], [])
    g.remove_edges([])
    g.remove_nodes([])
    assert F.shape(g.has_nodes([])) == (0,)
    assert F.shape(g.has_edges_between([], [])) == (0,)
    g.edge_ids([], [])
    g.edge_ids([], [], return_uv=True)
    g.find_edges([])

556
557
    assert F.shape(g.in_edges([], form="eid")) == (0,)
    u, v = g.in_edges([], form="uv")
558
559
    assert F.shape(u) == (0,)
    assert F.shape(v) == (0,)
560
    u, v, e = g.in_edges([], form="all")
561
562
563
564
    assert F.shape(u) == (0,)
    assert F.shape(v) == (0,)
    assert F.shape(e) == (0,)

565
566
    assert F.shape(g.out_edges([], form="eid")) == (0,)
    u, v = g.out_edges([], form="uv")
567
568
    assert F.shape(u) == (0,)
    assert F.shape(v) == (0,)
569
    u, v, e = g.out_edges([], form="all")
570
571
572
573
574
575
576
    assert F.shape(u) == (0,)
    assert F.shape(v) == (0,)
    assert F.shape(e) == (0,)

    assert F.shape(g.in_degrees([])) == (0,)
    assert F.shape(g.out_degrees([])) == (0,)

Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
    g = dgl.graph(([], []), idtype=idtype, device=F.ctx())
    error_thrown = True
    try:
        g.in_degrees([0])
        fail = False
    except:
        pass
    assert error_thrown
    error_thrown = True
    try:
        g.out_degrees([0])
        fail = False
    except:
        pass
    assert error_thrown

593
594
595
596

@unittest.skipIf(
    F._default_context_str == "gpu", reason="GPU does not have COO impl."
)
597
def _test_hypersparse():
598
    N1 = 1 << 50  # should crash if allocated a CSR
599
600
    N2 = 1 << 48

601
602
603
604
605
606
607
608
609
610
611
612
613
614
    g = dgl.heterograph(
        {
            ("user", "follows", "user"): (
                F.tensor([0], F.int64),
                F.tensor([1], F.int64),
            ),
            ("user", "plays", "game"): (
                F.tensor([0], F.int64),
                F.tensor([N2], F.int64),
            ),
        },
        {"user": N1, "game": N1},
        device=F.ctx(),
    )
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
615
616
617
618
    assert g.num_nodes("user") == N1
    assert g.num_nodes("game") == N1
    assert g.num_edges("follows") == 1
    assert g.num_edges("plays") == 1
619
620
621
622

    assert g.has_edges_between(0, 1, "follows")
    assert not g.has_edges_between(0, 0, "follows")
    mask = F.asnumpy(g.has_edges_between([0, 0], [0, 1], "follows")).tolist()
623
624
    assert mask == [0, 1]

625
626
627
    assert g.has_edges_between(0, N2, "plays")
    assert not g.has_edges_between(0, 0, "plays")
    mask = F.asnumpy(g.has_edges_between([0, 0], [0, N2], "plays")).tolist()
628
629
    assert mask == [0, 1]

630
631
632
633
    assert F.asnumpy(g.predecessors(0, "follows")).tolist() == []
    assert F.asnumpy(g.successors(0, "follows")).tolist() == [1]
    assert F.asnumpy(g.predecessors(1, "follows")).tolist() == [0]
    assert F.asnumpy(g.successors(1, "follows")).tolist() == []
634

635
636
637
638
    assert F.asnumpy(g.predecessors(0, "plays")).tolist() == []
    assert F.asnumpy(g.successors(0, "plays")).tolist() == [N2]
    assert F.asnumpy(g.predecessors(N2, "plays")).tolist() == [0]
    assert F.asnumpy(g.successors(N2, "plays")).tolist() == []
639

640
641
    assert g.edge_ids(0, 1, etype="follows") == 0
    assert g.edge_ids(0, N2, etype="plays") == 0
642

643
    u, v = g.find_edges([0], "follows")
644
645
    assert F.asnumpy(u).tolist() == [0]
    assert F.asnumpy(v).tolist() == [1]
646
    u, v = g.find_edges([0], "plays")
647
648
    assert F.asnumpy(u).tolist() == [0]
    assert F.asnumpy(v).tolist() == [N2]
649
    u, v, e = g.all_edges("all", "eid", "follows")
650
651
652
    assert F.asnumpy(u).tolist() == [0]
    assert F.asnumpy(v).tolist() == [1]
    assert F.asnumpy(e).tolist() == [0]
653
    u, v, e = g.all_edges("all", "eid", "plays")
654
655
656
657
    assert F.asnumpy(u).tolist() == [0]
    assert F.asnumpy(v).tolist() == [N2]
    assert F.asnumpy(e).tolist() == [0]

658
659
660
661
662
663
664
665
666
667
668
669
670
    assert g.in_degrees(0, "follows") == 0
    assert g.in_degrees(1, "follows") == 1
    assert F.asnumpy(g.in_degrees([0, 1], "follows")).tolist() == [0, 1]
    assert g.in_degrees(0, "plays") == 0
    assert g.in_degrees(N2, "plays") == 1
    assert F.asnumpy(g.in_degrees([0, N2], "plays")).tolist() == [0, 1]
    assert g.out_degrees(0, "follows") == 1
    assert g.out_degrees(1, "follows") == 0
    assert F.asnumpy(g.out_degrees([0, 1], "follows")).tolist() == [1, 0]
    assert g.out_degrees(0, "plays") == 1
    assert g.out_degrees(N2, "plays") == 0
    assert F.asnumpy(g.out_degrees([0, N2], "plays")).tolist() == [1, 0]

671

672
def _test_edge_ids():
673
    N1 = 1 << 50  # should crash if allocated a CSR
674
675
    N2 = 1 << 48

676
677
678
679
680
681
682
683
684
685
686
687
688
    g = dgl.heterograph(
        {
            ("user", "follows", "user"): (
                F.tensor([0], F.int64),
                F.tensor([1], F.int64),
            ),
            ("user", "plays", "game"): (
                F.tensor([0], F.int64),
                F.tensor([N2], F.int64),
            ),
        },
        {"user": N1, "game": N1},
    )
689
    with pytest.raises(DGLError):
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
        eid = g.edge_ids(0, 0, etype="follows")

    g2 = dgl.heterograph(
        {
            ("user", "follows", "user"): (
                F.tensor([0, 0], F.int64),
                F.tensor([1, 1], F.int64),
            ),
            ("user", "plays", "game"): (
                F.tensor([0], F.int64),
                F.tensor([N2], F.int64),
            ),
        },
        {"user": N1, "game": N1},
        device=F.cpu(),
    )

    eid = g2.edge_ids(0, 1, etype="follows")
708
    assert eid == 0
709

710

711
712
713
@pytest.mark.skipif(
    F.backend_name != "pytorch", reason="Only support PyTorch for now"
)
nv-dlasalle's avatar
nv-dlasalle committed
714
@parametrize_idtype
715
716
def test_adj(idtype):
    g = create_test_heterograph(idtype)
717
718
719
720
721
722
723
724
725
726
727
728
729
730
    adj = g.adj("follows")
    assert F.asnumpy(adj.indices()).tolist() == [[0, 1], [1, 2]]
    assert np.allclose(F.asnumpy(adj.val), np.array([1, 1]))
    g.edata["h"] = {("user", "plays", "game"): F.tensor([1, 2, 3, 4])}
    print(g.edata["h"])
    adj = g.adj("plays", "h")
    assert F.asnumpy(adj.indices()).tolist() == [[0, 1, 2, 1], [0, 0, 1, 1]]
    assert np.allclose(F.asnumpy(adj.val), np.array([1, 2, 3, 4]))


@parametrize_idtype
def test_adj_external(idtype):
    g = create_test_heterograph(idtype)
    adj = F.sparse_to_numpy(g.adj_external(transpose=True, etype="follows"))
Minjie Wang's avatar
Minjie Wang committed
731
    assert np.allclose(
732
733
        adj, np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])
    )
734
    adj = F.sparse_to_numpy(g.adj_external(transpose=False, etype="follows"))
Minjie Wang's avatar
Minjie Wang committed
735
    assert np.allclose(
736
737
        adj, np.array([[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 0.0]])
    )
738
    adj = F.sparse_to_numpy(g.adj_external(transpose=True, etype="plays"))
739
    assert np.allclose(adj, np.array([[1.0, 1.0, 0.0], [0.0, 1.0, 1.0]]))
740
    adj = F.sparse_to_numpy(g.adj_external(transpose=False, etype="plays"))
741
742
    assert np.allclose(adj, np.array([[1.0, 0.0], [1.0, 1.0], [0.0, 1.0]]))

743
    adj = g.adj_external(transpose=True, scipy_fmt="csr", etype="follows")
Minjie Wang's avatar
Minjie Wang committed
744
    assert np.allclose(
745
746
747
        adj.todense(),
        np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]),
    )
748
    adj = g.adj_external(transpose=True, scipy_fmt="coo", etype="follows")
Minjie Wang's avatar
Minjie Wang committed
749
    assert np.allclose(
750
751
752
        adj.todense(),
        np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]),
    )
753
    adj = g.adj_external(transpose=True, scipy_fmt="csr", etype="plays")
Minjie Wang's avatar
Minjie Wang committed
754
    assert np.allclose(
755
756
        adj.todense(), np.array([[1.0, 1.0, 0.0], [0.0, 1.0, 1.0]])
    )
757
    adj = g.adj_external(transpose=True, scipy_fmt="coo", etype="plays")
Minjie Wang's avatar
Minjie Wang committed
758
    assert np.allclose(
759
760
        adj.todense(), np.array([[1.0, 1.0, 0.0], [0.0, 1.0, 1.0]])
    )
761
    adj = F.sparse_to_numpy(g["follows"].adj_external(transpose=True))
Minjie Wang's avatar
Minjie Wang committed
762
    assert np.allclose(
763
764
765
        adj, np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])
    )

Minjie Wang's avatar
Minjie Wang committed
766

nv-dlasalle's avatar
nv-dlasalle committed
767
@parametrize_idtype
768
769
def test_inc(idtype):
    g = create_test_heterograph(idtype)
770
771
772
773
774
775
776
    adj = F.sparse_to_numpy(g["follows"].inc("in"))
    assert np.allclose(adj, np.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]]))
    adj = F.sparse_to_numpy(g["follows"].inc("out"))
    assert np.allclose(adj, np.array([[1.0, 0.0], [0.0, 1.0], [0.0, 0.0]]))
    adj = F.sparse_to_numpy(g["follows"].inc("both"))
    assert np.allclose(adj, np.array([[-1.0, 0.0], [1.0, -1.0], [0.0, 1.0]]))
    adj = F.sparse_to_numpy(g.inc("in", etype="plays"))
Minjie Wang's avatar
Minjie Wang committed
777
    assert np.allclose(
778
779
780
        adj, np.array([[1.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 1.0]])
    )
    adj = F.sparse_to_numpy(g.inc("out", etype="plays"))
Minjie Wang's avatar
Minjie Wang committed
781
    assert np.allclose(
782
783
784
785
786
787
788
789
        adj,
        np.array(
            [[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 1.0], [0.0, 0.0, 1.0, 0.0]]
        ),
    )
    adj = F.sparse_to_numpy(g.inc("both", etype="follows"))
    assert np.allclose(adj, np.array([[-1.0, 0.0], [1.0, -1.0], [0.0, 1.0]]))

790

nv-dlasalle's avatar
nv-dlasalle committed
791
@parametrize_idtype
792
def test_view(idtype):
793
    # test single node type
794
795
796
797
798
    g = dgl.heterograph(
        {("user", "follows", "user"): ([0, 1], [1, 2])},
        idtype=idtype,
        device=F.ctx(),
    )
799
    f1 = F.randn((3, 6))
800
801
    g.ndata["h"] = f1
    f2 = g.nodes["user"].data["h"]
802
803
804
    assert F.array_equal(f1, f2)
    fail = False
    try:
805
        g.ndata["h"] = {"user": f1}
806
807
808
809
810
811
    except Exception:
        fail = True
    assert fail

    # test single edge type
    f3 = F.randn((2, 4))
812
813
    g.edata["h"] = f3
    f4 = g.edges["follows"].data["h"]
814
815
816
    assert F.array_equal(f3, f4)
    fail = False
    try:
817
        g.edata["h"] = {"follows": f3}
818
819
820
821
    except Exception:
        fail = True
    assert fail

Minjie Wang's avatar
Minjie Wang committed
822
    # test data view
823
    g = create_test_heterograph(idtype)
824
825

    f1 = F.randn((3, 6))
826
827
    g.nodes["user"].data["h"] = f1  # ok
    f2 = g.nodes["user"].data["h"]
828
    assert F.array_equal(f1, f2)
829
830
    assert F.array_equal(g.nodes("user"), F.arange(0, 3, idtype))
    g.nodes["user"].data.pop("h")
831
832
833
834
835
836

    # multi type ndata
    f1 = F.randn((3, 6))
    f2 = F.randn((2, 6))
    fail = False
    try:
837
        g.ndata["h"] = f1
838
839
840
    except Exception:
        fail = True
    assert fail
841
842

    f3 = F.randn((2, 4))
843
844
845
    g.edges["user", "follows", "user"].data["h"] = f3
    f4 = g.edges["user", "follows", "user"].data["h"]
    f5 = g.edges["follows"].data["h"]
846
    assert F.array_equal(f3, f4)
Minjie Wang's avatar
Minjie Wang committed
847
    assert F.array_equal(f3, f5)
848
849
850
851
    assert F.array_equal(
        g.edges(etype="follows", form="eid"), F.arange(0, 2, idtype)
    )
    g.edges["follows"].data.pop("h")
852
853
854
855

    f3 = F.randn((2, 4))
    fail = False
    try:
856
        g.edata["h"] = f3
857
858
859
860
861
862
    except Exception:
        fail = True
    assert fail

    # test srcdata
    f1 = F.randn((3, 6))
863
864
    g.srcnodes["user"].data["h"] = f1  # ok
    f2 = g.srcnodes["user"].data["h"]
865
    assert F.array_equal(f1, f2)
866
867
    assert F.array_equal(g.srcnodes("user"), F.arange(0, 3, idtype))
    g.srcnodes["user"].data.pop("h")
868
869
870

    # test dstdata
    f1 = F.randn((3, 6))
871
872
    g.dstnodes["user"].data["h"] = f1  # ok
    f2 = g.dstnodes["user"].data["h"]
873
    assert F.array_equal(f1, f2)
874
875
876
    assert F.array_equal(g.dstnodes("user"), F.arange(0, 3, idtype))
    g.dstnodes["user"].data.pop("h")

877

nv-dlasalle's avatar
nv-dlasalle committed
878
@parametrize_idtype
879
def test_view1(idtype):
Minjie Wang's avatar
Minjie Wang committed
880
    # test relation view
881
    HG = create_test_heterograph(idtype)
882
    ntypes = ["user", "game", "developer"]
Minjie Wang's avatar
Minjie Wang committed
883
    canonical_etypes = [
884
885
886
887
888
889
        ("user", "follows", "user"),
        ("user", "plays", "game"),
        ("user", "wishes", "game"),
        ("developer", "develops", "game"),
    ]
    etypes = ["follows", "plays", "wishes", "develops"]
Minjie Wang's avatar
Minjie Wang committed
890
891
892
893
894
895
896

    def _test_query():
        for etype in etypes:
            utype, _, vtype = HG.to_canonical_etype(etype)
            g = HG[etype]
            srcs, dsts = edges[etype]
            for src, dst in zip(srcs, dsts):
897
                assert g.has_edges_between(src, dst)
Minjie Wang's avatar
Minjie Wang committed
898
899
900
901
            assert F.asnumpy(g.has_edges_between(srcs, dsts)).all()

            srcs, dsts = negative_edges[etype]
            for src, dst in zip(srcs, dsts):
902
                assert not g.has_edges_between(src, dst)
Minjie Wang's avatar
Minjie Wang committed
903
904
905
906
907
908
909
910
911
912
913
            assert not F.asnumpy(g.has_edges_between(srcs, dsts)).any()

            srcs, dsts = edges[etype]
            n_edges = len(srcs)

            # predecessors & in_edges & in_degree
            pred = [s for s, d in zip(srcs, dsts) if d == 0]
            assert set(F.asnumpy(g.predecessors(0)).tolist()) == set(pred)
            u, v = g.in_edges([0])
            assert F.asnumpy(v).tolist() == [0] * len(pred)
            assert set(F.asnumpy(u).tolist()) == set(pred)
914
            assert g.in_degrees(0) == len(pred)
Minjie Wang's avatar
Minjie Wang committed
915
916
917
918
919
920
921

            # successors & out_edges & out_degree
            succ = [d for s, d in zip(srcs, dsts) if s == 0]
            assert set(F.asnumpy(g.successors(0)).tolist()) == set(succ)
            u, v = g.out_edges([0])
            assert F.asnumpy(u).tolist() == [0] * len(succ)
            assert set(F.asnumpy(v).tolist()) == set(succ)
922
            assert g.out_degrees(0) == len(succ)
Minjie Wang's avatar
Minjie Wang committed
923

924
            # edge_ids
Minjie Wang's avatar
Minjie Wang committed
925
            for i, (src, dst) in enumerate(zip(srcs, dsts)):
926
927
928
                assert g.edge_ids(src, dst, etype=etype) == i
                _, _, eid = g.edge_ids(src, dst, etype=etype, return_uv=True)
                assert eid == i
929
930
931
            assert F.asnumpy(g.edge_ids(srcs, dsts)).tolist() == list(
                range(n_edges)
            )
932
            u, v, e = g.edge_ids(srcs, dsts, return_uv=True)
933
934
935
            u, v, e = F.asnumpy(u), F.asnumpy(v), F.asnumpy(e)
            assert u[e].tolist() == srcs
            assert v[e].tolist() == dsts
Minjie Wang's avatar
Minjie Wang committed
936
937
938
939
940
941
942

            # find_edges
            u, v = g.find_edges(list(range(n_edges)))
            assert F.asnumpy(u).tolist() == srcs
            assert F.asnumpy(v).tolist() == dsts

            # all_edges.
943
944
            for order in ["eid"]:
                u, v, e = g.all_edges(form="all", order=order)
Minjie Wang's avatar
Minjie Wang committed
945
946
947
948
949
950
951
952
953
                assert F.asnumpy(u).tolist() == srcs
                assert F.asnumpy(v).tolist() == dsts
                assert F.asnumpy(e).tolist() == list(range(n_edges))

            # in_degrees & out_degrees
            in_degrees = F.asnumpy(g.in_degrees())
            out_degrees = F.asnumpy(g.out_degrees())
            src_count = Counter(srcs)
            dst_count = Counter(dsts)
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
954
            for i in range(g.num_nodes(utype)):
Minjie Wang's avatar
Minjie Wang committed
955
                assert out_degrees[i] == src_count[i]
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
956
            for i in range(g.num_nodes(vtype)):
957
                assert in_degrees[i] == dst_count[i]
Minjie Wang's avatar
Minjie Wang committed
958
959

    edges = {
960
961
962
963
        "follows": ([0, 1], [1, 2]),
        "plays": ([0, 1, 2, 1], [0, 0, 1, 1]),
        "wishes": ([0, 2], [1, 0]),
        "develops": ([0, 1], [0, 1]),
Minjie Wang's avatar
Minjie Wang committed
964
965
966
    }
    # edges that does not exist in the graph
    negative_edges = {
967
968
969
970
        "follows": ([0, 1], [0, 1]),
        "plays": ([0, 2], [1, 0]),
        "wishes": ([0, 1], [0, 1]),
        "develops": ([0, 1], [1, 0]),
Minjie Wang's avatar
Minjie Wang committed
971
972
973
974
    }
    _test_query()
    etypes = canonical_etypes
    edges = {
975
976
977
978
        ("user", "follows", "user"): ([0, 1], [1, 2]),
        ("user", "plays", "game"): ([0, 1, 2, 1], [0, 0, 1, 1]),
        ("user", "wishes", "game"): ([0, 2], [1, 0]),
        ("developer", "develops", "game"): ([0, 1], [0, 1]),
Minjie Wang's avatar
Minjie Wang committed
979
980
981
    }
    # edges that does not exist in the graph
    negative_edges = {
982
983
984
985
986
        ("user", "follows", "user"): ([0, 1], [0, 1]),
        ("user", "plays", "game"): ([0, 2], [1, 0]),
        ("user", "wishes", "game"): ([0, 1], [0, 1]),
        ("developer", "develops", "game"): ([0, 1], [1, 0]),
    }
Minjie Wang's avatar
Minjie Wang committed
987
988
989
    _test_query()

    # test features
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
990
991
    HG.nodes["user"].data["h"] = F.ones((HG.num_nodes("user"), 5))
    HG.nodes["game"].data["m"] = F.ones((HG.num_nodes("game"), 3)) * 2
Minjie Wang's avatar
Minjie Wang committed
992
993

    # test only one node type
994
    g = HG["follows"]
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
995
    assert g.num_nodes() == 3
Minjie Wang's avatar
Minjie Wang committed
996
997
998

    # test ndata and edata
    f1 = F.randn((3, 6))
999
1000
    g.ndata["h"] = f1  # ok
    f2 = HG.nodes["user"].data["h"]
Minjie Wang's avatar
Minjie Wang committed
1001
    assert F.array_equal(f1, f2)
1002
    assert F.array_equal(g.nodes(), F.arange(0, 3, g.idtype))
Minjie Wang's avatar
Minjie Wang committed
1003
1004

    f3 = F.randn((2, 4))
1005
1006
    g.edata["h"] = f3
    f4 = HG.edges["follows"].data["h"]
Minjie Wang's avatar
Minjie Wang committed
1007
    assert F.array_equal(f3, f4)
1008
1009
    assert F.array_equal(g.edges(form="eid"), F.arange(0, 2, g.idtype))

Minjie Wang's avatar
Minjie Wang committed
1010

nv-dlasalle's avatar
nv-dlasalle committed
1011
@parametrize_idtype
1012
def test_flatten(idtype):
Minjie Wang's avatar
Minjie Wang committed
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
    def check_mapping(g, fg):
        if len(fg.ntypes) == 1:
            SRC = DST = fg.ntypes[0]
        else:
            SRC = fg.ntypes[0]
            DST = fg.ntypes[1]

        etypes = F.asnumpy(fg.edata[dgl.ETYPE]).tolist()
        eids = F.asnumpy(fg.edata[dgl.EID]).tolist()

        for i, (etype, eid) in enumerate(zip(etypes, eids)):
            src_g, dst_g = g.find_edges([eid], g.canonical_etypes[etype])
            src_fg, dst_fg = fg.find_edges([i])
            # TODO(gq): I feel this code is quite redundant; can we just add new members (like
            # "induced_srcid") to returned heterograph object and not store them as features?
1028
1029
1030
1031
1032
1033
            assert F.asnumpy(src_g) == F.asnumpy(
                F.gather_row(fg.nodes[SRC].data[dgl.NID], src_fg)[0]
            )
            tid = F.asnumpy(
                F.gather_row(fg.nodes[SRC].data[dgl.NTYPE], src_fg)
            ).item()
Minjie Wang's avatar
Minjie Wang committed
1034
            assert g.canonical_etypes[etype][0] == g.ntypes[tid]
1035
1036
1037
1038
1039
1040
            assert F.asnumpy(dst_g) == F.asnumpy(
                F.gather_row(fg.nodes[DST].data[dgl.NID], dst_fg)[0]
            )
            tid = F.asnumpy(
                F.gather_row(fg.nodes[DST].data[dgl.NTYPE], dst_fg)
            ).item()
Minjie Wang's avatar
Minjie Wang committed
1041
1042
1043
            assert g.canonical_etypes[etype][2] == g.ntypes[tid]

    # check for wildcard slices
1044
    g = create_test_heterograph(idtype)
1045
1046
1047
1048
1049
    g.nodes["user"].data["h"] = F.ones((3, 5))
    g.nodes["game"].data["i"] = F.ones((2, 5))
    g.edges["plays"].data["e"] = F.ones((4, 4))
    g.edges["wishes"].data["e"] = F.ones((2, 4))
    g.edges["wishes"].data["f"] = F.ones((2, 4))
Minjie Wang's avatar
Minjie Wang committed
1050

1051
    fg = g["user", :, "game"]  # user--plays->game and user--wishes->game
Minjie Wang's avatar
Minjie Wang committed
1052
    assert len(fg.ntypes) == 2
1053
1054
    assert fg.ntypes == ["user", "game"]
    assert fg.etypes == ["plays+wishes"]
1055
1056
    assert fg.idtype == g.idtype
    assert fg.device == g.device
1057
    etype = fg.etypes[0]
1058
    assert fg[etype] is not None  # Issue #2166
Minjie Wang's avatar
Minjie Wang committed
1059

1060
1061
1062
1063
    assert F.array_equal(fg.nodes["user"].data["h"], F.ones((3, 5)))
    assert F.array_equal(fg.nodes["game"].data["i"], F.ones((2, 5)))
    assert F.array_equal(fg.edata["e"], F.ones((6, 4)))
    assert "f" not in fg.edata
Minjie Wang's avatar
Minjie Wang committed
1064
1065
1066

    etypes = F.asnumpy(fg.edata[dgl.ETYPE]).tolist()
    eids = F.asnumpy(fg.edata[dgl.EID]).tolist()
1067
1068
1069
    assert set(zip(etypes, eids)) == set(
        [(3, 0), (3, 1), (2, 1), (2, 0), (2, 3), (2, 2)]
    )
Minjie Wang's avatar
Minjie Wang committed
1070
1071
1072

    check_mapping(g, fg)

1073
    fg = g["user", :, "user"]
1074
1075
    assert fg.idtype == g.idtype
    assert fg.device == g.device
Minjie Wang's avatar
Minjie Wang committed
1076
1077
    # NOTE(gq): The node/edge types from the parent graph is returned if there is only one
    # node/edge type.  This differs from the behavior above.
1078
1079
1080
1081
    assert fg.ntypes == ["user"]
    assert fg.etypes == ["follows"]
    u1, v1 = g.edges(etype="follows", order="eid")
    u2, v2 = fg.edges(etype="follows", order="eid")
Minjie Wang's avatar
Minjie Wang committed
1082
1083
1084
    assert F.array_equal(u1, u2)
    assert F.array_equal(v1, v2)

1085
    fg = g["developer", :, "game"]
1086
1087
    assert fg.idtype == g.idtype
    assert fg.device == g.device
1088
1089
1090
1091
    assert fg.ntypes == ["developer", "game"]
    assert fg.etypes == ["develops"]
    u1, v1 = g.edges(etype="develops", order="eid")
    u2, v2 = fg.edges(etype="develops", order="eid")
Minjie Wang's avatar
Minjie Wang committed
1092
1093
1094
1095
    assert F.array_equal(u1, u2)
    assert F.array_equal(v1, v2)

    fg = g[:, :, :]
1096
1097
    assert fg.idtype == g.idtype
    assert fg.device == g.device
1098
1099
    assert fg.ntypes == ["developer+user", "game+user"]
    assert fg.etypes == ["develops+follows+plays+wishes"]
Minjie Wang's avatar
Minjie Wang committed
1100
1101
1102
    check_mapping(g, fg)

    # Test another heterograph
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
    g = dgl.heterograph(
        {
            ("user", "follows", "user"): ([0, 1, 2], [1, 2, 3]),
            ("user", "knows", "user"): ([0, 2], [2, 3]),
        },
        idtype=idtype,
        device=F.ctx(),
    )
    g.nodes["user"].data["h"] = F.randn((4, 3))
    g.edges["follows"].data["w"] = F.randn((3, 2))
    g.nodes["user"].data["hh"] = F.randn((4, 5))
    g.edges["knows"].data["ww"] = F.randn((2, 10))

    fg = g["user", :, "user"]
1117
1118
    assert fg.idtype == g.idtype
    assert fg.device == g.device
1119
1120
    assert fg.ntypes == ["user"]
    assert fg.etypes == ["follows+knows"]
Minjie Wang's avatar
Minjie Wang committed
1121
1122
    check_mapping(g, fg)

1123
    fg = g["user", :, :]
1124
1125
    assert fg.idtype == g.idtype
    assert fg.device == g.device
1126
1127
    assert fg.ntypes == ["user"]
    assert fg.etypes == ["follows+knows"]
Minjie Wang's avatar
Minjie Wang committed
1128
1129
    check_mapping(g, fg)

1130
1131
1132
1133

@unittest.skipIf(
    F._default_context_str == "cpu", reason="Need gpu for this test"
)
nv-dlasalle's avatar
nv-dlasalle committed
1134
@parametrize_idtype
1135
1136
1137
1138
def test_to_device(idtype):
    # TODO: rewrite this test case to accept different graphs so we
    #  can test reverse graph and batched graph
    g = create_test_heterograph(idtype)
1139
1140
1141
    g.nodes["user"].data["h"] = F.ones((3, 5))
    g.nodes["game"].data["i"] = F.ones((2, 5))
    g.edges["plays"].data["e"] = F.ones((4, 4))
1142
1143
1144
    assert g.device == F.ctx()
    g = g.to(F.cpu())
    assert g.device == F.cpu()
1145
1146
1147
    assert F.context(g.nodes["user"].data["h"]) == F.cpu()
    assert F.context(g.nodes["game"].data["i"]) == F.cpu()
    assert F.context(g.edges["plays"].data["e"]) == F.cpu()
1148
1149
1150
1151
1152
    for ntype in g.ntypes:
        assert F.context(g.batch_num_nodes(ntype)) == F.cpu()
    for etype in g.canonical_etypes:
        assert F.context(g.batch_num_edges(etype)) == F.cpu()

1153
    if F.is_cuda_available():
1154
        g1 = g.to(F.cuda())
1155
        assert g1.device == F.cuda()
1156
1157
1158
        assert F.context(g1.nodes["user"].data["h"]) == F.cuda()
        assert F.context(g1.nodes["game"].data["i"]) == F.cuda()
        assert F.context(g1.edges["plays"].data["e"]) == F.cuda()
1159
1160
1161
1162
        for ntype in g1.ntypes:
            assert F.context(g1.batch_num_nodes(ntype)) == F.cuda()
        for etype in g1.canonical_etypes:
            assert F.context(g1.batch_num_edges(etype)) == F.cuda()
1163
1164
1165
        assert F.context(g.nodes["user"].data["h"]) == F.cpu()
        assert F.context(g.nodes["game"].data["i"]) == F.cpu()
        assert F.context(g.edges["plays"].data["e"]) == F.cpu()
1166
1167
1168
1169
1170
        for ntype in g.ntypes:
            assert F.context(g.batch_num_nodes(ntype)) == F.cpu()
        for etype in g.canonical_etypes:
            assert F.context(g.batch_num_edges(etype)) == F.cpu()
        with pytest.raises(DGLError):
1171
            g1.nodes["user"].data["h"] = F.copy_to(F.ones((3, 5)), F.cpu())
1172
        with pytest.raises(DGLError):
1173
            g1.edges["plays"].data["e"] = F.copy_to(F.ones((4, 4)), F.cpu())
1174

1175
1176
1177
1178

@unittest.skipIf(
    F._default_context_str == "cpu", reason="Need gpu for this test"
)
nv-dlasalle's avatar
nv-dlasalle committed
1179
@parametrize_idtype
1180
@pytest.mark.parametrize("g", get_cases(["block"]))
1181
1182
1183
1184
def test_to_device2(g, idtype):
    g = g.astype(idtype)
    g = g.to(F.cpu())
    assert g.device == F.cpu()
1185
1186
    if F.is_cuda_available():
        g1 = g.to(F.cuda())
1187
1188
1189
1190
        assert g1.device == F.cuda()
        assert g1.ntypes == g.ntypes
        assert g1.etypes == g.etypes
        assert g1.canonical_etypes == g.canonical_etypes
1191

1192
1193
1194
1195
1196
1197
1198
1199

@unittest.skipIf(
    F._default_context_str == "cpu", reason="Need gpu for this test"
)
@unittest.skipIf(
    dgl.backend.backend_name != "pytorch",
    reason="Pinning graph inplace only supported for PyTorch",
)
nv-dlasalle's avatar
nv-dlasalle committed
1200
@parametrize_idtype
1201
1202
1203
1204
def test_pin_memory_(idtype):
    # TODO: rewrite this test case to accept different graphs so we
    #  can test reverse graph and batched graph
    g = create_test_heterograph(idtype)
1205
1206
1207
    g.nodes["user"].data["h"] = F.ones((3, 5))
    g.nodes["game"].data["i"] = F.ones((2, 5))
    g.edges["plays"].data["e"] = F.ones((4, 4))
1208
1209
1210
    g = g.to(F.cpu())
    assert not g.is_pinned()

1211
1212
1213
1214
    # unpin an unpinned CPU graph, directly return
    g.unpin_memory_()
    assert not g.is_pinned()
    assert g.device == F.cpu()
1215

1216
1217
1218
1219
    # pin a CPU graph
    g.pin_memory_()
    assert g.is_pinned()
    assert g.device == F.cpu()
1220
1221
1222
1223
1224
1225
    assert g.nodes["user"].data["h"].is_pinned()
    assert g.nodes["game"].data["i"].is_pinned()
    assert g.edges["plays"].data["e"].is_pinned()
    assert F.context(g.nodes["user"].data["h"]) == F.cpu()
    assert F.context(g.nodes["game"].data["i"]) == F.cpu()
    assert F.context(g.edges["plays"].data["e"]) == F.cpu()
1226
1227
1228
1229
    for ntype in g.ntypes:
        assert F.context(g.batch_num_nodes(ntype)) == F.cpu()
    for etype in g.canonical_etypes:
        assert F.context(g.batch_num_edges(etype)) == F.cpu()
1230

1231
1232
1233
    # it's fine to clone with new formats, but new graphs are not pinned
    # >>> g.formats()
    # {'created': ['coo'], 'not created': ['csr', 'csc']}
1234
1235
    assert not g.formats("csc").is_pinned()
    assert not g.formats("csr").is_pinned()
1236
    # 'coo' formats is already created and thus not cloned
1237
    assert g.formats("coo").is_pinned()
1238
1239
1240
1241
1242

    # pin a pinned graph, directly return
    g.pin_memory_()
    assert g.is_pinned()
    assert g.device == F.cpu()
1243

1244
1245
1246
1247
    # unpin a pinned graph
    g.unpin_memory_()
    assert not g.is_pinned()
    assert g.device == F.cpu()
1248

1249
    g1 = g.to(F.cuda())
1250

1251
1252
1253
1254
    # unpin an unpinned GPU graph, directly return
    g1.unpin_memory_()
    assert not g1.is_pinned()
    assert g1.device == F.cuda()
1255

1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
    # error pinning a GPU graph
    with pytest.raises(DGLError):
        g1.pin_memory_()

    # test pin empty homograph
    g2 = dgl.graph(([], []))
    g2.pin_memory_()
    assert g2.is_pinned()
    g2.unpin_memory_()
    assert not g2.is_pinned()

    # test pin heterograph with 0 edge of one relation type
1268
1269
1270
    g3 = dgl.heterograph(
        {("a", "b", "c"): ([0, 1], [1, 2]), ("c", "d", "c"): ([], [])}
    ).astype(idtype)
1271
1272
1273
1274
    g3.pin_memory_()
    assert g3.is_pinned()
    g3.unpin_memory_()
    assert not g3.is_pinned()
1275

1276

nv-dlasalle's avatar
nv-dlasalle committed
1277
@parametrize_idtype
1278
def test_convert_bound(idtype):
1279
    def _test_bipartite_bound(data, card):
1280
        with pytest.raises(DGLError):
1281
1282
1283
1284
1285
1286
            dgl.heterograph(
                {("_U", "_E", "_V"): data},
                {"_U": card[0], "_V": card[1]},
                idtype=idtype,
                device=F.ctx(),
            )
1287
1288

    def _test_graph_bound(data, card):
1289
1290
        with pytest.raises(DGLError):
            dgl.graph(data, num_nodes=card, idtype=idtype, device=F.ctx())
1291

1292
1293
1294
1295
    _test_bipartite_bound(([1, 2], [1, 2]), (2, 3))
    _test_bipartite_bound(([0, 1], [1, 4]), (2, 3))
    _test_graph_bound(([1, 3], [1, 2]), 3)
    _test_graph_bound(([0, 1], [1, 3]), 3)
1296
1297


nv-dlasalle's avatar
nv-dlasalle committed
1298
@parametrize_idtype
1299
1300
def test_convert(idtype):
    hg = create_test_heterograph(idtype)
Minjie Wang's avatar
Minjie Wang committed
1301
1302
    hs = []
    for ntype in hg.ntypes:
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
1303
        h = F.randn((hg.num_nodes(ntype), 5))
1304
        hg.nodes[ntype].data["h"] = h
Minjie Wang's avatar
Minjie Wang committed
1305
        hs.append(h)
1306
    hg.nodes["user"].data["x"] = F.randn((3, 3))
Minjie Wang's avatar
Minjie Wang committed
1307
1308
    ws = []
    for etype in hg.canonical_etypes:
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
1309
        w = F.randn((hg.num_edges(etype), 5))
1310
        hg.edges[etype].data["w"] = w
Minjie Wang's avatar
Minjie Wang committed
1311
        ws.append(w)
1312
    hg.edges["plays"].data["x"] = F.randn((4, 3))
Minjie Wang's avatar
Minjie Wang committed
1313

1314
    g = dgl.to_homogeneous(hg, ndata=["h"], edata=["w"])
1315
1316
    assert g.idtype == idtype
    assert g.device == hg.device
1317
1318
1319
1320
    assert F.array_equal(F.cat(hs, dim=0), g.ndata["h"])
    assert "x" not in g.ndata
    assert F.array_equal(F.cat(ws, dim=0), g.edata["w"])
    assert "x" not in g.edata
Minjie Wang's avatar
Minjie Wang committed
1321

1322
    src, dst = g.all_edges(order="eid")
Minjie Wang's avatar
Minjie Wang committed
1323
1324
1325
1326
    src = F.asnumpy(src)
    dst = F.asnumpy(dst)
    etype_id, eid = F.asnumpy(g.edata[dgl.ETYPE]), F.asnumpy(g.edata[dgl.EID])
    ntype_id, nid = F.asnumpy(g.ndata[dgl.NTYPE]), F.asnumpy(g.ndata[dgl.NID])
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
1327
    for i in range(g.num_edges()):
Minjie Wang's avatar
Minjie Wang committed
1328
1329
1330
1331
        srctype = hg.ntypes[ntype_id[src[i]]]
        dsttype = hg.ntypes[ntype_id[dst[i]]]
        etype = hg.etypes[etype_id[i]]
        src_i, dst_i = hg.find_edges([eid[i]], (srctype, etype, dsttype))
1332
1333
        assert F.asnumpy(src_i).item() == nid[src[i]]
        assert F.asnumpy(dst_i).item() == nid[dst[i]]
Minjie Wang's avatar
Minjie Wang committed
1334

1335
1336
1337
1338
1339
1340
1341
1342
    mg = nx.MultiDiGraph(
        [
            ("user", "user", "follows"),
            ("user", "game", "plays"),
            ("user", "game", "wishes"),
            ("developer", "game", "develops"),
        ]
    )
Minjie Wang's avatar
Minjie Wang committed
1343
1344

    for _mg in [None, mg]:
1345
        hg2 = dgl.to_heterogeneous(
1346
1347
1348
1349
1350
1351
1352
            g,
            hg.ntypes,
            hg.etypes,
            ntype_field=dgl.NTYPE,
            etype_field=dgl.ETYPE,
            metagraph=_mg,
        )
1353
1354
        assert hg2.idtype == hg.idtype
        assert hg2.device == hg.device
Minjie Wang's avatar
Minjie Wang committed
1355
1356
1357
        assert set(hg.ntypes) == set(hg2.ntypes)
        assert set(hg.canonical_etypes) == set(hg2.canonical_etypes)
        for ntype in hg.ntypes:
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
1358
            assert hg.num_nodes(ntype) == hg2.num_nodes(ntype)
1359
1360
1361
            assert F.array_equal(
                hg.nodes[ntype].data["h"], hg2.nodes[ntype].data["h"]
            )
Minjie Wang's avatar
Minjie Wang committed
1362
        for canonical_etype in hg.canonical_etypes:
1363
1364
            src, dst = hg.all_edges(etype=canonical_etype, order="eid")
            src2, dst2 = hg2.all_edges(etype=canonical_etype, order="eid")
Minjie Wang's avatar
Minjie Wang committed
1365
1366
            assert F.array_equal(src, src2)
            assert F.array_equal(dst, dst2)
1367
1368
1369
1370
            assert F.array_equal(
                hg.edges[canonical_etype].data["w"],
                hg2.edges[canonical_etype].data["w"],
            )
Minjie Wang's avatar
Minjie Wang committed
1371
1372

    # hetero_from_homo test case 2
1373
    g = dgl.graph(([0, 1, 2, 0], [2, 2, 3, 3]), idtype=idtype, device=F.ctx())
Minjie Wang's avatar
Minjie Wang committed
1374
1375
    g.ndata[dgl.NTYPE] = F.tensor([0, 0, 1, 2])
    g.edata[dgl.ETYPE] = F.tensor([0, 0, 1, 2])
1376
    hg = dgl.to_heterogeneous(g, ["l0", "l1", "l2"], ["e0", "e1", "e2"])
1377
1378
    assert hg.idtype == idtype
    assert hg.device == g.device
Minjie Wang's avatar
Minjie Wang committed
1379
    assert set(hg.canonical_etypes) == set(
1380
1381
        [("l0", "e0", "l1"), ("l1", "e1", "l2"), ("l0", "e2", "l2")]
    )
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
1382
1383
1384
1385
1386
1387
    assert hg.num_nodes("l0") == 2
    assert hg.num_nodes("l1") == 1
    assert hg.num_nodes("l2") == 1
    assert hg.num_edges("e0") == 2
    assert hg.num_edges("e1") == 1
    assert hg.num_edges("e2") == 1
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
    assert F.array_equal(hg.ndata[dgl.NID]["l0"], F.tensor([0, 1], F.int64))
    assert F.array_equal(hg.ndata[dgl.NID]["l1"], F.tensor([2], F.int64))
    assert F.array_equal(hg.ndata[dgl.NID]["l2"], F.tensor([3], F.int64))
    assert F.array_equal(
        hg.edata[dgl.EID][("l0", "e0", "l1")], F.tensor([0, 1], F.int64)
    )
    assert F.array_equal(
        hg.edata[dgl.EID][("l0", "e2", "l2")], F.tensor([3], F.int64)
    )
    assert F.array_equal(
        hg.edata[dgl.EID][("l1", "e1", "l2")], F.tensor([2], F.int64)
    )
Minjie Wang's avatar
Minjie Wang committed
1400
1401

    # hetero_from_homo test case 3
1402
1403
1404
    mg = nx.MultiDiGraph(
        [("user", "movie", "watches"), ("user", "TV", "watches")]
    )
1405
    g = dgl.graph(((0, 0), (1, 2)), idtype=idtype, device=F.ctx())
Minjie Wang's avatar
Minjie Wang committed
1406
1407
1408
    g.ndata[dgl.NTYPE] = F.tensor([0, 1, 2])
    g.edata[dgl.ETYPE] = F.tensor([0, 0])
    for _mg in [None, mg]:
1409
1410
1411
        hg = dgl.to_heterogeneous(
            g, ["user", "TV", "movie"], ["watches"], metagraph=_mg
        )
1412
1413
        assert hg.idtype == g.idtype
        assert hg.device == g.device
Minjie Wang's avatar
Minjie Wang committed
1414
        assert set(hg.canonical_etypes) == set(
1415
1416
            [("user", "watches", "movie"), ("user", "watches", "TV")]
        )
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
1417
1418
1419
1420
1421
        assert hg.num_nodes("user") == 1
        assert hg.num_nodes("TV") == 1
        assert hg.num_nodes("movie") == 1
        assert hg.num_edges(("user", "watches", "TV")) == 1
        assert hg.num_edges(("user", "watches", "movie")) == 1
Minjie Wang's avatar
Minjie Wang committed
1422
1423
        assert len(hg.etypes) == 2

1424
    # hetero_to_homo test case 2
1425
1426
1427
1428
1429
1430
    hg = dgl.heterograph(
        {("_U", "_E", "_V"): ([0, 1], [0, 1])},
        {"_U": 2, "_V": 3},
        idtype=idtype,
        device=F.ctx(),
    )
1431
    g = dgl.to_homogeneous(hg)
1432
1433
    assert hg.idtype == g.idtype
    assert hg.device == g.device
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
1434
    assert g.num_nodes() == 5
1435

1436
    # hetero_to_subgraph_to_homo
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
    hg = dgl.heterograph(
        {
            ("user", "plays", "game"): ([0, 1, 1, 2], [0, 0, 2, 1]),
            ("user", "follows", "user"): ([0, 1, 1], [1, 2, 2]),
        },
        idtype=idtype,
        device=F.ctx(),
    )
    hg.nodes["user"].data["h"] = F.copy_to(
        F.tensor([[1, 0], [0, 1], [1, 1]], dtype=idtype), ctx=F.ctx()
    )
    sg = dgl.node_subgraph(hg, {"user": [1, 2]})
1449
1450
    assert len(sg.ntypes) == 2
    assert len(sg.etypes) == 2
1451
1452
1453
1454
    assert sg.num_nodes("user") == 2
    assert sg.num_nodes("game") == 0
    g = dgl.to_homogeneous(sg, ndata=["h"])
    assert "h" in g.ndata.keys()
1455
1456
    assert g.num_nodes() == 2

1457
1458
1459
1460

@unittest.skipIf(
    F._default_context_str == "gpu", reason="Test on cpu is enough"
)
nv-dlasalle's avatar
nv-dlasalle committed
1461
@parametrize_idtype
1462
1463
def test_to_homo_zero_nodes(idtype):
    # Fix gihub issue #2870
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
    g = dgl.heterograph(
        {
            ("A", "AB", "B"): (
                np.random.randint(0, 200, (1000,)),
                np.random.randint(0, 200, (1000,)),
            ),
            ("B", "BA", "A"): (
                np.random.randint(0, 200, (1000,)),
                np.random.randint(0, 200, (1000,)),
            ),
        },
        num_nodes_dict={"A": 200, "B": 200, "C": 0},
        idtype=idtype,
    )
    g.nodes["A"].data["x"] = F.randn((200, 3))
    g.nodes["B"].data["x"] = F.randn((200, 3))
    gg = dgl.to_homogeneous(g, ["x"])
    assert "x" in gg.ndata

1483

nv-dlasalle's avatar
nv-dlasalle committed
1484
@parametrize_idtype
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
def test_to_homo2(idtype):
    # test the result homogeneous graph has nodes and edges sorted by their types
    hg = create_test_heterograph(idtype)
    g = dgl.to_homogeneous(hg)
    ntypes = F.asnumpy(g.ndata[dgl.NTYPE])
    etypes = F.asnumpy(g.edata[dgl.ETYPE])
    p = 0
    for tid, ntype in enumerate(hg.ntypes):
        num_nodes = hg.num_nodes(ntype)
        for i in range(p, p + num_nodes):
            assert ntypes[i] == tid
        p += num_nodes
    p = 0
    for tid, etype in enumerate(hg.canonical_etypes):
        num_edges = hg.num_edges(etype)
        for i in range(p, p + num_edges):
            assert etypes[i] == tid
        p += num_edges
    # test store_type=False
    g = dgl.to_homogeneous(hg, store_type=False)
    assert dgl.NTYPE not in g.ndata
    assert dgl.ETYPE not in g.edata
    # test return_count=True
    g, ntype_count, etype_count = dgl.to_homogeneous(hg, return_count=True)
    for i, count in enumerate(ntype_count):
        assert count == hg.num_nodes(hg.ntypes[i])
    for i, count in enumerate(etype_count):
        assert count == hg.num_edges(hg.canonical_etypes[i])

1514

nv-dlasalle's avatar
nv-dlasalle committed
1515
@parametrize_idtype
1516
1517
1518
1519
1520
1521
1522
def test_invertible_conversion(idtype):
    # Test whether to_homogeneous and to_heterogeneous are invertible
    hg = create_test_heterograph(idtype)
    g = dgl.to_homogeneous(hg)
    hg2 = dgl.to_heterogeneous(g, hg.ntypes, hg.etypes)
    assert_is_identical_hetero(hg, hg2, True)

1523

nv-dlasalle's avatar
nv-dlasalle committed
1524
@parametrize_idtype
1525
1526
def test_metagraph_reachable(idtype):
    g = create_test_heterograph(idtype)
Mufei Li's avatar
Mufei Li committed
1527
    x = F.randn((3, 5))
1528
    g.nodes["user"].data["h"] = x
Mufei Li's avatar
Mufei Li committed
1529

1530
    new_g = dgl.metapath_reachable_graph(g, ["follows", "plays"])
1531
    assert new_g.idtype == idtype
1532
    assert new_g.ntypes == ["game", "user"]
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
1533
    assert new_g.num_edges() == 3
Mufei Li's avatar
Mufei Li committed
1534
1535
    assert F.asnumpy(new_g.has_edges_between([0, 0, 1], [0, 1, 1])).all()

1536
    new_g = dgl.metapath_reachable_graph(g, ["follows"])
1537
    assert new_g.idtype == idtype
1538
    assert new_g.ntypes == ["user"]
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
1539
    assert new_g.num_edges() == 2
Mufei Li's avatar
Mufei Li committed
1540
1541
    assert F.asnumpy(new_g.has_edges_between([0, 1], [1, 2])).all()

1542
1543
1544
1545
1546

@unittest.skipIf(
    dgl.backend.backend_name == "mxnet",
    reason="MXNet doesn't support bool tensor",
)
nv-dlasalle's avatar
nv-dlasalle committed
1547
@parametrize_idtype
1548
1549
def test_subgraph_mask(idtype):
    g = create_test_heterograph(idtype)
1550
1551
    g_graph = g["follows"]
    g_bipartite = g["plays"]
1552
1553
1554

    x = F.randn((3, 5))
    y = F.randn((2, 4))
1555
1556
    g.nodes["user"].data["h"] = x
    g.edges["follows"].data["h"] = y
1557
1558

    def _check_subgraph(g, sg):
1559
1560
        assert sg.idtype == g.idtype
        assert sg.device == g.device
1561
1562
1563
        assert sg.ntypes == g.ntypes
        assert sg.etypes == g.etypes
        assert sg.canonical_etypes == g.canonical_etypes
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
        assert F.array_equal(
            F.tensor(sg.nodes["user"].data[dgl.NID]), F.tensor([1, 2], idtype)
        )
        assert F.array_equal(
            F.tensor(sg.nodes["game"].data[dgl.NID]), F.tensor([0], idtype)
        )
        assert F.array_equal(
            F.tensor(sg.edges["follows"].data[dgl.EID]), F.tensor([1], idtype)
        )
        assert F.array_equal(
            F.tensor(sg.edges["plays"].data[dgl.EID]), F.tensor([1], idtype)
        )
        assert F.array_equal(
            F.tensor(sg.edges["wishes"].data[dgl.EID]), F.tensor([1], idtype)
        )
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
1579
1580
        assert sg.num_nodes("developer") == 0
        assert sg.num_edges("develops") == 0
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
        assert F.array_equal(
            sg.nodes["user"].data["h"], g.nodes["user"].data["h"][1:3]
        )
        assert F.array_equal(
            sg.edges["follows"].data["h"], g.edges["follows"].data["h"][1:2]
        )

    sg1 = g.subgraph(
        {
            "user": F.tensor([False, True, True], dtype=F.bool),
            "game": F.tensor([True, False, False, False], dtype=F.bool),
        }
    )
1594
    _check_subgraph(g, sg1)
1595
    if F._default_context_str != "gpu":
1596
        # TODO(minjie): enable this later
1597
1598
1599
1600
1601
1602
1603
        sg2 = g.edge_subgraph(
            {
                "follows": F.tensor([False, True], dtype=F.bool),
                "plays": F.tensor([False, True, False, False], dtype=F.bool),
                "wishes": F.tensor([False, True], dtype=F.bool),
            }
        )
1604
        _check_subgraph(g, sg2)
1605

1606

nv-dlasalle's avatar
nv-dlasalle committed
1607
@parametrize_idtype
1608
1609
def test_subgraph(idtype):
    g = create_test_heterograph(idtype)
1610
1611
    g_graph = g["follows"]
    g_bipartite = g["plays"]
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
1612

Minjie Wang's avatar
Minjie Wang committed
1613
1614
    x = F.randn((3, 5))
    y = F.randn((2, 4))
1615
1616
    g.nodes["user"].data["h"] = x
    g.edges["follows"].data["h"] = y
Minjie Wang's avatar
Minjie Wang committed
1617
1618

    def _check_subgraph(g, sg):
1619
1620
        assert sg.idtype == g.idtype
        assert sg.device == g.device
1621
1622
1623
        assert sg.ntypes == g.ntypes
        assert sg.etypes == g.etypes
        assert sg.canonical_etypes == g.canonical_etypes
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
        assert F.array_equal(
            F.tensor(sg.nodes["user"].data[dgl.NID]), F.tensor([1, 2], g.idtype)
        )
        assert F.array_equal(
            F.tensor(sg.nodes["game"].data[dgl.NID]), F.tensor([0], g.idtype)
        )
        assert F.array_equal(
            F.tensor(sg.edges["follows"].data[dgl.EID]), F.tensor([1], g.idtype)
        )
        assert F.array_equal(
            F.tensor(sg.edges["plays"].data[dgl.EID]), F.tensor([1], g.idtype)
        )
        assert F.array_equal(
            F.tensor(sg.edges["wishes"].data[dgl.EID]), F.tensor([1], g.idtype)
        )
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
1639
1640
        assert sg.num_nodes("developer") == 0
        assert sg.num_edges("develops") == 0
1641
1642
1643
1644
1645
1646
1647
1648
        assert F.array_equal(
            sg.nodes["user"].data["h"], g.nodes["user"].data["h"][1:3]
        )
        assert F.array_equal(
            sg.edges["follows"].data["h"], g.edges["follows"].data["h"][1:2]
        )

    sg1 = g.subgraph({"user": [1, 2], "game": [0]})
Minjie Wang's avatar
Minjie Wang committed
1649
    _check_subgraph(g, sg1)
1650
    if F._default_context_str != "gpu":
1651
        # TODO(minjie): enable this later
1652
        sg2 = g.edge_subgraph({"follows": [1], "plays": [1], "wishes": [1]})
1653
        _check_subgraph(g, sg2)
Minjie Wang's avatar
Minjie Wang committed
1654

1655
    # backend tensor input
1656
1657
1658
1659
1660
1661
    sg1 = g.subgraph(
        {
            "user": F.tensor([1, 2], dtype=idtype),
            "game": F.tensor([0], dtype=idtype),
        }
    )
1662
    _check_subgraph(g, sg1)
1663
    if F._default_context_str != "gpu":
1664
        # TODO(minjie): enable this later
1665
1666
1667
1668
1669
1670
1671
        sg2 = g.edge_subgraph(
            {
                "follows": F.tensor([1], dtype=idtype),
                "plays": F.tensor([1], dtype=idtype),
                "wishes": F.tensor([1], dtype=idtype),
            }
        )
1672
        _check_subgraph(g, sg2)
1673
1674

    # numpy input
1675
    sg1 = g.subgraph({"user": np.array([1, 2]), "game": np.array([0])})
1676
    _check_subgraph(g, sg1)
1677
    if F._default_context_str != "gpu":
1678
        # TODO(minjie): enable this later
1679
1680
1681
1682
1683
1684
1685
        sg2 = g.edge_subgraph(
            {
                "follows": np.array([1]),
                "plays": np.array([1]),
                "wishes": np.array([1]),
            }
        )
1686
        _check_subgraph(g, sg2)
1687

Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
1688
    def _check_subgraph_single_ntype(g, sg, preserve_nodes=False):
1689
1690
        assert sg.idtype == g.idtype
        assert sg.device == g.device
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
1691
1692
1693
        assert sg.ntypes == g.ntypes
        assert sg.etypes == g.etypes
        assert sg.canonical_etypes == g.canonical_etypes
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
1694
1695

        if not preserve_nodes:
1696
1697
1698
1699
            assert F.array_equal(
                F.tensor(sg.nodes["user"].data[dgl.NID]),
                F.tensor([1, 2], g.idtype),
            )
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
1700
1701
        else:
            for ntype in sg.ntypes:
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
1702
                assert g.num_nodes(ntype) == sg.num_nodes(ntype)
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
1703

1704
1705
1706
        assert F.array_equal(
            F.tensor(sg.edges["follows"].data[dgl.EID]), F.tensor([1], g.idtype)
        )
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
1707
1708

        if not preserve_nodes:
1709
1710
1711
1712
1713
1714
            assert F.array_equal(
                sg.nodes["user"].data["h"], g.nodes["user"].data["h"][1:3]
            )
        assert F.array_equal(
            sg.edges["follows"].data["h"], g.edges["follows"].data["h"][1:2]
        )
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
1715

Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
1716
    def _check_subgraph_single_etype(g, sg, preserve_nodes=False):
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
1717
1718
1719
        assert sg.ntypes == g.ntypes
        assert sg.etypes == g.etypes
        assert sg.canonical_etypes == g.canonical_etypes
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
1720
1721

        if not preserve_nodes:
1722
1723
1724
1725
1726
1727
1728
1729
            assert F.array_equal(
                F.tensor(sg.nodes["user"].data[dgl.NID]),
                F.tensor([0, 1], g.idtype),
            )
            assert F.array_equal(
                F.tensor(sg.nodes["game"].data[dgl.NID]),
                F.tensor([0], g.idtype),
            )
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
1730
1731
        else:
            for ntype in sg.ntypes:
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
1732
                assert g.num_nodes(ntype) == sg.num_nodes(ntype)
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
1733

1734
1735
1736
1737
        assert F.array_equal(
            F.tensor(sg.edges["plays"].data[dgl.EID]),
            F.tensor([0, 1], g.idtype),
        )
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
1738
1739
1740

    sg1_graph = g_graph.subgraph([1, 2])
    _check_subgraph_single_ntype(g_graph, sg1_graph)
1741
    if F._default_context_str != "gpu":
1742
1743
1744
        # TODO(minjie): enable this later
        sg1_graph = g_graph.edge_subgraph([1])
        _check_subgraph_single_ntype(g_graph, sg1_graph)
1745
        sg1_graph = g_graph.edge_subgraph([1], relabel_nodes=False)
1746
1747
1748
        _check_subgraph_single_ntype(g_graph, sg1_graph, True)
        sg2_bipartite = g_bipartite.edge_subgraph([0, 1])
        _check_subgraph_single_etype(g_bipartite, sg2_bipartite)
1749
        sg2_bipartite = g_bipartite.edge_subgraph([0, 1], relabel_nodes=False)
1750
        _check_subgraph_single_etype(g_bipartite, sg2_bipartite, True)
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
1751

1752
    def _check_typed_subgraph1(g, sg):
1753
1754
        assert g.idtype == sg.idtype
        assert g.device == sg.device
1755
1756
        assert set(sg.ntypes) == {"user", "game"}
        assert set(sg.etypes) == {"follows", "plays", "wishes"}
Minjie Wang's avatar
Minjie Wang committed
1757
        for ntype in sg.ntypes:
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
1758
            assert sg.num_nodes(ntype) == g.num_nodes(ntype)
Minjie Wang's avatar
Minjie Wang committed
1759
        for etype in sg.etypes:
1760
1761
            src_sg, dst_sg = sg.all_edges(etype=etype, order="eid")
            src_g, dst_g = g.all_edges(etype=etype, order="eid")
Minjie Wang's avatar
Minjie Wang committed
1762
1763
            assert F.array_equal(src_sg, src_g)
            assert F.array_equal(dst_sg, dst_g)
1764
1765
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
        assert F.array_equal(
            sg.nodes["user"].data["h"], g.nodes["user"].data["h"]
        )
        assert F.array_equal(
            sg.edges["follows"].data["h"], g.edges["follows"].data["h"]
        )
        g.nodes["user"].data["h"] = F.scatter_row(
            g.nodes["user"].data["h"], F.tensor([2]), F.randn((1, 5))
        )
        g.edges["follows"].data["h"] = F.scatter_row(
            g.edges["follows"].data["h"], F.tensor([1]), F.randn((1, 4))
        )
        assert F.array_equal(
            sg.nodes["user"].data["h"], g.nodes["user"].data["h"]
        )
        assert F.array_equal(
            sg.edges["follows"].data["h"], g.edges["follows"].data["h"]
        )
Minjie Wang's avatar
Minjie Wang committed
1782

1783
    def _check_typed_subgraph2(g, sg):
1784
1785
        assert set(sg.ntypes) == {"developer", "game"}
        assert set(sg.etypes) == {"develops"}
1786
        for ntype in sg.ntypes:
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
1787
            assert sg.num_nodes(ntype) == g.num_nodes(ntype)
1788
        for etype in sg.etypes:
1789
1790
            src_sg, dst_sg = sg.all_edges(etype=etype, order="eid")
            src_g, dst_g = g.all_edges(etype=etype, order="eid")
1791
1792
1793
            assert F.array_equal(src_sg, src_g)
            assert F.array_equal(dst_sg, dst_g)

1794
    sg3 = g.node_type_subgraph(["user", "game"])
1795
    _check_typed_subgraph1(g, sg3)
1796
    sg4 = g.edge_type_subgraph(["develops"])
1797
    _check_typed_subgraph2(g, sg4)
1798
    sg5 = g.edge_type_subgraph(["follows", "plays", "wishes"])
1799
    _check_typed_subgraph1(g, sg5)
1800

1801

nv-dlasalle's avatar
nv-dlasalle committed
1802
@parametrize_idtype
1803
def test_apply(idtype):
1804
    def node_udf(nodes):
1805
1806
        return {"h": nodes.data["h"] * 2}

1807
    def node_udf2(nodes):
1808
1809
        return {"h": F.sum(nodes.data["h"], dim=1, keepdims=True)}

1810
    def edge_udf(edges):
1811
        return {"h": edges.data["h"] * 2 + edges.src["h"]}
1812

1813
    g = create_test_heterograph(idtype)
1814
1815
1816
    g.nodes["user"].data["h"] = F.ones((3, 5))
    g.apply_nodes(node_udf, ntype="user")
    assert F.array_equal(g.nodes["user"].data["h"], F.ones((3, 5)) * 2)
Minjie Wang's avatar
Minjie Wang committed
1817

1818
1819
1820
    g["plays"].edata["h"] = F.ones((4, 5))
    g.apply_edges(edge_udf, etype=("user", "plays", "game"))
    assert F.array_equal(g["plays"].edata["h"], F.ones((4, 5)) * 4)
Minjie Wang's avatar
Minjie Wang committed
1821
1822

    # test apply on graph with only one type
1823
1824
    g["follows"].apply_nodes(node_udf)
    assert F.array_equal(g.nodes["user"].data["h"], F.ones((3, 5)) * 4)
1825

1826
1827
    g["plays"].apply_edges(edge_udf)
    assert F.array_equal(g["plays"].edata["h"], F.ones((4, 5)) * 12)
Minjie Wang's avatar
Minjie Wang committed
1828

1829
    # Test the case that feature size changes
1830
1831
1832
    g.nodes["user"].data["h"] = F.ones((3, 5))
    g.apply_nodes(node_udf2, ntype="user")
    assert F.array_equal(g.nodes["user"].data["h"], F.ones((3, 1)) * 5)
1833

Minjie Wang's avatar
Minjie Wang committed
1834
1835
    # test fail case
    # fail due to multiple types
1836
    with pytest.raises(DGLError):
Minjie Wang's avatar
Minjie Wang committed
1837
1838
        g.apply_nodes(node_udf)

1839
    with pytest.raises(DGLError):
Minjie Wang's avatar
Minjie Wang committed
1840
1841
        g.apply_edges(edge_udf)

1842

nv-dlasalle's avatar
nv-dlasalle committed
1843
@parametrize_idtype
1844
def test_level2(idtype):
1845
    # edges = {
Minjie Wang's avatar
Minjie Wang committed
1846
1847
1848
1849
    #    'follows': ([0, 1], [1, 2]),
    #    'plays': ([0, 1, 2, 1], [0, 0, 1, 1]),
    #    'wishes': ([0, 2], [1, 0]),
    #    'develops': ([0, 1], [0, 1]),
1850
    # }
1851
    g = create_test_heterograph(idtype)
1852

Minjie Wang's avatar
Minjie Wang committed
1853
    def rfunc(nodes):
1854
1855
        return {"y": F.sum(nodes.mailbox["m"], 1)}

Minjie Wang's avatar
Minjie Wang committed
1856
    def rfunc2(nodes):
1857
1858
        return {"y": F.max(nodes.mailbox["m"], 1)}

Minjie Wang's avatar
Minjie Wang committed
1859
    def mfunc(edges):
1860
1861
        return {"m": edges.src["h"]}

Minjie Wang's avatar
Minjie Wang committed
1862
    def afunc(nodes):
1863
        return {"y": nodes.data["y"] + 1}
Minjie Wang's avatar
Minjie Wang committed
1864
1865
1866
1867
1868

    #############################################################
    #  send_and_recv
    #############################################################

1869
1870
1871
1872
    g.nodes["user"].data["h"] = F.ones((3, 2))
    g.send_and_recv([2, 3], mfunc, rfunc, etype="plays")
    y = g.nodes["game"].data["y"]
    assert F.array_equal(y, F.tensor([[0.0, 0.0], [2.0, 2.0]]))
Minjie Wang's avatar
Minjie Wang committed
1873
1874

    # only one type
1875
1876
1877
    g["plays"].send_and_recv([2, 3], mfunc, rfunc)
    y = g.nodes["game"].data["y"]
    assert F.array_equal(y, F.tensor([[0.0, 0.0], [2.0, 2.0]]))
1878

Minjie Wang's avatar
Minjie Wang committed
1879
1880
    # test fail case
    # fail due to multiple types
1881
    with pytest.raises(DGLError):
Minjie Wang's avatar
Minjie Wang committed
1882
1883
        g.send_and_recv([2, 3], mfunc, rfunc)

1884
    g.nodes["game"].data.clear()
Minjie Wang's avatar
Minjie Wang committed
1885
1886
1887
1888
1889

    #############################################################
    #  pull
    #############################################################

1890
1891
1892
1893
    g.nodes["user"].data["h"] = F.ones((3, 2))
    g.pull(1, mfunc, rfunc, etype="plays")
    y = g.nodes["game"].data["y"]
    assert F.array_equal(y, F.tensor([[0.0, 0.0], [2.0, 2.0]]))
Minjie Wang's avatar
Minjie Wang committed
1894
1895

    # only one type
1896
1897
1898
    g["plays"].pull(1, mfunc, rfunc)
    y = g.nodes["game"].data["y"]
    assert F.array_equal(y, F.tensor([[0.0, 0.0], [2.0, 2.0]]))
Minjie Wang's avatar
Minjie Wang committed
1899
1900

    # test fail case
1901
    with pytest.raises(DGLError):
Minjie Wang's avatar
Minjie Wang committed
1902
1903
        g.pull(1, mfunc, rfunc)

1904
    g.nodes["game"].data.clear()
Minjie Wang's avatar
Minjie Wang committed
1905
1906
1907
1908
1909

    #############################################################
    #  update_all
    #############################################################

1910
1911
1912
1913
    g.nodes["user"].data["h"] = F.ones((3, 2))
    g.update_all(mfunc, rfunc, etype="plays")
    y = g.nodes["game"].data["y"]
    assert F.array_equal(y, F.tensor([[2.0, 2.0], [2.0, 2.0]]))
Minjie Wang's avatar
Minjie Wang committed
1914
1915

    # only one type
1916
1917
1918
    g["plays"].update_all(mfunc, rfunc)
    y = g.nodes["game"].data["y"]
    assert F.array_equal(y, F.tensor([[2.0, 2.0], [2.0, 2.0]]))
Minjie Wang's avatar
Minjie Wang committed
1919
1920
1921

    # test fail case
    # fail due to multiple types
1922
    with pytest.raises(DGLError):
Minjie Wang's avatar
Minjie Wang committed
1923
1924
1925
1926
        g.update_all(mfunc, rfunc)

    # test multi
    g.multi_update_all(
1927
1928
1929
1930
1931
1932
        {"plays": (mfunc, rfunc), ("user", "wishes", "game"): (mfunc, rfunc2)},
        "sum",
    )
    assert F.array_equal(
        g.nodes["game"].data["y"], F.tensor([[3.0, 3.0], [3.0, 3.0]])
    )
Minjie Wang's avatar
Minjie Wang committed
1933
1934
1935

    # test multi
    g.multi_update_all(
1936
1937
1938
1939
1940
1941
1942
1943
1944
1945
        {
            "plays": (mfunc, rfunc, afunc),
            ("user", "wishes", "game"): (mfunc, rfunc2),
        },
        "sum",
        afunc,
    )
    assert F.array_equal(
        g.nodes["game"].data["y"], F.tensor([[5.0, 5.0], [5.0, 5.0]])
    )
Minjie Wang's avatar
Minjie Wang committed
1946
1947

    # test cross reducer
1948
1949
    g.nodes["user"].data["h"] = F.randn((3, 2))
    for cred in ["sum", "max", "min", "mean", "stack"]:
Minjie Wang's avatar
Minjie Wang committed
1950
        g.multi_update_all(
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
            {"plays": (mfunc, rfunc, afunc), "wishes": (mfunc, rfunc2)},
            cred,
            afunc,
        )
        y = g.nodes["game"].data["y"]
        g["plays"].update_all(mfunc, rfunc, afunc)
        y1 = g.nodes["game"].data["y"]
        g["wishes"].update_all(mfunc, rfunc2)
        y2 = g.nodes["game"].data["y"]
        if cred == "stack":
1961
1962
1963
1964
            # stack has an internal order by edge type id
            yy = F.stack([y1, y2], 1)
            yy = yy + 1  # final afunc
            assert F.array_equal(y, yy)
Minjie Wang's avatar
Minjie Wang committed
1965
1966
1967
1968
1969
1970
1971
        else:
            yy = get_redfn(cred)(F.stack([y1, y2], 0), 0)
            yy = yy + 1  # final afunc
            assert F.array_equal(y, yy)

    # test fail case
    # fail because cannot infer ntype
1972
    with pytest.raises(DGLError):
Minjie Wang's avatar
Minjie Wang committed
1973
        g.update_all(
1974
1975
1976
1977
            {"plays": (mfunc, rfunc), "follows": (mfunc, rfunc2)}, "sum"
        )

    g.nodes["game"].data.clear()
Minjie Wang's avatar
Minjie Wang committed
1978

1979

nv-dlasalle's avatar
nv-dlasalle committed
1980
@parametrize_idtype
1981
1982
1983
@unittest.skipIf(
    F._default_context_str == "cpu", reason="Need gpu for this test"
)
1984
def test_more_nnz(idtype):
1985
1986
1987
1988
1989
1990
    g = dgl.graph(
        ([0, 0, 0, 0, 0], [1, 1, 1, 1, 1]), idtype=idtype, device=F.ctx()
    )
    g.ndata["x"] = F.copy_to(F.ones((2, 5)), ctx=F.ctx())
    g.update_all(fn.copy_u("x", "m"), fn.sum("m", "y"))
    y = g.ndata["y"]
1991
1992
1993
1994
1995
    ans = np.zeros((2, 5))
    ans[1] = 5
    ans = F.copy_to(F.tensor(ans, dtype=F.dtype(y)), ctx=F.ctx())
    assert F.array_equal(y, ans)

1996

nv-dlasalle's avatar
nv-dlasalle committed
1997
@parametrize_idtype
1998
def test_updates(idtype):
1999
    def msg_func(edges):
2000
2001
        return {"m": edges.src["h"]}

2002
    def reduce_func(nodes):
2003
2004
        return {"y": F.sum(nodes.mailbox["m"], 1)}

2005
    def apply_func(nodes):
2006
2007
        return {"y": nodes.data["y"] * 2}

2008
    g = create_test_heterograph(idtype)
2009
    x = F.randn((3, 5))
2010
    g.nodes["user"].data["h"] = x
2011
2012

    for msg, red, apply in itertools.product(
2013
2014
2015
2016
        [fn.copy_u("h", "m"), msg_func],
        [fn.sum("m", "y"), reduce_func],
        [None, apply_func],
    ):
2017
2018
        multiplier = 1 if apply is None else 2

2019
2020
        g["user", "plays", "game"].update_all(msg, red, apply)
        y = g.nodes["game"].data["y"]
2021
2022
        assert F.array_equal(y[0], (x[0] + x[1]) * multiplier)
        assert F.array_equal(y[1], (x[1] + x[2]) * multiplier)
2023
        del g.nodes["game"].data["y"]
2024

2025
2026
2027
2028
        g["user", "plays", "game"].send_and_recv(
            ([0, 1, 2], [0, 1, 1]), msg, red, apply
        )
        y = g.nodes["game"].data["y"]
2029
2030
        assert F.array_equal(y[0], x[0] * multiplier)
        assert F.array_equal(y[1], (x[1] + x[2]) * multiplier)
2031
        del g.nodes["game"].data["y"]
2032
2033

        # pulls from destination (game) node 0
2034
2035
        g["user", "plays", "game"].pull(0, msg, red, apply)
        y = g.nodes["game"].data["y"]
2036
        assert F.array_equal(y[0], (x[0] + x[1]) * multiplier)
2037
        del g.nodes["game"].data["y"]
2038
2039

        # pushes from source (user) node 0
2040
2041
        g["user", "plays", "game"].push(0, msg, red, apply)
        y = g.nodes["game"].data["y"]
2042
        assert F.array_equal(y[0], x[0] * multiplier)
2043
        del g.nodes["game"].data["y"]
Minjie Wang's avatar
Minjie Wang committed
2044

2045

nv-dlasalle's avatar
nv-dlasalle committed
2046
@parametrize_idtype
2047
2048
def test_backward(idtype):
    g = create_test_heterograph(idtype)
Minjie Wang's avatar
Minjie Wang committed
2049
2050
    x = F.randn((3, 5))
    F.attach_grad(x)
2051
    g.nodes["user"].data["h"] = x
Minjie Wang's avatar
Minjie Wang committed
2052
2053
    with F.record_grad():
        g.multi_update_all(
2054
2055
2056
2057
2058
2059
2060
            {
                "plays": (fn.copy_u("h", "m"), fn.sum("m", "y")),
                "wishes": (fn.copy_u("h", "m"), fn.sum("m", "y")),
            },
            "sum",
        )
        y = g.nodes["game"].data["y"]
Minjie Wang's avatar
Minjie Wang committed
2061
2062
        F.backward(y, F.ones(y.shape))
    print(F.grad(x))
2063
2064
2065
2066
2067
2068
2069
2070
2071
2072
    assert F.array_equal(
        F.grad(x),
        F.tensor(
            [
                [2.0, 2.0, 2.0, 2.0, 2.0],
                [2.0, 2.0, 2.0, 2.0, 2.0],
                [2.0, 2.0, 2.0, 2.0, 2.0],
            ]
        ),
    )
2073

2074

nv-dlasalle's avatar
nv-dlasalle committed
2075
@parametrize_idtype
2076
def test_empty_heterograph(idtype):
2077
    def assert_empty(g):
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2078
2079
2080
        assert g.num_nodes("user") == 0
        assert g.num_edges("plays") == 0
        assert g.num_nodes("game") == 0
2081
2082

    # empty src-dst pair
2083
    assert_empty(dgl.heterograph({("user", "plays", "game"): ([], [])}))
2084

2085
2086
2087
    g = dgl.heterograph(
        {("user", "follows", "user"): ([], [])}, idtype=idtype, device=F.ctx()
    )
2088
2089
    assert g.idtype == idtype
    assert g.device == F.ctx()
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2090
2091
    assert g.num_nodes("user") == 0
    assert g.num_edges("follows") == 0
2092
2093

    # empty relation graph with others
2094
2095
2096
2097
2098
2099
2100
2101
    g = dgl.heterograph(
        {
            ("user", "plays", "game"): ([], []),
            ("developer", "develops", "game"): ([0, 1], [0, 1]),
        },
        idtype=idtype,
        device=F.ctx(),
    )
2102
2103
    assert g.idtype == idtype
    assert g.device == F.ctx()
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2104
2105
2106
2107
2108
    assert g.num_nodes("user") == 0
    assert g.num_edges("plays") == 0
    assert g.num_nodes("game") == 2
    assert g.num_edges("develops") == 2
    assert g.num_nodes("developer") == 2
2109

2110

nv-dlasalle's avatar
nv-dlasalle committed
2111
@parametrize_idtype
2112
def test_types_in_function(idtype):
2113
    def mfunc1(edges):
2114
        assert edges.canonical_etype == ("user", "follow", "user")
2115
2116
2117
        return {}

    def rfunc1(nodes):
2118
        assert nodes.ntype == "user"
2119
2120
2121
        return {}

    def filter_nodes1(nodes):
2122
        assert nodes.ntype == "user"
2123
2124
2125
        return F.zeros((3,))

    def filter_edges1(edges):
2126
        assert edges.canonical_etype == ("user", "follow", "user")
2127
2128
2129
        return F.zeros((2,))

    def mfunc2(edges):
2130
        assert edges.canonical_etype == ("user", "plays", "game")
2131
2132
2133
        return {}

    def rfunc2(nodes):
2134
        assert nodes.ntype == "game"
2135
2136
2137
        return {}

    def filter_nodes2(nodes):
2138
        assert nodes.ntype == "game"
2139
2140
2141
        return F.zeros((3,))

    def filter_edges2(edges):
2142
        assert edges.canonical_etype == ("user", "plays", "game")
2143
2144
        return F.zeros((2,))

2145
2146
2147
2148
2149
    g = dgl.heterograph(
        {("user", "follow", "user"): ((0, 1), (1, 2))},
        idtype=idtype,
        device=F.ctx(),
    )
2150
2151
2152
2153
2154
2155
2156
2157
2158
    g.apply_nodes(rfunc1)
    g.apply_edges(mfunc1)
    g.update_all(mfunc1, rfunc1)
    g.send_and_recv([0, 1], mfunc1, rfunc1)
    g.push([0], mfunc1, rfunc1)
    g.pull([1], mfunc1, rfunc1)
    g.filter_nodes(filter_nodes1)
    g.filter_edges(filter_edges1)

2159
2160
2161
2162
2163
2164
    g = dgl.heterograph(
        {("user", "plays", "game"): ([0, 1], [1, 2])},
        idtype=idtype,
        device=F.ctx(),
    )
    g.apply_nodes(rfunc2, ntype="game")
2165
2166
2167
2168
2169
    g.apply_edges(mfunc2)
    g.update_all(mfunc2, rfunc2)
    g.send_and_recv([0, 1], mfunc2, rfunc2)
    g.push([0], mfunc2, rfunc2)
    g.pull([1], mfunc2, rfunc2)
2170
    g.filter_nodes(filter_nodes2, ntype="game")
2171
2172
    g.filter_edges(filter_edges2)

2173

nv-dlasalle's avatar
nv-dlasalle committed
2174
@parametrize_idtype
2175
def test_stack_reduce(idtype):
2176
    # edges = {
2177
2178
2179
2180
    #    'follows': ([0, 1], [1, 2]),
    #    'plays': ([0, 1, 2, 1], [0, 0, 1, 1]),
    #    'wishes': ([0, 2], [1, 0]),
    #    'develops': ([0, 1], [0, 1]),
2181
    # }
2182
    g = create_test_heterograph(idtype)
2183
2184
    g.nodes["user"].data["h"] = F.randn((3, 200))

2185
    def rfunc(nodes):
2186
2187
        return {"y": F.sum(nodes.mailbox["m"], 1)}

2188
    def rfunc2(nodes):
2189
2190
        return {"y": F.max(nodes.mailbox["m"], 1)}

2191
    def mfunc(edges):
2192
2193
        return {"m": edges.src["h"]}

2194
    g.multi_update_all(
2195
2196
2197
        {"plays": (mfunc, rfunc), "wishes": (mfunc, rfunc2)}, "stack"
    )
    assert g.nodes["game"].data["y"].shape == (
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2198
        g.num_nodes("game"),
2199
2200
2201
        2,
        200,
    )
2202
    # only one type-wise update_all, stack still adds one dimension
2203
2204
    g.multi_update_all({"plays": (mfunc, rfunc)}, "stack")
    assert g.nodes["game"].data["y"].shape == (
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2205
        g.num_nodes("game"),
2206
2207
2208
2209
        1,
        200,
    )

2210

nv-dlasalle's avatar
nv-dlasalle committed
2211
@parametrize_idtype
2212
def test_isolated_ntype(idtype):
2213
2214
2215
2216
2217
2218
    g = dgl.heterograph(
        {("A", "AB", "B"): ([0, 1, 2], [1, 2, 3])},
        num_nodes_dict={"A": 3, "B": 4, "C": 4},
        idtype=idtype,
        device=F.ctx(),
    )
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2219
2220
2221
    assert g.num_nodes("A") == 3
    assert g.num_nodes("B") == 4
    assert g.num_nodes("C") == 4
2222
2223
2224
2225
2226
2227
2228

    g = dgl.heterograph(
        {("A", "AC", "C"): ([0, 1, 2], [1, 2, 3])},
        num_nodes_dict={"A": 3, "B": 4, "C": 4},
        idtype=idtype,
        device=F.ctx(),
    )
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2229
2230
2231
    assert g.num_nodes("A") == 3
    assert g.num_nodes("B") == 4
    assert g.num_nodes("C") == 4
2232
2233
2234
2235
2236
2237
2238

    G = dgl.graph(
        ([0, 1, 2], [4, 5, 6]), num_nodes=11, idtype=idtype, device=F.ctx()
    )
    G.ndata[dgl.NTYPE] = F.tensor(
        [0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2], dtype=F.int64
    )
2239
    G.edata[dgl.ETYPE] = F.tensor([0, 0, 0], dtype=F.int64)
2240
    g = dgl.to_heterogeneous(G, ["A", "B", "C"], ["AB"])
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2241
2242
2243
    assert g.num_nodes("A") == 3
    assert g.num_nodes("B") == 4
    assert g.num_nodes("C") == 4
2244

2245

nv-dlasalle's avatar
nv-dlasalle committed
2246
@parametrize_idtype
2247
def test_ismultigraph(idtype):
2248
2249
2250
2251
2252
2253
    g1 = dgl.heterograph(
        {("A", "AB", "B"): ([0, 0, 1, 2], [1, 2, 5, 5])},
        {"A": 6, "B": 6},
        idtype=idtype,
        device=F.ctx(),
    )
2254
    assert g1.is_multigraph == False
2255
2256
2257
2258
2259
2260
    g2 = dgl.heterograph(
        {("A", "AC", "C"): ([0, 0, 0, 1], [1, 1, 2, 5])},
        {"A": 6, "C": 6},
        idtype=idtype,
        device=F.ctx(),
    )
2261
    assert g2.is_multigraph == True
2262
    g3 = dgl.graph(((0, 1), (1, 2)), num_nodes=6, idtype=idtype, device=F.ctx())
2263
    assert g3.is_multigraph == False
2264
2265
2266
    g4 = dgl.graph(
        ([0, 0, 1], [1, 1, 2]), num_nodes=6, idtype=idtype, device=F.ctx()
    )
2267
    assert g4.is_multigraph == True
2268
2269
2270
2271
2272
2273
2274
2275
2276
    g = dgl.heterograph(
        {
            ("A", "AB", "B"): ([0, 0, 1, 2], [1, 2, 5, 5]),
            ("A", "AA", "A"): ([0, 1], [1, 2]),
        },
        {"A": 6, "B": 6},
        idtype=idtype,
        device=F.ctx(),
    )
2277
    assert g.is_multigraph == False
2278
2279
2280
2281
2282
2283
2284
2285
2286
    g = dgl.heterograph(
        {
            ("A", "AB", "B"): ([0, 0, 1, 2], [1, 2, 5, 5]),
            ("A", "AC", "C"): ([0, 0, 0, 1], [1, 1, 2, 5]),
        },
        {"A": 6, "B": 6, "C": 6},
        idtype=idtype,
        device=F.ctx(),
    )
2287
    assert g.is_multigraph == True
2288
2289
2290
2291
2292
2293
2294
2295
2296
    g = dgl.heterograph(
        {
            ("A", "AB", "B"): ([0, 0, 1, 2], [1, 2, 5, 5]),
            ("A", "AA", "A"): ([0, 0, 1], [1, 1, 2]),
        },
        {"A": 6, "B": 6},
        idtype=idtype,
        device=F.ctx(),
    )
2297
    assert g.is_multigraph == True
2298
2299
2300
2301
2302
2303
2304
2305
2306
    g = dgl.heterograph(
        {
            ("A", "AC", "C"): ([0, 0, 0, 1], [1, 1, 2, 5]),
            ("A", "AA", "A"): ([0, 1], [1, 2]),
        },
        {"A": 6, "C": 6},
        idtype=idtype,
        device=F.ctx(),
    )
2307
2308
    assert g.is_multigraph == True

2309
2310
2311

@parametrize_idtype
def test_graph_index_is_unibipartite(idtype):
2312
2313
2314
2315
2316
    g1 = dgl.heterograph(
        {("A", "AB", "B"): ([0, 0, 1], [1, 2, 5])},
        idtype=idtype,
        device=F.ctx(),
    )
2317
2318
2319
    assert g1._graph.is_metagraph_unibipartite()

    # more complicated bipartite
2320
2321
2322
2323
2324
2325
2326
2327
    g2 = dgl.heterograph(
        {
            ("A", "AB", "B"): ([0, 0, 1], [1, 2, 5]),
            ("A", "AC", "C"): ([1, 0], [0, 0]),
        },
        idtype=idtype,
        device=F.ctx(),
    )
2328
2329
    assert g2._graph.is_metagraph_unibipartite()

2330
2331
2332
2333
2334
2335
2336
2337
2338
    g3 = dgl.heterograph(
        {
            ("A", "AB", "B"): ([0, 0, 1], [1, 2, 5]),
            ("A", "AC", "C"): ([1, 0], [0, 0]),
            ("A", "AA", "A"): ([0, 1], [0, 1]),
        },
        idtype=idtype,
        device=F.ctx(),
    )
2339
2340
    assert not g3._graph.is_metagraph_unibipartite()

2341
2342
2343
2344
2345
2346
2347
2348
    g4 = dgl.heterograph(
        {
            ("A", "AB", "B"): ([0, 0, 1], [1, 2, 5]),
            ("C", "CA", "A"): ([1, 0], [0, 0]),
        },
        idtype=idtype,
        device=F.ctx(),
    )
2349
2350
2351
2352

    assert not g4._graph.is_metagraph_unibipartite()


nv-dlasalle's avatar
nv-dlasalle committed
2353
@parametrize_idtype
2354
def test_bipartite(idtype):
2355
2356
2357
2358
2359
    g1 = dgl.heterograph(
        {("A", "AB", "B"): ([0, 0, 1], [1, 2, 5])},
        idtype=idtype,
        device=F.ctx(),
    )
2360
2361
    assert g1.is_unibipartite
    assert len(g1.ntypes) == 2
2362
2363
2364
    assert g1.etypes == ["AB"]
    assert g1.srctypes == ["A"]
    assert g1.dsttypes == ["B"]
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2365
2366
    assert g1.num_nodes("A") == 2
    assert g1.num_nodes("B") == 6
2367
    assert g1.number_of_src_nodes("A") == 2
2368
    assert g1.number_of_src_nodes() == 2
2369
    assert g1.number_of_dst_nodes("B") == 6
2370
    assert g1.number_of_dst_nodes() == 6
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2371
    assert g1.num_edges() == 3
2372
2373
2374
2375
2376
2377
2378
2379
    g1.srcdata["h"] = F.randn((2, 5))
    assert F.array_equal(g1.srcnodes["A"].data["h"], g1.srcdata["h"])
    assert F.array_equal(g1.nodes["A"].data["h"], g1.srcdata["h"])
    assert F.array_equal(g1.nodes["SRC/A"].data["h"], g1.srcdata["h"])
    g1.dstdata["h"] = F.randn((6, 3))
    assert F.array_equal(g1.dstnodes["B"].data["h"], g1.dstdata["h"])
    assert F.array_equal(g1.nodes["B"].data["h"], g1.dstdata["h"])
    assert F.array_equal(g1.nodes["DST/B"].data["h"], g1.dstdata["h"])
2380
2381

    # more complicated bipartite
2382
2383
2384
2385
2386
2387
2388
2389
    g2 = dgl.heterograph(
        {
            ("A", "AB", "B"): ([0, 0, 1], [1, 2, 5]),
            ("A", "AC", "C"): ([1, 0], [0, 0]),
        },
        idtype=idtype,
        device=F.ctx(),
    )
2390
2391

    assert g2.is_unibipartite
2392
2393
    assert g2.srctypes == ["A"]
    assert set(g2.dsttypes) == {"B", "C"}
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2394
2395
2396
    assert g2.num_nodes("A") == 2
    assert g2.num_nodes("B") == 6
    assert g2.num_nodes("C") == 1
2397
    assert g2.number_of_src_nodes("A") == 2
2398
    assert g2.number_of_src_nodes() == 2
2399
2400
2401
2402
2403
2404
2405
2406
2407
2408
2409
2410
2411
2412
2413
2414
    assert g2.number_of_dst_nodes("B") == 6
    assert g2.number_of_dst_nodes("C") == 1
    g2.srcdata["h"] = F.randn((2, 5))
    assert F.array_equal(g2.srcnodes["A"].data["h"], g2.srcdata["h"])
    assert F.array_equal(g2.nodes["A"].data["h"], g2.srcdata["h"])
    assert F.array_equal(g2.nodes["SRC/A"].data["h"], g2.srcdata["h"])

    g3 = dgl.heterograph(
        {
            ("A", "AB", "B"): ([0, 0, 1], [1, 2, 5]),
            ("A", "AC", "C"): ([1, 0], [0, 0]),
            ("A", "AA", "A"): ([0, 1], [0, 1]),
        },
        idtype=idtype,
        device=F.ctx(),
    )
2415
    assert not g3.is_unibipartite
2416

2417
2418
2419
2420
2421
2422
2423
2424
    g4 = dgl.heterograph(
        {
            ("A", "AB", "B"): ([0, 0, 1], [1, 2, 5]),
            ("C", "CA", "A"): ([1, 0], [0, 0]),
        },
        idtype=idtype,
        device=F.ctx(),
    )
2425
2426
2427

    assert not g4.is_unibipartite

2428

nv-dlasalle's avatar
nv-dlasalle committed
2429
@parametrize_idtype
2430
def test_dtype_cast(idtype):
2431
    g = dgl.graph(([0, 1, 0, 2], [0, 1, 1, 0]), idtype=idtype, device=F.ctx())
2432
    assert g.idtype == idtype
2433
2434
    g.ndata["feat"] = F.tensor([3, 4, 5])
    g.edata["h"] = F.tensor([3, 4, 5, 6])
2435
    if idtype == "int32":
2436
        g_cast = g.long()
2437
        assert g_cast.idtype == F.int64
2438
2439
    else:
        g_cast = g.int()
2440
        assert g_cast.idtype == F.int32
2441
    check_graph_equal(g, g_cast, check_idtype=False)
2442

2443

2444
2445
2446
def test_float_cast():
    for t in [F.float16, F.float32, F.float64]:
        idtype = F.int32
2447
2448
2449
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460
        g = dgl.heterograph(
            {
                ("user", "follows", "user"): (
                    F.tensor([0, 1, 1, 2, 2, 3], dtype=idtype),
                    F.tensor([0, 0, 1, 1, 2, 2], dtype=idtype),
                ),
                ("user", "plays", "game"): (
                    F.tensor([0, 1, 1], dtype=idtype),
                    F.tensor([0, 0, 1], dtype=idtype),
                ),
            },
            idtype=idtype,
            device=F.ctx(),
        )
2461
2462
2463
2464
2465
        uvalues = [1, 2, 3, 4]
        gvalues = [5, 6]
        fvalues = [7, 8, 9, 10, 11, 12]
        pvalues = [13, 14, 15]
        dataNamesTypes = [
2466
2467
2468
2469
2470
2471
2472
2473
2474
2475
2476
2477
2478
2479
2480
2481
2482
2483
2484
2485
2486
2487
            ("a", F.float16),
            ("b", F.float32),
            ("c", F.float64),
            ("d", F.int32),
            ("e", F.int64),
        ]
        for name, type in dataNamesTypes:
            g.nodes["user"].data[name] = F.copy_to(
                F.tensor(uvalues, dtype=type), ctx=F.ctx()
            )
        for name, type in dataNamesTypes:
            g.nodes["game"].data[name] = F.copy_to(
                F.tensor(gvalues, dtype=type), ctx=F.ctx()
            )
        for name, type in dataNamesTypes:
            g.edges["follows"].data[name] = F.copy_to(
                F.tensor(fvalues, dtype=type), ctx=F.ctx()
            )
        for name, type in dataNamesTypes:
            g.edges["plays"].data[name] = F.copy_to(
                F.tensor(pvalues, dtype=type), ctx=F.ctx()
            )
2488
2489
2490
2491
2492
2493
2494
2495

        if t == F.float16:
            g = dgl.transforms.functional.to_half(g)
        if t == F.float32:
            g = dgl.transforms.functional.to_float(g)
        if t == F.float64:
            g = dgl.transforms.functional.to_double(g)

2496
        for name, origType in dataNamesTypes:
2497
            # integer tensors shouldn't be converted
2498
2499
2500
2501
2502
            reqType = (
                t
                if (origType in [F.float16, F.float32, F.float64])
                else origType
            )
2503

2504
            values = g.nodes["user"].data[name]
2505
2506
2507
2508
            assert values.dtype == reqType
            assert len(values) == len(uvalues)
            assert F.allclose(values, F.tensor(uvalues), 0, 0)

2509
            values = g.nodes["game"].data[name]
2510
2511
2512
2513
            assert values.dtype == reqType
            assert len(values) == len(gvalues)
            assert F.allclose(values, F.tensor(gvalues), 0, 0)

2514
            values = g.edges["follows"].data[name]
2515
2516
2517
2518
            assert values.dtype == reqType
            assert len(values) == len(fvalues)
            assert F.allclose(values, F.tensor(fvalues), 0, 0)

2519
            values = g.edges["plays"].data[name]
2520
2521
2522
2523
            assert values.dtype == reqType
            assert len(values) == len(pvalues)
            assert F.allclose(values, F.tensor(pvalues), 0, 0)

2524

nv-dlasalle's avatar
nv-dlasalle committed
2525
@parametrize_idtype
2526
def test_format(idtype):
2527
    # single relation
2528
    g = dgl.graph(([0, 1, 0, 2], [0, 1, 1, 0]), idtype=idtype, device=F.ctx())
2529
2530
2531
    assert g.formats()["created"] == ["coo"]
    g1 = g.formats(["coo", "csr", "csc"])
    assert len(g1.formats()["created"]) + len(g1.formats()["not created"]) == 3
2532
    g1.create_formats_()
2533
2534
    assert len(g1.formats()["created"]) == 3
    assert g.formats()["created"] == ["coo"]
2535
2536

    # multiple relation
2537
2538
2539
2540
2541
2542
2543
2544
2545
2546
2547
2548
    g = dgl.heterograph(
        {
            ("user", "follows", "user"): ([0, 1], [1, 2]),
            ("user", "plays", "game"): ([0, 1, 1, 2], [0, 0, 1, 1]),
            ("developer", "develops", "game"): ([0, 1], [0, 1]),
        },
        idtype=idtype,
        device=F.ctx(),
    )
    user_feat = F.randn((g["follows"].number_of_src_nodes(), 5))
    g["follows"].srcdata["h"] = user_feat
    g1 = g.formats("csc")
2549
    # test frame
2550
    assert F.array_equal(g1["follows"].srcdata["h"], user_feat)
2551
    # test each relation graph
2552
2553
    assert g1.formats()["created"] == ["csc"]
    assert len(g1.formats()["not created"]) == 0
2554

2555
2556
2557
2558
2559
2560
    # in_degrees
    g = dgl.rand_graph(100, 2340).to(F.ctx())
    ind_arr = []
    for vid in range(0, 100):
        ind_arr.append(g.in_degrees(vid))
    in_degrees = g.in_degrees()
2561
    g = g.formats("coo")
2562
2563
2564
2565
    for vid in range(0, 100):
        assert g.in_degrees(vid) == ind_arr[vid]
    assert F.array_equal(in_degrees, g.in_degrees())

2566

nv-dlasalle's avatar
nv-dlasalle committed
2567
@parametrize_idtype
2568
def test_edges_order(idtype):
2569
    # (0, 2), (1, 2), (0, 1), (0, 1), (2, 1)
2570
2571
2572
2573
2574
    g = dgl.graph(
        (np.array([0, 1, 0, 0, 2]), np.array([2, 2, 1, 1, 1])),
        idtype=idtype,
        device=F.ctx(),
    )
2575

2576
    print(g.formats())
2577
    src, dst = g.all_edges(order="srcdst")
2578
2579
    assert F.array_equal(src, F.tensor([0, 0, 0, 1, 2], dtype=idtype))
    assert F.array_equal(dst, F.tensor([1, 1, 2, 2, 1], dtype=idtype))
2580

2581

nv-dlasalle's avatar
nv-dlasalle committed
2582
@parametrize_idtype
2583
def test_reverse(idtype):
2584
2585
2586
2587
2588
2589
2590
2591
2592
2593
    g = dgl.heterograph(
        {
            ("user", "follows", "user"): (
                [0, 1, 2, 4, 3, 1, 3],
                [1, 2, 3, 2, 0, 0, 1],
            )
        },
        idtype=idtype,
        device=F.ctx(),
    )
2594
    gidx = g._graph
2595
    r_gidx = gidx.reverse()
2596

Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2597
2598
    assert gidx.num_nodes(0) == r_gidx.num_nodes(0)
    assert gidx.num_edges(0) == r_gidx.num_edges(0)
2599
2600
    g_s, g_d, _ = gidx.edges(0)
    rg_s, rg_d, _ = r_gidx.edges(0)
2601
2602
    assert F.array_equal(g_s, rg_d)
    assert F.array_equal(g_d, rg_s)
2603
2604

    # force to start with 'csr'
2605
2606
    gidx = gidx.formats("csr")
    gidx = gidx.formats(["coo", "csr", "csc"])
2607
    r_gidx = gidx.reverse()
2608
2609
    assert "csr" in gidx.formats()["created"]
    assert "csc" in r_gidx.formats()["created"]
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2610
2611
    assert gidx.num_nodes(0) == r_gidx.num_nodes(0)
    assert gidx.num_edges(0) == r_gidx.num_edges(0)
2612
2613
    g_s, g_d, _ = gidx.edges(0)
    rg_s, rg_d, _ = r_gidx.edges(0)
2614
2615
    assert F.array_equal(g_s, rg_d)
    assert F.array_equal(g_d, rg_s)
2616
2617

    # force to start with 'csc'
2618
2619
    gidx = gidx.formats("csc")
    gidx = gidx.formats(["coo", "csr", "csc"])
2620
    r_gidx = gidx.reverse()
2621
2622
    assert "csc" in gidx.formats()["created"]
    assert "csr" in r_gidx.formats()["created"]
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2623
2624
    assert gidx.num_nodes(0) == r_gidx.num_nodes(0)
    assert gidx.num_edges(0) == r_gidx.num_edges(0)
2625
2626
    g_s, g_d, _ = gidx.edges(0)
    rg_s, rg_d, _ = r_gidx.edges(0)
2627
2628
    assert F.array_equal(g_s, rg_d)
    assert F.array_equal(g_d, rg_s)
2629

2630
2631
2632
2633
2634
2635
2636
2637
2638
2639
2640
2641
2642
2643
2644
    g = dgl.heterograph(
        {
            ("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],
            ),
            ("developer", "develops", "game"): ([0, 1, 1, 2], [0, 0, 1, 1]),
        },
        idtype=idtype,
        device=F.ctx(),
    )
2645
    gidx = g._graph
2646
2647
2648
2649
2650
2651
2652
2653
    r_gidx = gidx.reverse()

    # metagraph
    mg = gidx.metagraph
    r_mg = r_gidx.metagraph
    for etype in range(3):
        assert mg.find_edge(etype) == r_mg.find_edge(etype)[::-1]

2654
    # three node types and three edge types
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2655
2656
2657
2658
2659
2660
    assert gidx.num_nodes(0) == r_gidx.num_nodes(0)
    assert gidx.num_nodes(1) == r_gidx.num_nodes(1)
    assert gidx.num_nodes(2) == r_gidx.num_nodes(2)
    assert gidx.num_edges(0) == r_gidx.num_edges(0)
    assert gidx.num_edges(1) == r_gidx.num_edges(1)
    assert gidx.num_edges(2) == r_gidx.num_edges(2)
2661
2662
    g_s, g_d, _ = gidx.edges(0)
    rg_s, rg_d, _ = r_gidx.edges(0)
2663
2664
    assert F.array_equal(g_s, rg_d)
    assert F.array_equal(g_d, rg_s)
2665
2666
    g_s, g_d, _ = gidx.edges(1)
    rg_s, rg_d, _ = r_gidx.edges(1)
2667
2668
    assert F.array_equal(g_s, rg_d)
    assert F.array_equal(g_d, rg_s)
2669
2670
    g_s, g_d, _ = gidx.edges(2)
    rg_s, rg_d, _ = r_gidx.edges(2)
2671
2672
    assert F.array_equal(g_s, rg_d)
    assert F.array_equal(g_d, rg_s)
2673
2674

    # force to start with 'csr'
2675
2676
    gidx = gidx.formats("csr")
    gidx = gidx.formats(["coo", "csr", "csc"])
2677
    r_gidx = gidx.reverse()
2678
    # three node types and three edge types
2679
2680
    assert "csr" in gidx.formats()["created"]
    assert "csc" in r_gidx.formats()["created"]
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2681
2682
2683
2684
2685
2686
    assert gidx.num_nodes(0) == r_gidx.num_nodes(0)
    assert gidx.num_nodes(1) == r_gidx.num_nodes(1)
    assert gidx.num_nodes(2) == r_gidx.num_nodes(2)
    assert gidx.num_edges(0) == r_gidx.num_edges(0)
    assert gidx.num_edges(1) == r_gidx.num_edges(1)
    assert gidx.num_edges(2) == r_gidx.num_edges(2)
2687
2688
    g_s, g_d, _ = gidx.edges(0)
    rg_s, rg_d, _ = r_gidx.edges(0)
2689
2690
    assert F.array_equal(g_s, rg_d)
    assert F.array_equal(g_d, rg_s)
2691
2692
    g_s, g_d, _ = gidx.edges(1)
    rg_s, rg_d, _ = r_gidx.edges(1)
2693
2694
    assert F.array_equal(g_s, rg_d)
    assert F.array_equal(g_d, rg_s)
2695
2696
    g_s, g_d, _ = gidx.edges(2)
    rg_s, rg_d, _ = r_gidx.edges(2)
2697
2698
    assert F.array_equal(g_s, rg_d)
    assert F.array_equal(g_d, rg_s)
2699
2700

    # force to start with 'csc'
2701
2702
    gidx = gidx.formats("csc")
    gidx = gidx.formats(["coo", "csr", "csc"])
2703
    r_gidx = gidx.reverse()
2704
    # three node types and three edge types
2705
2706
    assert "csc" in gidx.formats()["created"]
    assert "csr" in r_gidx.formats()["created"]
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2707
2708
2709
2710
2711
2712
    assert gidx.num_nodes(0) == r_gidx.num_nodes(0)
    assert gidx.num_nodes(1) == r_gidx.num_nodes(1)
    assert gidx.num_nodes(2) == r_gidx.num_nodes(2)
    assert gidx.num_edges(0) == r_gidx.num_edges(0)
    assert gidx.num_edges(1) == r_gidx.num_edges(1)
    assert gidx.num_edges(2) == r_gidx.num_edges(2)
2713
2714
    g_s, g_d, _ = gidx.edges(0)
    rg_s, rg_d, _ = r_gidx.edges(0)
2715
2716
    assert F.array_equal(g_s, rg_d)
    assert F.array_equal(g_d, rg_s)
2717
2718
    g_s, g_d, _ = gidx.edges(1)
    rg_s, rg_d, _ = r_gidx.edges(1)
2719
2720
    assert F.array_equal(g_s, rg_d)
    assert F.array_equal(g_d, rg_s)
2721
2722
    g_s, g_d, _ = gidx.edges(2)
    rg_s, rg_d, _ = r_gidx.edges(2)
2723
2724
2725
    assert F.array_equal(g_s, rg_d)
    assert F.array_equal(g_d, rg_s)

2726

nv-dlasalle's avatar
nv-dlasalle committed
2727
@parametrize_idtype
2728
2729
def test_clone(idtype):
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
2730
2731
    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())
2732
2733

    new_g = g.clone()
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2734
2735
    assert g.num_nodes() == new_g.num_nodes()
    assert g.num_edges() == new_g.num_edges()
2736
2737
    assert g.device == new_g.device
    assert g.idtype == new_g.idtype
2738
2739
    assert F.array_equal(g.ndata["h"], new_g.ndata["h"])
    assert F.array_equal(g.edata["h"], new_g.edata["h"])
2740
    # data change
2741
2742
2743
2744
    new_g.ndata["h"] = F.copy_to(F.tensor([2, 2, 2], dtype=idtype), ctx=F.ctx())
    assert F.array_equal(g.ndata["h"], new_g.ndata["h"]) == False
    g.edata["h"] = F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())
    assert F.array_equal(g.edata["h"], new_g.edata["h"]) == False
2745
2746
    # graph structure change
    g.add_nodes(1)
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2747
    assert g.num_nodes() != new_g.num_nodes()
2748
    new_g.add_edges(1, 1)
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2749
    assert g.num_edges() != new_g.num_edges()
2750
2751

    # zero data graph
2752
    g = dgl.graph(([], []), num_nodes=0, idtype=idtype, device=F.ctx())
2753
    new_g = g.clone()
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2754
2755
    assert g.num_nodes() == new_g.num_nodes()
    assert g.num_edges() == new_g.num_edges()
2756
2757

    # heterograph
2758
    g = create_test_heterograph3(idtype)
2759
2760
2761
    g.edges["plays"].data["h"] = F.copy_to(
        F.tensor([1, 2, 3, 4], dtype=idtype), ctx=F.ctx()
    )
2762
    new_g = g.clone()
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2763
2764
2765
2766
2767
    assert g.num_nodes("user") == new_g.num_nodes("user")
    assert g.num_nodes("game") == new_g.num_nodes("game")
    assert g.num_nodes("developer") == new_g.num_nodes("developer")
    assert g.num_edges("plays") == new_g.num_edges("plays")
    assert g.num_edges("develops") == new_g.num_edges("develops")
2768
2769
2770
2771
2772
2773
2774
2775
2776
    assert F.array_equal(
        g.nodes["user"].data["h"], new_g.nodes["user"].data["h"]
    )
    assert F.array_equal(
        g.nodes["game"].data["h"], new_g.nodes["game"].data["h"]
    )
    assert F.array_equal(
        g.edges["plays"].data["h"], new_g.edges["plays"].data["h"]
    )
2777
2778
    assert g.device == new_g.device
    assert g.idtype == new_g.idtype
2779
2780
    u, v = g.edges(form="uv", order="eid", etype="plays")
    nu, nv = new_g.edges(form="uv", order="eid", etype="plays")
2781
2782
2783
2784
2785
    assert F.array_equal(u, nu)
    assert F.array_equal(v, nv)
    # graph structure change
    u = F.tensor([0, 4], dtype=idtype)
    v = F.tensor([2, 6], dtype=idtype)
2786
2787
    g.add_edges(u, v, etype="plays")
    u, v = g.edges(form="uv", order="eid", etype="plays")
2788
2789
    assert u.shape[0] != nu.shape[0]
    assert v.shape[0] != nv.shape[0]
2790
2791
2792
2793
2794
2795
2796
2797
2798
2799
2800
2801
    assert (
        g.nodes["user"].data["h"].shape[0]
        != new_g.nodes["user"].data["h"].shape[0]
    )
    assert (
        g.nodes["game"].data["h"].shape[0]
        != new_g.nodes["game"].data["h"].shape[0]
    )
    assert (
        g.edges["plays"].data["h"].shape[0]
        != new_g.edges["plays"].data["h"].shape[0]
    )
2802
2803


nv-dlasalle's avatar
nv-dlasalle committed
2804
@parametrize_idtype
2805
2806
2807
2808
2809
2810
2811
def test_add_edges(idtype):
    # homogeneous graph
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    u = 0
    v = 1
    g.add_edges(u, v)
    assert g.device == F.ctx()
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2812
2813
    assert g.num_nodes() == 3
    assert g.num_edges() == 3
2814
2815
2816
2817
    u = [0]
    v = [1]
    g.add_edges(u, v)
    assert g.device == F.ctx()
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2818
2819
    assert g.num_nodes() == 3
    assert g.num_edges() == 4
2820
2821
2822
2823
    u = F.tensor(u, dtype=idtype)
    v = F.tensor(v, dtype=idtype)
    g.add_edges(u, v)
    assert g.device == F.ctx()
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2824
2825
    assert g.num_nodes() == 3
    assert g.num_edges() == 5
2826
    u, v = g.edges(form="uv", order="eid")
2827
2828
2829
2830
2831
2832
2833
2834
    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
    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.add_edges(u, v)
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2835
2836
    assert g.num_nodes() == 4
    assert g.num_edges() == 4
2837
    u, v = g.edges(form="uv", order="eid")
2838
2839
2840
2841
2842
    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())
2843
2844
    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())
2845
2846
    u = F.tensor([0, 1], dtype=idtype)
    v = F.tensor([2, 3], dtype=idtype)
2847
2848
2849
2850
    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()),
    }
2851
    g.add_edges(u, v, e_feat)
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2852
2853
    assert g.num_nodes() == 4
    assert g.num_edges() == 4
2854
    u, v = g.edges(form="uv", order="eid")
2855
2856
    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))
2857
2858
2859
    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))
2860
2861

    # zero data graph
2862
    g = dgl.graph(([], []), num_nodes=0, idtype=idtype, device=F.ctx())
2863
2864
    u = F.tensor([0, 1], dtype=idtype)
    v = F.tensor([2, 2], dtype=idtype)
2865
2866
2867
2868
    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()),
    }
2869
    g.add_edges(u, v, e_feat)
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2870
2871
    assert g.num_nodes() == 3
    assert g.num_edges() == 2
2872
    u, v = g.edges(form="uv", order="eid")
2873
2874
    assert F.array_equal(u, F.tensor([0, 1], dtype=idtype))
    assert F.array_equal(v, F.tensor([2, 2], dtype=idtype))
2875
2876
    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))
2877
2878

    # bipartite graph
2879
2880
2881
2882
2883
    g = dgl.heterograph(
        {("user", "plays", "game"): ([0, 1], [1, 2])},
        idtype=idtype,
        device=F.ctx(),
    )
2884
2885
2886
2887
    u = 0
    v = 1
    g.add_edges(u, v)
    assert g.device == F.ctx()
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2888
2889
2890
    assert g.num_nodes("user") == 2
    assert g.num_nodes("game") == 3
    assert g.num_edges() == 3
2891
2892
2893
2894
    u = [0]
    v = [1]
    g.add_edges(u, v)
    assert g.device == F.ctx()
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2895
2896
2897
    assert g.num_nodes("user") == 2
    assert g.num_nodes("game") == 3
    assert g.num_edges() == 4
2898
2899
2900
2901
    u = F.tensor(u, dtype=idtype)
    v = F.tensor(v, dtype=idtype)
    g.add_edges(u, v)
    assert g.device == F.ctx()
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2902
2903
2904
    assert g.num_nodes("user") == 2
    assert g.num_nodes("game") == 3
    assert g.num_edges() == 5
2905
    u, v = g.edges(form="uv")
2906
2907
2908
2909
    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
2910
2911
2912
2913
2914
    g = dgl.heterograph(
        {("user", "plays", "game"): ([0, 1], [1, 2])},
        idtype=idtype,
        device=F.ctx(),
    )
2915
2916
2917
2918
    u = F.tensor([0, 2], dtype=idtype)
    v = F.tensor([2, 3], dtype=idtype)
    g.add_edges(u, v)
    assert g.device == F.ctx()
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2919
2920
2921
    assert g.num_nodes("user") == 3
    assert g.num_nodes("game") == 4
    assert g.num_edges() == 4
2922
    u, v = g.edges(form="uv", order="eid")
2923
2924
2925
2926
    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
2927
2928
2929
2930
2931
2932
2933
2934
2935
2936
2937
2938
    g = dgl.heterograph(
        {("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()
    )
    g.edata["h"] = F.copy_to(F.tensor([1, 1], dtype=idtype), ctx=F.ctx())
2939
2940
    u = F.tensor([0, 2], dtype=idtype)
    v = F.tensor([2, 3], dtype=idtype)
2941
2942
2943
2944
    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()),
    }
2945
    g.add_edges(u, v, e_feat)
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2946
2947
2948
    assert g.num_nodes("user") == 3
    assert g.num_nodes("game") == 4
    assert g.num_edges() == 4
2949
    u, v = g.edges(form="uv", order="eid")
2950
2951
    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))
2952
2953
2954
2955
2956
2957
2958
2959
    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)
    )
    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))
2960
2961

    # heterogeneous graph
2962
    g = create_test_heterograph3(idtype)
2963
2964
    u = F.tensor([0, 2], dtype=idtype)
    v = F.tensor([2, 3], dtype=idtype)
2965
    g.add_edges(u, v, etype="plays")
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2966
2967
2968
2969
2970
    assert g.num_nodes("user") == 3
    assert g.num_nodes("game") == 4
    assert g.num_nodes("developer") == 2
    assert g.num_edges("plays") == 6
    assert g.num_edges("develops") == 2
2971
    u, v = g.edges(form="uv", order="eid", etype="plays")
2972
2973
    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))
2974
2975
2976
2977
2978
2979
2980
2981
2982
    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)
    )
2983
2984

    # add with feature
2985
    e_feat = {"h": F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())}
2986
2987
    u = F.tensor([0, 2], dtype=idtype)
    v = F.tensor([2, 3], dtype=idtype)
2988
2989
2990
2991
    g.nodes["game"].data["h"] = F.copy_to(
        F.tensor([2, 2, 1, 1], dtype=idtype), ctx=F.ctx()
    )
    g.add_edges(u, v, data=e_feat, etype="develops")
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
2992
2993
2994
2995
2996
    assert g.num_nodes("user") == 3
    assert g.num_nodes("game") == 4
    assert g.num_nodes("developer") == 3
    assert g.num_edges("plays") == 6
    assert g.num_edges("develops") == 4
2997
    u, v = g.edges(form="uv", order="eid", etype="develops")
2998
2999
    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))
3000
3001
3002
3003
3004
3005
3006
3007
3008
3009
    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)
    )

3010

nv-dlasalle's avatar
nv-dlasalle committed
3011
@parametrize_idtype
3012
3013
3014
def test_add_nodes(idtype):
    # homogeneous Graphs
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
3015
    g.ndata["h"] = F.copy_to(F.tensor([1, 1, 1], dtype=idtype), ctx=F.ctx())
3016
    g.add_nodes(1)
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
3017
    assert g.num_nodes() == 4
3018
    assert F.array_equal(g.ndata["h"], F.tensor([1, 1, 1, 0], dtype=idtype))
3019
3020

    # zero node graph
3021
    g = dgl.graph(([], []), num_nodes=3, idtype=idtype, device=F.ctx())
3022
3023
3024
3025
    g.ndata["h"] = F.copy_to(F.tensor([1, 1, 1], dtype=idtype), ctx=F.ctx())
    g.add_nodes(
        1, data={"h": F.copy_to(F.tensor([2], dtype=idtype), ctx=F.ctx())}
    )
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
3026
    assert g.num_nodes() == 4
3027
    assert F.array_equal(g.ndata["h"], F.tensor([1, 1, 1, 2], dtype=idtype))
3028
3029

    # bipartite graph
3030
3031
3032
3033
3034
3035
3036
3037
3038
3039
    g = dgl.heterograph(
        {("user", "plays", "game"): ([0, 1], [1, 2])},
        idtype=idtype,
        device=F.ctx(),
    )
    g.add_nodes(
        2,
        data={"h": F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())},
        ntype="user",
    )
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
3040
    assert g.num_nodes("user") == 4
3041
3042
3043
3044
    assert F.array_equal(
        g.nodes["user"].data["h"], F.tensor([0, 0, 2, 2], dtype=idtype)
    )
    g.add_nodes(2, ntype="game")
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
3045
    assert g.num_nodes("game") == 5
3046
3047

    # heterogeneous graph
3048
    g = create_test_heterograph3(idtype)
3049
3050
3051
3052
3053
3054
3055
    g.add_nodes(1, ntype="user")
    g.add_nodes(
        2,
        data={"h": F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())},
        ntype="game",
    )
    g.add_nodes(0, ntype="developer")
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
3056
3057
3058
    assert g.num_nodes("user") == 4
    assert g.num_nodes("game") == 4
    assert g.num_nodes("developer") == 2
3059
3060
3061
3062
3063
3064
3065
3066
3067
3068
3069
3070
    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)
    )


@unittest.skipIf(
    dgl.backend.backend_name == "mxnet",
    reason="MXNet has error with (0,) shape tensor.",
)
nv-dlasalle's avatar
nv-dlasalle committed
3071
@parametrize_idtype
3072
3073
3074
3075
3076
def test_remove_edges(idtype):
    # homogeneous Graphs
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    e = 0
    g.remove_edges(e)
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
3077
    assert g.num_edges() == 1
3078
    u, v = g.edges(form="uv", order="eid")
3079
3080
3081
3082
3083
    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.remove_edges(e)
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
3084
    assert g.num_edges() == 1
3085
    u, v = g.edges(form="uv", order="eid")
3086
3087
3088
3089
    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.remove_edges(e)
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
3090
    assert g.num_edges() == 0
3091
3092
3093

    # has node data
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
3094
    g.ndata["h"] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
3095
    g.remove_edges(1)
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
3096
    assert g.num_edges() == 1
3097
    assert F.array_equal(g.ndata["h"], F.tensor([1, 2, 3], dtype=idtype))
3098
3099
3100

    # has edge data
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
3101
    g.edata["h"] = F.copy_to(F.tensor([1, 2], dtype=idtype), ctx=F.ctx())
3102
    g.remove_edges(0)
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
3103
    assert g.num_edges() == 1
3104
    assert F.array_equal(g.edata["h"], F.tensor([2], dtype=idtype))
3105
3106
3107
3108
3109
3110
3111
3112
3113
3114

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

    # bipartite graph
3115
3116
3117
3118
3119
    g = dgl.heterograph(
        {("user", "plays", "game"): ([0, 1], [1, 2])},
        idtype=idtype,
        device=F.ctx(),
    )
3120
3121
    e = 0
    g.remove_edges(e)
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
3122
    assert g.num_edges() == 1
3123
    u, v = g.edges(form="uv", order="eid")
3124
3125
    assert F.array_equal(u, F.tensor([1], dtype=idtype))
    assert F.array_equal(v, F.tensor([2], dtype=idtype))
3126
    g = dgl.heterograph(
3127
3128
3129
3130
        {("user", "plays", "game"): ([0, 1], [1, 2])},
        idtype=idtype,
        device=F.ctx(),
    )
3131
3132
    e = [0]
    g.remove_edges(e)
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
3133
    assert g.num_edges() == 1
3134
    u, v = g.edges(form="uv", order="eid")
3135
3136
3137
3138
    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.remove_edges(e)
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
3139
    assert g.num_edges() == 0
3140
3141

    # has data
3142
    g = dgl.heterograph(
3143
3144
3145
3146
3147
3148
3149
3150
3151
3152
3153
        {("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()
    )
    g.edata["h"] = F.copy_to(F.tensor([1, 2], dtype=idtype), ctx=F.ctx())
3154
    g.remove_edges(1)
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
3155
    assert g.num_edges() == 1
3156
3157
3158
3159
3160
3161
3162
    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)
    )
    assert F.array_equal(g.edata["h"], F.tensor([1], dtype=idtype))
3163
3164

    # heterogeneous graph
3165
    g = create_test_heterograph3(idtype)
3166
3167
3168
3169
    g.edges["plays"].data["h"] = F.copy_to(
        F.tensor([1, 2, 3, 4], dtype=idtype), ctx=F.ctx()
    )
    g.remove_edges(1, etype="plays")
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
3170
    assert g.num_edges("plays") == 3
3171
    u, v = g.edges(form="uv", order="eid", etype="plays")
3172
3173
    assert F.array_equal(u, F.tensor([0, 1, 2], dtype=idtype))
    assert F.array_equal(v, F.tensor([0, 1, 1], dtype=idtype))
3174
3175
3176
    assert F.array_equal(
        g.edges["plays"].data["h"], F.tensor([1, 3, 4], dtype=idtype)
    )
3177
    # remove all edges of 'develops'
3178
    g.remove_edges([0, 1], etype="develops")
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
3179
    assert g.num_edges("develops") == 0
3180
3181
3182
3183
3184
3185
3186
3187
3188
3189
    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)
    )

3190

nv-dlasalle's avatar
nv-dlasalle committed
3191
@parametrize_idtype
3192
3193
3194
3195
3196
def test_remove_nodes(idtype):
    # homogeneous Graphs
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    n = 0
    g.remove_nodes(n)
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
3197
3198
    assert g.num_nodes() == 2
    assert g.num_edges() == 1
3199
    u, v = g.edges(form="uv", order="eid")
3200
3201
3202
3203
3204
    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.remove_nodes(n)
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
3205
3206
    assert g.num_nodes() == 2
    assert g.num_edges() == 0
3207
3208
3209
    g = dgl.graph(([0, 1], [1, 2]), idtype=idtype, device=F.ctx())
    n = F.tensor([2], dtype=idtype)
    g.remove_nodes(n)
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
3210
3211
    assert g.num_nodes() == 2
    assert g.num_edges() == 1
3212
    u, v = g.edges(form="uv", order="eid")
3213
3214
3215
3216
3217
3218
3219
3220
3221
3222
3223
3224
3225
    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())
3226
3227
    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())
3228
    g.remove_nodes(F.tensor([0], dtype=idtype))
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
3229
3230
    assert g.num_nodes() == 2
    assert g.num_edges() == 1
3231
    u, v = g.edges(form="uv", order="eid")
3232
3233
    assert F.array_equal(u, F.tensor([1], dtype=idtype))
    assert F.array_equal(v, F.tensor([1], dtype=idtype))
3234
3235
    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))
3236
3237

    # node id larger than current max node id
3238
    g = dgl.heterograph(
3239
3240
3241
3242
        {("user", "plays", "game"): ([0, 1], [1, 2])},
        idtype=idtype,
        device=F.ctx(),
    )
3243
    n = 0
3244
    g.remove_nodes(n, ntype="user")
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
3245
3246
3247
    assert g.num_nodes("user") == 1
    assert g.num_nodes("game") == 3
    assert g.num_edges() == 1
3248
    u, v = g.edges(form="uv", order="eid")
3249
3250
    assert F.array_equal(u, F.tensor([0], dtype=idtype))
    assert F.array_equal(v, F.tensor([2], dtype=idtype))
3251
    g = dgl.heterograph(
3252
3253
3254
3255
        {("user", "plays", "game"): ([0, 1], [1, 2])},
        idtype=idtype,
        device=F.ctx(),
    )
3256
    n = [1]
3257
    g.remove_nodes(n, ntype="user")
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
3258
3259
3260
    assert g.num_nodes("user") == 1
    assert g.num_nodes("game") == 3
    assert g.num_edges() == 1
3261
    u, v = g.edges(form="uv", order="eid")
3262
3263
    assert F.array_equal(u, F.tensor([0], dtype=idtype))
    assert F.array_equal(v, F.tensor([1], dtype=idtype))
3264
    g = dgl.heterograph(
3265
3266
3267
3268
        {("user", "plays", "game"): ([0, 1], [1, 2])},
        idtype=idtype,
        device=F.ctx(),
    )
3269
    n = F.tensor([0], dtype=idtype)
3270
    g.remove_nodes(n, ntype="game")
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
3271
3272
3273
    assert g.num_nodes("user") == 2
    assert g.num_nodes("game") == 2
    assert g.num_edges() == 2
3274
    u, v = g.edges(form="uv", order="eid")
3275
    assert F.array_equal(u, F.tensor([0, 1], dtype=idtype))
3276
    assert F.array_equal(v, F.tensor([0, 1], dtype=idtype))
3277
3278

    # heterogeneous graph
3279
    g = create_test_heterograph3(idtype)
3280
3281
3282
3283
    g.edges["plays"].data["h"] = F.copy_to(
        F.tensor([1, 2, 3, 4], dtype=idtype), ctx=F.ctx()
    )
    g.remove_nodes(0, ntype="game")
Hongzhi (Steve), Chen's avatar
Hongzhi (Steve), Chen committed
3284
3285
3286
3287
3288
    assert g.num_nodes("user") == 3
    assert g.num_nodes("game") == 1
    assert g.num_nodes("developer") == 2
    assert g.num_edges("plays") == 2
    assert g.num_edges("develops") == 1
3289
3290
3291
3292
3293
3294
3295
3296
    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], dtype=idtype))
    assert F.array_equal(
        g.nodes["developer"].data["h"], F.tensor([3, 3], dtype=idtype)
    )
    u, v = g.edges(form="uv", order="eid", etype="plays")
3297
3298
    assert F.array_equal(u, F.tensor([1, 2], dtype=idtype))
    assert F.array_equal(v, F.tensor([0, 0], dtype=idtype))
3299
3300
3301
3302
    assert F.array_equal(
        g.edges["plays"].data["h"], F.tensor([3, 4], dtype=idtype)
    )
    u, v = g.edges(form="uv", order="eid", etype="develops")
3303
3304
    assert F.array_equal(u, F.tensor([1], dtype=idtype))
    assert F.array_equal(v, F.tensor([0], dtype=idtype))
3305

3306

nv-dlasalle's avatar
nv-dlasalle committed
3307
@parametrize_idtype
3308
3309
def test_frame(idtype):
    g = dgl.graph(([0, 1, 2], [1, 2, 3]), idtype=idtype, device=F.ctx())
3310
3311
    g.ndata["h"] = F.copy_to(F.tensor([0, 1, 2, 3], dtype=idtype), ctx=F.ctx())
    g.edata["h"] = F.copy_to(F.tensor([0, 1, 2], dtype=idtype), ctx=F.ctx())
3312
3313
3314
3315

    # remove nodes
    sg = dgl.remove_nodes(g, [3])
    # check for lazy update
3316
3317
3318
3319
    assert F.array_equal(sg._node_frames[0]._columns["h"].storage, g.ndata["h"])
    assert F.array_equal(sg._edge_frames[0]._columns["h"].storage, g.edata["h"])
    assert sg.ndata["h"].shape[0] == 3
    assert sg.edata["h"].shape[0] == 2
3320
    # update after read
3321
3322
3323
3324
3325
3326
3327
    assert F.array_equal(
        sg._node_frames[0]._columns["h"].storage,
        F.tensor([0, 1, 2], dtype=idtype),
    )
    assert F.array_equal(
        sg._edge_frames[0]._columns["h"].storage, F.tensor([0, 1], dtype=idtype)
    )
3328
3329

    ng = dgl.add_nodes(sg, 1)
3330
3331
3332
3333
3334
    assert ng.ndata["h"].shape[0] == 4
    assert F.array_equal(
        ng._node_frames[0]._columns["h"].storage,
        F.tensor([0, 1, 2, 0], dtype=idtype),
    )
3335
    ng = dgl.add_edges(ng, [3], [1])
3336
3337
3338
3339
3340
    assert ng.edata["h"].shape[0] == 3
    assert F.array_equal(
        ng._edge_frames[0]._columns["h"].storage,
        F.tensor([0, 1, 0], dtype=idtype),
    )
3341
3342
3343

    # multi level lazy update
    sg = dgl.remove_nodes(g, [3])
3344
3345
    assert F.array_equal(sg._node_frames[0]._columns["h"].storage, g.ndata["h"])
    assert F.array_equal(sg._edge_frames[0]._columns["h"].storage, g.edata["h"])
3346
    ssg = dgl.remove_nodes(sg, [1])
3347
3348
3349
3350
3351
3352
    assert F.array_equal(
        ssg._node_frames[0]._columns["h"].storage, g.ndata["h"]
    )
    assert F.array_equal(
        ssg._edge_frames[0]._columns["h"].storage, g.edata["h"]
    )
3353
    # ssg is changed
3354
3355
3356
3357
3358
3359
    assert ssg.ndata["h"].shape[0] == 2
    assert ssg.edata["h"].shape[0] == 0
    assert F.array_equal(
        ssg._node_frames[0]._columns["h"].storage,
        F.tensor([0, 2], dtype=idtype),
    )
3360
    # sg still in lazy model
3361
3362
3363
3364
3365
3366
3367
3368
3369
3370
3371
3372
    assert F.array_equal(sg._node_frames[0]._columns["h"].storage, g.ndata["h"])
    assert F.array_equal(sg._edge_frames[0]._columns["h"].storage, g.edata["h"])


@unittest.skipIf(
    dgl.backend.backend_name == "tensorflow",
    reason="TensorFlow always create a new tensor",
)
@unittest.skipIf(
    F._default_context_str == "cpu",
    reason="cpu do not have context change problem",
)
nv-dlasalle's avatar
nv-dlasalle committed
3373
@parametrize_idtype
3374
def test_frame_device(idtype):
3375
3376
3377
3378
    g = dgl.graph(([0, 1, 2], [2, 3, 1]))
    g.ndata["h"] = F.copy_to(F.tensor([1, 1, 1, 2], dtype=idtype), ctx=F.cpu())
    g.ndata["hh"] = F.copy_to(F.ones((4, 3), dtype=idtype), ctx=F.cpu())
    g.edata["h"] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.cpu())
3379
3380
3381

    g = g.to(F.ctx())
    # lazy device copy
3382
3383
3384
3385
3386
3387
    assert F.context(g._node_frames[0]._columns["h"].storage) == F.cpu()
    assert F.context(g._node_frames[0]._columns["hh"].storage) == F.cpu()
    print(g.ndata["h"])
    assert F.context(g._node_frames[0]._columns["h"].storage) == F.ctx()
    assert F.context(g._node_frames[0]._columns["hh"].storage) == F.cpu()
    assert F.context(g._edge_frames[0]._columns["h"].storage) == F.cpu()
3388
3389

    # lazy device copy in subgraph
3390
3391
3392
3393
3394
3395
3396
    sg = dgl.node_subgraph(g, [0, 1, 2])
    assert F.context(sg._node_frames[0]._columns["h"].storage) == F.ctx()
    assert F.context(sg._node_frames[0]._columns["hh"].storage) == F.cpu()
    assert F.context(sg._edge_frames[0]._columns["h"].storage) == F.cpu()
    print(sg.ndata["hh"])
    assert F.context(sg._node_frames[0]._columns["hh"].storage) == F.ctx()
    assert F.context(sg._edge_frames[0]._columns["h"].storage) == F.cpu()
3397
3398
3399

    # back to cpu
    sg = sg.to(F.cpu())
3400
3401
3402
3403
3404
3405
3406
3407
3408
    assert F.context(sg._node_frames[0]._columns["h"].storage) == F.ctx()
    assert F.context(sg._node_frames[0]._columns["hh"].storage) == F.ctx()
    assert F.context(sg._edge_frames[0]._columns["h"].storage) == F.cpu()
    print(sg.ndata["h"])
    print(sg.ndata["hh"])
    print(sg.edata["h"])
    assert F.context(sg._node_frames[0]._columns["h"].storage) == F.cpu()
    assert F.context(sg._node_frames[0]._columns["hh"].storage) == F.cpu()
    assert F.context(sg._edge_frames[0]._columns["h"].storage) == F.cpu()
3409
3410
3411

    # set some field
    sg = sg.to(F.ctx())
3412
3413
3414
3415
3416
    assert F.context(sg._node_frames[0]._columns["h"].storage) == F.cpu()
    sg.ndata["h"][0] = 5
    assert F.context(sg._node_frames[0]._columns["h"].storage) == F.ctx()
    assert F.context(sg._node_frames[0]._columns["hh"].storage) == F.cpu()
    assert F.context(sg._edge_frames[0]._columns["h"].storage) == F.cpu()
3417
3418
3419

    # add nodes
    ng = dgl.add_nodes(sg, 3)
3420
3421
3422
3423
    assert F.context(ng._node_frames[0]._columns["h"].storage) == F.ctx()
    assert F.context(ng._node_frames[0]._columns["hh"].storage) == F.ctx()
    assert F.context(ng._edge_frames[0]._columns["h"].storage) == F.cpu()

3424

nv-dlasalle's avatar
nv-dlasalle committed
3425
@parametrize_idtype
3426
def test_create_block(idtype):
3427
3428
3429
    block = dgl.create_block(
        ([0, 1, 2], [1, 2, 3]), idtype=idtype, device=F.ctx()
    )
3430
3431
3432
3433
3434
3435
3436
3437
3438
3439
3440
3441
3442
3443
    assert block.num_src_nodes() == 3
    assert block.num_dst_nodes() == 4
    assert block.num_edges() == 3

    block = dgl.create_block(([], []), idtype=idtype, device=F.ctx())
    assert block.num_src_nodes() == 0
    assert block.num_dst_nodes() == 0
    assert block.num_edges() == 0

    block = dgl.create_block(([], []), 3, 4, idtype=idtype, device=F.ctx())
    assert block.num_src_nodes() == 3
    assert block.num_dst_nodes() == 4
    assert block.num_edges() == 0

3444
3445
3446
    block = dgl.create_block(
        ([0, 1, 2], [1, 2, 3]), 4, 5, idtype=idtype, device=F.ctx()
    )
3447
3448
3449
3450
3451
3452
3453
    assert block.num_src_nodes() == 4
    assert block.num_dst_nodes() == 5
    assert block.num_edges() == 3

    sx = F.randn((4, 5))
    dx = F.randn((5, 6))
    ex = F.randn((3, 4))
3454
3455
3456
    block.srcdata["x"] = sx
    block.dstdata["x"] = dx
    block.edata["x"] = ex
3457
3458
3459
3460
3461

    g = dgl.block_to_graph(block)
    assert g.num_src_nodes() == 4
    assert g.num_dst_nodes() == 5
    assert g.num_edges() == 3
3462
3463
3464
3465
3466
3467
3468
3469
3470
3471
3472
3473
3474
3475
3476
3477
3478
3479
3480
3481
3482
3483
3484
3485
3486
3487
3488
3489
3490
3491
3492
3493
3494
3495
3496
3497
3498
3499
3500
3501
3502
3503
3504
3505
3506
3507
3508
3509
3510
3511
3512
3513
3514
3515
3516
3517
3518
3519
3520
3521
3522
    assert g.srcdata["x"] is sx
    assert g.dstdata["x"] is dx
    assert g.edata["x"] is ex

    block = dgl.create_block(
        {
            ("A", "AB", "B"): ([1, 2, 3], [2, 1, 0]),
            ("B", "BA", "A"): ([2, 3], [3, 4]),
        },
        idtype=idtype,
        device=F.ctx(),
    )
    assert block.num_src_nodes("A") == 4
    assert block.num_src_nodes("B") == 4
    assert block.num_dst_nodes("B") == 3
    assert block.num_dst_nodes("A") == 5
    assert block.num_edges("AB") == 3
    assert block.num_edges("BA") == 2

    block = dgl.create_block(
        {("A", "AB", "B"): ([], []), ("B", "BA", "A"): ([], [])},
        idtype=idtype,
        device=F.ctx(),
    )
    assert block.num_src_nodes("A") == 0
    assert block.num_src_nodes("B") == 0
    assert block.num_dst_nodes("B") == 0
    assert block.num_dst_nodes("A") == 0
    assert block.num_edges("AB") == 0
    assert block.num_edges("BA") == 0

    block = dgl.create_block(
        {("A", "AB", "B"): ([], []), ("B", "BA", "A"): ([], [])},
        num_src_nodes={"A": 5, "B": 5},
        num_dst_nodes={"A": 6, "B": 4},
        idtype=idtype,
        device=F.ctx(),
    )
    assert block.num_src_nodes("A") == 5
    assert block.num_src_nodes("B") == 5
    assert block.num_dst_nodes("B") == 4
    assert block.num_dst_nodes("A") == 6
    assert block.num_edges("AB") == 0
    assert block.num_edges("BA") == 0

    block = dgl.create_block(
        {
            ("A", "AB", "B"): ([1, 2, 3], [2, 1, 0]),
            ("B", "BA", "A"): ([2, 3], [3, 4]),
        },
        num_src_nodes={"A": 5, "B": 5},
        num_dst_nodes={"A": 6, "B": 4},
        idtype=idtype,
        device=F.ctx(),
    )
    assert block.num_src_nodes("A") == 5
    assert block.num_src_nodes("B") == 5
    assert block.num_dst_nodes("B") == 4
    assert block.num_dst_nodes("A") == 6
    assert block.num_edges(("A", "AB", "B")) == 3
    assert block.num_edges(("B", "BA", "A")) == 2
3523
3524
3525
3526
3527
3528
3529

    sax = F.randn((5, 3))
    sbx = F.randn((5, 4))
    dax = F.randn((6, 5))
    dbx = F.randn((4, 6))
    eabx = F.randn((3, 7))
    ebax = F.randn((2, 8))
3530
3531
3532
3533
3534
3535
    block.srcnodes["A"].data["x"] = sax
    block.srcnodes["B"].data["x"] = sbx
    block.dstnodes["A"].data["x"] = dax
    block.dstnodes["B"].data["x"] = dbx
    block.edges["AB"].data["x"] = eabx
    block.edges["BA"].data["x"] = ebax
3536
3537

    hg = dgl.block_to_graph(block)
3538
3539
3540
3541
3542
3543
3544
3545
3546
3547
3548
3549
3550
    assert hg.num_nodes("A_src") == 5
    assert hg.num_nodes("B_src") == 5
    assert hg.num_nodes("A_dst") == 6
    assert hg.num_nodes("B_dst") == 4
    assert hg.num_edges(("A_src", "AB", "B_dst")) == 3
    assert hg.num_edges(("B_src", "BA", "A_dst")) == 2
    assert hg.nodes["A_src"].data["x"] is sax
    assert hg.nodes["B_src"].data["x"] is sbx
    assert hg.nodes["A_dst"].data["x"] is dax
    assert hg.nodes["B_dst"].data["x"] is dbx
    assert hg.edges["AB"].data["x"] is eabx
    assert hg.edges["BA"].data["x"] is ebax

3551

nv-dlasalle's avatar
nv-dlasalle committed
3552
@parametrize_idtype
3553
@pytest.mark.parametrize("fmt", ["coo", "csr", "csc"])
3554
def test_adj_tensors(idtype, fmt):
3555
    if fmt == "coo":
3556
3557
3558
3559
3560
        A = ssp.random(10, 10, 0.2).tocoo()
        A.data = np.arange(20)
        row = F.tensor(A.row, idtype)
        col = F.tensor(A.col, idtype)
        g = dgl.graph((row, col))
3561
    elif fmt == "csr":
3562
3563
3564
3565
        A = ssp.random(10, 10, 0.2).tocsr()
        A.data = np.arange(20)
        indptr = F.tensor(A.indptr, idtype)
        indices = F.tensor(A.indices, idtype)
3566
        g = dgl.graph(("csr", (indptr, indices, [])))
3567
        with pytest.raises(DGLError):
3568
3569
            g2 = dgl.graph(("csr", (indptr[:-1], indices, [])), num_nodes=10)
    elif fmt == "csc":
3570
3571
3572
3573
        A = ssp.random(10, 10, 0.2).tocsc()
        A.data = np.arange(20)
        indptr = F.tensor(A.indptr, idtype)
        indices = F.tensor(A.indices, idtype)
3574
        g = dgl.graph(("csc", (indptr, indices, [])))
3575
        with pytest.raises(DGLError):
3576
            g2 = dgl.graph(("csr", (indptr[:-1], indices, [])), num_nodes=10)
3577
3578
3579
3580

    A_coo = A.tocoo()
    A_csr = A.tocsr()
    A_csc = A.tocsc()
3581
    row, col = g.adj_tensors("coo")
3582
3583
3584
    assert np.array_equal(F.asnumpy(row), A_coo.row)
    assert np.array_equal(F.asnumpy(col), A_coo.col)

3585
    indptr, indices, eids = g.adj_tensors("csr")
3586
    assert np.array_equal(F.asnumpy(indptr), A_csr.indptr)
3587
    if fmt == "csr":
3588
3589
3590
3591
3592
3593
3594
3595
3596
        assert len(eids) == 0
        assert np.array_equal(F.asnumpy(indices), A_csr.indices)
    else:
        indices_sorted = F.zeros(len(indices), idtype)
        indices_sorted = F.scatter_row(indices_sorted, eids, indices)
        indices_sorted_np = np.zeros(len(indices), dtype=A_csr.indices.dtype)
        indices_sorted_np[A_csr.data] = A_csr.indices
        assert np.array_equal(F.asnumpy(indices_sorted), indices_sorted_np)

3597
    indptr, indices, eids = g.adj_tensors("csc")
3598
    assert np.array_equal(F.asnumpy(indptr), A_csc.indptr)
3599
    if fmt == "csc":
3600
3601
3602
3603
3604
3605
3606
3607
3608
        assert len(eids) == 0
        assert np.array_equal(F.asnumpy(indices), A_csc.indices)
    else:
        indices_sorted = F.zeros(len(indices), idtype)
        indices_sorted = F.scatter_row(indices_sorted, eids, indices)
        indices_sorted_np = np.zeros(len(indices), dtype=A_csc.indices.dtype)
        indices_sorted_np[A_csc.data] = A_csc.indices
        assert np.array_equal(F.asnumpy(indices_sorted), indices_sorted_np)

3609

3610
3611
3612
def _test_forking_pickler_entry(g, q):
    q.put(g.formats())

3613
3614
3615
3616

@unittest.skipIf(
    dgl.backend.backend_name == "mxnet", reason="MXNet doesn't support spawning"
)
3617
def test_forking_pickler():
3618
3619
    ctx = mp.get_context("spawn")
    g = dgl.graph(([0, 1, 2], [1, 2, 3]))
3620
3621
3622
3623
    g.create_formats_()
    q = ctx.Queue(1)
    proc = ctx.Process(target=_test_forking_pickler_entry, args=(g, q))
    proc.start()
3624
    fmt = q.get()["created"]
3625
    proc.join()
3626
3627
3628
    assert "coo" in fmt
    assert "csr" in fmt
    assert "csc" in fmt
3629
3630


3631
if __name__ == "__main__":
3632
3633
3634
3635
3636
    # test_create()
    # test_query()
    # test_hypersparse()
    # test_adj("int32")
    # test_inc()
3637
    # test_view("int32")
3638
    # test_view1("int32")
3639
    # test_flatten(F.int32)
3640
3641
    # test_convert_bound()
    # test_convert()
3642
    # test_to_device("int32")
3643
    # test_transform("int32")
3644
3645
    # test_subgraph("int32")
    # test_subgraph_mask("int32")
3646
3647
3648
3649
3650
    # test_apply()
    # test_level1()
    # test_level2()
    # test_updates()
    # test_backward()
3651
    # test_empty_heterograph('int32')
3652
3653
3654
3655
    # test_types_in_function()
    # test_stack_reduce()
    # test_isolated_ntype()
    # test_bipartite()
3656
    # test_dtype_cast()
3657
    # test_float_cast()
3658
    # test_reverse("int32")
3659
    # test_format()
3660
3661
3662
3663
3664
3665
3666
3667
3668
    # test_add_edges(F.int32)
    # test_add_nodes(F.int32)
    # test_remove_edges(F.int32)
    # test_remove_nodes(F.int32)
    # test_clone(F.int32)
    # test_frame(F.int32)
    # test_frame_device(F.int32)
    # test_empty_query(F.int32)
    # test_create_block(F.int32)
3669
    pass