Creating polytomies by removing edges

[hyanwong]

To test semi-realistic tree sequences with polytomies, it helps to be able to create polytomies in a tree sequence by removing edges (and grafting their descendant edges onto the parent of the removed edge). It turns out that the algorithm to do this over an entire tree sequence has some quite tricky logic, but after chatting to @a-ignatieva, I think the code below should do it. At some point we should test this more thoroughly and maybe even incorporate it into tskit? It turns out to be rather a core function for testing how to deal with polytomies

import itertools
import collections

def remove_edges(ts, edge_id_remove_list):
    edges_to_remove_by_child = collections.defaultdict(list)
    edge_id_remove_list = set(edge_id_remove_list)
    for m in ts.mutations():
        if m.edge in edge_id_remove_list:
            # If we remove this edge, we will remove the associated mutation
            # as the child node won't have ancestral material in this region.
            # So we force the user to explicitly (re)move the mutations beforehand
            raise ValueError("Cannot remove edges that have associated mutations")
    for remove_edge in edge_id_remove_list:
        e = ts.edge(remove_edge)
        edges_to_remove_by_child[e.child].append(e)

    # sort left-to-right for each child
    for k, v in edges_to_remove_by_child.items():
        edges_to_remove_by_child[k] = sorted(v, key=lambda e: e.left)
        # check no overlaps
        for e1, e2 in zip(edges_to_remove_by_child[k], edges_to_remove_by_child[k][1:]):
            assert e1.right <= e2.left

    # Sanity check: this means the topmost node will deal with modified edges left at the end
    assert ts.edge(-1).parent not in edges_to_remove_by_child
    
    new_edges = collections.defaultdict(list)
    tables = ts.dump_tables()
    tables.edges.clear()
    samples = set(ts.samples())
    # Edges are sorted by parent time, youngest first, so we can iterate over
    # nodes-as-parents visiting children before parents by using itertools.groupby
    for parent_id, ts_edges in itertools.groupby(ts.edges(), lambda e: e.parent):
        # Iterate through the ts edges *plus* the polytomy edges we created in previous steps.
        # This allows us to re-edit polytomy edges when the edges_to_remove are stacked
        edges = list(ts_edges)
        if parent_id in new_edges:
             edges += new_edges.pop(parent_id)
        if parent_id in edges_to_remove_by_child:
            for e in edges:
                assert parent_id == e.parent
                l = -1
                if e.id in edge_id_remove_list:
                    continue
                # NB: we go left to right along the target edges, reducing edge e as required
                for target_edge in edges_to_remove_by_child[parent_id]:
                    # As we go along the target_edges, gradually split e into chunks.
                    # If edge e is in the target_edge region, change the edge parent
                    assert target_edge.left > l
                    l = target_edge.left
                    if e.left >= target_edge.right:
                        # This target edge is entirely to the LHS of edge e, with no overlap
                        continue
                    elif e.right <= target_edge.left:
                        # This target edge is entirely to the RHS of edge e with no overlap.
                        # Since target edges are sorted by left coord, all other target edges
                        # are to RHS too, and we are finished dealing with edge e
                        tables.edges.append(e)
                        e = None
                        break
                    else:
                        # Edge e must overlap with current target edge somehow
                        if e.left < target_edge.left:
                            # Edge had region to LHS of target
                            # Add the left hand section (change the edge right coord)
                            tables.edges.add_row(left=e.left, right=target_edge.left, parent=e.parent, child=e.child)
                            e = e.replace(left=target_edge.left)
                        if e.right > target_edge.right:
                            # Edge continues after RHS of target
                            assert e.left < target_edge.right
                            new_edges[target_edge.parent].append(
                                e.replace(right=target_edge.right, parent=target_edge.parent)
                            )
                            e = e.replace(left=target_edge.right)
                        else:
                            # No more of edge e to RHS
                            assert e.left < e.right
                            new_edges[target_edge.parent].append(e.replace(parent=target_edge.parent))
                            e = None
                            break
                if e is not None:
                    # Need to add any remaining regions of edge back in 
                    tables.edges.append(e)
        else:
            # NB: sanity check at top means that the oldest node will have no edges above,
            # so the last iteration should hit this branch
            for e in edges:
                if e.id not in edge_id_remove_list:
                    tables.edges.append(e)
    assert len(new_edges) == 0
    tables.sort()
    return tables.tree_sequence()

