Support haploid / hemizygous samples (sex chromosomes); stop silently mishandling them

[andrewkern]

Summary

pg_gpu has no notion of haploid / hemizygous samples. When a VCF (or zarr store) contains haploid or pseudo-diploid calls -- the typical situation for the X/Y in males, the whole genome of haplodiploid or haploid organisms, or organellar contigs -- the loaders silently produce a haplotype matrix with the wrong number of haplotypes and biased statistics, with no warning. This was reported by @jiseonmin in #120 (male-only flies on the X: "the output looks weird since it assumes everyone is homozygous everywhere").

This issue proposes (1) making the failure loud across every loader, and (2) adding first-class opt-in haploid support.

The problem

A user loads an X-chromosome VCF containing only males. Two encodings are common in the wild, and pg_gpu mishandles both:

• Pseudo-diploid homozygous (GT=1/1, 0/0) -- what GATK and many pipelines emit for haploid regions. Every male is read as a homozygous diploid, so each contributes two identical haplotypes. The reported n_haplotypes is doubled, within-individual LD is perfect, heterozygosity is zero everywhere, and the SFS/pi/Tajima's D are computed over 2N pairwise-identical pseudo-haplotypes -- biased and wrong, but no error.
• True haploid (GT=1, GT=0) -- a single allele per call.

The reporter ran utils/genome_scan.py, which routes VCFs straight through HaplotypeMatrix.from_vcf (utils/genome_scan.py:121).

Current state of the codebase

HaplotypeMatrix.from_vcf (pg_gpu/haplotype_matrix.py:480-487) is the path the reporter hit:

genotypes = allel.GenotypeArray(vcf['calldata/GT'])
num_variants, num_samples, ploidy = genotypes.shape
assert ploidy == 2

haplotypes = np.empty((num_variants, 2 * num_samples), dtype=genotypes.dtype)
haplotypes[:, :num_samples] = genotypes[:, :, 0]
haplotypes[:, num_samples:] = genotypes[:, :, 1]

The assert ploidy == 2 is effectively dead code, because allel.read_vcf defaults to a fixed ploidy-2 layout regardless of the VCF's actual ploidy. I verified both cases against scikit-allel:

VCF encoding	calldata/GT returned	ploidy axis	assert
true haploid GT=1	[[1, -1], ...] (2nd allele padded with -1)	2	passes
pseudo-diploid GT=1/1	[[1, 1], ...]	2	passes

So:

• For true haploid input the assert never fires; instead haplotypes[:, num_samples:] = genotypes[:, :, 1] fills the entire second haplotype block with -1 (missing). The matrix is double-width, half of it spurious missing data, and num_haplotypes reports 2x the true sample count.
• For pseudo-diploid input the assert never fires either; each individual becomes two identical haplotypes. This is the reporter's "everyone is homozygous everywhere."

In short, the one guard that exists provides false reassurance: it cannot catch either real-world haploid encoding, and on the rare path where it would fire it raises a bare AssertionError with no message.

Other entry points have the same gap:

• GenotypeMatrix.from_vcf (pg_gpu/genotype_matrix.py:387-406) has no ploidy guard at all. It computes dosage with np.sum(gt, axis=2), so a pseudo-diploid 1/1 male becomes dosage 2 (homozygous alt) everywhere -- same over-homozygosity bias -- while a true haploid (padded with -1) is flagged missing in every cell.
• HaplotypeMatrix.from_zarr / _build_eager (pg_gpu/haplotype_matrix.py:668-675) does raise a clear ValueError when the store's ploidy axis is not 2 -- but a pseudo-diploid VCZ store has a ploidy-2 axis, so it slips through with the same silent bias.
• HaplotypeMatrix.from_ts (pg_gpu/haplotype_matrix.py:899-900) is haploid-native (tskit genotypes are one row per sample node), so it is already correct, but there is no ploidy concept to make that explicit or to mix with diploid data.

There is no haploid/hemizygous handling, no ploidy= parameter, and no test covering haploid, pseudo-diploid, or mixed-ploidy inputs anywhere in pg_gpu/ or tests/.

Proposed solution

Two layers -- the first is a small, safe correctness fix; the second is the actual feature.

1. Make the failure loud (minimum viable). Replace the dead assert with real detection that inspects the data, not scikit-allel's padded shape:

• Detect true-haploid padding: a sample/site whose second allele column is -1 while the first is called. If an entire region is haploid-padded, either load it correctly as single haplotypes (preferred) or raise a clear, actionable error naming the affected samples.
• Detect suspected hemizygosity: if some samples are homozygous (or padded) at every called site across a region, emit a warning that the data look haploid and that statistics will be biased unless the samples are declared hemizygous. (Homozygous-everywhere cannot be proven to be haploid vs. a truly inbred diploid, so this stays a warning, not a hard error.)
• Apply the same guard/warning consistently to GenotypeMatrix.from_vcf and the zarr _build_eager paths so no loader is silently wrong.

2. First-class haploid support (the enhancement). Add an opt-in way to declare ploidy so hemizygous samples each contribute one haplotype instead of two:

• A ploidy= / haploid_samples= parameter on from_vcf / from_zarr (e.g. ploidy=1 for an all-haploid contig, or haploid_samples=[...] to mark hemizygous individuals on a sex chromosome).
• For declared-haploid samples, collapse the pseudo-diploid pair to a single haplotype (and for true-haploid input, drop the padded -1 column) so n_haplotypes, the per-site valid counts, and the SFS are correct.
• Support mixed ploidy within one matrix (e.g. autosome + X, or males hemizygous + females diploid on the X) by per-sample collapse, since downstream kernels already track per-site valid sample counts via the 'include' missing-data mode.
• Document from_ts as already haploid-native and wire any new ploidy metadata through consistently.

3. Tests. Add coverage for true-haploid VCFs, pseudo-diploid homozygous VCFs, and mixed-ploidy inputs across HaplotypeMatrix/GenotypeMatrix from_vcf/from_zarr, asserting correct n_haplotypes and unbiased pi/SFS versus a hand-constructed expectation.

Recommendation

Ship layer 1 first as a focused correctness fix -- it stops silent wrong answers immediately and is low-risk -- then build layer 2 as the real feature. The reporter's exact case (genome_scan.py on a male X) should go from "silently wrong" to either "correct" (with ploidy/haploid_samples declared) or "loudly warned" (with defaults).

Reported in #120 (thanks @jiseonmin). cc @nspope @andrewkern

[nspope]

Totally agree on getting it to error out, asap, then working on the enhancement.

For the enhancement, I'm not sure about splitting the API into two arguments. How about ploidy=... where this can be either an integer (1 or 2, applies globally) or a dict that maps sample names to 1 or 2? (One could also pass a defaultdict with a default ploidy, if there's only a few haploid samples).

[andrewkern]

that shape sounds smart. one ploidy= arg, two potential input shapes, int or dict