test_frame.py

import torch as th
from torch.autograd import Variable
import numpy as np
from dgl.frame import Frame, FrameRef
from dgl.utils import Index, toindex
import utils as U

N = 10
D = 5

def check_fail(fn):
    try:
        fn()
        return False
    except:
        return True

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(num_rows=N)
    for k, v in data.items():
        f1.update_column(k, v)
    print(f1.schemes)
    assert f1.keys() == set(data.keys())
    assert f1.num_columns == 3
    assert f1.num_rows == N
    f2 = Frame(data)
    assert f2.keys() == set(data.keys())
    assert f2.num_columns == 3
    assert f2.num_rows == N
    f1.clear()
    assert len(f1.schemes) == 0
    assert f1.num_rows == 0

def test_column1():
    # Test frame column getter/setter
    data = create_test_data()
    f = Frame(data)
    assert f.num_rows == N
    assert len(f) == 3
    assert U.allclose(f['a1'].data, data['a1'].data)
    f['a1'] = data['a2']
    assert U.allclose(f['a2'].data, data['a2'].data)
    # 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 len(f) == 0

def test_column2():
    # Test frameref column getter/setter
    data = Frame(create_test_data())
    f = FrameRef(data, toindex([3, 4, 5, 6, 7]))
    assert f.num_rows == 5
    assert len(f) == 3
    assert U.allclose(f['a1'], data['a1'].data[3:8])
    # set column should reflect on the referenced data
    f['a1'] = th.zeros([5, D])
    assert U.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)

def test_append1():
    # test append API on Frame
    data = create_test_data()
    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']
    assert c1.data.shape == (2 * N, D)
    truth = th.cat([data['a1'], data['a1']])
    assert U.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)

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))
    assert th.all(f._index.tousertensor() == th.tensor(new_idx, dtype=th.int64))
    assert data.num_rows == 4 * N

def test_append3():
    # test append on empty frame
    f = Frame(num_rows=5)
    data = {'h' : th.ones((3, 2))}
    f.append(data)
    assert f.num_rows == 8
    ans = th.cat([th.zeros((5, 2)), th.ones((3, 2))], dim=0)
    assert U.allclose(f['h'].data, ans)
    # test append with new column
    data = {'h' : 2 * th.ones((3, 2)), 'w' : 2 * th.ones((3, 2))}
    f.append(data)
    assert f.num_rows == 11
    ans1 = th.cat([ans, 2 * th.ones((3, 2))], 0)
    ans2 = th.cat([th.zeros((8, 2)), 2 * th.ones((3, 2))], 0)
    assert U.allclose(f['h'].data, ans1)
    assert U.allclose(f['w'].data, ans2)

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

    # getter
    # test non-duplicate keys
    rowid = Index(th.tensor([0, 2]))
    rows = f[rowid]
    for k, v in rows.items():
        assert v.shape == (len(rowid), D)
        assert U.allclose(v, data[k][rowid])
    # test duplicate keys
    rowid = Index(th.tensor([8, 2, 2, 1]))
    rows = f[rowid]
    for k, v in rows.items():
        assert v.shape == (len(rowid), D)
        assert U.allclose(v, data[k][rowid])

    # setter
    rowid = Index(th.tensor([0, 2, 4]))
    vals = {'a1' : th.zeros((len(rowid), D)),
            'a2' : th.zeros((len(rowid), D)),
            'a3' : th.zeros((len(rowid), D)),
            }
    f[rowid] = vals
    for k, v in f[rowid].items():
        assert U.allclose(v, th.zeros((len(rowid), D)))

    # 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)

