SimurgArray/src/simurg/pipeline/train.py

"""``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