Tskit: sample frequency spectrum

[jeromekelleher]

Add support for site frequency spectrum. Needs a specific C algorithm for efficiency.

I think we should call it the "sample_frequency_spectrum", since it's counting samples. site_frequency_spectrum is actively misleading I think, as well as leading to absurd function names like site_site_frequency_spectrum and branch_site_frequency_spectrum. I'm open to other names though.

See also #196.

[petrelharp]

Another name is "allele frequency spectrum", which is more what we are computing (the distribution of frequencies of alleles).

From email:

Re the SFS, it wasn't obvious to me why you left out the total branch length over 0 samples - this is still a useful thing, right?

Isn't this always zero?

[jeromekelleher]

Isn't this always zero?

Not if your sample set size is < n - isn't it the branch length over "the complement of this sample set" then?

[jeromekelleher]

Another name is "allele frequency spectrum", which is more what we are computing (the distribution of frequencies of alleles).

I like it...

[petrelharp]

Not if your sample set is not all the samples - isn't it the branch length over "the complement of this sample set" then?

Oh, sorry - I misread that. Right, the 0th entry of the AFS. Yes, that could be useful (although since it depends on what else is in the tree sequence, caution is needed).

[jeromekelleher]

I've worked through the algorithm for the AFS @petrelharp, and the zeroth element works like you would expect almost all the time. The only time we get weird results where the obvious definition algorithm differs from the efficient edge-based algorithm is when we have edges lying around that are not ancestry to any sample (e.g., unsimplified forward sim). In this case, to output the same thing as the tree based definition algorithm we'd need to keep track of the count of all samples that are under each node as well as the sample sets, and only update the AFS when this is > 0. The definition of AFS[0] is then, "the span weighted total branch length along branches ancestral to some sample that is not in the current sample set."

This might not be worth the bother though, and we may be better off just setting AFS[0] = 0 always and documenting this. I do think we should keep things simple and make AFS[k] refer to the total branch length over k samples rather than trying to save space by making it k - 1 (or k+1, whichever). I think this would be a source of persistent off by one errors that would make people hate us forever.

[petrelharp]

The only time we get weird results where the obvious definition algorithm differs from the efficient edge-based algorithm is when we have edges lying around that are not ancestry to any sample

Ah, right, because for the 0th element (and also the nth element) we're effectively using a summary function f( ) that doesn't have f(0) = 0 (or, for the nth slot, f(n) = 0).

Hm: if there is a 0th element of the AFS, for site statistics it should be the proportion of sites that aren't segregating in the sample set. (so then sum(AFS) = 1.0 or sequence_length depending on normalise). But we don't even have a notion of this since we're working with continuous sequence. So, I don't think we should return a value for the 0th element at all, really. I guess we could put a 0 (or a nan) there, but that seems more confusing. It's not too bad, I think, that AFS[k] corresponds to the alleles with k+1 samples below, because it is k+1st element of AFS; and also people are used to only getting values for the AFS for 1...n-1?

[jeromekelleher]

@molpopgen, what's your take on the zeroth entry in the AFS? Should we make AFS[0] equal to the count for singletons, or fill it with some value (either something meaningful or NaN/0 or whatever), or something else?

[molpopgen]

I think we need to be a but more precise perhaps? If site is an entry in the site table, then one can define a zero class as the "proportion of sites that are mutated but invariant in the sample". However, that is not a typical definition, which would be the fraction/number of monomorphic sites in a sequence, which is undefined for models like infinite-sites.

[petrelharp]

If site is an entry in the site table, then one can define a zero class as the "proportion of sites that are mutated but invariant in the sample". However, that is not a typical definition, which would be the fraction/number of monomorphic sites in a sequence, which is undefined for models like infinite-sites.

That was my view. So, if there's anything there it should be NaN; the question is whether to start out the AFS vector with the singletons (so, singletons in the 0th slot) or put the singletons in the next slot (in the 1th slot, which is attractive since singletons are seen one=1 time).

[molpopgen]

I prefer that element 0 be the singletons, as that is the more idiomatic implementation, b/c Python is zero-offset by default.

[jeromekelleher]

OK, happy to be out voted here; makes my life easier! Thanks for the input.

[jeromekelleher]

OK, I'm really sorry to keep digging this up, but are we sure that throwing away the zeroth element of the AFS is the right thing to do? It makes the code for computing it sooo much uglier! Suppose we have a single tree covering the sequence (of length 1), then to compute the K-d branch-wise AFS is:

