Conversation
tajimas_d, normalized_fay_wu_h, zeng_e, and zeng_dh computed their null variance from the nominal haplotype count in every windowed engine, while the scalar diversity functions use the harmonic mean of per-site valid counts. Under missing data this made windowed values diverge from the scalar reference; with complete data the two conventions coincide, which is why existing tests never caught it. All four windowed engines (_windowed_thetas_scatter, windowed_statistics, windowed_statistics_fused, windowed_statistics_fused_chunked) now compute a per-window effective sample size the same way diversity._compute_thetas does, and use it in place of n_hap. Windows with an effective sample size below 3 now report NaN, matching the scalar guard. The two fused engines share a common helper for this instead of duplicating the formula.
21aa184 to
2fe4099
Compare
andrewkern
left a comment
There was a problem hiding this comment.
I reviewed this locally on poppy against a checkout of the branch, running the windowed and parity suites on an A100 and measuring the fused path's memory directly. Line numbers refer to the branch as of the current head. Three things need to change before merging; the rest are worth doing but could be a follow-up.
Blocking
1. The parity suite fails on a GPU host. tests/test_implementation_parity.py lines 408, 411 and 414 still carry the strict xfail rules for the scatter engine's neutrality tests under missing data (reason _SCATTER_VARIANCE). They describe exactly the divergence this PR fixes, so they now XPASS:
pytest tests/test_implementation_parity.py -k missing_include
3 failed, 43 passed, 5 xfailed
FAILED test_path_matches_scalar[missing_include-tajimas_d-scatter] [XPASS(strict)]
FAILED test_path_matches_scalar[missing_include-normalized_fay_wu_h-scatter] [XPASS(strict)]
FAILED test_path_matches_scalar[missing_include-zeng_e-scatter] [XPASS(strict)]
CI has no GPU, so the green run does not cover this. Delete the three scatter rules and the _SCATTER_VARIANCE constant at line 378. Keep the fused rule at line 402 (see point 3).
2. _fused_tajimas_d builds temporaries proportional to sites x overlap depth, inside the engine that exists to bound memory. The site_win membership map at pg_gpu/windowed_analysis.py:1532 and the two cp.broadcast_to(...)[site_ok] gathers allocate O(n_per_var * n_var) on the GPU, plus a second allele_counts pass over the whole matrix. windowed_statistics_fused_chunked calls this at line 2712 on the un-chunked matrix, so the post-pass lands exactly where memory is tightest. Measured with 64 haplotypes x 2M variants:
| overlap depth (window / step) | extra pool peak |
|---|---|
| 1 | 0.00 GB |
| 10 | 0.16 GB |
| 100 | 8.78 GB (44 bytes per site-window slot) |
At 20M variants a request like statistics=['tajimas_d', 'garud_h12'], window_size=100_000, step_size=1_000 needs about 88 GB of transients and OOMs on an 80 GB card, where main completes. The fix at the right depth is to accumulate sum(1/n_valid) and the valid-site count inside _fused_windowed_kernel_v2 (two more per-window outputs, accumulated per chunk like out_seg) and derive n_harm_w from those; no membership map and no second matrix pass. A smaller fallback is prefix sums over inv_valid so that sum_inv_w = cs[win_stop] - cs[win_start], which is O(n_var).
3. The fused D numerator mixes two effective sample sizes, and the docs claim more than the PR delivers. At line 1562, d_num = mpd_sum - S / a1 now uses a1(n_harm_w) for the Watterson term. That matches neither the scalar's per-site sum nor the engine's own theta_w column, which still uses seg_count / a1(n_hap) at lines 2180 and 2696. On an 8-haplotype matrix with 30% missing calls, window 0 gives scalar D = 0.91181 (scatter and generic engines agree) but the fused engine gives 0.59828, and the (tajimas_d, fused) parity cell still xfails. Before this PR the fused result dict was at least internally consistent (pi - theta_w equalled the D numerator); now one row mixes two conventions, so anyone re-deriving D from the engine's pi and theta_w gets a different number. The PR description defers this as #304, but the changelog (docs/source/changelog.rst lines 180-185, "All four windowed engines now use the same per-window effective sample size") reads as a complete fix. Either fix the numerator here (the same kernel edit as point 2 can accumulate (n_present - 1) * a1_inv[nv], which closes #304 too) or narrow the changelog to the variance term and say that fused-engine D still differs.
Should fix before merge
4. docs/source/features.rst lines 629-636 still document the old behaviour ("the windowed neutrality tests use the full sample size in their variance formula, while the plain versions use an average of the per-site sample sizes"). That is now false for the scatter and generic engines and contradicts the changelog in the same PR. Rewrite or delete it; if the fused numerator stays as is, replace it with the narrower caveat.
5. Nothing tests the fused or chunked engines under missing data. The only un-xfailed test (tests/test_windowed_analysis.py:1226) calls windowed_statistics, the generic engine. The fused missing-data check is the strict xfail at test_implementation_parity.py:402, which only asserts that D does not match the scalar, and the chunked engine has no missing-data D test at all. My own check (fused and forced-chunked, overlapping grid, 34k variants, 10% missing) matched a hand-built reference in every window, so the code is right today, but an off-by-one in k_lo, a wrong member_win ordering, or a broken (S < 3) | (n_harm_w < 3) mask would pass both CI and the local suite. Please add a fused and chunked analogue of test_per_variant_missing_data_tajimas_d_matches_scalar, forcing chunking by monkeypatching _memutil.estimate_fused_chunk_size.
6. Coefficient evaluation now scales with the number of distinct effective n, and the cache thrashes. _achaz_coeffs_per_window (line 703) calls _achaz_variance_coefficients once per distinct n_harm_w value per test, and the 128-entry lru_cache evicts everything once 4 * distinct_n > 128. Measured: a cold call at n = 5000 costs 13.8 ms (a warm hit is 1 us); 200 distinct n over 10k windows costs 0.56 s per test, and a repeat call costs the same again because the cache no longer holds the entries. On a large panel with patchy missingness a 22-chromosome scan spends tens of seconds building coefficients where main spent milliseconds, on the scatter engine that is the default path. Hoist np.unique out of windowed_test, evaluate all requested weight pairs per n in one loop, and set maxsize=None (each entry is two floats).
Cleanups, fine as a follow-up
7. Membership map and allele_counts pass duplicated. _fused_tajimas_d (line 1496) rebuilds the site-to-window membership as a copy of the _PER_SITE_SCATTER_STATS block at lines 2330-2344 and re-runs allele_counts, which line 2355 already runs for daf_hist/mu_sfs, with a redundant hap.max() reduction the cap filter already guarantees. The docstring at line 2037 ("Reads the haplotype matrix once") no longer holds. The kernel fix in point 2 removes all of this; failing that, extract a _site_window_membership(...) helper and call allele_counts(hap, n_alleles=_FUSED_MAX_ALLELES).
8. Two formulas for one statistic. _fused_tajimas_d (line 1553) hand-rolls b1/b2/c1/c2/e1/e2 while the scatter closure (lines 862-872) and the generic engine (lines 3414-3421) go through _achaz_coeffs_per_window. Achaz alpha/beta for ('pi', 'watterson') equal c1/c2 to 1e-14 for n = 4..1000, so the fused block is a duplicate that only looks different, and it already differs in form (sqrt(max(var, 0)) with a post-hoc NaN versus the where guard). One _windowed_neutrality_test(w1, w2, n_harm_w, S, num) helper and one _windowed_harmonic_n(n_valid, scatter) for the recipe now written at lines 856-860, 1536-1548 and 3405-3412 would leave a single site to change.
9. Two extra scatter passes to compute a constant on complete data. With complete data or missing_data='exclude', n_harm_w is n_hap in every non-empty window, yet the scatter engine (line 856), the generic engine and _fused_tajimas_d each run two more scatter passes over (n_var * n_per_var) elements, each materializing a broadcast float64 temporary (about 8 GB transient and 1e9 atomic adds at depth 100 over 5M variants). bool((n_valid == n_hap).all()) is one reduction with n_valid already in hand; when true, set n_harm_w = np.full(n_windows, n_hap) and skip both.
10. Nondeterministic rounding at exact half-integer harmonic means (rare). Line 860 rounds a sum accumulated by atomic scatter-add. A window whose valid counts are {8, 24, 4, 6, 12} has harmonic mean exactly 7.5; float64 sums of the reciprocals in different orders round to 7 or 8 (3 of 14 exact-tie multisets flipped in a randomized check), so the coefficients for that window can change between identical calls. Needs an exact tie, so low severity, but it is new nondeterminism; a deterministic segmented reduction or a tolerance-aware round removes it.
11. _harmonic_mean_n (pg_gpu/diversity.py:373) returns a Python int for scalars and a float64 array otherwise, so all three windowed callers append .astype(np.int64) before using it as a table index and cache key. Return int64 unconditionally and write int(_harmonic_mean_n(...)) in _compute_thetas.
12. _harmonic_a1_a2_array (diversity.py:395): the clip upper bound is dead by construction, _harmonic_sums_sq is _harmonic_sums with the arange squared, and the n < 3 rule is enforced in three stacked layers (NaN for n < 2 here, NaN for n < 3 in _achaz_coeffs_per_window, and the n_harm_w >= 3 mask in every engine). A threshold change can be applied to some layers and not others without a test noticing; one guard in the shared helper would do.
13. Stale docstring. tests/test_windowed_analysis.py:1195 still says tajimas_d is excluded because its effective n is a separate issue with an xfail below; the xfail is gone and the next test asserts equality.
|
geez, these all seem to be real issues---thanks! |
Closes #134
tajimas_d,normalized_fay_wu_h,zeng_e, andzeng_dhcomputed their null variance from the nominal haplotype count in every windowed engine, while the scalar diversity functions use the harmonic mean of per-site valid counts.All four windowed engines (
_windowed_thetas_scatter,windowed_statistics,windowed_statistics_fused,windowed_statistics_fused_chunked) now compute a per-window effective sample size the same waydiversity._compute_thetasdoes, and use it in place ofn_hap. Windows with an effective sample size below 3 now report NaN, matching the scalar guard. The two fused engines share a common helper for this instead of duplicating the formula. A previously xfailed test pinning the bug now passes.This fix surfaced a separate inconsistency, #304, which is deferred