diff --git a/scripts/make_kfold_splits.py b/scripts/make_kfold_splits.py new file mode 100644 index 0000000..9b7c43e --- /dev/null +++ b/scripts/make_kfold_splits.py @@ -0,0 +1,250 @@ +"""Create reproducible patient-level k-fold splits for the paired LVEF cohort. + +E13 requirements: +- Split by unique subject_id, never by study/record. +- Each subject belongs to exactly one outer test fold. +- Within each fold, assign every row to train/val/test. +- Verify zero subject overlap within every fold. +- Preserve approximately the canonical 70/10/20 train/val/test proportions. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +import numpy as np +import pandas as pd +from check_ef40_prevalence import normalize_ef_le_40 +from make_splits import ( + file_sha256, + hash_values, + read_cohort, + verify_no_subject_overlap, + write_cohort, +) + + +def make_subject_folds( + subjects: list[Any], + *, + n_folds: int = 5, + val_frac: float = 0.10, + seed: int = 42, +) -> list[dict[str, list[Any]]]: + """Create patient-level outer folds with an inner validation holdout.""" + if n_folds < 2: + raise ValueError("n_folds must be at least 2.") + + if not 0.0 < val_frac < 1.0: + raise ValueError("val_frac must be between 0 and 1.") + + if val_frac >= 1.0 - (1.0 / n_folds): + raise ValueError( + "val_frac is too large for the requested number of folds; " + "the training split would be empty." + ) + subjects = sorted(subjects, key=lambda x: str(x)) + + if len(subjects) < n_folds: + raise ValueError( + f"Need at least {n_folds} subjects for {n_folds}-fold CV; got {len(subjects)}." + ) + + rng = np.random.default_rng(seed) + shuffled = list(rng.permutation(subjects)) + + test_folds = [list(x) for x in np.array_split(shuffled, n_folds)] + folds: list[dict[str, list[Any]]] = [] + + for fold_idx, test_subjects in enumerate(test_folds): + test_set = set(test_subjects) + remaining = [s for s in shuffled if s not in test_set] + + # val_frac is expressed relative to the full cohort. + # With 5 folds, test is ~20%; selecting 10% of the full cohort + # for validation leaves approximately 70% for training. + n_val = max(1, int(round(len(subjects) * val_frac))) + + fold_rng = np.random.default_rng(seed + fold_idx + 1) + remaining_shuffled = list(fold_rng.permutation(remaining)) + + val_subjects = remaining_shuffled[:n_val] + train_subjects = remaining_shuffled[n_val:] + + split = { + "train": train_subjects, + "val": val_subjects, + "test": test_subjects, + } + verify_no_subject_overlap(split) + folds.append(split) + + return folds + + +def verify_test_fold_coverage( + subjects: list[Any], + folds: list[dict[str, list[Any]]], +) -> None: + """Verify that every subject appears in exactly one outer test fold.""" + test_subjects = [subject for fold in folds for subject in fold["test"]] + + if len(test_subjects) != len(set(test_subjects)): + raise AssertionError("A subject appears in more than one outer test fold.") + + if set(test_subjects) != set(subjects): + raise AssertionError("Outer test folds do not cover every subject exactly once.") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--input", + type=Path, + default=Path("data/processed/echo_hubert_manifest.parquet"), + help="Joined EchoJEPA + HuBERT manifest.", + ) + parser.add_argument( + "--out-dir", + type=Path, + default=Path("data/processed/kfold"), + help="Directory where fold-specific manifests and metadata are written.", + ) + parser.add_argument("--n-folds", type=int, default=5) + parser.add_argument("--val-frac", type=float, default=0.10) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + + cohort = read_cohort(args.input) + + if "subject_id" not in cohort.columns: + raise KeyError("Input cohort must contain a subject_id column.") + + if cohort["subject_id"].isna().any(): + raise ValueError("Input cohort contains missing subject_id values.") + + subjects = cohort["subject_id"].drop_duplicates().tolist() + + folds = make_subject_folds( + subjects, + n_folds=args.n_folds, + val_frac=args.val_frac, + seed=args.seed, + ) + + verify_test_fold_coverage(subjects, folds) + + args.out_dir.mkdir(parents=True, exist_ok=True) + fold_metadata = [] + + for fold_idx, splits in enumerate(folds): + verify_no_subject_overlap(splits) + + split_lookup = { + subject_id: split for split, subject_ids in splits.items() for subject_id in subject_ids + } + + fold_df = cohort.copy() + + if "split" in fold_df.columns: + fold_df["split_canonical"] = fold_df["split"] + + fold_df["split"] = fold_df["subject_id"].map(split_lookup) + + if fold_df["split"].isna().any(): + raise AssertionError(f"Fold {fold_idx}: some cohort rows were not assigned a split.") + + output_splits = { + split: (fold_df.loc[fold_df["split"] == split, "subject_id"].drop_duplicates().tolist()) + for split in ("train", "val", "test") + } + verify_no_subject_overlap(output_splits) + + fold_path = args.out_dir / f"fold_{fold_idx}.parquet" + write_cohort(fold_df, fold_path) + + subject_split_df = pd.DataFrame( + [ + { + "subject_id": subject_id, + "split": split, + } + for split, subject_ids in splits.items() + for subject_id in subject_ids + ] + ).sort_values(["split", "subject_id"]) + + subject_split_path = args.out_dir / f"fold_{fold_idx}_subjects.csv" + subject_split_df.to_csv(subject_split_path, index=False) + + row_counts = { + split: int((fold_df["split"] == split).sum()) for split in ("train", "val", "test") + } + + subject_counts = { + split: int(fold_df.loc[fold_df["split"] == split, "subject_id"].nunique()) + for split in ("train", "val", "test") + } + + ef40_bool = normalize_ef_le_40(fold_df["ef_le_40"]) + + ef40_counts = { + split: int(ef40_bool.loc[fold_df["split"] == split].sum()) + for split in ("train", "val", "test") + } + + fold_metadata.append( + { + "fold": fold_idx, + "manifest_path": str(fold_path), + "subject_splits_path": str(subject_split_path), + "row_counts": row_counts, + "subject_counts": subject_counts, + "ef_le_40_counts": ef40_counts, + "split_subject_id_hashes": { + split: hash_values(subject_ids) for split, subject_ids in splits.items() + }, + "leakage_check": { + "train_val_overlap": 0, + "train_test_overlap": 0, + "val_test_overlap": 0, + }, + } + ) + + print(f"Fold {fold_idx}: rows={row_counts}, subjects={subject_counts}") + + metadata = { + "task": "E13_patient_level_kfold", + "input_path": str(args.input), + "input_sha256": file_sha256(args.input), + "seed": args.seed, + "n_folds": args.n_folds, + "val_frac": args.val_frac, + "n_rows": int(len(cohort)), + "n_subjects": int(len(subjects)), + "subject_id_hash": hash_values(subjects), + "outer_test_coverage": { + "n_unique_test_subjects": int( + len({subject for fold in folds for subject in fold["test"]}) + ), + "expected_subjects": int(len(subjects)), + "each_subject_tested_once": True, + }, + "folds": fold_metadata, + } + + metadata_path = args.out_dir / "kfold_manifest.json" + metadata_path.write_text(json.dumps(metadata, indent=2), encoding="utf-8") + + print(f"Wrote {args.n_folds} fold manifests to: {args.out_dir}") + print(f"Wrote k-fold metadata to: {metadata_path}") + print("Leakage check passed for every fold.") + print("Every subject appears in exactly one outer test fold.") + + +if __name__ == "__main__": + main() diff --git a/scripts/run_kfold_cv.py b/scripts/run_kfold_cv.py new file mode 100644 index 0000000..af710d8 --- /dev/null +++ b/scripts/run_kfold_cv.py @@ -0,0 +1,292 @@ +#!/usr/bin/env python3 +"""Run patient-level k-fold cross-validation for E13. + +For each fold, this script will: +1. Train the existing ECG, echo, concat, and fused probes. +2. Evaluate the fused checkpoint under full, echo_dropped, and ecg_dropped conditions. +3. Save fold-specific metrics and predictions. +4. Aggregate results across folds. + +The existing train_probes.py and evaluate_missing_modality.py paths are reused so +the k-fold experiment stays comparable with the canonical single-split run. +""" + +from __future__ import annotations + +import argparse +import json +import subprocess +import sys +from pathlib import Path +from typing import Any + +import numpy as np + +from primed_ai.probes import manifest as manifest_io +from primed_ai.probes.common import auroc, regression_metrics + +CONDITIONS = ("full", "echo_dropped", "ecg_dropped") + + +def run_command(command: list[str]) -> None: + """Run one pipeline command and stop immediately if it fails.""" + print("\nRunning:") + print(" ".join(command)) + subprocess.run(command, check=True) + + +def read_json(path: Path) -> dict[str, Any]: + """Read a JSON result file.""" + with path.open("r", encoding="utf-8") as f: + return json.load(f) + + +def train_fold( + fold_manifest: Path, + fold_dir: Path, + *, + epochs: int, + seed: int, + fusion_dim: int, +) -> Path: + """Train all four probes for one fold and return the fused checkpoint path.""" + probes_dir = fold_dir / "probes" + + command = [ + sys.executable, + "scripts/train_probes.py", + "--manifest", + str(fold_manifest), + "--probe", + "all", + "--out-dir", + str(probes_dir), + "--epochs", + str(epochs), + "--fusion-dim", + str(fusion_dim), + "--seed", + str(seed), + ] + + run_command(command) + + checkpoint = probes_dir / "fused" / "cross_attn_fused.pt" + + if not checkpoint.is_file(): + raise FileNotFoundError(f"Expected fused checkpoint was not created: {checkpoint}") + + return checkpoint + + +def evaluate_fold( + fold_manifest: Path, + checkpoint: Path, + fold_dir: Path, + *, + fusion_dim: int, + seed: int, + n_bootstrap: int, +) -> Path: + """Evaluate one fold's fused checkpoint under all modality conditions.""" + results_dir = fold_dir / "results" + results_dir.mkdir(parents=True, exist_ok=True) + + output_path = results_dir / "missing_modality.json" + + manifest_df = manifest_io.load(fold_manifest) + echo_dim, ecg_dim = manifest_io.dims(manifest_df) + + command = [ + sys.executable, + "scripts/evaluate_missing_modality.py", + "--manifest", + str(fold_manifest), + "--checkpoint", + str(checkpoint), + "--output", + str(output_path), + "--embed-dim", + str(fusion_dim), + "--echo-dim", + str(echo_dim), + "--ecg-dim", + str(ecg_dim), + "--seed", + str(seed), + "--n-bootstrap", + str(n_bootstrap), + ] + + run_command(command) + + if not output_path.is_file(): + raise FileNotFoundError(f"Expected missing-modality result was not created: {output_path}") + + return output_path + + +def aggregate_results(result_paths: list[Path]) -> dict[str, Any]: + """Combine per-fold metrics and out-of-fold predictions.""" + fold_payloads = [read_json(path) for path in result_paths] + + summary: dict[str, Any] = { + "n_folds": len(fold_payloads), + "conditions": {}, + } + + for condition in CONDITIONS: + per_fold = [] + pooled_lvef = [] + pooled_prediction = [] + pooled_ef40 = [] + + for fold_idx, payload in enumerate(fold_payloads): + metrics = payload["test"][condition] + predictions = payload["predictions"][condition] + + per_fold.append( + { + "fold": fold_idx, + "n_test": len(predictions["lvef"]), + "mae": metrics["mae"], + "ef40_auroc": metrics["ef40_auroc"], + "bootstrap": payload["bootstrap"].get(condition, {}), + } + ) + + pooled_lvef.extend(predictions["lvef"]) + pooled_prediction.extend(predictions["prediction"]) + pooled_ef40.extend(predictions["ef_le_40"]) + + fold_mae = np.asarray([row["mae"] for row in per_fold], dtype=float) + fold_auroc = np.asarray( + [row["ef40_auroc"] for row in per_fold], + dtype=float, + ) + + pooled_lvef_array = np.asarray(pooled_lvef, dtype=float) + pooled_prediction_array = np.asarray(pooled_prediction, dtype=float) + pooled_ef40_array = np.asarray(pooled_ef40, dtype=bool) + + pooled_metrics = regression_metrics( + pooled_lvef_array, + pooled_prediction_array, + ) + pooled_auroc = auroc( + pooled_ef40_array, + -pooled_prediction_array, + ) + + summary["conditions"][condition] = { + "per_fold": per_fold, + "across_fold": { + "mae_mean": round(float(np.mean(fold_mae)), 4), + "mae_std": round(float(np.std(fold_mae, ddof=1)), 4), + "ef40_auroc_mean": round(float(np.mean(fold_auroc)), 4), + "ef40_auroc_std": round(float(np.std(fold_auroc, ddof=1)), 4), + }, + "pooled_out_of_fold": { + "n": len(pooled_lvef_array), + "mae": pooled_metrics["mae"], + "ef40_auroc": round(float(pooled_auroc), 4), + }, + } + + return summary + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--folds-dir", + type=Path, + required=True, + help="Directory containing fold_0.parquet, fold_1.parquet, etc.", + ) + parser.add_argument( + "--out-dir", + type=Path, + default=Path("results/kfold_cv"), + help="Directory for fold checkpoints and evaluation results.", + ) + parser.add_argument("--n-folds", type=int, default=5) + parser.add_argument("--epochs", type=int, default=50) + parser.add_argument("--fusion-dim", type=int, default=256) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--n-bootstrap", type=int, default=1000) + args = parser.parse_args() + + args.out_dir.mkdir(parents=True, exist_ok=True) + + result_paths: list[Path] = [] + + for fold_idx in range(args.n_folds): + print(f"\n{'=' * 60}") + print(f"Fold {fold_idx}") + print(f"{'=' * 60}") + + fold_manifest = args.folds_dir / f"fold_{fold_idx}.parquet" + + if not fold_manifest.is_file(): + raise FileNotFoundError(f"Fold manifest not found: {fold_manifest}") + + fold_dir = args.out_dir / f"fold_{fold_idx}" + fold_dir.mkdir(parents=True, exist_ok=True) + + fold_seed = args.seed + + checkpoint = train_fold( + fold_manifest, + fold_dir, + epochs=args.epochs, + seed=fold_seed, + fusion_dim=args.fusion_dim, + ) + + result_path = evaluate_fold( + fold_manifest, + checkpoint, + fold_dir, + fusion_dim=args.fusion_dim, + seed=fold_seed, + n_bootstrap=args.n_bootstrap, + ) + + result_paths.append(result_path) + + summary = aggregate_results(result_paths) + + summary["config"] = { + "folds_dir": str(args.folds_dir), + "out_dir": str(args.out_dir), + "n_folds": args.n_folds, + "epochs": args.epochs, + "fusion_dim": args.fusion_dim, + "seed": args.seed, + "n_bootstrap": args.n_bootstrap, + } + + summary_path = args.out_dir / "kfold_results.json" + + with summary_path.open("w", encoding="utf-8") as f: + json.dump(summary, f, indent=2) + + print(f"\nWrote k-fold results to: {summary_path}") + + for condition in CONDITIONS: + results = summary["conditions"][condition] + across = results["across_fold"] + pooled = results["pooled_out_of_fold"] + + print(f"\n{condition}") + print(f" Across-fold MAE: {across['mae_mean']:.4f} ± {across['mae_std']:.4f}") + print( + f" Across-fold AUROC: {across['ef40_auroc_mean']:.4f} ± {across['ef40_auroc_std']:.4f}" + ) + print(f" Pooled OOF MAE: {pooled['mae']:.4f}") + print(f" Pooled OOF AUROC: {pooled['ef40_auroc']:.4f}") + + +if __name__ == "__main__": + main() diff --git a/tests/test_kfold_cv.py b/tests/test_kfold_cv.py new file mode 100644 index 0000000..81864cc --- /dev/null +++ b/tests/test_kfold_cv.py @@ -0,0 +1,193 @@ +import json +import sys +from pathlib import Path + +SCRIPTS_DIR = Path(__file__).resolve().parents[1] / "scripts" +sys.path.insert(0, str(SCRIPTS_DIR)) + +import run_kfold_cv as kfold_cv # noqa: E402 + +CONDITIONS = ("full", "echo_dropped", "ecg_dropped") + + +def _payload(lvef, prediction): + ef_le_40 = [value <= 40 for value in lvef] + + return { + "test": { + condition: { + "mae": 5.0, + "ef40_auroc": 1.0, + } + for condition in CONDITIONS + }, + "bootstrap": { + condition: { + "mae_ci_low": 4.0, + "mae_ci_high": 6.0, + } + for condition in CONDITIONS + }, + "predictions": { + condition: { + "lvef": lvef, + "prediction": prediction, + "ef_le_40": ef_le_40, + } + for condition in CONDITIONS + }, + } + + +def test_aggregate_results_combines_out_of_fold_predictions(tmp_path): + fold_0 = tmp_path / "fold_0.json" + fold_1 = tmp_path / "fold_1.json" + + fold_0.write_text( + json.dumps(_payload([30.0, 60.0], [35.0, 55.0])), + encoding="utf-8", + ) + fold_1.write_text( + json.dumps(_payload([20.0, 70.0], [25.0, 65.0])), + encoding="utf-8", + ) + + result = kfold_cv.aggregate_results([fold_0, fold_1]) + + assert result["n_folds"] == 2 + + for condition in CONDITIONS: + condition_result = result["conditions"][condition] + + assert len(condition_result["per_fold"]) == 2 + + across = condition_result["across_fold"] + assert across["mae_mean"] == 5.0 + assert across["mae_std"] == 0.0 + assert across["ef40_auroc_mean"] == 1.0 + assert across["ef40_auroc_std"] == 0.0 + + pooled = condition_result["pooled_out_of_fold"] + assert pooled["n"] == 4 + assert pooled["mae"] == 5.0 + assert pooled["ef40_auroc"] == 1.0 + + +def test_main_runs_each_fold_without_real_training(tmp_path, monkeypatch): + folds_dir = tmp_path / "folds" + out_dir = tmp_path / "results" + folds_dir.mkdir() + + for fold_idx in range(2): + (folds_dir / f"fold_{fold_idx}.parquet").touch() + + trained = [] + evaluated = [] + + def fake_train_fold( + fold_manifest, + fold_dir, + *, + epochs, + seed, + fusion_dim, + ): + trained.append( + { + "manifest": fold_manifest.name, + "fold_dir": fold_dir.name, + "epochs": epochs, + "seed": seed, + "fusion_dim": fusion_dim, + } + ) + return fold_dir / "probes" / "fused" / "cross_attn_fused.pt" + + def fake_evaluate_fold( + fold_manifest, + checkpoint, + fold_dir, + *, + fusion_dim, + seed, + n_bootstrap, + ): + evaluated.append( + { + "manifest": fold_manifest.name, + "fold_dir": fold_dir.name, + "seed": seed, + "fusion_dim": fusion_dim, + "n_bootstrap": n_bootstrap, + } + ) + return fold_dir / "results" / "missing_modality.json" + + fake_summary = { + "n_folds": 2, + "conditions": { + condition: { + "per_fold": [], + "across_fold": { + "mae_mean": 0.0, + "mae_std": 0.0, + "ef40_auroc_mean": 0.0, + "ef40_auroc_std": 0.0, + }, + "pooled_out_of_fold": { + "n": 0, + "mae": 0.0, + "ef40_auroc": 0.0, + }, + } + for condition in kfold_cv.CONDITIONS + }, + } + + monkeypatch.setattr(kfold_cv, "train_fold", fake_train_fold) + monkeypatch.setattr(kfold_cv, "evaluate_fold", fake_evaluate_fold) + monkeypatch.setattr( + kfold_cv, + "aggregate_results", + lambda result_paths: fake_summary, + ) + + monkeypatch.setattr( + sys, + "argv", + [ + "run_kfold_cv.py", + "--folds-dir", + str(folds_dir), + "--out-dir", + str(out_dir), + "--n-folds", + "2", + "--epochs", + "3", + "--fusion-dim", + "16", + "--seed", + "42", + "--n-bootstrap", + "10", + ], + ) + + kfold_cv.main() + + assert [row["manifest"] for row in trained] == [ + "fold_0.parquet", + "fold_1.parquet", + ] + assert len(evaluated) == 2 + + assert all(row["seed"] == 42 for row in trained) + assert all(row["seed"] == 42 for row in evaluated) + + summary_path = out_dir / "kfold_results.json" + assert summary_path.is_file() + + summary = json.loads(summary_path.read_text(encoding="utf-8")) + assert summary["config"]["n_folds"] == 2 + assert summary["config"]["seed"] == 42 diff --git a/tests/test_kfold_splits.py b/tests/test_kfold_splits.py new file mode 100644 index 0000000..b735cea --- /dev/null +++ b/tests/test_kfold_splits.py @@ -0,0 +1,170 @@ +import json +import sys + +import pandas as pd +import pytest +from make_kfold_splits import ( + main, + make_subject_folds, + verify_no_subject_overlap, + verify_test_fold_coverage, +) + + +def test_five_fold_sizes_match_70_10_20(): + subjects = list(range(100)) + + folds = make_subject_folds( + subjects, + n_folds=5, + val_frac=0.10, + seed=42, + ) + + assert len(folds) == 5 + + for fold in folds: + assert len(fold["train"]) == 70 + assert len(fold["val"]) == 10 + assert len(fold["test"]) == 20 + + +def test_no_subject_leakage_within_folds(): + subjects = list(range(100)) + + folds = make_subject_folds( + subjects, + n_folds=5, + val_frac=0.10, + seed=42, + ) + + for fold in folds: + verify_no_subject_overlap(fold) + + +def test_every_subject_is_tested_once(): + subjects = list(range(100)) + + folds = make_subject_folds( + subjects, + n_folds=5, + val_frac=0.10, + seed=42, + ) + + verify_test_fold_coverage(subjects, folds) + + test_subjects = [subject for fold in folds for subject in fold["test"]] + + assert len(test_subjects) == 100 + assert len(set(test_subjects)) == 100 + + +def test_same_seed_is_reproducible(): + subjects = list(range(100)) + + folds_a = make_subject_folds( + subjects, + n_folds=5, + val_frac=0.10, + seed=42, + ) + + folds_b = make_subject_folds( + subjects, + n_folds=5, + val_frac=0.10, + seed=42, + ) + + assert folds_a == folds_b + + +def test_different_seed_changes_assignment(): + subjects = list(range(100)) + + folds_a = make_subject_folds( + subjects, + n_folds=5, + val_frac=0.10, + seed=42, + ) + + folds_b = make_subject_folds( + subjects, + n_folds=5, + val_frac=0.10, + seed=43, + ) + + assert folds_a != folds_b + + +def test_val_frac_cannot_empty_training_split(): + subjects = list(range(100)) + + with pytest.raises(ValueError, match="training split would be empty"): + make_subject_folds( + subjects, + n_folds=2, + val_frac=0.5, + seed=42, + ) + + +def test_main_writes_ef40_counts_and_preserves_canonical_split( + tmp_path, + monkeypatch, +): + input_path = tmp_path / "cohort.parquet" + out_dir = tmp_path / "kfold" + + cohort = pd.DataFrame( + { + "subject_id": list(range(100)), + "split": ["train"] * 70 + ["val"] * 10 + ["test"] * 20, + "ef_le_40": [i % 4 == 0 for i in range(100)], + } + ) + cohort.to_parquet(input_path, index=False) + + monkeypatch.setattr( + sys, + "argv", + [ + "make_kfold_splits.py", + "--input", + str(input_path), + "--out-dir", + str(out_dir), + "--n-folds", + "5", + "--val-frac", + "0.10", + "--seed", + "42", + ], + ) + + main() + + metadata = json.loads((out_dir / "kfold_manifest.json").read_text(encoding="utf-8")) + + fold_df = pd.read_parquet(out_dir / "fold_0.parquet") + fold_metadata = metadata["folds"][0] + + assert "split_canonical" in fold_df.columns + + expected_canonical = cohort.sort_values("subject_id")["split"].tolist() + actual_canonical = fold_df.sort_values("subject_id")["split_canonical"].tolist() + assert actual_canonical == expected_canonical + + for split in ("train", "val", "test"): + expected_count = int( + fold_df.loc[ + fold_df["split"] == split, + "ef_le_40", + ].sum() + ) + assert fold_metadata["ef_le_40_counts"][split] == expected_count