Fetching all populations where an allele is present

[clawsoon]

Hello again,

I believe that I just used a bunch of treeseq data incorrectly. What I want to find is all of the populations where an allele is currently found, along with associated info about when it first appeared, etc. Before I regenerate everything, I'm hoping someone can confirm that I was, in fact, doing it wrong the first time, and that I am doing it correctly now.

First attempt, which I'm pretty sure is wrong:

for tree in ts.trees():
  for mutation in tree.mutations():
    # This is wrong because leaves will contain both samples with the
    # mutation (which we want) and samples which have reverted to the
    # ancestral state or mutated to something else (which we don't want).
    populations = set(ts.node(leaf).population for leaf in tree.leaves(mutation.node))

Second attempt, which I'm hoping is right:

for variant in ts.variants():

  for allele_index, allele in enumerate(variant.alleles):

    if allele_index == 0 or not allele:
      continue

    sample_nodes = numpy.where(variant.genotypes == allele_index)[0]
    full_nodes = [ts.node(node) for node in sample_nodes]

    # Get all individuals with the allele.
    individuals = set(node.individual for node in full_nodes)

    # Get all populations with the allele.
    populations = set(node.population for node in full_nodes)

    # Get date when allele first appeared (whether or not it disappeared or
    # reappeared again later or elsewhere).
    # Is there a more efficient way to do this?
    mutation_nodes = [ts.node(mutation.node) for mutation in variant.site.mutations if mutation.derived_state == allele]
    mutation_nodes.sort(reverse=True, key=lambda node: node.time)
    oldest = mutation_nodes[0]
    first_seen = oldest.time

    # Get site info.
    position = variant.site.position
    ancestral_state = variant.site.ancestral_state
    derived_state = allele

Any feedback would be appreciated. Thanks!

[benjeffery]

Many tree sequences you'll find in the wild have only bi-allelic snps with no back mutations, and under that assumption your first code is correct, so worth checking that first.
At a quick glance your second code looks correct, but as you say will be slow. It may be quicker to first make a dict of mutations indexed by their parents, and use that to go down the mutation chain, and identity when you can safely call leaves. Let me know if I'm making sense here!

[clawsoon]

Thanks, @benjeffery. I think I'm wrapping my head around what you're saying, though my brain is still a little foggy this morning.

The second code finished running while I was sleeping, so hopefully it's correct, since I'm about to launch a few hundred hours of processing based on it. :-) I guess the core of it that I'd want to confirm for sure is these bits:

for variant in ts.variants():
  for allele_index, allele in enumerate(variant.alleles):

    # Skip ancestral alleles and missing data.
    if allele_index == 0 or not allele:
      continue

    # I think(??) this gives me indices that I can use in ts.node()?
    sample_nodes = numpy.where(variant.genotypes == allele_index)[0]

    full_nodes = [ts.node(node) for node in sample_nodes]
    populations = set(node.population for node in full_nodes)

    # I think(??) this gives me all the mutations which ended up
    # at the target state for the site?
    mutation_nodes = [ts.node(mutation.node) for mutation in variant.site.mutations if mutation.derived_state == allele]

One thing that's throwing me a bit is that for chr22_q this method finds 1,134,248 non-ancestral alleles, while ts.num_mutations is 3,257,559. Are there really nearly 3 times as many mutations as there are non-ancestral alleles? Summing up the mutations per site using this next bit of code does add up to 3,257,559 (and gives me a max of 530 mutations for a single site), so I guess there truly are nearly 3 times as many mutations as sites overall?

counter = collections.Counter()
for site in ts.sites():
  counter[len(site.mutations)] += 1
sum(k*v for k, v in counter.items())
# 3,257,559

Where I hesitated at trying to make the tree method work was not feeling smart enough to confirm that I'd be keeping track of everything correctly in the case of multiple nesting. It felt like there'd be a lot of tracking of which set of leaves I'd have to subtract from which other set of leaves to figure out which state a leaf finally ended up in. My brain gave me a busy signal when I tried to contemplate it, so I didn't pursue it further. It's possible that it's simpler than I'm guessing.

[clawsoon]

Another difference I'm seeing between the two approaches: I'm doing simple stuff with "how many populations does this appear in?" and "how many individuals does this appear in?"

The first approach with walking the trees assigns many more "how many does this appear in" to a single individual, since each "this" is a mutation. If the same allele appears in multiple people as a result of separate mutations, each appearance ends up in a different bucket.

The second approach with walking the variants assigns many more "how many does this appear in" to multiple individuals, since each "this" is an allele. If multiple separate mutation events resulted in the same allele, it doesn't care; it puts them all in the same bucket.

So now I just have to decide which of the two I care about. :-)

[benjeffery]

I guess there truly are nearly 3 times as many mutations as sites overall?

Yes, some sites have many mutations as they are either error or other artefacts where the site isn't in linkage with it's surrounding region.