Fix regressions in test_input.

tests/icosagon/test_input.py

@@ -9,28 +9,17 @@ def _some_data():
     d = Data()
     d.add_node_type('Gene', 1000)
     d.add_node_type('Drug', 100)
-    d.add_relation_type('Target', 1, 0, torch.rand(100, 1000))
-    d.add_relation_type('Interaction', 0, 0, torch.rand(1000, 1000))
-    d.add_relation_type('Side Effect: Nausea', 1, 1, torch.rand(100, 100))
-    d.add_relation_type('Side Effect: Infertility', 1, 1, torch.rand(100, 100))
-    d.add_relation_type('Side Effect: Death', 1, 1, torch.rand(100, 100))
-    return d
 
+    fam = d.add_relation_family('Drug-Gene', 1, 0, False)
+    fam.add_relation_type('Target', 1, 0, torch.rand(100, 1000))
 
-def _some_data_with_interactions():
-    d = Data()
-    d.add_node_type('Gene', 1000)
-    d.add_node_type('Drug', 100)
-    d.add_relation_type('Target', 1, 0,
-        torch.rand((100, 1000), dtype=torch.float32).round())
-    d.add_relation_type('Interaction', 0, 0,
-        torch.rand((1000, 1000), dtype=torch.float32).round())
-    d.add_relation_type('Side Effect: Nausea', 1, 1,
-        torch.rand((100, 100), dtype=torch.float32).round())
-    d.add_relation_type('Side Effect: Infertility', 1, 1,
-        torch.rand((100, 100), dtype=torch.float32).round())
-    d.add_relation_type('Side Effect: Death', 1, 1,
-        torch.rand((100, 100), dtype=torch.float32).round())
+    fam = d.add_relation_family('Gene-Gene', 0, 0, False)
+    fam.add_relation_type('Interaction', 0, 0, torch.rand(1000, 1000))
+
+    fam = d.add_relation_family('Drug-Drug', 1, 1, False)
+    fam.add_relation_type('Side Effect: Nausea', 1, 1, torch.rand(100, 100))
+    fam.add_relation_type('Side Effect: Infertility', 1, 1, torch.rand(100, 100))
+    fam.add_relation_type('Side Effect: Death', 1, 1, torch.rand(100, 100))
     return d