Tskit: Method to 'prune back' a tree sequence

[jeromekelleher]

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.

[petrelharp]

This is great. prune_leaves is OK, but it's not specific to leaves. Other ideas:

• truncate
• pollard
• time_slice

[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

[gtsambos]

Maybe I'm not fully understanding here, but isn't this basically what map_ancestors/find_ancestors does? You'd basically just need to add a sample on each lineage at the right time (the census event you've talked about with me in person before - btw I'm happy to start working on this soon if it will be useful), and then apply map_ancestors with the ancestors being all nodes at this census time or earlier, and the samples being all the leaves.

[jeromekelleher]

You're right, it's definitely closely related @gtsambos. It would be excellent if we can implement it using map/find_ancestors!

ps. Let's talk about the census event soon --- I do actually need this as it happens.

[gtsambos]

Hi @jeromekelleher, I just did a mock-up of this idea in this gist

then apply map_ancestors with the ancestors being all nodes at this census time or earlier, and the samples being all the leaves.

It's actually even simpler than this - you just need to set the census nodes as your samples, and all older nodes as ancestors. Then each leaf node in the output corresponds to a lineage present at the specified census time.

Btw, here's a Python function that puts nodes on trees at a given 'census' time:
(Ni = input node table, Ei = input edge table, L = sequence length)

def add_census(Ni, Ei, census_time, L):
    No = tskit.NodeTable()
    No.append_columns(time=Ni.time, flags=Ni.flags)
    Eo = tskit.EdgeTable()

    edge_list = []
    for ind, e in enumerate(Ei):
        edge_list.append(e)
        if ind + 1 == len(Ei) or Ei[ind + 1].parent != Ei[ind].parent:
            for e in edge_list:
                if Ni.time[e.parent] > census_time and Ni.time[e.child] < census_time:
                    v = No.add_row(time=census_time)
                    Eo.add_row(e.left, e.right, v, e.child)
                    Eo.add_row(e.left, e.right, e.parent, v)
                else:
                    Eo.add_row(e.left, e.right, e.parent, e.child)
            edge_list = []

    # Sort output. 
    new_tables = tskit.TableCollection(sequence_length=L)
    new_tables.nodes.append_columns(time=No.time,flags=No.flags)
    new_tables.edges.append_columns(left=Eo.left, right=Eo.right, 
                                    parent=Eo.parent, child=Eo.child)
    new_tables.sort()
    return(new_tables.tree_sequence())

ps. Let's talk about the census event soon --- I do actually need this as it happens.

Yes! 💻

[gtsambos]

The major difference between my implementation and yours is that you treat any (parent, child) combo as the same lineage,

            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]

even where the edges with this combo span non-adjacent intervals of the sequences. Whereas my add_census function above treats each edge as a distinct lineage.

I guess there might be contextual reasons to prefer one or the other, but my way makes a bit more sense to me because it ensures that each distinct 'leaf' corresponds to no more than one distinct historical chromosome. Ie. a particular ancestor may have contributed two distinct segments to a particular descendant's chromosome, but that doesn't necessarily mean that both segments have passed through the same sequence of individuals (which is my usual interpretation of lineage)

[molpopgen]

This would be cool to have implemented. There are some subtleties that arise for handling mutations that don't occur in the case of decapitation. For any parent/child edge that you truncate to a time "in between", some of the mutations on the child nodes may need ramapping to to the new node. I'm thinking about forward simulations here, where the mutation's origin time may be a part of its metadata, and so you may know that it predates your census time but post-dates the first parent ancestral to that time. For many coalescent sims, (I think) you can justify simply binomially sampling any mutations and remap them.

[petrelharp]

It's actually even simpler than this - you just need to set the census nodes as your samples, and all older nodes as ancestors.

Shouldn't you just set all the census-and-older nodes to be your samples?

handling mutations

Oh, this is a good point.

[jeromekelleher]

So, perhaps the thing to do is to add in the 'census' nodes separately then, as I can imagine this being useful without the pollarding step. Handling the mutations is tricky, and I think we'll have to make the handling pluggable somehow (a callback that decides if a mutation goes above or below the census node, or something), as there's no general way for tskit to know what time the mutation happens at.

[molpopgen]

Exactly. For each parent/child combo spanning the census time, you can pass the parent time, child time, census time, and the index of the mutation on the child node to a function returning True if it should be "pushed up" to the new census node. You could even just pass in all indexes for all mutations on the child node, which may be more efficient. tskit could provide default using numpy to make a random decision based on the proportion of time above the census node.

[petrelharp]

So, perhaps the thing to do is to add in the 'census' nodes separately then, as I can imagine this being useful without the pollarding step.

Right; then we could just simplify.

[jeromekelleher]

Related to #382