def test_row2():
    # test row getter/setter autograd compatibility
    data = create_test_data(grad=True)
    f = FrameRef(Frame(data))

    # getter
    c1 = f['a1']
    # test non-duplicate keys
    rowid = Index(th.tensor([0, 2]))
    rows = f[rowid]
    rows['a1'].backward(th.ones((len(rowid), D)))
    assert U.allclose(c1.grad[:,0], th.tensor([1., 0., 1., 0., 0., 0., 0., 0., 0., 0.]))
    c1.grad.data.zero_()
    # test duplicate keys
    rowid = Index(th.tensor([8, 2, 2, 1]))
    rows = f[rowid]
    rows['a1'].backward(th.ones((len(rowid), D)))
    assert U.allclose(c1.grad[:,0], th.tensor([0., 1., 2., 0., 0., 0., 0., 0., 1., 0.]))
    c1.grad.data.zero_()

    # setter
    c1 = f['a1']
    rowid = Index(th.tensor([0, 2, 4]))
    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)))
    assert U.allclose(c1.grad[:,0], th.tensor([0., 1., 0., 1., 0., 1., 1., 1., 1., 1.]))
    assert U.allclose(vals['a1'].grad, th.ones((len(rowid), D)))
    assert vals['a2'].grad is None

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[toindex(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)
    newidx = toindex(newidx)
    for k, v in f.items():
        assert U.allclose(v, data[k][newidx])

def test_row4():
    # test updating row with empty frame but has preset num_rows
    f = FrameRef(Frame(num_rows=5))
    rowid = Index(th.tensor([0, 2, 4]))
    f[rowid] = {'h' : th.ones((3, 2))}
    ans = th.zeros((5, 2))
    ans[th.tensor([0, 2, 4])] = th.ones((3, 2))
    assert U.allclose(f['h'], ans)

def test_sharing():
    data = Frame(create_test_data())
    f1 = FrameRef(data, index=toindex([0, 1, 2, 3]))
    f2 = FrameRef(data, index=toindex([2, 3, 4, 5, 6]))
    # test read
    for k, v in f1.items():
        assert U.allclose(data[k].data[0:4], v)
    for k, v in f2.items():
        assert U.allclose(data[k].data[2:7], v)
    f2_a1 = f2['a1'].data
    # test write
    # update own ref should not been seen by the other.
    f1[Index(th.tensor([0, 1]))] = {
            'a1' : th.zeros([2, D]),
            'a2' : th.zeros([2, D]),
            'a3' : th.zeros([2, D]),
            }
    assert U.allclose(f2['a1'], f2_a1)
    # update shared space should been seen by the other.
    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] = th.ones([2, D])
    assert U.allclose(f2['a1'], f2_a1)

def test_slicing():
    data = Frame(create_test_data(grad=True))
    f1 = FrameRef(data, index=toindex(slice(1, 5)))
    f2 = FrameRef(data, index=toindex(slice(3, 8)))
    # test read
    for k, v in f1.items():
        assert U.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 U.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[toindex(slice(0,2))] = 1
    assert U.allclose(f2['a1'], f2_a1)

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

def test_add_rows():
    data = Frame()
    f1 = FrameRef(data)
    f1.add_rows(4)
    x = th.randn(1, 4)
    f1[Index(th.tensor([0]))] = {'x': x}
    ans = th.cat([x, th.zeros(3, 4)])
    assert U.allclose(f1['x'], ans)
    f1.add_rows(4)
    f1[toindex(slice(4,8))] = {'x': th.ones(4, 4), 'y': th.ones(4, 5)}
    ans = th.cat([ans, th.ones(4, 4)])
    assert U.allclose(f1['x'], ans)
    ans = th.cat([th.zeros(4, 5), th.ones(4, 5)])
    assert U.allclose(f1['y'], ans)

def test_inplace():
    f = FrameRef(Frame(create_test_data()))
    print(f.schemes)
    a1addr = f['a1'].data.data_ptr()
    a2addr = f['a2'].data.data_ptr()
    a3addr = f['a3'].data.data_ptr()

    # column updates are always out-of-place
    f['a1'] = th.ones((N, D))
    newa1addr = f['a1'].data.data_ptr()
    assert a1addr != newa1addr
    a1addr = newa1addr
    # full row update that becomes column update
    f[toindex(slice(0, N))] = {'a1' : th.ones((N, D))}
    assert f['a1'].data.data_ptr() != a1addr

    # row update (outplace) w/ slice
    f[toindex(slice(1, 4))] = {'a2' : th.ones((3, D))}
    newa2addr = f['a2'].data.data_ptr()
    assert a2addr != newa2addr
    a2addr = newa2addr
    # row update (outplace) w/ list
    f[toindex([1, 3, 5])] = {'a2' : th.ones((3, D))}
    newa2addr = f['a2'].data.data_ptr()
    assert a2addr != newa2addr
    a2addr = newa2addr

    # row update (inplace) w/ slice
    f.update_data(toindex(slice(1, 4)), {'a2' : th.ones((3, D))}, True)
    newa2addr = f['a2'].data.data_ptr()
    assert a2addr == newa2addr
    # row update (inplace) w/ list
    f.update_data(toindex([1, 3, 5]), {'a2' : th.ones((3, D))}, True)
    newa2addr = f['a2'].data.data_ptr()
    assert a2addr == newa2addr

if __name__ == '__main__':
    test_create()
    test_column1()
    test_column2()
    test_append1()
    test_append2()
    test_append3()
    test_row1()
    test_row2()
    test_row3()
    test_row4()
    test_sharing()
    test_slicing()
    test_add_rows()
    test_inplace()