Skip to content

Use pairwise-complete allele frequencies for naive-r2 under missing data - #312

Open
nspope wants to merge 6 commits into
mainfrom
nsp-220-pairwise-complete-r2
Open

nspope wants to merge 6 commits into
mainfrom
nsp-220-pairwise-complete-r2

Conversation

@nspope

@nspope nspope commented Sep 18, 2026 •

Copy link
Copy Markdown
Collaborator

Closes #220

_pairwise_ld_core computed D = p_AB - p_A*p_B where p_AB (the joint 1,1-frequency) was summed over the pairwise-complete sample set (gametes valid at both sites) but p_A/p_B were each site's own marginal frequency (summed over all of that site's valid gametes). Under structured missing data those are different sample sets, so D and r2 came out biased -- in one case, a mathematically impossible r2 > 1. missing_data='exclude' never exercised this, since dropping every site with any missingness makes the marginal and joint sets trivially coincide.

p_A/p_B are now conditioned on the pairwise-joint-valid set for each pair, matching the standard 2x2-table formula _tile_counts/_tile_sigma_d2 already use elsewhere in the LD module. pairwise_r2, pairwise_LD_v, and locate_unlinked (the three callers of the shared core) all pick up the fix; the tiled naive-r2 path used by zns/omega (_tile_r2_naive) is rebuilt on the same _tile_counts helper the projection estimator already uses, dropping its now-unnecessary marginal-frequency parameters.

sigma_d2 is untouched: it's a different, deliberately unbiased estimator -- it was just never affected by it, being built entirely from pairwise joint counts with no marginal-frequency vector anywhere in its own construction.

