test_frame.py 7.52 KB
Newer Older
1
2
3
import torch as th
from torch.autograd import Variable
import numpy as np
Minjie Wang's avatar
Minjie Wang committed
4
from dgl.frame import Frame, FrameRef
Minjie Wang's avatar
Minjie Wang committed
5
from dgl.utils import Index, toindex
6
7

N = 10
Minjie Wang's avatar
Minjie Wang committed
8
D = 5
9

Minjie Wang's avatar
Minjie Wang committed
10
11
12
13
14
15
def check_fail(fn):
    try:
        fn()
        return False
    except:
        return True
16
17
18
19
20
21
22
23
24
25
26

def create_test_data(grad=False):
    c1 = Variable(th.randn(N, D), requires_grad=grad)
    c2 = Variable(th.randn(N, D), requires_grad=grad)
    c3 = Variable(th.randn(N, D), requires_grad=grad)
    return {'a1' : c1, 'a2' : c2, 'a3' : c3}

def test_create():
    data = create_test_data()
    f1 = Frame()
    for k, v in data.items():
Minjie Wang's avatar
Minjie Wang committed
27
28
29
        f1.update_column(k, v)
    print(f1.schemes)
    assert f1.keys() == set(data.keys())
30
31
32
    assert f1.num_columns == 3
    assert f1.num_rows == N
    f2 = Frame(data)
Minjie Wang's avatar
Minjie Wang committed
33
    assert f2.keys() == set(data.keys())
34
35
36
37
38
39
    assert f2.num_columns == 3
    assert f2.num_rows == N
    f1.clear()
    assert len(f1.schemes) == 0
    assert f1.num_rows == 0

Minjie Wang's avatar
Minjie Wang committed
40
41
def test_column1():
    # Test frame column getter/setter
42
43
    data = create_test_data()
    f = Frame(data)
Minjie Wang's avatar
Minjie Wang committed
44
45
    assert f.num_rows == N
    assert len(f) == 3
Minjie Wang's avatar
Minjie Wang committed
46
    assert th.allclose(f['a1'].data, data['a1'].data)
47
    f['a1'] = data['a2']
Minjie Wang's avatar
Minjie Wang committed
48
    assert th.allclose(f['a2'].data, data['a2'].data)
Minjie Wang's avatar
Minjie Wang committed
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
    # add a different length column should fail
    def failed_add_col():
        f['a4'] = th.zeros([N+1, D])
    assert check_fail(failed_add_col)
    # delete all the columns
    del f['a1']
    del f['a2']
    assert len(f) == 1
    del f['a3']
    assert f.num_rows == 0
    assert len(f) == 0
    # add a different length column should succeed
    f['a4'] = th.zeros([N+1, D])
    assert f.num_rows == N+1
    assert len(f) == 1

def test_column2():
    # Test frameref column getter/setter
    data = Frame(create_test_data())
    f = FrameRef(data, [3, 4, 5, 6, 7])
    assert f.num_rows == 5
    assert len(f) == 3
Minjie Wang's avatar
Minjie Wang committed
71
    assert th.allclose(f['a1'], data['a1'].data[3:8])
Minjie Wang's avatar
Minjie Wang committed
72
73
    # set column should reflect on the referenced data
    f['a1'] = th.zeros([5, D])
Minjie Wang's avatar
Minjie Wang committed
74
75
76
77
78
79
    assert th.allclose(data['a1'].data[3:8], th.zeros([5, D]))
    # add new partial column should fail with error initializer
    f.set_initializer(lambda shape, dtype : assert_(False))
    def failed_add_col():
        f['a4'] = th.ones([5, D])
    assert check_fail(failed_add_col)
80

Minjie Wang's avatar
Minjie Wang committed
81
82
def test_append1():
    # test append API on Frame
83
    data = create_test_data()
Minjie Wang's avatar
Minjie Wang committed
84
85
86
87
88
89
90
    f1 = Frame()
    f2 = Frame(data)
    f1.append(data)
    assert f1.num_rows == N
    f1.append(f2)
    assert f1.num_rows == 2 * N
    c1 = f1['a1']
Minjie Wang's avatar
Minjie Wang committed
91
    assert c1.data.shape == (2 * N, D)
Minjie Wang's avatar
Minjie Wang committed
92
    truth = th.cat([data['a1'], data['a1']])
