alnn_rewrite/tasks/evaluate.py

"""Evaluation tasks (pipeline steps 4-6).

All three tasks share one core, :func:`_evaluate_ensemble`, which loads each
ensemble member one at a time, runs Monte-Carlo forward passes over the (fixed,
unshuffled) test set at each requested noise level, and writes the results to a
netCDF file in the locked-down evaluation schema (:mod:`evaluation.schema`).

    step 4  evaluate_regular  -> normal ensemble, noise=[0.0], K=1
    step 5  evaluate_bayesian -> Bayesian ensemble, noise=[0.0], K=mc_passes
    step 6  evaluate_noisy    -> both ensembles across config noise_levels

Every evaluator is parameterized by dataset *split*, so the same code produces
test-split and validation-split evaluations (``*_val_task`` variants). Both
loaders are unshuffled, so sample order is stable in each. Finally,
``combine_evaluations_task`` pools the per-split files into ``combined_*.nc``
(test + val by default) -- the splits are patient-disjoint, so pooling simply
enlarges the evaluation set that analysis draws on.

Determinism
-----------
The test loader is unshuffled, so sample order is stable and aligns with the
netCDF ``sample`` axis. Added Gaussian noise is drawn from a generator seeded by
``derive_seed(base_seed, noise_level_index, stream=_NOISE_SEED_STREAM)`` and
re-seeded at the start of every MC pass, so:
  * the noise perturbation for a given (sample, level) is identical across the K
    MC passes (only the sampled BNN weights vary), and
  * every model sees the exact same noised inputs at a given level (fair
    comparison across the ensemble).
"""

import pathlib as pl
from typing import Any, Dict, List, Sequence

import numpy as np
import toml
import torch

from evaluation.schema import (
    CLASS_NAMES,
    EvaluationAccumulator,
    concat_evaluations,
    load_evaluation,
    save_evaluation,
)
from model.factory import KIND_BAYESIAN, KIND_NORMAL, load_model
from util.control import check_control_events, release_torch_memory
from util.progress import ProgressTracker
from util.seeding import derive_seed
from util.ui_logger import PipelineLogger

# RNG stream id for input noise (0=normal training, 1=Bayesian training).
_NOISE_SEED_STREAM = 2

# Subdirectory (under work_dir) where evaluation netCDF files are written.
_EVAL_SUBDIR = "evaluations"

# Evaluable dataset splits -> the ``state`` key holding their DataLoader. Both
# val and test loaders are unshuffled, so their sample order is stable and can be
# aligned with the netCDF ``sample`` axis. (``train`` is available for
# completeness but is NOT part of the default pipeline: its samples were fit on,
# so its accuracy is optimistic and not a generalization estimate.)
SPLIT_TEST = "test"
SPLIT_VAL = "val"
SPLIT_TRAIN = "train"
_SPLIT_LOADERS = {
    SPLIT_TEST: "test_loader",
    SPLIT_VAL: "val_loader",
    SPLIT_TRAIN: "train_loader",
}

# Fallbacks if a config predates the [evaluation] section.
_DEFAULT_NOISE_LEVELS: List[float] = [0.0, 0.02, 0.05, 0.1, 0.2]
_DEFAULT_MC_PASSES = 30


def _events(state: Dict[str, Any]):
    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
    return stop_event, pause_event


def _collect_sample_metadata(loader, image_to_ptid: Dict[int, str]):
    """Return ordered ``(image_ids, ptids, true_classes)`` for the test set.

    The order matches iteration of the (unshuffled) loader and therefore the
    order used when accumulating model outputs.
    """
    image_ids: List[int] = []
    true_classes: List[str] = []
    for _mri, _xls, targets, ids in loader:
        for i in range(int(targets.shape[0])):
            image_ids.append(int(ids[i]))
            true_classes.append(CLASS_NAMES[int(targets[i].argmax())])
    ptids = [str(image_to_ptid.get(iid, "unknown")) for iid in image_ids]
    return image_ids, ptids, true_classes


