test_function.py 1.6 KB
Newer Older
1
2
3
import torch as th
import dgl
import dgl.function as fn
4
import utils as U
5
6
7

def generate_graph():
    g = dgl.DGLGraph()
Minjie Wang's avatar
Minjie Wang committed
8
    g.add_nodes(10) # 10 nodes.
Lingfan Yu's avatar
Lingfan Yu committed
9
    h = th.arange(1, 11, dtype=th.float)
10
    g.ndata['h'] = h
11
12
13
14
15
16
17
18
    # 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)
    h = th.tensor([1., 2., 1., 3., 1., 4., 1., 5., 1., 6.,\
            1., 7., 1., 8., 1., 9., 10.])
19
    g.edata['h'] = h
20
21
    return g

22
23
def reducer_both(nodes):
    return {'h' : th.sum(nodes.mailbox['m'], 1)}
24
25
26
27

def test_copy_src():
    # copy_src with both fields
    g = generate_graph()
Minjie Wang's avatar
Minjie Wang committed
28
29
    g.register_message_func(fn.copy_src(src='h', out='m'))
    g.register_reduce_func(reducer_both)
30
    g.update_all()
31
    assert U.allclose(g.ndata['h'],
32
33
34
35
36
            th.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))

def test_copy_edge():
    # copy_edge with both fields
    g = generate_graph()
Minjie Wang's avatar
Minjie Wang committed
37
38
    g.register_message_func(fn.copy_edge(edge='h', out='m'))
    g.register_reduce_func(reducer_both)
39
    g.update_all()
40
    assert U.allclose(g.ndata['h'],
41
42
43
44
45
            th.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))

def test_src_mul_edge():
    # src_mul_edge with all fields
    g = generate_graph()
Minjie Wang's avatar
Minjie Wang committed
46
47
    g.register_message_func(fn.src_mul_edge(src='h', edge='h', out='m'))
    g.register_reduce_func(reducer_both)
48
    g.update_all()
49
    assert U.allclose(g.ndata['h'],
50
51
52
53
54
55
            th.tensor([100., 1., 1., 1., 1., 1., 1., 1., 1., 284.]))

if __name__ == '__main__':
    test_copy_src()
    test_copy_edge()
    test_src_mul_edge()