diversity.pi gives wrong answer on multi-allelic sites (sum-of-int-values DAC)

[andrewkern]

Summary

diversity.pi (and the pi path through windowed_analysis) computes the
derived-allele count per variant as d = haplotypes.sum(axis=0). That works
when haplotype values are strictly {0, 1} (biallelic), but on a site whose
haplotype values include a 2 (triallelic under e.g. msprime's default
Jukes-Cantor mutation model), the int sum inflates d and the per-variant
pi falls back to a biallelic formula on incorrect counts.

Comparison to scikit-allel's allel.sequence_diversity shows ~1% per-window
discrepancy on data containing triallelic sites; on model="binary"
msprime data the two libraries agree to FP precision.

How to reproduce

import msprime
import allel
import cupy as cp
from pg_gpu import HaplotypeMatrix, diversity

# Default mutation model = jc69 -> can produce triallelic sites
ts = msprime.sim_ancestry(samples=20, sequence_length=100_000,
                          recombination_rate=1e-4, random_seed=42, ploidy=2)
ts = msprime.sim_mutations(ts, rate=1e-3, random_seed=42)
hm = HaplotypeMatrix.from_ts(ts)

# pg_gpu
pi_pg = diversity.pi(hm)

# scikit-allel
haps = cp.asnumpy(hm.haplotypes)
pos = cp.asnumpy(hm.positions)
ac = allel.HaplotypeArray(haps.T).count_alleles()
pi_a = allel.sequence_diversity(pos, ac)

print(f"pg_gpu: {pi_pg}, allel: {pi_a}, ratio: {pi_pg / pi_a}")
# observed: pg_gpu and allel disagree by ~1%

The same call with model="binary" produces matching values to FP precision.

What's actually wrong

pg_gpu/diversity.py:_site_contribution's 'pi' branch is:

return cp.where(seg, 2 * d * (n_safe - d) / (n_safe * (n_safe - 1)), 0.0)

This assumes d = count_of_derived_alleles and n_safe = total_haps. With
binary data, d = haps.sum(axis=0) is the derived-allele count. With
triallelic data, haps.sum(axis=0) mixes alt1 (+1) and alt2 (+2)
contributions and the resulting d no longer corresponds to a single
allele's count.

allel.mean_pairwise_difference(ac) handles the general case correctly
because ac carries one column per allele and the pairwise sum is over
all allele pairs.

Why it has been hiding

tests/test_scikit_allel_comparison.py constructs haplotype matrices by
hand with values restricted to {0, 1}, so count_alleles() returns
2-column ac and the formulas trivially agree. The discrepancy only
shows up against msprime data with the default JC mutation model, which
came up while building a test_streaming_vs_allel.py parity test for
streaming-from-zarr work.

Suggested fix

Replace the haps.sum(axis=0) DAC with a per-allele count from the
underlying (n_var, n_dip, 2) genotype block (or extend
HaplotypeMatrix to track per-allele counts directly). The pi formula
becomes a sum-over-allele-pairs the same way
allel.mean_pairwise_difference handles it. Same change applies to
other diversity / divergence statistics that derive a "derived-allele
count" from the int sum (theta_h, theta_l, fay_wu_h,
zeng_e, etc.).

[andrewkern]

let's come up with a plan to squash this next week @nspope

[nspope]

There's multiple mechanisms whereby multiallelic sites have an impact. My prior is to try to match allel where possible: this will largely be "folded" statistics like diversity, Taj D, etc. The unfolded statistics that pg_gpu supports are a little more complicated to handle, but we don't have a prerogative to match allel so can be a bit more flexible.

Did a trace with an LLM to get a better picture (I still need to go through an verify these):

mechanisms

A. Index-as-count inflation. Several paths derive allele counts with cp.sum(haplotypes, axis=0), which treats the integer allele index as a count: a third allele stored as 2 contributes 2, a fourth allele stored as 3 contributes 3. The clearest symptom is that admixture._allele_freq returns an alternate-allele "frequency" of up to 3.000 on multiallelic data — a frequency above 1 is impossible. This silently corrupts the Patterson f-statistics with no error raised. Note that the main diversity.pi path no longer does this: its fused CUDA kernel _memutil.dac_and_n counts v > 0 rather than summing, so the literal "sum-of-int-values" framing in the issue title now applies mainly to admixture, to the exclude-mode branches of allele_frequency_spectrum/segregating_sites/singleton_count, and to several windowed_analysis/ld_statistics paths.

