Implementation of K-fold for better analysis

config.toml

@@ -3,6 +3,7 @@
 # "scen_load_all"  - Load and evaluate model, with noise
 # "scen_load_eval" - Load and evaluate model, without noise
 # "scen_analyze"   - Analyze existing evaluation files only (no training/eval)
+# "scen_kfold"     - Patient-grouped k-fold CV: train+eval per fold, then OOF + analyze
 scenario = "scen_train_all"
 
 
@@ -38,7 +39,16 @@ deterministic = false # force deterministic cuDNN kernels (slower, exact reprodu
 [evaluation]
 # Gaussian noise standard deviations applied to images during noisy evaluation
 # (step 6). 0.0 is the clean baseline and should be kept first.
-noise_levels = [0.0, 0.02, 0.05, 0.1, 0.2]
+noise_levels = [0.0, 0.05, 0.1, 0.2, 0.3, 0.5, 0.7, 1.0, 1.5, 2.0, 3.0, 5.0]
 # Monte-Carlo forward passes used to estimate Bayesian predictive/model
 # uncertainty (step 5/6). Ignored for the deterministic ensemble.
 mc_passes = 30
+
+
+[kfold]
+# Used only by the "scen_kfold" scenario. Each fold trains a full ensemble on
+# ~(k-1)/k of the patients and evaluates on the held-out 1/k; out-of-fold results
+# are pooled to cover every sample. NOTE: cost is ~k * a normal train+eval run.
+k_folds = 4          # number of patient-grouped folds (each ~1/k held out)
+val_fraction = 0.15  # fraction of each fold's non-test patients used for validation
+include_noise = false # also run the noise sweep within each fold (much slower)

main.py

