new (?) statistic to efficiently calculate shared times between every pair of samples at every tree

[mmosmond]

Hi all, I've been messing around with different ways to calculate the shared time between every pair of samples in many trees, which is a useful metric because it describes the covariance of characteristics of the samples (e.g., traits, locations) under Brownian motion. At first I couldn't figure out how to use an existing statistic to do this, so I used the general_stat function to make one of my own:

import numpy as np
def shared_times(ts):
    """Use general_stat function to calculate shared branch lengths between all samples at all trees.
    
    Parameters
    ts: tskit tree sequence
    
    Returns
    A (n x k x k) numpy array, with n the number of trees in the sequence and k the number of sample nodes. 
    At each tree the (k x k) matrix gives the shared evolutionary times between each pair of sample nodes.
    """
    
    k = ts.num_samples #number of samples
    W = np.identity(k) #each node i in [0,1,...,k] given row vector of weights with 1 in column i and 0's elsewhere 
    def f(x): return (x.reshape(-1,1) * x).flatten() #determine which pairs of samples share the branch above a node, flattened

    return ts.general_stat(
        W, f, k**2, mode='branch', windows='trees', polarised=True, strict=False
    ).reshape(ts.num_trees, k, k)

For example:

import msprime
ts = msprime.sim_ancestry(samples=5, population_size=1e4, sequence_length=1e4, recombination_rate=1e-8, random_seed=1)
sts = shared_times(ts)

The result can also be achieved by simply looping through trees and samples to calculate pairwise TMRCAs:

def shared_times(ts):
    k = ts.num_samples
    sts = np.zeros((ts.num_trees,k,k))
    for t,tree in enumerate(ts.trees()):
        T = tree.time(tree.root)
        for i in range(k):
            for j in range(i):
                st = T - tree.tmrca(i,j)
                sts[t,i,j] = st
                sts[t,j,i] = st
    return sts

but this is considerably slower when there are many trees (although it seems to be faster when there are few trees with many samples, e.g., ts = msprime.sim_ancestry(samples=2e2, population_size=1e4, random_seed=1) -- not sure why).

Coming back to this much later (today), I realized you can use the divergence statistic, with a few tweaks, to do the same:

def shared_times(ts):
    
    k = ts.num_samples
    
    # get 2*tmrcas
    sample_sets = [[i] for i in ts.samples()]
    indexes = [(i,j) for i in ts.samples() for j in range(i,k)] #compare each sample with each other only once (entries of upper triangular)
    divs = ts.divergence(sample_sets=sample_sets, indexes=indexes, mode='branch', windows='trees')
    divs[np.isnan(divs)] = 0 #tmrcas with self are said to be nan, so convert to 0

    # convert to matrices
    divs_mat = np.zeros((ts.num_trees,k,k)) 
    for ix,(i,j) in enumerate(indexes):
        divs_mat[:,i,j] = divs[:,ix] #convert list to upper triangle
        divs_mat[:,j,i] = divs[:,ix] #fill in lower triangle symmetrically
    
    # convert to shared times
    sts = np.zeros(divs_mat.shape)
    for i,div in enumerate(divs_mat):
        sts[i] = (np.max(div) - div)/2 #convert from 2*tmrcas to shared times
    
    return sts

This latter method seems to be the fastest for small tree sequences (eg, ts = msprime.sim_ancestry(samples=5, population_size=1e4, sequence_length=1e6, recombination_rate=1e-8, random_seed=1)) but loses this advantage as tree sequences get larger (eg, msprime.sim_ancestry(samples=1e2, population_size=1e4, sequence_length=1e6, recombination_rate=1e-8, random_seed=1)). It also uses >2x the RAM of the first method, which quickly becomes important with larger tree sequences.

I'm wondering if my general_stat matrix formulation would be useful as a new statistic of its own: related to the divergence statistic but calculated differently (I think), sometimes faster, sometimes less memory intensive, and semantically simpler (for creating matrices). If so, I guess I'd have to also think about what happens as you vary the options in general_stat, e.g., mode='node'. Anyway, just wanted to put this out there in case it is helpful/useful. Thanks for building all this stuff!

(This is my first time writing an issue so sorry if I've messed up somehow -- I'd be keen to hear how to do this better!)