out_dim = [len(sample_set) for sample_set in sample_sets]
S = np.zeros(out_dim)
trees = [
      next(ts.trees(tracked_samples=sample_set, sample_counts=True))
      for sample_set in sample_sets]
for node in trees[0].nodes():
      # Note this must be a tuple for indexing to work as we want here.
      x = tuple([tree.num_tracked_samples(node) for tree in trees])
      S[x] += trees[0].branch_length(node)

if we want to throw away the zeroth element, then we'd do something like

out_dim = [len(sample_set) - 1 for sample_set in sample_sets]
S = np.zeros(out_dim)
trees = [
      next(ts.trees(tracked_samples=sample_set, sample_counts=True))
      for sample_set in sample_sets]
for node in trees[0].nodes():
      x = [tree.num_tracked_samples(node) - 1 for tree in trees]
      if np.all(x >= 0): 
          S[tuple(x)] += trees[0].branch_length(node)

The first version looks much more elegant to me, and it feels wrong to throw away all information when any of the sample sets have zero elements under a node. This is going to be a common case if the sample sets are disjoint, isn't it?

I prefer that element 0 be the singletons, as that is the more idiomatic implementation, b/c Python is zero-offset by default.

Surely it's more idiomatic to have the total branch length over k samples be S[k] rather than S[k - 1], for exactly this reason?

Hm: if there is a 0th element of the AFS, for site statistics it should be the proportion of sites that aren't segregating in the sample set. (so then sum(AFS) = 1.0 or sequence_length depending on normalise). But we don't even have a notion of this since we're working with continuous sequence.

I don't understand why this is a problem - what about sites in the table that have no mutations? Or, if we think of a small sample set in a big tree, sites that are fixed within this sample set?

As I say, I'm sorry to keep banging on at this one - I promise this is the last time I'll ask! This is the sort of thing we just can't change after it's shipped, so I just want to be triply sure that we're making the right decision. I agree that people are used to the AFS being defined for only k = 1, ..., n - 1 (and there will be people tripped up by the zeroth element) but people are also not used to the branch length definition in K dimensions, where defining the zeroth element may turn out to be useful. If you stop people doing stupid things you also stop them doing clever things...

[jeromekelleher]

@GertjanBisschop and @KLohse, any thoughts on what the right thing to do with the zeroth element of the AFS is (see last few comments above). I notice that we defined the zeroth element for the get_joint_site_frequency_spectra function in the book chapter Konrad. Feel free to email me for some context if you have no idea what's going on here!

[petrelharp]

it feels wrong to throw away all information when any of the sample sets have zero elements under a node. This is going to be a common case if the sample sets are disjoint, isn't it?

Whoops - you are right! - we should not throw out the cases where a site is monomorphic in one sample set but polymorphic in the other. Ok, I retract my opinion: we should keep the 0th elements, but set the entry in the AFS that corresponds to "monomorphic in all sample sets" to nan.

[KLohse]

Sorry for the slow response. Your suspicion about the ugly code was of course correct and I completely second @petrelharp: monomorphic sites in a pop should be included as the 0th element of the jSFS which is what we did in the book chapter. Only when they are invariant in all pops can they be set to nan. It's useful to think about the limit: in the case of a single sample from each of i pops the jSFS will have 2^i entries: presence/absence in each pop.

[jeromekelleher]

Thanks Konrad, this is very helpful!

[jeromekelleher]

Thinking about this more, I think we should remove the restriction on forcing afs[0, ..., 0] = 0. My argument works on a single sample set, so I'll keep it to the 1D SFS. Consider the following example:

nodes_ex = io.StringIO("""\
id      is_sample   population      individual      time    metadata
0       1       0       -1      0.00000000000000
1       1       0       -1      0.00000000000000
2       1       0       -1      0.00000000000000
3       0       0       -1      1.0
4       0       0       -1      2.0
""")
edges_ex = io.StringIO("""\
id      left            right           parent  child
0       0.00000000      1.00000000      3       1
1       0.00000000      1.00000000      3       2
2       0.00000000      1.00000000      4       0
3       0.00000000      1.00000000      4       3
""")

ts = tskit.load_text(nodes=nodes_ex, edges=edges_ex, strict=False)

tree = ts.first()
print(tree.draw_text())

afs = ts.allele_frequency_spectrum(mode="branch", polarised=True)
print(afs, np.sum(afs), tree.total_branch_length, sep="\t")

