test_frame.py 11.7 KB
Newer Older
1
import numpy as np
Minjie Wang's avatar
Minjie Wang committed
2
from dgl.frame import Frame, FrameRef
Minjie Wang's avatar
Minjie Wang committed
3
from dgl.utils import Index, toindex
4
import backend as F
VoVAllen's avatar
VoVAllen committed
5
6
import dgl
import unittest
7
8

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

Minjie Wang's avatar
Minjie Wang committed
11
12
13
14
15
16
def check_fail(fn):
    try:
        fn()
        return False
    except:
        return True
17
18

def create_test_data(grad=False):
19
20
21
22
23
24
25
    c1 = F.randn((N, D))
    c2 = F.randn((N, D))
    c3 = F.randn((N, D))
    if grad:
        c1 = F.attach_grad(c1)
        c2 = F.attach_grad(c2)
        c3 = F.attach_grad(c3)
26
27
28
29
    return {'a1' : c1, 'a2' : c2, 'a3' : c3}

def test_create():
    data = create_test_data()
30
    f1 = Frame(num_rows=N)
31
    for k, v in data.items():
Minjie Wang's avatar
Minjie Wang committed
32
33
34
        f1.update_column(k, v)
    print(f1.schemes)
    assert f1.keys() == set(data.keys())
35
36
37
    assert f1.num_columns == 3
    assert f1.num_rows == N
    f2 = Frame(data)
Minjie Wang's avatar
Minjie Wang committed
38
    assert f2.keys() == set(data.keys())
39
40
41
42
43
44
    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
45
46
def test_column1():
    # Test frame column getter/setter
47
48
    data = create_test_data()
    f = Frame(data)
Minjie Wang's avatar
Minjie Wang committed
49
50
    assert f.num_rows == N
    assert len(f) == 3
51
    assert F.allclose(f['a1'].data, data['a1'])
52
    f['a1'] = data['a2']
53
    assert F.allclose(f['a2'].data, data['a2'])
Minjie Wang's avatar
Minjie Wang committed
54
55
    # add a different length column should fail
    def failed_add_col():
56
        f['a4'] = F.zeros([N+1, D])
Minjie Wang's avatar
Minjie Wang committed
57
58
59
60
61
62
63
64
65
66
67
    assert check_fail(failed_add_col)
    # delete all the columns
    del f['a1']
    del f['a2']
    assert len(f) == 1
    del f['a3']
    assert len(f) == 0

def test_column2():
    # Test frameref column getter/setter
    data = Frame(create_test_data())
Minjie Wang's avatar
Minjie Wang committed
68
    f = FrameRef(data, toindex([3, 4, 5, 6, 7]))
Minjie Wang's avatar
Minjie Wang committed
69
70
    assert f.num_rows == 5
    assert len(f) == 3
71
    assert F.allclose(f['a1'], F.narrow_row(data['a1'].data, 3, 8))
Minjie Wang's avatar
Minjie Wang committed
72
    # set column should reflect on the referenced data
73
74
    f['a1'] = F.zeros([5, D])
    assert F.allclose(F.narrow_row(data['a1'].data, 3, 8), F.zeros([5, D]))
Minjie Wang's avatar
Minjie Wang committed
75
76
77
    # add new partial column should fail with error initializer
    f.set_initializer(lambda shape, dtype : assert_(False))
    def failed_add_col():
78
        f['a4'] = F.ones([5, D])
Minjie Wang's avatar
Minjie Wang committed
79
    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']
91
92
93
    assert tuple(F.shape(c1.data)) == (2 * N, D)
    truth = F.cat([data['a1'], data['a1']], 0)
    assert F.allclose(truth, c1.data)
Minjie Wang's avatar
Minjie Wang committed
94
    # append dict of different length columns should fail
95
    f3 = {'a1' : F.zeros((3, D)), 'a2' : F.zeros((3, D)), 'a3' : F.zeros((2, D))}
Minjie Wang's avatar
Minjie Wang committed
96
97
98
    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))
118
    assert F.array_equal(f._index.tousertensor(), F.copy_to(F.tensor(new_idx, dtype=F.int64), F.cpu()))
Minjie Wang's avatar
Minjie Wang committed
119
120
    assert data.num_rows == 4 * N

121
122
123
def test_append3():
    # test append on empty frame
    f = Frame(num_rows=5)
124
    data = {'h' : F.ones((3, 2))}
125
126
    f.append(data)
    assert f.num_rows == 8
