engine/pipeline.py

"""Top-level run: prefilter -> GP -> baseline -> permutation -> final TEST eval.

Strict invariants:
- The engine never sees feature names other than the opaque IDs passed in;
  this module asserts every column matches `^g\\d+$`.
- The TEST split is touched once per program: for the winner and for the
  baseline, after the GP is done. The permutation null uses TEST only as
  the score target with shuffled-label TRAIN.

Two entry points:
- ``run_gp_pipeline`` — the existing one-shot batch path used by
  ``scripts/run_h2.py``. Backwards-compatible (binary AUROC default).
- ``run_gp_pipeline_streaming`` — the live-streaming path used by the
  Lab's FastAPI worker thread; per-generation callback for SSE.
"""

from __future__ import annotations

import re
import time
from typing import Callable

import numpy as np
import pandas as pd

from engine.baseline import BASELINE_K, make_baseline
from engine.fitness import holdout_score
from engine.gp import run_gp
from engine.objectives import BinaryAUROCObjective, Objective
from engine.permutation import permutation_null, permutation_p_value
from engine.prefilter import top_n_features
from engine.split import make_split

_OPAQUE_ID_RE = re.compile(r"^g\d+$")


def _check_opaque_only(M: pd.DataFrame) -> None:
    bad = [c for c in M.columns if not _OPAQUE_ID_RE.match(str(c))]
    if bad:
        raise ValueError(
            "engine: matrix columns must be opaque IDs (^g\\d+$); got "
            f"non-conforming columns e.g. {bad[:5]}"
        )


def _prepare(M: pd.DataFrame, y: np.ndarray, *, seed: int, test_size: float,
             prefilter_n: int | None, objective: Objective,
             ):
    """Returns ``(M, M_train, M_test, gp_pool, baseline_genes, n_input,
    n_dropped, split)``. When ``prefilter_n`` is ``None`` the GP samples
    from the full opaque-ID column set; the baseline still uses the
    univariate top-K so it remains a meaningful sanity check."""
    _check_opaque_only(M)
    n_genes_input = int(M.shape[1])
    nan_cols = M.columns[M.isna().any(axis=0)]
    if len(nan_cols):
        M = M.drop(columns=nan_cols)
    split = make_split(
        M.index, y, test_size=test_size, random_state=seed,
        stratify=objective.binary,
    )
    M_train = M.loc[split.train_ids]
    M_test = M.loc[split.test_ids]
    if prefilter_n is None:
        gp_pool = list(M_train.columns)
        baseline_genes, _ = top_n_features(
            M_train, split.y_train, n=BASELINE_K, objective=objective,
        )
    else:
        shortlist, _ = top_n_features(
            M_train, split.y_train, n=prefilter_n, objective=objective,
        )
        gp_pool = shortlist
        baseline_genes = shortlist[:BASELINE_K]
    return (M, M_train, M_test, gp_pool, baseline_genes,
            n_genes_input, int(len(nan_cols)), split)


def _build_artefacts(
    *,
    M: pd.DataFrame,
    split,
    n_genes_input: int,
    n_genes_dropped_nan: int,
    seed: int,
    test_size: float,
    prefilter_n: int | None,
    cv_folds: int,
    n_permutations: int,
    population_size: int,
    n_generations: int,
    objective: Objective,
    gp_log: list[dict],
    winner,
    winner_cv_fitness: float,
    winner_holdout: float,
    baseline,
    baseline_holdout: float,
    null: list[float],
    p_value: float,
    gp_seconds: float,
) -> tuple[dict, dict]:
    fitness_label = objective.fitness_label()
    evolution_log = {
        "run": {
            "seed": int(seed),
            "objective_spec": objective.to_dict(),
            "fitness_label": fitness_label,
            "params": {
                "population_size": population_size,
                "n_generations": n_generations,
                "test_size": test_size,
                "prefilter_n": prefilter_n,
                "cv_folds": cv_folds,
                "n_permutations": n_permutations,
                "baseline_k": BASELINE_K,
            },
            "n_train": int(len(split.train_ids)),
            "n_test": int(len(split.test_ids)),
            "n_genes": int(M.shape[1]),
            "n_genes_input": n_genes_input,
            "n_genes_dropped_nan": n_genes_dropped_nan,
            "prefilter_N": None if prefilter_n is None else int(prefilter_n),
            "prefilter_note": (
                f"Top-N features by the objective's univariate signal "
                f"({fitness_label}), computed on TRAIN only, name-blind."
                if prefilter_n is not None
                else "Prefilter off: GP samples from the full set of "
                     "opaque feature IDs (name-blind). Baseline still uses "
                     "the univariate top-K for sanity."
            ),
            "gp_seconds": round(gp_seconds, 2),
        },
        "generations": gp_log,
    }
    result = {
        "objective_spec": objective.to_dict(),
        "fitness_label": fitness_label,
        "winning": {
            "id": winner.program_id,
            "gene_ids": list(winner.gene_ids),
            "feature_sets": [list(s) for s in winner.feature_sets],
            "program_repr": winner.program_repr(),
            "cv_fitness": float(winner_cv_fitness),
            "holdout_auroc": float(winner_holdout),
            "holdout_score": float(winner_holdout),
            "permutation_p": float(p_value),
        },
        "baseline": {
            "id": baseline.program_id,
            "gene_ids": list(baseline.gene_ids),
            "feature_sets": [list(s) for s in baseline.feature_sets],
            "program_repr": baseline.program_repr(),
            "holdout_auroc": float(baseline_holdout),
            "holdout_score": float(baseline_holdout),
        },
        "permutation_summary": {
            "n_permutations": n_permutations,
            "null_auroc_mean": float(np.mean(null)),
            "null_score_mean": float(np.mean(null)),
            "null_auroc_p95": float(np.quantile(null, 0.95)),
            "null_score_p95": float(np.quantile(null, 0.95)),
        },
    }
    return evolution_log, result


