scripts/tmb_resid_gp.py

"""Blind GP run on MSI-residualized TMB — do gene COMBINATIONS beat
the best single gene?

The residual is the "leftover TMB" you get after taking MSI out: each
tumor compared to its own MSI class's normal (log1p(TMB), then
z-scored within each MSI group; see ``validate/tmb_resid_rank.py``).
The diagnostic already showed no single gene predicts this leftover
strongly on colorectal — best |Spearman| ≈ 0.35, borderline. The
question here is whether the engine can find COMBINATIONS that beat
that single-gene ceiling on held-out patients.

Isolated: ONE new file under ``scripts/`` (allowed to be biology-
aware). READS the existing processed matrix, REUSES the sealed map
via ``airgap.anonymise`` / ``airgap.reveal``, and IMPORTS the residual
straight from ``validate.tmb_resid_rank.build_residual_cohort`` so the
target the engine sees is byte-identical to what the diagnostic
scored. WRITES nothing to disk, adds no dataset, no API route, no UI,
does not touch ``engine_v2/``, ``dsl/``, ``api/``, or ``web/``. Deleting
this file leaves zero trace.

Held-out honesty: the pipeline builds a train/test split from
(M.index, seed, test_size, stratify=False) since TMB_OBJECTIVE is
continuous. The single-gene ceiling is computed on the SAME test rows
via ``validate.tmb_rank._spearman_per_column``, so ``combined`` and
``single_ceiling`` are apples-to-apples on the held-out split.

Airgap: the engine sees only opaque IDs. Symbols are revealed ONCE
per seed at the end, bounded to the winner's opaque IDs — same
discipline ``/evaluate`` uses.

Run:
    python -m scripts.tmb_resid_gp
    python -m scripts.tmb_resid_gp --seeds 1 7 13 --pop 300 --gens 50
"""

from __future__ import annotations

import argparse
import sys
import warnings
from dataclasses import dataclass
from typing import Sequence

import numpy as np
import pandas as pd

# Silence scipy's ConstantInputWarning noise from degenerate program
# evaluations — the engine already floors those to WORST_FITNESS.
warnings.filterwarnings(
    "ignore", message="An input array is constant", category=RuntimeWarning,
)
try:
    from scipy.stats import ConstantInputWarning as _CIW
    warnings.filterwarnings("ignore", category=_CIW)
except Exception:  # noqa: BLE001 — best-effort; older scipy lacks it
    pass

from airgap import anonymise, reveal
from engine.split import make_split
from engine_v2 import run_v2_pipeline
from engine_v2.fitness import TMB_OBJECTIVE
from validate.tmb_rank import _spearman_per_column
from validate.tmb_resid_rank import build_residual_cohort


@dataclass
class SeedResult:
    seed: int
    program_repr: str
    winner_symbols: list[str]
    combined: float
    single_ceiling: float
    synergy: float
    permutation_p: float
    n_test: int


def _test_ids_for(
    M: pd.DataFrame, y: np.ndarray, *, seed: int, test_size: float = 0.3,
) -> pd.Index:
    """Reproduce the pipeline's held-out split byte-identically.
    ``TMB_OBJECTIVE.binary`` is False, so ``make_split`` picks a plain
    random partition — same call the pipeline makes, so the test rows
    match exactly."""
    split = make_split(
        M.index, y, test_size=test_size, random_state=int(seed), stratify=False,
    )
    return split.test_ids


def _single_gene_ceiling_on_test(
    X_named: pd.DataFrame,
    residual: pd.Series,
    test_ids: pd.Index,
) -> tuple[float, str]:
    """Fair single-gene ceiling ON THE SAME held-out test rows.
    Returns (|spearman|_max, best_symbol). Uses the diagnostic's
    ``_spearman_per_column`` verbatim so the metric is byte-identical
    to what the ranking would give if it were run only on the test
    slice."""
    X_test = X_named.loc[test_ids]
    y_test = residual.loc[test_ids].to_numpy()
    corr = _spearman_per_column(X_test, y_test)
    valid = corr.dropna()
    if valid.empty:
        return 0.0, "(none)"
    best_sym = str(valid.abs().idxmax())
    return float(abs(valid.loc[best_sym])), best_sym


