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
004739da8e
commit
35d1e08621
9 files changed
+618
No files matched your search
@@ -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)
|
||||
Reference in new issue
Block a user