test_graphbolt_utils.py 17 KB
Newer Older
1
import backend as F
2
import dgl.graphbolt as gb
3
import pytest
4
5
6
import torch


7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
def test_find_reverse_edges_homo():
    edges = (torch.tensor([1, 3, 5]), torch.tensor([2, 4, 5]))
    edges = gb.add_reverse_edges(edges)
    expected_edges = (
        torch.tensor([1, 3, 5, 2, 4, 5]),
        torch.tensor([2, 4, 5, 1, 3, 5]),
    )
    assert torch.equal(edges[0], expected_edges[0])
    assert torch.equal(edges[1], expected_edges[1])


def test_find_reverse_edges_hetero():
    edges = {
        "A:r:B": (torch.tensor([1, 5]), torch.tensor([2, 5])),
        "B:rr:A": (torch.tensor([3]), torch.tensor([3])),
    }
    edges = gb.add_reverse_edges(edges, {"A:r:B": "B:rr:A"})
    expected_edges = {
        "A:r:B": (torch.tensor([1, 5]), torch.tensor([2, 5])),
        "B:rr:A": (torch.tensor([3, 2, 5]), torch.tensor([3, 1, 5])),
    }
    assert torch.equal(edges["A:r:B"][0], expected_edges["A:r:B"][0])
    assert torch.equal(edges["A:r:B"][1], expected_edges["A:r:B"][1])
    assert torch.equal(edges["B:rr:A"][0], expected_edges["B:rr:A"][0])
    assert torch.equal(edges["B:rr:A"][1], expected_edges["B:rr:A"][1])


def test_find_reverse_edges_bi_reverse_types():
    edges = {
        "A:r:B": (torch.tensor([1, 5]), torch.tensor([2, 5])),
        "B:rr:A": (torch.tensor([3]), torch.tensor([3])),
    }
    edges = gb.add_reverse_edges(edges, {"A:r:B": "B:rr:A", "B:rr:A": "A:r:B"})
    expected_edges = {
        "A:r:B": (torch.tensor([1, 5, 3]), torch.tensor([2, 5, 3])),
        "B:rr:A": (torch.tensor([3, 2, 5]), torch.tensor([3, 1, 5])),
    }
    assert torch.equal(edges["A:r:B"][0], expected_edges["A:r:B"][0])
    assert torch.equal(edges["A:r:B"][1], expected_edges["A:r:B"][1])
    assert torch.equal(edges["B:rr:A"][0], expected_edges["B:rr:A"][0])
    assert torch.equal(edges["B:rr:A"][1], expected_edges["B:rr:A"][1])


def test_find_reverse_edges_circual_reverse_types():
    edges = {
        "A:r1:B": (torch.tensor([1]), torch.tensor([1])),
        "B:r2:C": (torch.tensor([2]), torch.tensor([2])),
        "C:r3:A": (torch.tensor([3]), torch.tensor([3])),
    }
    edges = gb.add_reverse_edges(
        edges, {"A:r1:B": "B:r2:C", "B:r2:C": "C:r3:A", "C:r3:A": "A:r1:B"}
    )
    expected_edges = {
        "A:r1:B": (torch.tensor([1, 3]), torch.tensor([1, 3])),
        "B:r2:C": (torch.tensor([2, 1]), torch.tensor([2, 1])),
        "C:r3:A": (torch.tensor([3, 2]), torch.tensor([3, 2])),
    }
    assert torch.equal(edges["A:r1:B"][0], expected_edges["A:r1:B"][0])
    assert torch.equal(edges["A:r1:B"][1], expected_edges["A:r1:B"][1])
    assert torch.equal(edges["B:r2:C"][0], expected_edges["B:r2:C"][0])
    assert torch.equal(edges["B:r2:C"][1], expected_edges["B:r2:C"][1])
    assert torch.equal(edges["A:r1:B"][0], expected_edges["A:r1:B"][0])
    assert torch.equal(edges["A:r1:B"][1], expected_edges["A:r1:B"][1])
    assert torch.equal(edges["C:r3:A"][0], expected_edges["C:r3:A"][0])
    assert torch.equal(edges["C:r3:A"][1], expected_edges["C:r3:A"][1])


