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
1 parent
adad0efb08
commit
1eb7ea8154
6 files changed
+304
-58
No files matched your search
@@ -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
|
||||
|
||||
@@ -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`). |
|
||||
|
||||
@@ -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:
|
||||
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
@@ -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)
|
||||
|
||||
@@ -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
|
||||
Reference in new issue
Block a user