"""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()