[petrelharp]

Hi, @mmosmond! This is a great issue, thanks! Fortunately or unfortunately, we might have already done this - I haven't done a detailed comparison, but from your description I think this is the same thing as the genetic relatedness method that @brieuclehmann implemented in #898. (see the lengthy thread in that issue for discussion about what exactly it should be called, etcetera)

Have a look and see what you think?

[mmosmond]

Amazing, thanks @petrelharp! I should have done a better search of the existing issues -- of course you've already done this! You are right, genetic_relatedness can produce the mean-centred shared-time matrix:

import msprime
import numpy as np

def shared_times_gstat(ts):
    """general_stat method, making use of flattened matrices"""
    
    k = ts.num_samples #number of sampled genomes
    W = np.identity(k) # each node i in [0,1,...,k] given row vector of weights with 1 in column i and 0's elsewhere 
    def f(x): return (x.reshape(-1,1) * x).flatten() # matrix with 1's where branch above a node with value x contributes to shared time between samples 

    return ts.general_stat(
        W, f, k**2, mode='branch', windows='trees', polarised=True, strict=False, span_normalise=True
    ).reshape(ts.num_trees, k, k)

def shared_times_rel(ts):
    """relatedness method (note the matrices are mean centred)"""
    
    n = ts.num_samples
    sample_sets = [[i] for i in ts.samples()]
    indexes = [(i,j) for i in ts.samples() for j in ts.samples()]
    
    return ts.genetic_relatedness(
        sample_sets, indexes, mode='branch', windows='trees', span_normalise=True, proportion=False
    ).reshape(ts.num_trees,n,n)

def mean_centre(square_matrix):
    """function to mean centre a square matrix"""
    
    k,k = square_matrix.shape
    Tmat = np.identity(k) - [[1/k for _ in range(k)] for _ in range(k)] 
    
    return np.matmul(Tmat, np.matmul(square_matrix, np.transpose(Tmat)))  

# test it out
k = 100
ts = msprime.simulate(k, recombination_rate = 1, random_seed=1, length=10)

sts_rel = shared_times_rel(ts) #genetic_relatedness method (mean centred)
sts = shared_times_gstat(ts) #general_stat method (not mean centred)
sts_mc = np.array(list(map(mean_centre, sts))) #mean centre each matrix to compare

np.all(sts_mc - sts_rel < 1e-12) #check the mean centred matrices are essentially the same

True

And a quick look at the performance shows that the two methods take about the same time:

import time
import tracemalloc

def test_fun(fun, ts, nreps=1, **kwargs):
    """function to test these functions"""

    print("\n",fun)
    tracemalloc.start()
    t0 = time.time()
    for _ in range(nreps):
        fun(ts,**kwargs)
    t1 = time.time()
    current, peak = tracemalloc.get_traced_memory()
    print(f"Max mem: {peak / 10**6}MB")
    tracemalloc.stop()
    print("seconds:", t1 - t0)

for fun in [shared_times_gstat, shared_times_rel]:
    test_fun(fun, ts, nreps=1)

 <function shared_times_gstat at 0x110bf8ee0>
Max mem: 12.620492MB
seconds: 0.2728910446166992

 <function shared_times_rel at 0x10632ac10>
Max mem: 25.42874MB
seconds: 0.32077693939208984

The main difference in performance is that the genetic_relatedness method uses twice the memory, but this can be fixed by just computing the upper diagonal (i.e., using indexes = [(i,j) for i in ts.samples() for j in range(i,n)] so that we do each pairwise comparison just once) and then filling in the full matrix from that if needed.

The main difference in output is that it looks like genetic_relatedness cannot output the non-mean-centred matrices. As you mentioned in #898, people will often want the mean-centred matrices in practice, although perhaps not always. For instance, I am interested in calculating the shared times between sample lineages in multiple epochs (and then multiplying the times in each epoch by a constant, summing over epochs, and then mean-centring), but I don't think this can be done starting with only the mean-centred matrices. Is there still talk of introducing a centre=False option? I used the sample_count_stat method (as in #898) to get the non-mean-centred matrices but this was much, much slower.

I don't know how things are working under the hood here, but it seems like the (flattened-)matrix-style of the general_stat method above could be faster/simpler than the pairwise approach of the genetic_relatedness method?

