engine_v2/permutation.py

"""Permutation nulls for engine_v2.

- ``permutation_null`` (MSI / TMB): hold the winning program FIXED,
  shuffle the target N times, re-evaluate against each shuffled target.
- ``unsup_random_null`` (unsupervised): there's no target to shuffle, so
  the null is "random Vector-only programs scored on the same held-out
  context". p = fraction of random programs with structure ≥ winner.
"""

from __future__ import annotations

import random
from typing import Sequence

import numpy as np
import pandas as pd

from engine_v2.fitness import V2Objective, evaluate_holdout
from engine_v2.nodes import ExecContext, Node
from engine_v2.types import TType


def permutation_null(
    program: Node,
    ctx: ExecContext,
    y: np.ndarray,
    *,
    objective: V2Objective,
    n_permutations: int = 200,
    seed: int = 0,
) -> list[float]:
    rng = np.random.default_rng(seed)
    nulls: list[float] = []
    label_key = objective.target  # "msi" or "tmb"
    for _ in range(n_permutations):
        y_shuf = rng.permutation(y)
        shuf_labels = {**ctx.labels, label_key: y_shuf}
        shuf_ctx = ExecContext(M=ctx.M, clinical=ctx.clinical, labels=shuf_labels)
        nulls.append(evaluate_holdout(program, shuf_ctx, y_shuf, objective=objective))
    return nulls


def unsup_random_null(
    ctx_test: ExecContext,
    pool: Sequence[str],
    *,
    objective: V2Objective,
    n_permutations: int = 200,
    rates: dict | None = None,
    max_depth: int = 4,
    max_genes_per_set: int = 8,
    seed: int = 0,
    ctx_train: ExecContext | None = None,
) -> list[float]:
    """Random-program null for the unsupervised objective.

    Sample ``n_permutations`` Vector-only programs (no target binding),
    score each on ``ctx_test`` (with ``ctx_train`` so the silhouette is
    out-of-sample, matching how the WINNER is now scored — otherwise the
    null and the observation are not on the same scale and p is
    deflated).
    """
    from engine_v2.synthesize import random_program

    py_rng = random.Random(seed + 9001)
    nulls: list[float] = []
    for _ in range(n_permutations):
        prog = random_program(
            py_rng,
            pool,
            objective_target=objective.target,
            max_depth=max_depth,
            max_genes_per_set=max_genes_per_set,
            rates=rates,
            return_type=TType.VECTOR,
        )
        nulls.append(
            evaluate_holdout(
                prog, ctx_test, None,
                objective=objective, ctx_train=ctx_train,
            )
        )
    return nulls


def permutation_p_value(observed: float, nulls: list[float]) -> float:
    arr = np.asarray(nulls)
    return float(((arr >= observed).sum() + 1) / (len(arr) + 1))