Use per-window harmonic-mean n for windowed neutrality tests - #306
Conversation
21aa184 to
2fe4099
Compare
andrewkern
left a comment
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! |
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.
… them The scatter engine's per-window neutrality-test variance already uses the per-site harmonic-mean n this PR introduces, so the three strict xfails for tajimas_d/normalized_fay_wu_h/zeng_e on the scatter path now XPASS instead of xfailing.
_fused_tajimas_d rebuilt a site-to-window membership map independently of the kernel's own per-window variant ranges (searchsorted against bp bounds, then a broadcast+boolean-index gather), allocating O(n_per_var * n_var) transients -- 8.78 GB at 64 haplotypes x 2M variants at 100x overlap depth, enough to OOM on an 80 GB card at larger scale. The kernel already loops each window's own variant range for every other output, so this was pure duplication of work already done by the caller. _fused_windowed_kernel_v2 now accumulates sum(1/n_valid) and the valid-site count directly, per window, as two more outputs (n_valid >= 2, same gate the kernel already uses). _fused_tajimas_d no longer needs the haplotype matrix, positions, or window bounds at all -- just the kernel's own per-window sums. Separately, the kernel also now accumulates each site's own (n_present - 1) * a1_inv[n_valid] as a third new output -- the same per-site Watterson formula the scalar path already uses in diversity._ac_contribution, via a new a1_inv lookup table input (diversity._get_a1_inv, already cached and GPU-resident). Before this, Tajima's D's numerator used a1 at the per-window harmonic-mean n while the standalone theta_w output used a1 at the nominal n_hap -- two different effective sample sizes computing what should be the same quantity, and neither used the scalar's actual per-site-weighted formula. Both now use the kernel-accumulated value directly, so the fused engine's theta_w and tajimas_d actually match the scalar under missing data instead of only trying to.
…tats Only the per-variant engine had missing-data test coverage for pi/theta_w/segregating_sites/singletons/tajimas_d; the fused and chunked engines had multiallelic-only coverage with no missing data. Add both, forcing the chunked engine's small-chunk path via monkeypatching _memutil.estimate_fused_chunk_size, using the same window-subset-vs-scalar comparison the per-variant test already uses. Also drops a stale docstring/comment referencing a tajimas_d xfail that no longer exists (the test right below it already asserts equality).
_achaz_variance_coefficients was @lru_cache(maxsize=128), keyed on (w1_name, w2_name, n); each of up to 4 weight-pair combinations across the neutrality-test family populates its own share of that cache with one entry per distinct per-window harmonic-mean n. More than ~32 distinct n in one call started evicting entries from earlier pairs before the later ones were done, so a repeat call recomputed everything from cold instead of hitting the cache -- costly on a large panel with patchy missingness, where hundreds of distinct n aren't unusual. Each cache entry is two floats, so there's no real memory pressure argument for a bound: raise maxsize to unlimited. Separately, _achaz_coeffs_per_window recomputed np.unique(n_harm_w) on every call; a request for several neutrality tests at once (e.g. tajimas_d and zeng_e) shares the same n_harm_w across up to 3 calls in the scatter engine. Hoist that computation to the caller, once, and pass the (uniq_n, inv_idx) through.
_fused_tajimas_d hand-rolled Tajima's (1989) b1/b2/c1/c2/e1/e2 variance
coefficients inline instead of using _achaz_coeffs_per_window, which the
scatter and generic engines already use. Achaz's alpha/beta for the
('pi', 'watterson') weight pair are algebraically Tajima's original c1/c2,
so this was a second, differently-shaped implementation of the same
formula (and guarded differently: sqrt(max(var, 0)) with a post-hoc NaN
mask, instead of checking var > 0 before the sqrt). Now shares the one
helper and the one guard style with the other two engines.
The scatter and generic engines always ran two extra scatter-add passes to compute the per-window harmonic-mean effective sample size, even when every site in the matrix has full data -- in which case that mean is trivially n_hap everywhere. Short-circuit on a single reduction (n_valid == n_hap for every site) and skip straight to n_harm_w = n_hap. 'exclude' mode already guarantees this is always true (it drops any site with missingness first), so it gets the fast path automatically too.
bool(cp.all(...)) at its point of use forced the host to wait on a reduction kernel that hadn't even started yet. Moved the cp.all(...) call up to right after n_valid/n_hap are available, well before the rest of each engine's GPU work, so it runs concurrently instead of stalling the branch decision.
windowed_statistics_fused and its chunked variant each called .get() nine times in a row on the kernel's output arrays, one host sync per array. Stack them into a single array first and transfer that once, matching the pattern the scatter and generic engines already use for their own multi-output transfers.
e10eddc to
1793829
Compare
_harmonic_mean_n returned a plain Python int for scalar input but a float64 array otherwise, so every windowed caller appended .astype(np.int64) to normalize it. The one scalar caller (_compute_thetas) relied on the int branch implicitly, since its result feeds _achaz_variance_coefficients's lru_cache, which needs a hashable key -- a 0-d ndarray isn't. Now the function always returns an int64 array, and _compute_thetas wraps it in int(...) explicitly.
65cd53b to
06cf193
Compare
|
OK I think this is ready @andrewkern. Everything in the review should be addressed except for the non-deterministic rounding and unnecessary clamp nits (points 10, 12); which are bit more involved than they appear and are largely cosmetic (so I'm opting to skip). |
|
will look this over today |
|
okay this is GTG! I'm gonna merge this now |
2 participants
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