SimurgArray/src/simurg/pgx/star.py

"""Star-allele / haplotype calling from array genotypes — spec PGX-012 … PGX-020. Pure functions, no I/O.

Evidence model (PGX-012) [OURS]: each defining variant v of a gene is ``called`` (alt dosage g, GenCall score GS,
BAF), ``no_call``, ``masked`` or ``not_covered``; only called variants contribute. With diploid CN 2 the genotype
likelihood is P(g | d) = 1 − ε for the expected alt dosage d and ε/2 otherwise, with ε = 1e-3 for GS ≥ 0.15 and
0.05 below (calibrated: after probe masking, 0 of 9740 GSA v4 PGx calls at DB variants disagreed with 1000G 30x in any
GS bin — errors are probe-specific and handled by masks, not by GS; K-029). When the gene's copy number n ≠ 2
(targeted CN), the BAF likelihood N(b; k/n, σ_b) replaces the
genotype term (k = expected alt copies). Haplotype positions with an IUPAC value (CPIC core definitions) are wildcards:
the best-fitting dosage is used.

Candidates (PGX-013) [DIFF: solver-free]: every unordered pair of defined haplotypes, and for n ≠ 2 the structural
forms of CYP2D6-like genes: n = 1 → one whole-gene deletion haplotype; n = 0 → two; n ≥ 3 → one haplotype duplicated
(``*1x2``, n − 1 copies). Diplotypes that no called variant separates are array-equivalent (PGX-017) and form one
solution, named after the member with the highest prior. Scores (PGX-014): raw_score = softmax(LL); score =
softmax(LL + log π), π = Σ over members of f(h1)·f(h2) (×2 for h1 ≠ h2), f floored at 1e-5.
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field

import numpy as np
import numpy.typing as npt

F64 = npt.NDArray[np.float64]
FREQ_FLOOR = 1e-5
CALLED, NO_CALL, MASKED, NOT_COVERED = 0, 1, 2, 3
STATE_NAMES = {CALLED: "called", NO_CALL: "no_call", MASKED: "masked", NOT_COVERED: "not_on_array"}


@dataclass
class GeneModel:
    gene: str
    variants: list[tuple[str, int, str, str]]  # (chrom, pos, ref, alt), one per defining variant
    variant_ids: list[tuple[str, str | None]]  # (hgvs, rsid)
    haplotypes: list[str]
    carries: npt.NDArray[np.int8]  # (H, V): 1 carries alt, 0 reference, −1 wildcard
    freq: F64  # (H,) population frequency of each haplotype (prior)
    deletion_allele: str | None = None  # named whole-gene deletion (CYP2D6 *5), or "*0"
    reference: str = ""


@dataclass
class Observations:
    state: npt.NDArray[np.int8]  # (V,)
    dosage: npt.NDArray[np.int8]  # alt dosage (0 … ploidy), −1 when not called
    gs: F64
    baf: F64


@dataclass
class Solution:
    haplotypes: tuple[str, ...]  # representative, e.g. ("*1", "*4") or ("*2x2", "*4")
    ll: float
    prior: float
    members: list[tuple[str, ...]] = field(default_factory=list)
    raw_score: float = 0.0
    score: float = 0.0
    copy_number: int = 2
    copies: tuple[int, ...] = ()  # haplotype index of every gene copy of the representative (SNV model)

    @property
    def genotype(self) -> str:
        return "/".join(self.haplotypes)


ERR_CALLED = 1e-3
ERR_LOW_GS = 0.05
LOW_GS = 0.15


def gs_error(gs: F64) -> F64:
    """PGX-012: per-call error rate from the GenCall score (calibrated, see the module docstring)."""
    return np.where(np.nan_to_num(gs, nan=0.0) >= LOW_GS, ERR_CALLED, ERR_LOW_GS)


def _dosage_ll(obs: Observations, expected: npt.NDArray[np.int16], wild: npt.NDArray[np.bool_]) -> F64:
    """(D, V) log-likelihood of each candidate's expected dosage (CN 2 genotype model)."""
    eps = gs_error(obs.gs)[None, :]
    g = obs.dosage[None, :].astype(np.int16)
    hit = (expected == g) | wild
    return np.where(hit, np.log(1 - eps), np.log(eps / 2))


