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_no_sample_subtrees(self):
tables = tskit.TableCollection(sequence_length=1.0)
c1 = tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
c2 = tables.nodes.add_row(time=0)
c3 = tables.nodes.add_row(time=0)
p1 = tables.nodes.add_row(time=1)
p2 = tables.nodes.add_row(time=1)
tables.edges.add_row(left=0, right=1, parent=p1, child=c2)
tables.edges.add_row(left=0, right=1, parent=p1, child=c3)
tables.edges.add_row(left=0, right=1, parent=p2, child=c1)
expected = comb.TopologyCounter()
expected[0] = collections.Counter({(0, 0): 1})
self.verify_topologies(tables.tree_sequence(), expected=[expected])
def test_multiple_roots(self):
tables = tskit.TableCollection(sequence_length=1.0)
tables.populations.add_row()
tables.populations.add_row()
tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, population=0)
tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, population=1)
# Not samples so they are ignored
tables.nodes.add_row(time=1)
tables.nodes.add_row(time=1, population=1)
expected = comb.TopologyCounter()
expected[0] = collections.Counter({(0, 0): 1})
expected[1] = collections.Counter({(0, 0): 1})
self.verify_topologies(tables.tree_sequence(), expected=[expected])
def test_ignores_non_sample_leaves(self):
nodes = io.StringIO(
"""\
id is_sample time population individual metadata
0 1 0.000000 0 -1
1 0 0.000000 0 -1
2 1 0.000000 0 -1
3 0 0.000000 0 -1
4 1 0.000000 0 -1
5 0 1.000000 0 -1
6 0 1.000000 0 -1
7 0 2.000000 0 -1
8 0 3.000000 0 -1
"""
)
edges = io.StringIO(
"""\
left right parent child
0.000000 1.000000 5 0
0.000000 1.000000 5 1
0.000000 1.000000 6 2
0.000000 1.000000 6 3
0.000000 1.000000 7 5
0.000000 1.000000 7 6
0.000000 1.000000 8 4
0.000000 1.000000 8 7
"""
)
ts = tskit.load_text(
nodes, edges, sequence_length=1, strict=False, base64_metadata=False
)
sample_sets = [[0], [2], [4]]
expected = comb.TopologyCounter()
expected[0] = collections.Counter({(0, 0): 1})
expected[1] = collections.Counter({(0, 0): 1})
expected[2] = collections.Counter({(0, 0): 1})
expected[0, 1] = collections.Counter({(0, 0): 1})
expected[0, 2] = collections.Counter({(0, 0): 1})
expected[1, 2] = collections.Counter({(0, 0): 1})
expected[0, 1, 2] = collections.Counter({(1, 2): 1})
tree_topologies = ts.first().count_topologies(sample_sets)
treeseq_topologies = list(ts.count_topologies(sample_sets))
self.assertEqual(tree_topologies, expected)
self.assertEqual(treeseq_topologies, [expected])
def test_polytomies(self):
tables = tskit.TableCollection(sequence_length=1.0)
tables.populations.add_row()
tables.populations.add_row()
tables.populations.add_row()
c1 = tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, population=0)
c2 = tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, population=1)
c3 = tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, population=2)
p = tables.nodes.add_row(time=1)
tables.edges.add_row(left=0, right=1, parent=p, child=c1)
tables.edges.add_row(left=0, right=1, parent=p, child=c2)
tables.edges.add_row(left=0, right=1, parent=p, child=c3)
expected = comb.TopologyCounter()
for pop_combos in [0, 1, 2, (0, 1), (0, 2), (1, 2), (0, 1, 2)]:
expected[pop_combos] = collections.Counter({(0, 0): 1})
self.verify_topologies(tables.tree_sequence(), expected=[expected])
def test_three_populations(self):
nodes = io.StringIO(
"""\
id is_sample time population individual metadata
0 1 0.000000 0 -1
1 1 0.000000 1 -1
2 1 0.000000 1 -1
3 1 0.000000 2 -1
4 1 0.000000 2 -1
5 1 0.000000 0 -1
6 0 1.000000 0 -1
7 0 2.000000 0 -1
8 0 2.000000 0 -1
9 0 3.000000 0 -1
10 0 4.000000 0 -1
"""
)
edges = io.StringIO(
"""\
left right parent child
0.000000 1.000000 6 4
0.000000 1.000000 6 5
0.000000 1.000000 7 1
0.000000 1.000000 7 2
0.000000 1.000000 8 3
0.000000 1.000000 8 6
0.000000 1.000000 9 7
0.000000 1.000000 9 8
0.000000 1.000000 10 0
0.000000 1.000000 10 9
"""
)
ts = tskit.load_text(
nodes, edges, sequence_length=1, strict=False, base64_metadata=False
)
expected = comb.TopologyCounter()
expected[0] = collections.Counter({(0, 0): 2})
expected[1] = collections.Counter({(0, 0): 2})
expected[2] = collections.Counter({(0, 0): 2})
expected[0, 1] = collections.Counter({(0, 0): 4})
expected[0, 2] = collections.Counter({(0, 0): 4})
expected[1, 2] = collections.Counter({(0, 0): 4})
expected[0, 1, 2] = collections.Counter({(1, 0): 4, (1, 1): 4})
self.verify_topologies(ts, expected=[expected])
def test_no_full_topology(self):
tables = tskit.TableCollection(sequence_length=1.0)
tables.populations.add_row()
tables.populations.add_row()
tables.populations.add_row()
child1 = tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, population=0)
child2 = tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, population=1)
parent = tables.nodes.add_row(time=1)
tables.edges.add_row(left=0, right=1, parent=parent, child=child1)
tables.edges.add_row(left=0, right=1, parent=parent, child=child2)
# Left as root so there is no topology with all three populations
tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, population=2)
expected = comb.TopologyCounter()
for pop_combo in [(0,), (1,), (2,), (0, 1)]:
expected[pop_combo] = collections.Counter({(0, 0): 1})
self.verify_topologies(tables.tree_sequence(), expected=[expected])
def test_single_population(self):
n = 10
ts = msprime.simulate(n, recombination_rate=10)
expected = comb.TopologyCounter()
expected[0] = collections.Counter({(0, 0): n})
self.verify_topologies(ts, expected=[expected] * ts.num_trees)
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)