_pairwise_ld_core computed D = p_AB - p_A*p_B where p_AB (the joint
11-frequency) was summed over the pairwise-complete sample set (gametes
valid at both sites) but p_A/p_B were each site's own marginal frequency
(summed over all of that site's valid gametes). Under structured missing
data those are different sample sets, so D and r2 came out biased --
in one case, a mathematically impossible r2 > 1. missing_data='exclude'
never exercised this, since dropping every site with any missingness
makes the marginal and joint sets trivially coincide.

p_A/p_B are now conditioned on the pairwise-joint-valid set for each pair,
matching the standard 2x2-table formula _tile_counts/_tile_sigma_d2
already use elsewhere in the LD module. pairwise_r2, pairwise_LD_v, and
locate_unlinked (the three callers of the shared core) all pick up the
fix; the tiled naive-r2 path used by zns/omega (_tile_r2_naive) is rebuilt
on the same _tile_counts helper the projection estimator already uses,
dropping its now-unnecessary marginal-frequency parameters.

sigma_d2 is untouched: it's a different, deliberately unbiased estimator
(Ragsdale & Gravel 2019), not a fix for this bug applied under another
name -- it was just never affected by it, being built entirely from
pairwise joint counts with no marginal-frequency vector anywhere in its
own construction.
@andrewkern

Copy link
Copy Markdown
Member

this looks right for the haplotype path, but two things to fix here before we merge:

  1. we need the same fix on the genotype path
  2. mem usage looks like it goes way up. At 8,000 sites the memory pool goes from 2.66 GB to 4.2 GB, from the new sum_i, p_i and p_j m x m arrays. Since joint_n is symmetric, p_j is exactly p_i.T, so it can be a view, not a new array. And p_i and p_AB can be computed in place into sum_i and joint_11.

also looks like the changelog is now in conflict

…lete-r2

# Conflicts:
#	docs/source/changelog.rst
Review comment on PR #312: at 8,000 sites the memory pool went from
2.66 GB to 4.2 GB from the new sum_i, p_i, p_j m x m arrays. joint_n
is a self-product of the whole matrix against itself (not a tile), so
p_j is exactly p_i's transpose -- a view, not a fresh computation.
p_i and p_AB are now computed in place into sum_i and joint_11 via a
safe (zero-substituted) denominator, since cupy's divide has no
where= to mask the division directly (verified: TypeError, "Wrong
arguments"). The numerator is provably 0 wherever the denominator is,
so this is exact, not an approximation.

Measured at 8,000 sites (40 haplotypes, 5% missing): memory pool
high-water mark drops from 4.168 GB to 3.144 GB for pairwise_r2().
Values are unchanged -- existing LD test suite passes without
modification.
Review comment on PR #312, item 1: "we need the same fix on the
genotype path". _r2_matrix_diploid (the naive-r2 path for a
GenotypeMatrix, reached via zns/omega(estimator='r2') and the public
r2_matrix_diploid alias) computed each site's mean/variance from its
own marginal valid count, while the covariance cross term was already
implicitly pairwise-joint (zeroed wherever either site is invalid at a
gamete). Same mismatched sample sets as the haplotype bug, same bias,
same reason missing_data='exclude' never exercises it.

Mean/variance are now conditioned on the pairwise-joint-valid set,
using the standard pairwise-complete-observations correlation formula.
A site's own dosage variance (used to NaN a monomorphic site's whole
row/column, per _drop_undefined_sites's contract) stays a marginal,
per-site property -- not the pairwise-conditioned one -- so a site
can't be NaN for some pairs and not others. Applies the same memory
discipline as the haplotype-path fix from the start: sum_j/var_j are
transpose views (joint_n is a self-product of the whole matrix against
itself, not a tile), and the pairwise cross/variance terms are computed
in place into the base sum arrays via a safe denominator (cupy's divide
has no where=).

Verified against a naive pairwise-complete-obs reference implementation
under structured (per-site-varying) missing data -- confirmed the old
code actually disagreed with it (e.g. 0.02667 vs 0.02413 for one pair),
and the fix matches exactly, including the monomorphic-site whole-row/
column NaN and no-missing-data no-op cases.

Also resolves the changelog merge conflict against main (item 3 in the
same review comment): both sides had added independent bullets to the
Bug fixes section.
…_core

The safe_n array (joint_n with zeros substituted, to divide in place
without cupy's divide where= support) cost a full float64 (m, m)
array. Dividing directly by joint_n and masking the result afterward
needs only a 1-byte-per-element bool mask instead: IEEE-754 float
division never traps (verified directly against numpy: cupy raises no
warning for this, numpy raises two by default), it just produces
0/0 -> nan at the zero positions -- always nan here, never inf, since
the numerator is provably 0 wherever joint_n is -- and the explicit
zeroing overwrites every nan before anything else reads these arrays.

This relies on joint_n/sum_i/joint_11 staying float64: integer
divide-by-zero is undefined behavior on GPU, unlike float.

Measured at 8,000 sites: 3.144 GB -> 2.888 GB (pre-PR-312 baseline is
2.632 GB). Same test results as the safe_n version.
sum_i = hap_clean.T @ valid_mask

# In-place divide by joint_n (zeros included): numerator is provably 0
# wherever joint_n is, so this is always 0/0 -> nan, overwritten below.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is a bit of a cludge, as it relies on cupy replacing 0/0 with NaN and being silent about it. The alternative is to explicitly index on ~zero, but this creates intermediate copies which inflates memory usage (and the whole scheme here is designed to limit memory usage).

_r2_matrix_diploid returned 0.0 for a pair with zero jointly-valid
individuals, while pairwise_r2 (the haplotype analogue) returns nan
for the identical case. Root cause: the diploid path's undefined-pair
handling was modeled on _tile_r2_naive's 0.0 convention, which is for
a different context (zns/omega's scalar tiled aggregation) -- the
right reference is pairwise_r2, since both build a full (m, m) matrix
meant to be inspected directly or fed to _drop_undefined_sites, which
relies on undefined entries being nan.

0.0 isn't just inconsistent, it's not mathematically derivable here:
r2 = cov^2/(var_i*var_j) is 0/0 when the pairwise-restricted sample has
no variance at one site, and 0/0 doesn't resolve to 0 just because the
numerator is also forced to 0 by the same degeneracy -- there's no
limit to take, var_i is exactly 0, not approaching it. The zero
variance is a property of which individuals survived the joint-validity
filter, not evidence the sites are uncorrelated.

Also removes the separate marginal site_defined/n_valid/marginal_mean
check: a globally invariant site has zero variance in every pairwise-
restricted subsample too (a subset of constant values is still
constant), so denom > 0 alone already NaNs its whole row/column --
the marginal check was redundant with it, not an independent case.
Net: simpler code, three fewer (m, m) temporaries, and it now matches
pairwise_r2's rule exactly.
@nspope

nspope commented Sep 30, 2026

Copy link
Copy Markdown
Collaborator Author

take another look @andrewkern ? Memory usage is now not much higher for your example than the original implementation (but see caveat above), and the genotype path is implemented.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

missing_data='include' r2 mixes pairwise-complete counts with per-site frequencies

2 participants