Test it by removing all edges that have no associated mutation. Note that there are two ways of doing this: we could either remove every interval that has no mutations, or we could remove edges with no mutations (i.e. if a mutation exists on any interval with parent=p and child=c, we keep all intervals with parent=p and child=c). The edge removal approach therefore removes fewer edges:

import numpy as np
import msprime

def unsupported_edges(ts, per_interval=False):
    """
    Return the internal edges that are unsupported by a mutation.
    If ``per_interval`` is True, each interval needs to be supported,
    otherwise, a mutation on an edge (even if there are multiple intervals
    per edge) will result in all intervals on that edge being treated
    as supported.
    """
    edges_to_remove = np.ones(ts.num_edges, dtype="bool")
    edges_to_remove[[m.edge for m in ts.mutations()]] = False
    # We don't remove edges above samples
    edges_to_remove[np.isin(ts.edges_child, ts.samples())] = False

    if per_interval:
        return np.where(edges_to_remove)[0]
    else:
        keep = (edges_to_remove == False)
        for p, c in zip(ts.edges_parent[keep], ts.edges_child[keep]):
            edges_to_remove[np.logical_and(ts.edges_parent == p, ts.edges_child == c)] = False
        return np.where(edges_to_remove)[0]


for rep, ts in enumerate(
    msprime.sim_ancestry(50, sequence_length=1e6, recombination_rate=1e-4, num_replicates=20, random_seed=321)
):
    ts = msprime.sim_mutations(ts, rate=1e-3, random_seed=1)
    del_by_interval = unsupported_edges(ts, per_interval=True)
    del_by_edge = unsupported_edges(ts, per_interval=False)
    
    poly_ts_by_interval = remove_edges(ts, del_by_interval)
    poly_ts_by_edge = remove_edges(ts, del_by_edge)

    assert ts.num_trees == poly_ts_by_interval.num_trees == poly_ts_by_edge.num_trees
    for t1, t2, t3 in zip(ts.trees(), poly_ts_by_interval.trees(), poly_ts_by_edge.trees()):
        assert set(t1.roots) == set(t2.roots) == set(t3.roots)
        for node in t2.nodes():
            # Every node in comparable trees should have the same set of descendant samples
            assert set(t1.samples(node)) == set(t2.samples(node))
            # Every node in tree2 should have a subset of descendant nodes in tree 1
            assert set(t1.nodes(node)) >= set(t2.nodes(node))
        for node in t3.nodes():
            # Every node in comparable trees should have the same set of descendant samples
            assert set(t1.samples(node)) == set(t3.samples(node))
            # Every node in tree2 should have a subset of descendant nodes in tree 1
            assert set(t1.nodes(node)) >= set(t3.nodes(node))
    
    print(f"Replicate {rep}: {ts.num_trees} trees")
    print(f" {len(del_by_interval) / ts.num_edges * 100}% of intervals to remove")
    print(f" {len(del_by_edge) / ts.num_edges * 100}% of edges to remove")

[nspope]

Excellent, this'll be really useful for testing dating.

[hyanwong]

Yep. See tskit-dev/tsdate#359 (comment)

[nspope]

Just a note that this implementation keeps completely disconnected nodes in the "polytomified" tree sequence; this is nice because it keeps node IDs consistent, but might require special handling in some cases. The alternative would be to purge these and return a node mapping along with the modified tree sequence.

[hyanwong]

Yep. You can then "simplify keeping unary" to get rid of them. If you do that with the tables, you should get a node mapping back.