Creating a mutation of a given frequency.

[NW-Anderson]

Hi all

I am wanting to begin a simulation in slim with ~50 selected snps at frequency 0.05.

Currently, I begin by following along the pyslim vignette "starting with diversity generated by coalescent simulation". I'll simulate a coalescent history, and then add mutations at a low rate. Next I search for mutations with a frequency between 0.049 - 0.051, sample ~ 50 of these to assign selection coefficients and remove all other mutations from the mutation table. Finally, I output the tree sequence and move into SLiM. My issue is that this is not efficient. There is a lot of simulating that is only removed from the final tree sequence.

So my question is: Is there a way to drop mutations onto a tree sequence such that they have a given frequency at t = 0?

Thanks in advance!
Nathan

[petrelharp]

Hm, good question. It would be possible to do this - by searching through the trees for branches that were above 50% of the samples and then picking some proportional to their length - but it would be a good deal slower than the procedure you're outlining. Adding mutations is very fast, and I think the other steps should be do-able reasonably quickly; have you measured the time taken in these various steps to make sure it's actually slowing you down substantially?

[petrelharp]

How much time do you spend in step 2, and what mutation rate are you using?

[NW-Anderson]

Hi Dr Ralph

Sorry for the delay.

I had been simulating 10,000,000 sites sampled from along one chromosome arm and adding mutations at a rate of 1.2e-9. With these parameters sim_mutations() could take ~20 minutes to complete and there was a couple hours total run time, most of that spent in sim_ancestry(). But, of these many sites, and thousands of mutations added, I am only interested in keeping about 50 of them. Ideally, I'd like to simulate the ancestry only for 50 sites. And then I could drop a mutation at frequency on 0.05 frequency at each site.

To try and reduce runtime, I've tried simply decreasing the number of sites or decreasing the mutation rate. An issue I've run into with this method is there would not exist enough mutations in the 0.049 - 0.051 range. I have had some luck today by decreasing the number of sites and increasing the mutation rate, both by a factor of 100. Doing this decreased the total run time to a bit less than an hour and ~5 minutes was spent in sim_mutations(), which is a significant improvement. However, I am not sure whether messing with the mutation rate like this defeats the purpose of this simulation - to get a bunch of mutations ~ 0.05 frequency with a "realistic" ancestry.

[NW-Anderson]

I realized my above code only works in the special case where no two mutations occur in the same site. Because we are calculating the mutation frequency at each site, this can mess up which rows we are selecting from the mutation table. Also, calling ots.tables.mutations[k], rather than ots.mutation(k) will gobble up your memory.

I've gotten around the first issue by simply removing mutations from the table that whose site has already appeared in the table.`

# loading mutated tree sequence
ots = tskit.load("seed" + str(seed) + "_nofilter.trees")
print(f"The tree sequence now has {ots.num_mutations} mutations, at "
      f"{ots.num_sites} distinct sites.")


# because our p calculating later on reports the frequency of **any**
# mutation at each site, and we are interested in the frequency of specific
# mutations...and I cant figure out how to find that straightforward...
# a brute workaround is to remove all the mutations that occur on sites already 
# containing a mutation

tables = ots.tables
tables.mutations.clear()
taken_sites = []
for mut in ots.mutations():
    if mut.site not in taken_sites:
        taken_sites.append(mut.site)
        _ = tables.mutations.append(mut)
ts = tables.tree_sequence()
print(f"The tree sequence now has {ts.num_mutations} mutations, at "
      f"{ts.num_sites} distinct sites.")

# now that we only have one mutation per site. We can calculate mutation freqeuncies 
currentGenomes = ts.samples(time=0)
num_nodes = np.array(len(currentGenomes))
p = ts.sample_count_stat([currentGenomes], lambda x: x/num_nodes, 1, windows='sites',
        strict=False, span_normalise=False, polarised=True)

# find out which mutations are around 0.05 frequency. I want to keep them
keep = np.where(np.logical_and(p < 0.075, p > 0.025))[0]
tables = ts.tables
tables.mutations.clear()
# there shouldnt be any more mutations sharing the same site
# but make double sure
taken_sites = []
for k in keep:
    mut = ts.mutation(k)
    if mut.site not in taken_sites:
        taken_sites.append(mut.site)
        _ = tables.mutations.append(ts.mutation(k))
ts = tables.tree_sequence()
ts = ts.simplify()
print(f"The tree sequence now has {ts.num_mutations} mutations, at "
      f"{ts.num_sites} distinct sites.")

