test_frame.py 8.42 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
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
def test_slicing():
    data = Frame(create_test_data(grad=True))
    f1 = FrameRef(data, index=slice(1, 5))
    f2 = FrameRef(data, index=slice(3, 8))
    # test read
    for k, v in f1.items():
        assert th.allclose(data[k].data[1:5], v)
    f2_a1 = f2['a1'].data
    # test write
    f1[Index(th.tensor([0, 1]))] = {
            'a1': th.zeros([2, D]),
            'a2': th.zeros([2, D]),
            'a3': th.zeros([2, D]),
            }
    assert th.allclose(f2['a1'], f2_a1)
    
    f1[Index(th.tensor([2, 3]))] = {
            'a1': th.ones([2, D]),
            'a2': th.ones([2, D]),
            'a3': th.ones([2, D]),
            }
    f2_a1[0:2] = 1
    assert th.allclose(f2['a1'], f2_a1)

    f1[2:4] = {
            'a1': th.zeros([2, D]),
            'a2': th.zeros([2, D]),
            'a3': th.zeros([2, D]),
            }
    f2_a1[0:2] = 0
    assert th.allclose(f2['a1'], f2_a1)

271
272
if __name__ == '__main__':
    test_create()
Minjie Wang's avatar
Minjie Wang committed
273
274
275
276
277
278
279
280
    test_column1()
    test_column2()
    test_append1()
    test_append2()
    test_row1()
    test_row2()
    test_row3()
    test_sharing()
281
    test_slicing()