engine/permutation.py

"""Permutation null distribution via the fast baseline procedure.

Shuffles the TRAIN labels (via the objective's ``permute``), re-runs the
prefilter on the shuffled labels, fits the baseline (top-K) on
shuffled-label TRAIN, and scores it against the TRUE TEST labels with
the same objective. p = fraction of nulls whose score is at least the
observed value — "could a separation this strong arise by chance?".

The matrix is ranked once and reused for every permutation, so the loop
is fast even at n=200 over ~20k features.
"""

from __future__ import annotations

import numpy as np
import pandas as pd

from engine.baseline import BASELINE_K
from engine.objectives import BinaryAUROCObjective, Objective
from engine.prefilter import precompute_ranks


def permutation_null(
    M_train: pd.DataFrame,
    y_train: np.ndarray,
    M_test: pd.DataFrame,
    y_test: np.ndarray,
    *,
    objective: Objective | None = None,
    n_permutations: int = 200,
    seed: int = 0,
    k: int = BASELINE_K,
) -> list[float]:
    obj = objective or BinaryAUROCObjective()
    rng = np.random.default_rng(seed)
    ranks = precompute_ranks(M_train)
    cols = list(M_train.columns)
    null_scores: list[float] = []
    for _ in range(n_permutations):
        y_shuf = obj.permute(y_train, rng)
        per_feature = obj.prefilter_score_per_feature(ranks, y_shuf)
        scores = pd.Series(per_feature, index=cols)
        top_k = scores.nlargest(k).index.tolist()
        states_tr = M_train[top_k].mean(axis=1).values.reshape(-1, 1)
        states_te = M_test[top_k].mean(axis=1).values.reshape(-1, 1)
        null_scores.append(
            float(obj.holdout_score(states_tr, y_shuf, states_te, y_test))
        )
    return null_scores


def permutation_p_value(observed: float, null_dist: list[float]) -> float:
    """Fraction of nulls >= observed, with a +1 / +1 add-one correction."""
    null_arr = np.asarray(null_dist)
    return float(((null_arr >= observed).sum() + 1) / (len(null_arr) + 1))