test: cover chrX haploid-male counting and MNV overlap branches

divref/divref/haplotype.py

@@ -166,7 +166,9 @@ def haplo_coordinates(
     Returns:
         Hail struct expression with int32 fields `start` (inclusive) and `end` (exclusive).
     """
-    sorted_variants = hl.sorted(variants, key=lambda x: x.locus.position)
+    # Same sort key as get_haplo_sequence so the two stay in lockstep; only the minimum position
+    # (`[0].locus.position`) is read here, and `_max_reference_end` is order-independent.
+    sorted_variants = hl.sorted(variants, key=lambda x: (x.locus.position, hl.len(x.alleles[0])))
     min_variant = sorted_variants[0]
     return hl.struct(
         start=min_variant.locus.position - 1 - window_size,

divref/divref/tools/compute_haplotypes.py

@@ -13,6 +13,33 @@
 logger = logging.getLogger(__name__)
 
 
+def _haploid_adjusted_call(
+    locus: hl.Expression,
+    gt: hl.Expression,
+    sex_karyotype: hl.Expression,
+) -> hl.Expression:
+    """
+    Collapse chrX non-PAR male genotypes to haploid; pass every other call through unchanged.
+
+    Males (`sex_karyotype == "XY"`) at chrX non-PAR loci are encoded as pseudo-homozygous diploid
+    (0|0 / 1|1) in the SHAPEIT5 BCFs, but gnomAD reports chrX non-PAR allele numbers with males
+    counted as haploid. Replacing such a call with `hl.call(gt[0])` makes the downstream
+    `call_stats` AN/AC match the gnomAD convention. It is lossless given the verified encoding
+    (`GT[0] == GT[1]` there) but would undercount a heterozygous male non-PAR call if one were
+    ever emitted. Autosomes, PAR, and chrY are unaffected.
+
+    Args:
+        locus: The variant locus expression.
+        gt: The diploid genotype call expression.
+        sex_karyotype: The sample's sex-karyotype string expression.
+
+    Returns:
+        A call expression: haploid `call(gt[0])` for chrX non-PAR males, else `gt` unchanged.
+    """
+    is_male_nonpar = locus.in_x_nonpar() & (sex_karyotype == "XY")
+    return hl.if_else(is_male_nonpar, hl.call(gt[0]), gt)
+
+
 def _compute_locus_groups(
     variants_ht: hl.Table,
     window_size: int,
@@ -37,6 +64,7 @@ def _compute_locus_groups(
         `n_groups` is the total number of groups.
     """
     rows = variants_ht.key_by().select("row_idx", "locus", "ref_len").collect()
+    logger.info("compute_locus_groups: collected %d variants to the driver", len(rows))
     rows.sort(key=lambda v: (v.locus.contig, v.locus.position))
     group_of: dict[int, int] = {}
     current_group = -1
@@ -636,17 +664,9 @@ def compute_haplotypes(
     mt = mt.add_row_index().add_col_index()
     mt = mt.filter_entries(mt.freq[mt.pop_int].AF >= variant_freq_threshold)
 
-    # Males in chrX non-PAR are encoded as pseudo-homozygous diploid (0|0 or 1|1) in the
-    # SHAPEIT5 BCFs, but gnomAD reports chrX non-PAR allele numbers with males counted as
-    # haploid. Treat males as haploid here so that empirical AC/AN match the gnomAD HT
-    # convention; everywhere else (autosomes, PAR, chrY) keep the original diploid call.
-    #
-    # `hl.call(mt.GT[0])` discards the second allele. This is lossless given the verified
-    # SHAPEIT5 encoding (males are pseudo-homozygous in non-PAR, so GT[0] == GT[1]) but
-    # would silently undercount carriers if upstream ever emitted heterozygous male calls
-    # in non-PAR.
-    is_male_nonpar = mt.locus.in_x_nonpar() & (mt.sex_karyotype == "XY")
-    adjusted_gt = hl.if_else(is_male_nonpar, hl.call(mt.GT[0]), mt.GT)
+    # Count chrX non-PAR males as haploid so empirical AC/AN match gnomAD's convention; see
+    # `_haploid_adjusted_call`. Autosomes, PAR, and chrY keep the original diploid call.
+    adjusted_gt = _haploid_adjusted_call(mt.locus, mt.GT, mt.sex_karyotype)
     mt = mt.annotate_rows(
         frequencies_by_pop=hl.agg.group_by(mt.pop_int, hl.agg.call_stats(adjusted_gt, 2)),
         ref_len=hl.len(mt.alleles[0]),

divref/tests/test_haplotype_compat.py

@@ -61,6 +61,10 @@ def test_variant_distance(a: str, b: str, expected: int) -> None:
         ("chr1:100:CA:C", "chr1:100:C:CGG", "same_position_insertion_deletion"),  # "A" != "GG"
         # Insertion anchored at a deleted base: contested anchor -> conflict.
         ("chr1:400:TGG:T", "chr1:402:G:GTTTT", "insertion_in_deletion"),
+        # MNV overlaps (rare in practice) fall through to the catch-all reasons.
+        ("chr1:100:AT:GC", "chr1:100:AT:CG", "same_position_other"),  # two MNVs at one site
+        ("chr1:100:ATGC:A", "chr1:101:TG:CA", "other_overlap"),  # MNV inside a deletion
+        ("chr1:100:ATG:CGT", "chr1:101:T:A", "other_overlap"),  # MNV spanning a downstream SNP
         # Composable overlaps -> None.
         ("chr1:100:C:A", "chr1:100:C:CAA", None),  # snp + insertion at one site
         ("chr1:100:C:A", "chr1:100:CA:C", None),  # snp + deletion at one site

divref/tests/tools/test_compute_haplotypes.py

@@ -14,6 +14,7 @@
 from divref.tools.compute_haplotypes import _compute_metrics
 from divref.tools.compute_haplotypes import _enumerate_subfragments
 from divref.tools.compute_haplotypes import _form_parent_blocks
+from divref.tools.compute_haplotypes import _haploid_adjusted_call
 from divref.tools.compute_haplotypes import compute_haplotypes
 
 
@@ -1123,3 +1124,33 @@ def test_compute_haplotypes_chrx_nonpar(
     results: list[hl.Struct] = result.collect()
     assert all(len(r.haplotype) >= 2 for r in results)
     assert all(len(r.variants) == len(r.haplotype) for r in results)
+
+
+def test_haploid_adjusted_call_counts_chrx_nonpar_males_once(
+    hail_context: None,  # noqa: ARG001
+) -> None:
+    """A chrX non-PAR male alt-carrier contributes one allele to call_stats; everyone else two."""
+    # Two samples: col 0 is male (XY), col 1 is female (XX). Two sites: a chrX non-PAR locus and an
+    # autosome locus. Every entry is a (pseudo-)homozygous alt call (1|1).
+    mt = hl.utils.range_matrix_table(n_rows=2, n_cols=2)
+    mt = mt.annotate_cols(sex_karyotype=hl.if_else(mt.col_idx == 0, "XY", "XX"))
+    mt = mt.annotate_rows(
+        locus=hl.if_else(
+            mt.row_idx == 0,
+            hl.locus("chrX", 50_000_000, reference_genome="GRCh38"),  # well inside non-PAR
+            hl.locus("chr1", 1_000_000, reference_genome="GRCh38"),
+        ),
+        alleles=["A", "T"],
+    )
+    mt = mt.annotate_entries(GT=hl.call(1, 1))
+
+    adjusted = _haploid_adjusted_call(mt.locus, mt.GT, mt.sex_karyotype)
+    mt = mt.annotate_rows(cs=hl.agg.call_stats(adjusted, 2))
+    by_contig = {row.locus.contig: row.cs for row in mt.rows().collect()}
+
+    # chrX non-PAR: male counted haploid (1 allele) + female diploid (2) -> AN 3, all alt.
+    assert by_contig["chrX"].AN == 3
+    assert by_contig["chrX"].AC[1] == 3
+    # Autosome: both diploid -> AN 4 (the correction does not apply off chrX non-PAR).
+    assert by_contig["chr1"].AN == 4
+    assert by_contig["chr1"].AC[1] == 4