74
def test_unique_and_compact_hetero():
75
76
77
78
79
    N1 = torch.tensor(
        [0, 5, 2, 7, 12, 7, 9, 5, 6, 2, 3, 4, 1, 0, 9], device=F.ctx()
    )
    N2 = torch.tensor([0, 3, 3, 5, 2, 7, 2, 8, 4, 9, 2, 3], device=F.ctx())
    N3 = torch.tensor([1, 2, 6, 6, 1, 8, 3, 6, 3, 2], device=F.ctx())
80
    expected_unique = {
81
82
83
        "n1": torch.tensor([0, 5, 2, 7, 12, 9, 6, 3, 4, 1], device=F.ctx()),
        "n2": torch.tensor([0, 3, 5, 2, 7, 8, 4, 9], device=F.ctx()),
        "n3": torch.tensor([1, 2, 6, 8, 3], device=F.ctx()),
84
    }
85
86
87
88
89
90
91
92
93
94
    if N1.is_cuda:
        expected_reverse_id = {
            k: v.sort()[1] for k, v in expected_unique.items()
        }
        expected_unique = {k: v.sort()[0] for k, v in expected_unique.items()}
    else:
        expected_reverse_id = {
            k: torch.arange(0, v.shape[0], device=F.ctx())
            for k, v in expected_unique.items()
        }
95
96
97
98
99
    nodes_dict = {
        "n1": N1.split(5),
        "n2": N2.split(4),
        "n3": N3.split(2),
    }
100
101
    expected_nodes_dict = {
        "n1": [
102
103
104
            torch.tensor([0, 1, 2, 3, 4], device=F.ctx()),
            torch.tensor([3, 5, 1, 6, 2], device=F.ctx()),
            torch.tensor([7, 8, 9, 0, 5], device=F.ctx()),
105
106
        ],
        "n2": [
107
108
109
            torch.tensor([0, 1, 1, 2], device=F.ctx()),
            torch.tensor([3, 4, 3, 5], device=F.ctx()),
            torch.tensor([6, 7, 3, 1], device=F.ctx()),
110
111
        ],
        "n3": [
112
113
114
115
116
            torch.tensor([0, 1], device=F.ctx()),
            torch.tensor([2, 2], device=F.ctx()),
            torch.tensor([0, 3], device=F.ctx()),
            torch.tensor([4, 2], device=F.ctx()),
            torch.tensor([4, 1], device=F.ctx()),
117
118
        ],
    }
119
120
121
122

    unique, compacted = gb.unique_and_compact(nodes_dict)
    for ntype, nodes in unique.items():
        expected_nodes = expected_unique[ntype]
123
        assert torch.equal(nodes, expected_nodes)
124
125

    for ntype, nodes in compacted.items():
126
        expected_nodes = expected_nodes_dict[ntype]
127
        assert isinstance(nodes, list)
128
        for expected_node, node in zip(expected_nodes, nodes):
129
            node = expected_reverse_id[ntype][node]
130
131
132
133
            assert torch.equal(expected_node, node)


def test_unique_and_compact_homo():
134
135
136
137
138
139
140
141
142
143
144
145
146
    N = torch.tensor(
        [0, 5, 2, 7, 12, 7, 9, 5, 6, 2, 3, 4, 1, 0, 9], device=F.ctx()
    )
    expected_unique_N = torch.tensor(
        [0, 5, 2, 7, 12, 9, 6, 3, 4, 1], device=F.ctx()
    )
    if N.is_cuda:
        expected_reverse_id_N = expected_unique_N.sort()[1]
        expected_unique_N = expected_unique_N.sort()[0]
    else:
        expected_reverse_id_N = torch.arange(
            0, expected_unique_N.shape[0], device=F.ctx()
        )
147
    nodes_list = N.split(5)
