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

docs/messages.md

@@ -75,6 +75,7 @@ Generated from `src/simurg/core/errors.py` — do not edit by hand.
 | `E_SHEET_NO_KEY_COLUMNS` | error | Sample sheet lacks SentrixBarcode_A/SentrixPosition_A | IN-002 |
 | `E_SHEET_RESERVED` | error | Sample sheet column collides with an output column | IN-005 |
 | `E_SHEET_TOO_LONG` | error | Sample sheet name or value longer than 500 characters | IN-006 |
+| `E_TRAIN_TOO_SMALL` | error | Too few GTCs to train a cluster file | plan §12 |
 | `I_CN_SEX_RESOLVED` | info | Sample sex updated from U by the caller | CYTO-024 |
 | `I_UNLISTED_IDATS` | info | IDAT pairs not listed in the sample sheet are ignored | IN-016 |
 | `W_CN_NO_FASTA` | warning | No FASTA: VCF REF padding bases written as N | CYTO-029 |

scripts/train/concordance_chr.py

@@ -0,0 +1,91 @@
+#!/usr/bin/env python3
+"""Array genotypes vs 1000 Genomes 30x on one chromosome (ArrayTrain validation, plan §12).
+
+    python scripts/train/concordance_chr.py --vcfs VCF_DIR --identity identity.tsv --truth chr22.vcf.gz \
+        --bcftools BCFTOOLS --chrom 22 --out summary.json [--samples-glob '*.snv.vcf.gz']
+
+Per array sample (identity table: sample_id → individual): biallelic SNVs of the chromosome whose (POS, REF, ALT)
+is a record of the phased 3202-sample panel; genotype concordance and non-reference discordance (NRD = discordant
+/ (discordant + concordant non-ref); plan §16.5) over sites called on the array, plus the array call rate there.
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import subprocess
+from pathlib import Path
+
+import numpy as np
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--vcfs", required=True, type=Path)
+    ap.add_argument("--identity", required=True, type=Path)
+    ap.add_argument("--truth", required=True, type=Path)
+    ap.add_argument("--bcftools", default="bcftools")
+    ap.add_argument("--chrom", default="22")
+    ap.add_argument("--out", required=True, type=Path)
+    a = ap.parse_args()
+    with a.identity.open() as fh:
+        ident = {r["sample_id"]: r["individual"] for r in csv.DictReader((ln for ln in fh if not ln.startswith("#")),
+                                                                           delimiter="\t")}  # fmt: skip
+    vcfs = {p.name.split(".")[0]: p for p in sorted(a.vcfs.rglob("*.snv.vcf.gz"))}
+    vcfs = {k: v for k, v in vcfs.items() if k in ident}
+    inds = sorted({ident[k] for k in vcfs})
+    # array sites (from the first VCF) → truth genotypes of the individuals at those sites
+    first = next(iter(vcfs.values()))
+    q = subprocess.run([a.bcftools, "query", "-r", a.chrom, "-i", 'TYPE="snp" && N_ALT=1', "-f", "%POS\t%REF\t%ALT\n",
+                        str(first)], capture_output=True, text=True, check=True).stdout  # fmt: skip
+    sites = {(int(p), r, al) for p, r, al in (ln.split("\t") for ln in q.splitlines())}
+    reg = a.out.with_suffix(".sites.tsv")
+    chrom_t = f"chr{a.chrom}"
+    reg.write_text("".join(f"{chrom_t}\t{p}\n" for p in sorted({s[0] for s in sites})))
+    tq = subprocess.run([a.bcftools, "query", "-R", str(reg), "-s", ",".join(inds), "--force-samples", "-i",
+                         'TYPE="snp"', "-f", "%POS\t%REF\t%ALT[\t%GT]\n", str(a.truth)], capture_output=True, text=True,
+                        check=True).stdout  # fmt: skip
+    head = subprocess.run([a.bcftools, "query", "-l", "-s", ",".join(inds), "--force-samples", str(a.truth)],
+                          capture_output=True, text=True, check=True).stdout.split()  # fmt: skip
+    truth: dict[tuple[int, str, str], dict[str, int]] = {}
+    for ln in tq.splitlines():
+        f = ln.split("\t")
+        k = (int(f[0]), f[1], f[2])
+        if k in sites:
+            truth[k] = {
+                s: sum(x == "1" for x in g.replace("|", "/").split("/")) for s, g in zip(head, f[3:], strict=True)
+            }
+    per = []
+    for sid, path in sorted(vcfs.items()):
+        ind = ident[sid]
+        aq = subprocess.run([a.bcftools, "query", "-r", a.chrom, "-f", "%POS\t%REF\t%ALT[\t%GT]\n", str(path)],
+                            capture_output=True, text=True, check=True).stdout  # fmt: skip
+        n = called = conc = nrd_bad = nrd_ok = 0
+        for ln in aq.splitlines():
+            p, r, al, gt = ln.split("\t")
+            k = (int(p), r, al)
+            t = truth.get(k, {}).get(ind)
+            if t is None:
+                continue
+            n += 1
+            if "." in gt:
+                continue
+            called += 1
+            d = sum(x == "1" for x in gt.replace("|", "/").split("/"))
+            conc += d == t
+            if d > 0 or t > 0:
+                nrd_ok += d == t
+                nrd_bad += d != t
+        per.append({"sample": sid, "individual": ind, "sites": n, "call_rate": called / n if n else None,
+                    "concordance": conc / called if called else None,
+                    "nrd": nrd_bad / (nrd_bad + nrd_ok) if nrd_bad + nrd_ok else None})  # fmt: skip
+    summ = {k: float(np.mean([p[k] for p in per if p[k] is not None])) for k in ("call_rate", "concordance", "nrd")}
+    a.out.write_text(json.dumps({"chrom": a.chrom, "samples": len(per), "truth_sites": len(truth), **summ,
+                                 "per_sample": per}, indent=1))  # fmt: skip
+    print(json.dumps({"samples": len(per), "truth_sites": len(truth), **summ}))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

src/simurg/cli/main.py

@@ -803,6 +803,27 @@ def cn_annotate_cmd(
     run_guarded(rep, lambda: run_cn_annotate(opts, rep))
 
 
+# --------------------------------------------------------------------------------------------------------
+@app.command("train")
+def train_cmd(
+    gtcs: Annotated[Path, typer.Option("--gtcs", help="Folder searched recursively for training GTCs (required).")],
+    bpm: Annotated[Path, typer.Option("--bpm", help="BPM manifest (required).")],
+    out: Annotated[Path, typer.Option("--out", help="Output cluster file (.egt, required).")],
+    min_samples: Annotated[int, typer.Option("--min-samples", min=1)] = 100,
+    max_samples: Annotated[int, typer.Option("--max-samples", min=1)] = 2000,
+    force: ForceOpt = False,
+    log_json: LogJson = False,
+    debug: Debug = False,
+    quiet: Quiet = False,
+) -> None:
+    """ArrayTrain: train an EGT cluster file from a cohort's GTCs (plan §12)."""
+    from simurg.pipeline.train import TrainOptions, run_train
+
+    rep = _reporter("train", out.parent, log_json, debug, quiet)
+    o = TrainOptions(gtcs=gtcs, bpm=bpm, out=out, min_samples=min_samples, max_samples=max_samples, force=force)
+    run_guarded(rep, lambda: run_train(o, rep))
+
+
 # --------------------------------------------------------------------------------------------------------
 pgx_app = typer.Typer(help="Pharmacogenomics: targeted copy number, star alleles, phenotypes (v2).")
 app.add_typer(pgx_app, name="pgx")

