First draft at polytomy breaking function

Author: hyanwong

python/CHANGELOG.rst

@@ -4,6 +4,9 @@
 
 **Features**
 
+- Add ``randomly_split_polytomies`` methods for tables and tree sequences
+  (:user:`hyanwong`, :issue:`809`)
+
 - Genomic intervals returned by python functions are now namedtuples, allowing ``.left``
   ``.right`` and ``.span`` usage (:user:`hyanwong`, :issue:`784`, :pr:`786`, :pr:`811`)
 

python/tests/test_topology.py

@@ -7857,3 +7857,56 @@ def test_failure_with_migrations(self):
         self.assertRaises(ValueError, ts.ltrim)
         self.assertRaises(ValueError, ts.rtrim)
         self.assertRaises(ValueError, ts.trim)
+
+
+class TestPolytomySplitting(unittest.TestCase):
+    "Test the ability to randomly split polytomies"
+    nodes_polytomy_4 = """\
+    id      is_sample   population      time
+    0       1       0               0.00000000000000
+    1       1       0               0.00000000000000
+    2       1       0               0.00000000000000
+    3       1       0               0.00000000000000
+    4       0       0               1.00000000000000
+    """
+    edges_polytomy_4 = """\
+    id      left            right           parent  child
+    0       0.00000000      1.00000000      4       0,1,2,3
+    """
+    # def test_equiprobable(self):
+    #    pass
+
+    def test_nonbinary_tree(self):
+        ts = tskit.load_text(
+            nodes=io.StringIO(self.nodes_polytomy_4),
+            edges=io.StringIO(self.edges_polytomy_4),
+            strict=False,
+        )
+        ts_binary1 = ts.randomly_split_polytomies(squash_edges=True, random_seed=123)
+        # Should be no extra edges to squash
+        ts_binary2 = ts.randomly_split_polytomies(squash_edges=False, random_seed=123)
+        self.assertTrue(ts_equal(ts_binary1, ts_binary2))
+        for tree in ts_binary1.trees():
+            for node in tree.nodes():
+                assert tree.num_children(node) < 3
+
+    def test_nonbinary_ts(self):
+        demographic_events = [
+            msprime.SimpleBottleneck(time=1.0, population=0, proportion=0.95)
+        ]
+        ts = msprime.simulate(
+            20,
+            recombination_rate=10,
+            mutation_rate=5,
+            demographic_events=demographic_events,
+            random_seed=7,
+        )
+        n_poly = 0
+        for e in ts.edgesets():
+            if len(e.children) > 2:
+                n_poly += 1
+        self.assertGreater(n_poly, 3)
+        ts_binary = ts.randomly_split_polytomies(random_seed=123)
+        for tree in ts_binary.trees():
+            for node in tree.nodes():
+                assert tree.num_children(node) < 3

python/tskit/tables.py

