Slicing a tree (not a tree sequence) at a specific time

[hyanwong]

I wanted to be able to take a time (e.g. of a mutation) and slice one tree in a tree sequence at that time, to find the nodes below each of the sliced lineages. E.g. given a tree like this, with a focal mutation 3 in red, sliced at the dotted line

I want to identify the nodes below the lineages that exist in the tree at that time, i.e. nodes 20, 2, 21, 24, 26, 29, 31, 33, 17, 34.

We can do that using a "timedesc" traversal, like this. I hope that's as efficient as it could be, since we only descend as far as we need to go.

time = ts.mutation(focal_mut).time
lineages = {child: tskit.NULL for child in tree.roots} # map child->parent, but we don't need parent here
for node in tree.nodes(order="timedesc"):
    if nodes_time[node] <= time:
        break
    else:
        del lineages[node]
        lineages.update({child: node for child in tree.children(node)})

For my purposes, I then wanted to take each of these identified nodes and look how the number of lineages increases through time. We can do that by taking each of these nodes and carrying on the descent from there:

samples = set(ts.samples())
for child, _ in lineages.items():
    steps = [time]
    count = [1]
    for node in tree.nodes(child, order="timedesc"):
        nc = tree.num_children(node)
        # For this stat, even if a lineage terminates in a leaf, we carry on counting it.
        if nc > 0:
            if nodes_time[node] == steps[-1]:  # child nodes at same time
                count[-1] += nc - 1
            else:
                steps.append(nodes_time[node])
                count.append(count[-1] + nc - 1)
    # Save steps & count here, or plot them out