127
128
    ans = F.cat([F.zeros((5, 2)), F.ones((3, 2))], 0)
    assert F.allclose(f['h'].data, ans)
129
    # test append with new column
130
    data = {'h' : 2 * F.ones((3, 2)), 'w' : 2 * F.ones((3, 2))}
131
132
    f.append(data)
    assert f.num_rows == 11
133
134
135
136
    ans1 = F.cat([ans, 2 * F.ones((3, 2))], 0)
    ans2 = F.cat([F.zeros((8, 2)), 2 * F.ones((3, 2))], 0)
    assert F.allclose(f['h'].data, ans1)
    assert F.allclose(f['w'].data, ans2)
137

Minjie Wang's avatar
Minjie Wang committed
138
139
140
141
def test_row1():
    # test row getter/setter
    data = create_test_data()
    f = FrameRef(Frame(data))
142
143
144

    # getter
    # test non-duplicate keys
145
    rowid = Index(F.tensor([0, 2]))
146
    rows = f[rowid]
Minjie Wang's avatar
Minjie Wang committed
147
    for k, v in rows.items():
148
        assert tuple(F.shape(v)) == (len(rowid), D)
149
        assert F.allclose(v, F.gather_row(data[k], F.tensor(rowid.tousertensor())))
150
    # test duplicate keys
151
    rowid = Index(F.tensor([8, 2, 2, 1]))
152
    rows = f[rowid]
Minjie Wang's avatar
Minjie Wang committed
153
    for k, v in rows.items():
154
        assert tuple(F.shape(v)) == (len(rowid), D)
155
        assert F.allclose(v, F.gather_row(data[k], F.tensor(rowid.tousertensor())))
156
157

    # setter
158
159
160
161
    rowid = Index(F.tensor([0, 2, 4]))
    vals = {'a1' : F.zeros((len(rowid), D)),
            'a2' : F.zeros((len(rowid), D)),
            'a3' : F.zeros((len(rowid), D)),
162
163
            }
    f[rowid] = vals
Minjie Wang's avatar
Minjie Wang committed
164
    for k, v in f[rowid].items():
165
        assert F.allclose(v, F.zeros((len(rowid), D)))
166

Minjie Wang's avatar
Minjie Wang committed
167
168
169
    # setting rows with new column should raise error with error initializer
    f.set_initializer(lambda shape, dtype : assert_(False))
    def failed_update_rows():
170
        vals['a4'] = F.ones((len(rowid), D))
Minjie Wang's avatar
Minjie Wang committed
171
172
        f[rowid] = vals
    assert check_fail(failed_update_rows)
173

Minjie Wang's avatar
Minjie Wang committed
174
175
def test_row2():
    # test row getter/setter autograd compatibility
176
    data = create_test_data(grad=True)
Minjie Wang's avatar
Minjie Wang committed
177
    f = FrameRef(Frame(data))
178

179
180
181
182
183
184
185
    with F.record_grad():
        # getter
        c1 = f['a1']
        # test non-duplicate keys
        rowid = Index(F.tensor([0, 2]))
        rows = f[rowid]
        y = rows['a1']
VoVAllen's avatar
VoVAllen committed
186
        F.backward(y, F.ones((len(rowid), D)))
187
    assert F.allclose(F.grad(c1)[:,0], F.tensor([1., 0., 1., 0., 0., 0., 0., 0., 0., 0.]))
188

189
190
191
192
193
194
195
    f['a1'] = F.attach_grad(f['a1'])
    with F.record_grad():
        c1 = f['a1']
        # test duplicate keys
        rowid = Index(F.tensor([8, 2, 2, 1]))
        rows = f[rowid]
        y = rows['a1']
VoVAllen's avatar
VoVAllen committed
196
        F.backward(y, F.ones((len(rowid), D)))
197
198
199
200
201
202
203
204
205
206
207
208
209
    assert F.allclose(F.grad(c1)[:,0], F.tensor([0., 1., 2., 0., 0., 0., 0., 0., 1., 0.]))

    f['a1'] = F.attach_grad(f['a1'])
    with F.record_grad():
        # setter
        c1 = f['a1']
        rowid = Index(F.tensor([0, 2, 4]))
        vals = {'a1' : F.attach_grad(F.zeros((len(rowid), D))),
                'a2' : F.attach_grad(F.zeros((len(rowid), D))),
                'a3' : F.attach_grad(F.zeros((len(rowid), D))),
                }
        f[rowid] = vals
        c11 = f['a1']
