SimurgArray/src/simurg/qc/locus_summary.py

"""Per-locus summary statistics across samples (spec OUT-072, CLI-025; plan §7.9).

[OSS: BeadArrayFiles ``examples/locus_summary.py`` (BSD-2) — formulas: call/genotype frequencies, minor frequency,
k-regularised heterozygote excess with k = (0.013/q)^4, ChiTest_P100 = 2·Φ(±dh; 0, 0.1), GenomeStudio percentile of
the called GenCall scores.] Vectorised over a block of loci; the caller streams blocks for large cohorts.
"""

from __future__ import annotations

import math

import numpy as np
import numpy.typing as npt

COLUMNS = [
    "locus_name",
    "locus_id",
    "n_nocalls",
    "n_calls",
    "call_freq",
    "aa_freq",
    "ab_freq",
    "bb_freq",
    "minor_freq",
    "gentrain_score",
    "gc50",
    "gc10",
    "het_excess",
    "chi_test_p100",
    "cluster_sep",
    "aa_t_mean",
    "aa_t_sd",
    "ab_t_mean",
    "ab_t_sd",
    "bb_t_mean",
    "bb_t_sd",
    "aa_r_mean",
    "aa_r_sd",
    "ab_r_mean",
    "ab_r_sd",
    "bb_r_mean",
    "bb_r_sd",
    "plus_minus_strand",
    "het_excess_simple",
]
HET_EXCESS_EPS = 1e-9  # OQ-030: denominator guard of the simple form
M_SQRT1_2 = 0.70710678118654752440  # C literal (1/sqrt(2) in Python is 1 ULP lower)


def _norm_cdf_sd01(x: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
    """Φ(x; 0, 0.1) with the branch structure of the cephes ``ndtr`` used by the reference."""
    z = (x / 0.1) * M_SQRT1_2
    out = np.empty_like(z)
    for i, v in enumerate(z.tolist()):
        a = abs(v)
        if a < M_SQRT1_2:
            y = 0.5 + 0.5 * math.erf(v)
        else:
            y = 0.5 * math.erfc(a)
            if v > 0:
                y = 1.0 - y
        out[i] = y
    return out


def gs_percentile_columns(
    sorted_scores: npt.NDArray[np.float32], counts: npt.NDArray[np.int64], p: int
) -> npt.NDArray[np.float64]:
    """GenomeStudio percentile of the first ``counts[j]`` values of each column of ``sorted_scores`` (double math)."""
    n = counts.astype(np.int64)
    pos = n * p / 100.0
    idx = pos.astype(np.int64)  # int(num*p/100) truncates like the reference
    idx = np.where(pos - idx < 0.5, idx - 1, idx)
    cols = np.arange(sorted_scores.shape[1])
    last = np.maximum(n - 1, 0)
    lo = np.clip(idx, 0, np.maximum(last - 1, 0))
    y1 = sorted_scores[lo, cols].astype(np.float64)
    y2 = sorted_scores[np.minimum(lo + 1, last), cols].astype(np.float64)
    with np.errstate(invalid="ignore", divide="ignore"):
        x1 = 100 * (lo + 0.5) / n
        x2 = 100 * (lo + 1 + 0.5) / n
        mid = y1 + (y2 - y1) / (x2 - x1) * (p - x1)
    first = sorted_scores[0, cols].astype(np.float64)
    lastv = sorted_scores[last, cols].astype(np.float64)
    out = np.where(idx < 0, first, np.where(idx >= n - 1, lastv, mid))
    return np.where(n == 0, np.nan, out)


def summarize_block(genotypes: npt.NDArray[np.uint8], scores: npt.NDArray[np.float32]) -> dict[str, npt.NDArray]:  # type: ignore[type-arg]
    """Statistics for a block of loci; inputs are (samples × loci)."""
    g = np.asarray(genotypes)
    nc = (g == 0).sum(axis=0).astype(np.int64)
    aa = (g == 1).sum(axis=0).astype(np.int64)
    ab = (g == 2).sum(axis=0).astype(np.int64)
    bb = (g == 3).sum(axis=0).astype(np.int64)
    calls = aa + ab + bb
    with np.errstate(invalid="ignore", divide="ignore"):
        call_freq = np.where(calls + nc > 0, calls / (calls + nc), np.nan)
        cf = calls.astype(np.float64)
        aa_f = np.where(calls > 0, aa / cf, np.nan)
        ab_f = np.where(calls > 0, ab / cf, np.nan)
        bb_f = np.where(calls > 0, bb / cf, np.nan)
        a_freq = np.where(calls > 0, (aa * 2 + ab) / (2 * cf), np.nan)
        q = np.minimum(a_freq, 1.0 - a_freq)
        p = 1 - q
        temp = 0.013 / q
        k = temp * temp * temp * temp
        dh = ((ab / cf + k) / (2 * p * q + k)) - 1
        simple = (ab_f - 2 * p * q) / (2 * p * q + HET_EXCESS_EPS)
    mono = (aa + ab == 0) | (ab + bb == 0)
    regular = (calls > 0) & ~mono
    hw = np.zeros(len(calls))
    het = np.zeros(len(calls))
    if regular.any():
        d = dh[regular]
        cdf = _norm_cdf_sd01(d)
        hw[regular] = np.where(d < 0, 2 * cdf, 2 * (1 - cdf))
        het[regular] = d
    hw[(calls > 0) & mono] = 1.0
    called_scores = np.where(g != 0, scores.astype(np.float32), np.float32(np.inf))
    called_scores.sort(axis=0)
    return {
        "n_nocalls": nc,
        "n_calls": calls,
        "call_freq": call_freq,
        "aa_freq": aa_f,
        "ab_freq": ab_f,
        "bb_freq": bb_f,
        "minor_freq": q,
        "gc50": gs_percentile_columns(called_scores, calls, 50),
        "gc10": gs_percentile_columns(called_scores, calls, 10),
        "het_excess": het,
        "chi_test_p100": hw,
        "het_excess_simple": np.where(calls > 0, simple, np.nan),
    }