Method to 'prune back' a tree sequence

I want a function that will remove the topology after (forwards in time) a specific time, and mark each extant lineage with as a sample. The intuition is that we’re drawing a line accross the tees, and sort of “pruning back” to this line. (This will allow us to generate the haplotypes extant at this time.)

Here’s a simple implementation:

def prune_leaves(ts, time):
    """
    Returns a tree sequence in which all the topology and mutation
    information after the specified time is remove. That is, the
    samples in the returned tree sequence will be the lineages
    extant at the specified time.
    """
    tables = ts.dump_tables()
    t = tables.nodes.time
        
    tables.edges.clear()
    samples = []
    edge_map = {}
    for edge in ts.edges():
        if t[edge.child] > time:
            tables.edges.add_row(edge.left, edge.right, edge.parent, edge.child)
        elif t[edge.child] <= time and t[edge.parent] > time:
            key = (edge.child, edge.parent)
            if key in edge_map:
                # If the same parent/child combination exists, then we should map
                # to the same ancestor node.
                u = edge_map[key]
            else:
                u = tables.nodes.add_row(
                    flags=tskit.NODE_IS_SAMPLE, time=time)
                samples.append(u)
                edge_map[key] = u
            tables.edges.add_row(edge.left, edge.right, edge.parent, u)
    tables.sort()
    tables.simplify(samples)
    return tables.tree_sequence()

There’s probably some subtleties and corner case I’ve missed out here, but this probably the basic idea.

I think this would be a useful addition to tskit. I don’t know what we’d call it though, and I think it’s worth keeping this in mind with the discussion going on over in #261 where we’re chopping up the tree sequence in the space rather than time direction.

[molpopgen]

If removing all nodes ancestral to a certain time is decapitating, then this would be to amputate or dismember the tree:

import tskit.gory_ops

[jeromekelleher]

Yep, agreed. That’s basically what I’ve done in #2240