Improved algorithm for site-mode divergence matrix

[jeromekelleher]

The divergence matrix code merged in #2736 gives a reasonably efficient way to compute the pairwise branch-mode divergence among a set of samples, but it's very poor for site mode. This is probably because we're not using the fast MRCA calculations in the internal loop.

The simplest way to do this may be to keep a count of the total number of mutations that are ancestral to each node as we to along the genome, and then we should be able to use the same framework as the branch stat calculation.

This will be complicated by keeping track of windows also though.

Doing this will need an incremental algorithm that can start and stop at a given point efficiently, so we should probably get some C code in for doing this in a generic way first, as dicussed in #2778

[jeromekelleher]

Actually, I don't see how we can make this much more efficient than the naive approach currently implemented. My assertion above is wrong:

The core of the current branch mode operation is to do this:

                    w = tree.mrca(u, v)                  
                    tu = ts.nodes_time[w] - ts.nodes_time[u]
                    tv = ts.nodes_time[w] - ts.nodes_time[v]                  
                    d = (tu + tv) * span
                    D[i, j, k] += d

for pairs of samples u, v and their indexes j, k.

Above I implied we could do roughly the same thing for the site mode if we kept track of the total number of mutations descending from each node, i.e. replacing ts.nodes_time with num_descending_mutations. However this doesn't work because the num_descending_mutations for the MRCA isn't just the sum of the num_descending_mutations of the two incoming paths for the target samples u and v, it also includes all mutations for any other children (i.e., if it's a polytomy). So, to use this logic we need to keep track of which child of the MRCA u and v came in from, in which case there's nothing to be gained from using the Schieber Vishkin algorithm for computing MRCAs.

More generally, it's not clear to me that there is much to be gained from using more complex algorithms either. Fundamentally, each mutation defines a partition of the samples into the set of those that descend from the mutation and those that don't (say A and B), and we need to update the divergence matrix by adding one to each cell for the cartesian product of A and B. In principle, we can avoid incrementing-by-one each time if two or mutation mutations share the same sample partition. However, running some simple coalescent simulations (human-like) I get

n       #muts   singleton fraction      mean
10      5532    0.024945770065075923    9.250836120401338
100     10564   0.32677016281711474     2.04570100697134
250     12369   0.3879860942679279      1.795731707317073
500     13814   0.4302881135080353      1.6597380752132644
1000    14731   0.47987237797841287     1.533041939848059
10000   19604   0.6191083452356662      1.3052799786936546

So, for n=10k, 61% of the mutations have a unique sample set, and the mean number of mutations per sample set is 1.3. So, you're not going to gain much by trying to gather together mutations that just update exactly the same parts of the pairwise divergence matrix.

Code:

  for n in [10, 100, 250, 500, 1000, 10000]:
        ts = msprime.sim_ancestry(
            n,
            ploidy=1,
            population_size=10**4,
            sequence_length=1e7,
            recombination_rate=1e-8,
            random_seed=seed,
        )
        ts = msprime.sim_mutations(ts, rate=1e-8, random_seed=seed)
        import collections
        counter = collections.Counter()
        for tree in ts.trees():
            for site in tree.sites():
                for mut in site.mutations:
                    ss = frozenset(tree.samples(mut.node))
                    counter[ss] += 1

        vals = np.array(list(counter.values()))
        ones = np.sum(vals == 1)
        print(n, ts.num_mutations, ones / ts.num_mutations, np.mean(vals), sep="\t")

So, I can see two ways of doing this: mutation-by-mutation and sample-by-sample.

For the mutation-by-mutation approach we look down from each mutation to construct the sets A and B (which descend and do not descend) from the mutation. In worst case this is O(num_mutations * n^2), assuming each mutation exactly splits the set of samples (which they won't).

For the sample-by-sample approach, we go tree-by-tree and sample and look at each pair of samples in turn. This is O(num_trees * n^2 * tree_height), because we have to climb the tree to compute the number of mutations on the path for each pair (see above).

It seems then like the mutation-by-mutation approach is a winner then in most cases?

[jeromekelleher]

Huh, a quick numba implementation of the mutation x mutation approach is much faster than both the current divergence_matrix code and the original tree-stats based approach. So, it's worth implementing this as a next step I think, and putting off better algorithms to later.