148
    expected_nodes_list = [
149
150
151
        torch.tensor([0, 1, 2, 3, 4], device=F.ctx()),
        torch.tensor([3, 5, 1, 6, 2], device=F.ctx()),
        torch.tensor([7, 8, 9, 0, 5], device=F.ctx()),
152
    ]
153
154
155

    unique, compacted = gb.unique_and_compact(nodes_list)

156
    assert torch.equal(unique, expected_unique_N)
157
    assert isinstance(compacted, list)
158
    for expected_node, node in zip(expected_nodes_list, compacted):
159
        node = expected_reverse_id_N[node]
160
161
162
        assert torch.equal(expected_node, node)


163
def test_unique_and_compact_node_pairs_hetero():
164
165
    node_pairs = {
        "n1:e1:n2": (
166
167
            torch.tensor([1, 3, 4, 6, 2, 7, 9, 4, 2, 6], device=F.ctx()),
            torch.tensor([2, 2, 2, 4, 1, 1, 1, 3, 3, 3], device=F.ctx()),
168
169
        ),
        "n1:e2:n3": (
170
171
            torch.tensor([5, 2, 6, 4, 7, 2, 8, 1, 3, 0], device=F.ctx()),
            torch.tensor([1, 3, 3, 3, 2, 2, 2, 7, 7, 7], device=F.ctx()),
172
173
        ),
        "n2:e3:n3": (
174
175
            torch.tensor([2, 5, 4, 1, 4, 3, 6, 0], device=F.ctx()),
            torch.tensor([1, 1, 3, 3, 2, 2, 7, 7], device=F.ctx()),
176
177
178
        ),
    }

179
    expected_unique_nodes = {
180
181
182
        "n1": torch.tensor([1, 3, 4, 6, 2, 7, 9, 5, 8, 0], device=F.ctx()),
        "n2": torch.tensor([1, 2, 3, 4, 5, 6, 0], device=F.ctx()),
        "n3": torch.tensor([1, 2, 3, 7], device=F.ctx()),
183
    }
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
    if expected_unique_nodes["n1"].is_cuda:
        expected_reverse_id = {
            "n1": expected_unique_nodes["n1"].sort()[1],
            "n2": torch.tensor([0, 1, 2, 3, 6, 4, 5], device=F.ctx()),
            "n3": expected_unique_nodes["n3"].sort()[1],
        }
        expected_unique_nodes = {
            "n1": expected_unique_nodes["n1"].sort()[0],
            "n2": torch.tensor([1, 2, 3, 4, 0, 5, 6], device=F.ctx()),
            "n3": expected_unique_nodes["n3"].sort()[0],
        }
    else:
        expected_reverse_id = {
            k: torch.arange(0, v.shape[0], device=F.ctx())
            for k, v in expected_unique_nodes.items()
        }
200
    expected_node_pairs = {
201
        "n1:e1:n2": (
202
203
            torch.tensor([0, 1, 2, 3, 4, 5, 6, 2, 4, 3], device=F.ctx()),
            torch.tensor([1, 1, 1, 3, 0, 0, 0, 2, 2, 2], device=F.ctx()),
204
        ),
205
        "n1:e2:n3": (
206
207
            torch.tensor([7, 4, 3, 2, 5, 4, 8, 0, 1, 9], device=F.ctx()),
            torch.tensor([0, 2, 2, 2, 1, 1, 1, 3, 3, 3], device=F.ctx()),
208
        ),
209
        "n2:e3:n3": (
210
211
            torch.tensor([1, 4, 3, 0, 3, 2, 5, 6], device=F.ctx()),
            torch.tensor([0, 0, 2, 2, 1, 1, 3, 3], device=F.ctx()),
212
213
214
215
216
217
        ),
    }

    unique_nodes, compacted_node_pairs = gb.unique_and_compact_node_pairs(
        node_pairs
    )
218
219
    for ntype, nodes in unique_nodes.items():
        expected_nodes = expected_unique_nodes[ntype]
220
        assert torch.equal(nodes, expected_nodes)
