diff --git a/tests/icosagon/test_trainprep.py b/tests/icosagon/test_trainprep.py index 79a8aee..eb78859 100644 --- a/tests/icosagon/test_trainprep.py +++ b/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)