validate/transfer_gse65858.py

"""Transfer test on the GSE65858 independent HPV cohort.

Score a WINNER's revealed gene symbols on GSE65858 (an INDEPENDENT
HPV validation cohort from GEO; Illumina HumanHT-12 v4 microarray)
and return the orientation-agnostic AUROC + a permutation p-value.

Airgap discipline
-----------------
GSE65858 is NAMED-SIDE ONLY. It lives in `validate/`, is keyed by
real gene symbols, is NEVER anonymised, and MUST NEVER be passed to
`engine`/`engine_v2`. The only thing that crosses from the blind
side to this cohort is the winning program's REVEALED gene symbols
(the bounded reveal `/evaluate` already uses). This module never
sends any gene NAMES back to the blind side.

The function is pure and testable — no network, no engine import.
Tests use a synthetic GSE65858-shaped fixture.
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from pathlib import Path
from typing import Sequence

import numpy as np
import pandas as pd

from data_pipeline import schema


# ---------------------------------------------------------------------------
# Cohort loading (cached per process — the cohort doesn't change between
# calls). Lazy so importing this module is cheap.
# ---------------------------------------------------------------------------

_COHORT_CACHE: dict = {}


@dataclass
class _Cohort:
    expression: pd.DataFrame  # symbols × samples
    hpv: pd.Series            # per-sample, values in {"HPV+", "HPV-"}


def _load_cohort(processed_dir: Path | None = None) -> _Cohort:
    """Load the processed GSE65858 cohort. Uses
    ``schema.GSE65858_PROCESSED_DIR`` by default. Raises FileNotFoundError
    with a clear message if the parquets aren't built yet."""
    pd_dir = Path(processed_dir) if processed_dir else schema.GSE65858_PROCESSED_DIR
    key = str(pd_dir.resolve())
    if key in _COHORT_CACHE:
        return _COHORT_CACHE[key]
    expr_path = pd_dir / "expression.parquet"
    clin_path = pd_dir / "clinical.parquet"
    if not expr_path.exists() or not clin_path.exists():
        raise FileNotFoundError(
            f"GSE65858 processed parquets not found in {pd_dir}. "
            "Run: python -m data_pipeline.download_gse65858 && "
            "python -m data_pipeline.build_gse65858"
        )
    expr = pd.read_parquet(expr_path)
    clin = pd.read_parquet(clin_path)
    if "sample_id" not in clin.columns:
        raise ValueError(
            f"{clin_path.name} missing 'sample_id' column."
        )
    hpv = (
        clin.set_index("sample_id")["hpv_status"]
        .reindex(expr.columns)
    )
    coh = _Cohort(expression=expr, hpv=hpv)
    _COHORT_CACHE[key] = coh
    return coh


# ---------------------------------------------------------------------------
# Orientation-agnostic AUROC via the rank-sum formula. Same shape as
# validate/hpv_rank._auroc_per_column, but for a single 1-D score.
# ---------------------------------------------------------------------------


def _omni_auroc(score: np.ndarray, y: np.ndarray) -> float:
    """orientation-agnostic AUROC of a 1-D score vs a binary label.
    Returns 0.5 for degenerate inputs (no variance, single class).
    """
    if score.size == 0 or y.size == 0 or score.size != y.size:
        return 0.5
    finite = np.isfinite(score) & np.isfinite(y)
    if int(finite.sum()) < 2:
        return 0.5
    s = score[finite].astype(float)
    yy = y[finite].astype(int)
    n_pos = int((yy == 1).sum())
    n_neg = int((yy == 0).sum())
    if n_pos == 0 or n_neg == 0:
        return 0.5
    ranks = pd.Series(s).rank(method="average").to_numpy()
    S_pos = ranks[yy == 1].sum()
    auroc = (S_pos - n_pos * (n_pos + 1) / 2.0) / (n_pos * n_neg)
    omni = max(auroc, 1.0 - auroc)
    if not np.isfinite(omni):
        return 0.5
    return float(omni)


def _finite_or_none(x: float) -> float | None:
    if x is None:
        return None
    try:
        f = float(x)
    except (TypeError, ValueError):
        return None
    return f if math.isfinite(f) else None


# ---------------------------------------------------------------------------
# Public API.
# ---------------------------------------------------------------------------