221
222
    for etype, pair in compacted_node_pairs.items():
        u, v = pair
223
224
225
        ntype1, _, ntype2 = etype.split(":")
        u = expected_reverse_id[ntype1][u]
        v = expected_reverse_id[ntype2][v]
226
        expected_u, expected_v = expected_node_pairs[etype]
227
228
229
230
231
        assert torch.equal(u, expected_u)
        assert torch.equal(v, expected_v)


def test_unique_and_compact_node_pairs_homo():
232
233
234
    dst_nodes = torch.tensor([1, 1, 3, 3, 5, 5, 2, 6, 6, 6, 6], device=F.ctx())
    src_nodes = torch.tensor([2, 3, 1, 4, 5, 2, 5, 1, 4, 4, 6], device=F.ctx())
    node_pairs = (src_nodes, dst_nodes)
235

236
237
238
239
240
241
242
    expected_unique_nodes = torch.tensor([1, 2, 3, 5, 6, 4], device=F.ctx())
    expected_dst_nodes = torch.tensor(
        [0, 0, 2, 2, 3, 3, 1, 4, 4, 4, 4], device=F.ctx()
    )
    expected_src_ndoes = torch.tensor(
        [1, 2, 0, 5, 3, 1, 3, 0, 5, 5, 4], device=F.ctx()
    )
243
244
245
    unique_nodes, compacted_node_pairs = gb.unique_and_compact_node_pairs(
        node_pairs
    )
246
    assert torch.equal(unique_nodes, expected_unique_nodes)
247

248
    u, v = compacted_node_pairs
249
250
    assert torch.equal(u, expected_src_ndoes)
    assert torch.equal(v, expected_dst_nodes)
251
252
253
    assert torch.equal(
        unique_nodes[:5], torch.tensor([1, 2, 3, 5, 6], device=F.ctx())
    )
254
255
256


def test_incomplete_unique_dst_nodes_():
257
258
259
260
261
    node_pairs = (
        torch.arange(0, 50, device=F.ctx()),
        torch.arange(100, 150, device=F.ctx()),
    )
    unique_dst_nodes = torch.arange(150, 200, device=F.ctx())
262
263
    with pytest.raises(IndexError):
        gb.unique_and_compact_node_pairs(node_pairs, unique_dst_nodes)
264
265


266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
def test_unique_and_compact_csc_formats_hetero():
    dst_nodes = {
        "n2": torch.tensor([2, 4, 1, 3]),
        "n3": torch.tensor([1, 3, 2, 7]),
    }
    csc_formats = {
        "n1:e1:n2": gb.CSCFormatBase(
            indptr=torch.tensor([0, 3, 4, 7, 10]),
            indices=torch.tensor([1, 3, 4, 6, 2, 7, 9, 4, 2, 6]),
        ),
        "n1:e2:n3": gb.CSCFormatBase(
            indptr=torch.tensor([0, 1, 4, 7, 10]),
            indices=torch.tensor([5, 2, 6, 4, 7, 2, 8, 1, 3, 0]),
        ),
        "n2:e3:n3": gb.CSCFormatBase(
            indptr=torch.tensor([0, 2, 4, 6, 8]),
            indices=torch.tensor([2, 5, 4, 1, 4, 3, 6, 0]),
        ),
    }

    expected_unique_nodes = {
        "n1": torch.tensor([1, 3, 4, 6, 2, 7, 9, 5, 8, 0]),
        "n2": torch.tensor([2, 4, 1, 3, 5, 6, 0]),
        "n3": torch.tensor([1, 3, 2, 7]),
    }
    expected_csc_formats = {
        "n1:e1:n2": gb.CSCFormatBase(
            indptr=torch.tensor([0, 3, 4, 7, 10]),
            indices=torch.tensor([0, 1, 2, 3, 4, 5, 6, 2, 4, 3]),
        ),
        "n1:e2:n3": gb.CSCFormatBase(
            indptr=torch.tensor([0, 1, 4, 7, 10]),
            indices=torch.tensor([7, 4, 3, 2, 5, 4, 8, 0, 1, 9]),
        ),
        "n2:e3:n3": gb.CSCFormatBase(
            indptr=torch.tensor([0, 2, 4, 6, 8]),
            indices=torch.tensor([0, 4, 1, 2, 1, 3, 5, 6]),
        ),
    }

    unique_nodes, compacted_csc_formats = gb.unique_and_compact_csc_formats(
        csc_formats, dst_nodes
    )

    for ntype, nodes in unique_nodes.items():
        expected_nodes = expected_unique_nodes[ntype]
        assert torch.equal(nodes, expected_nodes)
    for etype, pair in compacted_csc_formats.items():
        indices = pair.indices
        indptr = pair.indptr
        expected_indices = expected_csc_formats[etype].indices
        expected_indptr = expected_csc_formats[etype].indptr
        assert torch.equal(indices, expected_indices)
        assert torch.equal(indptr, expected_indptr)