@@ -24,6 +24,7 @@
 Tree sequence IO via the tables API.
 """
 import base64
+import collections
 import datetime
 import itertools
 import json
@@ -2781,3 +2782,204 @@ def union(
             self.provenances.add_row(
                 record=json.dumps(provenance.get_provenance_dict(parameters))
             )
+
+    def randomly_split_polytomies(
+        self,
+        *,
+        epsilon=None,
+        squash_edges=True,
+        record_provenance=True,
+        random_seed=None,
+    ):
+        """
+        Modifies the table collection in place, adding extra nodes and edges
+        so that any node with greater than 2 children (i.e. a multifurcation
+        or "polytomy") is resolved into successive bifurcations. This is identical
+        to :meth:`TreeSequence.randomly_split_polytomies` but acts *in place* to
+        alter the data in this :class:`TableCollection`. Please see
+        :meth:`TreeSequence.randomly_split_polytomies` for a fuller description,
+        and details of parameters.
+        """
+        if epsilon is None:
+            epsilon = 1e-10
+        rng = np.random.default_rng(seed=random_seed)
+
+        def is_unknown_time_array(a):
+            np_unknown_time = np.float64(UNKNOWN_TIME)
+            return a.view(np.uint64) == np_unknown_time.view(np.uint64)
+
+        def resolve_polytomy(parent_node_id, child_ids, new_nodes_by_time_desc):
+            """
+            For a polytomy and list of child node ids, return a list of (child, parent)
+            tuples, describing a bifurcating tree, rooted at parent_node_id, where the
+            new_nodes_by_time_desc have been used to break polytomies. All possible
+            topologies should be equiprobable.
+            """
+            nonlocal rng
+            assert len(child_ids) == len(new_nodes_by_time_desc) + 2
+            # Polytomies broken by sequentially splicing onto edges, so an initial edge
+            # is required. This will always remain above the top node & is removed later
+            edges = [
+                [child_ids[0], None],
+            ]
+            # We know beforehand how many random ints are needed: generate them all now
+            edge_choice = rng.integers(0, np.arange(1, len(child_ids) * 2 - 1, 2))
+            tmp_new_node_lab = [parent_node_id] + new_nodes_by_time_desc
+            assert len(edge_choice) == len(child_ids) - 1
+            for node_lab, child_id, target_edge_id in zip(
+                tmp_new_node_lab, child_ids[1:], edge_choice
+            ):
+                target_edge = edges[target_edge_id]
+                # Insert in the right place, to keep edges in parent time order
+                edges.insert(target_edge_id, [child_id, node_lab])
+                edges.insert(target_edge_id, [target_edge[0], node_lab])
+                target_edge[0] = node_lab
+            top_edge = edges.pop()  # remove the edge above the top node
+            assert top_edge[1] is None
+
+            # Re-map the internal nodes IDs so they are used in time order
+            real_node = iter(new_nodes_by_time_desc)
+            node_map = {c: c for c in child_ids}
+            node_map[edges[-1][1]] = parent_node_id  # last edge == oldest parent
+            for e in reversed(edges):
+                # Reversing along the edges, parents are in inverse time order
+                for idx in (1, 0):  # look at parent (1) then child (0)
+                    if e[idx] not in node_map:
+                        node_map[e[idx]] = next(real_node)
+                    e[idx] = node_map[e[idx]]
+            assert len(node_map) == len(new_nodes_by_time_desc) + len(child_ids) + 1
+            return edges
+
+        edge_table = self.edges
+        node_table = self.nodes
+        # Store existing left, so we can change it if the edge is split
+        existing_edges_left = edge_table.left
+        # Keep other edge arrays etc. for fast read access
+        existing_edges_right = edge_table.right
+        existing_edges_parent = edge_table.parent
+        existing_edges_child = edge_table.child
+        existing_node_time = node_table.time
+
+        # We can save a lot of effort if we don't need to check the time of mutations
+        # We definitely don't need to check on the first iteration, a
+        check_mutations = np.any(
+            np.logical_not(is_unknown_time_array(self.mutations.time))
+        )
+        ts = self.tree_sequence()  # Only needed to check mutations
+        tree_iter = ts.trees()  # ditto
+
+        edge_table.clear()
+
+        edges_from_node = collections.defaultdict(set)  # Active descendant edge ids
+        nodes_changed = set()
+
+        for interval, e_out, e_in in ts.edge_diffs(include_terminal=True):
+            pos = interval[0]
+            prev_tree = None if pos == 0 else next(tree_iter)
+
+            for edge in itertools.chain(e_out, e_in):
+                if edge.parent != tskit.NULL:
+                    nodes_changed.add(edge.parent)
+
+            if check_mutations and prev_tree is not None:
+                # This is grim. There must be a more efficient way.
+                # It would also help if mutations were sorted such that all mutations
+                # above the same node appeared consecutively, with oldest first.
+                oldest_mutation_for_node = {}
+                for site in prev_tree.sites():
+                    for mutation in site.mutations:
+                        if not util.is_unknown_time(mutation.time):
+                            oldest_mutation_for_node[mutation.node] = max(
+                                oldest_mutation_for_node[mutation.node], mutation.time
+                            )
+
+            for parent_node in nodes_changed:
+                child_edge_ids = edges_from_node[parent_node]
+                if len(child_edge_ids) >= 3:
+                    # We have a previous polytomy to break
+                    parent_time = existing_node_time[parent_node]
+                    new_nodes = []
+                    child_ids = existing_edges_child[list(child_edge_ids)]
+                    left = None
+                    max_time = 0
+                    # Split existing edges
+                    for edge_id, child_id in zip(child_edge_ids, child_ids):
+                        max_time = max(max_time, existing_node_time[child_id])
+                        if check_mutations and child_id in oldest_mutation_for_node:
+                            max_time = max(max_time, oldest_mutation_for_node[child_id])
+                        if left is None:
+                            left = existing_edges_left[edge_id]
+                        else:
+                            assert left == existing_edges_left[edge_id]
+                        if existing_edges_right[edge_id] > interval[0]:
+                            # make sure we carry on the edge after this polytomy
+                            existing_edges_left[edge_id] = pos
+                    # Arbitrarily, if epsilon is not small enough, use half the min dist
+                    dt = min((parent_time - max_time) / (len(child_ids) * 2), epsilon)
+                    # Break this N-degree polytomy. This requires N-2 extra nodes to be
+                    # introduced: create them here in order of decreasing time
+                    new_nodes = [
+                        node_table.add_row(time=parent_time - (i * dt))
+                        for i in range(1, len(child_ids) - 1)
+                    ]
+                    # print("New nodes:", new_nodes, node_table.time[new_nodes])
+                    for new_edge in resolve_polytomy(parent_node, child_ids, new_nodes):
+                        edge_table.add_row(
+                            left=left, right=pos, child=new_edge[0], parent=new_edge[1],
+                        )
+                        # print("new_edge: left={}, right={}, child={}, parent={}"
+                        #    .format(left, pos, new_edge[0], new_edge[1]))
+                else:
+                    # Previous node was not a polytomy - just add the edges_out
+                    for edge_id in child_edge_ids:
+                        if existing_edges_right[edge_id] == pos:  # is an out edge
+                            edge_table.add_row(
+                                left=existing_edges_left[edge_id],
+                                right=pos,
+                                parent=parent_node,
+                                child=existing_edges_child[edge_id],
+                            )
+
+            for edge in e_out:
+                if edge.parent != tskit.NULL:
+                    # print("REMOVE", edge.id)
+                    edges_from_node[edge.parent].remove(edge.id)
+            for edge in e_in:
+                if edge.parent != tskit.NULL:
+                    # print("ADD", edge.id)
+                    edges_from_node[edge.parent].add(edge.id)
+
+            # Chop if we have created a polytomy: the polytomy itself will be resolved
+            # at a future iteration, when any edges move into or out of the polytomy
+            while nodes_changed:
+                node = nodes_changed.pop()
+                edge_ids = edges_from_node[node]
+                # print("Looking at", node)
+
+                if len(edge_ids) == 0:
+                    del edges_from_node[node]
+                # if this node has changed *to* a polytomy, we need to cut all of the
+                # child edges that were previously present by adding the previous
+                # segment and left-truncating
+                elif len(edge_ids) >= 3:
+                    for edge_id in edge_ids:
+                        if existing_edges_left[edge_id] < interval[0]:
+                            self.edges.add_row(
+                                left=existing_edges_left[edge_id],
+                                right=interval[0],
+                                parent=existing_edges_parent[edge_id],
+                                child=existing_edges_child[edge_id],
+                            )
+                        existing_edges_left[edge_id] = interval[0]
+        assert len(edges_from_node) == 0
+        self.sort()
+
+        if squash_edges:
+            self.edges.squash()
+            self.sort()  # Bug: https://github.com/tskit-dev/tskit/issues/808
+
+        if record_provenance:
+            parameters = {"command": "randomly_split_polytomies"}
+            self.provenances.add_row(
+                record=json.dumps(provenance.get_provenance_dict(parameters))
+            )

python/tskit/trees.py

@@ -6219,6 +6219,62 @@ def count_topologies(self, sample_sets=None):
 
         yield from combinatorics.treeseq_count_topologies(self, sample_sets)
 
+    def randomly_split_polytomies(
+        self,
+        *,
+        epsilon=None,
+        squash_edges=True,
+        record_provenance=True,
+        random_seed=None,
+    ):
+        """
+        Return a tree sequence with extra nodes and edges
+        so that any node with greater than 2 children (i.e. a multifurcation
+        or "polytomy") is resolved into successive bifurcations. For any
+        multifucating node ``u`` with ``n`` children, the :math:`(2n - 3)!!`
+        possible bifurcating topologies are produced with equal probability.
+
+        Polytomies are split per node, not per tree, so that if an identical
+        polytomy spans several trees, it will be randomly resolved into a single
+        set of bifurcating splits. However, if on shifting to a new genomic
+        region, the children of node ``u`` change, either being added or removed,
+        an entirely new random resolution of the node will be applied to that
+        region.
+
+        Because a tree sequence requires that
+        :ref:`parents be older than children<sec_valid_tree_sequence_requirements>`,
+        the newly added nodes are inserted at a time fractionally younger than
+        than the time of node ``u``. This can be controlled by the ``epsilon``
+        parameter.
+
+        :param epsilon: The maximum time between each newly inserted node. For a
+            given polytomy, if possible, the ``n-2`` extra
+            nodes are inserted at time ``epsilon`` apart from each other and from the
+            time of node ``u``. By default, this is set to a very small time value
+            (:math:`1e-10`). However, if there is a child node or a mutation above
+            a child node whose time is very close to the time of ``u``, ``epsilon``
+            may not be small enough. In this case, a smaller time interval is used
+            so that the tables will still encode a valid tree sequence.
+        :param bool squash_edges: If True (default), run :meth:`.squash()` at the
+            end of the process. This can help to reduce the total number of extra
+            edges produced.
+        :param bool record_provenance: If True, add details of this operation to the
+            provenance information of the returned tree sequence. (Default: True).
+        :param int random_seed: The random seed. If this is None, a random seed will
+            be automatically generated. Valid random seeds must be between 1 and
+            :math:`2^32 − 1`.
+        :return: A new tree sequence with polytomies split into random bifurcations.
+        :rtype: .TreeSequence
+        """
+        tables = self.dump_tables()
+        tables.randomly_split_polytomies(
+            epsilon=epsilon,
+            squash_edges=squash_edges,
+            record_provenance=record_provenance,
+            random_seed=random_seed,
+        )
+        return tables.tree_sequence()
+
     ############################################
     #
     # Deprecated APIs. These are either already unsupported, or will be unsupported in a