def _run_one_seed(
    M: pd.DataFrame,
    residual: pd.Series,
    X_named: pd.DataFrame,
    seed: int,
    *,
    population: int,
    generations: int,
    permutations: int,
) -> SeedResult:
    y = residual.to_numpy()
    _log, result = run_v2_pipeline(
        M, y,
        objective=TMB_OBJECTIVE,
        seed=int(seed),
        test_size=0.3,
        prefilter_n=None,
        population_size=int(population),
        n_generations=int(generations),
        n_permutations=int(permutations),
        tournament_k=2,
        p_mutate=0.85,
        immigrant_fraction=0.10,
        coherence_weight=0.0,
        scalar_share_override=0.0,
    )
    winning = result.get("winning", {}) or {}
    combined = abs(float(winning.get("holdout_score") or 0.0))
    program_repr = str(winning.get("program_repr") or "")
    permutation_p = float(winning.get("permutation_p") or 1.0)
    winner_ids = list(winning.get("gene_ids") or [])
    winner_symbols = reveal(winner_ids) if winner_ids else []

    test_ids = _test_ids_for(M, y, seed=int(seed))
    single_ceiling, _best_sym = _single_gene_ceiling_on_test(
        X_named, residual, test_ids,
    )
    synergy = combined - single_ceiling

    return SeedResult(
        seed=int(seed),
        program_repr=program_repr,
        winner_symbols=winner_symbols,
        combined=combined,
        single_ceiling=single_ceiling,
        synergy=synergy,
        permutation_p=permutation_p,
        n_test=int(len(test_ids)),
    )