def test_unique_and_compact_csc_formats_homo():
    seeds = torch.tensor([1, 3, 5, 2, 6])
324
    indptr = torch.tensor([0, 2, 4, 6, 7, 11])
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
    indices = torch.tensor([2, 3, 1, 4, 5, 2, 5, 1, 4, 4, 6])
    csc_formats = gb.CSCFormatBase(indptr=indptr, indices=indices)

    expected_unique_nodes = torch.tensor([1, 3, 5, 2, 6, 4])
    expected_indptr = indptr
    expected_indices = torch.tensor([3, 1, 0, 5, 2, 3, 2, 0, 5, 5, 4])

    unique_nodes, compacted_csc_formats = gb.unique_and_compact_csc_formats(
        csc_formats, seeds
    )

    indptr = compacted_csc_formats.indptr
    indices = compacted_csc_formats.indices
    assert torch.equal(indptr, expected_indptr)
    assert torch.equal(indices, expected_indices)
    assert torch.equal(unique_nodes, expected_unique_nodes)


343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
def test_unique_and_compact_incorrect_indptr():
    seeds = torch.tensor([1, 3, 5, 2, 6, 7])
    indptr = torch.tensor([0, 2, 4, 6, 7, 11])
    indices = torch.tensor([2, 3, 1, 4, 5, 2, 5, 1, 4, 4, 6])
    csc_formats = gb.CSCFormatBase(indptr=indptr, indices=indices)

    # The number of seeds is not corresponding to indptr.
    with pytest.raises(AssertionError):
        gb.unique_and_compact_csc_formats(csc_formats, seeds)

    seeds = torch.tensor([1, 3, 5, 2, 6])
    indptr = torch.tensor([0, 2, 4, 6, 7, 11])
    indices = torch.tensor([2, 3, 1, 4, 5, 2, 5, 1, 4, 4])
    csc_formats = gb.CSCFormatBase(indptr=indptr, indices=indices)
    # The value of last element in indptr is not corresponding to indices.
    with pytest.raises(AssertionError):
        gb.unique_and_compact_csc_formats(csc_formats, seeds)


362
def test_compact_csc_format_hetero():
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
    dst_nodes = {
        "n2": torch.tensor([2, 4, 1, 3]),
        "n3": torch.tensor([1, 3, 2, 7]),
    }
    csc_formats = {
        "n1:e1:n2": gb.CSCFormatBase(
            indptr=torch.tensor([0, 3, 4, 7, 10]),
            indices=torch.tensor([1, 3, 4, 6, 2, 7, 9, 4, 2, 6]),
        ),
        "n1:e2:n3": gb.CSCFormatBase(
            indptr=torch.tensor([0, 1, 4, 7, 10]),
            indices=torch.tensor([5, 2, 6, 4, 7, 2, 8, 1, 3, 0]),
        ),
        "n2:e3:n3": gb.CSCFormatBase(
            indptr=torch.tensor([0, 2, 4, 6, 8]),
            indices=torch.tensor([2, 5, 4, 1, 4, 3, 6, 0]),
        ),
    }