def _run_mc_passes(
    model: torch.nn.Module,
    loader,
    device: str,
    noise_level: float,
    n_passes: int,
    noise_seed: int,
    n_classes: int,
    parent_tracker: ProgressTracker,
    stop_event,
    pause_event,
    label: str,
) -> np.ndarray:
    """Run ``n_passes`` forward passes over the test set, returning ``(K, S, C)``.

    For a deterministic model use ``n_passes == 1``. Bayesian models sample fresh
    weights each pass (their reparameterization layers stay stochastic in eval
    mode), producing the Monte-Carlo ensemble the uncertainty measures need.
    """
    n_samples = len(loader.dataset)
    mc = np.zeros((n_passes, n_samples, n_classes), dtype=np.float32)

    pass_tracker = parent_tracker.get_sub_tracker(label)
    pass_tracker.update(total=n_passes, advance=0)

    with torch.no_grad():
        for k in range(n_passes):
            # Re-seed per pass so the noise realization is identical across MC
            # passes (and across models) for a given (sample, noise_level).
            noise_gen = (
                torch.Generator(device=device).manual_seed(int(noise_seed))
                if noise_level > 0.0
                else None
            )

            batch_tracker = pass_tracker.get_sub_tracker("Batches")
            batch_tracker.update(total=len(loader), advance=0)

            idx = 0
            for mri, xls, _targets, _ids in loader:
                check_control_events(stop_event, pause_event)
                mri = mri.to(device, non_blocking=True)
                xls = xls.to(device, non_blocking=True)

                if noise_gen is not None:
                    noise = torch.randn(
                        mri.shape, generator=noise_gen, device=device, dtype=mri.dtype
                    )
                    mri = mri + noise * noise_level

                outputs = model((mri, xls))
                bs = int(outputs.shape[0])
                mc[k, idx : idx + bs] = outputs.detach().cpu().numpy()
                idx += bs
                batch_tracker.update(advance=1)

            pass_tracker.update(advance=1)

    return mc


def _evaluate_ensemble(
    *,
    model_paths: Sequence[pl.Path],
    kind: str,
    config: Dict[str, Any],
    state: Dict[str, Any],
    log: PipelineLogger,
    track: ProgressTracker,
    noise_levels: Sequence[float],
    n_passes: int,
    out_path: pl.Path,
    split: str = SPLIT_TEST,
) -> None:
    """Evaluate one ensemble and write a schema-compliant netCDF to ``out_path``."""
    device = config["training"]["device"]
    n_classes = int(config["data"]["num_classes"])
    base_seed = int(config["data"]["seed"])
    stop_event, pause_event = _events(state)

    loader = state[_SPLIT_LOADERS[split]]
    image_to_ptid: Dict[int, str] = state.get("image_to_ptid", {})

    image_ids, ptids, true_classes = _collect_sample_metadata(loader, image_to_ptid)
    log.info(
        f"Evaluating {len(model_paths)} {kind} model(s) on {len(image_ids)} "
        f"{split} samples, {len(noise_levels)} noise level(s), "
        f"{n_passes} MC pass(es) each."
    )

    accumulator = EvaluationAccumulator(
        image_ids, ptids, true_classes, noise_levels=list(noise_levels)
    )

    models_tracker = track.get_sub_tracker(f"{kind.capitalize()} models")
    models_tracker.update(total=len(model_paths), advance=0)

    for model_index, path in enumerate(model_paths):
        check_control_events(stop_event, pause_event)
        log.info(f"Loading {kind} model {model_index + 1}/{len(model_paths)}: {path.name}")
        model = load_model(path, kind, config, device)

        for level_index, noise in enumerate(noise_levels):
            noise_seed = derive_seed(base_seed, level_index, stream=_NOISE_SEED_STREAM)
            mc_probs = _run_mc_passes(
                model=model,
                loader=loader,
                device=device,
                noise_level=float(noise),
                n_passes=n_passes,
                noise_seed=noise_seed,
                n_classes=n_classes,
                parent_tracker=models_tracker,
                stop_event=stop_event,
                pause_event=pause_event,
                label=f"model {model_index + 1} | sigma={noise}",
            )
            accumulator.add(
                model_index=model_index,
                model_kind=kind,
                noise_level=float(noise),
                mc_probs=mc_probs,
            )

        del model
        release_torch_memory(device)
        models_tracker.update(advance=1)

    dataset = accumulator.to_dataset(
        attrs={
            "seed": base_seed,
            "n_mc": n_passes,
            "ensemble_kind": kind,
            "split": split,
            "config": _config_snapshot(config),
        }
    )
    out_path.parent.mkdir(parents=True, exist_ok=True)
    save_evaluation(dataset, str(out_path))
    log.info(
        f"Saved {kind} {split} evaluation to {out_path} "
        f"(models={dataset.sizes['model']}, samples={dataset.sizes['sample']}, "
        f"noise_levels={dataset.sizes['noise_level']})."
    )


