SimurgArray/scripts/pgx/validate_star_getrm.py

#!/usr/bin/env python3
"""Compare `simurg pgx star-allele call` with GeT-RM consensus diplotypes (PGX-032).

    python scripts/pgx/validate_star_getrm.py --calls STAR_DIR/pgx_calls.json --truth getrm_truth.tsv \
        --identity scripts/pgx/demo36_identity.tsv --out summary.json

GeT-RM notation is normalised per allele: "*2(*35)" / "*2 (*35)" accept *2 or *35; a trailing sub-allele letter
("*1F") also matches the core allele ("*1") when the call has no letter; "(*68)+*4" tandems and "xN" duplications
reduce to the core allele (structural forms are reported separately); "(*28 + *60)" keeps the star alleles the DB
defines. "A/B or C/D" accepts either. UGT1A1 and SLCO1B1 are reported but left out of the total: GeT-RM uses pre-2023
names (UGT1A1 *60, SLCO1B1 *1B/*17/*21) that the CPIC definitions no longer carry. A sample's gene counts as **top**
concordant when the top solution matches, and as **equivalent**
concordant when the truth is among the array-equivalent members of any feasible solution. Replicate arrays of one
individual are counted once (the first).
"""

from __future__ import annotations

import argparse
import csv
import json
import re
from collections import defaultdict
from pathlib import Path


def truth_options(text: str) -> list[set[str]] | None:
    """Per haplotype, the set of acceptable core alleles; None when the text is not a diplotype."""
    t = text.replace(" ", "")
    if "/" not in t:
        return None
    out = []
    for side in t.split("/", 1):
        opts = set(re.findall(r"\*\d+[A-Z]?", side))
        if not opts:
            return None
        cores = {re.sub(r"x\d+|xN", "", o) for o in opts}
        cores |= {re.sub(r"[A-Z]$", "", o) for o in cores}
        if "+" in side:  # tandem / hybrid: the non-parenthesised last unit is the core allele
            last = re.findall(r"\*\d+[A-Z]?", side.split("+")[-1])
            cores |= set(last)
        out.append(cores)
    return out


LEGACY_NOMENCLATURE = {"UGT1A1", "SLCO1B1"}  # GeT-RM uses pre-2023 names (*60, SLCO1B1 *1B/*17/*21) not in CPIC


def truth_alternatives(text: str) -> list[list[set[str]]]:
    """GeT-RM "A/B or C/D" → one option list per alternative diplotype."""
    alts = [truth_options(x) for x in re.split(r"\bor\b", text)]
    return [x for x in alts if x is not None]


def matches(call: tuple[str, str], truth: list[list[set[str]]]) -> bool:
    a, b = (re.sub(r"x\d+$", "", x) for x in call)
    return any((a in t[0] and b in t[1]) or (a in t[1] and b in t[0]) for t in truth)


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--calls", required=True, type=Path)
    ap.add_argument("--truth", required=True, type=Path)
    ap.add_argument("--identity", required=True, type=Path)
    ap.add_argument("--out", required=True, type=Path)
    a = ap.parse_args()
    with a.identity.open() as fh:
        ident = {r["sample_id"]: r["individual"] for r in csv.DictReader((ln for ln in fh if not ln.startswith("#")),
                                                                           delimiter="\t")}  # fmt: skip
    truth: dict[tuple[str, str], str] = {}
    with a.truth.open() as fh:
        for r in csv.DictReader(fh, delimiter="\t"):
            truth[(r["individual"], r["gene"])] = r["diplotype"]
    calls = json.loads(a.calls.read_text())
    per_gene: dict[str, dict[str, int]] = defaultdict(lambda: defaultdict(int))
    detail = []
    seen = set()
    for s in calls["samples"]:
        ind = ident.get(s["sample_id"])
        if ind is None or ind in seen:
            continue
        seen.add(ind)
        for g in s["genes"]:
            t = truth.get((ind, g["gene"]))
            if t is None:
                continue
            opts = truth_alternatives(t)
            st = per_gene[g["gene"]]
            if not opts:
                st["not_comparable"] += 1
                continue
            st["n"] += 1
            top = g["candidates"][0]["genotype"].split("/") if g["candidates"] else ["?", "?"]
            ok_top = len(top) == 2 and matches((top[0], top[1]), opts)
            feas = [c for c in g["candidates"] if c["genotype"] in g["feasible_solutions"]]
            ok_eq = ok_top or any(matches(tuple(e.split("/")[:2]), opts) for c in feas for e in c["equivalents"]
                                  if len(e.split("/")) >= 2)  # fmt: skip
            st["top"] += ok_top
            st["equivalent"] += ok_eq
            if not ok_top:
                detail.append({"individual": ind, "gene": g["gene"], "truth": t, "call": g["genotype"],
                               "within_equivalents": ok_eq})  # fmt: skip
    comparable = {k: v for k, v in per_gene.items() if k not in LEGACY_NOMENCLATURE}
    summary = {"individuals": len(seen), "genes": {k: dict(v) for k, v in sorted(per_gene.items())},
               "legacy_nomenclature_genes": sorted(LEGACY_NOMENCLATURE),
               "total": {k: sum(v.get(k, 0) for v in comparable.values()) for k in ("n", "top", "equivalent")},
               "discordant": detail}  # fmt: skip
    a.out.write_text(json.dumps(summary, indent=1))
    for g, v in sorted(per_gene.items()):
        print(f"{g:8s} top {v['top']}/{v['n']}  within equivalents {v['equivalent']}/{v['n']}"
              + (f"  (not comparable {v['not_comparable']})" if v.get("not_comparable") else "")
              + ("  [legacy nomenclature]" if g in LEGACY_NOMENCLATURE else ""))  # fmt: skip
    print("total (excluding legacy-nomenclature genes)", summary["total"])
    return 0


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