Multiallelic-correct statistics: integrate nsp-multiallelic-trunk into main - #154
Closed
andrewkern wants to merge 46 commits into
Closed
Multiallelic-correct statistics: integrate nsp-multiallelic-trunk into main#154andrewkern wants to merge 46 commits into
andrewkern wants to merge 46 commits into
Conversation
- add per-allele allele_counts primitive - scalar diversity uses per-allele primitive - scalar theta uses per-allele primitive - up test coverage for missing-data branches of segregating_sites/pi
- guard that the achaz test fixture stays biallelic - per-allele SFS + FrequencySpectrum, matching tskit - SFS: exclude both fixed classes (bins 0 and n)
- route sfs.sfs() through the per-allele allele_counts primitive - per-allele folded/scaled SFS (matches tskit polarised=False) - change _stream_sum to preserve the kernel's own dtype instead of forcing int64
- per-allele joint_sfs + project_joint_sfs - _joint_global_k / _joint_per_allele_counts helpers - rewrite joint_sfs_folded / joint_sfs_folded_scaled to match tskit's polarised=False 2-sample-set AFS - NB: this implies a behaviour change with biallelic only-data (weird behaviour on part of allel) - add _joint_folded_cells helper
Multiallelic handling for scalar statistics and SFS
…s primitive - add _aligned_pop_counts helper - add TestDivergenceVsTskit in test_multiallelic.py
Multiallelic handling in count-based divergence stats
- rewrite _pairwise_diffs_matrix_gpu - additionally fixes missing-data handling bug - add TestPairwiseDistanceKernel - add TestDistanceStatsConsumers
Correct pairwise distances for multiallelic data
- per-allele patterson_f2 / patterson_f3 (match tskit) - _aligned_allele_counts helper - add patterson_f4 and biallelic-restrict patterson_d - remove dead index-inflated helpers and bump consumer test coverage
- each window matches the per-allele scalar diversity functions on its subset, biallelic output is unchanged - in _fused_windowed_kernel_v2 replace the single `dac` accumulator with a per-allele cnt[MAX_ALLELES] - MAX_ALLELES cap drops sites with a greater allele index before launch - refresh stale single-pop docstrings/comments
- rewrite _twopop_site_components to per-allele counting - replace formulas to use per allele sums - rewrite _fused_windowed_twopop_kernel to per-allele counting - fst_wc left unchanged as there seems to be a mismatch with scalar path - equivalence testing between scalar path and windowed statistics path
consistent with scalar path (with some behaviour changes): - singletons were folded (dac==1 | dac==n-1) now are unfolded (dac==1) matching scalar - theta_w / tajimas_d numerator: per-site a1(n_valid) instead of a1(n_hap) - Tajima variance still uses n_hap and doesn't match scalar path - add TestMultiallelicSinglePop.test_per_variant_matches_scalar
…_statistics; remove dead code paths
Multiallelic admixture statistics
Fix windowed analysis for multiallelic data (with some exceptions)
Add tests/test_implementation_parity.py, which asserts that the same population-genetic statistic agrees across the code paths that compute it: scalar functions, FrequencySpectrum, and the windowed scatter, fused-CUDA, and python-loop engines. Each statistic is described once in a declarative _STATS registry; the parametrized test drives every supported path across four data conditions (clean, missing/include, missing/exclude, multiallelic) and asserts it matches the scalar reference. GPU-only, skipped without CUDA. 16 statistics, ~180 cases. Divergences between paths are recorded as strict xfails in _XFAILS with a mechanism-level reason; a (stat, path) may carry several condition-scoped rules. The remaining divergences are: - fst_wc (fused, all conditions): the fused kernel was not converted to the per-allele one-vs-rest Weir-Cockerham ANOVA the scalar/python-loop path uses - theta_w, tajimas_d (fused, missing/include): the fused kernel's per-site sample-size handling under missing data still diverges from the scalar - tajimas_d, normalized_fay_wu_h, zeng_e (fs and scatter, missing/include): the Achaz neutrality-test variance uses a modal / full-sample size rather than per-site valid counts - pi, theta_w, tajimas_d, segregating_sites, fay_wu_h, normalized_fay_wu_h, zeng_e (fs, multiallelic): FrequencySpectrum builds a per-allele SFS whose segregating-class count and pi differ from the scalar per-allele counts - da (scatter, missing/exclude): the within-population pi terms use a different site set than the scalar (dxy agrees) Refs #135
Pin statistics across implementations in a parity test suite
iHS and nSL assumed a biallelic focal SNP: the ancestral/derived split dropped (iHS) or merged (nSL) haplotypes carrying an allele index >= 2, and the allele-count/MAF filter counted only alleles 0 and 1, hiding a common high-index allele from min_maf. The shared-haplotype homozygosity core is already allele-agnostic (it decays on h1 != h2), so only the focal grouping changes: score each derived allele one-vs-ancestral. Both focal CUDA kernels gain a focal_allele parameter (class 0 = ancestral, class 1 = the focal derived allele; a homozygous pair of any other allele extends the shared run but scores neither), and the wrappers loop over derived alleles. nsl/ihs return the same (n_variants,) array as before on biallelic data and (n_variants, K-1) when the matrix has K > 2 alleles, with column j the score for derived allele j+1 (NaN where that allele or the ancestral allele is absent, or the allele is MAF-filtered). This reduces exactly to the classic biallelic result. The MAF filter now counts per allele via _memutil.allele_counts (replacing the 0/1-only _count_alleles_kernel): min_maf gates on each derived allele's own frequency min(freq_a, 1-freq_a), reducing to min(n0,n1)/(n0+n1) on biallelic sites. Windowed mean_nsl averages every finite (site, allele) score in a window. xpEHH, XP-nSL, EHH decay, and Garud H are allele-agnostic and untouched. There is no external oracle (scikit-allel ihs/nsl are biallelic-only; tskit has no iHS/nSL), so tests/test_selection_multiallelic.py validates via biallelic reduction against scikit-allel, relabel invariance (recoding derived 1 -> 2 shifts the score to column 1 and matches bit-for-bit), a brute-force pure-numpy nSL reference on triallelic data, the per-allele MAF fix, reference-absent sites, and first-ever mean_nsl coverage. Follow-up: the per-allele wrapper reruns the full SSL scan once per derived allele; a single-pass kernel depositing into all allele classes at once would remove that factor when multiallelic scans become a hot path. The biallelic path (one iteration) is unaffected. Closes #140
nSL and iHS scan the shared-haplotype length across every variant, so their scores depend on the full region. Subsetting or chunking the variants truncates the scan and changes the values. In particular, windowed mean_nsl over a StreamingHaplotypeMatrix runs per chunk and truncates the scan at chunk boundaries, so the streaming result approximates rather than reproduces the eager whole-region computation. Document this in the nsl/ihs docstrings, at the mean_nsl computation site, and in the _stream_windowed_analysis caveats (alongside the existing local_pca and sliding-window notes). No behavior change.
Multiallelic support for iHS and nSL selection statistics
- previously for multiallelic site the PCA input was the arbitrary integer allele label - rewrite so that biallelic sites go through the existing path, multiallelic sites go through a new one-hot encoded path - the two contributions are summed (covariance is additive) - the frequencies are sourced from allele_counts not the old dac_and_n - some finagling with host-device syncing - tests against an independent numpy host reference - validate pca/randomized_pca on multiallelic in test suite
Add _prepare_centered: the all-allele centered one-hot standardization (all alleles including the reference, centered by allele frequency, no variance scaling, missing imputed to p_a so it contributes 0). Its Gram X @ X.T is the site-mode genetic_relatedness matrix at haplotype granularity. Rewrite eager pca around it: drop the scaler argument, reject a GenotypeMatrix (pointing to pca_dosage), and normalize the Gram by the segregating-site count (proportion=True), then eigendecompose. This makes pca the tskit PCA of the GRM, pinned to eigh(ts.genetic_relatedness(mode='site', centre=True, polarised=False, proportion=True)). randomized_pca, local_pca, lostruct, and the old Patterson standardization helpers are untouched here and move to the new convention in a later change.
Migrate the windowed PCA family -- local_pca (dense and streaming engines), lostruct, and the jackknife paths -- to _prepare_centered, the all-allele centered (tskit) standardization, matching the eager pca. Drop the scaler menu from their signatures and from LocalPCAResult, and stop the windowed_analysis local_pca dispatch from forwarding a scaler. lostruct keeps its own per-window covariance /(ncol-1) (Li & Ralph); only the standardization changed. Add _sign_align_windows: each window's top-k eigenvectors are sign-aligned to the previous valid window (flip when the dot product is negative), so the returned per-window PC coordinates are continuous along the genome instead of flipping arbitrarily. lostruct's distance is sign-invariant and unaffected. Keep the per-window sync footprint at the pre-reframe level: _prepare_centered gains a need_segregating flag (windowed callers skip the segregating-site sync) and an n_alleles arg, and the streaming dispatcher computes the allele width once and passes it, avoiding a per-window int(hap.max()) sync. randomized_pca and the old Patterson standardization helpers are untouched here (a later change migrates randomized_pca and deletes them).
randomized_pca now uses the same all-allele centered standardization as pca: it rejects GenotypeMatrix (pointing to randomized_pca_dosage), drops the scaler argument, builds X via _prepare_centered, and scales coords by 1/sqrt(n_segregating) so the randomized approximation matches the proportion-normalized Gram that pca decomposes. The explained-variance ratio is scale-invariant and unchanged. Removes the memory-constrained _DeferredPCA branch. Deletes the now-unused Patterson-standardization helpers: _prepare_matrix, _multiallelic_onehot_dense, _biallelic_standardize_dense, _DeferredPCA, _DeferredPCATranspose, _materialize_prepared. The multiallelic per-allele one-hot Patterson standardization no longer exists; the diploid biallelic Patterson/GCTA path returns as pca_dosage on a GenotypeMatrix.
Adds the Patterson/GCTA PCA of biallelic diploid dosages as pca_dosage and randomized_pca_dosage on GenotypeMatrix, complementing the tskit genetic_relatedness PCA (pca / randomized_pca on HaplotypeMatrix). _prepare_dosage standardizes each variant's alt-allele dosage (0/1/2) by centering on its mean m and scaling by sqrt(p (1 - p)), p = m / 2, imputing missing to m so it contributes zero after centering; monomorphic sites are scale-guarded. pca_dosage eigendecomposes the individual Gram; randomized_pca_dosage uses a randomized SVD. Both reject the wrong matrix type (pointing to the tskit functions) and reject dosages > 2 (non-biallelic). This reproduces scikit-allel's diploid Patterson PCA: explained-variance ratios match to 1e-8 and coordinates up to a per-component sign. Population subsetting is done inline (the shared population helper is haplotype granularity on this branch); this should be unified in the type-domain-consistency follow-up.
Make the PCA standardization multiallelic-correct.
- add _relatedness_columns(hap, sample_sets, n_alleles) that is substrate for tskit genetic_relatedness - add TestRelatednessColumns - genetic_relatedness(haplotype_matrix, sample_sets, indexes, *, centre, polarised, span_normalize, missing_data) - add proportion option to genetic_relatedness - add streaming genetic_relatedness - add TestGeneticRelatednessDispatch
…Matrix Fix a pre-existing gap this surfaced: grm(GenotypeMatrix, population=...) failed because get_population_matrix only indexed .haplotypes.
…umers)
Eager genetic_relatedness normalizes by get_span('auto') (accessible-mask count
> n_total_sites > raw per-base span) and, because an accessible mask filters the
HaplotypeMatrix's variant view, computes over accessible sites only. The
streaming path did neither: accessible_bed was dropped by from_zarr's streaming
branches, StreamingHaplotypeMatrix carried no mask, and _stream_genetic_relatedness
hardcoded the raw chunk-grid span. So the same data loaded eagerly vs streamed
diverged on span_normalize=True whenever a mask was present.
Give the streaming matrix the same accessible-span contract as the eager one:
- _StreamingMatrixBase gains accessible_mask / n_total_sites and a get_span
mirroring HaplotypeMatrix.get_span, using the source's variant-position bounds
(mappable_lo/hi) so the span matches an eager matrix from the same store.
- from_zarr forwards accessible_bed to _build_streaming, which resolves the mask
once and attaches it.
- iter_gpu_chunks attaches the mask to every chunk it yields, at the single
chunk-production point, so every streaming consumer (genetic_relatedness now,
and windowed_analysis) sees accessible-filtered variants by construction --
rather than each reduction re-applying the mask itself.
- _stream_genetic_relatedness takes its span from get_span and drops its own
per-chunk attach.
Tests: eager/streaming genetic_relatedness parity with and without an
accessible_bed (filter + span), and that iter_gpu_chunks filters chunk variants.
The sibling StreamingGenotypeMatrix path still drops accessible_bed, so streaming
grm/ibs ignore a mask; tracked separately.
Closes #151
Streaming genetic_relatedness: honor the accessible mask (span + variant filtering)
Member
Author
|
we're 46 commits aheads of main here @nspope -- should we merge this in? or do you want to keep the campaign going on this branch? i opened a PR just to offer it - feel free to close |
Collaborator
|
there's still LD stats and a couple other things to get in ... though if you'd rather work off main, that's fine too |
Collaborator
|
I also wanted to do a general docs/cleanup pass, so actually let's wait until then. I think it's better to be sure the docs are lined up with the implementation. I can get back to this Wed, I think it's just 2-3 more PRs |
Member
Author
|
nah it's fine to keep working off this branch. i'll close this PR. I'm going to keep working on the accessibility mask / span stuff on this branch then. #152 is another similar class of bug. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Integrates the multiallelic-correctness effort (
nsp-multiallelic-trunk) intomain. Previously most statistics assumed a biallelic 0/1 encoding, so a site with an allele index >= 2 gave wrong results — folded/collapsed derived alleles, dosages that overflowed the [0, 2] range, PCs that depended on arbitrary allele numbering, and mis-scoped diploid estimators. This branch reframes each family around per-allele counting (matching tskit / scikit-allel where an oracle exists) and pins the behavior with new tests.mainis a strict ancestor of this branch (46 commits ahead, 0 behind), so it merges cleanly. 42 files, ~+5,000/−1,500.What changed, by family
allele_countsprimitive; scalar diversity, theta estimators, SFS, andFrequencySpectrumoperate per derived allele (matching tskitpolarised=False).patterson_dbiallelic-restricted.MultiallelicCapWarningfor the fixed-capacity fused kernels.genetic_relatednesspinned totskit.genetic_relatedness(per-set per-allele centered covariance, all groupings);grm/ibsscoped toGenotypeMatrix; fixedget_population_matrixfor genotypes.genetic_relatedness; accessible-mask span + variant filtering unified atiter_gpu_chunksso all streaming consumers honor a load-time mask.Testing
tests/test_implementation_parity.py) — asserts the same statistic agrees across the scalar / FrequencySpectrum / scatter / fused / python-loop paths over clean, missing/include, missing/exclude, and multiallelic data; remaining known divergences are recorded as strict xfails (the follow-up punch list).tests/test_multiallelic.py,test_decomposition_multiallelic.py,test_selection_multiallelic.py) — biallelic reduction against tskit / scikit-allel oracles, brute-force numpy references for genuinely multiallelic sites, index-independence, and missing-data handling.Known follow-up
GenotypeMatrixstreaming loader still dropsaccessible_bed, so streaminggrm/ibsignore an accessible mask (variant filtering, not span). Scoped, not blocking; theiter_gpu_chunkschoke-point here makes the fix straightforward.