Implement algorithm for calculating species tree topology distributions.

Author: daniel-goldstein

python/tests/test_combinatorics.py

@@ -23,8 +23,13 @@
 """
 Test cases for combinatorial algorithms.
 """
+import io
 import itertools
 import unittest
+from collections import Counter
+from collections import defaultdict
+
+import msprime
 
 import tskit
 import tskit.combinatorics as comb
@@ -413,3 +418,166 @@ def test_is_symmetrical(self):
         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())
+
+
+class TestCountTopologies(unittest.TestCase):
+    def verify_topologies(self, ts, expected=None):
+        populations = [pop.id for pop in ts.populations()]
+        topologies = [comb.count_topologies(t) for t in ts.trees()]
+        inc_topologies = list(comb.count_topologies_incremental(ts))
+        for num_pops in range(1, ts.num_populations + 1):
+            for i, t in enumerate(ts.trees()):
+                just_t = ts.keep_intervals([t.interval])
+                for pops in itertools.combinations(populations, num_pops):
+                    actual_topologies = topologies[i][frozenset(pops)]
+                    actual_inc_topologies = inc_topologies[i][frozenset(pops)]
+                    if len(t.roots) == 1:
+                        subsampled_topologies = self.subsample_topologies(just_t, pops)
+                        self.assertEqual(actual_topologies, subsampled_topologies)
+                    if expected is not None:
+                        self.assertEqual(
+                            actual_topologies, expected[i][frozenset(pops)]
+                        )
+                    self.assertEqual(actual_topologies, actual_inc_topologies)
+
+    def subsample_topologies(self, ts, populations):
+        samples_per_pop = [ts.samples(population=p) for p in populations]
+        topologies = Counter()
+        for subsample in itertools.product(*samples_per_pop):
+            for pop_tree in ts.simplify(samples=subsample).trees():
+                # regions before and after keep interval have all samples as roots
+                # so don't count those
+                # The single tree of interest should have one root
+                if len(pop_tree.roots) == 1:
+                    topologies[pop_tree.rank()] += 1
+        return topologies
+
+    def test_single_population(self):
+        n = 10
+        ts = msprime.simulate(n, recombination_rate=10)
+        expected = defaultdict(Counter)
+        expected[frozenset([0])] = Counter({(0, 0): n})
+        self.verify_topologies(ts, [expected] * ts.num_trees)
+
+    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 = defaultdict(Counter)
+        expected[frozenset([0])] = Counter({(0, 0): 2})
+        expected[frozenset([1])] = Counter({(0, 0): 2})
+        expected[frozenset([2])] = Counter({(0, 0): 2})
+        expected[frozenset([0, 1])] = Counter({(0, 0): 4})
+        expected[frozenset([0, 2])] = Counter({(0, 0): 4})
+        expected[frozenset([1, 2])] = Counter({(0, 0): 4})
+        expected[frozenset([0, 1, 2])] = Counter({(1, 0): 4, (1, 1): 4})
+        self.verify_topologies(ts, [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 = defaultdict(Counter)
+        expected[frozenset([0])] = Counter({(0, 0): 1})
+        expected[frozenset([1])] = Counter({(0, 0): 1})
+        self.verify_topologies(tables.tree_sequence(), [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 = defaultdict(Counter)
+        for pop_combo in [[0], [1], [2], [0, 1]]:
+            expected[frozenset(pop_combo)] = Counter({(0, 0): 1})
+        self.verify_topologies(tables.tree_sequence(), [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 = defaultdict(Counter)
+        for pop_combos in [[0], [1], [2], [0, 1], [0, 2], [1, 2], [0, 1, 2]]:
+            expected[frozenset(pop_combos)] = Counter({(0, 0): 1})
+        self.verify_topologies(tables.tree_sequence(), [expected])
+
+    def test_msprime_migrations(self):
+        for num_populations in range(2, 5):
+            samples = [5] * num_populations
+            ts = self.simulate_multiple_populations(samples)
+            self.verify_topologies(ts)
+
+    def simulate_multiple_populations(self, sample_sizes):
+        d = len(sample_sizes)
+        M = 0.2
+        m = M / (2 * (d - 1))
+
+        migration_matrix = [
+            [m if k < d and k == i + 1 else 0 for k in range(d)] for i in range(d)
+        ]
+
+        pop_configurations = [
+            msprime.PopulationConfiguration(sample_size=size) for size in sample_sizes
+        ]
+        return msprime.simulate(
+            population_configurations=pop_configurations,
+            migration_matrix=migration_matrix,
+            recombination_rate=0.05,
+        )

python/tskit/combinatorics.py

@@ -26,11 +26,155 @@
 """
 import heapq
 import itertools
+from collections import Counter
+from collections import defaultdict
 from functools import lru_cache
 
+import numpy as np
+
 import tskit
 
 
+def count_topologies_incremental(ts, key=None):
+    def get_population(u):
+        return ts.node(u).population
+
+    if key is None:
+        key = get_population
+
+    topology_tree = [None for _ in range(ts.num_nodes)]
+    parent = np.full(ts.num_nodes, -1)
+
+    def update_state(tree, u):
+        stack = [u]
+        while len(stack) > 0:
+            v = stack.pop()
+            children = []
+            c = tree.left_child(v)
+            while c != tskit.NULL:
+                if topology_tree[c] is not None:
+                    children.append(topology_tree[c])
+                c = tree.right_sib(c)
+            if len(children) > 0:
+                topology_tree[v] = TopologyTree(children)
+            else:
+                topology_tree[v] = None
+            p = parent[v]
+            if p != -1:
+                stack.append(p)
+
+    for u in ts.samples():
+        topology_tree[u] = TopologyTree(children=[], label=key(u))
+
+    for tree, (_, edges_out, edges_in) in zip(ts.trees(), ts.edge_diffs()):
+        # Avoid recomputing anything for the parent until all child edges
+        # for that parent are inserted/removed
+        for p, sibling_edges in itertools.groupby(edges_out, key=lambda e: e.parent):
+            for e in sibling_edges:
+                parent[e.child] = -1
+            update_state(tree, p)
+        for p, sibling_edges in itertools.groupby(edges_in, key=lambda e: e.parent):
+            for e in sibling_edges:
+                parent[e.child] = p
+            update_state(tree, p)
+
+        all_topologies = defaultdict(Counter)
+        for root in tree.roots:
+            # None if there are no samples under the root
+            if topology_tree[root] is not None:
+                for labels, counter in topology_tree[root].topologies.items():
+                    all_topologies[labels] += counter
+        yield all_topologies
+
+
+def count_topologies(tree, key=None):
+    def pop_key(u):
+        return tree.tree_sequence.node(u).population
+
+    if key is None:
+        key = pop_key
+
+    all_topologies = defaultdict(Counter)
+    for root in tree.roots:
+        topo_tree = TopologyTree.from_tsk_tree(tree, root, key)
+        for labels, counter in topo_tree.topologies.items():
+            all_topologies[labels] += counter
+
+    return all_topologies
+
+
+class TopologyTree:
+    def __init__(self, children, label=None):
+        self.children = children
+        if len(children) == 0:
+            assert label is not None
+        self.label = label
+
+        self.topologies = self._topologies()
+
+    def _topologies(self):
+        if self.is_leaf():
+            rank_tree = RankTree(children=[], label=self.label)
+            labels = frozenset([self.label])
+            return {labels: Counter({rank_tree.rank(): 1})}
+
+        # Create partial trees (incomplete sets of child nodes)
+        partials = defaultdict(Counter)
+        for child in self.children:
+            merged = TopologyTree.merge_child_topologies(partials, child)
+            for labels, counter in merged.items():
+                partials[labels] += counter
+
+        # Join sets of child nodes together and add to existing topologies
+        topologies = defaultdict(Counter)
+        for root_labels, sibling_topologies in partials.items():
+            for k, count in sibling_topologies.items():
+                # A node must have at least two children
+                if len(k) >= 2:
+                    children = []
+                    for labels, rank in k:
+                        labels_list = list(sorted(labels))
+                        n = len(labels_list)
+                        children.append(RankTree.unrank(rank, n, labels_list))
+                    children.sort(key=RankTree.canonical_order)
+                    tree = RankTree(children)
+                    topologies[root_labels][tree.rank()] += count
+                else:
+                    # Pass on the single tree without adding a parent
+                    for _, rank in k:
+                        topologies[root_labels][rank] += count
+
+        return topologies
+
+    @staticmethod
+    def merge_child_topologies(partials, child):
+        merged = defaultdict(Counter)
+        for labels, topologies in child.topologies.items():
+            for rank, count in topologies.items():
+                topology = frozenset([(labels, rank)])
+                for sibling_labels, children in partials.items():
+                    if labels.isdisjoint(sibling_labels):
+                        for sib_topologies, sib_count in children.items():
+                            merged_topologies = sib_topologies.union(topology)
+                            merged_labels = sibling_labels.union(labels)
+                            merged[merged_labels][merged_topologies] += (
+                                count * sib_count
+                            )
+                merged[labels][topology] += count
+        return merged
+
+    def is_leaf(self):
+        return len(self.children) == 0
+
+    @staticmethod
+    def from_tsk_tree(tree, u, key=None):
+        if tree.is_leaf(u):
+            return TopologyTree([], label=key(u))
+
+        children = [TopologyTree.from_tsk_tree(tree, c, key) for c in tree.children(u)]
+        return TopologyTree(children=children)
+
+
 def all_trees(num_leaves):
     """
     Generates all unique leaf-labelled trees with ``num_leaves``
@@ -210,12 +354,12 @@ def label_rank(self):
         return self._label_rank
 
     @staticmethod
-    def unrank(rank, num_leaves):
+    def unrank(rank, num_leaves, labels=None):
         shape_rank, label_rank = rank
         if shape_rank < 0 or label_rank < 0:
             raise ValueError("Rank is out of bounds.")
         unlabelled = RankTree.shape_unrank(shape_rank, num_leaves)
-        return unlabelled.label_unrank(label_rank)
+        return unlabelled.label_unrank(label_rank, labels)
 
     @staticmethod
     def shape_unrank(shape_rank, n):