ArrayTrain (plan §12): simurg train — EGT cluster files from a cohort's GTCs
- simurg/train/clusters.py: vectorised per-locus θ mixture (AA/AB/BB + outlier component, θ-ordered bounds), cluster selection and merging, lone clusters homozygous, missing clusters imputed by cross-locus regression per assay type, cluster_sep and train (GenTrain-field) scores. - formats/egt.py: write_egt (EGT v3, data block 9, round-trips through the reader). - pipeline/train.py: normalised θ/R from GTC raw X/Y + transforms, sex-aware (X females, Y males, MT haploid), intensity-only loci as one R cluster. - scripts/train/concordance_chr.py: array vs 1000G 30x on one chromosome. Demo (36 GSA v4 PGx arrays, chr22 vs 1000G): trained 99.29 % / NRD 2.6 % vs vendor EGT 99.57 % / 1.55 %.
This commit is contained in:
kemal committed
1 parent
004739da8e
commit
35d1e08621
9 files changed
+618
No files matched your search
@@ -0,0 +1 @@
|
||||
"""ArrayTrain: cluster-file training (plan §12)."""
|
||||
@@ -0,0 +1,233 @@
|
||||
"""ArrayTrain: per-locus genotype cluster training in (θ, R) — plan §12. Pure functions, no I/O.
|
||||
|
||||
Explicitly **not** a reproduction of the vendor's clustering [OURS; LIT: Teo et al. 2007 Bioinformatics 23:2741
|
||||
(Illuminus: mixture models in (θ, R)); Carvalho et al. 2007 Biostatistics 8:485 (CRLMM); Illumina GenTrain 2/3
|
||||
technical notes (cluster record fields)]:
|
||||
|
||||
1. **Mixture per locus**, vectorised over a block of loci: a three-component mixture on θ (AA, AB, BB) plus a small
|
||||
uniform outlier component, fitted by EM from fixed starts (θ ≈ 0.08 / 0.5 / 0.92) with the θ ordering enforced by
|
||||
bounded means (AA ≤ 0.4, 0.2 ≤ AB ≤ 0.8, BB ≥ 0.6). Haploid loci (male X, Y, MT) have no AB component.
|
||||
Robustness comes from the outlier component (the plan's Student-t mixture is replaced by it; documented).
|
||||
2. **Model selection**: a component keeps its cluster only with ≥ max(3, 1 % of the samples) hard-assigned
|
||||
samples, and adjacent clusters closer than ``MERGE_GAP`` in θ are merged (keeps 1-, 2- and 3-cluster models).
|
||||
3. **Missing clusters** are imputed from relationships learned over well-clustered loci of the same assay type
|
||||
(Infinium I / II): least-squares lines predicting a cluster's θ from its neighbour's (θ_AB from θ_AA, …), median
|
||||
R ratios R_AA/R_AB and R_BB/R_AB, median SDs — the empirical replacement for the vendor's neural network. A lone
|
||||
cluster is homozygous (AA or BB by its θ).
|
||||
4. **Scores**: ``cluster_sep`` = the smallest standardised θ gap between adjacent present clusters,
|
||||
gap / (sd₁ + sd₂), squashed to 0–1 by (gap_std − 1) / 3; ``train_score`` (the EGT GenTrain field, which scales
|
||||
every call's GenCall score) = mean posterior of the assigned cluster × (0.6 + 0.4 × cluster_sep) × fraction of
|
||||
samples with intensity, capped at 0.9 [OURS].
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
|
||||
F32 = npt.NDArray[np.float32]
|
||||
F64 = npt.NDArray[np.float64]
|
||||
INIT_MEANS = np.array([0.08, 0.5, 0.92])
|
||||
MEAN_BOUNDS = ((0.0, 0.4), (0.2, 0.8), (0.6, 1.0))
|
||||
SD_BOUNDS = (0.008, 0.15)
|
||||
OUTLIER_W = 0.02
|
||||
MERGE_GAP = 0.08
|
||||
EM_ITERS = 40
|
||||
MIN_FRAC = 0.01
|
||||
MIN_N = 3
|
||||
SCORE_CAP = 0.9
|
||||
|
||||
|
||||
@dataclass
|
||||
class BlockFit:
|
||||
"""Per-locus cluster parameters for a block of L loci; (L, 3) arrays are (AA, AB, BB)."""
|
||||
|
||||
n: npt.NDArray[np.int32]
|
||||
theta_mean: F32
|
||||
theta_dev: F32
|
||||
r_mean: F32
|
||||
r_dev: F32
|
||||
cluster_sep: F32
|
||||
train_score: F32
|
||||
conf: F32 # mean max posterior
|
||||
valid_frac: F32
|
||||
|
||||
|
||||
def fit_block(theta: F32, r: F32, use: npt.NDArray[np.bool_], haploid: npt.NDArray[np.bool_]) -> BlockFit:
|
||||
"""Fit a block: ``theta``/``r`` (L, S) normalised intensities, ``use`` (L, S) samples that inform each locus
|
||||
(sex-aware selection done by the caller), ``haploid`` (L,) loci without a heterozygous cluster."""
|
||||
L, S = theta.shape
|
||||
t = np.where(use & np.isfinite(theta) & np.isfinite(r) & (r > 0), theta.astype(np.float64), np.nan)
|
||||
ok = np.isfinite(t)
|
||||
nv = ok.sum(axis=1).astype(np.float64)
|
||||
tz = np.nan_to_num(t, nan=0.5)
|
||||
m = np.tile(INIT_MEANS, (L, 1))
|
||||
sd = np.full((L, 3), 0.04)
|
||||
w = np.full((L, 3), (1 - OUTLIER_W) / 3)
|
||||
w[haploid, 1] = 0.0
|
||||
w[haploid, 0] = w[haploid, 2] = (1 - OUTLIER_W) / 2
|
||||
lo = np.array([b[0] for b in MEAN_BOUNDS])
|
||||
hi = np.array([b[1] for b in MEAN_BOUNDS])
|
||||
resp = np.zeros((L, S, 3))
|
||||
for _ in range(EM_ITERS):
|
||||
dens = w[:, None, :] * np.exp(-0.5 * ((tz[:, :, None] - m[:, None, :]) / sd[:, None, :]) ** 2) / (
|
||||
sd[:, None, :] * np.sqrt(2 * np.pi)) # fmt: skip
|
||||
tot = dens.sum(axis=2) + OUTLIER_W # uniform outlier density on [0, 1]
|
||||
resp = dens / tot[:, :, None]
|
||||
resp[~ok] = 0.0
|
||||
nk = resp.sum(axis=1)
|
||||
with np.errstate(invalid="ignore", divide="ignore"):
|
||||
mk = (resp * tz[:, :, None]).sum(axis=1) / nk
|
||||
vk = (resp * (tz[:, :, None] - mk[:, None, :]) ** 2).sum(axis=1) / nk
|
||||
m = np.where(nk > 0.5, np.clip(mk, lo, hi), m)
|
||||
sd = np.where(nk > 1.5, np.clip(np.sqrt(np.maximum(vk, 0)), *SD_BOUNDS), sd)
|
||||
w = np.where(nv[:, None] > 0, nk / np.maximum(nv[:, None], 1) * (1 - OUTLIER_W), w)
|
||||
w[haploid, 1] = 0.0
|
||||
# hard assignment (outliers and missing → −1)
|
||||
best = resp.argmax(axis=2)
|
||||
pbest = resp.max(axis=2)
|
||||
assign = np.where(ok & (pbest > 0.5), best, -1)
|
||||
counts = np.stack([(assign == k).sum(axis=1) for k in range(3)], axis=1)
|
||||
keep = counts >= np.maximum(MIN_N, np.ceil(MIN_FRAC * nv))[:, None]
|
||||
# merge adjacent kept clusters that sit too close (the larger keeps the samples)
|
||||
for a_, b_ in ((0, 1), (1, 2), (0, 2)):
|
||||
close = keep[:, a_] & keep[:, b_] & (np.abs(m[:, b_] - m[:, a_]) < MERGE_GAP)
|
||||
if close.any():
|
||||
big_a = counts[:, a_] >= counts[:, b_]
|
||||
src = np.where(big_a, b_, a_)
|
||||
dst = np.where(big_a, a_, b_)
|
||||
for i in np.flatnonzero(close):
|
||||
assign[i, assign[i] == src[i]] = dst[i]
|
||||
keep[i, src[i]] = False
|
||||
counts = np.stack([(assign == k).sum(axis=1) for k in range(3)], axis=1)
|
||||
keep &= counts > 0
|
||||
# a locus with a single cluster is homozygous: an AB-labelled lone cluster moves to AA or BB by position
|
||||
lone_ab = keep[:, 1] & ~keep[:, 0] & ~keep[:, 2]
|
||||
if lone_ab.any():
|
||||
for i in np.flatnonzero(lone_ab).tolist():
|
||||
hom = 0 if m[i, 1] < 0.5 else 2
|
||||
assign[i, assign[i] == 1] = hom
|
||||
keep[i, 1], keep[i, hom] = False, True
|
||||
counts = np.stack([(assign == k).sum(axis=1) for k in range(3)], axis=1)
|
||||
tm = np.full((L, 3), np.nan)
|
||||
td = np.full((L, 3), np.nan)
|
||||
rm = np.full((L, 3), np.nan)
|
||||
rd = np.full((L, 3), np.nan)
|
||||
rr = np.where(ok, r.astype(np.float64), np.nan)
|
||||
for k in range(3):
|
||||
sel = assign == k
|
||||
with np.errstate(invalid="ignore", divide="ignore"):
|
||||
import warnings
|
||||
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("ignore", RuntimeWarning)
|
||||
tm[:, k] = np.nanmean(np.where(sel, t, np.nan), axis=1)
|
||||
td[:, k] = np.nanstd(np.where(sel, t, np.nan), axis=1)
|
||||
rm[:, k] = np.nanmean(np.where(sel, rr, np.nan), axis=1)
|
||||
rd[:, k] = np.nanstd(np.where(sel, rr, np.nan), axis=1)
|
||||
tm[~keep] = np.nan
|
||||
n_out = np.where(keep, counts, 0).astype(np.int32)
|
||||
# separation and training score
|
||||
gaps = []
|
||||
for a_, b_ in ((0, 1), (1, 2), (0, 2)):
|
||||
both = keep[:, a_] & keep[:, b_]
|
||||
if a_ == 0 and b_ == 2:
|
||||
both &= ~keep[:, 1] # AA–BB gap only counts when AB is absent
|
||||
g = np.abs(tm[:, b_] - tm[:, a_]) / (np.maximum(td[:, a_], SD_BOUNDS[0]) + np.maximum(td[:, b_], SD_BOUNDS[0]))
|
||||
gaps.append(np.where(both, g, np.inf))
|
||||
gap_std = np.min(np.stack(gaps, axis=1), axis=1)
|
||||
sep = np.where(np.isfinite(gap_std), np.clip((gap_std - 1.0) / 3.0, 0.0, 1.0), 1.0)
|
||||
with np.errstate(invalid="ignore", divide="ignore"):
|
||||
conf = np.where(nv > 0, np.nansum(np.where(assign >= 0, pbest, 0.0), axis=1) / np.maximum(nv, 1), 0.0)
|
||||
valid_frac = nv / np.maximum(use.sum(axis=1), 1)
|
||||
score = np.clip(conf * (0.6 + 0.4 * sep) * valid_frac, 0.0, SCORE_CAP)
|
||||
score = np.where(n_out.sum(axis=1) > 0, score, 0.0)
|
||||
return BlockFit(n_out, tm.astype(np.float32), td.astype(np.float32), rm.astype(np.float32),
|
||||
rd.astype(np.float32), sep.astype(np.float32), score.astype(np.float32),
|
||||
conf.astype(np.float32), valid_frac.astype(np.float32)) # fmt: skip
|
||||
|
||||
|
||||
@dataclass
|
||||
class Relations:
|
||||
"""Cross-locus relationships of well-clustered loci (step 3): medians, and least-squares lines that predict one
|
||||
cluster's θ from a neighbour's (θ_AB = ab_from_aa[0] + ab_from_aa[1]·θ_AA, …)."""
|
||||
|
||||
d_aa: float # θ_AB − θ_AA
|
||||
d_bb: float # θ_BB − θ_AB
|
||||
r_aa: float # R_AA / R_AB
|
||||
r_bb: float # R_BB / R_AB
|
||||
td: tuple[float, float, float]
|
||||
rcv: tuple[float, float, float] # R SD / R mean
|
||||
ab_from_aa: tuple[float, float] = (0.4, 1.0)
|
||||
ab_from_bb: tuple[float, float] = (-0.4, 1.0)
|
||||
bb_from_ab: tuple[float, float] = (0.4, 1.0)
|
||||
aa_from_ab: tuple[float, float] = (-0.4, 1.0)
|
||||
bb_from_aa: tuple[float, float] = (0.8, 1.0)
|
||||
aa_from_bb: tuple[float, float] = (-0.8, 1.0)
|
||||
|
||||
|
||||
def learn_relations(fit: BlockFit, min_n: int = 10) -> Relations:
|
||||
good = (fit.n >= min_n).all(axis=1)
|
||||
if good.sum() < 20:
|
||||
return Relations(0.4, 0.4, 1.0, 1.0, (0.03, 0.05, 0.03), (0.12, 0.12, 0.12))
|
||||
tm, rm = fit.theta_mean[good].astype(np.float64), fit.r_mean[good].astype(np.float64)
|
||||
td, rd = fit.theta_dev[good].astype(np.float64), fit.r_dev[good].astype(np.float64)
|
||||
|
||||
def line(x: F64, y: F64) -> tuple[float, float]:
|
||||
b, a = np.polyfit(x, y, 1)
|
||||
return float(a), float(b)
|
||||
|
||||
return Relations(float(np.median(tm[:, 1] - tm[:, 0])), float(np.median(tm[:, 2] - tm[:, 1])),
|
||||
float(np.median(rm[:, 0] / rm[:, 1])), float(np.median(rm[:, 2] / rm[:, 1])),
|
||||
tuple(float(np.median(td[:, k])) for k in range(3)), # type: ignore[arg-type]
|
||||
tuple(float(np.median(rd[:, k] / rm[:, k])) for k in range(3)), # type: ignore[arg-type]
|
||||
line(tm[:, 0], tm[:, 1]), line(tm[:, 2], tm[:, 1]), line(tm[:, 1], tm[:, 2]),
|
||||
line(tm[:, 1], tm[:, 0]), line(tm[:, 0], tm[:, 2]), line(tm[:, 2], tm[:, 0])) # fmt: skip
|
||||
|
||||
|
||||
def impute_missing(fit: BlockFit, rel: Relations) -> None:
|
||||
"""Step 3, in place: fill θ/R means and SDs of clusters without samples (their n stays 0)."""
|
||||
tm, rm = fit.theta_mean.astype(np.float64), fit.r_mean.astype(np.float64)
|
||||
td, rd = fit.theta_dev.astype(np.float64), fit.r_dev.astype(np.float64)
|
||||
have = fit.n > 0
|
||||
for i in np.flatnonzero(~have.all(axis=1)).tolist():
|
||||
h = have[i]
|
||||
t, rr = tm[i].copy(), rm[i].copy()
|
||||
if not h.any():
|
||||
t[:] = (0.05, 0.5, 0.95)
|
||||
rr[:] = 1.0
|
||||
else:
|
||||
|
||||
def ln(c: tuple[float, float], x: float) -> float:
|
||||
return c[0] + c[1] * x
|
||||
|
||||
if not h[1]:
|
||||
if h[0] and h[2]:
|
||||
t[1] = (ln(rel.ab_from_aa, t[0]) + ln(rel.ab_from_bb, t[2])) / 2
|
||||
rr[1] = (rr[0] / rel.r_aa + rr[2] / rel.r_bb) / 2
|
||||
elif h[0]:
|
||||
t[1] = ln(rel.ab_from_aa, t[0])
|
||||
rr[1] = rr[0] / rel.r_aa
|
||||
else:
|
||||
t[1] = ln(rel.ab_from_bb, t[2])
|
||||
rr[1] = rr[2] / rel.r_bb
|
||||
if not h[0]:
|
||||
t[0] = ln(rel.aa_from_bb, t[2]) if (h[2] and not h[1]) else ln(rel.aa_from_ab, t[1])
|
||||
rr[0] = rr[1] * rel.r_aa
|
||||
if not h[2]:
|
||||
t[2] = ln(rel.bb_from_aa, t[0]) if (h[0] and not h[1]) else ln(rel.bb_from_ab, t[1])
|
||||
rr[2] = rr[1] * rel.r_bb
|
||||
t = np.clip(t, 0.0, 1.0)
|
||||
t[1] = min(max(t[1], t[0] + 0.02), t[2] - 0.02) if t[2] - t[0] > 0.04 else (t[0] + t[2]) / 2
|
||||
for k in range(3):
|
||||
if not h[k]:
|
||||
tm[i, k], rm[i, k] = t[k], max(rr[k], 1e-3)
|
||||
td[i, k] = rel.td[k]
|
||||
rd[i, k] = rel.rcv[k] * rm[i, k]
|
||||
td = np.where(np.isfinite(td) & (td > 0), td, np.array(rel.td)[None, :])
|
||||
rd = np.where(np.isfinite(rd) & (rd > 0), rd, np.array(rel.rcv)[None, :] * rm)
|
||||
fit.theta_mean, fit.r_mean = tm.astype(np.float32), rm.astype(np.float32)
|
||||
fit.theta_dev = np.maximum(td, SD_BOUNDS[0]).astype(np.float32)
|
||||
fit.r_dev = rd.astype(np.float32)
|
||||
Reference in new issue
Block a user