381
382

    expected_original_row_ids = {
383
384
385
386
387
        "n1": torch.tensor(
            [1, 3, 4, 6, 2, 7, 9, 4, 2, 6, 5, 2, 6, 4, 7, 2, 8, 1, 3, 0]
        ),
        "n2": torch.tensor([2, 4, 1, 3, 2, 5, 4, 1, 4, 3, 6, 0]),
        "n3": torch.tensor([1, 3, 2, 7]),
388
    }
389
    expected_csc_formats = {
390
        "n1:e1:n2": gb.CSCFormatBase(
391
392
            indptr=torch.tensor([0, 3, 4, 7, 10]),
            indices=torch.arange(0, 10),
393
394
        ),
        "n1:e2:n3": gb.CSCFormatBase(
395
396
            indptr=torch.tensor([0, 1, 4, 7, 10]),
            indices=torch.arange(0, 10) + 10,
397
398
        ),
        "n2:e3:n3": gb.CSCFormatBase(
399
400
            indptr=torch.tensor([0, 2, 4, 6, 8]),
            indices=torch.arange(0, 8) + 4,
401
402
403
404
405
406
407
408
409
410
411
412
        ),
    }
    original_row_ids, compacted_csc_formats = gb.compact_csc_format(
        csc_formats, dst_nodes
    )

    for ntype, nodes in original_row_ids.items():
        expected_nodes = expected_original_row_ids[ntype]
        assert torch.equal(nodes, expected_nodes)
    for etype, csc_format in compacted_csc_formats.items():
        indptr = csc_format.indptr
        indices = csc_format.indices
413
414
        expected_indptr = expected_csc_formats[etype].indptr
        expected_indices = expected_csc_formats[etype].indices
415
416
417
418
419
        assert torch.equal(indptr, expected_indptr)
        assert torch.equal(indices, expected_indices)


def test_compact_csc_format_homo():
420
421
422
423
    seeds = torch.tensor([1, 3, 5, 2, 6])
    indptr = torch.tensor([0, 2, 4, 6, 7, 11])
    indices = torch.tensor([2, 3, 1, 4, 5, 2, 5, 1, 4, 4, 6])
    csc_formats = gb.CSCFormatBase(indptr=indptr, indices=indices)
424

425
426
    expected_original_row_ids = torch.tensor(
        [1, 3, 5, 2, 6, 2, 3, 1, 4, 5, 2, 5, 1, 4, 4, 6]
427
    )
428
429
430
    expected_indptr = indptr
    expected_indices = torch.arange(0, len(indices)) + 5

431
    original_row_ids, compacted_csc_formats = gb.compact_csc_format(
432
        csc_formats, seeds
433
434
435
    )

    indptr = compacted_csc_formats.indptr
436
437
    indices = compacted_csc_formats.indices

438
439
440
    assert torch.equal(indptr, expected_indptr)
    assert torch.equal(indices, expected_indices)
    assert torch.equal(original_row_ids, expected_original_row_ids)
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459


def test_compact_incorrect_indptr():
    seeds = torch.tensor([1, 3, 5, 2, 6, 7])
    indptr = torch.tensor([0, 2, 4, 6, 7, 11])
    indices = torch.tensor([2, 3, 1, 4, 5, 2, 5, 1, 4, 4, 6])
    csc_formats = gb.CSCFormatBase(indptr=indptr, indices=indices)

    # The number of seeds is not corresponding to indptr.
    with pytest.raises(AssertionError):
        gb.compact_csc_format(csc_formats, seeds)

    seeds = torch.tensor([1, 3, 5, 2, 6])
    indptr = torch.tensor([0, 2, 4, 6, 7, 11])
    indices = torch.tensor([2, 3, 1, 4, 5, 2, 5, 1, 4, 4])
    csc_formats = gb.CSCFormatBase(indptr=indptr, indices=indices)
    # The value of last element in indptr is not corresponding to indices.
    with pytest.raises(AssertionError):
        gb.compact_csc_format(csc_formats, seeds)