@@ -24,7 +24,7 @@ from textual.widgets import (
 
 from tasks import PIPELINE_TASKS, SCENARIOS, describe_protected_conflicts
 from util.config_manager import handle_directory_config
-from util.progress import ProgressTracker
+from util.progress import MAX_LEVELS, ProgressTracker
 from util.screens import OverwriteConfirmScreen
 from util.ui_logger import PipelineLogger
 
@@ -89,7 +89,7 @@ class CNNHarnessApp(App[None]):
 
                     # --- Progress Trackers ---
                     with Vertical(id="progress_area"):
-                        for i in range(5):
+                        for i in range(MAX_LEVELS):
                             with Horizontal(
                                 id=f"progress_row_{i}", classes="progress_row"
                             ):

model/dataset.py

@@ -285,32 +285,14 @@ def initalize_dataloaders(
     return loaders
 
 
-def divide_dataset_by_patient_id(
-    dataset: ADNIDataset,
-    ptids: List[Tuple[int, str]],
-    ratios: Tuple[float, float, float],
-    seed: int,
-) -> List[data.Subset[ADNIDataset]]:
-    """
-    Divides the dataset into training, validation, and test sets based on patient IDs.
-    Ensures that all samples from the same patient are in the same set.
+def _patient_index(
+    dataset: ADNIDataset, ptids: List[Tuple[int, str]]
+) -> Dict[str, List[int]]:
+    """Map patient id -> dataset sample indices, validating PTID consistency.
 
-    Args:
-        dataset (ADNIDataset): The dataset to divide.
-        ptids (List[Tuple[int, str]]): A list of tuples containing image file IDs and their corresponding patient IDs.
-        ratios (Tuple[float, float, float]): The ratios for training, validation, and test sets.
-        seed (int): The random seed for reproducibility.
-
-    Returns:
-        List[data.Subset[ADNIDataset]]: A list of subsets for training, validation, and test sets.
-
-    Notes:
-        This split is grouped by PTID, so all images from the same patient are assigned
-        to exactly one partition to avoid patient-level leakage across train/val/test.
+    Shared by the single-split and k-fold partitioners so both enforce the same
+    patient-grouping and leakage checks.
     """
-    if not math.isclose(sum(ratios), 1.0, rel_tol=1e-9, abs_tol=1e-9):
-        raise ValueError(f"Ratios must sum to 1.0, got {ratios} (sum={sum(ratios)}).")
-
     if not ptids:
         raise ValueError("ptids list cannot be empty.")
 
@@ -320,7 +302,6 @@ def divide_dataset_by_patient_id(
         patient_id_str = str(patient_id).strip()
         if not patient_id_str or patient_id_str.lower() == "nan":
             raise ValueError(f"Invalid PTID for Image Data ID {image_id_int}.")
-
         if (
             image_id_int in image_to_patient
             and image_to_patient[image_id_int] != patient_id_str
@@ -329,7 +310,6 @@ def divide_dataset_by_patient_id(
                 f"Conflicting PTIDs for Image Data ID {image_id_int}: "
                 f"{image_to_patient[image_id_int]} vs {patient_id_str}."
             )
-
         image_to_patient[image_id_int] = patient_id_str
 
     patient_to_indices: Dict[str, List[int]] = {}
@@ -338,11 +318,38 @@ def divide_dataset_by_patient_id(
             raise ValueError(
                 f"Missing PTID mapping for dataset Image Data ID {image_id}."
             )
+        patient_to_indices.setdefault(image_to_patient[image_id], []).append(idx)
+
+    return patient_to_indices
+
+
+def divide_dataset_by_patient_id(
+    dataset: ADNIDataset,
+    ptids: List[Tuple[int, str]],
+    ratios: Tuple[float, float, float],
+    seed: int,
+) -> List[data.Subset[ADNIDataset]]:
+    """
+    Divides the dataset into training, validation, and test sets based on patient IDs.
+    Ensures that all samples from the same patient are in the same set.
 
-        patient_id = image_to_patient[image_id]
-        if patient_id not in patient_to_indices:
-            patient_to_indices[patient_id] = []
-        patient_to_indices[patient_id].append(idx)
+    Args:
+        dataset (ADNIDataset): The dataset to divide.
+        ptids (List[Tuple[int, str]]): A list of tuples containing image file IDs and their corresponding patient IDs.
+        ratios (Tuple[float, float, float]): The ratios for training, validation, and test sets.
+        seed (int): The random seed for reproducibility.
+
+    Returns:
+        List[data.Subset[ADNIDataset]]: A list of subsets for training, validation, and test sets.
+
+    Notes:
+        This split is grouped by PTID, so all images from the same patient are assigned
+        to exactly one partition to avoid patient-level leakage across train/val/test.
+    """
+    if not math.isclose(sum(ratios), 1.0, rel_tol=1e-9, abs_tol=1e-9):
+        raise ValueError(f"Ratios must sum to 1.0, got {ratios} (sum={sum(ratios)}).")
+
+    patient_to_indices = _patient_index(dataset, ptids)
 
     shuffled_patients = list(patient_to_indices.keys())
     random.Random(seed).shuffle(shuffled_patients)
@@ -402,3 +409,83 @@ def divide_dataset_by_patient_id(
         Subset(dataset, val_indices),
         Subset(dataset, test_indices),
     ]
+
+
+def kfold_split_by_patient_id(
+    dataset: ADNIDataset,
+    ptids: List[Tuple[int, str]],
+    k_folds: int,
+    fold: int,
+    seed: int,
+    val_fraction: float = 0.15,
+) -> List[data.Subset[ADNIDataset]]:
+    """Patient-grouped k-fold partition for a single fold.
+
+    Patients are shuffled once (deterministically from ``seed``) and dealt
+    round-robin into ``k_folds`` folds, so every patient is in exactly one fold
+    regardless of which fold is requested. For the requested ``fold``:
+      - test  = that fold's patients (the held-out slice),
+      - the remaining patients are split into val/train by ``val_fraction``.
+    All splits are patient-disjoint (no leakage), and across all folds every
+    sample is a test point exactly once.
+
+    Args:
+        dataset: The dataset to partition.
+        ptids: (Image Data ID, PTID) pairs covering the dataset.
+        k_folds: Number of folds (>= 2).
+        fold: Which fold to hold out as test (0 <= fold < k_folds).
+        seed: Seed for the (shared across folds) patient shuffle.
+        val_fraction: Fraction of the NON-test patients used for validation.
+
+    Returns:
+        ``[train, val, test]`` subsets.
+    """
+    if k_folds < 2:
+        raise ValueError(f"k_folds must be >= 2, got {k_folds}.")
+    if not 0 <= fold < k_folds:
+        raise ValueError(f"fold must be in [0, {k_folds}), got {fold}.")
+    if not 0.0 <= val_fraction < 1.0:
+        raise ValueError(f"val_fraction must be in [0.0, 1.0), got {val_fraction}.")
+
+    patient_to_indices = _patient_index(dataset, ptids)
+    patients = list(patient_to_indices.keys())
+    if len(patients) < k_folds:
+        raise ValueError(
+            f"Cannot make {k_folds} folds from {len(patients)} patients."
+        )
+    # Deterministic shuffle shared across all folds, then round-robin deal so
+    # fold sizes are balanced and the same partition is reproduced for every fold.
+    random.Random(seed).shuffle(patients)
+    fold_patients = [patients[i::k_folds] for i in range(k_folds)]
+
+    test_patients = fold_patients[fold]
+    rest = [p for f in range(k_folds) if f != fold for p in fold_patients[f]]
+    n_val = int(round(len(rest) * val_fraction))
+    val_patients = rest[:n_val]
+    train_patients = rest[n_val:]
+
+    def _indices(group: List[str]) -> List[int]:
+        return [idx for patient in group for idx in patient_to_indices[patient]]
+
+    train_indices = _indices(train_patients)
+    val_indices = _indices(val_patients)
+    test_indices = _indices(test_patients)
+
+    # Safety: patient-disjoint and full, exact coverage of the dataset.
+    train_set, val_set, test_set = (
+        set(train_indices),
+        set(val_indices),
+        set(test_indices),
+    )
+    if train_set & val_set or train_set & test_set or val_set & test_set:
+        raise ValueError("Sample index overlap detected across k-fold splits.")
+    if train_set | val_set | test_set != set(range(len(dataset))):
+        raise ValueError(
+            "k-fold coverage check failed: not all samples assigned exactly once."
+        )
+
+    return [
+        Subset(dataset, train_indices),
+        Subset(dataset, val_indices),
+        Subset(dataset, test_indices),
+    ]

tasks/__init__.py

@@ -27,6 +27,7 @@ from . import load_data
 from . import load_models
 from . import train_bayesian
 from . import train_normal
+from . import kfold  # imported last: depends on the sibling task modules above
 
 
 class TaskEntry(TypedDict):
@@ -72,6 +73,10 @@ PIPELINE_TASKS: Dict[str, TaskEntry] = {
         "task_name": "Analyze Evaluations",
         "task_func": analysis.run_analysis_task,
     },
+    "run_kfold": {
+        "task_name": "K-Fold Cross-Validation",
+        "task_func": kfold.run_kfold_task,
+    },
 }
 
 # --------------------------------------------------------------------------
@@ -98,6 +103,11 @@ ARTIFACT_CATEGORIES: Dict[str, ArtifactCategory] = {
         "protected": True,
         "desc": "model evaluation outputs (netCDF)",
     },
+    "kfold": {
+        "dirs": ("folds",),
+        "protected": True,
+        "desc": "per-fold k-fold models & evaluations",
+    },
     "analysis": {
         "dirs": (analysis.ANALYSIS_SUBDIR,),
         "protected": False,
@@ -114,6 +124,9 @@ TASK_WRITES = {
     "evaluate_bayesian": ("evaluations",),
     "evaluate_noisy": ("evaluations",),
     "run_analysis": ("analysis",),
+    # k-fold writes per-fold models+evals under folds/ and the aggregated OOF
+    # files under evaluations/; both are protected (slow to regenerate).
+    "run_kfold": ("kfold", "evaluations"),
 }
 
 
@@ -192,4 +205,11 @@ SCENARIOS: Dict[str, ScenarioEntry] = {
             "run_analysis",
         ],
     },
+    "scen_kfold": {
+        "label": "5. K-Fold Cross-Validation (train + OOF eval + analyze)",
+        "tasks": [
+            "run_kfold",
+            "run_analysis",
+        ],
+    },
 }

tasks/kfold.py

@@ -0,0 +1,151 @@
+"""Patient-grouped k-fold cross-validation driver (single process, sequential).
+
+Runs the ordinary train->evaluate pipeline once per fold, each on a different
+patient-grouped partition, writing per-fold artifacts under
+``<work_dir>/folds/fold_<k>/``. It then concatenates every fold's held-out
+evaluation along the ``sample`` axis into out-of-fold (OOF) files at
+``<work_dir>/evaluations/oof_*.nc`` -- a single evaluation covering EVERY sample,
+each predicted by the fold-ensemble that never trained on it. The existing
+analyses then run unchanged on the OOF files (via ``scen_analyze`` or the
+``run_analysis`` step appended to ``scen_kfold``).
+
+Reuses the existing task functions verbatim: each fold just gets its own
+dataloaders in ``state`` and a per-fold ``work_dir`` in a shallow config copy.
+The full dataset is loaded ONCE and every fold is a set of Subset indices over
+it, so memory use matches a single run.
+"""
+
+import pathlib as pl
+from typing import Any, Dict, List
+
+import numpy as np
+import xarray as xr
+
+import model.dataset as ds
+from evaluation import CLASS_NAMES, load_evaluation, save_evaluation
+from util.control import check_control_events
+from util.progress import ProgressTracker
+from util.ui_logger import PipelineLogger
+
+from . import evaluate, load_data, load_models, train_bayesian, train_normal
+
+# Inner per-fold pipeline (load_data is handled by the driver; noise optional).
+_BASE_INNER_TASKS = [
+    ("Train Regular", train_normal.train_normal_task),
+    ("Train Bayesian", train_bayesian.train_bayesian_task),
+    ("Load Models", load_models.load_models_task),
+    ("Evaluate Regular", evaluate.evaluate_normal_task),
+    ("Evaluate Bayesian", evaluate.evaluate_bayesian_task),
+]
+
+_FOLDS_SUBDIR = "folds"
+_EVAL_SUBDIR = "evaluations"
+
+
+def _kfold_cfg(config: Dict[str, Any]):
+    cfg = config.get("kfold", {})
+    k_folds = int(cfg.get("k_folds", 4))
+    val_fraction = float(cfg.get("val_fraction", 0.15))
+    include_noise = bool(cfg.get("include_noise", False))
+    return k_folds, val_fraction, include_noise
+
+
+def _aggregate_oof(
+    work_dir: pl.Path, fold_dirs: List[pl.Path], log: PipelineLogger
+) -> None:
+    """Concatenate each fold's held-out evaluations into OOF files."""
+    out_dir = work_dir / _EVAL_SUBDIR
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    # Evaluation file names are identical across folds (normal.nc, bayesian.nc,
+    # and noisy_* when enabled); discover them from the folds that produced them.
+    stems = sorted(
+        {p.stem for d in fold_dirs for p in (d / _EVAL_SUBDIR).glob("*.nc")}
+    )
+    for stem in stems:
+        parts = [
+            load_evaluation(str(d / _EVAL_SUBDIR / f"{stem}.nc"))
+            for d in fold_dirs
+            if (d / _EVAL_SUBDIR / f"{stem}.nc").exists()
+        ]
+        if not parts:
+            continue
+        # Folds hold disjoint samples; concatenating along `sample` yields full,
+        # non-overlapping out-of-fold coverage. Reset the (per-fold) sample index.
+        oof = xr.concat(parts, dim="sample")
+        oof = oof.assign_coords(sample=np.arange(oof.sizes["sample"]))
+        oof.attrs["n_folds"] = len(parts)
+        oof.attrs["kfold"] = 1
+        save_evaluation(oof, str(out_dir / f"oof_{stem}.nc"))
+        for part in parts:
+            part.close()
+        log.info(
+            f"Aggregated OOF: oof_{stem}.nc "
+            f"({oof.sizes['sample']} samples from {len(parts)} folds)."
+        )
+
+
+def run_kfold_task(
+    track: ProgressTracker,
+    log: PipelineLogger,
+    config: Dict[str, Any],
+    state: Dict[str, Any],
+) -> None:
+    events = state.get("events")
+    stop_event = events.get("stop") if isinstance(events, dict) else None
+    pause_event = events.get("pause") if isinstance(events, dict) else None
+
+    base_seed = load_data.seed_pipeline(config, log)
+    k_folds, val_fraction, include_noise = _kfold_cfg(config)
+
+    dataset, ptids, image_to_ptid = load_data.load_full_dataset(config, log)
+    state["image_to_ptid"] = image_to_ptid
+
+    inner_tasks = list(_BASE_INNER_TASKS)
+    if include_noise:
+        inner_tasks.append(("Evaluate Noisy", evaluate.evaluate_noisy_task))
+
+    work_dir = pl.Path(config["work_dir"])
+    fold_dirs: List[pl.Path] = []
+
+    log.info(
+        f"Starting {k_folds}-fold cross-validation "
+        f"(val_fraction={val_fraction}, include_noise={include_noise})."
+    )
+    track.set_title("K-Fold Cross-Validation")
+    track.update(total=k_folds, advance=0)
+
+    for fold in range(k_folds):
+        check_control_events(stop_event, pause_event)
+
+        subsets = ds.kfold_split_by_patient_id(
+            dataset, ptids, k_folds, fold, base_seed, val_fraction
+        )
+        train_loader, val_loader, test_loader = ds.initalize_dataloaders(
+            subsets, batch_size=config["training"]["batch_size"], seed=base_seed
+        )
+        state["train_loader"] = train_loader
+        state["val_loader"] = val_loader
+        state["test_loader"] = test_loader
+
+        fold_dir = work_dir / _FOLDS_SUBDIR / f"fold_{fold}"
+        fold_dirs.append(fold_dir)
+        # Per-fold config: same everything, but outputs go under the fold dir.
+        fold_config = {**config, "work_dir": str(fold_dir)}
+
+        log.info(
+            f"===== Fold {fold + 1}/{k_folds} "
+            f"(train={len(subsets[0])}, val={len(subsets[1])}, "
+            f"test={len(subsets[2])} scans) ====="
+        )
+        for task_label, task_func in inner_tasks:
+            check_control_events(stop_event, pause_event)
+            fold_track = track.get_sub_tracker(
+                f"Fold {fold + 1}/{k_folds}: {task_label}"
+            )
+            task_func(fold_track, log, fold_config, state)
+
+        track.update(advance=1)
+
+    log.info("All folds complete; aggregating out-of-fold evaluations.")
+    _aggregate_oof(work_dir, fold_dirs, log)

tasks/load_data.py

@@ -1,5 +1,5 @@
 import pathlib as pl
-from typing import Any, Dict
+from typing import Any, Dict, List, Tuple
 
 import pandas as pd
 
@@ -9,20 +9,11 @@ from util.seeding import seed_everything
 from util.ui_logger import PipelineLogger
 
 
-def load_data_task(
-    track: ProgressTracker,
-    log: PipelineLogger,
-    config: Dict[str, Any],
-    state: Dict[str, Any],
-):
-
-    if config["data"]["seed"] is None:
+def seed_pipeline(config: Dict[str, Any], log: PipelineLogger) -> int:
+    """Resolve the base seed (default 0) and seed all RNGs once. Returns it."""
+    if config["data"].get("seed") is None:
         log.info("Seed is not defined - using default seed of 0")
         config["data"]["seed"] = 0
-
-    # Seed everything once at the start of the pipeline so data loading, weight
-    # init, and shuffling are reproducible. Per-member seeds are derived from
-    # this base inside the training tasks.
     base_seed = int(config["data"]["seed"])
     deterministic = bool(config.get("training", {}).get("deterministic", False))
     applied = seed_everything(base_seed, deterministic=deterministic)
@@ -30,7 +21,17 @@ def load_data_task(
         f"Seeded RNGs with base seed {applied}"
         + (" (deterministic cuDNN)" if deterministic else "")
     )
+    return base_seed
+
+
+def load_full_dataset(
+    config: Dict[str, Any], log: PipelineLogger
+) -> Tuple["ds.ADNIDataset", List[Tuple[int, str]], Dict[int, str]]:
+    """Load the full ADNI dataset and its patient mapping (no split).
 
+    Shared by the single-split loader and the k-fold driver, which partition the
+    returned dataset differently.
+    """
     log.info("Loading Files")
     mri_files = pl.Path(config["data"]["mri_files_path"]).glob("*.nii")
     xls_file = pl.Path(config["data"]["xls_file_path"])
@@ -44,7 +45,6 @@ def load_data_task(
 
     ptid_df = pd.read_csv(xls_file)
     ptid_df.columns = ptid_df.columns.str.strip()
-
     ptid_df = ptid_df[["Image Data ID", "PTID"]].dropna(  # type: ignore
         subset=["Image Data ID", "PTID"]
     )
@@ -53,9 +53,25 @@ def load_data_task(
     ptid_df = ptid_df[ptid_df["PTID"] != ""]
 
     ptids = list(zip(ptid_df["Image Data ID"].tolist(), ptid_df["PTID"].tolist()))
+    image_to_ptid = {int(iid): str(pid) for iid, pid in ptids}
+    return dataset, ptids, image_to_ptid
+
+
+def load_data_task(
+    track: ProgressTracker,
+    log: PipelineLogger,
+    config: Dict[str, Any],
+    state: Dict[str, Any],
+):
+    # Seed everything once at the start of the pipeline so data loading, weight
+    # init, and shuffling are reproducible. Per-member seeds are derived from
+    # this base inside the training tasks.
+    base_seed = seed_pipeline(config, log)
+
+    dataset, ptids, image_to_ptid = load_full_dataset(config, log)
 
     # Mapping used later by evaluation to attach patient ids to the sample axis.
-    state["image_to_ptid"] = {int(iid): str(pid) for iid, pid in ptids}
+    state["image_to_ptid"] = image_to_ptid
 
     # Split is grouped by PTID to prevent patient-level leakage across partitions.
     datasets = ds.divide_dataset_by_patient_id(

util/progress.py

@@ -6,6 +6,12 @@ from textual.widgets import Label, ProgressBar
 
 logger = logging.getLogger(__name__)
 
+# Number of stacked progress-bar rows (levels 0 .. MAX_LEVELS-1). Deep enough for
+# the k-fold nesting: root(0) -> run_kfold(1) -> fold task(2) -> train/eval(3) ->
+# epochs/passes(4) -> batches(5). main.py builds exactly this many rows.
+MAX_LEVELS = 6
+
+
 class ProgressTracker:
     """Manages hierarchical progress bars for the UI thread."""
 
@@ -15,7 +21,7 @@ class ProgressTracker:
 
     def get_sub_tracker(self, title: str = "") -> "ProgressTracker":
         """Returns a sub-tracker for child processes."""
-        if self.level >= 4:
+        if self.level >= MAX_LEVELS - 1:
             return self  # Max depth reached
         sub = ProgressTracker(self.app, self.level + 1)
         if title:
@@ -88,7 +94,7 @@ class ProgressTracker:
 
     def _clear_subs(self) -> None:
         """Clears and hides all child progress bars (must run in UI thread)."""
-        for l in range(self.level + 1, 5):
+        for l in range(self.level + 1, MAX_LEVELS):
             try:
                 bar = self.app.query_one(f"#progress_bar_{l}", ProgressBar)
                 row = self.app.query_one(f"#progress_row_{l}", Horizontal)