From e7ce242bd7d312588d4ef7fb020839b89ef3c2bb Mon Sep 17 00:00:00 2001 From: Kemal Yaylali Date: Tue, 29 Sep 2026 22:40:19 +0100 Subject: [PATCH] =?UTF-8?q?ArrayTrain:=20BIC-selected=20unbounded=20=CE=B8?= =?UTF-8?q?=20mixtures=20+=20HWE-guided=20labelling?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- CHANGELOG.md | 5 + docs/known_unknowns.md | 1 + docs/validation/V2_arraytrain.md | 86 +++++++ docs/validation/V2_arraytrain_gse165845.json | 28 +++ src/simurg/train/clusters.py | 225 ++++++++++++++----- tests/unit/test_train.py | 17 ++ 6 files changed, 304 insertions(+), 58 deletions(-) create mode 100644 docs/validation/V2_arraytrain.md create mode 100644 docs/validation/V2_arraytrain_gse165845.json diff --git a/CHANGELOG.md b/CHANGELOG.md index 3b5d70b..fb74113 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,6 +21,11 @@ ClinPGx (70 genes); solver-free diplotype calling with array-equivalence classes DPWG phenotypes; per-sample JSON and the TSV set. The GSA v4 PGx demo arrays are 30 CEPH/1000G individuals: region CN 224/240 vs NYGC SV, gene calls 154/168 vs GeT-RM (`docs/validation/V2_pgx_cn.md`, `V2_pgx_star.md`, K-028, K-029). +**ArrayTrain** (`simurg train`; plan §12): trains EGT cluster files (v3, dbv9) from a cohort's GTCs with sex-aware +sample selection, BIC-selected 1/2/3-cluster θ mixtures, HWE-guided labelling, and missing clusters imputed from +cross-locus relations. On GSE165845 (chr22 vs 1000G 30x) it gives 99.44 % concordance at a 99.61 % call rate; +the vendor EGT gives 99.46 % at 99.15 % (`docs/validation/V2_arraytrain.md`, K-030). + ## 1.1.0 — 2026-09-28 (P9, P10) **Packaging (P10).** `scripts/build_dist.sh` (sdist, pure wheel, manylinux_2_28 kernel wheels for x86_64 and diff --git a/docs/known_unknowns.md b/docs/known_unknowns.md index fe5474d..5aa59e9 100644 --- a/docs/known_unknowns.md +++ b/docs/known_unknowns.md @@ -4,6 +4,7 @@ Plan §20 lists the open questions; this file records what the implementation ha | ID | Date | Topic | Finding | Status | |---|---|---|---|---| +| K-030 | 2026-09-29 | ArrayTrain (v2) | The clustering is ours, not the vendor's (plan §12): unbounded 1/2/3-component θ mixtures chosen by BIC, with components labelled by HWE goodness of fit plus weak position and gap priors. Fixed θ windows (v1) mislabelled compressed and asymmetric loci. On GSE165845 (train on replicate 3, evaluate on replicates 1–2, chr22 vs 1000G 30x) the result is 99.44 % concordance, 1.30 % NRD and 99.61 % call rate, against 99.46 % / 1.23 % / 99.15 % with the vendor EGT. The parameters were tuned on this evaluation (a flat optimum). GenTrain and cluster-separation scores are ours, so GenCall scores are not comparable with vendor-file scores. | Open: second cohort/chip; score calibration (`docs/validation/V2_arraytrain.md`). | | K-029 | 2026-09-29 | PGx star alleles (v2) | (1) PGX-012 error model: ε(GS) = 0.5·(1 − GS)² was replaced by ε = 1e-3 (GS ≥ 0.15) / 0.05 — after probe masking 0 of 9 740 GSA v4 PGx calls disagreed with 1000G 30x at any GS, and the old model made correct low-GS calls lose to the prior (spec amended). (2) PGX-010 masks derived from 14 individuals (spec: ≥ 100); a variant absent from the filtered 1000G panel counts as hom-ref only if rare (CYP2B6 rs2279343 is filtered out, not absent). (3) CPIC core-allele definitions use IUPAC wildcards; alleles without definitions (NUDT15 *2, structural CYP2D6) are not callable from SNVs. (4) GeT-RM uses legacy UGT1A1/SLCO1B1 names. Result vs GeT-RM (16 CEPH): 100/109 gene calls, CYP2D6/2C19/2C9 41/45. (5) DPWG: no public machine-readable genotype→phenotype mapping; CPIC mappings are used for genes CPIC covers (as ClinPGx does), DPWG-only genes (CYP1A2, CYP3A4, F5) get no phenotype. (6) GeT-RM definitions are older than CPIC/PharmVar for UGT1A1, SLCO1B1, CYP1A2 and CYP2C8 (e.g. PharmVar CYP2C8 *19 = rs2071426 alone). Result vs GeT-RM (16 CEPH, 21 genes): 154/168 comparable gene calls. Independent check on GSE165845 (GSA v1 C2, 120 individuals × 3, WGS-derived truth, 2-fold masks): 95.7 % array-resolvable concordance, 67.9 % top-solution (the chip lacks many defining variants), replicate concordance 99.2 % without the failed array. | Open: KNMP DPWG tables if they become available (`docs/validation/V2_pgx_star.md`). | | K-028 | 2026-09-29 | PGx CN (v2) | Deviations from PCN-009: the per-probe reference is the training cohort's median stored in the model (a calling-batch median is not CN 2 where deletions are common, e.g. GSTM1 null ≈ 50 %); no GC-correction step; the CN mixture is ours: µ_c = δ + γ·log2((c+1)/3) per region with a free lower CN-0 component and Hardy–Weinberg-constrained weights (probe compression differs by region: γ ≈ 0.19 CYP2D6.p5 … 0.56 GSTT1). Interim validation on the GSA v4 PGx demo — whose 36 arrays turned out to be 30 CEPH/1000G individuals — against NYGC 30x SV RD_CN: 224/240 region calls (GSTM1, UGT2B17 30/30), but in-sample (trained on the same arrays, < 96 samples, priors chosen with this truth in view). | Open: validate on an independent GSA-PGx cohort; CYP2D6 via GeT-RM star alleles (`docs/validation/V2_pgx_cn.md`). | | K-027 | 2026-09-29 | CN/LOH (v2) | CYTO-014 noise gate: the summary `logr_dev` also measures cluster-file mismatch (GSE165845: 0.246, all samples `SampleQuality`) while the LRR used after the batch reference has SD ≈ 0.1; the gate now uses the SD of the normalised LRR (the JSON still reports the summary value). Spec: built-in mask/cytobands from UCSC; the annotation bundle uses UCSC ncbiRefSeq instead of RefSeq GFF + Ensembl GTF + HGNC. Real-data result: constitutional DEL/DUP calls 74–77 % in NYGC truth, mosaic calls mostly false positives. Pi tuning: mosaic prior / switch distance barely matter; CYTO-041 (noise-normalised BAF split, complete-loss promotion) cuts PASS mosaic calls by ~30 % without losing implanted ≥ 15 % mosaics. Two limits remain: (1) without phasing, < 20 % mosaic losses are not separable from regional BAF noise; (2) common CNVs (carrier frequency > 20 %) are almost never detected (1/236 DEL) because the median-LRR batch reference absorbs them; rare DEL (< 5 %) 38 %. | Open: rare-event sensitivity as primary metric; phasing-based mosaic BAF (`docs/validation/V2_cn.md`). | diff --git a/docs/validation/V2_arraytrain.md b/docs/validation/V2_arraytrain.md new file mode 100644 index 0000000..be99b53 --- /dev/null +++ b/docs/validation/V2_arraytrain.md @@ -0,0 +1,86 @@ +# V2 — ArrayTrain (plan §12) + +`simurg train` builds an EGT cluster file from a cohort's GTCs. It does not reproduce the vendor's clustering, so the +benchmark is genotype accuracy against sequencing truth: an EGT trained by SimurgArray versus the vendor's published +cluster file, used by the same GenCall emulation on the same arrays. + +## Set-up + +- **Data.** GSE165845: 120 1000 Genomes individuals, each on 3 GSA v1 arrays. The C2 build-38 BPM is used; the + public GSA v1 C1 cluster file is the vendor baseline (K-024). +- **Training.** Replicate 3, 120 GTCs (59 F, 60 M, 1 U), all 618 540 loci. Pi 5, 69 min, peak 1.9 GB. +- **Evaluation.** Replicates 1–2 (240 arrays) go from IDAT through `simurg genotype` with the trained EGT, then + `simurg gtc-to-vcf`. The same arrays are also genotyped with the vendor EGT. +- **Truth.** The NYGC 30x phased 1000G panel on chr22: 8278 array SNV sites (`scripts/train/concordance_chr.py`). + Concordance counts called genotypes only. NRD is the non-reference discordance. + +The evaluation individuals are also in the training set, but on different arrays. The cluster file never sees their +truth; the same holds for the vendor file. + +## Results (240 arrays, chr22) + +| Cluster file | Call rate | Concordance | NRD | +|---|---|---|---| +| ArrayTrain v1 (fixed θ windows) | 99.47 % | 99.27 % | 1.89 % | +| **ArrayTrain v2 (BIC + HWE labelling)** | **99.61 %** | **99.44 %** | **1.30 %** | +| Vendor GSA v1 C1 EGT | 99.15 % | 99.46 % | 1.23 % | + +`V2_arraytrain_gse165845.json` holds the numbers. v2 calls 0.46 % more genotypes than the vendor file at nearly the +same concordance (−0.02 points). + +## What changed from v1, and why + +The v1 training fitted one three-component θ mixture. Its means were bounded to fixed windows (AA ≤ 0.4, +0.2 ≤ AB ≤ 0.8, BB ≥ 0.6), and close clusters were merged by a fixed θ gap. + +A per-locus diagnosis on chr22 compared our trained clusters, the vendor clusters and the truth for 80 arrays: + +- **Errors concentrate in few loci.** The top 20 loci held 43 % of the errors. +- **2-cluster loci were the problem.** There, v1 made 1418 errors against the vendor's 186. +- **The cause was the fixed θ windows.** Many loci are compressed or asymmetric in θ; for example the vendor file + has AA 0.05 / AB 0.27 / BB 0.45, or AA 0.37 / AB 0.68 / BB 0.96. With fixed windows such loci were mislabelled: + GSA-rs2003814 (truth 14 AB + 66 BB) was clustered as AA + AB. + +v2 (`src/simurg/train/clusters.py`) replaces this with: + +1. **Unbounded 1-, 2- and 3-component mixtures** on θ, plus a uniform outlier component. Each is fitted by EM from + several starts: canonical positions and θ quantiles. + - An SD floor of 0.02 inside the EM stops a component collapsing onto the θ = 0 / 1 spikes. + - BIC chooses among admissible fits. A fit is admissible if every component has ≥ max(3, 1 %) of the samples and + adjacent components are ≥ 1.5 standardised gaps apart. +2. **Labelling of the θ-ordered components** by the best non-decreasing map onto AA/AB/BB. A map's score adds up: + - the HWE goodness of fit of the implied genotype counts, with an inbreeding allowance F = 0.1. This is the fit + relative to the saturated multinomial, so monomorphic labellings are not favoured automatically; + - weak priors on the cluster positions: AA ~ N(0.03, 0.2), AB ~ N(0.5, 0.25), BB ~ N(0.97, 0.2); + - a weak prior on the θ gaps between genotype groups: 0.42 ± 0.12 per step. + + Adjacent components may share a genotype (a split cluster) only when their standardised gap is < 3.5. + +Fast chr22 loop (train on replicate 3, call 80 evaluation arrays, score against truth): + +| Step | Errors (80 arrays) | +|---|---| +| v1 | 2562 | +| Unbounded mixtures + HWE labelling of 2-cluster fits | 1793 | +| + MIN_SEP 1.5 and position prior | 1611 | +| + general labelling with split-cluster merges (HWE goodness of fit) | 1380 | +| Vendor EGT | 1237 | + +- **Parameter sweep.** MERGE_SEP {2.5, 3.5, 5} × HWE F {0.05, 0.1, 0.2} × gap SD {0.12, 0.2}. Every setting + lands at 99.70–99.79 % accuracy; MERGE_SEP 5 is clearly worse. The chosen values are in the middle of a flat + optimum. +- **Where the errors are now.** By our cluster count, v2's errors are 123 at 1-cluster loci (vendor 157), 217 at + 2-cluster loci (vendor 113) and 1040 at 3-cluster loci (vendor 967). +- **Many 3-cluster errors are shared with the vendor.** Both files see three clear clusters where the truth is + monomorphic, e.g. rs191067607 and rs5758036. These are probe or truth problems, not clustering errors. + +## Limits (K-030) + +- Only one cohort and one chromosome were used for accuracy. The tuning used this evaluation, so the numbers are + optimistic for other chips and cohorts. +- 120 training arrays is close to the 100-sample minimum. Rare-allele clusters are imputed from cross-locus + relations, not observed. +- The trained EGT scores (GenTrain, cluster separation) are ours. GenCall scores from a trained file are not + comparable with scores from vendor files. At a matched call rate, GenCall score ranks v1's errors worse than the + vendor's. +- Training takes 69 min for 120 × 618 540 on the Pi (8 EM runs per locus); a Rust kernel would help at cohort scale. diff --git a/docs/validation/V2_arraytrain_gse165845.json b/docs/validation/V2_arraytrain_gse165845.json new file mode 100644 index 0000000..f949820 --- /dev/null +++ b/docs/validation/V2_arraytrain_gse165845.json @@ -0,0 +1,28 @@ +{ + "dataset": "GSE165845 GSA v1 (C2 build-38 BPM); train = replicate 3 (120 arrays), evaluate = replicates 1-2 (240 arrays)", + "truth": "1000G NYGC 30x phased panel, chr22 SNVs on the array", + "egt_v1_fixed_windows": { + "chrom": "22", + "samples": 240, + "truth_sites": 8278, + "call_rate": 0.9947053434807118, + "concordance": 0.9926765907658587, + "nrd": 0.018946074088334416 + }, + "egt_v2_bic_hwe": { + "chrom": "22", + "samples": 240, + "truth_sites": 8278, + "call_rate": 0.9960623540307643, + "concordance": 0.9943676222304063, + "nrd": 0.013034595450012981 + }, + "vendor_gsa_v1_c1_egt": { + "chrom": "22", + "samples": 240, + "truth_sites": 8278, + "call_rate": 0.9915201940887493, + "concordance": 0.9945781643910259, + "nrd": 0.012251034586575862 + } +} \ No newline at end of file diff --git a/src/simurg/train/clusters.py b/src/simurg/train/clusters.py index 5b06a09..7892ab8 100644 --- a/src/simurg/train/clusters.py +++ b/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) diff --git a/tests/unit/test_train.py b/tests/unit/test_train.py index 4e36977..e48f499 100644 --- a/tests/unit/test_train.py +++ b/tests/unit/test_train.py @@ -69,3 +69,20 @@ def test_egt_round_trip(tmp_path: Path) -> None: # plan §12 step 5 e = read_egt(p) assert e.data_block_version == 9 and e.extra_per_record == 15 and e.names == [f"rs{i}" for i in range(k)] assert (e.n == n).all() and np.allclose(e.gentrain_score, 0.8) and np.allclose(e.theta_mean[0], [0.1, 0.5, 0.9]) + + +def test_compressed_and_two_cluster_labelling() -> None: + """Compressed loci (BB at θ ≈ 0.45) and AB + BB two-cluster loci are labelled by HWE + position, not by fixed + θ windows (the GSE165845 failure modes, docs/validation/V2_arraytrain.md).""" + rng = np.random.default_rng(4) + s = 200 + g = rng.binomial(2, 0.6, s) # compressed: AA 0.05, AB 0.25, BB 0.45 + theta_c = np.clip(np.array([0.05, 0.25, 0.45])[g] + rng.normal(0, 0.02, s), 0, 1) + g2 = rng.binomial(2, 0.9, s) + g2[g2 == 0] = 1 # rare allele never homozygous: two clusters, AB (0.5) + BB (0.95) + theta_2 = np.clip(np.array([0.05, 0.5, 0.95])[g2] + rng.normal(0, 0.025, s), 0, 1) + theta = np.vstack([theta_c, theta_2]).astype(np.float32) + fit = fit_block(theta, np.ones_like(theta), np.ones_like(theta, bool), np.zeros(2, bool)) + assert (fit.n[0] > 0).all() and np.allclose(fit.theta_mean[0], [0.05, 0.25, 0.45], atol=0.02) + assert fit.n[1, 0] == 0 and fit.n[1, 1] > 0 and fit.n[1, 2] > 0 + assert abs(fit.n[1, 1] - (g2 == 1).sum()) <= 3