|
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316 |
- from triacontagon.batch import _same_data_org, \
- DualBatcher, \
- Batcher
- from triacontagon.data import Data
- from triacontagon.decode import dedicom_decoder
- import torch
-
-
- def test_same_data_org_01():
- data = Data()
- assert _same_data_org(data, data)
-
- data.add_vertex_type('Foo', 10)
- assert _same_data_org(data, data)
-
- data.add_vertex_type('Bar', 10)
- assert _same_data_org(data, data)
-
- data_1 = Data()
- assert not _same_data_org(data, data_1)
-
- data_1.add_vertex_type('Foo', 10)
- assert not _same_data_org(data, data_1)
-
- data_1.add_vertex_type('Bar', 10)
- assert _same_data_org(data, data_1)
-
-
- def test_same_data_org_02():
- data = Data()
- data.add_vertex_type('Foo', 4)
- data.add_edge_type('Foo-Foo', 0, 0, [
- torch.tensor([
- [0, 0, 0, 1],
- [1, 0, 0, 0],
- [0, 1, 1, 0],
- [1, 0, 1, 0]
- ]).to_sparse()
- ], dedicom_decoder)
-
- assert _same_data_org(data, data)
-
- data_1 = Data()
- data_1.add_vertex_type('Foo', 4)
- data_1.add_edge_type('Foo-Foo', 0, 0, [
- torch.tensor([
- [0, 0, 0, 1],
- [1, 0, 0, 0],
- [0, 1, 1, 0],
- [1, 0, 0, 0]
- ]).to_sparse()
- ], dedicom_decoder)
-
- assert not _same_data_org(data, data_1)
-
-
- def test_batcher_01():
- d = Data()
- d.add_vertex_type('Gene', 5)
-
- d.add_edge_type('Gene-Gene', 0, 0, [
- torch.tensor([
- [0, 1, 0, 1, 0],
- [0, 0, 0, 0, 1],
- [1, 0, 0, 0, 0],
- [0, 0, 1, 0, 0],
- [0, 0, 0, 1, 0]
- ]).to_sparse()
- ], dedicom_decoder)
-
- b = Batcher(d, batch_size=1)
-
- visited = set()
- for t in b:
- print(t)
- k = tuple(t.edges[0].tolist())
- visited.add(k)
-
- assert visited == { (0, 1), (0, 3),
- (1, 4), (2, 0), (3, 2), (4, 3) }
-
-
- def test_batcher_02():
- d = Data()
- d.add_vertex_type('Gene', 5)
-
- d.add_edge_type('Gene-Gene', 0, 0, [
- torch.tensor([
- [0, 1, 0, 1, 0],
- [0, 0, 0, 0, 1],
- [1, 0, 0, 0, 0],
- [0, 0, 1, 0, 0],
- [0, 0, 0, 1, 0]
- ]).to_sparse(),
-
- torch.tensor([
- [1, 0, 1, 0, 0],
- [0, 0, 0, 1, 0],
- [0, 0, 0, 0, 1],
- [0, 1, 0, 0, 0],
- [0, 0, 1, 0, 0]
- ]).to_sparse()
- ], dedicom_decoder)
-
- b = Batcher(d, batch_size=1)
-
- visited = set()
- for t in b:
- print(t)
- k = (t.relation_type_index,) + \
- tuple(t.edges[0].tolist())
- visited.add(k)
-
- assert visited == { (0, 0, 1), (0, 0, 3),
- (0, 1, 4), (0, 2, 0), (0, 3, 2), (0, 4, 3),
- (1, 0, 0), (1, 0, 2), (1, 1, 3), (1, 2, 4),
- (1, 3, 1), (1, 4, 2) }
-
-
- def test_batcher_03():
- d = Data()
- d.add_vertex_type('Gene', 5)
- d.add_vertex_type('Drug', 4)
-
- d.add_edge_type('Gene-Gene', 0, 0, [
- torch.tensor([
- [0, 1, 0, 1, 0],
- [0, 0, 0, 0, 1],
- [1, 0, 0, 0, 0],
- [0, 0, 1, 0, 0],
- [0, 0, 0, 1, 0]
- ]).to_sparse(),
-
- torch.tensor([
- [1, 0, 1, 0, 0],
- [0, 0, 0, 1, 0],
- [0, 0, 0, 0, 1],
- [0, 1, 0, 0, 0],
- [0, 0, 1, 0, 0]
- ]).to_sparse()
- ], dedicom_decoder)
-
- d.add_edge_type('Gene-Drug', 0, 1, [
- torch.tensor([
- [0, 1, 0, 0],
- [1, 0, 0, 1],
- [0, 1, 0, 0],
- [0, 0, 1, 0],
- [0, 1, 1, 0]
- ]).to_sparse()
- ], dedicom_decoder)
-
- b = Batcher(d, batch_size=1)
-
- visited = set()
- for t in b:
- print(t)
- k = (t.vertex_type_row, t.vertex_type_column,
- t.relation_type_index,) + \
- tuple(t.edges[0].tolist())
- visited.add(k)
-
- assert visited == { (0, 0, 0, 0, 1), (0, 0, 0, 0, 3),
- (0, 0, 0, 1, 4), (0, 0, 0, 2, 0), (0, 0, 0, 3, 2), (0, 0, 0, 4, 3),
- (0, 0, 1, 0, 0), (0, 0, 1, 0, 2), (0, 0, 1, 1, 3), (0, 0, 1, 2, 4),
- (0, 0, 1, 3, 1), (0, 0, 1, 4, 2),
- (0, 1, 0, 0, 1), (0, 1, 0, 1, 0), (0, 1, 0, 1, 3),
- (0, 1, 0, 2, 1), (0, 1, 0, 3, 2), (0, 1, 0, 4, 1),
- (0, 1, 0, 4, 2) }
-
-
- def test_batcher_04():
- d = Data()
- d.add_vertex_type('Gene', 5)
-
- d.add_edge_type('Gene-Gene', 0, 0, [
- torch.tensor([
- [0, 1, 0, 1, 0],
- [0, 0, 0, 0, 1],
- [1, 0, 0, 0, 0],
- [0, 0, 1, 0, 0],
- [0, 0, 0, 1, 0]
- ]).to_sparse()
- ], dedicom_decoder)
-
- b = Batcher(d, batch_size=3)
-
- visited = set()
- for t in b:
- print(t)
- for e in t.edges:
- k = tuple(e.tolist())
- visited.add(k)
-
- assert visited == { (0, 1), (0, 3),
- (1, 4), (2, 0), (3, 2), (4, 3) }
-
-
- def test_batcher_05():
- d = Data()
- d.add_vertex_type('Gene', 5)
- d.add_vertex_type('Drug', 4)
-
- d.add_edge_type('Gene-Gene', 0, 0, [
- torch.tensor([
- [0, 1, 0, 1, 0],
- [0, 0, 0, 0, 1],
- [1, 0, 0, 0, 0],
- [0, 0, 1, 0, 0],
- [0, 0, 0, 1, 0]
- ]).to_sparse(),
-
- torch.tensor([
- [1, 0, 1, 0, 0],
- [0, 0, 0, 1, 0],
- [0, 0, 0, 0, 1],
- [0, 1, 0, 0, 0],
- [0, 0, 1, 0, 0]
- ]).to_sparse()
- ], dedicom_decoder)
-
- d.add_edge_type('Gene-Drug', 0, 1, [
- torch.tensor([
- [0, 1, 0, 0],
- [1, 0, 0, 1],
- [0, 1, 0, 0],
- [0, 0, 1, 0],
- [0, 1, 1, 0]
- ]).to_sparse()
- ], dedicom_decoder)
-
- b = Batcher(d, batch_size=5)
-
- visited = set()
- for t in b:
- print(t)
- for e in t.edges:
- k = (t.vertex_type_row, t.vertex_type_column,
- t.relation_type_index,) + \
- tuple(e.tolist())
- visited.add(k)
-
- assert visited == { (0, 0, 0, 0, 1), (0, 0, 0, 0, 3),
- (0, 0, 0, 1, 4), (0, 0, 0, 2, 0), (0, 0, 0, 3, 2), (0, 0, 0, 4, 3),
- (0, 0, 1, 0, 0), (0, 0, 1, 0, 2), (0, 0, 1, 1, 3), (0, 0, 1, 2, 4),
- (0, 0, 1, 3, 1), (0, 0, 1, 4, 2),
- (0, 1, 0, 0, 1), (0, 1, 0, 1, 0), (0, 1, 0, 1, 3),
- (0, 1, 0, 2, 1), (0, 1, 0, 3, 2), (0, 1, 0, 4, 1),
- (0, 1, 0, 4, 2) }
-
-
- def test_dual_batcher_01():
- d = Data()
- d.add_vertex_type('Gene', 5)
- d.add_vertex_type('Drug', 4)
-
- d.add_edge_type('Gene-Gene', 0, 0, [
- torch.tensor([
- [0, 1, 0, 1, 0],
- [0, 0, 0, 0, 1],
- [1, 0, 0, 0, 0],
- [0, 0, 1, 0, 0],
- [0, 0, 0, 1, 0]
- ]).to_sparse(),
-
- torch.tensor([
- [1, 0, 1, 0, 0],
- [0, 0, 0, 1, 0],
- [0, 0, 0, 0, 1],
- [0, 1, 0, 0, 0],
- [0, 0, 1, 0, 0]
- ]).to_sparse()
- ], dedicom_decoder)
-
- d.add_edge_type('Gene-Drug', 0, 1, [
- torch.tensor([
- [0, 1, 0, 0],
- [1, 0, 0, 1],
- [0, 1, 0, 0],
- [0, 0, 1, 0],
- [0, 1, 1, 0]
- ]).to_sparse()
- ], dedicom_decoder)
-
- b = DualBatcher(d, d, batch_size=5)
-
- visited_pos = set()
- visited_neg = set()
- for t_pos, t_neg in b:
- assert t_pos.vertex_type_row == t_neg.vertex_type_row
- assert t_pos.vertex_type_column == t_neg.vertex_type_column
- assert t_pos.relation_type_index == t_neg.relation_type_index
- assert len(t_pos.edges) == len(t_neg.edges)
-
- for e in t_pos.edges:
- k = (t_pos.vertex_type_row, t_pos.vertex_type_column,
- t_pos.relation_type_index,) + \
- tuple(e.tolist())
- visited_pos.add(k)
-
- for e in t_neg.edges:
- k = (t_neg.vertex_type_row, t_neg.vertex_type_column,
- t_neg.relation_type_index,) + \
- tuple(e.tolist())
- visited_neg.add(k)
-
- expected = { (0, 0, 0, 0, 1), (0, 0, 0, 0, 3),
- (0, 0, 0, 1, 4), (0, 0, 0, 2, 0), (0, 0, 0, 3, 2), (0, 0, 0, 4, 3),
- (0, 0, 1, 0, 0), (0, 0, 1, 0, 2), (0, 0, 1, 1, 3), (0, 0, 1, 2, 4),
- (0, 0, 1, 3, 1), (0, 0, 1, 4, 2),
- (0, 1, 0, 0, 1), (0, 1, 0, 1, 0), (0, 1, 0, 1, 3),
- (0, 1, 0, 2, 1), (0, 1, 0, 3, 2), (0, 1, 0, 4, 1),
- (0, 1, 0, 4, 2) }
-
- assert visited_pos == expected
- assert visited_neg == expected
|