#  make sure that we only have the frequency range we want
currentGenomes = ts.samples(time=0)
num_nodes = np.array(len(currentGenomes))
p = ts.sample_count_stat([currentGenomes], lambda x: x/num_nodes, 1, windows='sites',
        strict=False, span_normalise=False, polarised=True)
max(p)
min(p)
# dumping the filtered tree
ts.dump("seed" + str(seed) + ".trees")

[petrelharp]

Good point about multiple mutations. It's on our list to implement a "mutation frequency" function as well as this "derived allele frequency" one you are using: see #504.

Instead of your code here, though - you might use your code above to look for sites with derived frequency at 50%, then subset to those with only a single mutation? (You're right, though, about ots.tables.mutations[k] being very bad for memory!) That'll get you mutations conditioned to be the only one at a site and at 50% frequency.

My strong suggestion is that since sim_mutations is a relatively small fraction of your runtime, you should just leave well enough alone. However... if you want to really lay down only 50 mutations, what you need to do is (a) iterate over the trees, and keep track of a list of those branches that are subtended by 50% of the samples, and (b) once you've got the full list, pick 50 of them, drawing with probability proportional to length * span. I don't know of any shortcuts you can take here, and it seems likely to me to be slower than subsetting sim_mutations... but, who knows?

[NW-Anderson]

Hi Peter

Thank you for your suggestion to change the order of filtering in my original code!

I've been playing around with both methods and I've found looping through trees and nodes is slightly faster than calling sim_mutations() and filtering by frequency. For example, I've been profiling a 100 Mb region with a sample of 20,000 individuals from a population of 100,000. And I want 20 mutations in the 0.049 - 0.051 range. All in all, I got the 20 mutations I wanted for this example in ~ 4 minutes by looping through each tree, vs ~ 7 minutes the filtering method.

I've found the major difference between the two methods is the number of sites needed. Using the filtering method, I need to simulate ancestry for 100,000 sites in order for there to be > 20 mutations in the range I want. This is using a mutation rate of 1.2e-7, I'm sure that there are other combinations of number of sites and mutation rate that work faster, but I got tired of profiling. Alternatively, I've been able to simulate ancestry for as few as 50 sites and still had plenty of branches at unique sites to drop mutations on. You're correct that looping through every node in every tree takes some time. But, less than is spent in sim_ancestry() + sim_mutations() using the filtering method.

It sounds like Dr Thornton has found a method to find derived node counts that is much faster than our idea of looping over nodes in each tree. I'm going to mess around with his edge differences method and get back to everyone!

In the meantime, I'm attaching my code for the functions to filter a mutated tree sequence for mutations of a given frequency, trim_mutations() and to add mutations of a given frequency to an unmutated tree sequence, add_mutations(). In case anyone is interested.

Thanks again for your thoughts!

EDIT: I realized my previous version of add_mutations(), calculated probabilities of mutations falling on each branch, but never used this probability for dropping mutations. I think this version is correct.

def weighted_shuffle(items, weights):
    order = sorted(range(len(items)), key=lambda i: -random.random() ** (1.0 / weights[i]))
    return [items[i] for i in order]

def add_mutations(ts, allele_frequency, num_mutations=1, atol=0.001):
    """
    :param ts: the un-mutated tree sequence
    :param allele_frequency: the desired allele frequency
    :param num_mutations: the number of mutations to add at this freq
    :param atol: absolute tolerance (allele_frequency =/- atol), setting
        to zero would mean all mutations are exactly the given frequency
    """
    # get the branches satisfying the desired allele frequency
    branches = []
    for tree in ts.trees():
        tree_interval = (tree.interval.left, tree.interval.right)
        for n in tree.nodes():
            samples = tree.num_samples(n)
            f = samples / ts.num_samples
            if np.isclose(f, allele_frequency, atol=atol):
                time_span = (tree.time(n), tree.time(tree.parent(n)))
                branches.append((n, time_span, tree_interval))
    
    # how many branches do we have to choose from to drop mutations on
    num_branches = len(branches)
    if num_branches == 0:
        raise ValueError("Didn't find any branches with given conditions")
    
    # finding probability of a mutation happening on each branch
    probs = np.zeros(num_branches)
    for i, b in enumerate(branches):
        # branch is weighted by the product of its time and sequence spans
        probs[i] = (b[1][1] - b[1][0]) * (b[2][1] - b[2][0])
    # needs to be probability masses
    probs /= np.sum(probs)
    
    # copying ts tables to work with
    tables = ts.tables
    
    # this will hold the positions that have been assigned mutations 
    #   so we can ensure no >1 mutation per site (biallelic loci)
    taken_pos = []
    
    # we want to loop through the branches in a random order and start adding mutations
    branch_id = [i for i in range(num_branches)]
    # random.shuffle(branch_id)
    branch_id = weighted_shuffle(branch_id, probs)
    
    for current_branch in branch_id:
        if tables.mutations.num_rows < num_mutations:
            branch = branches[current_branch]
            left,right = branch[2]
            
            pos = math.floor(np.random.rand() * (right - left) + left)
            if pos not in taken_pos:
                taken_pos.append(pos)
                site_id = tables.sites.add_row(position=pos, ancestral_state="0")
                bottom, top = branch[1]
                time = np.random.rand() * (top-bottom) + bottom
                tables.mutations.add_row(site_id, branch[0], "1", time=time)
        else:
            break
            
    tables.sort()
    ts = tables.tree_sequence()

    
    
    if num_mutations > ts.num_mutations:
        print("Didn't find enough branches at distinct sites with given conditions "
              "returning all we can.")
    else:
    
        # slim down the site table etc
        ts = ts.simplify()
        print(f"The tree sequence now has {ts.num_mutations} mutations, at "
              f"{ts.num_sites} distinct sites.")

    # return a tree sequence with those mutations
    return ts

def trim_mutations(ts, allele_frequency, num_mutations=1, atol=0.001):
    """
    :param ts: the un-mutated tree sequence
    :param allele_frequency: the desired allele frequency
    :param num_mutations: the number of mutations to add at this freq
    :param atol: absolute tolerance (allele_frequency =/- atol), setting
        to zero would mean all mutations are exactly the given frequency
    """
    print(f"The original tree sequence  has {ts.num_mutations} mutations, at "
          f"{ts.num_sites} distinct sites.")
    
    # First, filter for sites that have mutant frequency near where we want
    currentGenomes = ts.samples(time=0)
    num_nodes = np.array(len(currentGenomes))
    p = ts.sample_count_stat([currentGenomes], lambda x: x/num_nodes, 1, windows='sites',
            strict=False, span_normalise=False, polarised=True)

    keep_sites = np.where(np.logical_and(p < (allele_frequency + atol), 
                                   p > (allele_frequency - atol)))[0]
    keep_muts = np.where(np.isin(ts.tables.mutations.site, keep_sites))[0]
    
    tables = ts.tables
    tables.mutations.clear()
    for k in keep_muts:
        mut = ts.mutation(k)
        # these next couple lines keep parent ids from referring
        #   to mutations that are no longer in the mutation table.
        #   were about to filter for sites that only have a sinlge
        #   mutation. So it doesnt matter what we do with them
        #   just that it doesnt throw a mutation when we try to call
        #   tree_sequence() in a few lines.
        if mut.parent != -1:
            mut.parent = -1
        
        _ = tables.mutations.append(mut)

    ts = tables.tree_sequence()
    ts = ts.simplify()
    print(f"After filtering for frequency, the tree sequence now has {ts.num_mutations} mutations, at "
          f"{ts.num_sites} distinct sites.")
    
    # next filter for sites with only a single mutation
    sites = ts.tables.mutations.site
    not_dup_sites = [x for x, count in collections.Counter(sites).items() if count == 1]
    not_dup_muts = np.where(np.isin(ts.tables.mutations.site, not_dup_sites))[0]
    
    tables = ts.tables   
    tables.mutations.clear()
    for k in not_dup_muts:
        _ = tables.mutations.append(ts.mutation(k))
    ts = tables.tree_sequence()
    ts = ts.simplify()
    print("After filtering for biallelic sites (those with only one mutation) "
          f"the tree sequence now has {ts.num_mutations} mutations, at "
          f"{ts.num_sites} distinct sites.")
    
    # finally filter down to the number of mutations we want
    if ts.num_mutations < num_mutations:
        print("There are not num_mutations in the given frequency range in "
              "the original tree sequence. Returning all mutations present "
              "after filtering for biallelic sites.")
    else: 
        tables = ts.tables
        tables.mutations.clear()
        keep = np.random.choice(range(ts.num_mutations),  
                                  size= num_mutations,
                                  replace = False)
        keep.sort()
        
        for k in keep:
            _ = tables.mutations.append(ts.mutation(k))
            
        ts = tables.tree_sequence()
        ts = ts.simplify()
        print(f"The tree sequence now has {ts.num_mutations} mutations, at "
              f"{ts.num_sites} distinct sites.")
        
    
    return ts

[molpopgen]

I have been working on this myself lately. Using edge differences to get at derived node counts is about two orders of magnitude faster than tree traversal. If you need the total times on the trees for any reason, then tree traversal is faster than edge differences.

Output:

num trees = 3096
Function 'times_from_trees' took 0.030463218688964844
Function 'times_from_edge_differences' took 0.11167454719543457
Function 'sample_counts_from_trees' took 74.86195659637451
Function 'sample_counts_from_edge_differences' took 0.07632184028625488

Code:

from time import time
import msprime
import tskit
import numpy as np


def timer(func):
    def wrapper(ts):
        t1 = time()
        rv = func(ts)
        t2 = time()
        print(f"Function {func.__name__!r} took {t2 - t1}")
        return rv
    return wrapper


@timer
def times_from_trees(ts: tskit.TreeSequence):
    rv = np.zeros(ts.num_trees)
    for (i, t) in enumerate(ts.trees()):
        rv[i] = t.total_branch_length*t.span

    return rv


@timer
def sample_counts_from_trees(ts: tskit.TreeSequence):
    rv = []
    for (i, t) in enumerate(ts.trees()):
        counts = np.zeros(ts.num_nodes, dtype=np.uint32)
        for node in t.nodes():
            counts[node] = len([i for i in t.leaves(node)])

        rv.append(counts)

    return rv


@timer
def times_from_edge_differences(ts: tskit.TreeSequence):
    rv = np.zeros(ts.num_trees)
    branch_lengths = np.zeros(ts.num_nodes)

    for (i, e) in enumerate(ts.edge_diffs()):
        for edge_out in e.edges_out:
            if edge_out.parent != tskit.NULL:
                branch_lengths[edge_out.child] = 0.0
        for edge_in in e.edges_in:
            if edge_in.parent != tskit.NULL:
                ptime = ts.node(edge_in.parent).time
                ctime = ts.node(edge_in.child).time
                branch_lengths[edge_in.child] = ptime - ctime

        rv[i] = np.sum(branch_lengths)*(e.interval.right - e.interval.left)

    return rv


@timer
def sample_counts_from_edge_differences(ts: tskit.TreeSequence):
    rv = []
    counts = np.zeros(ts.num_nodes, dtype=np.int32)
    for s in ts.samples():
        counts[s] = 1
    parents = [tskit.NULL for i in range(ts.num_nodes)]
    for tree_index, e in enumerate(ts.edge_diffs()):
        for o in e.edges_out:
            assert counts[o.child] <= counts[
                o.parent], f"{o.child} {counts[o.child]} {counts[o.parent]}, {counts}"
            lc = counts[o.child]
            p = parents[o.child]
            while p != tskit.NULL:
                counts[p] -= lc
                p = parents[p]
            parents[o.child] = tskit.NULL

        for edge_in in e.edges_in:
            lc = counts[edge_in.child]
            parents[edge_in.child] = edge_in.parent
            p = parents[edge_in.child]
            while p != tskit.NULL:
                counts[p] += lc
                p = parents[p]

        rv.append(np.copy(counts))

    return rv


ts = msprime.sim_ancestry(1000, recombination_rate=1e-7,
                          sequence_length=1e6, random_seed=666*42, population_size=1000)

assert ts.num_trees > 1

# print(ts.draw_text())

print(f"num trees = {ts.num_trees}")

times_trees = times_from_trees(ts)

times_edge_diffs = times_from_edge_differences(ts)

counts_from_trees = sample_counts_from_trees(ts)

counts_from_edges = sample_counts_from_edge_differences(ts)

assert len(counts_from_trees) == len(counts_from_edges)

for i, j in zip(counts_from_trees, counts_from_edges):
    assert np.all(i - j == 0), f"{np.unique(i - j, return_counts=True)}"

[NW-Anderson]

Hi Dr Thornton

This is a really clever idea, thank you for sharing!

From my latest reply to Dr Ralph, I think identifying branches with the correct proportion of derived nodes in the present is the way to go.

If you don't mind, I'll soup up my add_mutations() function using your edge differences idea.

I'll get back to everyone once I have any updates. Thanks again for both of your help!

[molpopgen]

Go ahead!

For my application, I can't put mutations down as a plain function of (branch length)*span, so I have to use a procedure like this.