def transfer_score(
    symbols: Sequence[str],
    *,
    n_permutations: int = 1000,
    seed: int = 0,
    processed_dir: Path | None = None,
) -> dict:
    """Score the winner's gene symbols on GSE65858.

    Aggregation matches the module-ranking "Combined AUROC" (per-patient
    mean of the found genes) BUT standardises each found gene within
    GSE65858 first — the cross-platform fix so TCGA RNA-seq scale vs
    Illumina array intensity doesn't drown the signal.

    Parameters
    ----------
    symbols : Sequence[str]
        Winner's revealed gene symbols. This is the ONLY input from the
        blind side.
    n_permutations : int, default 1000
        How many label shuffles form the null distribution.
    seed : int, default 0
        RNG seed for the permutation null.
    processed_dir : Path, optional
        Override the cohort location (used by tests).

    Returns
    -------
    dict
        ``{
            auroc, p, n, n_pos, n_neg,
            n_found, n_missing,
            found_symbols, missing_symbols,
        }``.
        All numeric fields are finite-guarded (``None`` if
        non-finite / not computable). Symbols are the ONLY gene NAMES
        in the payload — same discipline as ``/evaluate``.
    """
    coh = _load_cohort(processed_dir)
    dedup_symbols = list(dict.fromkeys(str(s) for s in symbols if s))
    found = [s for s in dedup_symbols if s in coh.expression.index]
    missing = [s for s in dedup_symbols if s not in coh.expression.index]

    # Restrict to called-HPV samples with complete expression for the
    # found genes.
    called = coh.hpv.isin([
        schema.GSE65858_HPV_POS_LABEL, schema.GSE65858_HPV_NEG_LABEL,
    ])
    kept_samples = coh.expression.columns[called.fillna(False).values]
    if len(found) == 0 or len(kept_samples) == 0:
        return {
            "auroc": None,
            "p": None,
            "n": int(len(kept_samples)),
            "n_pos": int((coh.hpv.reindex(kept_samples) == schema.GSE65858_HPV_POS_LABEL).sum()),
            "n_neg": int((coh.hpv.reindex(kept_samples) == schema.GSE65858_HPV_NEG_LABEL).sum()),
            "n_found": int(len(found)),
            "n_missing": int(len(missing)),
            "found_symbols": found,
            "missing_symbols": missing,
        }

    sub = coh.expression.loc[found, kept_samples]
    # Drop samples with any NaN across the found genes (rare, defensive).
    keep_mask = ~sub.isna().any(axis=0)
    kept_samples = kept_samples[keep_mask.values]
    sub = sub.loc[:, kept_samples]
    if sub.shape[1] < 2:
        return {
            "auroc": None,
            "p": None,
            "n": int(sub.shape[1]),
            "n_pos": 0,
            "n_neg": 0,
            "n_found": int(len(found)),
            "n_missing": int(len(missing)),
            "found_symbols": found,
            "missing_symbols": missing,
        }

    # Standardise each found gene ACROSS the kept patients (z-score).
    # Cross-platform fix: TCGA RNA-seq and Illumina array intensity live
    # on different absolute scales; z-scoring per gene removes the
    # scale mismatch so the per-patient mean is comparable.
    mu = sub.mean(axis=1).to_numpy()
    sd = sub.std(axis=1, ddof=0).to_numpy()
    sd_safe = np.where(sd > 0, sd, 1.0)
    z = (sub.to_numpy() - mu[:, None]) / sd_safe[:, None]
    # Per-patient mean across the standardised found genes.
    score = z.mean(axis=0)

    y_ser = coh.hpv.reindex(kept_samples)
    y = (y_ser == schema.GSE65858_HPV_POS_LABEL).astype(int).to_numpy()

    n = int(len(y))
    n_pos = int((y == 1).sum())
    n_neg = int((y == 0).sum())
    if n_pos == 0 or n_neg == 0:
        return {
            "auroc": None,
            "p": None,
            "n": n,
            "n_pos": n_pos,
            "n_neg": n_neg,
            "n_found": int(len(found)),
            "n_missing": int(len(missing)),
            "found_symbols": found,
            "missing_symbols": missing,
        }

    auroc = _omni_auroc(score, y)

    # Permutation null: shuffle y, recompute AUROC on the SAME score.
    rng = np.random.default_rng(int(seed))
    hits = 0
    n_perm = max(1, int(n_permutations))
    y_shuf = y.copy()
    for _ in range(n_perm):
        rng.shuffle(y_shuf)
        null_auroc = _omni_auroc(score, y_shuf)
        if null_auroc >= auroc:
            hits += 1
    # Standard permutation-p with a +1 numerator + denominator to avoid
    # p = 0 with finite null samples.
    p = (hits + 1) / (n_perm + 1)

    return {
        "auroc": _finite_or_none(auroc),
        "p": _finite_or_none(p),
        "n": n,
        "n_pos": n_pos,
        "n_neg": n_neg,
        "n_found": int(len(found)),
        "n_missing": int(len(missing)),
        "found_symbols": found,
        "missing_symbols": missing,
    }