Minjie Wang's avatar
Minjie Wang committed
93
94
95
96
97
98
    assert th.allclose(truth, c1.data)
    # append dict of different length columns should fail
    f3 = {'a1' : th.zeros((3, D)), 'a2' : th.zeros((3, D)), 'a3' : th.zeros((2, D))}
    def failed_append():
        f1.append(f3)
    assert check_fail(failed_append)
Minjie Wang's avatar
Minjie Wang committed
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117

def test_append2():
    # test append on FrameRef
    data = Frame(create_test_data())
    f = FrameRef(data)
    assert f.is_contiguous()
    assert f.is_span_whole_column()
    assert f.num_rows == N
    # append on the underlying frame should not reflect on the ref
    data.append(data)
    assert f.is_contiguous()
    assert not f.is_span_whole_column()
    assert f.num_rows == N
    # append on the FrameRef should work
    f.append(data)
    assert not f.is_contiguous()
    assert not f.is_span_whole_column()
    assert f.num_rows == 3 * N
    new_idx = list(range(N)) + list(range(2*N, 4*N))
Minjie Wang's avatar
Minjie Wang committed
118
    assert th.all(f.index().tousertensor() == th.tensor(new_idx, dtype=th.int64))
Minjie Wang's avatar
Minjie Wang committed
119
120
121
122
123
124
    assert data.num_rows == 4 * N

def test_row1():
    # test row getter/setter
    data = create_test_data()
    f = FrameRef(Frame(data))
125
126
127

    # getter
    # test non-duplicate keys
Minjie Wang's avatar
Minjie Wang committed
128
    rowid = Index(th.tensor([0, 2]))
129
    rows = f[rowid]
Minjie Wang's avatar
Minjie Wang committed
130
    for k, v in rows.items():
131
        assert v.shape == (len(rowid), D)
Minjie Wang's avatar
Minjie Wang committed
132
        assert th.allclose(v, data[k][rowid])
133
    # test duplicate keys
Minjie Wang's avatar
Minjie Wang committed
134
    rowid = Index(th.tensor([8, 2, 2, 1]))
135
    rows = f[rowid]
Minjie Wang's avatar
Minjie Wang committed
136
    for k, v in rows.items():
137
        assert v.shape == (len(rowid), D)
Minjie Wang's avatar
Minjie Wang committed
138
        assert th.allclose(v, data[k][rowid])
139
140

    # setter
Minjie Wang's avatar
Minjie Wang committed
141
    rowid = Index(th.tensor([0, 2, 4]))
142
143
144
145
146
    vals = {'a1' : th.zeros((len(rowid), D)),
            'a2' : th.zeros((len(rowid), D)),
            'a3' : th.zeros((len(rowid), D)),
            }
    f[rowid] = vals
Minjie Wang's avatar
Minjie Wang committed
147
    for k, v in f[rowid].items():
Minjie Wang's avatar
Minjie Wang committed
148
        assert th.allclose(v, th.zeros((len(rowid), D)))
149

Minjie Wang's avatar
Minjie Wang committed
150
151
152
153
154
155
    # setting rows with new column should raise error with error initializer
    f.set_initializer(lambda shape, dtype : assert_(False))
    def failed_update_rows():
        vals['a4'] = th.ones((len(rowid), D))
        f[rowid] = vals
    assert check_fail(failed_update_rows)
156

Minjie Wang's avatar
Minjie Wang committed
157
158
def test_row2():
    # test row getter/setter autograd compatibility
159
    data = create_test_data(grad=True)
Minjie Wang's avatar
Minjie Wang committed
160
    f = FrameRef(Frame(data))
161
162
163
164

    # getter
    c1 = f['a1']
    # test non-duplicate keys
Minjie Wang's avatar
Minjie Wang committed
165
    rowid = Index(th.tensor([0, 2]))
166
167
    rows = f[rowid]
    rows['a1'].backward(th.ones((len(rowid), D)))
Minjie Wang's avatar
Minjie Wang committed
168
    assert th.allclose(c1.grad[:,0], th.tensor([1., 0., 1., 0., 0., 0., 0., 0., 0., 0.]))
169
170
    c1.grad.data.zero_()
    # test duplicate keys
Minjie Wang's avatar
Minjie Wang committed
171
    rowid = Index(th.tensor([8, 2, 2, 1]))
