|
- from triacontagon.batch import Batcher
- from triacontagon.data import Data
- from triacontagon.decode import dedicom_decoder
- import torch
-
-
- 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) }
|