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.
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

No files matched your search

+5
View File
@@ -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
+1
View File
@@ -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`). |
+86
View File
@@ -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.
@@ -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
}
}
+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)
+17
View File
@@ -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