kemal
b96ae21c70
Sub-region CNs (p5, intron 2, exon 9) resolve full copies and *36/*68/*13 units; candidates carry the tandem names, background alleles, and CPIC activity sums. GeT-RM: CYP2D6 13/15 top (NA12878 *3/*68+*4 and NA11832 *1/*68+*4 exact), 15/15 within feasible; all genes 132/144.
244 lines
11 KiB
Python
244 lines
11 KiB
Python
"""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 = ""
|
||
basis: str = "star_allele" # PGX-002 result basis: star_allele | haplotype_named | variant_based | tagging_variant
|
||
activity: dict[str, float] = field(default_factory=dict) # allele → activity value (variant-based ordering)
|
||
function: dict[str, str] = field(default_factory=dict)
|
||
|
||
|
||
@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
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class Cyp2d6Structure:
|
||
"""CYP2D6 copies from the sub-region CNs (PGX-013) [OURS; DOC: PharmVar CYP2D6 structural variation]:
|
||
``full`` complete CYP2D6 copies, ``u36`` *36-type units (CYP2D7 exon 9), ``u68`` *68-type units (CYP2D6 exon 1
|
||
only), ``u13`` *13-type units (CYP2D7 5′ part). CN(p5) = full + u36 + u68, CN(intron 2) = full + u36 + u13,
|
||
CN(exon 9) = full + u13; with u68 · u13 = 0 this has one solution."""
|
||
|
||
full: int
|
||
u36: int
|
||
u68: int
|
||
u13: int
|
||
|
||
|
||
def cyp2d6_structure(p5: int | None, intron2: int | None, exon9: int | None) -> Cyp2d6Structure | None:
|
||
if p5 is None:
|
||
return None
|
||
ni = p5 if intron2 is None else intron2
|
||
n9 = ni if exon9 is None else exon9
|
||
u36 = max(0, ni - n9)
|
||
x = p5 - u36 - n9
|
||
if x >= 0:
|
||
full, u68, u13 = n9, x, 0
|
||
else:
|
||
full, u68, u13 = p5 - u36, 0, -x
|
||
if full < 0:
|
||
return Cyp2d6Structure(max(0, p5), 0, 0, 0)
|
||
return Cyp2d6Structure(full, u36, u68, u13)
|
||
|
||
|
||
def structure_candidates(m: GeneModel, st: Cyp2d6Structure) -> list[tuple[tuple[int, ...], tuple[str, ...]]]:
|
||
"""Candidates for a CYP2D6 structure: SNV copies = full + u36 (a *36 unit carries the *10 sequence up to exon 8);
|
||
*36 units attach to a *10 copy (``*36+*10``), *68 units to *4 when present (``*68+*4``), *13 units are separate
|
||
haplotypes."""
|
||
snv_copies = st.full + st.u36
|
||
base = candidates(m, snv_copies)
|
||
i10 = m.haplotypes.index("*10") if "*10" in m.haplotypes else None
|
||
out = []
|
||
for ix, names in base:
|
||
names_l = list(names)
|
||
if st.u36:
|
||
if i10 is None or ix.count(i10) < st.u36 + 1: # the *36 unit(s) plus the *10 they precede
|
||
continue
|
||
k = next((n for n, x in enumerate(names_l) if x.startswith("*10")), None)
|
||
if k is None:
|
||
continue
|
||
unit = "*36" if st.u36 == 1 else f"*36x{st.u36}"
|
||
names_l[k] = f"{unit}+*10"
|
||
if st.u68:
|
||
unit = "*68" if st.u68 == 1 else f"*68x{st.u68}"
|
||
k = next((n for n, x in enumerate(names_l) if x.split("+")[-1].split("x")[0] == "*4"), 0)
|
||
names_l[k] = f"{unit}+{names_l[k]}"
|
||
if st.u13:
|
||
if len(names_l) == 2 and names_l[1] == (m.deletion_allele or "*0"):
|
||
names_l[1] = "*13" if st.u13 == 1 else f"*13x{st.u13}"
|
||
else:
|
||
names_l[0] = f"*13+{names_l[0]}"
|
||
out.append((ix, tuple(names_l)))
|
||
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, structure: Cyp2d6Structure | None = None,
|
||
) -> tuple[list[Solution], list[Solution]]: # fmt: skip
|
||
"""All scored solutions (ranked) and the feasible ones (PGX-014/015)."""
|
||
called = obs.state == CALLED
|
||
haploid: list[tuple[tuple[int, ...], tuple[str, ...]]] = [((i,), (m.haplotypes[i],)) for i in
|
||
range(len(m.haplotypes))] # fmt: skip
|
||
if ploidy == 1:
|
||
cands = haploid
|
||
elif structure is not None and (structure.u36 or structure.u68 or structure.u13):
|
||
cands = structure_candidates(m, structure)
|
||
cn = structure.full + structure.u36
|
||
else:
|
||
cands = 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
|