rpotter6298
198 lines
6.9 KiB
Python
198 lines
6.9 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Any, Iterable, Optional
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
from sklearn.model_selection import KFold, StratifiedKFold
|
|
|
|
from .network_manager import PatientSplit
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SplitPlan:
|
|
train_patient_ids: set[Any]
|
|
val_patient_ids: set[Any]
|
|
holdout_patient_ids: set[Any]
|
|
|
|
|
|
def build_patient_split_plans(
|
|
patient_ids: Iterable[Any],
|
|
patient_labels: Iterable[Any],
|
|
*,
|
|
n_splits: int,
|
|
seed: int,
|
|
holdout_per_class: int = 0,
|
|
holdout_seed: int = 123,
|
|
) -> list[SplitPlan]:
|
|
"""
|
|
Core vector-based splitter.
|
|
|
|
Inputs are one row per patient:
|
|
- patient_ids: unique patient IDs
|
|
- patient_labels: one label per patient
|
|
"""
|
|
ids = np.asarray(list(patient_ids))
|
|
labels = np.asarray(list(patient_labels))
|
|
if ids.ndim != 1 or labels.ndim != 1:
|
|
raise ValueError("patient_ids and patient_labels must be 1D arrays")
|
|
if ids.size != labels.size:
|
|
raise ValueError(f"Length mismatch: ids={ids.size}, labels={labels.size}")
|
|
if ids.size == 0:
|
|
raise ValueError("No patients available for splitting")
|
|
if len(set(ids.tolist())) != ids.size:
|
|
raise ValueError("patient_ids must be unique (one label per patient)")
|
|
if n_splits < 2:
|
|
raise ValueError("n_splits must be >= 2")
|
|
|
|
holdout_ids: set[Any] = set()
|
|
if holdout_per_class > 0:
|
|
rng = np.random.default_rng(holdout_seed)
|
|
for label in np.unique(labels):
|
|
idx = np.where(labels == label)[0]
|
|
if idx.size == 0:
|
|
continue
|
|
n = min(holdout_per_class, idx.size)
|
|
chosen = rng.choice(idx, size=n, replace=False)
|
|
holdout_ids.update(ids[chosen].tolist())
|
|
|
|
keep_mask = ~np.isin(ids, list(holdout_ids))
|
|
cv_ids = ids[keep_mask]
|
|
cv_labels = labels[keep_mask]
|
|
if cv_ids.size < n_splits:
|
|
raise ValueError(
|
|
f"Not enough patients ({cv_ids.size}) for n_splits={n_splits} after holdout removal"
|
|
)
|
|
|
|
use_stratified = _can_stratify(cv_labels, n_splits)
|
|
if use_stratified:
|
|
splitter = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=seed)
|
|
splits = list(splitter.split(cv_ids, cv_labels))
|
|
else:
|
|
splitter = KFold(n_splits=n_splits, shuffle=True, random_state=seed)
|
|
splits = list(splitter.split(cv_ids))
|
|
|
|
plans: list[SplitPlan] = []
|
|
for train_idx, val_idx in splits:
|
|
plans.append(
|
|
SplitPlan(
|
|
train_patient_ids=set(cv_ids[train_idx].tolist()),
|
|
val_patient_ids=set(cv_ids[val_idx].tolist()),
|
|
holdout_patient_ids=set(holdout_ids),
|
|
)
|
|
)
|
|
return plans
|
|
|
|
|
|
class PatientFirstSplitManager:
|
|
"""
|
|
Patient-level splitter for V2.
|
|
|
|
Behavior:
|
|
- Optional binary filtering happens first (labels in {0,1} only).
|
|
- Optional holdout is sampled at the patient level (never per-eye rows).
|
|
- K-fold split is built on remaining patients.
|
|
- Returned dataframes contain all rows for each selected patient.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
patient_col: str = "Patient ID",
|
|
label_col: Optional[str] = None,
|
|
) -> None:
|
|
self.patient_col = patient_col
|
|
self.label_col = label_col
|
|
|
|
def build_plans(
|
|
self,
|
|
*,
|
|
clinical: Any,
|
|
args: Any,
|
|
profile: Optional[Any] = None,
|
|
) -> list[PatientSplit]:
|
|
profile_label_col = getattr(profile, "label_col", None) if profile is not None else None
|
|
profile_patient_col = getattr(profile, "patient_col", None) if profile is not None else None
|
|
patient_col = profile_patient_col or self.patient_col
|
|
label_col = self.label_col or profile_label_col or getattr(clinical, "label_col", None)
|
|
if label_col is None:
|
|
raise ValueError("Could not resolve label column from SplitManager or clinical.label_col")
|
|
|
|
if not hasattr(clinical, "df"):
|
|
raise ValueError("Clinical object must expose a dataframe at .df")
|
|
df_full = clinical.df.copy()
|
|
self._validate_columns(df_full, label_col, patient_col=patient_col)
|
|
|
|
eval_mode = str(getattr(args, "eval_mode", "multiclass")).lower()
|
|
if eval_mode == "binary":
|
|
df_full = df_full[df_full[label_col].isin([0, 1])].reset_index(drop=True)
|
|
|
|
holdout_per_class = int(getattr(args, "holdout_per_class", 0) or 0)
|
|
holdout_seed = int(getattr(args, "holdout_seed", 123))
|
|
n_splits = int(getattr(args, "n_splits", 5))
|
|
fold_seed = int(getattr(args, "fold_seed", 42))
|
|
|
|
patient_table = self._patient_label_table(df_full, label_col, patient_col=patient_col)
|
|
plans = build_patient_split_plans(
|
|
patient_ids=patient_table[patient_col].to_numpy(),
|
|
patient_labels=patient_table["_label"].to_numpy(),
|
|
n_splits=n_splits,
|
|
seed=fold_seed,
|
|
holdout_per_class=holdout_per_class,
|
|
holdout_seed=holdout_seed,
|
|
)
|
|
|
|
out: list[PatientSplit] = []
|
|
for plan in plans:
|
|
train_df = (
|
|
df_full[df_full[patient_col].isin(plan.train_patient_ids)]
|
|
.reset_index(drop=True)
|
|
)
|
|
val_df = (
|
|
df_full[df_full[patient_col].isin(plan.val_patient_ids)]
|
|
.reset_index(drop=True)
|
|
)
|
|
holdout_df = None
|
|
if plan.holdout_patient_ids:
|
|
holdout_df = (
|
|
df_full[df_full[patient_col].isin(plan.holdout_patient_ids)]
|
|
.reset_index(drop=True)
|
|
)
|
|
out.append(PatientSplit(train=train_df, val=val_df, holdout=holdout_df))
|
|
return out
|
|
|
|
def _validate_columns(self, df: pd.DataFrame, label_col: str, patient_col: Optional[str] = None) -> None:
|
|
pcol = patient_col or self.patient_col
|
|
if pcol not in df.columns:
|
|
raise ValueError(f"Missing required patient column: {pcol!r}")
|
|
if label_col not in df.columns:
|
|
raise ValueError(f"Missing required label column: {label_col!r}")
|
|
|
|
def _patient_label_table(
|
|
self,
|
|
df: pd.DataFrame,
|
|
label_col: str,
|
|
patient_col: Optional[str] = None,
|
|
) -> pd.DataFrame:
|
|
pcol = patient_col or self.patient_col
|
|
grouped = (
|
|
df.groupby(pcol, as_index=False)[label_col]
|
|
.agg(lambda x: x.mode().iloc[0] if not x.mode().empty else x.iloc[0])
|
|
.rename(columns={label_col: "_label"})
|
|
.sort_values(pcol)
|
|
.reset_index(drop=True)
|
|
)
|
|
if grouped.empty:
|
|
raise ValueError("No patients available for splitting")
|
|
return grouped
|
|
|
|
|
|
def _can_stratify(labels: np.ndarray, n_splits: int) -> bool:
|
|
if labels.size == 0:
|
|
return False
|
|
unique, counts = np.unique(labels, return_counts=True)
|
|
if len(unique) < 2:
|
|
return False
|
|
return bool(np.all(counts >= n_splits))
|