decagon-pytorch/tests/triacontagon/test_split.py

from triacontagon.split import split_adj_mat, \
    split_edge_type, \
    split_data
from triacontagon.util import _equal
from triacontagon.data import EdgeType, \
    Data
from triacontagon.decode import dedicom_decoder
import torch


def test_split_adj_mat_01():
    adj_mat = torch.tensor([
        [0, 1, 0, 0, 1],
        [0, 0, 1, 0, 1],
        [1, 0, 0, 1, 0],
        [0, 0, 1, 1, 0]
    ]).to_sparse()

    (res,) = split_adj_mat(adj_mat, (1.,))
    assert torch.all(_equal(res, adj_mat))


def test_split_adj_mat_02():
    adj_mat = torch.tensor([
        [0, 1, 0, 0, 1],
        [0, 0, 1, 0, 1],
        [1, 0, 0, 1, 0],
        [0, 0, 1, 1, 0]
    ]).to_sparse()

    a, b = split_adj_mat(adj_mat, ( .5, .5 ))
    assert torch.all(_equal(a+b, adj_mat))


def test_split_adj_mat_03():
    adj_mat = torch.tensor([
        [0, 1, 0, 0, 1],
        [0, 0, 1, 0, 1],
        [1, 0, 0, 1, 0],
        [0, 0, 1, 1, 0]
    ]).to_sparse()

    a, b, c = split_adj_mat(adj_mat, ( .8, .1, .1 ))
    print('a:', a.to_dense(), 'b:', b.to_dense(), 'c:', c.to_dense())

    assert torch.all(_equal(a+b+c, adj_mat))


def test_split_edge_type_01():
    et = EdgeType('Dummy', 0, 1, [
        torch.tensor([
            [0, 1, 0, 0, 0],
            [0, 0, 1, 0, 1],
            [1, 0, 0, 0, 1],
            [0, 1, 0, 1, 0]
        ]).to_sparse()
    ], None, None)

    res = split_edge_type(et, (1.,))

    assert torch.all(_equal(et.adjacency_matrices[0],
        res[0].adjacency_matrices[0]))


def test_split_edge_type_02():
    et = EdgeType('Dummy', 0, 1, [
        torch.tensor([
            [0, 1, 0, 0, 0],
            [0, 0, 1, 0, 1],
            [1, 0, 0, 0, 1],
            [0, 1, 0, 1, 0]
        ]).to_sparse()
    ], None, None)

    res = split_edge_type(et, (.5, .5))

    assert torch.all(_equal(et.adjacency_matrices[0],
        res[0].adjacency_matrices[0] + \
        res[1].adjacency_matrices[0]))


def test_split_edge_type_03():
    et = EdgeType('Dummy', 0, 1, [
        torch.tensor([
            [0, 1, 0, 0, 0],
            [0, 0, 1, 0, 1],
            [1, 0, 0, 0, 1],
            [0, 1, 0, 1, 0]
        ]).to_sparse()
    ], None, None)

    res = split_edge_type(et, (.4, .4, .2))

    assert torch.all(_equal(et.adjacency_matrices[0],
        res[0].adjacency_matrices[0] + \
        res[1].adjacency_matrices[0] + \
        res[2].adjacency_matrices[0]))


def test_split_edge_type_04():
    et = EdgeType('Dummy', 0, 1, [
        torch.tensor([
            [0, 1, 0, 0, 0],
            [0, 0, 1, 0, 1],
            [1, 0, 0, 0, 1],
            [0, 1, 0, 1, 0]
        ]).to_sparse(),

        torch.tensor([
            [1, 0, 0, 0, 0],
            [0, 1, 0, 1, 0],
            [0, 0, 1, 1, 0],
            [1, 0, 1, 0, 0]
        ]).to_sparse()
    ], None, None)

    res = split_edge_type(et, (.4, .4, .2))

    assert torch.all(_equal(et.adjacency_matrices[0],
        res[0].adjacency_matrices[0] + \
        res[1].adjacency_matrices[0] + \
        res[2].adjacency_matrices[0]))

    assert torch.all(_equal(et.adjacency_matrices[1],
        res[0].adjacency_matrices[1] + \
        res[1].adjacency_matrices[1] + \
        res[2].adjacency_matrices[1]))


