197 lines
8.4 KiB
Python
197 lines
8.4 KiB
Python
"""Genome-wide CN/LOH/mosaic hidden Markov model — spec CYTO-004/005/012/018, plan §10.2. Pure functions, no I/O.
|
||
|
||
[LIT: Wang et al. 2007 Genome Res 17:1665 (PennCNV: LRR/BAF emissions weighted by population B-allele frequency,
|
||
distance-dependent transitions); Loh et al. 2018 Nature 559:350 (MoChA: mosaic BAF split); DOC: cnvPartition
|
||
technical note (reference LRR state means CN0 −5, CN1 −0.45, CN2 0, CN3 0.3, CN4 0.75)] [OURS: state grid, mixture
|
||
weights, switch model]
|
||
|
||
* Diploid regions: CN0, CN1, CN2, CN2-LOH, CN3, CN3-LOH, CN4 and mosaic loss / gain / CN-LOH at cell fractions
|
||
0.2, 0.4, 0.6, 0.8. Haploid regions (male X outside PAR, male Y): CN0, CN1, CN2 and mosaic gain/loss relative to
|
||
the region's median LRR.
|
||
* LRR ~ (1−ε)·N(µ_state, σ) + ε·U(−5, 3); σ = 1.4826·MAD of the sample's autosomal LRR (floor 0.08), scaled per
|
||
probe by the reference's LRR spread (CYTO-042); CN0 σ = 2.
|
||
Mosaic means interpolate linearly between the constitutional means by the cell fraction.
|
||
* BAF ~ (1−ε)·Σ_g w_g N(b_g, σ_b) + ε·U(0, 1); genotype weights from the probe's PFB p (binomial in the copy
|
||
number; LOH: {0: 1−p, 1: p}); mosaic heterozygotes split into the two bands 0.5 ± d(f).
|
||
* Transitions: switch probability 1 − exp(−d / D) for inter-probe distance d, D = 5·10⁸ bp; a switch lands on state
|
||
j with probability π_j (the same target distribution from every state).
|
||
* Viterbi path for segmentation; forward–backward posteriors for QUAL.
|
||
"""
|
||
|
||
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]
|
||
FRACTIONS = (0.2, 0.4, 0.6, 0.8)
|
||
LRR_MEAN = {0: -5.0, 1: -0.45, 2: 0.0, 3: 0.3, 4: 0.75}
|
||
EPS = 0.01
|
||
MOSAIC_PRIOR = 0.004 # total switch-target mass of the 12 diploid mosaic states
|
||
SWITCH_D = 5e8
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class State:
|
||
name: str
|
||
cn: int # copy number (baseline-relative for haploid regions: 0, 1, 2)
|
||
loh: bool
|
||
mosaic: str # "" | "loss" | "gain" | "cnloh"
|
||
fraction: float # cell fraction of mosaic states, 1.0 otherwise
|
||
call: str # "" (baseline) | DEL | DUP | LOH | GAINLOH | MDEL | MDUP | MLOH
|
||
lrr_mu: float
|
||
baf: tuple[tuple[str, float], ...] # (genotype-weight key, BAF mean)
|
||
prior: float
|
||
|
||
|
||
def _d_split(kind: str, f: float) -> float:
|
||
if kind == "loss":
|
||
return f / (2 * (2 - f))
|
||
if kind == "gain":
|
||
return f / (2 * (2 + f))
|
||
return f / 2
|
||
|
||
|
||
def diploid_states(mosaic_prior: float = MOSAIC_PRIOR) -> list[State]:
|
||
s = [
|
||
State("CN0", 0, False, "", 1.0, "DEL", LRR_MEAN[0], (), 0.01),
|
||
State("CN1", 1, False, "", 1.0, "DEL", LRR_MEAN[1], (("h0", 0.0), ("h1", 1.0)), 0.06),
|
||
State("CN2", 2, False, "", 1.0, "", 0.0, (("aa", 0.0), ("ab", 0.5), ("bb", 1.0)), 0.60),
|
||
State("CN2LOH", 2, True, "", 1.0, "LOH", 0.0, (("h0", 0.0), ("h1", 1.0)), 0.06),
|
||
State("CN3", 3, False, "", 1.0, "DUP", LRR_MEAN[3],
|
||
(("t0", 0.0), ("t1", 1 / 3), ("t2", 2 / 3), ("t3", 1.0)), 0.06),
|
||
State("CN3LOH", 3, True, "", 1.0, "GAINLOH", LRR_MEAN[3], (("h0", 0.0), ("h1", 1.0)), 0.02),
|
||
State("CN4", 4, False, "", 1.0, "DUP", LRR_MEAN[4],
|
||
(("q0", 0.0), ("q1", 0.25), ("q2", 0.5), ("q3", 0.75), ("q4", 1.0)), 0.03),
|
||
] # fmt: skip
|
||
for kind, call, mu_end in (("loss", "MDEL", LRR_MEAN[1]), ("gain", "MDUP", LRR_MEAN[3]), ("cnloh", "MLOH", 0.0)):
|
||
for f in FRACTIONS:
|
||
d = _d_split(kind, f)
|
||
cn = {"loss": 1, "gain": 3, "cnloh": 2}[kind]
|
||
bands = (("aa", 0.0), ("ab-", 0.5 - d), ("ab+", 0.5 + d), ("bb", 1.0))
|
||
s.append(State(f"M{kind}{f:.1f}", cn, kind == "cnloh", kind, f, call, mu_end * f, bands, mosaic_prior / 12))
|
||
return s
|
||
|
||
|
||
def haploid_states(mosaic_prior: float = MOSAIC_PRIOR) -> list[State]:
|
||
s = [
|
||
State("H0", 0, False, "", 1.0, "DEL", LRR_MEAN[0], (), 0.02),
|
||
State("H1", 1, False, "", 1.0, "", 0.0, (("h0", 0.0), ("h1", 1.0)), 0.80),
|
||
State("H2", 2, False, "", 1.0, "DUP", 0.45, (("h0", 0.0), ("h1", 1.0)), 0.06),
|
||
]
|
||
for f in FRACTIONS:
|
||
s.append(State(f"Hgain{f:.1f}", 2, False, "gain", f, "MDUP", 0.45 * f, (("h0", 0.0), ("h1", 1.0)),
|
||
mosaic_prior / 8)) # fmt: skip
|
||
s.append(State(f"Hloss{f:.1f}", 0, False, "loss", f, "MDEL", -1.0 * f, (("h0", 0.0), ("h1", 1.0)),
|
||
mosaic_prior / 8)) # fmt: skip
|
||
return s
|
||
|
||
|
||
def _geno_weights(key: str, p: F64) -> F64:
|
||
q = 1 - p
|
||
table = {
|
||
"aa": q * q, "ab": 2 * p * q, "bb": p * p, "ab-": p * q, "ab+": p * q,
|
||
"h0": q, "h1": p,
|
||
"t0": q**3, "t1": 3 * p * q * q, "t2": 3 * p * p * q, "t3": p**3,
|
||
"q0": q**4, "q1": 4 * p * q**3, "q2": 6 * p * p * q * q, "q3": 4 * p**3 * q, "q4": p**4,
|
||
} # fmt: skip
|
||
return table[key]
|
||
|
||
|
||
def _norm(x: F64, mu: float, sd: float | F64) -> F64:
|
||
return np.exp(-0.5 * ((x - mu) / sd) ** 2) / (sd * math.sqrt(2 * math.pi))
|
||
|
||
|
||
@dataclass
|
||
class Params:
|
||
lrr_sd: float | F64 # scalar or per probe (CYTO-042)
|
||
baf_sd: float
|
||
lrr_offset: float = 0.0 # added to every state mean (haploid regions: the region's median LRR)
|
||
switch_d: float = SWITCH_D
|
||
eps: float = EPS
|
||
extra: dict[str, float] = field(default_factory=dict)
|
||
|
||
|
||
def emissions(states: list[State], lrr: F64, baf: F64, pfb: F64, par: Params, het_sd: F64 | None = None) -> F64:
|
||
"""Per-probe, per-state emission likelihood (n × S); missing LRR/BAF contribute 1. ``het_sd``: per-probe spread
|
||
of heterozygous BAF in the reference; heterozygous components use max(σ_b, het_sd)."""
|
||
n = len(lrr)
|
||
out = np.ones((n, len(states)), dtype=np.float64)
|
||
lok = np.isfinite(lrr)
|
||
bok = np.isfinite(baf)
|
||
x = np.where(lok, lrr, 0.0)
|
||
b = np.clip(np.where(bok, baf, 0.5), 0.0, 1.0)
|
||
p = np.clip(np.where(np.isfinite(pfb), pfb, 0.5), 0.01, 0.99)
|
||
sd_het: float | F64 = (
|
||
np.maximum(par.baf_sd, np.nan_to_num(het_sd, nan=par.baf_sd)) if het_sd is not None else par.baf_sd
|
||
)
|
||
for j, st in enumerate(states):
|
||
sd = 2.0 if st.cn == 0 and not st.mosaic else par.lrr_sd
|
||
le = (1 - par.eps) * _norm(x, st.lrr_mu + par.lrr_offset, sd) + par.eps / 8.0
|
||
if st.baf:
|
||
be = np.zeros(n)
|
||
for key, mean in st.baf:
|
||
bsd = par.baf_sd if mean in (0.0, 1.0) or het_sd is None else sd_het
|
||
be += _geno_weights(key, p) * _norm(b, mean, bsd)
|
||
be = (1 - par.eps) * be + par.eps
|
||
else:
|
||
be = np.ones(n)
|
||
out[:, j] = np.where(lok, le, 1.0) * np.where(bok, be, 1.0)
|
||
return out
|
||
|
||
|
||
def switch_probs(pos: npt.NDArray[np.int64], switch_d: float = SWITCH_D) -> F64:
|
||
"""Switch probability before each probe (the first probe gets 1: it draws from the prior)."""
|
||
d = np.diff(np.asarray(pos, dtype=np.float64), prepend=np.nan)
|
||
p = 1.0 - np.exp(-np.maximum(d, 0.0) / switch_d)
|
||
p[0] = 1.0
|
||
return np.clip(np.nan_to_num(p, nan=1.0), 1e-12, 1.0)
|
||
|
||
|
||
def forward_backward(e: F64, sw: F64, prior: F64) -> F64:
|
||
"""Posterior state probabilities (n × S) under the switch model."""
|
||
n, s = e.shape
|
||
alpha = np.empty((n, s))
|
||
c = np.empty(n)
|
||
a = prior * e[0]
|
||
c[0] = a.sum() or 1e-300
|
||
alpha[0] = a / c[0]
|
||
for t in range(1, n):
|
||
a = ((1 - sw[t]) * alpha[t - 1] + sw[t] * prior) * e[t]
|
||
c[t] = a.sum() or 1e-300
|
||
alpha[t] = a / c[t]
|
||
beta = np.ones(s)
|
||
post = np.empty((n, s))
|
||
post[n - 1] = alpha[n - 1]
|
||
for t in range(n - 1, 0, -1):
|
||
eb = e[t] * beta
|
||
beta = ((1 - sw[t]) * eb + sw[t] * float(prior @ eb)) / c[t]
|
||
g = alpha[t - 1] * beta
|
||
post[t - 1] = g / (g.sum() or 1e-300)
|
||
return post
|
||
|
||
|
||
def viterbi(e: F64, sw: F64, prior: F64) -> npt.NDArray[np.int32]:
|
||
"""Most likely state path under the switch model (log space)."""
|
||
n, s = e.shape
|
||
le = np.log(np.maximum(e, 1e-300))
|
||
lp = np.log(np.maximum(prior, 1e-300))
|
||
back = np.empty((n, s), dtype=np.int32)
|
||
ld = lp + le[0]
|
||
idx = np.arange(s, dtype=np.int32)
|
||
for t in range(1, n):
|
||
stay = np.log(np.maximum(1 - sw[t] + sw[t] * prior, 1e-300)) + ld
|
||
best = int(np.argmax(ld))
|
||
jump = math.log(max(sw[t], 1e-300)) + lp + ld[best]
|
||
take_jump = jump > stay
|
||
back[t] = np.where(take_jump, best, idx)
|
||
ld = np.where(take_jump, jump, stay) + le[t]
|
||
path = np.empty(n, dtype=np.int32)
|
||
path[-1] = int(np.argmax(ld))
|
||
for t in range(n - 1, 0, -1):
|
||
path[t - 1] = back[t, path[t]]
|
||
return path
|