SimurgArray/scripts/validate_p6_vcf.py

#!/usr/bin/env python3
"""Phase P6 acceptance: SimurgArray VCFs vs a reference converter's multi-sample VCF (gtc2vcf or GTCtoVCF).

Both sides are left-aligned with ``bcftools norm -f`` and compared per probe (ID) and sample as unordered allele
sets (a haploid call ``a`` equals ``a/a``). Every difference is put into a category; the acceptance criterion
(plan §18 P6) is zero *unexplained* genotype differences.

Usage: validate_p6_vcf.py --ours DIR --oracle F.vcf.gz --fasta F --bcftools PATH --out summary.json
"""

from __future__ import annotations

import argparse
import json
import subprocess
import sys
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any


def query(bcftools: str, fasta: str, path: Path) -> list[list[str]]:
    norm = subprocess.Popen([bcftools, "norm", "-f", fasta, "-c", "w", "-Ou", str(path)], stdout=subprocess.PIPE,
                            stderr=subprocess.DEVNULL)  # fmt: skip
    out = subprocess.run(
        [bcftools, "query", "-f", "%ID\t%CHROM\t%POS\t%REF\t%ALT[\t%GT]\n", "-"],
        stdin=norm.stdout, capture_output=True, text=True, check=True,
    ).stdout  # fmt: skip
    norm.wait()
    return [ln.split("\t") for ln in out.splitlines()]


def alleles_of(gt: str, ref: str, alt: str) -> tuple[str, ...] | None:
    if gt in (".", "./.", ".|."):
        return None
    al = [ref, *alt.split(",")]
    idx = gt.replace("|", "/").split("/")
    if any(i == "." for i in idx):
        return None
    got = tuple(sorted(al[int(i)] for i in idx))
    return got * 2 if len(got) == 1 else got


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--ours", required=True, type=Path)
    ap.add_argument("--oracle", required=True, type=Path)
    ap.add_argument("--fasta", required=True)
    ap.add_argument("--bcftools", required=True)
    ap.add_argument("--out", required=True, type=Path)
    ap.add_argument("--name", default="oracle")
    a = ap.parse_args()

    samples = subprocess.run([a.bcftools, "query", "-l", str(a.oracle)], capture_output=True, text=True,
                             check=True).stdout.split()  # fmt: skip
    oracle: dict[str, tuple[str, str, str, str, list[str]]] = {}
    for f in query(a.bcftools, a.fasta, a.oracle):
        for pid in f[0].replace(";", ",").split(","):
            oracle[pid] = (f[1], f[2], f[3], f[4], f[5:])
    cats: Counter[str] = Counter()
    examples: dict[str, list[Any]] = defaultdict(list)
    ours_ids: set[str] = set()
    probes_in_merged = 0
    for si, sample in enumerate(samples):
        vcf = a.ours / f"{sample}.snv.vcf.gz"
        if not vcf.exists():
            cats["sample_missing"] += 1
            continue
        for f in query(a.bcftools, a.fasta, vcf):
            ids = f[0].replace(";", ",").split(",")
            merged = len(ids) > 1
            ours_al = alleles_of(f[5], f[3], f[4])
            haploid = "/" not in f[5] and f[5] != "."
            for pid in ids:
                if si == 0:
                    ours_ids.add(pid)
                    probes_in_merged += merged
                o = oracle.get(pid)
                if o is None:
                    cats["only_ours"] += 1
                    continue
                if (o[0], o[1], o[2]) != (f[1], f[2], f[3]):
                    cats["representation_diff"] += 1
                    if len(examples["representation_diff"]) < 20:
                        examples["representation_diff"].append([pid, f[1:5], list(o[:4])])
                    continue
                theirs = alleles_of(o[4][si], o[2], o[3])
                if ours_al == theirs:
                    cats["match_called" if ours_al else "match_missing"] += 1
                    continue
                if ours_al is None:
                    if merged:
                        key = "ours_missing_merged_site"
                    elif f[1] in ("Y", "chrY") and f[5] == ".":
                        key = "ours_missing_haploid_Y"
                    elif theirs is not None and theirs[0] != theirs[1] and (f[5] == "." or haploid):
                        key = "ours_missing_het_in_haploid"
                    else:
                        key = "ours_missing_UNEXPLAINED"
                elif theirs is None:
                    key = "oracle_missing_ours_called_merged" if merged else "oracle_missing_ours_called"
                else:
                    key = "diff_merged_site" if merged else "diff_UNEXPLAINED"
                cats[key] += 1
                if len(examples[key]) < 20:
                    examples[key].append([sample, pid, f[1:6], o[4][si], list(o[:4])])
    only_oracle = sorted(set(oracle) - ours_ids)
    summary = {
        "oracle": a.name,
        "samples": len(samples),
        "probes_oracle": len(oracle),
        "probes_ours": len(ours_ids),
        "probes_in_merged_records": probes_in_merged,
        "probes_only_in_oracle": len(only_oracle),
        "only_in_oracle_examples": only_oracle[:20],
        "categories": dict(sorted(cats.items())),
        "examples": examples,
    }
    unexplained = sum(v for k, v in cats.items() if "UNEXPLAINED" in k)
    summary["unexplained"] = unexplained
    a.out.parent.mkdir(parents=True, exist_ok=True)
    a.out.write_text(json.dumps(summary, indent=2) + "\n")
    print(json.dumps({k: summary[k] for k in ("probes_oracle", "probes_ours", "probes_only_in_oracle", "categories",
                                                "unexplained")}))  # fmt: skip
    return 0 if unexplained == 0 else 1


if __name__ == "__main__":
    sys.exit(main())