Add tests for genotype imputation

[szhan]

Description

Add toy examples for testing imputation. Results obtained by running BEAGLE 4.1 are stored for comparison.

Fixes #2802

PR Checklist:

• Implement BEAGLE's interpolation-style imputation algorithm in Python.
• Implement a simpler version of the BEAGLE algorithm using functionalities in _tskit (tsimpute).
• Develop toy datasets for comparing BEAGLE and tskit-based imputation.
• Record BEAGLE results (including parameters; forward, backward, and state prob. matrices; allele probabilities).
• Implement functions allowing for convenient access to the toy datasets and BEAGLE results.
• Do test runs to profile performance.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 76.97%. Comparing base (8b1be4f) to head (24d7a49).

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #2815       +/-   ##
===========================================
- Coverage   89.79%   76.97%   -12.82%     
===========================================
  Files          30       30               
  Lines       30399    30399               
  Branches     5909     5643      -266     
===========================================
- Hits        27296    23400     -3896     
- Misses       1778     5614     +3836     
- Partials     1325     1385       +60     

Flag	Coverage Δ
c-tests	86.21% <ø> (ø)
lwt-tests	80.78% <ø> (ø)
python-c-tests	67.71% <ø> (ø)
python-tests	?

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

see 13 files with indirect coverage changes

---

Continue to review full report in Codecov by Sentry.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 8b1be4f...24d7a49. Read the comment docs.

[jeromekelleher]

Excellent!

And the $10M question - do we compute the same thing with tskit???

[szhan]

I think the tricky is getting the normalization factors right.

[jeromekelleher]

So we might not get exactly the same values, but I guess would expect the output matrices to be proportional (since normalisation factors can be arbitrary)

[jeromekelleher]

I've updated this to get the tskit code to run @szhan, and to remove the dependence on pandas for CSV parsing (which would have been a pain to update all the requirements.txt files). We can do the same thing (more or less) with numpy.

[szhan]

Here is the first stab at the comparison.

Rows = sites
Columns = ref. haplotypes

Forward probability matrix
tskit

[[0.23333333 0.03333333 0.23333333 0.23333333 0.03333333 0.23333333]
 [0.24637681 0.00724638 0.24637681 0.24637681 0.00724638 0.24637681]
 [0.2384058  0.02318841 0.2384058  0.2384058  0.02318841 0.2384058 ]
 [0.15942246 0.18115508 0.15942246 0.15942246 0.18115508 0.15942246]
 [0.05073599 0.39852802 0.05073599 0.05073599 0.39852802 0.05073599]]

beagle

[[9.9990e-01 1.0000e-04 9.9990e-01 9.9990e-01 1.0000e-04 9.9990e-01]
 [1.6665e-01 1.7000e-05 1.6665e-01 1.6665e-01 1.7000e-05 1.6665e-01]
 [1.7000e-05 1.6665e-01 1.7000e-05 1.7000e-05 1.6665e-01 1.7000e-05]
 [1.7000e-05 1.6665e-01 1.7000e-05 1.7000e-05 1.6665e-01 1.7000e-05]]

Backward probability matrix
tskit

[[0.97527473 1.34615385 0.97527473 0.97527473 1.34615385 0.97527473]
 [0.88461538 8.84615385 0.88461538 0.88461538 8.84615385 0.88461538]
 [0.58974359 9.43589744 0.58974359 0.58974359 9.43589744 0.58974359]
 [0.38017094 2.09094017 0.38017094 0.38017094 2.09094017 0.38017094]
 [1.         1.         1.         1.         1.         1.        ]]

beagle

[[0.166667 0.166667 0.166667 0.166667 0.166667 0.166667]
 [0.166667 0.166667 0.166667 0.166667 0.166667 0.166667]
 [0.166667 0.166667 0.166667 0.166667 0.166667 0.166667]
 [0.166667 0.166667 0.166667 0.166667 0.166667 0.166667]]

The forward values look pretty proportional to me. But I don't get why the values in the BEAGLE backward matrix is so uniform.

[szhan]

I suspect it has to do with my compiled copy of BEAGLE 4.1. BEAGLE's forward-backward algorithm only keeps track of the values for one site/marker in order to reduce memory requirements. For some reasons, I'm not seeing it being updated across the iterations.

[szhan]

Also, I'm seeing that BEAGLE is initialising the backward probability matrix to 1/n, where n = number of haplotypes in the reference panel.

See in LSHapBaum.java.

public HapAlleleProbs randomHapSample(int hap) {
        Arrays.fill(alleleProbs, 0f);
        int nMarkers = impData.nClusters();
        windowIndex = 0;
        arrayIndex = -1;
        setForwardValues(0, nMarkers, hap);
        Arrays.fill(bwdVal, 1.0f/n);
        setStateProbs(nMarkers-1, currentIndex());
        for (int m=nMarkers-2; m>=0; --m) {
            setBwdValue(m, hap);
            setStateProbs(m, previousIndex(hap));
        }
        setAlleleProbs(alleleProbs);
        return new LowMemHapAlleleProbs(refMarkers, impData.targetSamples(),
                hap, alleleProbs);
    }

See this line Arrays.fill(bwdVal, 1.0f/n);.

[jeromekelleher]

OK, I guess we'll need to figure this out. But, we should be able to check that the final imputed value (info scores?) are equal, easily enough?

[szhan]

There's also the matrix that combines the forward and backward values. I think BEAGLE gets the allele probabilities out of that.

[szhan]

I just discussed with Duncan. He spotted that the uniform backward values has to do with parametrization.

[jeromekelleher]

This looks really excellent @szhan! I'm not sure whether the Beagle code belongs here or in the tsimpute repo, but it's really good progress.

I guess ultimately what we want in tskit is some intermediate results that we compute using the Beagle algorithms that we can compare with what we compute in tskit. The actual job of imputation has a lot of moving parts, which we probably don't want to put together directly in tskit.

[szhan]

I think we should just focus on the interpolated allele probabilities at all the imputed markers on query haplotypes. We can compare these allele probabilities from BEAGLE and our implementation here.

[szhan]

Just to summarise the outputs from BEAGLE we can probably use for comparison here:

• Forward probability matrix at genotyped markers.
• Backward probability matrix at genotyped markers.
• HMM state probability matrix at genotyped markers.
• Interpolated allele probabilities at imputed markers.

It is hard to replicate exactly what BEAGLE is doing, so at best we will be able to look at correlations between the values (as discussed above).

[benjeffery]

Hi @szhan - what's the plan with this work?

[szhan]

Will circle back to this. Please keep.