def test_to_tsk_tree_errors(self):
alpha_tree = RankTree.unrank((0, 0), 3, ["A", "B", "C"])
out_of_bounds_tree = RankTree.unrank((0, 0), 3, [2, 3, 4])
with self.assertRaises(ValueError):
alpha_tree.to_tsk_tree()
with self.assertRaises(ValueError):
out_of_bounds_tree.to_tsk_tree()
def test_big_trees(self):
n = 14
shape = 22
labelling = 0
tree = RankTree.unrank((shape, labelling), n)
tsk_tree = tskit.Tree.unrank((shape, labelling), n)
self.assertEqual(tree.rank(), tsk_tree.rank())
n = 10
shape = 95
labelling = comb.num_labellings(shape, n) // 2
tree = RankTree.unrank((shape, labelling), n)
tsk_tree = tskit.Tree.unrank((shape, labelling), n)
self.assertEqual(tree.rank(), tsk_tree.rank())
def test_unrank(self):
for n in range(6):
for shape_rank, t in enumerate(RankTree.all_unlabelled_trees(n)):
for label_rank, labelled_tree in enumerate(RankTree.all_labellings(t)):
unranked = RankTree.unrank((shape_rank, label_rank), n)
self.assertTrue(labelled_tree == unranked)
# The number of labelled trees gets very big quickly
for n in range(6, 10):
for shape_rank in range(comb.num_shapes(n)):
rank = (shape_rank, 0)
unranked = RankTree.unrank(rank, n)
self.assertTrue(rank, unranked.rank())
rank = (shape_rank, comb.num_labellings(shape_rank, n) - 1)
unranked = RankTree.unrank(rank, n)
self.assertTrue(rank, unranked.rank())
def test_all_labellings_roundtrip(self):
n = 5
rank_tree = RankTree.unrank((comb.num_shapes(n) - 1, 0), n)
tsk_tree = rank_tree.to_tsk_tree()
rank_tree_labellings = RankTree.all_labellings(rank_tree)
tsk_tree_labellings = tskit.all_tree_labellings(tsk_tree)
for rank_t, tsk_t in zip(rank_tree_labellings, tsk_tree_labellings):
self.assertEqual(rank_t, RankTree.from_tsk_tree(tsk_t))
def verify_unrank_errors(self, rank, n):
with self.assertRaises(ValueError):
RankTree.unrank(rank, n)
with self.assertRaises(ValueError):
tskit.Tree.unrank(rank, n)
def verify_unrank_errors(self, rank, n):
with pytest.raises(ValueError):
RankTree.unrank(rank, n)
with pytest.raises(ValueError):
tskit.Tree.unrank(rank, n)
