Count extant lineages at a given time

[jeromekelleher]

An operation that came up as being a useful feature in the recent PopSim meeting is to be able to count the number of extant lineages in each tree at a given time. This is equivalent to (a) deleting all edges older than t and truncating all edges intersecting with it (this would require the planned time argument to delete_interval in #382); and (b) counting the number of roots in each tree. We could define the operation by just doing this, but it seems a bit inefficient, especially if we do this for lots of different times.

Here's a way we can do it more efficiently:

def count_roots(ts, t):
    time = ts.tables.nodes.time    
    num_children = np.zeros(ts.num_nodes, dtype=int)
    num_roots = ts.num_samples
    ret = np.zeros(ts.num_trees, dtype=int)

    j = 0
    for _, edges_out, edges_in in ts.edge_diffs():
        for edge in edges_out:
            if time[edge.parent] <= t:
                num_children[edge.parent] -= 1
                if num_children[edge.parent] == 1:
                    num_roots += 1
        for edge in edges_in:
            if time[edge.parent] <= t:                
                num_children[edge.parent] += 1
                if num_children[edge.parent] == 2:
                    num_roots -= 1
        ret[j] = num_roots
        j += 1
    return ret

We can code this up pretty efficiently in C, and return the numpy array. Should we do it this way? Should we call the method TreeSequence.num_roots(time=None), which returns a numpy array like the above? (If time is None, it defaults to the oldest time in the ts)

[petrelharp]

We definately want to do this, somehow! Note that another way to do it would be to remove everything below t and count the number of tips.

And gee, your algorithm is very nice. num_roots is a good name; but it might be nice to connote somehow the fact that these aren't actually roots in the trees unless time=None, maybe num_lineages_at( ) or num_roots_at( )?

p.s. Asking about number of lineages in the trees irregardless of whether they are above samples is even easier. Here is, for instance, a function that I think tells you how many lineages there are in the tree at position pos at time t:

def num_lineages(ts, pos, t):
   left = ts.tables.edges.left
   right = ts.tables.edges.right
   a = ts.tables.nodes.time[ts.tables.edges.parent]
   b = ts.tables.nodes.times[ts.tables.edges.child]
   return np.sum((left <= pos) * (right > pos) * (a < t) * (b >= t))

[jeromekelleher]

Very nice @petrelharp, that's another great way to look at it. It does assume that the trees have one root though (what happens in the degenerate case of no edges?). The algorithm I have above isn't complete either though, I think I'm assuming that all leaves are samples at the moment (but, I'm sure that can be fixed).

I agree num_roots will confuse people, so num_lineages is a better choice. I'm not crazy about the _at suffix though --- isn't that a bit like putting a _where at the end of a bunch of other methods that query stuff?

[jeromekelleher]

Resurrecting this one - @petrelharp, does this do what we need for the Relate selection stat? It's probably something we should just stick in, since we know how to do it. TreeSequence.num_lineages(time), where time is a mandatory positional argument seems like a good option?

def num_lineages(time):
    """
     Returns the number of lineages ancestral to samples in each tree along the genome
     at the specified time.   
    """

It's not the only way that we'll want to use this type of information, but it's a useful operation to add to the suite I think?

[jeromekelleher]

I suppose we should add a Tree.num_lineages(time) operation as well, which computes the same thing. This would probably need to do a top-down traversal of some sort I guess (but we could imagine adding some options to make it faster).

[jeromekelleher]

Ah, what we really want to do here is to return the number of lineages at a vector of time points. So, we'd have something like

def num_lineages(self, time, windows=None):
    """
     Returns the number of lineages ancestral to samples in each tree along the genome
     at the specified time points. The windows parameter is interpreted in a same way as 
     for the stats API. The returned value is a (num_windows, num_times) array, where 
     dimensions are dropped in the same ways as the stats API, when defaults are used.
    """

