ArrayTrain (plan §12): simurg train — EGT cluster files from a cohort's GTCs

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

Co-Authored-By: Claude Opus 5.5 <[email protected]>
Claude-Session: https://claude.ai/code/session_01814GKVzZVL3FuF3DeiZXhv

tests/unit/test_train.py

@@ -0,0 +1,71 @@
+"""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])