alnn_rewrite/tasks/__init__.py

"""Pipeline task registry and scenario definitions.

A *task* is a callable ``task(tracker, logger, config, state) -> None``. A
*scenario* is an ordered list of task ids the pipeline runs. Tasks communicate
through the mutable ``state`` dict (dataloaders, model paths, control events).

Pipeline stages (see ai/ARCHITECTURE.md):
    1. load_data        -> implemented
    2. train_regular    -> implemented
    3. train_bayesian   -> implemented
    4. evaluate_regular -> implemented (TEST split; *_val variants do VALIDATION)
    5. evaluate_bayesian-> implemented (TEST split; *_val variants do VALIDATION)
    6. evaluate_noisy   -> implemented (TEST split; *_val variant does VALIDATION)
       combine_evaluations -> pools the per-split evaluations into combined_*.nc
    7. run_analysis     -> reads evaluations (prefers combined_*), writes analysis/.
    load_models         -> implemented; discovers saved .pt ensembles on disk and
                           records their paths in ``state`` so evaluation is
                           decoupled from training (training frees models from VRAM
                           after saving). Evaluation loads one model at a time.
"""

import os
from typing import Any, Callable, Dict, List, Tuple, TypedDict

from . import analyze
from . import evaluate
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):
    task_name: str
    task_func: Callable[..., None]


class ScenarioEntry(TypedDict):
    label: str
    tasks: List[str]


PIPELINE_TASKS: Dict[str, TaskEntry] = {
    "load_data": {
        "task_name": "Load Image and ADNIMERGE",
        "task_func": load_data.load_data_task,
    },
    "train_regular": {
        "task_name": "Train Regular Models",
        "task_func": train_normal.train_normal_task,
    },
    "train_bayesian": {
        "task_name": "Train Bayesian Models",
        "task_func": train_bayesian.train_bayesian_task,
    },
    "load_models": {
        "task_name": "Load Saved Models",
        "task_func": load_models.load_models_task,
    },
    "evaluate_regular": {
        "task_name": "Evaluate Regular Models",
        "task_func": evaluate.evaluate_normal_task,
    },
    "evaluate_bayesian": {
        "task_name": "Evaluate Bayesian Models",
        "task_func": evaluate.evaluate_bayesian_task,
    },
    "evaluate_noisy": {
        "task_name": "Evaluate Models on Noised Data",
        "task_func": evaluate.evaluate_noisy_task,
    },
    "evaluate_regular_val": {
        "task_name": "Evaluate Regular Models (Validation)",
        "task_func": evaluate.evaluate_normal_val_task,
    },
    "evaluate_bayesian_val": {
        "task_name": "Evaluate Bayesian Models (Validation)",
        "task_func": evaluate.evaluate_bayesian_val_task,
    },
    "evaluate_noisy_val": {
        "task_name": "Evaluate Models on Noised Data (Validation)",
        "task_func": evaluate.evaluate_noisy_val_task,
    },
    "combine_evaluations": {
        "task_name": "Combine Split Evaluations",
        "task_func": evaluate.combine_evaluations_task,
    },
    "run_analysis": {
        "task_name": "Analyze Evaluations",
        "task_func": analyze.run_analysis_task,
    },
    "run_kfold": {
        "task_name": "K-Fold Cross-Validation",
        "task_func": kfold.run_kfold_task,
    },
}

# --------------------------------------------------------------------------
# Output artifact categories and overwrite protection
# --------------------------------------------------------------------------
# Each category maps to the work_dir subdirectories it owns and whether it is
# "protected" -- i.e. expensive/slow to regenerate, so overwriting it warrants a
# confirmation prompt. Analysis artifacts are cheap to regenerate and therefore
# unprotected: reruns of an analysis-only scenario proceed without a warning.
class ArtifactCategory(TypedDict):
    dirs: Tuple[str, ...]
    protected: bool
    desc: str


ARTIFACT_CATEGORIES: Dict[str, ArtifactCategory] = {
    "models": {
        "dirs": ("normal_models", "bayesian_models"),
        "protected": True,
        "desc": "trained model checkpoints",
    },
    "evaluations": {
        "dirs": ("evaluations",),
        "protected": True,
        "desc": "model evaluation outputs (netCDF)",
    },
    "kfold": {
        "dirs": ("folds",),
        "protected": True,
        "desc": "per-fold k-fold models & evaluations",
    },
    "analysis": {
        "dirs": (analyze.ANALYSIS_SUBDIR,),
        "protected": False,
        "desc": "analysis artifacts (regeneratable)",
    },
}

