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

src/simurg/pipeline/train.py

@@ -0,0 +1,150 @@
+"""``simurg train`` — ArrayTrain: an EGT cluster file from a cohort's GTCs (plan §12). I/O and orchestration; the
+numerics live in ``simurg.train.clusters``.
+
+Normalised intensities come from each GTC's raw X/Y and normalisation transforms (independent of the cluster file the
+GTC was called with). Sex-aware training: X loci (outside the PAR) from female samples, Y loci from males (haploid),
+MT haploid, PAR and autosomes from everyone; the sex is the GTC's ``gender`` field. Intensity-only loci get a single
+R cluster.
+"""
+
+from __future__ import annotations
+
+import datetime as _dt
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import numpy as np
+import numpy.typing as npt
+
+from simurg import __version__
+from simurg.core.errors import ExitCode, SimurgError
+from simurg.core.log import Reporter
+
+BLOCK = 2000
+
+
+@dataclass
+class TrainOptions:
+    gtcs: Path
+    bpm: Path
+    out: Path
+    min_samples: int = 100
+    max_samples: int = 2000
+    force: bool = False
+    argv: list[str] = field(default_factory=list)
+
+
+def _load(paths: list[Path], lookup: npt.NDArray[np.int32], rep: Reporter) -> tuple[np.ndarray, np.ndarray, list[str]]:
+    from simurg.formats.gtc import read_gtc
+    from simurg.signal.normalization.apply import normalize, polar
+
+    n = len(lookup)
+    theta = np.full((len(paths), n), np.nan, dtype=np.float32)
+    r = np.full((len(paths), n), np.nan, dtype=np.float32)
+    sexes = []
+    for i, p in enumerate(paths):
+        g = read_gtc(p)
+        if len(g.raw_x) != n:
+            raise SimurgError("E_EGT_MANIFEST_MISMATCH", f"{p} has {len(g.raw_x)} loci, the manifest {n}")
+        nx, ny = normalize(g.raw_x, g.raw_y, lookup, g.normalization_transforms)
+        t, rr = polar(nx, ny)
+        zero = (g.raw_x == 0) & (g.raw_y == 0)
+        t[zero] = np.nan
+        rr[zero] = np.nan
+        theta[i], r[i] = t, rr
+        sexes.append(g.gender.decode(errors="replace")[:1].upper() if g.gender else "U")
+        if (i + 1) % 50 == 0:
+            rep.info(f"loaded {i + 1}/{len(paths)} GTC(s)")
+    return theta, r, sexes
+
+
+def run_train(o: TrainOptions, rep: Reporter) -> int:
+    """``simurg train`` (plan §12)."""
+    from simurg.formats.bpm import read_bpm
+    from simurg.formats.discovery import discover_gtcs
+    from simurg.formats.egt import write_egt
+    from simurg.genotype.sex import chromosome_classes
+    from simurg.train.clusters import BlockFit, fit_block, impute_missing, learn_relations
+
+    if o.out.exists() and not o.force:
+        raise SimurgError("E_OUTPUT_EXISTS", f"{o.out} exists", hint="give --force to overwrite")
+    import re
+
+    if not re.fullmatch(r"[A-Za-z0-9_-]+", o.out.stem):
+        rep.warn("W_EGT_NAME", f"{o.out.name}: characters outside [A-Za-z0-9_-] may trouble other tools (IN-045)")
+    man = read_bpm(o.bpm)
+    paths = sorted(discover_gtcs(o.gtcs).values())
+    if len(paths) > o.max_samples:
+        idx = np.random.default_rng(42).choice(len(paths), o.max_samples, replace=False)
+        paths = [paths[i] for i in sorted(idx)]
+    if len(paths) < o.min_samples:
+        raise SimurgError("E_TRAIN_TOO_SMALL", f"{len(paths)} GTC(s); training needs at least {o.min_samples}",
+                          hint="lower --min-samples only for research use")  # fmt: skip
+    theta, r, sexes = _load(paths, man.normalization_lookups(), rep)
+    sx = np.array(sexes)
+    is_x, is_y, _auto = chromosome_classes(list(man.chroms))
+    core = np.array([c.upper().removeprefix("CHR") for c in man.chroms])
+    is_mt = np.isin(core, ["MT", "M"])
+    females, males = sx == "F", sx == "M"
+    rep.info(f"training on {len(paths)} GTC(s): {int(females.sum())} F, {int(males.sum())} M, "
+             f"{int((~females & ~males).sum())} U")  # fmt: skip
+    nl = man.num_loci
+    haploid = is_y | is_mt
+    use_rows = {"all": np.ones(len(paths), bool), "F": females, "M": males}
+    parts: list[BlockFit] = []
+    for s in range(0, nl, BLOCK):
+        e = min(nl, s + BLOCK)
+        t, rr = theta[:, s:e].T, r[:, s:e].T
+        use = np.ones_like(t, dtype=bool)
+        use[is_x[s:e]] = use_rows["F"][None, :] if females.sum() >= 10 else True
+        use[is_y[s:e]] = use_rows["M"][None, :]
+        parts.append(fit_block(t, rr, use, haploid[s:e]))
+        if (s // BLOCK) % 50 == 0:
+            rep.info(f"clustered {e}/{nl} loci")
+    fit = BlockFit(*(np.concatenate([getattr(p, f) for p in parts]) for f in BlockFit.__dataclass_fields__))
+    # missing clusters, per assay type (Infinium II vs I)
+    for inf1 in (False, True):
+        sel = (man.assay_types > 0) == inf1
+        if not sel.any():
+            continue
+        sub = BlockFit(*(getattr(fit, f)[sel] for f in BlockFit.__dataclass_fields__))
+        rel = learn_relations(sub)
+        impute_missing(sub, rel)
+        for f in BlockFit.__dataclass_fields__:
+            getattr(fit, f)[sel] = getattr(sub, f)
+        rep.info(f"Infinium {'I' if inf1 else 'II'}: θ AA→AB {rel.d_aa:.3f}, AB→BB {rel.d_bb:.3f}, "
+                 f"R AA/AB {rel.r_aa:.2f}, BB/AB {rel.r_bb:.2f}")  # fmt: skip
+    # intensity-only loci: one R cluster
+    io = man.intensity_only
+    if io.any():
+        import warnings
+
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", RuntimeWarning)
+            med_r = np.nanmedian(r[:, io], axis=0)
+            sd_r = np.nanstd(r[:, io], axis=0)
+        nv = np.isfinite(r[:, io]).sum(axis=0)
+        fit.n[io] = np.stack([nv, np.zeros_like(nv), np.zeros_like(nv)], axis=1)
+        fit.theta_mean[io] = np.array([0.0, 0.5, 1.0], np.float32)
+        fit.theta_dev[io] = 0.05
+        fit.r_mean[io] = np.nan_to_num(med_r, nan=1.0)[:, None]
+        fit.r_dev[io] = np.nan_to_num(sd_r, nan=0.1)[:, None]
+        fit.train_score[io] = 0.0
+        fit.cluster_sep[io] = 0.0
+    o.out.parent.mkdir(parents=True, exist_ok=True)
+    now = _dt.datetime.now()
+    write_egt(o.out, list(man.names), man.address_a, fit.n, fit.theta_mean, fit.theta_dev, fit.r_mean, fit.r_dev,
+              fit.cluster_sep, fit.train_score, man.manifest_name, cluster_version=f"SimurgArray {__version__}",
+              date_created=f"{now.month}/{now.day}/{now.year} {now:%I:%M:%S %p}")  # fmt: skip
+    ncl = (fit.n > 0).sum(axis=1)[~io]
+    stats = {"samples": len(paths), "loci": nl, "intensity_only": int(io.sum()),
+             "clusters_3": int((ncl == 3).sum()), "clusters_2": int((ncl == 2).sum()),
+             "clusters_1": int((ncl == 1).sum()), "clusters_0": int((ncl == 0).sum()),
+             "median_train_score": float(np.median(fit.train_score[~io])),
+             "median_cluster_sep": float(np.median(fit.cluster_sep[~io]))}  # fmt: skip
+    Path(str(o.out) + ".json").write_text(json.dumps({"software": f"SimurgArray {__version__}",
+                                                      "manifest": man.manifest_name, **stats}, indent=1))  # fmt: skip
+    rep.info("ArrayTrain: " + ", ".join(f"{k} {v:.3f}" if isinstance(v, float) else f"{k} {v}"
+                                        for k, v in stats.items()) + f" → {o.out}")  # fmt: skip
+    return ExitCode.OK