SimurgArray/scripts/pgx/wgs_star_truth.py

#!/usr/bin/env python3
"""Star-allele truth from 1000 Genomes 30x phased genotypes with the PGx DB definitions (SNV/indel alleles only).

    python scripts/pgx/wgs_star_truth.py --db DB.sqlite --individuals ids.txt --out wgs_truth.tsv [--genes G,G]

For every star-allele gene of the DB the defining variants are fetched from the NYGC phased panel (pysam over HTTPS,
one window per gene; ``uv pip install pysam``). Each phased haplotype gets the allele whose defining variants it
carries completely and that explains most of the haplotype's carried variants — its defining ones plus those at its
wildcard (IUPAC) positions, so a *4 haplotype carrying the *10 variants stays *4 (ties: more defining variants,
higher CPIC/PharmVar frequency, natural order); a haplotype carrying none is the reference allele. A variant absent
from the filtered panel counts as reference only when rare (max allele frequency < 5 %); a gene with a common
defining variant missing from the panel is written with ``incomplete`` (not comparable). Structural alleles (CYP2D6
*5, xN, hybrids; CYP2A6 *4) are invisible to this truth: the diplotypes are SNV-level only.
"""

from __future__ import annotations

import argparse
import os
import re
import sqlite3
from collections import defaultdict
from pathlib import Path

PANEL = (
    "https://ftp.1000genomes.ebi.ac.uk/vol1/ftp/data_collections/1000G_2504_high_coverage/working/"
    "20220422_3202_phased_SNV_INDEL_SV/1kGP_high_coverage_Illumina.{c}.filtered.SNV_INDEL_SV_phased_panel.vcf.gz"
)


def natural(a: str) -> tuple[object, ...]:
    return tuple(int(x) if x.isdigit() else x for x in re.split(r"(\d+)", a))


def main() -> int:
    import pysam

    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--db", required=True, type=Path)
    ap.add_argument("--individuals", required=True, type=Path)
    ap.add_argument("--out", required=True, type=Path)
    ap.add_argument("--genes", default="")
    ap.add_argument("--cache", type=Path, default=Path("data/pgx/cache"))
    a = ap.parse_args()
    inds = sorted(set(a.individuals.read_text().split()))
    out_path = a.out.resolve()
    con = sqlite3.connect(a.db)
    genes = [g for (g,) in con.execute("SELECT gene FROM genes WHERE result_basis = 'star_allele' ORDER BY gene")]
    if a.genes:
        genes = [g for g in genes if g in a.genes.split(",")]
    a.cache.mkdir(parents=True, exist_ok=True)
    os.chdir(a.cache)
    rows = []
    for gene in genes:
        q_alle = "SELECT allele FROM alleles WHERE gene = ? AND defined = 1 AND structural = 0"
        alle = [x for (x,) in con.execute(q_alle, (gene,))]
        q_ref = "SELECT allele FROM alleles WHERE gene = ? AND kind = 'reference'"
        ref_allele = next((x for (x,) in con.execute(q_ref, (gene,))), "*1")
        defs: dict[str, set[tuple[int, str, str]]] = defaultdict(set)
        chrom = None
        loc_keys: dict[int, set[tuple[int, str, str]]] = defaultdict(set)
        q_var = "SELECT allele, chrom, pos, ref, alt, location_id FROM allele_variants WHERE gene = ?"
        for allele, c, p, r, al, lid in con.execute(q_var, (gene,)):
            defs[allele].add((p, r, al))
            loc_keys[lid].add((p, r, al))
            chrom = c
        wild: dict[str, set[tuple[int, str, str]]] = defaultdict(set)
        for allele, lid in con.execute("SELECT allele, location_id FROM allele_wildcards WHERE gene = ?", (gene,)):
            wild[allele] |= loc_keys.get(lid, set())
        if chrom is None:
            continue
        freq: dict[str, float] = defaultdict(float)
        for allele, f in con.execute("SELECT allele, AVG(freq) FROM frequencies WHERE gene = ? GROUP BY allele",
                                     (gene,)):  # fmt: skip
            freq[allele] = float(f or 0)
        vmax: dict[tuple[int, str, str], float] = defaultdict(float)
        for allele, vs in defs.items():
            for v in vs:
                vmax[v] = max(vmax[v], freq.get(allele, 0.0))
        keys = set().union(*defs.values()) if defs else set()
        core = chrom.removeprefix("chr")
        vf = pysam.VariantFile(PANEL.format(c=f"chr{core}"))
        present = [s for s in inds if s in vf.header.samples]
        vf.subset_samples(present)
        lo = min(p for p, _, _ in keys) - 10
        hi = max(p + len(r) for p, r, _ in keys) + 10
        hap_alts: dict[tuple[str, int], set[tuple[int, str, str]]] = defaultdict(set)
        seen = set()
        for rec in vf.fetch(f"chr{core}", lo, hi):
            for ai, alt in enumerate(rec.alts or (), start=1):
                k = (rec.pos, rec.ref, alt)
                if k not in keys:
                    continue
                seen.add(k)
                for s in present:
                    gt = rec.samples[s]["GT"]
                    for h, x in enumerate(gt[:2]):
                        if x == ai:
                            hap_alts[(s, h)].add(k)
        missing_common = [k for k in keys - seen if vmax[k] >= 0.05]
        for s in present:
            haps = []
            for h in (0, 1):
                carried = hap_alts[(s, h)]
                # the allele that explains most carried variants (its defining ones plus those at its wildcard
                # positions), then the most specific, then the most frequent
                best = (0, 0, 0.0)
                pick = ref_allele
                for x in alle:
                    d = defs.get(x, set())
                    if not d or not d <= carried:
                        continue
                    score = (len(carried & (d | wild.get(x, set()))), len(d), freq.get(x, 0.0))
                    if score > best or (score == best and natural(x) < natural(pick)):
                        best, pick = score, x
                haps.append(pick)
            dip = "/".join(sorted(haps, key=natural))
            rows.append((s, gene, dip, "incomplete" if missing_common else ""))
        print(f"{gene}: {len(keys)} defining variant(s), {len(seen)} in the panel, "
              f"{len(missing_common)} common missing", flush=True)  # fmt: skip
    with out_path.open("w") as fh:
        fh.write("individual\tgene\tdiplotype\tflag\n")
        for r in sorted(rows):
            fh.write("\t".join(r) + "\n")
    print(f"{len(rows)} gene call(s) -> {out_path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())