def run_gp_pipeline(
    M: pd.DataFrame,
    y: np.ndarray,
    *,
    objective: Objective | None = None,
    seed: int = 42,
    test_size: float = 0.3,
    prefilter_n: int | None = 2000,
    population_size: int = 150,
    n_generations: int = 30,
    n_permutations: int = 200,
    cv_folds: int = 5,
) -> tuple[dict, dict]:
    """Full GP pipeline. Inputs are name-blind (opaque feature IDs only).

    ``prefilter_n=None`` skips the prefilter — the GP samples from the full
    column set (still all opaque IDs). The baseline keeps using the
    univariate top-K so it stays an apples-to-apples sanity check.
    """
    obj = objective or BinaryAUROCObjective()
    (M, M_train, M_test, gp_pool, baseline_genes,
     n_genes_input, n_dropped, split) = _prepare(
        M, y, seed=seed, test_size=test_size,
        prefilter_n=prefilter_n, objective=obj,
    )

    t0 = time.time()
    gp_log, winner, winner_cv_fitness = run_gp(
        M_train, split.y_train, gp_pool,
        objective=obj,
        population_size=population_size,
        n_generations=n_generations,
        cv_folds=cv_folds,
        seed=seed,
    )
    gp_seconds = time.time() - t0

    winner_holdout = holdout_score(
        M_train, split.y_train, M_test, split.y_test, winner,
        objective=obj,
    )

    baseline = make_baseline(baseline_genes, k=BASELINE_K)
    baseline_holdout = holdout_score(
        M_train, split.y_train, M_test, split.y_test, baseline,
        objective=obj,
    )

    null = permutation_null(
        M_train, split.y_train, M_test, split.y_test,
        objective=obj,
        n_permutations=n_permutations, seed=seed,
    )
    p_value = permutation_p_value(winner_holdout, null)

    return _build_artefacts(
        M=M, split=split,
        n_genes_input=n_genes_input, n_genes_dropped_nan=n_dropped,
        seed=seed, test_size=test_size, prefilter_n=prefilter_n,
        cv_folds=cv_folds, n_permutations=n_permutations,
        population_size=population_size, n_generations=n_generations,
        objective=obj,
        gp_log=gp_log, winner=winner, winner_cv_fitness=winner_cv_fitness,
        winner_holdout=winner_holdout,
        baseline=baseline, baseline_holdout=baseline_holdout,
        null=null, p_value=p_value, gp_seconds=gp_seconds,
    )


def run_gp_pipeline_streaming(
    M: pd.DataFrame,
    y: np.ndarray,
    *,
    objective: Objective | None = None,
    on_generation: Callable[[dict], None],
    seed: int = 42,
    test_size: float = 0.3,
    prefilter_n: int | None = 2000,
    population_size: int = 150,
    n_generations: int = 30,
    n_permutations: int = 200,
    cv_folds: int = 5,
) -> dict:
    """Same as ``run_gp_pipeline`` but invokes ``on_generation`` per generation.

    Returns the ``result`` dict only — the caller has been receiving every
    generation already via the callback, so the full evolution log is
    rebuilt by the API layer from those events.
    """
    obj = objective or BinaryAUROCObjective()
    (M, M_train, M_test, gp_pool, baseline_genes,
     n_genes_input, n_dropped, split) = _prepare(
        M, y, seed=seed, test_size=test_size,
        prefilter_n=prefilter_n, objective=obj,
    )

    t0 = time.time()
    gp_log, winner, winner_cv_fitness = run_gp(
        M_train, split.y_train, gp_pool,
        objective=obj,
        population_size=population_size,
        n_generations=n_generations,
        cv_folds=cv_folds,
        seed=seed,
        on_generation=on_generation,
    )
    gp_seconds = time.time() - t0

    winner_holdout = holdout_score(
        M_train, split.y_train, M_test, split.y_test, winner,
        objective=obj,
    )

    baseline = make_baseline(baseline_genes, k=BASELINE_K)
    baseline_holdout = holdout_score(
        M_train, split.y_train, M_test, split.y_test, baseline,
        objective=obj,
    )

    null = permutation_null(
        M_train, split.y_train, M_test, split.y_test,
        objective=obj,
        n_permutations=n_permutations, seed=seed,
    )
    p_value = permutation_p_value(winner_holdout, null)

    _evolution_log, result = _build_artefacts(
        M=M, split=split,
        n_genes_input=n_genes_input, n_genes_dropped_nan=n_dropped,
        seed=seed, test_size=test_size, prefilter_n=prefilter_n,
        cv_folds=cv_folds, n_permutations=n_permutations,
        population_size=population_size, n_generations=n_generations,
        objective=obj,
        gp_log=gp_log, winner=winner, winner_cv_fitness=winner_cv_fitness,
        winner_holdout=winner_holdout,
        baseline=baseline, baseline_holdout=baseline_holdout,
        null=null, p_value=p_value, gp_seconds=gp_seconds,
    )
    return result