validate/hpv_rank.py

"""HPV-marker-rank diagnostic — does CDKN2A separate HPV+/HPV− on its own?

Mirrors ``validate/tmb_rank.py``: for the HNSC HPV+/HPV− cohort the
engine optimises against (``Load(schema.HNSC_PROCESSED_DIR)`` filtered
to ``hpv_status ∈ {"HPV+","HPV−"}`` AND complete expression — exact
mirror of ``api._prepare_lab_data`` for ``dataset="hnsc", target="hpv"``),
rank EVERY gene by its single-gene orientation-agnostic AUROC vs the
HPV label, computed on the engine's TRAIN split (so the diagnostic
never reads test data, same discipline the engine itself uses). Rank
descending: rank 1 = best single-gene separator.

Caption for readers: *Where each known HPV marker sits as a
single-gene HPV+/HPV− separator — rank near 1 = recoverable; high
rank = out-competed.*

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 data_pipeline import schema
from dsl import Cohort, Load
from engine.split import make_split


# Verbatim mirror of api/app.py::HPV_P16_GENES + HPV_CELL_CYCLE_GENES.
# Duplicated here (rather than imported from api/) so validate/ stays
# independent of the API server.
HPV_P16_GENES: list[str] = ["CDKN2A"]
HPV_CELL_CYCLE_GENES: list[str] = [
    "MCM2", "MCM3", "MCM4", "MCM5", "MCM6", "MCM7",
    "PCNA", "CDK1", "CCNE1", "CCNB1", "CDC6", "CDC20",
    "MKI67", "TOP2A", "RRM2", "TYMS", "FOXM1", "E2F1",
    "BUB1", "AURKB",
]

DEFAULT_SEED = 42
DEFAULT_TEST_SIZE = 0.3
COHORT_NAME = (
    "TCGA HNSC PanCancer — Load(HNSC_PROCESSED_DIR) filtered to "
    "hpv_status ∈ {HPV+, HPV-} & expression-complete rows; ranks "
    "computed on the engine's TRAIN split (mirrors "
    "api._prepare_lab_data dataset='hnsc' target='hpv')"
)


@dataclass
class GeneRank:
    symbol: str
    present: bool
    corr: float | None = None        # carries the omni-AUROC value
    rank: int | None = None
    percentile: float | None = None


@dataclass
class HPVRankDiagnostic:
    cohort: str
    seed: int
    n_samples: int        # TRAIN cohort size after the split
    n_pos: int            # HPV+ count in TRAIN
    n_neg: int            # HPV- count in TRAIN
    n_genes: int          # ranked (i.e. non-constant) genes
    p16_rows: list[GeneRank] = field(default_factory=list)
    cell_cycle_rows: list[GeneRank] = field(default_factory=list)
    top_separators: list[GeneRank] = field(default_factory=list)


# ---------------------------------------------------------------------------
# Cohort prep — same cohort + split the engine uses
# ---------------------------------------------------------------------------


def _hpv_cohort_named(cohort: Cohort | None) -> tuple[pd.DataFrame, np.ndarray]:
    if cohort is None:
        cohort = Load(schema.HNSC_PROCESSED_DIR)
    hpv_status = cohort.labels.get("hpv_status")
    if hpv_status is None:
        raise ValueError(
            "HNSC processed cohort is missing the hpv_status label. "
            "Re-run: python -m data_pipeline.build_hnsc"
        )
    usable = hpv_status.isin(["HPV+", "HPV-"]) & (
        ~cohort.expression.isna().any(axis=1)
    )
    ids = cohort.sample_ids[usable]
    X = cohort.expression.loc[ids]
    y = (hpv_status.reindex(ids) == "HPV+").astype(int).to_numpy()
    return X, y


def _train_slice(
    X: pd.DataFrame,
    y: np.ndarray,
    *,
    seed: int,
    test_size: float,
) -> tuple[pd.DataFrame, np.ndarray]:
    """Cut the engine's TRAIN side so the diagnostic stays on the same
    rows the GP actually fits against."""
    split = make_split(
        X.index, y, test_size=test_size, random_state=seed, stratify=True,
    )
    X_train = X.loc[split.train_ids]
    return X_train, np.asarray(split.y_train, dtype=int)


# ---------------------------------------------------------------------------
# Vectorised single-gene orientation-agnostic AUROC
# ---------------------------------------------------------------------------


def _auroc_per_column(X: pd.DataFrame, y: np.ndarray) -> pd.Series:
    """Single-gene orientation-agnostic AUROC for every column.

    Rank-sum / Mann-Whitney formula:
        AUROC = (S_pos − n_pos·(n_pos+1)/2) / (n_pos·n_neg)
    where S_pos is the sum of column ranks among the positive class.
    Returned value is ``max(AUROC, 1 − AUROC)`` so a gene that goes
    DOWN in HPV+ is rewarded the same as one that goes UP — matches
    the engine's ``V2Objective.score_vector`` for HPV. Zero-variance
    columns return NaN and are dropped by the caller.

    Vectorised over all gene columns; ~20k genes × ~340 train samples
    runs sub-second.
    """
    n = X.shape[0]
    if n == 0:
        return pd.Series(np.nan, index=X.columns)
    n_pos = int((y == 1).sum())
    n_neg = int((y == 0).sum())
    if n_pos == 0 or n_neg == 0:
        return pd.Series(np.nan, index=X.columns)

    R = X.rank(axis=0).to_numpy(dtype=float)
    pos_mask = (y == 1)
    S_pos = R[pos_mask].sum(axis=0)
    auroc = (S_pos - n_pos * (n_pos + 1) / 2.0) / (n_pos * n_neg)

    # Zero-variance columns: rank-sum on a constant column gives a
    # degenerate AUROC of 0.5 anyway, but flag them as NaN so they
    # don't pollute the ranking with a flat sea of "0.5"s.
    std = X.std(axis=0, ddof=0).to_numpy(dtype=float)
    omni = np.maximum(auroc, 1.0 - auroc)
    omni = np.where(std == 0, np.nan, omni)
    return pd.Series(omni, index=X.columns)


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


def hpv_rank_diagnostic(
    cohort: Cohort | None = None,
    *,
    p16: Optional[list[str]] = None,
    cell_cycle: Optional[list[str]] = None,
    top_n: int = 10,
    seed: int = DEFAULT_SEED,
    test_size: float = DEFAULT_TEST_SIZE,
) -> HPVRankDiagnostic:
    """Rank every gene by single-gene HPV+/HPV− omni-AUROC on the
    engine's TRAIN split; report where p16 (CDKN2A) and each cell-cycle
    gene land plus the ``top_n`` strongest single-gene separators.

    ``cohort`` defaults to ``Load(HNSC_PROCESSED_DIR)`` so the API
    endpoint can just call this without arguments; tests pass a
    synthetic Cohort.
    """
    p16_list = HPV_P16_GENES if p16 is None else p16
    cc_list = HPV_CELL_CYCLE_GENES if cell_cycle is None else cell_cycle

    X, y = _hpv_cohort_named(cohort)
    X_train, y_train = _train_slice(X, y, seed=seed, test_size=test_size)
    aurocs = _auroc_per_column(X_train, y_train)
    valid = aurocs.dropna()
    # Descending: highest omni-AUROC = rank 1. method="min" so ties share.
    ranks = valid.rank(method="min", ascending=False).astype(int)
    n_genes = int(len(valid))
    n_pos = int((y_train == 1).sum())
    n_neg = int((y_train == 0).sum())

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

    p16_rows = [row_for(s) for s in p16_list]
    cell_cycle_rows = [row_for(s) for s in cc_list]

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

    return HPVRankDiagnostic(
        cohort=COHORT_NAME,
        seed=int(seed),
        n_samples=int(X_train.shape[0]),
        n_pos=n_pos,
        n_neg=n_neg,
        n_genes=n_genes,
        p16_rows=p16_rows,
        cell_cycle_rows=cell_cycle_rows,
        top_separators=top_rows,
    )


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


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


def main() -> None:
    d = hpv_rank_diagnostic()
    print(f"HPV-rank diagnostic\n  cohort: {d.cohort}")
    print(
        f"  N (TRAIN) = {d.n_samples} samples   "
        f"(HPV+ {d.n_pos} / HPV- {d.n_neg})   "
        f"{d.n_genes} genes ranked   seed={d.seed}\n"
    )
    print("p16 (CDKN2A) — the canonical HPV+ surrogate marker:")
    for r in d.p16_rows:
        print(_fmt(r, d.n_genes))
    print()
    print("Cell-cycle / E2F-target genes (HPV-E7 releases this program):")
    for r in d.cell_cycle_rows:
        print(_fmt(r, d.n_genes))
    print()
    print(f"Top {len(d.top_separators)} single-gene HPV+/HPV- separators:")
    for r in d.top_separators:
        print(_fmt(r, d.n_genes))


if __name__ == "__main__":
    main()