Two locus general matrix

[lkirk]

Description

This is the last of the required components for the LD matrix methods. I wanted feedback on the API before I add documentation, but this method is complete and tested. The final things to do are to leak check the cpython code and add some documentation.

This feature enables a user to implement their own two-locus count statistic in python, similar to ts.sample_count_stat. User functions take two arguments, the first is a matrix of haplotype counts and the second is a vector of sample set sizes. For instance, this is how we would implement D with this api:

def D(X, n):
    pAB, pAb, paB = X / n
    pA = pAb + pAB
    pB = paB + pAB
    return pAB - (pA * pB)

Since this API supports multiallelic sites, the user can also pass a normalisation function to control how the data is normalised across multiple alleles. The normalisation function is only run when computing over multiallelic sites. I've set the default to be $1/(n_A n_B)$, which is simply the arithmetic mean of the alleles in a given pair of sites. This will suffice in the majority of cases (the only outlier is $r^2$, for which there is already a python API). We also support computing statistics between sample sets.

The user would use the above summary function like this:

ts.two_locus_count_stat(ts.samples(), D, 1)

Where 1 specifies the length of the output array, we always require 1 dimension -- same as the ts.sample_count_stat function.

PR Checklist:

• Tests that fully cover new/changed functionality.
• Documentation including tutorial content if appropriate.
• Changelogs, if there are API changes.

[codecov]

Codecov Report

❌ Patch coverage is 72.38095% with 58 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.79%. Comparing base (1835ea3) to head (e5a105f).
⚠️ Report is 2 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3426      +/-   ##
==========================================
- Coverage   91.92%   91.79%   -0.13%     
==========================================
  Files          37       37              
  Lines       32153    32352     +199     
  Branches     5143     5144       +1     
==========================================
+ Hits        29556    29699     +143     
- Misses       2264     2320      +56     
  Partials      333      333              

