engine/fitness.py

"""Cross-validated fitness, plus held-out evaluation.

Dispatches via an Objective so the same GP machinery can target either a
binary classification (MSI-H vs MSS, AUROC) or a continuous regression
(TMB, signed Spearman). Existing single-objective callers keep working
via the thin ``cv_auroc`` / ``holdout_auroc`` wrappers below.
"""

from __future__ import annotations

import numpy as np
import pandas as pd
from sklearn.linear_model import LogisticRegression

from engine.objectives import BinaryAUROCObjective, Objective
from engine.program import Program


def program_state(M: pd.DataFrame, program: Program) -> np.ndarray:
    """Compute the program's per-patient score(s). Returns (n_samples, n_sets)."""
    cols = [M[fs].mean(axis=1).values for fs in program.feature_sets]
    return np.column_stack(cols)


def cv_score(
    M_train: pd.DataFrame,
    y_train: np.ndarray,
    program: Program,
    *,
    objective: Objective,
    n_folds: int = 5,
    random_state: int = 0,
) -> float:
    states = program_state(M_train, program)
    return objective.cv_score(
        states, y_train, n_folds=n_folds, random_state=random_state,
    )


def fitness_fn(
    M_train: pd.DataFrame,
    y_train: np.ndarray,
    program: Program,
    *,
    objective: Objective | None = None,
    lambda_size: float = 0.005,
    n_folds: int = 5,
    random_state: int = 0,
) -> float:
    """CV score with a small linear penalty on the number of genes used."""
    obj = objective or BinaryAUROCObjective()
    score = cv_score(
        M_train, y_train, program,
        objective=obj, n_folds=n_folds, random_state=random_state,
    )
    return score - lambda_size * program.n_genes


def holdout_score(
    M_train: pd.DataFrame,
    y_train: np.ndarray,
    M_test: pd.DataFrame,
    y_test: np.ndarray,
    program: Program,
    *,
    objective: Objective | None = None,
) -> float:
    obj = objective or BinaryAUROCObjective()
    states_tr = program_state(M_train, program)
    states_te = program_state(M_test, program)
    return obj.holdout_score(states_tr, y_train, states_te, y_test)


# --- back-compat wrappers ---------------------------------------------------

def cv_auroc(
    M_train: pd.DataFrame,
    y_train: np.ndarray,
    program: Program,
    *,
    n_folds: int = 5,
    random_state: int = 0,
) -> float:
    """Legacy: 5-fold CV AUROC for binary y. Used by the H2 MSI run + tests."""
    return cv_score(
        M_train, y_train, program,
        objective=BinaryAUROCObjective(),
        n_folds=n_folds, random_state=random_state,
    )


def holdout_auroc(
    M_train: pd.DataFrame,
    y_train: np.ndarray,
    M_test: pd.DataFrame,
    y_test: np.ndarray,
    program: Program,
) -> tuple[float, LogisticRegression]:
    """Legacy: refit LR on full TRAIN, score AUROC on TEST. Returns (auroc, model)."""
    states_tr = program_state(M_train, program)
    states_te = program_state(M_test, program)
    lr = LogisticRegression(max_iter=1000)
    lr.fit(states_tr, y_train)
    from sklearn.metrics import roc_auc_score
    proba = lr.predict_proba(states_te)[:, 1]
    return float(roc_auc_score(y_test, proba)), lr