def _baf_ll(obs: Observations, frac_lo: F64, frac_hi: F64, sigma: float) -> F64:
    """(D, V) BAF log-likelihood for expected alt fractions (a wildcard spans [lo, hi]: the closest value is used)."""
    b = np.nan_to_num(obs.baf, nan=0.5)[None, :]
    near = np.clip(b, frac_lo, frac_hi)
    return np.log(0.98 * np.exp(-0.5 * ((b - near) / sigma) ** 2) / (sigma * math.sqrt(2 * math.pi)) + 0.02)


def candidates(m: GeneModel, cn: int | None) -> list[tuple[tuple[int, ...], tuple[str, ...]]]:
    """(haplotype index per gene copy, display names) for every candidate diplotype (PGX-013)."""
    h = len(m.haplotypes)
    out: list[tuple[tuple[int, ...], tuple[str, ...]]] = []
    if cn is None or cn == 2:
        for i in range(h):
            for j in range(i, h):
                out.append(((i, j), (m.haplotypes[i], m.haplotypes[j])))
        return out
    dele = m.deletion_allele or "*0"
    if cn == 0:
        return [((), (dele, dele))]
    if cn == 1:
        return [((i,), (m.haplotypes[i], dele)) for i in range(h)]
    extra = cn - 2  # copies beyond two: one haplotype duplicated (e.g. CN 3 → *2x2/*4)
    for i in range(h):
        for j in range(h):
            out.append(((i,) * (extra + 1) + (j,), (f"{m.haplotypes[i]}x{extra + 1}", m.haplotypes[j])))
    return out


def call_gene(
    m: GeneModel, obs: Observations, cn: int | None = None, baf_sigma: float = 0.06, ploidy: int = 2,
    max_candidates: int = 10, feasible_ratio: float = 0.5,
) -> tuple[list[Solution], list[Solution]]:  # fmt: skip
    """All scored solutions (ranked) and the feasible ones (PGX-014/015)."""
    called = obs.state == CALLED
    haploid = [((i,), (m.haplotypes[i],)) for i in range(len(m.haplotypes))]
    cands = haploid if ploidy == 1 else candidates(m, cn)
    if not cands:
        return [], []
    carries = m.carries.astype(np.int16)
    lo = np.zeros((len(cands), carries.shape[1]))
    hi = np.zeros_like(lo)
    for k, (ix, _) in enumerate(cands):
        for i in ix:
            c = carries[i]
            lo[k] += np.where(c < 0, 0, c)
            hi[k] += np.where(c < 0, 1, c)
    ncp = np.array([max(1, len(ix)) for ix, _ in cands], dtype=np.float64)[:, None]
    if ploidy == 1 or cn is None or cn == 2:
        wild = hi != lo
        ll_v = _dosage_ll(obs, lo.astype(np.int16), wild)
        if wild.any():  # a wildcard may also take the higher dosage
            ll_v = np.maximum(ll_v, _dosage_ll(obs, hi.astype(np.int16), wild))
    else:
        ll_v = _baf_ll(obs, lo / ncp, hi / ncp, baf_sigma)
    ll = (ll_v * called[None, :]).sum(axis=1)
    f = np.maximum(m.freq, FREQ_FLOOR)
    # array-equivalence classes: identical expected dosages at every called variant (PGX-017)
    keys: dict[bytes, list[int]] = {}
    for k in range(len(cands)):
        sig = np.round(np.concatenate([lo[k, called], hi[k, called]]) / ncp[k, 0], 6).tobytes()
        keys.setdefault(sig + str(len(cands[k][0])).encode(), []).append(k)
    sols = []
    for members in keys.values():
        pri = []
        for k in members:
            ix = cands[k][0]
            if not ix:
                pri.append(1e-6)
                continue
            p = float(np.prod([f[i] for i in ix]))
            if len(set(ix)) > 1:
                p *= 2
            pri.append(p)
        best = members[int(np.argmax(pri))]
        sols.append(Solution(cands[best][1], float(ll[best]), float(sum(pri)), [cands[k][1] for k in members],
                             copy_number=cn if cn is not None else 2, copies=cands[best][0]))  # fmt: skip
    lls = np.array([s.ll for s in sols])
    lp = lls + np.log([s.prior for s in sols])
    raw = np.exp(lls - lls.max())
    sc = np.exp(lp - lp.max())
    raw, sc = raw / raw.sum(), sc / sc.sum()
    for s, r, v in zip(sols, raw, sc, strict=True):
        s.raw_score, s.score = float(r), float(v)
    sols.sort(key=lambda s: (-s.score, -s.raw_score, s.genotype))
    top = sols[0].score
    feasible = [s for s in sols if s.score >= feasible_ratio * top]
    return sols[: max(max_candidates, len(feasible))], feasible