kemal
7cfb0fd962
- simurg/train/clusters.py: vectorised per-locus θ mixture (AA/AB/BB + outlier component, θ-ordered bounds), cluster selection and merging, lone clusters homozygous, missing clusters imputed by cross-locus regression per assay type, cluster_sep and train (GenTrain-field) scores. - formats/egt.py: write_egt (EGT v3, data block 9, round-trips through the reader). - pipeline/train.py: normalised θ/R from GTC raw X/Y + transforms, sex-aware (X females, Y males, MT haploid), intensity-only loci as one R cluster. - scripts/train/concordance_chr.py: array vs 1000G 30x on one chromosome. Demo (36 GSA v4 PGx arrays, chr22 vs 1000G): trained 99.29 % / NRD 2.6 % vs vendor EGT 99.57 % / 1.55 %.
72 lines
3.2 KiB
Python
72 lines
3.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])
|