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

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.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
Claude-Session: https://claude.ai/code/session_01814GKVzZVL3FuF3DeiZXhv

src/simurg/train/clusters.py

@@ -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)