def test_split_data_01():
    data = Data()
    data.add_vertex_type('Foo', 5)
    data.add_vertex_type('Bar', 4)

    foo_foo = torch.tensor([
        [0, 1, 0, 1, 0],
        [0, 0, 0, 1, 0],
        [0, 1, 0, 0, 1],
        [0, 1, 0, 0, 0],
        [1, 0, 0, 1, 0]
    ], dtype=torch.float32)
    foo_foo = (foo_foo + foo_foo.transpose(0, 1)) / 2

    foo_bar = torch.tensor([
        [0, 1, 0, 1],
        [0, 0, 0, 1],
        [0, 1, 0, 0],
        [1, 0, 0, 0],
        [0, 0, 1, 1]
    ], dtype=torch.float32)
    bar_foo = foo_bar.transpose(0, 1)

    bar_bar = torch.tensor([
        [0, 0, 1, 0],
        [1, 0, 0, 0],
        [0, 1, 0, 1],
        [0, 1, 0, 0],
    ], dtype=torch.float32)
    bar_bar = (bar_bar + bar_bar.transpose(0, 1)) / 2

    data.add_edge_type('Foo-Foo', 0, 0, [
        foo_foo.to_sparse().coalesce()
    ], dedicom_decoder)
    data.add_edge_type('Foo-Bar', 0, 1, [
        foo_bar.to_sparse().coalesce()
    ], dedicom_decoder)
    data.add_edge_type('Bar-Foo', 1, 0, [
        bar_foo.to_sparse().coalesce()
    ], dedicom_decoder)
    data.add_edge_type('Bar-Bar', 1, 1, [
        bar_bar.to_sparse().coalesce()
    ], dedicom_decoder)

    (res,) = split_data(data, (1.,))

    assert torch.all(_equal(res.edge_types[0, 0].adjacency_matrices[0],
        data.edge_types[0, 0].adjacency_matrices[0]))

    assert torch.all(_equal(res.edge_types[0, 1].adjacency_matrices[0],
        data.edge_types[0, 1].adjacency_matrices[0]))

    assert torch.all(_equal(res.edge_types[1, 0].adjacency_matrices[0],
        data.edge_types[1, 0].adjacency_matrices[0]))

    assert torch.all(_equal(res.edge_types[1, 1].adjacency_matrices[0],
        data.edge_types[1, 1].adjacency_matrices[0]))


def test_split_data_02():
    data = Data()
    data.add_vertex_type('Foo', 5)
    data.add_vertex_type('Bar', 4)

    foo_foo = torch.tensor([
        [0, 1, 0, 1, 0],
        [0, 0, 0, 1, 0],
        [0, 1, 0, 0, 1],
        [0, 1, 0, 0, 0],
        [1, 0, 0, 1, 0]
    ], dtype=torch.float32)
    foo_foo = (foo_foo + foo_foo.transpose(0, 1)) / 2

    foo_bar = torch.tensor([
        [0, 1, 0, 1],
        [0, 0, 0, 1],
        [0, 1, 0, 0],
        [1, 0, 0, 0],
        [0, 0, 1, 1]
    ], dtype=torch.float32)
    bar_foo = foo_bar.transpose(0, 1)

    bar_bar = torch.tensor([
        [0, 0, 1, 0],
        [1, 0, 0, 0],
        [0, 1, 0, 1],
        [0, 1, 0, 0],
    ], dtype=torch.float32)
    bar_bar = (bar_bar + bar_bar.transpose(0, 1)) / 2

    data.add_edge_type('Foo-Foo', 0, 0, [
        foo_foo.to_sparse().coalesce()
    ], dedicom_decoder)
    data.add_edge_type('Foo-Bar', 0, 1, [
        foo_bar.to_sparse().coalesce()
    ], dedicom_decoder)
    data.add_edge_type('Bar-Foo', 1, 0, [
        bar_foo.to_sparse().coalesce()
    ], dedicom_decoder)
    data.add_edge_type('Bar-Bar', 1, 1, [
        bar_bar.to_sparse().coalesce()
    ], dedicom_decoder)

    a, b = split_data(data, (.5,.5))

    assert torch.all(_equal(a.edge_types[0, 0].adjacency_matrices[0] + \
        b.edge_types[0, 0].adjacency_matrices[0],
            data.edge_types[0, 0].adjacency_matrices[0]))

    assert torch.all(_equal(a.edge_types[0, 1].adjacency_matrices[0] + \
        b.edge_types[0, 1].adjacency_matrices[0],
            data.edge_types[0, 1].adjacency_matrices[0]))

    assert torch.all(_equal(a.edge_types[1, 0].adjacency_matrices[0] + \
        b.edge_types[1, 0].adjacency_matrices[0],
            data.edge_types[1, 0].adjacency_matrices[0]))

    assert torch.all(_equal(a.edge_types[1, 1].adjacency_matrices[0] + \
        b.edge_types[1, 1].adjacency_matrices[0],
            data.edge_types[1, 1].adjacency_matrices[0]))