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.

[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?

[petrelharp]

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

[petrelharp]

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