ld_statistics.pi2 returns an inconsistent quantity across index patterns #279

Encountered during #278 , which xfails a test on the behavior described here.

The counts-based public API ld_statistics.pi2(counts, populations) dispatches through _pi2_multi to a per-index-pattern formula. Those formulas do not follow a single convention: the two-locus iikk branch returns the symmetrized statistic, while every other branch returns a single un-symmetrized index term. So the same public function returns a different kind of quantity depending on the population-index pattern.

The inconsistency

In _pi2_multi, the i == j and k == l (iikk) branch averages the two locus orderings:

numer1 = (c11 + c12) * (c13 + c14) * (c21 + c23) * (c22 + c24)
numer2 = (c21 + c22) * (c23 + c24) * (c11 + c13) * (c12 + c14)
result[valid_mask] = 0.5 * (numer1 + numer2)[valid_mask] / denom[valid_mask]

Every other pi2 branch (iikl, ijkk, the 3-distinct shared-pool case, and the 4-distinct all-different case in _pi2_iikl / _pi2_ijkk / _pi2_shared_pop / _pi2_all_different) computes a single numer with no such averaging. The fused kernels (compute_all_pi2_hap) likewise return the raw single term for every pattern and apply symmetrization uniformly one layer up (generate_stat_specs). So the counts-API and the fused kernels agree for every pattern except iikk, where the counts-API self-symmetrizes and the fused kernel does not:

import numpy as np, cupy as cp
from pg_gpu import ld_statistics
from pg_gpu.haplotype_kernels import compute_all_pi2_hap
from pg_gpu.ld_pipeline import PopData

rng = np.random.default_rng(0)
N, P = 8, 4
counts = np.zeros((N, 4 * P))
for p in range(P):
    for r in range(N):
        counts[r, p*4:p*4+4] = rng.multinomial(16, [0.4, 0.2, 0.2, 0.2])
counts = cp.asarray(counts)
nv = cp.stack([counts[:, p*4:p*4+4].sum(1) for p in range(P)], axis=1)
pops = [PopData(counts[:, p*4:p*4+4], nv[:, p]) for p in range(P)]

for cfg in [(0, 0, 1, 2), (0, 1, 2, 3), (0, 0, 1, 1)]:
    api = cp.asnumpy(ld_statistics.pi2(counts, cfg, nv))
    fused = cp.asnumpy(compute_all_pi2_hap(pops, [cfg])[0])
    print(cfg, "match:", np.allclose(api, fused))
# (0, 0, 1, 2) match: True
# (0, 1, 2, 3) match: True
# (0, 0, 1, 1) match: False   <- iikk: counts-API self-symmetrizes, fused does not

Impact

No currently-computed statistic is wrong. The two-population counts pipeline (_compute_two_pop_statistics_batch) is hand-arranged around this: it assigns pi2_0_0_1_1 = pi2(counts, (0,0,1,1)) directly (relying on the internal self-symmetrization) but symmetrizes Dz externally (Dz_0_0_1 = 0.5 * (dz(0,0,1) + dz(0,1,0)), from single-term dz calls). Both the counts pipeline and the fused pipeline match moments.LD to ~1e-11.

But the public contract is incoherent, and it is a maintenance landmine: dz symmetrizes externally while pi2 iikk symmetrizes internally, so any future edit that treats the pi2 branches uniformly with dz (or vice versa) would double- or under-symmetrize a real statistic. A direct caller of ld_statistics.pi2 also gets values computed under two different conventions depending on the index pattern.

Suggested direction

Make ld_statistics.pi2 return one kind of quantity for all index patterns. The natural choice is the symmetrized statistic (matching the iikk branch, the moments definition, and what the two-population pipeline already treats pi2(counts, cfg) as): symmetrize the iikl / ijkk / shared / all-different branches the way iikk already does. The alternative is to make iikk return the single term like the others and move its symmetrization into the caller (as dz already does), which keeps dz and pi2 structurally parallel. Either removes the inconsistency; the first keeps the public return value equal to the named statistic.

A cross-check test pins this today: tests/test_ld_statistics_multipop_counts.py asserts the counts-API matches the fused raw term for every pattern, with the iikk case marked xfail(strict=True). Unifying the convention will trip that marker (either the iikk case starts passing, or the symmetrized-everything choice makes the other cases need a symmetrized oracle), which is the signal to update the test alongside the fix.