Thanks for getting back to me and pointing me in the right direction. Great stuff! Glad I didn't have to weigh in on the definition of "kinship"...

[petrelharp]

Gee, this is very interesting! I'm surprised that the general_stat version is faster than the just-C version - any insight why? Perhaps because it doesn't have to loop over the indices? Here's the summary function that it uses as opposed to just outputting all pairwise entries?

It'd be great to add centre=False; do you have any interest in helping make this happen (or, at least verifying what we should be doing)? There isn't much code: here's what got added to add the genetic_relatedness function (although you'll see in the link above that some of it has changed a bit since); I'd imagine adding a tsk_treeseq_relatedness_noncentred method, perhaps? I've forgotten the details of all this - do we just want to remove the - meanx terms in the summary function?

As for speed, I think that the indexed version is about the same speed, and more flexible, so I'm not motivated to switch it up. But for memory, there is an underlying issue that we'd also like to take care of that's making it use lots of memory for large matrices - basically, it's cacheing a vector of length equal to the output at each node; in cases like this where the output is O(n^2) this can be a bad idea.

[jeromekelleher]

Just one quick note here - @petrelharp, do we want a new function tsk_treeseq_relatedness_noncentred or could we use an options flag to `tsk_treeseq_relatedness? It would reduce a bit of boilerplate, wouldn't it?

[petrelharp]

Ah, good point - an options flag would be maybe better.

And, @mmosmond - how about I'll make a quick draft of this, since I'm more familiar with the workings?

[mmosmond]

Sounds good @petrelharp. I've made a start at this today, following the developer instructions -- I've got a Python function and test working, but I'm a C novice so I'm still wrapping my head around how to pass the option in trees.c. And yeah, I think you can just drop the - meanx terms (although I'm confused about the 1/2). If it helps, if you give me the rough idea of how to do this I can try and make a PR (seems like a good learning experience? :) ).

[petrelharp]

Awesome! Want to make a PR and I'll fill in / provide input over there? (I'm not sure which thing you want a rough idea of how to do?)

I'm not sure about the 1/2, but it might be because the summary function is applied to each allele: https://tskit.dev/tskit/docs/stable/stats.html#polarisation - so at binary alleles it effectively gets called with both p and 1-p as arguments.

[mmosmond]

Thanks @petrelharp. Just made the PR. Right, that 1/2 logic makes sense to me.

Yeah, my guess why the general_stat method is equally fast is that it is filling in all relevant entries of the matrix at the same time (i.e., if node 3 is above samples 0 and 1 but not 2, then the branch length above 3 gets added to positions (0,0), (0,1), (1,0), and (1,1) simultaneously through matrix multiplication). Whereas with the genetic_relatedness method you loop through all samples/matrix-entries, as you say. But there is a good chance this is obvious to you and my reasoning is still a few layers above (=behind) yours...

Yeah, the memory usage seems to get pretty big for the general_stat method, which is one of the reasons why I ended up using a different strategy in a recent project. I'll have to do some thinking about the flexibility of the general_stat and genetic_relatedness methods...

[petrelharp]

I'm not quite following about matrix multiplication - we should chat sometime?

[mmosmond]

Well I guess "matrix" is a bit misleading because we flatten the matrix into a vector, but what I mean is that the function we pass to the general_stat method is

def f(x): return (x.reshape(-1,1) * x).flatten()

The weight for a node is a vector of length k = num_samples with a 1 in index i if sample i is a descendant. So I think of f as making a (flattened-)matrix with entry (i,j) = 1 if both i and j are descendants, and 0 otherwise. So with mode=branch, as you move up the tree you multiply the length of a branch by this (flattened)matrix, which adds the branch length to the appropriate pairs of samples. But not sure this helps... I'd be happy to chat this week sometime

[petrelharp]

That vaguely makes sense, but I've not got the big picture of why that'd be more efficient...

[jeromekelleher]

Can this statistic you're interested in be computed in terms of the divergence matrix @mmosmond? We've just merged a reasonably efficient (well for branch stats anyway; we'll need to follow up with a more efficient way for site mode, #2779 ) approach for computing pairwise divergence in #2736 , it it would be good to know if there's other things we should be thinking about layering on top of it.