| @@ -68,52 +68,132 @@ def test_batcher_02(): | |||
| 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) | |||
| 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(t.edges[0].tolist()) | |||
| 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) } | |||
| 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) } | |||