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)