172
173
    rows = f[rowid]
    rows['a1'].backward(th.ones((len(rowid), D)))
Minjie Wang's avatar
Minjie Wang committed
174
    assert th.allclose(c1.grad[:,0], th.tensor([0., 1., 2., 0., 0., 0., 0., 0., 1., 0.]))
175
176
177
178
    c1.grad.data.zero_()

    # setter
    c1 = f['a1']
Minjie Wang's avatar
Minjie Wang committed
179
    rowid = Index(th.tensor([0, 2, 4]))
180
181
182
183
184
185
186
    vals = {'a1' : Variable(th.zeros((len(rowid), D)), requires_grad=True),
            'a2' : Variable(th.zeros((len(rowid), D)), requires_grad=True),
            'a3' : Variable(th.zeros((len(rowid), D)), requires_grad=True),
            }
    f[rowid] = vals
    c11 = f['a1']
    c11.backward(th.ones((N, D)))
Minjie Wang's avatar
Minjie Wang committed
187
188
    assert th.allclose(c1.grad[:,0], th.tensor([0., 1., 0., 1., 0., 1., 1., 1., 1., 1.]))
    assert th.allclose(vals['a1'].grad, th.ones((len(rowid), D)))
189
190
    assert vals['a2'].grad is None

Minjie Wang's avatar
Minjie Wang committed
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
def test_row3():
    # test row delete
    data = Frame(create_test_data())
    f = FrameRef(data)
    assert f.is_contiguous()
    assert f.is_span_whole_column()
    assert f.num_rows == N
    del f[th.tensor([2, 3])]
    assert not f.is_contiguous()
    assert not f.is_span_whole_column()
    # delete is lazy: only reflect on the ref while the
    # underlying storage should not be touched
    assert f.num_rows == N - 2
    assert data.num_rows == N
    newidx = list(range(N))
    newidx.pop(2)
    newidx.pop(2)
Minjie Wang's avatar
Minjie Wang committed
208
    newidx = toindex(newidx)
Minjie Wang's avatar
Minjie Wang committed
209
    for k, v in f.items():
Minjie Wang's avatar
Minjie Wang committed
210
        assert th.allclose(v, data[k][newidx])
Minjie Wang's avatar
Minjie Wang committed
211
212
213
214
215
216
217

def test_sharing():
    data = Frame(create_test_data())
    f1 = FrameRef(data, index=[0, 1, 2, 3])
    f2 = FrameRef(data, index=[2, 3, 4, 5, 6])
    # test read
    for k, v in f1.items():
Minjie Wang's avatar
Minjie Wang committed
218
        assert th.allclose(data[k].data[0:4], v)
Minjie Wang's avatar
Minjie Wang committed
219
    for k, v in f2.items():
Minjie Wang's avatar
Minjie Wang committed
220
221
        assert th.allclose(data[k].data[2:7], v)
    f2_a1 = f2['a1'].data
Minjie Wang's avatar
Minjie Wang committed
222
223
    # test write
    # update own ref should not been seen by the other.
Minjie Wang's avatar
Minjie Wang committed
224
    f1[Index(th.tensor([0, 1]))] = {
Minjie Wang's avatar
Minjie Wang committed
225
226
227
228
            'a1' : th.zeros([2, D]),
            'a2' : th.zeros([2, D]),
            'a3' : th.zeros([2, D]),
            }
Minjie Wang's avatar
Minjie Wang committed
229
    assert th.allclose(f2['a1'], f2_a1)
Minjie Wang's avatar
Minjie Wang committed
230
    # update shared space should been seen by the other.
Minjie Wang's avatar
Minjie Wang committed
231
    f1[Index(th.tensor([2, 3]))] = {
Minjie Wang's avatar
Minjie Wang committed
232
233
234
235
236
            'a1' : th.ones([2, D]),
            'a2' : th.ones([2, D]),
            'a3' : th.ones([2, D]),
            }
    f2_a1[0:2] = th.ones([2, D])
Minjie Wang's avatar
Minjie Wang committed
237
    assert th.allclose(f2['a1'], f2_a1)
238
239
240

if __name__ == '__main__':
    test_create()
Minjie Wang's avatar
Minjie Wang committed
241
242
243
244
245
246
247
248
    test_column1()
    test_column2()
    test_append1()
    test_append2()
    test_row1()
    test_row2()
    test_row3()
    test_sharing()