Add test_prep_rel_two_node_types_asym_01().

tests/icosagon/test_trainprep.py

@@ -10,7 +10,8 @@ from icosagon.trainprep import TrainValTest, \
     prepare_adj_mat, \
     prepare_relation_type, \
     prep_rel_one_node_type, \
-    prep_rel_two_node_types_sym
+    prep_rel_two_node_types_sym, \
+    prep_rel_two_node_types_asym
 import torch
 import pytest
 import numpy as np
@@ -168,6 +169,40 @@ def test_prep_rel_two_node_types_sym_01():
     assert len(prep_rel.edges_back_neg.test) == 0
 
 
+def test_prep_rel_two_node_types_asym_01():
+    adj_mat = torch.zeros(200)
+    perm = torch.randperm(100)
+    adj_mat[perm[:10]] = 1
+    adj_mat = adj_mat.view(10, 20)
+
+    adj_mat_back = torch.zeros(200)
+    perm = torch.randperm(100)
+    adj_mat_back[perm[:10]] = 1
+    adj_mat_back = adj_mat_back.view(20, 10)
+
+    rel = RelationType('Dummy Relation', 0, 1, adj_mat, adj_mat_back)
+    ratios = TrainValTest(.8, .1, .1)
+    prep_rel = prep_rel_two_node_types_asym(rel, ratios)
+    assert prep_rel.name == rel.name
+    assert prep_rel.node_type_row == rel.node_type_row
+    assert prep_rel.node_type_column == rel.node_type_column
+    assert prep_rel.adjacency_matrix.shape == rel.adjacency_matrix.shape
+    assert prep_rel.adjacency_matrix_backward.shape == rel.adjacency_matrix_backward.shape
+    assert len(prep_rel.edges_pos.train) == 8
+    assert len(prep_rel.edges_pos.val) == 1
+    assert len(prep_rel.edges_pos.test) == 1
+    assert len(prep_rel.edges_neg.train) == 8
+    assert len(prep_rel.edges_neg.val) == 1
+    assert len(prep_rel.edges_neg.test) == 1
+
+    assert len(prep_rel.edges_back_pos.train) == 8
+    assert len(prep_rel.edges_back_pos.val) == 1
+    assert len(prep_rel.edges_back_pos.test) == 1
+    assert len(prep_rel.edges_back_neg.train) == 8
+    assert len(prep_rel.edges_back_neg.val) == 1
+    assert len(prep_rel.edges_back_neg.test) == 1
+
+
 # def prepare_relation(r, ratios):
 #     adj_mat = r.adjacency_matrix
 #     adj_mat_train, edges_pos, edges_neg = prepare_adj_mat(adj_mat)