(This wouldn't be the real docstring, just a sketch)

What do you think @petrelharp?

I think we want to sample the number of lineages at specific times rather than to average them out in windows.

cc @percyfal

[petrelharp]

I like it! Very nice, clean.

Do you remember what the original motivation was?

One application would be to make a genome x time plot of number of lineages (This would be like the genome x time plot of branch-segregating-sites if we had time windows in for statistics.)

To compute the Relate selection statistic or the number of coalescences we need "number of lineages" information, but in the particular tree that the mutation appears in, not averaged across a window; so although this is related it's not exactly what we need for those, I think?

[jeromekelleher]

Here's some doodling for tree versions:

def num_tree_lineages_simple(tree, t):
    lineages = 0
    for u in tree.nodes():
        if tree.time(u) <= t < tree.time(tree.parent(u)):
            lineages += 1
    return lineages


def num_tree_lineages_preorder(tree, t):
    # Only descend the tree as far as we need to.
    lineages = 0
    stack = [tree.virtual_root]
    while len(stack) > 0:
        parent = stack.pop()
        for child in tree.children(parent):
            child_time = tree.time(child)
            if child_time > t:
                stack.append(child)
            elif child_time <= t:
                lineages += 1
    return lineages


def num_tree_lineages_vector_simple(tree, t):
    """
    Return a vector N counting the number of lineages present at each
    time in the specified vector t.
    """
    t = np.array(t)
    k = t.shape[0]
    N = np.zeros(k, dtype=int)
    for u in tree.nodes():
        for j in range(k):
            if tree.time(u) <= t[j] < tree.time(tree.parent(u)):
                N[j] += 1
    return N

Notes:

• This code isn't very robust as it's not taking into account what happens when we specify a time that's greater than the root(s) time. We'll currently get an error.
• We should be able to do a version with the vector of input times where we use the traversal stack to index into the time array. I think this is probably the most efficient way that it can be done, for a single tree. We'd required that the input vector of times is sorted in nondecreasing order, and then push (tree.virtual_root, len(t) - 1) on as the first element of the stack, like so

def num_tree_lineages_vector_stack(tree, t):
    """
    Return a vector N counting the number of lineages present at each
    time in the specified vector t.
    """
    t = np.array(t)
    k = t.shape[0]
    N = np.zeros(k, dtype=int)
    stack = [(tree.virtual_root, k - 1)]
    while len(stack) > 0:
        parent, j = stack.pop()
        for child in tree.children(parent):
            child_time = tree.time(child)
            if child_time > t[j]:
                stack.append((child, j)) # This is probably right?
            else:
                 # do the things
    return N

It's probably best to get a fully working version of the tree method before we start thinking about how to make this work across many trees.

[jeromekelleher]

@a-ignatieva would you mind taking a quick look at these functions in the comment just above to see if they do what you want? I'm trying to figure out what we want here in terms of implementing this operation as efficiently as we can.

[petrelharp]

Here's a note spurred on by #1314. I think there are two general things we want to do:

• Make a plot of "number of extant lineages through time". We could do that by taking the average number of lineages across a time window and across a window of the genome, which would be what we get by adding time windows to segregating_sites.

• Compute things that require knowing the number of extant lineages at a bunch of specific places on the genome and times (e.g., at all the mutations, for the Relate statistic). So, this requires knowing how to keep track of that sort of information, but isn't necessarily part of the python API, as we'd do those things in C. But, now that I think about it, maybe adding the ability to the .trees( ) operator to be able to return that sort of information would be the way to go? Or, I suppose that your code above, @jeromekelleher, is proposing a method of the Tree object that does this?

Things that I can think of that fall in the second category are:

• the Relate statistic
• estimates of coalescence rates
• whatever SIA needs (I forget?)

[jeromekelleher]

Thanks @petrelharp, that does clarify things. I think it's worth having the method to return the number of lineages at a set of times as part of the Python API anyway, even if it's just something we use as a building block later for testing things like the Relate statistic.

We can build other methods up from there once this basic building block is in place I think (even if they don't use it, under the hood).