From dd6b6d79e8e72223b1bddd4337454ffad16f7113 Mon Sep 17 00:00:00 2001 From: Kemal Yaylali Date: Tue, 29 Sep 2026 19:53:53 +0100 Subject: [PATCH] =?UTF-8?q?ArrayTrain=20(plan=20=C2=A712):=20simurg=20trai?= =?UTF-8?q?n=20=E2=80=94=20EGT=20cluster=20files=20from=20a=20cohort's=20G?= =?UTF-8?q?TCs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 %. --- docs/messages.md | 1 + scripts/train/concordance_chr.py | 91 ++++++++++++ src/simurg/cli/main.py | 21 +++ src/simurg/core/errors.py | 1 + src/simurg/formats/egt.py | 49 +++++++ src/simurg/pipeline/train.py | 150 ++++++++++++++++++++ src/simurg/train/__init__.py | 1 + src/simurg/train/clusters.py | 233 +++++++++++++++++++++++++++++++ tests/unit/test_train.py | 71 ++++++++++ 9 files changed, 618 insertions(+) create mode 100644 scripts/train/concordance_chr.py create mode 100644 src/simurg/pipeline/train.py create mode 100644 src/simurg/train/__init__.py create mode 100644 src/simurg/train/clusters.py create mode 100644 tests/unit/test_train.py diff --git a/docs/messages.md b/docs/messages.md index 17635b5..6fe0b43 100644 --- a/docs/messages.md +++ b/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 | diff --git a/scripts/train/concordance_chr.py b/scripts/train/concordance_chr.py new file mode 100644 index 0000000..8ca78ee --- /dev/null +++ b/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()) diff --git a/src/simurg/cli/main.py b/src/simurg/cli/main.py index 685af70..868a272 100644 --- a/src/simurg/cli/main.py +++ b/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") diff --git a/src/simurg/core/errors.py b/src/simurg/core/errors.py index dffe020..b57ca79 100644 --- a/src/simurg/core/errors.py +++ b/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"), diff --git a/src/simurg/formats/egt.py b/src/simurg/formats/egt.py index 66352c1..54e4d44 100644 --- a/src/simurg/formats/egt.py +++ b/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), " 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 diff --git a/src/simurg/train/__init__.py b/src/simurg/train/__init__.py new file mode 100644 index 0000000..a985d19 --- /dev/null +++ b/src/simurg/train/__init__.py @@ -0,0 +1 @@ +"""ArrayTrain: cluster-file training (plan §12).""" diff --git a/src/simurg/train/clusters.py b/src/simurg/train/clusters.py new file mode 100644 index 0000000..6a87e62 --- /dev/null +++ b/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) diff --git a/tests/unit/test_train.py b/tests/unit/test_train.py new file mode 100644 index 0000000..4e36977 --- /dev/null +++ b/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])