Per-site decomposition for Weir-Cockerham FST - #309
Conversation
| 'theta_h', 'theta_l', 'fay_wu_h', 'singletons', | ||
| 'normalized_fay_wu_h', 'zeng_e', 'zeng_dh', 'max_daf'} | ||
| scatter_twopop = {'fst', 'fst_hudson', 'dxy', 'da'} | ||
| scatter_twopop = {'fst', 'fst_hudson', 'fst_wc', 'dxy', 'da'} |
There was a problem hiding this comment.
so in looking at this closely, i think he scatter check runs before the fused check, so every fst_wc request under 'include' now goes through scatter, not the fused path. The fused path chunks when the matrix is big, and this won't. obviously preexisting, but maybe we should file an issue here? or fix it on this PR?
There was a problem hiding this comment.
Let's defer this as it's pre-existing, i'll file a followup shortly.
| results['da'] = ((between_sum - (pi1_sum + pi2_sum) / 2.0) / spans_gpu).get() | ||
|
|
||
| if 'fst_wc' in stats_set: | ||
| results['fst_wc'] = cp.where(wc_abc_sum > 0, wc_a_sum / wc_abc_sum, |
There was a problem hiding this comment.
Under 'exclude' this used to go through the scalar loop, which gives 0.0 when a+b+c is zero. Now it gives NaN. That matches the fused kernel and Hudson, so I think NaN is right, but it is a behavior change. prob worth adding a line to the changelog.
The scalar still returns 0.0 for the same case. So windowed and scalar now disagree on empty or monomorphic windows. Could be a follow-up, but worth an issue.
| stats_set = set(statistics) | ||
| pop1_name, pop2_name = populations[0], populations[1] | ||
|
|
||
| if 'fst_wc' in stats_set: |
There was a problem hiding this comment.
This block is the same as the one in windowed_statistics_fused, and the chunked engine has a third. One small helper in _warnings.py that takes the matrix and the two pops would replace all three.
| @@ -1002,7 +1017,11 @@ def scatter_sum(values): | |||
| # Compute per-site components (single pass over the data) | |||
| mpd1, mpd2, between = _twopop_site_components(hap1, hap2) | |||
There was a problem hiding this comment.
This runs even when the request is only fst_wc. we should limit it to the stats that use it. And _wc_site_components recomputes k with two more .max() calls. Pass k in.
| rtol=1e-9, atol=1e-12) | ||
| wc_ref = divergence.fst_weir_cockerham(sub, 'p1', 'p2') | ||
| if np.isnan(wc_w) or np.isnan(wc_ref): | ||
| assert np.isnan(wc_w) and np.isnan(wc_ref) |
There was a problem hiding this comment.
The "both NaN" branch can never pass. fst_weir_cockerham never returns NaN. So if a window ever comes out NaN this fails with a confusing message. Either make the scalar return NaN too or drop the branch. Same at line 693.
| err_msg=f"Mismatch in {k}") | ||
|
|
||
| def test_two_pop_scatter_matches_fused_for_fst_wc(self, matrix_with_pops): | ||
| """The scatter engine (now handling fst_wc) must agree exactly with |
There was a problem hiding this comment.
"now handling fst_wc" is history. Say what the test checks, not what changed. Same for "now that fst_wc is scatter-eligible" at line 663.
| s_squared[vt] = (n1t * (p1[vt] - p_bar[vt])**2 | ||
| + n2t * (p2[vt] - p_bar[vt])**2) / ((r - 1) * (nt / r)[:, None]) | ||
| h_bar[vt] = (het1[vt] + het2[vt]) / nt[:, None] # per-allele obs het | ||
| h_bar[vt] = (het1[vt] + het2[vt]) / nt[:, None] |
There was a problem hiding this comment.
this dropped the comments that link the code to the paper. better to keep them in i reckon
|
small nit picks here. |
|
So I think there's some overlap here with #311 --- it fixes the NaN issues you raise. I'll deal with the other issues now, I think. |
fst_weir_cockerham already built per-site-per-allele variance components (a, b, c) before collapsing them to two floats inline, but nothing else could reuse that decomposition. Extracted _wc_site_components (mirroring diversity._ac_contribution's shared-decomposition pattern) returning the per-site (a, a+b+c) sums; the scalar function is now sum(a_site) / sum(abc_site). The scatter engine (_windowed_twopop_scatter) gains an fst_wc branch built on the same function, plus the paired-rows warning check the fused kernel already had (its other statistics never needed one). fst_wc joins scatter_twopop, so windowed fst_wc now goes through the fast scatter path under both missing_data modes instead of only the fused kernel under 'include' and a slow per-window scalar loop under 'exclude'.
- Deduplicate the scatter/fused check_paired_rows blocks into a shared check_paired_rows_for_populations helper in _warnings.py. - Skip _twopop_site_components in the scatter engine when only fst_wc is requested; it's only needed for the gamete statistics. - Pass k into _wc_site_components instead of recomputing it internally, so a caller that already knows it doesn't pay for a repeated host sync. - Restore the Weir & Cockerham 1984 Eqs 2/3/4 comment dropped during the _wc_site_components extraction. - Fix two dead "both NaN" test assertions that could only fail if triggered (fst_weir_cockerham never returns NaN today); they now check the known 0.0-vs-NaN sentinel mismatch directly. - Trim history-narration from two test docstrings.
a732fa4 to
e02ab96
Compare
|
OK, I think is good, with deferrals mentioned above. Maybe take a look at #311 next given the overlap? |
The scatter engine now serves fst_wc, so the parity suite checks it against the scalar under the missing-data and multiallelic conditions like the other two-population statistics.
|
looks good. i'm pushing one small change to close out all of 187, then will merge and move to 311. |
2 participants
Fixes #187
fst_weir_cockerham already built per-site-per-allele variance components (a, b, c) before collapsing them to two floats inline, but nothing else could reuse that decomposition. Extracted _wc_site_components (mirroring diversity._ac_contribution's shared-decomposition pattern) returning the per-site (a, a+b+c) sums; the scalar function is now sum(a_site) / sum(abc_site).
The scatter engine (_windowed_twopop_scatter) gains an fst_wc branch built on the same function, plus the paired-rows warning check the fused kernel already had (its other statistics never needed one). fst_wc joins scatter_twopop, so windowed fst_wc now goes through the fast scatter path under both missing_data modes instead of only the fused kernel under 'include' and a slow per-window scalar loop under 'exclude'.