def test_all_shapes_roundtrip(self):
n = 5
all_rank_tree_shapes = RankTree.all_unlabelled_trees(n)
all_tsk_tree_shapes = tskit.all_tree_shapes(n)
for rank_tree, tsk_tree in zip(all_rank_tree_shapes,
all_tsk_tree_shapes):
assert rank_tree.shape_equal(RankTree.from_tsk_tree(tsk_tree))
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_num_labellings(self):
for n in range(2, 8):
for tree in RankTree.all_unlabelled_trees(n):
tree = tree.label_unrank(0)
tree2 = tree.to_tsk_tree()
n_labellings = sum(1 for _ in RankTree.all_labellings(tree))
assert n_labellings == RankTree.from_tsk_tree(
tree2).num_labellings()
def test_to_from_tsk_tree(self):
for n in range(5):
for tree in RankTree.all_labelled_trees(n):
self.assertTrue(tree.is_canonical())
tsk_tree = tree.to_tsk_tree()
reconstructed = RankTree.from_tsk_tree(tsk_tree)
self.assertTrue(tree.is_canonical())
self.assertEqual(tree, reconstructed)
def test_to_from_tsk_tree(self):
for n in range(5):
for tree in RankTree.all_labelled_trees(n):
assert tree.is_canonical()
tsk_tree = tree.to_tsk_tree()
reconstructed = RankTree.from_tsk_tree(tsk_tree)
assert tree.is_canonical()
assert tree == reconstructed
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 test_from_unary_tree(self):
tables = tskit.TableCollection(sequence_length=1)
c = tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
p = tables.nodes.add_row(time=1)
tables.edges.add_row(left=0, right=1, parent=p, child=c)
t = tables.tree_sequence().first()
with self.assertRaises(ValueError):
RankTree.from_tsk_tree(t)
def test_label_unrank(self):
for n in range(7):
for shape_rank, tree in enumerate(RankTree.all_unlabelled_trees(n)):
for label_rank, labelled_tree in enumerate(
RankTree.all_labellings(tree)
):
rank = (shape_rank, label_rank)
unranked = tree.label_unrank(label_rank)
self.assertEqual(labelled_tree.rank(), rank)
self.assertEqual(unranked.rank(), rank)
def test_unranking_is_canonical(self):
for n in range(7):
for shape_rank in range(comb.num_shapes(n)):
for label_rank in range(comb.num_labellings(shape_rank, n)):
t = RankTree.shape_unrank(shape_rank, n)
self.assertTrue(t.is_canonical())
t = t.label_unrank(label_rank)
self.assertTrue(t.is_canonical())
t = tskit.Tree.unrank((shape_rank, label_rank), n)
self.assertTrue(RankTree.from_tsk_tree(t).is_canonical())
def test_is_symmetrical(self):
unlabelled_leaf = RankTree(children=[])
assert unlabelled_leaf.is_symmetrical()
three_leaf_asym = RankTree(children=[
unlabelled_leaf,
RankTree(children=[unlabelled_leaf, unlabelled_leaf]),
])
assert not three_leaf_asym.is_symmetrical()
six_leaf_sym = RankTree(children=[three_leaf_asym, three_leaf_asym])
assert six_leaf_sym.is_symmetrical()
def test_is_symmetrical(self):
unlabelled_leaf = RankTree(children=[])
self.assertTrue(unlabelled_leaf.is_symmetrical())
three_leaf_asym = RankTree(children=[
unlabelled_leaf,
RankTree(children=[unlabelled_leaf, unlabelled_leaf]),
])
self.assertFalse(three_leaf_asym.is_symmetrical())
six_leaf_sym = RankTree(children=[three_leaf_asym, three_leaf_asym])
self.assertTrue(six_leaf_sym.is_symmetrical())
def test_shape_unrank(self):
for n in range(6):
for rank, tree in enumerate(RankTree.all_unlabelled_trees(n)):
t = RankTree.shape_unrank(rank, n)
self.assertTrue(tree.shape_equal(t))
for n in range(2, 9):
for shape_rank, tree in enumerate(RankTree.all_unlabelled_trees(n)):
tsk_tree = tskit.Tree.unrank((shape_rank, 0), n)
self.assertEqual(shape_rank, tree.shape_rank())
shape_rank, _ = tsk_tree.rank()
self.assertEqual(shape_rank, tree.shape_rank())
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_generate_treeseq_roundtrip(self):
n = 5
span = 9
all_rank_trees = RankTree.all_labelled_trees(n)
all_tsk_trees = tskit.all_trees(n, span=span)
for rank_tree, tsk_tree in zip(all_rank_trees, all_tsk_trees):
ts1 = tsk_tree.tree_sequence
ts2 = rank_tree.to_tsk_tree(span=span).tree_sequence
assert ts1.tables.equals(ts2.tables, ignore_provenance=True)
def test_unrank_rank_round_trip(self):
for n in range(6): # Can do more but gets slow pretty quickly after 6
for shape_rank in range(comb.num_shapes(n)):
tree = RankTree.shape_unrank(shape_rank, n)
tree = tree.label_unrank(0)
self.assertEqual(tree.shape_rank(), shape_rank)
for label_rank in range(tree.num_labellings()):
tree = tree.label_unrank(label_rank)
self.assertEqual(tree.label_rank(), label_rank)
tsk_tree = tree.label_unrank(label_rank).to_tsk_tree()
_, tsk_label_rank = tsk_tree.rank()
self.assertEqual(tsk_label_rank, label_rank)
def test_all_labelled_trees_4(self):
expected = [
# 1 + 1 + 1 + 1 (partition of num leaves)
"(0,1,2,3)",
# 1 + 1 + 2
"(0,1,(2,3))",
"(0,2,(1,3))",
"(0,3,(1,2))",
"(1,2,(0,3))",
"(1,3,(0,2))",
"(2,3,(0,1))",
# 1 + 3
# partition of 3 = 1 + 1 + 1
"(0,(1,2,3))",
"(1,(0,2,3))",
"(2,(0,1,3))",
"(3,(0,1,2))",
# partition of 3 = 1 + 2
"(0,(1,(2,3)))",
"(0,(2,(1,3)))",
"(0,(3,(1,2)))",
"(1,(0,(2,3)))",
"(1,(2,(0,3)))",
"(1,(3,(0,2)))",
"(2,(0,(1,3)))",
"(2,(1,(0,3)))",
"(2,(3,(0,1)))",
"(3,(0,(1,2)))",
"(3,(1,(0,2)))",
"(3,(2,(0,1)))",
# 2 + 2
"((0,1),(2,3))",
"((0,2),(1,3))",
"((0,3),(1,2))",
]
actual = [t.newick() for t in RankTree.all_labelled_trees(4)]
self.assertEqual(expected, actual)
def test_add_sibling_topologies_simple(self):
a = RankTree(children=[], label="A")
b = RankTree(children=[], label="B")
ab = RankTree(children=[a, b])
a_counter = comb.TopologyCounter()
a_counter["A"][a.rank()] = 1
self.assertEqual(a_counter, comb.TopologyCounter.from_sample("A"))
b_counter = comb.TopologyCounter()
b_counter["B"][b.rank()] = 1
self.assertEqual(b_counter, comb.TopologyCounter.from_sample("B"))
partial_counter = comb.PartialTopologyCounter()
partial_counter.add_sibling_topologies(a_counter)
partial_counter.add_sibling_topologies(b_counter)
expected = comb.TopologyCounter()
expected["A"][a.rank()] = 1
expected["B"][b.rank()] = 1
expected["A", "B"][ab.rank()] = 1
joined_counter = partial_counter.join_all_combinations()
self.assertEqual(joined_counter, expected)
def test_equal(self):
unlabelled_leaf = RankTree(children=[])
self.assertEqual(unlabelled_leaf, unlabelled_leaf)
self.assertTrue(unlabelled_leaf.shape_equal(unlabelled_leaf))
leaf_zero = RankTree(children=[], label=0)
leaf_one = RankTree(children=[], label=1)
leaf_two = RankTree(children=[], label=2)
self.assertEqual(leaf_zero, leaf_zero)
self.assertNotEqual(leaf_zero, leaf_one)
self.assertTrue(leaf_zero.shape_equal(leaf_one))
tree1 = RankTree(children=[leaf_zero, leaf_one])
self.assertEqual(tree1, tree1)
self.assertNotEqual(tree1, unlabelled_leaf)
self.assertFalse(tree1.shape_equal(unlabelled_leaf))
tree2 = RankTree(children=[leaf_two, leaf_one])
self.assertNotEqual(tree1, tree2)
self.assertTrue(tree1.shape_equal(tree2))
def test_symmetrical_trees(self):
for n in range(2, 18, 2):
last_rank = comb.num_shapes(n) - 1
t = RankTree.shape_unrank(last_rank, n)
self.assertTrue(t.is_symmetrical())
def test_num_shapes(self):
for i in range(11):
all_trees = RankTree.all_unlabelled_trees(i)
self.assertEqual(len(list(all_trees)), comb.num_shapes(i))
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 test_all_labelled_trees_3(self):
expected = ["(0,1,2)", "(0,(1,2))", "(1,(0,2))", "(2,(0,1))"]
actual = [t.newick() for t in RankTree.all_labelled_trees(3)]
self.assertEqual(expected, actual)
def test_is_canonical(self):
for n in range(7):
for tree in RankTree.all_labelled_trees(n):
self.assertTrue(tree.is_canonical())
shape_not_canonical = RankTree(
children=[
RankTree(children=[], label=0),
RankTree(
children=[
RankTree(
children=[
RankTree(children=[], label=1),
RankTree(children=[], label=2),
]
),
RankTree(children=[], label=3),
]
),
]
)
self.assertFalse(shape_not_canonical.is_canonical())
labels_not_canonical = RankTree(
children=[
RankTree(children=[], label=0),
RankTree(
children=[
RankTree(
children=[
RankTree(children=[], label=2),
RankTree(children=[], label=3),
]
),
RankTree(
children=[
RankTree(children=[], label=1),
RankTree(children=[], label=4),
]
),
]
),
]
)
self.assertFalse(labels_not_canonical.is_canonical())
def test_label_rank(self):
for n in range(7):
for tree in RankTree.all_unlabelled_trees(n):
for rank, labelled_tree in enumerate(RankTree.all_labellings(tree)):
self.assertEqual(labelled_tree.label_rank(), rank)
def test_generate_trees_roundtrip(self):
n = 5
all_rank_trees = RankTree.all_labelled_trees(n)
all_tsk_trees = tskit.all_trees(n)
for rank_tree, tsk_tree in zip(all_rank_trees, all_tsk_trees):
self.assertEqual(rank_tree, RankTree.from_tsk_tree(tsk_tree))
def test_shape_rank(self):
for n in range(10):
for rank, tree in enumerate(RankTree.all_unlabelled_trees(n)):
self.assertEqual(tree.shape_rank(), rank)
def test_add_sibling_topologies_polytomy(self):
"""
Goes through the topology-merging step at the root
of this tree:
|
|
+----+-----+----+
| | | |
| | | |
| | | +---+
| | | | |
| | | | |
A A B A C
"""
partial_counter = comb.PartialTopologyCounter()
a = RankTree(children=[], label="A")
c = RankTree(children=[], label="C")
ac = RankTree(children=[a, c])
expected = collections.defaultdict(collections.Counter)
a_counter = comb.TopologyCounter.from_sample("A")
b_counter = comb.TopologyCounter.from_sample("B")
ac_counter = comb.TopologyCounter()
ac_counter["A"][a.rank()] = 1
ac_counter["C"][c.rank()] = 1
ac_counter["A", "C"][ac.rank()] = 1
partial_counter.add_sibling_topologies(a_counter)
expected[("A",)] = collections.Counter({((("A",), (0, 0)),): 1})
self.assertEqual(partial_counter.partials, expected)
partial_counter.add_sibling_topologies(a_counter)
expected[("A",)][((("A",), (0, 0)),)] += 1
self.assertEqual(partial_counter.partials, expected)
partial_counter.add_sibling_topologies(b_counter)
expected[("B",)][((("B",), (0, 0)),)] = 1
expected[("A", "B")][((("A",), (0, 0)), (("B",), (0, 0)))] = 2
self.assertEqual(partial_counter.partials, expected)
partial_counter.add_sibling_topologies(ac_counter)
expected[("A",)][((("A",), (0, 0)),)] += 1
expected[("C",)][((("C",), (0, 0)),)] = 1
expected[("A", "B")][((("A",), (0, 0)), (("B",), (0, 0)))] += 1
expected[("A", "C")][((("A",), (0, 0)), (("C",), (0, 0)))] = 2
expected[("A", "C")][((("A", "C"), (0, 0)),)] = 1
expected[("B", "C")][((("B",), (0, 0)), (("C",), (0, 0)))] = 1
expected[("A", "B", "C")][
((("A",), (0, 0)), (("B",), (0, 0)), (("C",), (0, 0)))
] = 2
expected[("A", "B", "C")][((("A", "C"), (0, 0)), (("B",), (0, 0)))] = 1
self.assertEqual(partial_counter.partials, expected)
expected_topologies = comb.TopologyCounter()
expected_topologies["A"][(0, 0)] = 3
expected_topologies["B"][(0, 0)] = 1
expected_topologies["C"][(0, 0)] = 1
expected_topologies["A", "B"][(0, 0)] = 3
expected_topologies["A", "C"][(0, 0)] = 3
expected_topologies["B", "C"][(0, 0)] = 1
expected_topologies["A", "B", "C"][(0, 0)] = 2
expected_topologies["A", "B", "C"][(1, 1)] = 1
joined_topologies = partial_counter.join_all_combinations()
self.assertEqual(joined_topologies, expected_topologies)
def test_join_topologies(self):
a = RankTree(children=[], label="A")
b = RankTree(children=[], label="B")
c = RankTree(children=[], label="C")
a_tuple = (("A"), a.rank())
b_tuple = (("B"), b.rank())
c_tuple = (("C"), c.rank())
ab_tuple = (("A", "B"), RankTree(children=[a, b]).rank())
ac_tuple = (("A", "C"), RankTree(children=[a, c]).rank())
bc_tuple = (("B", "C"), RankTree(children=[b, c]).rank())
self.verify_join_topologies((a_tuple, b_tuple), (0, 0))
self.verify_join_topologies((b_tuple, a_tuple), (0, 0))
self.verify_join_topologies((b_tuple, c_tuple), (0, 0))
self.verify_join_topologies((a_tuple, b_tuple, c_tuple), (0, 0))
self.verify_join_topologies((a_tuple, bc_tuple), (1, 0))
self.verify_join_topologies((b_tuple, ac_tuple), (1, 1))
self.verify_join_topologies((c_tuple, ab_tuple), (1, 2))
def test_num_labelled_trees(self):
# Number of leaf-labelled trees with n leaves on OEIS
n_trees = [0, 1, 1, 4, 26, 236, 2752, 39208]
for i, expected in zip(range(len(n_trees)), n_trees):
actual = sum(1 for _ in RankTree.all_labelled_trees(i))
self.assertEqual(actual, expected)
