test_subgraph.py 2.2 KB
Newer Older
Minjie Wang's avatar
Minjie Wang committed
1
import numpy as np
2
import dgl
3
import backend as F
4
import unittest
Minjie Wang's avatar
Minjie Wang committed
5
6
7

D = 5

8
def generate_graph(grad=False, add_data=True):
9
    g = dgl.DGLGraph().to(F.ctx())
10
    g.add_nodes(10)
Minjie Wang's avatar
Minjie Wang committed
11
12
13
14
15
16
    # create a graph where 0 is the source and 9 is the sink
    for i in range(1, 9):
        g.add_edge(0, i)
        g.add_edge(i, 9)
    # add a back flow from 9 to 0
    g.add_edge(9, 0)
17
18
19
20
21
22
23
24
    if add_data:
        ncol = F.randn((10, D))
        ecol = F.randn((17, D))
        if grad:
            ncol = F.attach_grad(ncol)
            ecol = F.attach_grad(ecol)
        g.ndata['h'] = ncol
        g.edata['l'] = ecol
Minjie Wang's avatar
Minjie Wang committed
25
26
    return g

27
28
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
def test_edge_subgraph():
29
30
31
32
33
34
35
    # Test when the graph has no node data and edge data.
    g = generate_graph(add_data=False)
    eid = [0, 2, 3, 6, 7, 9]
    sg = g.edge_subgraph(eid)
    sg.ndata['h'] = F.arange(0, sg.number_of_nodes())
    sg.edata['h'] = F.arange(0, sg.number_of_edges())

36
def test_subgraph():
Minjie Wang's avatar
Minjie Wang committed
37
    g = generate_graph()
38
39
    h = g.ndata['h']
    l = g.edata['l']
Minjie Wang's avatar
Minjie Wang committed
40
41
    nid = [0, 2, 3, 6, 7, 9]
    sg = g.subgraph(nid)
42
    eid = {2, 3, 4, 5, 10, 11, 12, 13, 16}
43
44
    assert set(F.asnumpy(sg.edata[dgl.EID])) == eid
    eid = sg.edata[dgl.EID]
45
46
47
    # the subgraph is empty initially except for NID/EID field
    assert len(sg.ndata) == 2
    assert len(sg.edata) == 2
48
    sh = sg.ndata['h']
VoVAllen's avatar
VoVAllen committed
49
    assert F.allclose(F.gather_row(h, F.tensor(nid)), sh)
Minjie Wang's avatar
Minjie Wang committed
50
51
52
53
    '''
    s, d, eid
    0, 1, 0
    1, 9, 1
Minjie Wang's avatar
Minjie Wang committed
54
55
56
57
    0, 2, 2    1
    2, 9, 3    1
    0, 3, 4    1
    3, 9, 5    1
Minjie Wang's avatar
Minjie Wang committed
58
59
60
    0, 4, 6
    4, 9, 7
    0, 5, 8
Minjie Wang's avatar
Minjie Wang committed
61
62
63
64
65
    5, 9, 9       3
    0, 6, 10   1
    6, 9, 11   1  3
    0, 7, 12   1
    7, 9, 13   1  3
Minjie Wang's avatar
Minjie Wang committed
66
    0, 8, 14
Minjie Wang's avatar
Minjie Wang committed
67
68
    8, 9, 15      3
    9, 0, 16   1
Minjie Wang's avatar
Minjie Wang committed
69
    '''
70
    assert F.allclose(F.gather_row(l, eid), sg.edata['l'])
Minjie Wang's avatar
Minjie Wang committed
71
72
    # update the node/edge features on the subgraph should NOT
    # reflect to the parent graph.
73
74
    sg.ndata['h'] = F.zeros((6, D))
    assert F.allclose(h, g.ndata['h'])
Minjie Wang's avatar
Minjie Wang committed
75

76
77
def _test_map_to_subgraph():
    g = dgl.DGLGraph()
Quan (Andy) Gan's avatar
Quan (Andy) Gan committed
78
79
80
81
82
    g.add_nodes(10)
    g.add_edges(F.arange(0, 9), F.arange(1, 10))
    h = g.subgraph([0, 1, 2, 5, 8])
    v = h.map_to_subgraph_nid([0, 8, 2])
    assert np.array_equal(F.asnumpy(v), np.array([0, 4, 2]))