def _config_snapshot(config: Dict[str, Any]) -> str:
    try:
        return toml.dumps(config)
    except Exception:
        return repr(config)


def _eval_cfg(config: Dict[str, Any]):
    eval_cfg = config.get("evaluation", {})
    noise_levels = list(eval_cfg.get("noise_levels", _DEFAULT_NOISE_LEVELS))
    mc_passes = int(eval_cfg.get("mc_passes", _DEFAULT_MC_PASSES))
    return noise_levels, mc_passes


def _out_path(config: Dict[str, Any], name: str) -> pl.Path:
    return pl.Path(config["work_dir"]) / _EVAL_SUBDIR / name


def eval_filename(split: str, kind: str, noisy: bool = False) -> str:
    """Canonical evaluation filename for a (split, kind, noisy) combination.

    The test split keeps its historical un-prefixed names (``normal.nc``,
    ``noisy_bayesian.nc``); other splits are prefixed (``val_normal.nc``).
    """
    prefix = "" if split == SPLIT_TEST else f"{split}_"
    return f"{prefix}{'noisy_' if noisy else ''}{kind}.nc"


def _evaluate_clean(
    track: ProgressTracker,
    log: PipelineLogger,
    config: Dict[str, Any],
    state: Dict[str, Any],
    kind: str,
    split: str,
) -> None:
    """Evaluate one ensemble on clean data for ``split`` (steps 4/5)."""
    key = "normal_model_paths" if kind == KIND_NORMAL else "bayesian_model_paths"
    paths = state.get(key) or []
    if not paths:
        log.error(f"No {kind} models to evaluate (run training/load_models first).")
        return
    n_passes = 1 if kind == KIND_NORMAL else _eval_cfg(config)[1]
    _evaluate_ensemble(
        model_paths=paths,
        kind=kind,
        config=config,
        state=state,
        log=log,
        track=track,
        noise_levels=[0.0],
        n_passes=n_passes,
        out_path=_out_path(config, eval_filename(split, kind)),
        split=split,
    )


def _evaluate_noisy(
    track: ProgressTracker,
    log: PipelineLogger,
    config: Dict[str, Any],
    state: Dict[str, Any],
    split: str,
) -> None:
    """Evaluate both ensembles across the noise sweep for ``split`` (step 6)."""
    noise_levels, mc_passes = _eval_cfg(config)
    any_models = False
    for kind, key, n_passes in (
        (KIND_NORMAL, "normal_model_paths", 1),
        (KIND_BAYESIAN, "bayesian_model_paths", mc_passes),
    ):
        paths = state.get(key) or []
        if not paths:
            log.info(f"No {kind} models found; skipping {kind} noisy evaluation.")
            continue
        any_models = True
        _evaluate_ensemble(
            model_paths=paths,
            kind=kind,
            config=config,
            state=state,
            log=log,
            track=track,
            noise_levels=noise_levels,
            n_passes=n_passes,
            out_path=_out_path(config, eval_filename(split, kind, noisy=True)),
            split=split,
        )
    if not any_models:
        log.error("No models to evaluate on noised data.")