Flag	Coverage Δ
C	82.71% <100.00%> (+<0.01%)	⬆️
c-python	77.29% <68.04%> (-0.06%)	⬇️
python-tests	96.40% <100.00%> (+<0.01%)	⬆️
python-tests-no-jit	33.20% <18.75%> (-0.02%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Components	Coverage Δ
Python API	98.70% <100.00%> (+<0.01%)	⬆️
Python C interface	90.65% <68.81%> (-0.59%)	⬇️
C library	88.87% <100.00%> (+0.01%)	⬆️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[lkirk]

I see that the CPython code coverage is split from the rest of the python tests. I can't test this directly without a small multiallelic test case in test_python_c.py.

[petrelharp]

Yep, it looks that way. Can you set one up? You could do that by either having a small tree sequence with high enough mutation rate there are lots of multiple hits (we do use msprime in that file) or explicitly constructing the tables.

[lkirk]

I was thinking of something like this:

[ins] In [1]: import msprime

[ins] In [2]: ts = msprime.sim_ancestry(2, recombination_rate=0.1, sequence_length=100, random_seed=23)
         ...: ts = msprime.sim_mutations(ts, rate=0.1, discrete_genome=True, random_seed=23)
         ...: print(f"allele counts: {set(len(s.mutations) for s in ts.sites())}")
allele counts: {1, 2, 3}

[ins] In [3]: ts.num_samples, ts.num_trees, ts.num_sites, ts.num_edges
Out[3]: (4, 47, 46, 137)

Which gives us a small multiallelic test case

[lkirk]

This function now serves as an inner wrapper. The the general stat accepts the summary function params so that the CPython code can pass them directly. All of the specialized stats functions call this function.

[apragsdale]

Thank you for opening this, @lkirk. I'm excited to see this implemented! Do we need to do any testing to demonstrate that there are no memory leaks or anything like that, which wouldn't be included in the test suite?

I know it could be found in the tests, but it may be helpful to spell out here how the API would work for a two- or more-way stat? Would it be:

ts.two_locus_count_stat([sample_list_1, sample_list_2], two_way_stat_func, 2)

[lkirk]

Yes, this needs leak checking and documentation, which I can add. I mostly wanted to make sure the user interface made sense first.

I know it could be found in the tests, but it may be helpful to spell out here how the API would work for a two- or more-way stat? Would it be:

ts.two_locus_count_stat([sample_list_1, sample_list_2], two_way_stat_func, 2)

The result_dim argument to the function gives the dimensions of the summary function output. The output is required to be 1D, so result_dim really tells us the length of the returned vector. In most cases, it'll be 1. See this line in the tests. So, if your summary function returned 1 value for a pair of sites, then the function call will look like this:

ts.two_locus_count_stat([sample_list_1, sample_list_2], two_way_stat_func, 1)

two_locus_count_stat was designed to work like sample_count_stat

[petrelharp]

Hi, @lkirk! This looks pretty straightforward. To have a careful opinion about the API, I think I need to see a reasonably careful docstring? For instance, what exactly are the arguments to f and norm_f? I can probably figure it out by tracing through the code, but it'll be less error-prone if you write it down. I'll have a look through for other issues, but it will be much easier to have the description in words of what exactly it's trying to do.

[petrelharp]

Does this runs the risk of not being run at all if there is no such example tree sequence?

[lkirk]

I don't think so, I'm selecting the index [0] out of this, which will fail for trying to select the 0th element out of an empty list. It might be better to use the pytest_params=False now that I look at get_example_tree_sequences again. I could avoid the .value[0]

[petrelharp]

insert some things like assert ts.num_sites > 0 and probably something asserting that this one is actually multiallelic here

[lkirk]

Good call, I will add this. These test tree sequences are shared by many tests and the docstring for this tree doesn't explicitly state that it has multiallelic sites, so it's good to assert in case it changes.

[petrelharp]

I think insert some asserts here that ts has the required properties for being a good test.

[lkirk]

Hi @petrelharp, thanks for taking a look. I didn't offer much of a description before, does this help?

Summary Function

In the sample_count_stat api, there are 3 required parameters:

1. sample_sets: List of lists of node ids.
2. f: A summary function that takes a one-dimensional array of length equal to the number of sample sets and returns a one-dimensional array.
3. output_dim: The length of the summary function's return value.

The two_locus_count_stat function (see function signature below) takes the same basic inputs except that f takes a matrix of haplotype counts and a vector of sample set sizes as input.

I've been writing summary functions with the signature f(X, n)

Where X is a matrix whose rows correspond to sample sets and columns correspond to haplotype counts ($w_{Ab}, w_{aB}, w_{AB}$). Here's a sample of what that looks like:

        sample_set 1, sample_set 2
w_Ab  [[9.            9.]
w_aB   [0.            0.]
w_AB   [0.            0.]]

Why lay the data out this way? Because numpy is row-major, iteration over the arrays gives us rows. That makes it easy to select the haplotype counts in one line:

AB, Ab, aB = X

Note: under the hood, the data is still laid out in an optimal way (at least as far as I can tell -- see this note)

In addition, the sample set sizes is shaped so that we can use it to normalize the counts in one go (n is a vector of sample set sizes)

pAB, pAb, paB = X / n

Finally, since we're no longer controlling the summary functions ourselves, the user has control over polarisation (see polarised parameter).

Normalisation

Perhaps the most clunky thing about this api is that the user will need a normalisation function for multiallelic data (norm_f parameter). Its function signature is the same as the code internal to ld_matrix(): f(X, n, nA, nB) -> [out_dims] where X and n are the same, nA and nB are the number of A and B alleles, respectively. Its output dimensions should be the same as our summary function (we also validate this just like we validate the output dims of our summary function).

The default normalization function of two_locus_count_stat is the total (or uniform) normalization, we average all results into one. Here's how I have defined it in the function signature

lambda X, n, nA, nB: np.expand_dims(1 / (nA * nB), axis=0),

I would have loved to be able to pass np.mean as a summary function instead, but we do need haplotype counts to norm $r^2$ and this keeps our original ld_matrix implementation untouched (we use the same exact machinery for this method, so we have to follow the C api).

[petrelharp]

I don't know why codecov says this line is not covered; do you? (sticking some prints or asserts in here can help track that down)

[petrelharp]

I've looked at the code pretty carefully, and things look good besides the comments. Nice work!