kemal/SimurgArray
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

ArrayTrain: BIC-selected unbounded θ mixtures + HWE-guided labelling

Browse Source 
Fixed θ windows mislabelled compressed/asymmetric loci (2-cluster loci: 1418 vs
186 errors). Now 1/2/3-component mixtures chosen by BIC, labelled by HWE
goodness of fit plus weak position and gap priors, split clusters merged.
GSE165845 chr22 vs 1000G 30x (240 arrays): concordance 99.44 % (v1 99.27 %,
vendor EGT 99.46 %), NRD 1.30 % (1.89 %, 1.23 %), call rate 99.61 %.
docs/validation/V2_arraytrain.md, K-030; plan §12, CLI-038.
This commit is contained in:
kemal committed 2026-09-29 22:40:19 +01:00
1 parent 2afd296776
commit e7ce242bd7
6 files changed
+304 -58
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/simurg/train/clusters.py
+167 -58
View File
@@ -4,12 +4,18 @@ Explicitly **not** a reproduction of the vendor's clustering [OURS; LIT: Teo et
(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).
1. **Mixtures per locus**, vectorised over a block of loci: unbounded 1-D Gaussian mixtures on θ with 1, 2 and 3
components plus a small uniform outlier component, each fitted by EM from several starts (canonical positions and
θ quantiles; SD floor ``EM_SD_MIN`` keeps components off the θ = 0 / 1 spikes). The number of components is
chosen by BIC among admissible fits (every component ≥ max(3, 1 % of the samples) hard-assigned samples and
adjacent components ≥ ``MIN_SEP`` standardised gaps apart). No fixed θ windows: compressed or asymmetric loci
(e.g. AA 0.05 / AB 0.27 / BB 0.45) are common. Robustness comes from the outlier component (the plan's Student-t
mixture is replaced by it; documented). Haploid loci (male X, Y, MT) take at most two components.
2. **Labelling**: the θ-ordered components map onto AA / AB / BB by the best non-decreasing labelling under the
HWE goodness of fit of the implied genotype counts (with an inbreeding allowance for population structure), weak
priors on each cluster's position (homozygotes towards the ends) and on the θ gaps between genotype groups; two
poorly separated adjacent components may share a genotype (a split cluster). This is what separates AA + AB
from AB + BB at two-cluster loci [OURS; HWE-guided labelling as in LIT: Teo et al. 2007].
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
@@ -22,6 +28,7 @@ technical notes (cluster record fields)]:
from __future__ import annotations
import itertools
from dataclasses import dataclass
import numpy as np
@@ -29,14 +36,19 @@ 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
EM_ITERS = 30
MIN_FRAC = 0.01
MIN_N = 3
EM_SD_MIN = 0.02 # SD floor inside the EM: keeps components off the θ = 0 / 1 spikes of saturated channels
MIN_SEP = 1.5 # standardised θ gap |Δμ| / (sd₁ + sd₂) below which two components are one cluster
HWE_F = 0.1 # inbreeding allowance in the labelling likelihood (population structure)
GAP_SD = 0.12 # spread of the expected cluster θ gaps in the labelling prior
GAP_AA, GAP_BB = 0.42, 0.42 # expected θ_AB − θ_AA and θ_BB − θ_AB
POS_MEAN = (0.03, 0.5, 0.97) # weak prior on where AA / AB / BB sit (homozygotes at the ends)
POS_SD = (0.2, 0.25, 0.2)
MERGE_SEP = 3.5 # adjacent components closer than this (standardised) may be one genotype cluster
SCORE_CAP = 0.9
@@ -55,6 +67,100 @@ class BlockFit:
valid_frac: F32
def _em(tz: F64, ok: npt.NDArray[np.bool_], m: F64, sd: F64) -> tuple[F64, F64, F64, F64]:
"""Unbounded 1-D Gaussian mixture EM (+ uniform outlier) vectorised over loci; ``m``/``sd`` (L, K) starts.
Returns means, SDs, responsibilities (L, S, K) and the log-likelihood (L,)."""
L, K = m.shape
nv = np.maximum(ok.sum(axis=1), 1).astype(np.float64)
w = np.full((L, K), (1 - OUTLIER_W) / K)
for it in range(EM_ITERS + 1):
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
if it == EM_ITERS:
break
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, 0.0, 1.0), m)
sd = np.where(nk > 1.5, np.clip(np.sqrt(np.maximum(vk, 0)), EM_SD_MIN, SD_BOUNDS[1]), sd)
w = nk / nv[:, None] * (1 - OUTLIER_W)
ll = np.where(ok, np.log(tot), 0.0).sum(axis=1)
return m, sd, resp, ll
def _starts(tz: F64, ok: npt.NDArray[np.bool_], k: int) -> list[F64]:
"""EM starts for k components: fixed canonical positions and the samples' θ quantiles."""
L = tz.shape[0]
fixed = {
1: [[0.5]],
2: [[0.1, 0.9], [0.1, 0.5], [0.5, 0.9]],
3: [[0.08, 0.5, 0.92], [0.05, 0.3, 0.6], [0.4, 0.7, 0.95]],
}[k]
out = [np.tile(np.array(f, np.float64), (L, 1)) for f in fixed]
qs = {1: [50], 2: [20, 80], 3: [10, 50, 90]}[k]
with np.errstate(all="ignore"):
import warnings
with warnings.catch_warnings():
warnings.simplefilter("ignore", RuntimeWarning)
q = np.nanpercentile(np.where(ok, tz, np.nan), qs, axis=1).T
out.append(np.where(np.isfinite(q), q, 0.5).reshape(L, k))
return out
def _hwe_gof(c: F64) -> F64:
"""HWE goodness of fit of genotype counts (…, 3): log-likelihood under HWE (with inbreeding allowance
``HWE_F``, allele-counting frequency) minus the saturated multinomial log-likelihood; ≤ 0, ≈ 0 when the counts
fit HWE whatever their number of genotype classes."""
nn = np.maximum(c.sum(axis=-1), 1)
p = (2 * c[..., 2] + c[..., 1] + 0.5) / (2 * nn + 1)
q = 1 - p
f = HWE_F
g = np.stack([q * q + f * p * q, 2 * p * q * (1 - f), p * p + f * p * q], axis=-1)
sat = np.where(c > 0, c * np.log(np.maximum(c, 1) / nn[..., None]), 0.0)
return np.asarray((c * np.log(np.maximum(g, 1e-300)) - sat).sum(axis=-1), np.float64)
def _labelings(k: int) -> npt.NDArray[np.int64]:
"""Non-decreasing maps of k θ-ordered components onto genotypes 0 (AA) / 1 (AB) / 2 (BB)."""
return np.array([c for c in itertools.combinations_with_replacement(range(3), k)], np.int64)
def _label(cnt: F64, m: F64, sd: F64, haploid: npt.NDArray[np.bool_]) -> npt.NDArray[np.int64]:
"""Genotype label of each θ-ordered component (L, k) [OURS]: every non-decreasing labelling is scored by the
HWE goodness of fit of the implied genotype counts, a weak Gaussian prior on each component's position
(``POS_MEAN``/``POS_SD``) and on the θ gaps between adjacent genotype groups (``GAP_AA``/``GAP_BB``). Two
adjacent components may share a genotype (a split cluster) only when they are poorly separated
(standardised gap < ``MERGE_SEP``); haploid loci take no AB and no HWE term."""
L, k = m.shape
labs = _labelings(k)
score = np.full((L, len(labs)), -np.inf)
pos_m, pos_s = np.array(POS_MEAN), np.array(POS_SD)
sep = np.diff(m, axis=1) / (sd[:, 1:] + sd[:, :-1]) if k > 1 else np.zeros((L, 0))
for j, lab in enumerate(labs):
gc = np.stack([np.where(lab == g, cnt, 0).sum(axis=1) for g in range(3)], axis=1)
sc = np.where(haploid, 0.0, _hwe_gof(gc))
sc = sc - 0.5 * (((m - pos_m[lab]) / pos_s[lab]) ** 2).sum(axis=1)
with np.errstate(invalid="ignore", divide="ignore"):
gm = np.stack([np.where(lab == g, cnt * m, 0).sum(axis=1) / gc[:, g] for g in range(3)], axis=1)
present = sorted(set(lab.tolist()))
for g1, g2 in itertools.pairwise(present):
expect = {(0, 1): GAP_AA, (1, 2): GAP_BB, (0, 2): GAP_AA + GAP_BB}[(g1, g2)]
sc = sc - 0.5 * ((gm[:, g2] - gm[:, g1] - expect) / GAP_SD) ** 2
ok = np.ones(L, bool)
for i in range(k - 1):
if lab[i] == lab[i + 1]:
ok &= sep[:, i] < MERGE_SEP
if 1 in present:
ok &= ~haploid
score[:, j] = np.where(ok, sc, -np.inf)
return np.asarray(labs[np.argmax(score, axis=1)], np.int64)
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."""
@@ -63,54 +169,57 @@ def fit_block(theta: F32, r: F32, use: npt.NDArray[np.bool_], haploid: npt.NDArr
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 j in np.flatnonzero(lone_ab).tolist():
hom = 0 if m[j, 1] < 0.5 else 2
assign[j, assign[j] == 1] = hom
keep[j, 1], keep[j, hom] = False, True
counts = np.stack([(assign == k).sum(axis=1) for k in range(3)], axis=1)
min_n = np.maximum(MIN_N, np.ceil(MIN_FRAC * nv))
# 1. mixtures with 1, 2 and 3 components, best start each; BIC among the admissible ones
best_bic = np.full(L, -np.inf)
k_best = np.ones(L, np.int64)
comp_of = np.full((L, S), -1) # hard component index (sorted by θ) of the chosen model
pbest = np.zeros((L, S)) # its max posterior per sample
means = np.full((L, 3), np.nan)
sds = np.full((L, 3), np.nan)
for k in (1, 2, 3):
ll_k = np.full(L, -np.inf)
m_k = np.zeros((L, k))
sd_k = np.zeros((L, k))
a_k = np.full((L, S), -1)
p_k = np.zeros((L, S))
for m0 in _starts(tz, ok, k):
m, sd, resp, ll = _em(tz, ok, m0, np.full((L, k), 0.04))
order = np.argsort(m, axis=1)
m = np.take_along_axis(m, order, axis=1)
sd = np.take_along_axis(sd, order, axis=1)
resp = np.take_along_axis(resp, order[:, None, :], axis=2)
pmax = resp.max(axis=2)
a = np.where(ok & (pmax > 0.5), resp.argmax(axis=2), -1)
better = ll > ll_k
ll_k = np.where(better, ll, ll_k)
m_k[better], sd_k[better], a_k[better], p_k[better] = m[better], sd[better], a[better], pmax[better]
cnt = np.stack([(a_k == j).sum(axis=1) for j in range(k)], axis=1)
valid = (cnt >= min_n[:, None]).all(axis=1)
if k > 1:
sep = np.diff(m_k, axis=1) / (sd_k[:, 1:] + sd_k[:, :-1])
valid &= (sep >= MIN_SEP).all(axis=1)
if k == 3:
valid &= ~haploid
bic = ll_k - 0.5 * (3 * k - 1) * np.log(np.maximum(nv, 2))
take = np.ones(L, bool) if k == 1 else valid & (bic > best_bic)
best_bic = np.where(take, bic, best_bic)
k_best[take] = k
comp_of[take], pbest[take] = a_k[take], p_k[take]
means[take] = np.nan
means[take, :k] = m_k[take]
sds[take] = np.nan
sds[take, :k] = sd_k[take]
# 2. component → genotype labels
lab = np.zeros((L, 3), np.int64) # lab[i, component] = genotype cluster (0 AA, 1 AB, 2 BB)
for k in (1, 2, 3):
sel = k_best == k
if sel.any():
cnt = np.stack([(comp_of[sel] == j).sum(axis=1) for j in range(k)], axis=1).astype(np.float64)
lab[sel, :k] = _label(cnt, means[sel, :k], sds[sel, :k], haploid[sel])
assign = np.where(comp_of >= 0, np.take_along_axis(lab, np.maximum(comp_of, 0), axis=1), -1)
counts = np.stack([(assign == j).sum(axis=1) for j in range(3)], axis=1)
keep = counts > 0
tm = np.full((L, 3), np.nan)
td = np.full((L, 3), np.nan)
rm = np.full((L, 3), np.nan)