Skip to content

Two-way two-locus statistics (python proto) #3179

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

lkirk
Copy link
Contributor

@lkirk lkirk commented May 28, 2025

This PR implements two-way LD statistics, specified between sample sets. During the development of this functionality, a number of issues with the designation of state_dims/result_dims were discovered. These have been resolved, testing clean for existing code and providing the proper behavior for this new code.

The mechanism by which users will specify a multi-population (or two-way) statistic is by providing the index argument. This helps us avoid creating another ld_matrix method for the TreeSequence object.

In other words, for a one-way statistic, a user would specify:

ts.ld_matrix(stat="D2", sample_sets=[ss1, ss2])

Which would output a 3D ndarray containing one LD matrix per sample set.

ts.ld_matrix(stat="D2", sample_sets=[ss1, ss2], indexes=[(0, 1)])

Which would output a 2D ndarray containing one LD matrix for the index pair. This would use our D2_ij_summary_func, instead of the D2_summary_func. Finally, if a user provided

ts.ld_matrix(stat="D2", sample_sets=[ss1, ss2], indexes=[(0, 1), (1, 1)])

We would output a 3D ndarray containing one LD matrix per index pair provided.

Since these are two-way statistics, the indexes must be length 2. We plan on enabling users to implement k-way via a "general_stat" api. We did not implement anything more than two-way statistics here because of the combinatoric explosion of logic required for indexes > 2.

I added some basic tests to demonstrate that things were working properly. If we compute two-way statistics on identical sample sets, they should be equal to the one-way statistics. Unfortunately, this does not apply to unbiased statistics, which I've validated manually.

I've also cleaned up the docstrings a bit and fixed a bug with the D_prime statistic, which should not be weighted by haplotype frequency.

Fixes #2832

@lkirk lkirk requested a review from jeromekelleher May 28, 2025 20:54
@lkirk
Copy link
Contributor Author

lkirk commented May 28, 2025

Since @apragsdale isn't in the tskit organization, pinging him here to add as reviewer.

Copy link

codecov bot commented May 28, 2025

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 89.63%. Comparing base (1c83828) to head (6274d10).

Additional details and impacted files
@@           Coverage Diff           @@
##             main    #3179   +/-   ##
=======================================
  Coverage   89.63%   89.63%           
=======================================
  Files          28       28           
  Lines       32004    32004           
  Branches     5879     5879           
=======================================
  Hits        28688    28688           
  Misses       1886     1886           
  Partials     1430     1430           
Flag Coverage Δ
c-tests 86.66% <ø> (ø)
lwt-tests 80.38% <ø> (ø)
python-c-tests 88.18% <ø> (ø)
python-tests 98.86% <ø> (ø)

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

Files with missing lines Coverage Δ
c/tskit/trees.c 90.61% <ø> (ø)
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@lkirk lkirk force-pushed the two-locus-multipopulation-py branch from b77bc1b to 9f8222e Compare May 29, 2025 05:24
Copy link

@apragsdale apragsdale left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @lkirk.

If we compute two-way statistics on identical sample sets, they should be equal to the one-way statistics. Unfortunately, this does not apply to unbiased statistics, which I've validated manually.

This is ok (even expected) for the unbiased statistics, since the formulation assumes sampling without replacement, but having overlapping sample sets would essentially copy those repeated sample nodes. I don't think we should worry about it.

@lkirk lkirk force-pushed the two-locus-multipopulation-py branch 5 times, most recently from 834f0c2 to dc86340 Compare June 2, 2025 07:52
Copy link

@apragsdale apragsdale left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for making that change. This looks good to me then.

Copy link
Member

@jeromekelleher jeromekelleher left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've had a look through and the code all LGTM. The changes seem sensible, but I'd need to spend a lot longer on it to really digest the details and understand the finer points.

I'm going to ping @petrelharp here for a review as he's much better at this sort of stuff

@jeromekelleher jeromekelleher requested a review from petrelharp June 5, 2025 11:01
check_set_indexes(num_sample_sets, tuple_size * num_index_tuples, index_tuples)


def ld_matrix(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is the proposed API

@petrelharp
Copy link
Contributor

Summary: this PR is

  1. adding new summary functions
  2. proposing how dimension-dropping in the result works (and fixing up previous stuff around this)

@petrelharp
Copy link
Contributor

petrelharp commented Jun 11, 2025

Input:

"one-way versus two-way" is determined by whether or not there are indexes.

sample_sets:

  • None: defaults to all samples
  • a list of samples (not allowed for two-way stats)
  • a list of lists of samples

indexes:

  • None: determines if the stat is one-way
  • a length-two tuple
  • a list of length-two tuples

sites or positions:

  • None (defaults to all sites)
  • a list of them (defaults to both lists the same)
  • a list of two lists of them

Output: full is a 3D array

  • first two dimensions are sites or positions; these are never dropped
  • third is indexes

The only question is whether the third dimension is dropped.

  • one-way: if sample_sets is None or a single list
  • two-way: if indexes is a single tuple

@lkirk lkirk force-pushed the two-locus-multipopulation-py branch 2 times, most recently from 67e7f01 to 92418ba Compare June 13, 2025 22:05
This PR implements two-way LD statistics, specified between sample sets.
During the development of this functionality, a number of issues with
the designation of state_dims/result_dims were discovered. These have
been resolved, testing clean for existing code and providing the proper
behavior for this new code.

The mechanism by which users will specify a multi-population (or
two-way) statistic is by providing the `indexes` argument. This helps us
avoid creating another `ld_matrix` method for the TreeSequence object.

In other words, for a one-way statistic, a user would specify:
```python
ts.ld_matrix(stat="D2", sample_sets=[ss1, ss2])
```
Which would output a 3D array containing one LD matrix per sample set.

```python
ts.ld_matrix(stat="D2", sample_sets=[ss1, ss2], indexes=(0, 1))
```
Will output a 2D array containing one LD matrix for the index
pair. This would use our `D2_ij_summary_func`, instead of the
`D2_summary_func`. Finally,
```python
ts.ld_matrix(stat="D2", sample_sets=[ss1, ss2], indexes=[(0, 1), (1, 1)])
```
will output a 3D array containing one LD matrix _per_ index pair
provided.

Tests have been added to validate the result dimension for all of the
possible input combinations.

Since these are two-way statistics, the indexes must be length 2. We
plan on enabling users to implement k-way via a "general_stat" api. We did
not implement anything more than two-way statistics here because of the
combinatoric explosion of logic required for indexes > 2.

I added some basic tests to demonstrate that things were working
properly. If we compute two-way statistics on identical sample sets,
they should be equal to the one-way statistics. This test does not work
on unbiased statistics unless the sample sets being tested are equal in
index. I've added another test for two-way unbiased statistics.

I've also cleaned up the docstrings a bit and fixed a bug with the
D_prime statistic, which should not be weighted by haplotype frequency.
@lkirk lkirk force-pushed the two-locus-multipopulation-py branch from 92418ba to 6274d10 Compare June 13, 2025 22:08
Copy link
Member

@jeromekelleher jeromekelleher left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm happy to merge this as a step in the right direction - are you ok to merge @petrelharp ?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Multipopulation two-site statistics
4 participants