engine/prefilter.py

"""Univariate prefilter: narrow ~20k features to top-N for the GP shortlist.

Computed name-blind on the TRAIN split only — the feature labels (opaque
IDs) are never used as anything but row keys; only the numeric scores and
y go in. Pre-ranking once lets the permutation null re-score after every
label shuffle very cheaply.

Dispatches via an Objective so the same prefilter machinery works for
both the binary AUROC and the continuous correlation paths.
"""

from __future__ import annotations

import numpy as np
import pandas as pd

from engine.objectives import BinaryAUROCObjective, Objective


def precompute_ranks(M: pd.DataFrame) -> pd.DataFrame:
    """Rank each column of M (samples × features). Reuse across permutations."""
    return M.rank(axis=0)


def auroc_per_feature(X_ranks: pd.DataFrame, y_bool: np.ndarray) -> np.ndarray:
    """Vectorised per-feature AUROC against a binary y.

    Mann-Whitney U / (n_pos * n_neg) — equivalent to scoring each column as
    a single-feature classifier of y. Kept exported because the permutation
    null in `engine/permutation.py` uses it directly for the binary path.
    """
    y_bool = np.asarray(y_bool, dtype=bool)
    n_pos = int(y_bool.sum())
    n_neg = int(len(y_bool) - n_pos)
    if n_pos == 0 or n_neg == 0:
        return np.full(X_ranks.shape[1], np.nan)
    sum_ranks_pos = X_ranks.iloc[y_bool].sum(axis=0).values
    U = sum_ranks_pos - n_pos * (n_pos + 1) / 2
    return U / (n_pos * n_neg)


def top_n_features(
    M_train: pd.DataFrame,
    y_train: np.ndarray,
    n: int = 2000,
    *,
    objective: Objective | None = None,
) -> tuple[list[str], pd.Series]:
    """Top-N feature IDs by the objective's univariate ranking, descending."""
    obj = objective or BinaryAUROCObjective()
    ranks = precompute_ranks(M_train)
    raw = obj.prefilter_score_per_feature(ranks, y_train)
    scores = pd.Series(raw, index=M_train.columns, name="prefilter_score")
    top = scores.nlargest(n).index.tolist()
    return top, scores