def _overall_verdict(rows: list[SeedResult]) -> str:
    """Plain-English overall call. Combines the median across seeds
    with p-value + consistency checks."""
    if not rows:
        return "No seeds completed."
    synergies = [r.synergy for r in rows]
    combineds = [r.combined for r in rows]
    ps = [r.permutation_p for r in rows]
    med_syn = float(np.median(synergies))
    med_combined = float(np.median(combineds))
    n_sig = sum(1 for p in ps if p < 0.05)
    n = len(rows)
    combinatorial = (
        med_syn >= 0.10
        and n_sig >= (n // 2 + 1)
        and med_combined >= 0.30
    )
    only_single = (
        abs(med_syn) < 0.05
        and med_combined < 0.50
    )
    unstable = (
        med_combined < 0.20
        or (max(combineds) - min(combineds)) >= 0.20
    )
    if combinatorial:
        return (
            "REAL COMBINATORIAL SIGNAL. Gene combinations predict the "
            "leftover better than any single gene — the DSL is finding "
            "real synergy on the MSI-residualized TMB target."
        )
    if only_single:
        return (
            "NO SYNERGY. Only the modest single-gene signal is there; the "
            "leftover isn't combinatorial on this cohort."
        )
    if unstable:
        return (
            "THE LEFTOVER IS LARGELY NOISE. Held-out scores are weak or "
            "swing seed-to-seed — nothing reliable to find."
        )
    return (
        "BORDERLINE. Some seeds see synergy, some don't — a bigger budget "
        "or a different objective might sharpen the read."
    )


def main(argv: Sequence[str] | None = None) -> int:
    ap = argparse.ArgumentParser(
        description=(
            "Blind GP run on MSI-residualized TMB. Reads the same "
            "residual validate/tmb_resid_rank.py ranks; reports per-seed "
            "combined held-out |spearman| vs the honest single-gene "
            "ceiling on the SAME test rows."
        ),
    )
    ap.add_argument("--seeds", type=int, nargs="+", default=[1, 7, 13])
    ap.add_argument("--pop", type=int, default=300,
                    help="population_size (default 300; slow but deep).")
    ap.add_argument("--gens", type=int, default=50)
    ap.add_argument("--perms", type=int, default=200)
    args = ap.parse_args(argv or sys.argv[1:])

    print("=" * 78)
    print("Blind GP run on MSI-residualized TMB — synergy check")
    print("=" * 78)
    print(
        "The residual is the 'leftover' each tumor's log(TMB) has after "
        "z-scoring\nwithin its MSI class. We ask the blind engine whether "
        "gene COMBINATIONS\nbeat the best single gene on the SAME held-out "
        "patients."
    )
    print("-" * 78)

    # Load real data via the diagnostic's helper — the residual and the
    # NAMED expression matrix here are byte-identical to what
    # validate/tmb_resid_rank scores.
    X_named, residual, group_stats = build_residual_cohort()
    print(f"Cohort           : {X_named.shape[0]} samples × "
          f"{X_named.shape[1]} genes (NAMED matrix)")
    for g in group_stats:
        print(
            f"  {g.label:>6s}  n={g.n:<4d}  "
            f"TMB̄ {g.tmb_mean:>8.2f}   log1p̄ {g.log1p_mean:+.3f}   "
            f"log1p σ {g.log1p_std:.3f}"
        )
    print(
        f"Target           : within-MSI-group standardised log1p(TMB) "
        f"(median {residual.median():+.3f}, σ {residual.std():.3f})"
    )

    # Anonymise ONCE. The engine sees only opaque IDs; symbols are
    # revealed at the end per seed, bounded to the winner's IDs.
    M = anonymise(X_named)
    print(f"Anonymised M     : {M.shape[0]} samples × {M.shape[1]} opaque cols")
    print()
    print(
        f"Objective        : TMB_OBJECTIVE  (target='tmb', binary=False) — "
        f"used as the\n                   correlation carrier for the "
        f"residual."
    )
    print(
        f"Config           : pop={args.pop}  gens={args.gens}  "
        f"perms={args.perms}  scalar_share=0  prefilter=off  "
        f"coherence=off\n                   diversity ON (k=2, "
        f"p_mutate=0.85, immigrants 10%)"
    )
    print(
        "Airgap           : engine sees only opaque IDs; reveal is "
        "bounded per seed\n                   to the WINNER's IDs at the "
        "end (same discipline as /evaluate)."
    )

    rows: list[SeedResult] = []
    for seed in args.seeds:
        print(f"\n[seed={seed}] running …", flush=True)
        r = _run_one_seed(
            M, residual, X_named, seed,
            population=int(args.pop),
            generations=int(args.gens),
            permutations=int(args.perms),
        )
        rows.append(r)
        symbols_short = ", ".join(r.winner_symbols[:8])
        if len(r.winner_symbols) > 8:
            symbols_short += f", … (+{len(r.winner_symbols) - 8})"
        print(f"  program            : {r.program_repr}")
        print(f"  revealed genes     : [{symbols_short}]")
        print(f"  combined |spearman|: {r.combined:.4f}   "
              f"(held-out; n_test={r.n_test})")
        print(f"  single-gene ceiling: {r.single_ceiling:.4f}   "
              f"(best gene on the SAME test rows)")
        marker = "+" if r.synergy >= 0 else ""
        print(f"  synergy            : {marker}{r.synergy:+.4f}   "
              f"(combined − ceiling)")
        print(f"  permutation p      : {r.permutation_p:.4f}   "
              f"(< 0.05 = beats random)")

    print("\n" + "-" * 78)
    print("RANGES ACROSS SEEDS")
    print("-" * 78)

    def _range(vs: list[float]) -> str:
        if not vs:
            return "—"
        return f"{min(vs):+.4f} … {max(vs):+.4f}   (median {np.median(vs):+.4f})"

    combineds = [r.combined for r in rows]
    ceilings = [r.single_ceiling for r in rows]
    synergies = [r.synergy for r in rows]
    ps = [r.permutation_p for r in rows]
    print(f"  combined |spearman| : {_range(combineds)}")
    print(f"  single-gene ceiling : {_range(ceilings)}")
    print(f"  synergy             : {_range(synergies)}")
    print(f"  permutation p       : {_range(ps)}")

    print("\n" + "=" * 78)
    print(f"OVERALL: {_overall_verdict(rows)}")
    print("=" * 78)
    return 0


if __name__ == "__main__":
    sys.exit(main())