| @@ -68,52 +68,132 @@ def test_batcher_02(): | |||||
| def test_batcher_03(): | 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, | k = (t.vertex_type_row, t.vertex_type_column, | ||||
| t.relation_type_index,) + \ | t.relation_type_index,) + \ | ||||
| tuple(t.edges[0].tolist()) | |||||
| tuple(e.tolist()) | |||||
| visited.add(k) | 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) } | |||||