VoVAllen's avatar
VoVAllen committed
210
        F.backward(c11, F.ones((N, D)))
211
212
213
    assert F.allclose(F.grad(c1)[:,0], F.tensor([0., 1., 0., 1., 0., 1., 1., 1., 1., 1.]))
    assert F.allclose(F.grad(vals['a1']), F.ones((len(rowid), D)))
    assert F.is_no_grad(vals['a2'])
214

Minjie Wang's avatar
Minjie Wang committed
215
216
217
218
219
220
221
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
222
    del f[toindex(F.tensor([2, 3]))]
Minjie Wang's avatar
Minjie Wang committed
223
224
225
226
227
228
229
230
231
    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
232
    newidx = toindex(newidx)
Minjie Wang's avatar
Minjie Wang committed
233
    for k, v in f.items():
234
        assert F.allclose(v, data[k][newidx])
Minjie Wang's avatar
Minjie Wang committed
235

VoVAllen's avatar
VoVAllen committed
236
237

@unittest.skipIf(dgl.backend.backend_name == "tensorflow", reason="TF doesn't support inplace update")
238
239
240
def test_row4():
    # test updating row with empty frame but has preset num_rows
    f = FrameRef(Frame(num_rows=5))
241
242
243
244
245
    rowid = Index(F.tensor([0, 2, 4]))
    f[rowid] = {'h' : F.ones((3, 2))}
    ans = F.zeros((5, 2))
    ans[F.tensor([0, 2, 4])] = F.ones((3, 2))
    assert F.allclose(f['h'], ans)
246

VoVAllen's avatar
VoVAllen committed
247
@unittest.skipIf(dgl.backend.backend_name == "tensorflow", reason="TF doesn't support inplace update")
Minjie Wang's avatar
Minjie Wang committed
248
249
def test_sharing():
    data = Frame(create_test_data())
Minjie Wang's avatar
Minjie Wang committed
250
251
    f1 = FrameRef(data, index=toindex([0, 1, 2, 3]))
    f2 = FrameRef(data, index=toindex([2, 3, 4, 5, 6]))
Minjie Wang's avatar
Minjie Wang committed
252
253
    # test read
    for k, v in f1.items():
254
        assert F.allclose(F.narrow_row(data[k].data, 0, 4), v)
Minjie Wang's avatar
Minjie Wang committed
255
    for k, v in f2.items():
256
257
        assert F.allclose(F.narrow_row(data[k].data, 2, 7), v)
    f2_a1 = f2['a1']
Minjie Wang's avatar
Minjie Wang committed
258
259
    # test write
    # update own ref should not been seen by the other.
260
261
262
263
    f1[Index(F.tensor([0, 1]))] = {
            'a1' : F.zeros([2, D]),
            'a2' : F.zeros([2, D]),
            'a3' : F.zeros([2, D]),
Minjie Wang's avatar
Minjie Wang committed
264
            }
265
    assert F.allclose(f2['a1'], f2_a1)
Minjie Wang's avatar
Minjie Wang committed
266
    # update shared space should been seen by the other.
267
268
269
270
    f1[Index(F.tensor([2, 3]))] = {
            'a1' : F.ones([2, D]),
            'a2' : F.ones([2, D]),
            'a3' : F.ones([2, D]),
Minjie Wang's avatar
Minjie Wang committed
271
            }
272
273
    F.narrow_row_set(f2_a1, 0, 2, F.ones([2, D]))
    assert F.allclose(f2['a1'], f2_a1)
274

VoVAllen's avatar
VoVAllen committed
275
@unittest.skipIf(dgl.backend.backend_name == "tensorflow", reason="TF doesn't support inplace update")
276
277
def test_slicing():
    data = Frame(create_test_data(grad=True))
Minjie Wang's avatar
Minjie Wang committed
278
279
    f1 = FrameRef(data, index=toindex(slice(1, 5)))
    f2 = FrameRef(data, index=toindex(slice(3, 8)))
280
281
    # test read
    for k, v in f1.items():
282
283
        assert F.allclose(F.narrow_row(data[k].data, 1, 5), v)
    f2_a1 = f2['a1']    # is a tensor
284
    # test write
285
286
287
288
    f1[Index(F.tensor([0, 1]))] = {
            'a1': F.zeros([2, D]),
            'a2': F.zeros([2, D]),
            'a3': F.zeros([2, D]),
289
            }