src/simurg/core/errors.py

@@ -177,6 +177,7 @@ CATALOGUE: dict[str, MessageDef] = dict(
         _m("E_PGX_GUIDELINE", "error", "Guideline tables not present in the PGx DB", "PGX-025"),
         _m("W_PGX_NO_PROBE_MAP", "warning", "The PGx DB has no probe map for this manifest", "PGX-035"),
         _m("E_PGX_PLATFORM_UNKNOWN", "error", "PGx CN platform (LCG/EX) unknown: give --csv or --platform", "CLI-032"),
+        _m("E_TRAIN_TOO_SMALL", "error", "Too few GTCs to train a cluster file", "plan §12"),
         _m("E_PGX_SOURCE", "error", "PGx source tables missing; run scripts/pgx/fetch_sources.sh", "PGX-001"),
         # --- methylation control QC (MET-001 .. MET-018) -------------------------------------------------
         _m("E_METH_NO_ASSAY", "error", "Probe section not found", "MET-002"),

src/simurg/formats/egt.py

@@ -233,3 +233,52 @@ def _parse_body(
         trailer_floats=trailer,
         extra_per_record=extra,
     )
+
+
+def write_egt(
+    path: str | Path,
+    names: list[str],
+    addresses: npt.NDArray[np.int32],
+    n: npt.NDArray[np.int32],
+    theta_mean: npt.NDArray[np.float32],
+    theta_dev: npt.NDArray[np.float32],
+    r_mean: npt.NDArray[np.float32],
+    r_dev: npt.NDArray[np.float32],
+    cluster_sep: npt.NDArray[np.float32],
+    total_score: npt.NDArray[np.float32],
+    manifest_name: str,
+    *,
+    cluster_version: str = "",
+    date_created: str = "",
+    normalization_version: str = "1.2.0",
+) -> None:
+    """Write an EGT v3, data-block version 9 file (plan §12 step 5): 15 extra floats per record (the first is the
+    intensity threshold, written as 0), per-record scores, genotypes, names, addresses, counts and the dbv-9 trailer.
+    Readable by this reader, BeadArrayFiles and gtc2vcf."""
+    from simurg.formats.binutil import BinaryWriter
+
+    k = len(names)
+    w = BinaryWriter()
+    w.i32(3)
+    for s in ("ArrayTrain", cluster_version, cluster_version, normalization_version, date_created):
+        w.string(s)
+    w.u8(1)
+    w.string(manifest_name)
+    w.i32(9)
+    w.string("")
+    w.i32(k)
+    rec = np.zeros(k, dtype=_record_dtype(15))
+    rec["n"], rec["r_dev"], rec["r_mean"] = n, r_dev, r_mean
+    rec["theta_dev"], rec["theta_mean"] = theta_dev, theta_mean
+    w.write(rec.tobytes())
+    sc = np.zeros(k, dtype=_SCORE_DTYPE)
+    sc["cluster_sep"], sc["total_score"], sc["original_score"] = cluster_sep, total_score, total_score
+    w.write(sc.tobytes())
+    for _ in range(k):
+        w.string("")
+    for nm in names:
+        w.string(nm)
+    w.array(np.asarray(addresses, np.int32), "<i4")
+    w.array(np.asarray(n, np.int32).ravel(), "<i4")
+    w.array(np.zeros(k, np.float32), "<f4")
+    Path(path).write_bytes(w.getvalue())

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

src/simurg/train/__init__.py

@@ -0,0 +1 @@
+"""ArrayTrain: cluster-file training (plan §12)."""

src/simurg/train/clusters.py

@@ -0,0 +1,233 @@
+"""ArrayTrain: per-locus genotype cluster training in (θ, R) — plan §12. Pure functions, no I/O.
+
+Explicitly **not** a reproduction of the vendor's clustering [OURS; LIT: Teo et al. 2007 Bioinformatics 23:2741
+(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).
+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
+   cluster is homozygous (AA or BB by its θ).
+4. **Scores**: ``cluster_sep`` = the smallest standardised θ gap between adjacent present clusters,
+   gap / (sd₁ + sd₂), squashed to 0–1 by (gap_std − 1) / 3; ``train_score`` (the EGT GenTrain field, which scales
+   every call's GenCall score) = mean posterior of the assigned cluster × (0.6 + 0.4 × cluster_sep) × fraction of
+   samples with intensity, capped at 0.9 [OURS].
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+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
+MIN_FRAC = 0.01
+MIN_N = 3
+SCORE_CAP = 0.9
+
+
+@dataclass
+class BlockFit:
+    """Per-locus cluster parameters for a block of L loci; (L, 3) arrays are (AA, AB, BB)."""
+
+    n: npt.NDArray[np.int32]
+    theta_mean: F32
+    theta_dev: F32
+    r_mean: F32
+    r_dev: F32
+    cluster_sep: F32
+    train_score: F32
+    conf: F32  # mean max posterior
+    valid_frac: F32
+
+
+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."""
+    L, S = theta.shape
+    t = np.where(use & np.isfinite(theta) & np.isfinite(r) & (r > 0), theta.astype(np.float64), np.nan)
+    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 i in np.flatnonzero(lone_ab).tolist():
+            hom = 0 if m[i, 1] < 0.5 else 2
+            assign[i, assign[i] == 1] = hom
+            keep[i, 1], keep[i, hom] = False, True
+        counts = np.stack([(assign == k).sum(axis=1) for k in range(3)], axis=1)
+    tm = np.full((L, 3), np.nan)
+    td = np.full((L, 3), np.nan)
+    rm = np.full((L, 3), np.nan)
+    rd = np.full((L, 3), np.nan)
+    rr = np.where(ok, r.astype(np.float64), np.nan)
+    for k in range(3):
+        sel = assign == k
+        with np.errstate(invalid="ignore", divide="ignore"):
+            import warnings
+
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore", RuntimeWarning)
+                tm[:, k] = np.nanmean(np.where(sel, t, np.nan), axis=1)
+                td[:, k] = np.nanstd(np.where(sel, t, np.nan), axis=1)
+                rm[:, k] = np.nanmean(np.where(sel, rr, np.nan), axis=1)
+                rd[:, k] = np.nanstd(np.where(sel, rr, np.nan), axis=1)
+    tm[~keep] = np.nan
+    n_out = np.where(keep, counts, 0).astype(np.int32)
+    # separation and training score
+    gaps = []
+    for a_, b_ in ((0, 1), (1, 2), (0, 2)):
+        both = keep[:, a_] & keep[:, b_]
+        if a_ == 0 and b_ == 2:
+            both &= ~keep[:, 1]  # AA–BB gap only counts when AB is absent
+        g = np.abs(tm[:, b_] - tm[:, a_]) / (np.maximum(td[:, a_], SD_BOUNDS[0]) + np.maximum(td[:, b_], SD_BOUNDS[0]))
+        gaps.append(np.where(both, g, np.inf))
+    gap_std = np.min(np.stack(gaps, axis=1), axis=1)
+    sep = np.where(np.isfinite(gap_std), np.clip((gap_std - 1.0) / 3.0, 0.0, 1.0), 1.0)
+    with np.errstate(invalid="ignore", divide="ignore"):
+        conf = np.where(nv > 0, np.nansum(np.where(assign >= 0, pbest, 0.0), axis=1) / np.maximum(nv, 1), 0.0)
+        valid_frac = nv / np.maximum(use.sum(axis=1), 1)
+    score = np.clip(conf * (0.6 + 0.4 * sep) * valid_frac, 0.0, SCORE_CAP)
+    score = np.where(n_out.sum(axis=1) > 0, score, 0.0)
+    return BlockFit(n_out, tm.astype(np.float32), td.astype(np.float32), rm.astype(np.float32),
+                    rd.astype(np.float32), sep.astype(np.float32), score.astype(np.float32),
+                    conf.astype(np.float32), valid_frac.astype(np.float32))  # fmt: skip
+
+
+@dataclass
+class Relations:
+    """Cross-locus relationships of well-clustered loci (step 3): medians, and least-squares lines that predict one
+    cluster's θ from a neighbour's (θ_AB = ab_from_aa[0] + ab_from_aa[1]·θ_AA, …)."""
+
+    d_aa: float  # θ_AB − θ_AA
+    d_bb: float  # θ_BB − θ_AB
+    r_aa: float  # R_AA / R_AB
+    r_bb: float  # R_BB / R_AB
+    td: tuple[float, float, float]
+    rcv: tuple[float, float, float]  # R SD / R mean
+    ab_from_aa: tuple[float, float] = (0.4, 1.0)
+    ab_from_bb: tuple[float, float] = (-0.4, 1.0)
+    bb_from_ab: tuple[float, float] = (0.4, 1.0)
+    aa_from_ab: tuple[float, float] = (-0.4, 1.0)
+    bb_from_aa: tuple[float, float] = (0.8, 1.0)
+    aa_from_bb: tuple[float, float] = (-0.8, 1.0)
+
+
+def learn_relations(fit: BlockFit, min_n: int = 10) -> Relations:
+    good = (fit.n >= min_n).all(axis=1)
+    if good.sum() < 20:
+        return Relations(0.4, 0.4, 1.0, 1.0, (0.03, 0.05, 0.03), (0.12, 0.12, 0.12))
+    tm, rm = fit.theta_mean[good].astype(np.float64), fit.r_mean[good].astype(np.float64)
+    td, rd = fit.theta_dev[good].astype(np.float64), fit.r_dev[good].astype(np.float64)
+
+    def line(x: F64, y: F64) -> tuple[float, float]:
+        b, a = np.polyfit(x, y, 1)
+        return float(a), float(b)
+
+    return Relations(float(np.median(tm[:, 1] - tm[:, 0])), float(np.median(tm[:, 2] - tm[:, 1])),
+                     float(np.median(rm[:, 0] / rm[:, 1])), float(np.median(rm[:, 2] / rm[:, 1])),
+                     tuple(float(np.median(td[:, k])) for k in range(3)),  # type: ignore[arg-type]
+                     tuple(float(np.median(rd[:, k] / rm[:, k])) for k in range(3)),  # type: ignore[arg-type]
+                     line(tm[:, 0], tm[:, 1]), line(tm[:, 2], tm[:, 1]), line(tm[:, 1], tm[:, 2]),
+                     line(tm[:, 1], tm[:, 0]), line(tm[:, 0], tm[:, 2]), line(tm[:, 2], tm[:, 0]))  # fmt: skip
+
+
+def impute_missing(fit: BlockFit, rel: Relations) -> None:
+    """Step 3, in place: fill θ/R means and SDs of clusters without samples (their n stays 0)."""
+    tm, rm = fit.theta_mean.astype(np.float64), fit.r_mean.astype(np.float64)
+    td, rd = fit.theta_dev.astype(np.float64), fit.r_dev.astype(np.float64)
+    have = fit.n > 0
+    for i in np.flatnonzero(~have.all(axis=1)).tolist():
+        h = have[i]
+        t, rr = tm[i].copy(), rm[i].copy()
+        if not h.any():
+            t[:] = (0.05, 0.5, 0.95)
+            rr[:] = 1.0
+        else:
+
+            def ln(c: tuple[float, float], x: float) -> float:
+                return c[0] + c[1] * x
+
+            if not h[1]:
+                if h[0] and h[2]:
+                    t[1] = (ln(rel.ab_from_aa, t[0]) + ln(rel.ab_from_bb, t[2])) / 2
+                    rr[1] = (rr[0] / rel.r_aa + rr[2] / rel.r_bb) / 2
+                elif h[0]:
+                    t[1] = ln(rel.ab_from_aa, t[0])
+                    rr[1] = rr[0] / rel.r_aa
+                else:
+                    t[1] = ln(rel.ab_from_bb, t[2])
+                    rr[1] = rr[2] / rel.r_bb
+            if not h[0]:
+                t[0] = ln(rel.aa_from_bb, t[2]) if (h[2] and not h[1]) else ln(rel.aa_from_ab, t[1])
+                rr[0] = rr[1] * rel.r_aa
+            if not h[2]:
+                t[2] = ln(rel.bb_from_aa, t[0]) if (h[0] and not h[1]) else ln(rel.bb_from_ab, t[1])
+                rr[2] = rr[1] * rel.r_bb
+        t = np.clip(t, 0.0, 1.0)
+        t[1] = min(max(t[1], t[0] + 0.02), t[2] - 0.02) if t[2] - t[0] > 0.04 else (t[0] + t[2]) / 2
+        for k in range(3):
+            if not h[k]:
+                tm[i, k], rm[i, k] = t[k], max(rr[k], 1e-3)
+                td[i, k] = rel.td[k]
+                rd[i, k] = rel.rcv[k] * rm[i, k]
+    td = np.where(np.isfinite(td) & (td > 0), td, np.array(rel.td)[None, :])
+    rd = np.where(np.isfinite(rd) & (rd > 0), rd, np.array(rel.rcv)[None, :] * rm)
+    fit.theta_mean, fit.r_mean = tm.astype(np.float32), rm.astype(np.float32)
+    fit.theta_dev = np.maximum(td, SD_BOUNDS[0]).astype(np.float32)
+    fit.r_dev = rd.astype(np.float32)

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])