# ---------------------------------------------------------------------------
# Pipeline task entry points
# ---------------------------------------------------------------------------
def evaluate_normal_task(track, log, config, state) -> None:
    """Step 4: normal ensemble on the clean TEST split."""
    _evaluate_clean(track, log, config, state, KIND_NORMAL, SPLIT_TEST)


def evaluate_bayesian_task(track, log, config, state) -> None:
    """Step 5: Bayesian ensemble (MC uncertainty) on the clean TEST split."""
    _evaluate_clean(track, log, config, state, KIND_BAYESIAN, SPLIT_TEST)


def evaluate_noisy_task(track, log, config, state) -> None:
    """Step 6: both ensembles across the noise sweep on the TEST split."""
    _evaluate_noisy(track, log, config, state, SPLIT_TEST)


def evaluate_normal_val_task(track, log, config, state) -> None:
    """Normal ensemble on the clean VALIDATION split."""
    _evaluate_clean(track, log, config, state, KIND_NORMAL, SPLIT_VAL)


def evaluate_bayesian_val_task(track, log, config, state) -> None:
    """Bayesian ensemble (MC uncertainty) on the clean VALIDATION split."""
    _evaluate_clean(track, log, config, state, KIND_BAYESIAN, SPLIT_VAL)


def evaluate_noisy_val_task(track, log, config, state) -> None:
    """Both ensembles across the noise sweep on the VALIDATION split."""
    _evaluate_noisy(track, log, config, state, SPLIT_VAL)


def combine_evaluations_task(
    track: ProgressTracker,
    log: PipelineLogger,
    config: Dict[str, Any],
    state: Dict[str, Any],
) -> None:
    """Pool per-split evaluations into ``combined_*.nc`` for analysis.

    Concatenates the configured splits (default test + val) along the ``sample``
    axis. The splits are patient-disjoint by construction, so pooling them yields
    a larger evaluation set with no duplicated samples -- which is the point:
    more evaluation samples means tighter accuracy/uncertainty estimates.
    Analysis prefers these pooled files over the single-split ones.
    """
    eval_dir = pl.Path(config["work_dir"]) / _EVAL_SUBDIR
    splits = list(
        config.get("evaluation", {}).get("combine_splits", [SPLIT_TEST, SPLIT_VAL])
    )
    if len(splits) < 2:
        log.info(
            f"combine_evaluations: only {splits} configured; nothing to pool."
        )
        return

    written = 0
    for kind in (KIND_NORMAL, KIND_BAYESIAN):
        for noisy in (False, True):
            paths = [
                eval_dir / eval_filename(split, kind, noisy=noisy) for split in splits
            ]
            present = [p for p in paths if p.exists()]
            if len(present) < 2:
                continue

            parts = [load_evaluation(str(p)) for p in present]
            try:
                # Guard against a mis-specified split: the same patient must not
                # appear in two pooled parts (that would double-count them).
                seen: set[str] = set()
                overlap: set[str] = set()
                for part in parts:
                    part_ptids = {str(p) for p in np.atleast_1d(part["ptid"].values)}
                    overlap |= seen & part_ptids
                    seen |= part_ptids
                if overlap:
                    log.error(
                        f"combine_evaluations: {len(overlap)} patient(s) appear in "
                        f"more than one split ({sorted(overlap)[:3]}...); refusing "
                        f"to pool {kind} (noisy={noisy})."
                    )
                    continue

                combined = concat_evaluations(
                    parts,
                    attrs={
                        "split": "+".join(splits),
                        "combined_from": ", ".join(p.name for p in present),
                        "n_sources": len(present),
                    },
                )
                out_name = f"combined_{'noisy_' if noisy else ''}{kind}.nc"
                save_evaluation(combined, str(eval_dir / out_name))
                written += 1
                log.info(
                    f"combine_evaluations: wrote {out_name} "
                    f"({combined.sizes['sample']} samples from "
                    f"{', '.join(p.name for p in present)})."
                )
            finally:
                for part in parts:
                    part.close()

    if written == 0:
        log.error(
            "combine_evaluations: found no split pairs to pool. Run the test and "
            "validation evaluation steps first."
        )