engine/split.py

"""Canonical TRAIN/TEST split used everywhere in H2.

Defined once here so the GP loop, the baseline, the permutation null, and
the final held-out evaluation all see the same patients on the same sides
of the wall. Stratified on the binary y. The orchestrator builds y from
the cohort labels; this module sees only the binary array.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split


@dataclass
class Split:
    train_ids: pd.Index
    test_ids: pd.Index
    y_train: np.ndarray
    y_test: np.ndarray


def make_split(
    sample_ids: pd.Index,
    y: np.ndarray,
    *,
    test_size: float = 0.3,
    random_state: int = 42,
    stratify: bool = True,
) -> Split:
    """TRAIN/TEST split returning original sample IDs on each side.

    Stratified on y when ``stratify`` (binary classification path); a plain
    random partition for continuous y.
    """
    if len(sample_ids) != len(y):
        raise ValueError("make_split: sample_ids and y length mismatch")
    pos = np.arange(len(sample_ids))
    pos_tr, pos_te, y_tr, y_te = train_test_split(
        pos, y, test_size=test_size, random_state=random_state,
        stratify=y if stratify else None,
    )
    return Split(
        train_ids=sample_ids[pos_tr],
        test_ids=sample_ids[pos_te],
        y_train=np.asarray(y_tr),
        y_test=np.asarray(y_te),
    )