290
    assert F.allclose(f2['a1'], f2_a1)
291

292
293
294
295
    f1[Index(F.tensor([2, 3]))] = {
            'a1': F.ones([2, D]),
            'a2': F.ones([2, D]),
            'a3': F.ones([2, D]),
296
            }
297
298
    F.narrow_row_set(f2_a1, 0, 2, 1)
    assert F.allclose(f2['a1'], f2_a1)
299

300
301
302
303
    f1[toindex(slice(2, 4))] = {
            'a1': F.zeros([2, D]),
            'a2': F.zeros([2, D]),
            'a3': F.zeros([2, D]),
304
            }
305
306
    F.narrow_row_set(f2_a1, 0, 2, 0)
    assert F.allclose(f2['a1'], f2_a1)
307

308
309
310
311
def test_add_rows():
    data = Frame()
    f1 = FrameRef(data)
    f1.add_rows(4)
312
313
314
315
    x = F.randn((1, 4))
    f1[Index(F.tensor([0]))] = {'x': x}
    ans = F.cat([x, F.zeros((3, 4))], 0)
    assert F.allclose(f1['x'], ans)
316
    f1.add_rows(4)
317
318
319
320
321
    f1[toindex(slice(4, 8))] = {'x': F.ones((4, 4)), 'y': F.ones((4, 5))}
    ans = F.cat([ans, F.ones((4, 4))], 0)
    assert F.allclose(f1['x'], ans)
    ans = F.cat([F.zeros((4, 5)), F.ones((4, 5))], 0)
    assert F.allclose(f1['y'], ans)
322

VoVAllen's avatar
VoVAllen committed
323
@unittest.skipIf(dgl.backend.backend_name == "tensorflow", reason="TF doesn't support inplace update")
Minjie Wang's avatar
Minjie Wang committed
324
325
326
def test_inplace():
    f = FrameRef(Frame(create_test_data()))
    print(f.schemes)
327
328
329
    a1addr = id(f['a1'])
    a2addr = id(f['a2'])
    a3addr = id(f['a3'])
Minjie Wang's avatar
Minjie Wang committed
330
331

    # column updates are always out-of-place
332
333
    f['a1'] = F.ones((N, D))
    newa1addr = id(f['a1'])
Minjie Wang's avatar
Minjie Wang committed
334
335
336
    assert a1addr != newa1addr
    a1addr = newa1addr
    # full row update that becomes column update
337
338
    f[toindex(slice(0, N))] = {'a1' : F.ones((N, D))}
    assert id(f['a1']) != a1addr
Minjie Wang's avatar
Minjie Wang committed
339
340

    # row update (outplace) w/ slice
341
342
    f[toindex(slice(1, 4))] = {'a2' : F.ones((3, D))}
    newa2addr = id(f['a2'])
Minjie Wang's avatar
Minjie Wang committed
343
344
345
    assert a2addr != newa2addr
    a2addr = newa2addr
    # row update (outplace) w/ list
346
347
    f[toindex([1, 3, 5])] = {'a2' : F.ones((3, D))}
    newa2addr = id(f['a2'])
Minjie Wang's avatar
Minjie Wang committed
348
349
350
351
    assert a2addr != newa2addr
    a2addr = newa2addr

    # row update (inplace) w/ slice
352
353
    f.update_data(toindex(slice(1, 4)), {'a2' : F.ones((3, D))}, True)
    newa2addr = id(f['a2'])
Minjie Wang's avatar
Minjie Wang committed
354
355
    assert a2addr == newa2addr
    # row update (inplace) w/ list
356
357
    f.update_data(toindex([1, 3, 5]), {'a2' : F.ones((3, D))}, True)
    newa2addr = id(f['a2'])
Minjie Wang's avatar
Minjie Wang committed
358
359
    assert a2addr == newa2addr

360
361
if __name__ == '__main__':
    test_create()
Minjie Wang's avatar
Minjie Wang committed
362
363
364
365
    test_column1()
    test_column2()
    test_append1()
    test_append2()
366
    test_append3()
Minjie Wang's avatar
Minjie Wang committed
367
368
369
    test_row1()
    test_row2()
    test_row3()
370
    test_row4()
Minjie Wang's avatar
Minjie Wang committed
371
    test_sharing()
372
    test_slicing()
373
    test_add_rows()
Minjie Wang's avatar
Minjie Wang committed
374
    test_inplace()