Make the PCA standardization multiallelic-correct.#143
Conversation
|
Looking at this again, I'm not sure matching scikit-allel is the right choice---that is, summing over all non-reference alleles. tskit doesn't have site-PCA implemented yet, but for the GRM it sums over all alleles including the reference. That'd be the more internally consistent option here (... this is unlikely to matter in practice, but might as well make the semantics all internally consistent) |
|
i think we should generally aim to match tskit behavior over scikit-allel here |
|
Sounds good. So I think this also means returning an unstandardized (but optionally centered) GRM for the relatedness issue ... unlike GCTA, which uses a 2p(1-p) standardization. If/when we decide on a multiallelic standardization we can add a |
|
sounds good to me |
- 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
bff75a1 to
50d4f80
Compare
|
so downside here: if we make everything consistent, then lostruct will differ from the R implementation (which "double-centers", as it centers each sample as well as each site). That's an operation that only really makes sense on the dosage matrix. The general (multiallelic) construction (which one-hot encode alleles) doesn't really have an analogue. I don't think this impacts the results much---it changes the eigenvalues by a scalar, eigenvectors are essentially the same, on the lostruct "example" we ship---but it does make the two forms numerically distinct. EDIT: that's not the full picture---it does look very similar on the genealogical test bed, but for the highly artificial test fixture with "frequency block structure" they look pretty different. Not sure what the best choice is here. |
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.
|
ok this is ready for another look @andrewkern --- I updated the PR description in the OP |
andrewkern
left a comment
andrewkern
left a comment
There was a problem hiding this comment.
this looks good to go Nate-- we'll want to add a bit in the docs explaining the choice to diverge from the lostruct R implementation, but we could do that on a separate PR if you like
|
I think we'll also need to benchmark the two implementations on some real data, unfortunately---these changes have had a bigger blast radius than anticipated. |
2 participants
EDIT: updated with the new tskit-analogous verison
decomposition's PCA standardization was wrong on multiallelic input: it derived per-site centering and scaling from a collapsed derived-allele count and centered the raw allele-index matrix, so a site with an allele index >= 2 fed the arbitrary integer allele label into the decomposition. Principal components, explained variance, and every local PCA window covariance became a function of arbitrary allele numbering, silently and non-reproducibly under relabeling.
Rather than patch the single standardization, this splits the family along the data type, because there are two established and genuinely different PCAs and they belong to different inputs:
Each entry point rejects the other matrix type with a TypeError that names its counterpart. local_pca, lostruct, and the local-PCA jackknife move to the same all-allele centered standardization as pca (they remain HaplotypeMatrix-only). The scaler menu (patterson / standard / None) is dropped in favor of the two fixed conventions.
What changed
pca and randomized_pca now build the all-allele centered one-hot (all present alleles, centered by frequency, no scaling, missing imputed to the frequency so it contributes zero) and normalize the Gram by the segregating-site count. That Gram is the genetic relatedness matrix, so pca is pinned against a tree sequence's genetic_relatedness output (all sample pairs, centered, proportion-normalized).
pca_dosage and randomized_pca_dosage are new. They standardize the biallelic diploid dosage by centering and binomial-variance scaling, then eigendecompose the individual Gram with no segregating-site normalization, so the coordinates match scikit-allel's magnitude as well as its subspace. They reject dosages > 2 (which cannot arise from a biallelic diploid site) and reject a HaplotypeMatrix. Population subsetting is done inline over individuals; the shared population helper is haplotype granularity on this branch, and unifying genotype/haplotype subsetting is a logged follow-up.
local_pca, lostruct, and local_pca_jackknife route each window through the all-allele centered standardization instead of the old per-variant standardization; they keep their own per-window covariance X @ X.T / (ncol - 1). Adjacent windows' per-PC eigenvectors are sign-aligned (flip when the dot with the previous valid window is negative) so windowed PC tracks are continuous.
The obsolete standardization machinery is deleted: the split biallelic/multiallelic Patterson-scaled helpers and the memory-constrained deferred-PCA path. The four public functions are exported from the package.
Behaviour
pca and randomized_pca change on all input, not just multiallelic: the standardization is now all-allele centered with no variance scaling and a segregating-site denominator, where before it was a per-variant Patterson standardization on haplotypes. Explained-variance ratios and the top subspace change accordingly; the output is the eigendecomposition of the genetic relatedness matrix.
The diploid dosage PCA (the previous default's biallelic behaviour) is still available, now as pca_dosage on a GenotypeMatrix.
local_pca and lostruct change on all input for the same reason, and no longer bit-match the published R local-PCA on biallelic data: the all-allele one-hot keeps a column per present allele where R keeps one dosage column per site, and it is single-centered (variant only) where R's classical-MDS double-centers (variant and sample). The two agree for genealogical structure but diverge for structure that aligns with average distance-from-reference (a per-group frequency shift); the reasoning is in the centering rationale note. The strict R-parity tests are xfailed with that reason.
Verification
The all-allele Gram equals a tree sequence's genetic_relatedness matrix, and pca's explained variance equals the eigenvalues of that matrix. randomized_pca matches full pca (explained variance plus top subspace). pca_dosage reproduces scikit-allel's diploid PCA explained-variance ratios to 1e-8 and coordinates up to a per-component sign; randomized_pca_dosage matches pca_dosage on structured data via canonical correlations. Type guards and the biallelic requirement are covered. Per-window local_pca / lostruct / jackknife match a direct per-window computation, adjacent windows are sign-aligned, and the R-parity suite is xfailed where the conventions genuinely differ.