Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
test_function.py 5.19 KB
Newer Older
Minjie Wang's avatar
Builtin function and API changes (#53)  
Minjie Wang committed
1
2
3
4
5
6
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
74
75
76
77
78
79
80
81
82
83
84
85
86
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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141 
import dgl
import dgl.function as fn
from dgl.graph import __REPR__

def generate_graph():
    g = dgl.DGLGraph()
    for i in range(10):
        g.add_node(i, h=i+1) # 10 nodes.
    # create a graph where 0 is the source and 9 is the sink
    for i in range(1, 9):
        g.add_edge(0, i, h=1)
        g.add_edge(i, 9, h=i+1)
    # add a back flow from 9 to 0
    g.add_edge(9, 0, h=10)
    return g

def check(g, h, fld):
    nh = [str(g.nodes[i][fld]) for i in range(10)]
    h = [str(x) for x in h]
    assert nh == h, "nh=[%s], h=[%s]" % (' '.join(nh), ' '.join(h))

def generate_graph1():
    """graph with anonymous repr"""
    g = dgl.DGLGraph()
    for i in range(10):
        g.add_node(i, __REPR__=i+1) # 10 nodes.
    # create a graph where 0 is the source and 9 is the sink
    for i in range(1, 9):
        g.add_edge(0, i, __REPR__=1)
        g.add_edge(i, 9, __REPR__=i+1)
    # add a back flow from 9 to 0
    g.add_edge(9, 0, __REPR__=10)
    return g

def test_copy_src():
    # copy_src with both fields
    g = generate_graph()
    g.register_message_func(fn.copy_src(src='h', out='m'), batchable=False)
    g.register_reduce_func(fn.sum(msgs='m', out='h'), batchable=False)
    check(g, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 'h')
    g.update_all()
    check(g, [10, 1, 1, 1, 1, 1, 1, 1, 1, 44], 'h')

    # copy_src with only src field; the out field should use anonymous repr
    g = generate_graph()
    g.register_message_func(fn.copy_src(src='h'), batchable=False)
    g.register_reduce_func(fn.sum(out='h'), batchable=False)
    check(g, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 'h')
    g.update_all()
    check(g, [10, 1, 1, 1, 1, 1, 1, 1, 1, 44], 'h')

    # copy_src with no src field; should use anonymous repr
    g = generate_graph1()
    g.register_message_func(fn.copy_src(out='m'), batchable=False)
    g.register_reduce_func(fn.sum(msgs='m', out='h'), batchable=False)
    check(g, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], __REPR__)
    g.update_all()
    check(g, [10, 1, 1, 1, 1, 1, 1, 1, 1, 44], 'h')

    # copy src with no fields;
    g = generate_graph1()
    g.register_message_func(fn.copy_src(), batchable=False)
    g.register_reduce_func(fn.sum(out='h'), batchable=False)
    check(g, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], __REPR__)
    g.update_all()
    check(g, [10, 1, 1, 1, 1, 1, 1, 1, 1, 44], 'h')

def test_copy_edge():
    # copy_edge with both fields
    g = generate_graph()
    g.register_message_func(fn.copy_edge(edge='h', out='m'), batchable=False)
    g.register_reduce_func(fn.sum(msgs='m', out='h'), batchable=False)
    check(g, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 'h')
    g.update_all()
    check(g, [10, 1, 1, 1, 1, 1, 1, 1, 1, 44], 'h')

    # copy_edge with only edge field; the out field should use anonymous repr
    g = generate_graph()
    g.register_message_func(fn.copy_edge(edge='h'), batchable=False)
    g.register_reduce_func(fn.sum(out='h'), batchable=False)
    check(g, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 'h')
    g.update_all()
    check(g, [10, 1, 1, 1, 1, 1, 1, 1, 1, 44], 'h')

    # copy_edge with no edge field; should use anonymous repr
    g = generate_graph1()
    g.register_message_func(fn.copy_edge(out='m'), batchable=False)
    g.register_reduce_func(fn.sum(msgs='m', out='h'), batchable=False)
    check(g, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], __REPR__)
    g.update_all()
    check(g, [10, 1, 1, 1, 1, 1, 1, 1, 1, 44], 'h')

    # copy edge with no fields;
    g = generate_graph1()
    g.register_message_func(fn.copy_edge(), batchable=False)
    g.register_reduce_func(fn.sum(out='h'), batchable=False)
    check(g, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], __REPR__)
    g.update_all()
    check(g, [10, 1, 1, 1, 1, 1, 1, 1, 1, 44], 'h')

def test_src_mul_edge():
    # src_mul_edge with all fields
    g = generate_graph()
    g.register_message_func(fn.src_mul_edge(src='h', edge='h', out='m'), batchable=False)
    g.register_reduce_func(fn.sum(msgs='m', out='h'), batchable=False)
    check(g, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 'h')
    g.update_all()
    check(g, [100, 1, 1, 1, 1, 1, 1, 1, 1, 284], 'h')

    g = generate_graph()
    g.register_message_func(fn.src_mul_edge(src='h', edge='h'), batchable=False)
    g.register_reduce_func(fn.sum(out='h'), batchable=False)
    check(g, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 'h')
    g.update_all()
    check(g, [100, 1, 1, 1, 1, 1, 1, 1, 1, 284], 'h')

    g = generate_graph1()
    g.register_message_func(fn.src_mul_edge(out='m'), batchable=False)
    g.register_reduce_func(fn.sum(msgs='m', out='h'), batchable=False)
    check(g, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], __REPR__)
    g.update_all()
    check(g, [100, 1, 1, 1, 1, 1, 1, 1, 1, 284], 'h')

    g = generate_graph1()
    g.register_message_func(fn.src_mul_edge(), batchable=False)
    g.register_reduce_func(fn.sum(out='h'), batchable=False)
    check(g, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], __REPR__)
    g.update_all()
    check(g, [100, 1, 1, 1, 1, 1, 1, 1, 1, 284], 'h')

    g = generate_graph1()
    g.register_message_func(fn.src_mul_edge(), batchable=False)
    g.register_reduce_func(fn.sum(), batchable=False)
    check(g, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], __REPR__)
    g.update_all()
    check(g, [100, 1, 1, 1, 1, 1, 1, 1, 1, 284], __REPR__)

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