validate/tmb_rank.py

"""TMB-rank diagnostic — is the MMR signal even findable from expression?

For each gene in the NAMED expression matrix, compute the signed Spearman
correlation with TMB on the SAME cohort the engine optimises against
(``Load("processed")`` filtered to ``tmb.notna() & no-NaN expression
rows`` — exact mirror of ``api._prepare_lab_data`` for ``target="tmb"``).
Rank ascending: rank 1 = most-negative correlation, since the TMB
objective rewards the most-negative association.

Caption for readers: *Where each known gene sits on the engine's TMB
target — rank near 1 = findable; high rank = the signal isn't in the
expression data.*

This module lives in ``validate/`` (allowed gene names) — never in
``engine/``, ``engine_v2/``, or ``dsl/``. The structural airgap test
scans only ``engine/``, so the invariant stays intact.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Optional

import numpy as np
import pandas as pd

from dsl import Cohort, Load
from validate.h1 import IMMUNE_GENES, MMR_GENES


COHORT_NAME = (
    "TCGA CRC PanCancer — Load('processed') filtered to tmb.notna() "
    "& expression-complete rows (mirrors api._prepare_lab_data target='tmb')"
)


@dataclass
class GeneRank:
    symbol: str
    present: bool
    corr: float | None = None
    rank: int | None = None
    percentile: float | None = None


@dataclass
class TMBRankDiagnostic:
    cohort: str
    n_samples: int
    n_genes: int
    mmr_rows: list[GeneRank] = field(default_factory=list)
    immune_rows: list[GeneRank] = field(default_factory=list)
    top_negative: list[GeneRank] = field(default_factory=list)


# ---------------------------------------------------------------------------
# Data prep — same cohort as the engine's TMB run, but kept as a NAMED matrix
# ---------------------------------------------------------------------------


def _tmb_cohort_named(cohort: Cohort | None) -> tuple[pd.DataFrame, np.ndarray]:
    if cohort is None:
        cohort = Load("processed")
    tmb = pd.to_numeric(cohort.labels["tmb"], errors="coerce")
    usable = tmb.notna() & (~cohort.expression.isna().any(axis=1))
    ids = cohort.sample_ids[usable]
    X = cohort.expression.loc[ids]
    y = tmb.loc[ids].astype(float).values
    return X, y


# ---------------------------------------------------------------------------
# Vectorised signed Spearman over all gene columns
# ---------------------------------------------------------------------------


def _spearman_per_column(X: pd.DataFrame, y: np.ndarray) -> pd.Series:
    """Pearson on ranks = Spearman. NaN for zero-variance columns.

    Vectorised over all gene columns; with ~20k genes × ~500 samples the
    naive scipy loop takes minutes; this is sub-second.
    """
    yr = pd.Series(y, index=X.index).rank().to_numpy(dtype=float)
    yr_c = yr - yr.mean()
    yr_norm = np.sqrt((yr_c * yr_c).sum())
    if yr_norm == 0:
        return pd.Series(np.nan, index=X.columns)

    Xr = X.rank(axis=0).to_numpy(dtype=float)
    Xr_c = Xr - Xr.mean(axis=0)
    Xr_norm = np.sqrt((Xr_c * Xr_c).sum(axis=0))

    with np.errstate(divide="ignore", invalid="ignore"):
        corr = (Xr_c * yr_c[:, None]).sum(axis=0) / (Xr_norm * yr_norm)
    corr = np.where(Xr_norm == 0, np.nan, corr)
    return pd.Series(corr, index=X.columns)


# ---------------------------------------------------------------------------
# Diagnostic
# ---------------------------------------------------------------------------


def tmb_rank_diagnostic(
    cohort: Cohort | None = None,
    *,
    mmr: Optional[list[str]] = None,
    immune: Optional[list[str]] = None,
    top_n: int = 10,
) -> TMBRankDiagnostic:
    """Rank every gene by signed Spearman with TMB; report where MMR /
    IMMUNE genes land and the ``top_n`` most-negative genes.

    ``cohort`` defaults to ``Load("processed")`` so the API endpoint can
    just call this without arguments; tests pass a synthetic Cohort.
    """
    mmr_list = MMR_GENES if mmr is None else mmr
    immune_list = IMMUNE_GENES if immune is None else immune

    X, y = _tmb_cohort_named(cohort)
    corr = _spearman_per_column(X, y)
    valid = corr.dropna()
    # Ascending: most-negative = rank 1. method="min" so ties share a rank.
    ranks = valid.rank(method="min", ascending=True).astype(int)
    n_genes = int(len(valid))

    def row_for(symbol: str) -> GeneRank:
        if symbol not in valid.index:
            return GeneRank(symbol=symbol, present=False)
        c = float(valid.loc[symbol])
        r = int(ranks.loc[symbol])
        return GeneRank(
            symbol=symbol,
            present=True,
            corr=c,
            rank=r,
            percentile=float(r) / n_genes if n_genes else None,
        )

    mmr_rows = [row_for(s) for s in mmr_list]
    immune_rows = [row_for(s) for s in immune_list]

    top = valid.sort_values(ascending=True).head(top_n)
    top_rows: list[GeneRank] = []
    for sym in top.index:
        c = float(top.loc[sym])
        r = int(ranks.loc[sym])
        top_rows.append(
            GeneRank(
                symbol=str(sym),
                present=True,
                corr=c,
                rank=r,
                percentile=float(r) / n_genes if n_genes else None,
            )
        )

    return TMBRankDiagnostic(
        cohort=COHORT_NAME,
        n_samples=int(X.shape[0]),
        n_genes=n_genes,
        mmr_rows=mmr_rows,
        immune_rows=immune_rows,
        top_negative=top_rows,
    )


# ---------------------------------------------------------------------------
# CLI: `python -m validate.tmb_rank`
# ---------------------------------------------------------------------------


def _fmt(r: GeneRank, total: int) -> str:
    if not r.present:
        return f"  {r.symbol:>10s}   not present"
    return (
        f"  {r.symbol:>10s}   corr {r.corr:+.4f}   "
        f"rank {r.rank:>5d} / {total}   (pct {r.percentile:.3f})"
    )


def main() -> None:
    d = tmb_rank_diagnostic()
    print(f"TMB-rank diagnostic\n  cohort: {d.cohort}")
    print(f"  N = {d.n_samples} samples, {d.n_genes} genes after filter\n")
    print("MMR genes (DNA spell-checkers — expected near rank 1 if findable):")
    for r in d.mmr_rows:
        print(_fmt(r, d.n_genes))
    print()
    print("IMMUNE genes (expected near the BOTTOM — they rise with TMB):")
    for r in d.immune_rows:
        print(_fmt(r, d.n_genes))
    print()
    print(f"Top {len(d.top_negative)} most-negatively-correlated genes:")
    for r in d.top_negative:
        print(_fmt(r, d.n_genes))


if __name__ == "__main__":
    main()