kemal
7c57eb5586
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.
89 lines
4.2 KiB
Python
89 lines
4.2 KiB
Python
"""ArrayTrain cluster training (plan §12)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
from simurg.formats.egt import read_egt, write_egt
|
|
from simurg.train.clusters import fit_block, impute_missing, learn_relations
|
|
|
|
|
|
def _block(rng: np.random.Generator, n_loci: int, s: int, freqs: np.ndarray) -> tuple[np.ndarray, ...]:
|
|
centres = np.array([0.05, 0.5, 0.95])
|
|
theta = np.empty((n_loci, s), np.float32)
|
|
r = np.empty((n_loci, s), np.float32)
|
|
truth = np.empty((n_loci, s), np.int64)
|
|
for i in range(n_loci):
|
|
p = freqs[i]
|
|
g = rng.binomial(2, p, s)
|
|
truth[i] = g
|
|
theta[i] = np.clip(centres[g] + rng.normal(0, 0.025, s), 0, 1)
|
|
r[i] = np.where(g == 1, 1.0, 0.85) + rng.normal(0, 0.08, s)
|
|
return theta, r, truth
|
|
|
|
|
|
def test_fit_three_clusters() -> None:
|
|
rng = np.random.default_rng(1)
|
|
theta, r, _ = _block(rng, 40, 200, np.full(40, 0.4))
|
|
fit = fit_block(theta, r, np.ones_like(theta, bool), np.zeros(40, bool))
|
|
assert (fit.n > 0).all()
|
|
assert np.allclose(fit.theta_mean.mean(axis=0), [0.05, 0.5, 0.95], atol=0.02)
|
|
assert np.all(fit.n.sum(axis=1) >= 190) # nearly every sample assigned
|
|
assert fit.train_score.min() > 0.7 and fit.cluster_sep.min() > 0.9
|
|
|
|
|
|
def test_monomorphic_and_haploid() -> None:
|
|
rng = np.random.default_rng(2)
|
|
theta = np.clip(0.3 + rng.normal(0, 0.02, (2, 150)), 0, 1).astype(np.float32) # one cluster at θ 0.3
|
|
r = np.ones_like(theta)
|
|
fit = fit_block(theta, r, np.ones_like(theta, bool), np.zeros(2, bool))
|
|
assert (fit.n[:, 1] == 0).all() and (fit.n[:, 0] > 0).all() # a lone cluster is homozygous (AA by θ < 0.5)
|
|
theta, r, _ = _block(rng, 3, 100, np.full(3, 0.5))
|
|
theta = np.where(np.abs(theta - 0.5) < 0.2, 0.95, theta).astype(np.float32) # haploid: no hets
|
|
fit = fit_block(theta, r, np.ones_like(theta, bool), np.ones(3, bool))
|
|
assert (fit.n[:, 1] == 0).all()
|
|
|
|
|
|
def test_impute_missing_clusters() -> None:
|
|
rng = np.random.default_rng(3)
|
|
theta, r, _ = _block(rng, 60, 200, np.full(60, 0.45))
|
|
rare, r2, _ = _block(rng, 5, 200, np.full(5, 0.0)) # only AA observed
|
|
fit = fit_block(np.vstack([theta, rare]), np.vstack([r, r2]), np.ones((65, 200), bool), np.zeros(65, bool))
|
|
rel = learn_relations(fit)
|
|
impute_missing(fit, rel)
|
|
assert (fit.n[60:, 1:] == 0).all()
|
|
assert np.all(np.abs(fit.theta_mean[60:, 1] - 0.5) < 0.08) and np.all(fit.theta_mean[60:, 2] > 0.85)
|
|
assert np.isfinite(fit.r_mean).all() and (fit.theta_dev > 0).all()
|
|
|
|
|
|
def test_egt_round_trip(tmp_path: Path) -> None: # plan §12 step 5
|
|
k = 4
|
|
n = np.array([[10, 5, 0]] * k, np.int32)
|
|
p = tmp_path / "t.egt"
|
|
write_egt(p, [f"rs{i}" for i in range(k)], np.arange(k, dtype=np.int32), n,
|
|
np.tile(np.float32([0.1, 0.5, 0.9]), (k, 1)), np.full((k, 3), 0.03, np.float32),
|
|
np.ones((k, 3), np.float32), np.full((k, 3), 0.1, np.float32), np.full(k, 0.5, np.float32),
|
|
np.full(k, 0.8, np.float32), "M.bpm") # fmt: skip
|
|
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
|