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 %.
This commit is contained in:
kemal committed
1 parent
5340b52ab6
commit
dd6b6d79e8
9 files changed
+618
No files matched your search
@@ -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 |
|
||||
|
||||
@@ -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())
|
||||
@@ -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")
|
||||
|
||||
@@ -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"),
|
||||
|
||||
@@ -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())
|
||||
@@ -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
|
||||
@@ -0,0 +1 @@
|
||||
"""ArrayTrain: cluster-file training (plan §12)."""
|
||||
@@ -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)
|
||||
@@ -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])
|
||||
Reference in new issue
Block a user