B. Multiallelic→biallelic collapse. dac_and_n (and _pop_dac_and_n) count every non-reference allele as a single "derived" class. Two haplotypes that carry different derived alleles (e.g. allele 1 and allele 2) are therefore scored as identical. This undercounts every pairwise-difference statistic. scikit-allel instead keeps per-allele counts and computes mean pairwise difference as 1 − Σ_a (ac_a choose 2) / (n choose 2), summing over distinct allele pairs. This is the mechanism that bites diversity.pi today and is the one the issue is really about.

C. Reference-absent sites. A site where allele 0 is absent but two or more derived alleles segregate has dac == n_valid, so the segregating test dac < n_valid is false and the site is dropped entirely from segregating_sites and Watterson's θ. We verified that the set of segregating sites pg_gpu misses is exactly the set of reference-absent sites.

by statistic

pg_gpu statistic	scikit-allel counterpart	current multiallelic behaviour	mechanism	recommended target
pi, windowed pi	sequence_diversity / mean_pairwise_difference	collapses derived alleles → undercounts (~1–1.6%)	B	per-allele mean pairwise difference (match allel)
theta_w	watterson_theta	drops reference-absent sites → slight undercount	C	count a site as segregating iff ≥2 distinct alleles (match allel)
segregating_sites	AlleleCountsArray.count_segregating	drops reference-absent sites	C	≥2 distinct alleles present
singleton_count	per-allele ac == 1 (via sfs)	counts total-derived==1, misses per-allele singletons at multiallelic sites	B	count singletons per derived allele (match allel)
tajimas_d	tajima_d	inherits π and θ_W bias	B+C	fixed once π and θ_W are fixed
allele_frequency_spectrum	sfs / sfs_folded	uses total-derived convention; exclude mode also has index-inflation	A (exclude mode)	decide convention (see below); fix exclude-mode counting
theta_h, theta_l, fay_wus_h, zeng_e, normalized_fay_wus_h	none (unfolded)	collapse perturbs each	B	design decision (see below)
dxy, da, dxy_min, gmin, distance-based	mean_pairwise_difference_between / sequence_divergence	collapse → undercounts (~1.6%)	B	per-allele between-population pairwise difference
fst (Hudson)	hudson_fst / average_hudson_fst	near-immune (+0.15%) but still collapse-based	B	per-allele numerator/denominator
fst_weir_cockerham	weir_cockerham_fst	collapse-based (not separately benchmarked here)	B	per-allele / genotype-based WC
fst_nei	(Nei’s Gst)	collapse-based	B	per-allele
patterson_f2 / average_patterson_f3 / patterson_d	patterson_f2 / average_patterson_f3 / patterson_d	index-inflated allele frequency (freq up to 3.0)	A	biallelic-restricted or per-allele decomposition (match allel)
grm, ibs	none with multiallelic support	dosage of allele index is ill-defined	—	require/auto biallelic input
LD r² family	rogers_huff_r (biallelic only)	dosage ill-defined	—	require/auto biallelic input

[andrewkern]

i believe scikit-allel just lumps 3rd or 4th alleles into the alt state. I wonder if we can do what tskit.statistics does instead?

[nspope]

I believe tskit is using the mutation parent hierarchy provided by the genealogies ... which is something we'd have to backfill if trying to take the same approach.

Oh, actually---the parent hierarchy is just for getting the right translation to allelic codes. So this could work. Let's see... I think tskit's strategy is better here (counting alleles as distinct mutational events). This distinction matters for segsites and the theta estimators/Tajima's D. For the pairwise statistics (pi/dxy), both tskit and allel do the same thing. For the admixture stats, it looks like allel forces biallelic, so we'd be better off matching tskit.

[nspope]

After looking this over---I agree that tskit is the right target; allel is all over the place in how it handles multiallelic sites. For the polarised SFS the tskit calculation is essentially flattening the dac array across alleles (getting allele-specific counts per site), then tallying counts into a histogram. There's some additional weighting needed for the folded case.

I think this could be fixed pretty simply by having dac_and_n emit a num_variants by num_alleles array instead of collapsing (and do something analogous for population-stratified counts). Then the individual statistics can ingest this array instead. Let's talk about it if you're around this afternoon.

[nspope]

We talked this over, and decided that tskit is the better target. I'll draft a PR with a new dac_and_n kernel emitting multiple allele counts per position (this also likely needs to emit the number of alleles per site) and switch over pi/dxy to use this, so we can see how it falls out before swapping everything over.