More batcher tests.

tests/triacontagon/test_batch.py

@@ -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) }