allele frequency function

[petrelharp]

We should implement the TreeSequence.allele_frequencies(sample_sets) function, which returns a numpy array of (non-ancestral allele frequencies) x (sample_sets).

Here's an implementation:

def allele_frequencies(ts, sample_sets=None):
    if sample_sets is None:
       sample_sets = [ts.samples()] 
    n = np.array([len(x) for x in sample_sets])
    def f(x):
       return x / n
    return ts.sample_count_stat(sample_sets, f, len(sample_sets), windows='sites', polarised=True, mode='site', strict=False, span_normalise=False)

Edit: originally this omitted span_normalise=False.

[jeromekelleher]

Nice, good idea.

[molpopgen]

It may be nice to optionally return the joint fs as a sparse matrix.

[petrelharp]

It may be nice to optionally return the joint fs as a sparse matrix.

I'm not sure what you mean here?

[molpopgen]

I assume your return value is the marginal fs in each sample set, but sometimes you want the joint fs to instead, such that fs[2,3] is the number of mutations present twice in sample set 0 and three times in set 1. Or maybe I'm not understanding what you are returning here.

[petrelharp]

That sounds like the AFS? Which we already have? Here I'm wanting to return the list of allele frequencies for each SNP.

[molpopgen]

Gotcha--my mistake.

[jeromekelleher]

That sounds like the AFS? Which we already have? Here I'm wanting to return the list of allele frequencies for each SNP.

Triggered my AFS implementation PTSD there - would not want to do that again!

[molpopgen]

I shouldn't comment when I haven't been sleeping. 😵

…

On Tue, Mar 31, 2020, 12:50 AM Jerome Kelleher ***@***.***> wrote: That sounds like the AFS? Which we already have? Here I'm wanting to return the list of allele frequencies for each SNP. Triggered my AFS implementation PTSD there - would *not* want to do that again! — You are receiving this because you commented. Reply to this email directly, view it on GitHub <#504 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ABQ6OHZCOYZBPTG4UORY6TDRKGOCJANCNFSM4LWYYVBA> .

[petrelharp]

This computes the total frequency of all mutations that are segregating in the samples, by site. People will also want to know the frequencies of mutations, separately: we should have another method that returns something of length equal to the number of mutations. I don't think we can compute this easily with the general stat framework, but here is a method to pull out this information for the mutations at a particular site:

def mut_counts(site):
    counts = {}
    t = ts.at(site.position)
    for m in site.mutations:
        counts[m.id] = t.num_samples(m.node)
        if m.parent != tskit.NULL:
            assert(m.parent in counts)
            counts[m.id] -= counts[m.parent]
    return counts