# Which artifact category each task writes to disk. Tasks not listed here (e.g.
# load_data, load_models) produce no persistent, overwriteable output.
TASK_WRITES = {
    "train_regular": ("models",),
    "train_bayesian": ("models",),
    "evaluate_regular": ("evaluations",),
    "evaluate_bayesian": ("evaluations",),
    "evaluate_noisy": ("evaluations",),
    "evaluate_regular_val": ("evaluations",),
    "evaluate_bayesian_val": ("evaluations",),
    "evaluate_noisy_val": ("evaluations",),
    "combine_evaluations": ("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"),
}


def scenario_output_categories(scenario_id: str) -> List[str]:
    """Return the artifact categories a scenario writes, in first-seen order."""
    categories: List[str] = []
    for task_id in SCENARIOS[scenario_id]["tasks"]:
        for category in TASK_WRITES.get(task_id, ()):
            if category not in categories:
                categories.append(category)
    return categories


def protected_output_conflicts(
    work_dir: str, scenario_id: str
) -> List[Tuple[str, str]]:
    """Return ``(category, subdir)`` for PROTECTED outputs the scenario writes
    that already exist and are non-empty. Empty result => safe to run without a
    destructive-overwrite prompt (only unprotected/regeneratable outputs, if any,
    would be replaced)."""
    conflicts: List[Tuple[str, str]] = []
    for category in scenario_output_categories(scenario_id):
        meta = ARTIFACT_CATEGORIES.get(category)
        if not meta or not meta["protected"]:
            continue
        for subdir in meta["dirs"]:
            path = os.path.join(work_dir, subdir)
            if os.path.isdir(path) and os.listdir(path):
                conflicts.append((category, subdir))
    return conflicts


def describe_protected_conflicts(work_dir: str, scenario_id: str) -> List[str]:
    """Human-readable one-liners for each protected overwrite conflict."""
    lines: List[str] = []
    for category, subdir in protected_output_conflicts(work_dir, scenario_id):
        desc = ARTIFACT_CATEGORIES[category]["desc"]
        lines.append(f"{subdir}/ — {desc}")
    return lines

# Evaluation of both held-out splits (test + val), pooled for analysis. Shared by
# the train and load scenarios so they stay in lockstep.
_EVAL_ALL_SPLITS = [
    "evaluate_regular",
    "evaluate_bayesian",
    "evaluate_regular_val",
    "evaluate_bayesian_val",
]
_EVAL_ALL_SPLITS_NOISY = ["evaluate_noisy", "evaluate_noisy_val"]

SCENARIOS: Dict[str, ScenarioEntry] = {
    "scen_train_all": {
        "label": "1. Train, Evaluate (test+val), Noise, & Analyze",
        "tasks": [
            "load_data",
            "train_regular",
            "train_bayesian",
            "load_models",
            *_EVAL_ALL_SPLITS,
            *_EVAL_ALL_SPLITS_NOISY,
            "combine_evaluations",
            "run_analysis",
        ],
    },
    "scen_load_all": {
        "label": "2. Load, Evaluate (test+val), Noise, & Analyze",
        "tasks": [
            "load_data",
            "load_models",
            *_EVAL_ALL_SPLITS,
            *_EVAL_ALL_SPLITS_NOISY,
            "combine_evaluations",
            "run_analysis",
        ],
    },
    "scen_load_eval": {
        "label": "3. Load & Evaluate (test+val, Skip Noise) & Analyze",
        "tasks": [
            "load_data",
            "load_models",
            *_EVAL_ALL_SPLITS,
            "combine_evaluations",
            "run_analysis",
        ],
    },
    "scen_analyze": {
        "label": "4. Analyze Existing Evaluations",
        "tasks": [
            "run_analysis",
        ],
    },
    # Set aside for now (kept working, not part of the default flow): full
    # patient-grouped cross-validation. Costs ~k x a normal train+eval run.
    "scen_kfold": {
        "label": "5. K-Fold Cross-Validation (train + OOF eval + analyze)",
        "tasks": [
            "run_kfold",
            "run_analysis",
        ],
    },
}