def test_gen_lists():
matrix = np.array([[1, 0, 0, 0, 0], [0, 1, 0, 0, 0],
[0, 0, 1, 0, 0], [0, 0, 0, 1, 0], [0, 0, 0, 1, 1]])
exp_pars_k_labels = np.array([0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 3, 3, 4])
exp_pars_t_labels = np.array([1, 2, 3, 4, 4, 2, 3, 4, 4, 3, 4, 4, 4, 4, 4])
exp_host_k_labels = np.array([0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 3, 3, 3])
exp_host_t_labels = np.array([1, 2, 3, 3, 4, 2, 3, 3, 4, 3, 3, 4, 3, 4, 4])
pars, hosts = np.nonzero(matrix)
obs_pars_k_labels, obs_pars_t_labels = _gen_lists(pars)
obs_hosts_k_labels, obs_hosts_t_labels = _gen_lists(hosts)
assert_allclose(exp_pars_k_labels, obs_pars_k_labels)
assert_allclose(exp_pars_t_labels, obs_pars_t_labels)
assert_allclose(exp_host_k_labels, obs_hosts_k_labels)
assert_allclose(exp_host_t_labels, obs_hosts_t_labels)
def test_get_dist():
labels = np.array([0, 1, 1, 2, 3])
k_labels, t_labels = _gen_lists(labels)
dists = np.array([[0, 2, 6, 3], [2, 0, 5, 4], [6, 5, 0, 7], [3, 4, 7, 0]])
index = np.array([2, 3, 1, 0])
expected_vec = np.array([7, 7, 5, 6, 0, 4, 3, 4, 3, 2])
actual_vec = _get_dist(k_labels, t_labels, dists, index)
assert_allclose(actual_vec, expected_vec)