afs = ts.allele_frequency_spectrum([1], mode="branch", polarised=True)
print(afs, np.sum(afs), tree.total_branch_length, sep="\t")

Which gives the output:

┏━┻┓
┃  3
┃ ┏┻┓
0 1 2

[[0. 4. 1. 0.]] 5.0     5.0
[[0. 2.]]       2.0     5.0

So, we have a tree with three leaves and a total branch length of 5. In the first case, we take the branch AFS over all samples, giving [0, 4, 1]: the total branch length subtending 1 sample is 4 and 2 samples is 1. Logically enough, the sum of this is equal to the total branch length of the tree.

In the second example, we compute the AFS of a subset of these samples, namely the set [1], and we compute the total branch length over 1 element of this set is 2. This is the sum of the branch length of node 1 and node 3. Notice, however, that now the sum of the AFS is now not equal to the total branch length of the tree.

The argument for not including the total branch length over zero elements of the sample set (as I understand it) is that we should not have to worry about what is not in the sample sets when looking at the AFS. But, this is inconsistent, because we are always taking the tree as a whole when we compute these values. If we were to simplify down to the topology just subtended by the sample set, then we would have a branch length of zero. The AFS is always computed in the context of the whole tree, and therefore excluding the branch length over samples that are not in the current sample set is both artificial and inconsistent.

I think that the zeroth element should be included, so that in this example we return [[3. 2.]], and we maintain the property that the sum(afs) == tree.total_branch_length.

[petrelharp]

I think that you are pointing out, basically, that the nth element depends on the whole tree in the same way that the 0th does. Good point. So, we could zero both of them out, or include them (carefully documenting what they are). I do like not throwing away information, so I'm OK with including them.

[jeromekelleher]

We already include the n'th element though: we have 1 element in the second example above, and afs[1] = 2.0.

In the first example, where the sample set is [0, 1, 2], afs[3] = 0. If there was a unary branch over the root though, presumably this would contain this branch length.

I haven't thought through the folded case as much, but it should be similar.

It all works very nicely I think, we just have to not make special cases. I'll push through these changes, document it as well as I can, and we can see if we all agree with the way it's done then.

[petrelharp]

We already include the n'th element though

Yes, I know, sorry - I meant to continue to include it.

[KLohse]

I agree that both the 0th and the nth need to be included for consistency; both categories become relevant for the general case of the AFS for multiple pops. Thinking about implementing the folding of the AFS in the most general way; it seems that one would have to specify a ref population with respect to which ASF the folding happens. In other words, with n pops, there are n ways of folding (and the various folded spectra are not equivalent).

[molpopgen]

Thinking about implementing the folding of the AFS in the most general way; it seems that one would have to specify a ref population with respect to which ASF the folding happens. In other words, with n pops, there are n ways of folding (and the various folded spectra are not equivalent).

I think the easiest thing here is to keep the definition of the folded AFS consistent. Since the folded AFS is the AFS of the minor allele for a single population, you just need to define the minor allele globally (e.g. with respect to all n sample lists) for n populations.

[petrelharp]

with n pops, there are n ways of folding (and the various folded spectra are not equivalent)

Uh-oh, you are right. I didn't think about this hard enough, and suggested above a folding method that collapses [i1,i2] + [n1-i1, i2] + [i1, n2-i2] + [n1-i1, n2-i2]; but this is not correct, because (for instance) [i,i] and [n1-i, i] are clearly different, since the two pops agree in the first case but disagree in the second.

But, about nonequivalence: if we don't know the ancestral allele, all that can be said is that we can't distinguish [i1, i2, ..., ik] from [n1-i1, n2-i2, ..., nk-ik]; so isn't the right thing to do to return a "lower triangular" array, e.g. with the upper triangle zeroed? I.e., the output has out[i1, ..., ik] = 0 if i1 + ... + ik > (n1 + ... + nk)/2? (need to check that)

Thanks for catching this, @KLohse!

[jeromekelleher]

On another topic, I merged the current version of the algorithms last week with some first draft method docs here and a tutorial here. It would be great if you guys could have a read through and open up some issues/PRs for problems/omissions that you spot.

[petrelharp]

Looks great. I've commented on that commit; in retrospect I should have made a PR. LMK if you want me to do that.

[jeromekelleher]

Looks great. I've commented on that commit; in retrospect I should have made a PR. LMK if you want me to do that.

That would be good, if you don't mind. The changes all look good.

[petrelharp]

Made changes over at #254.

[jeromekelleher]

AFS has been fully implemented.