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

import xarray as xr

from evaluation.schema import (
    CLASS_NAMES,
    EvaluationAccumulator,
    concat_evaluations,
    concat_noise_levels,
    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,
    noise_relative: bool = True,
) -> 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.

    Args:
        noise_relative: Scale the Gaussian noise by each image's own intensity
            standard deviation, so ``noise_level`` means "fraction of image std".
            Required for this dataset: the PET volumes are unnormalised (std in
            the thousands), so an absolute sigma of order 1 is numerically
            invisible. Per-image (not per-batch) scaling keeps the corruption
            level identical for every sample regardless of its intensity range,
            and independent of batch composition.
    """
    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
                    )
                    if noise_relative:
                        # Per-image std over the voxel dims -> sigma is a fraction
                        # of that image's own intensity spread.
                        scale = mri.std(
                            dim=tuple(range(1, mri.ndim)), keepdim=True
                        )
                    else:
                        scale = 1.0
                    mri = mri + noise * (noise_level * scale)

                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)

    noise_relative = bool(
        config.get("evaluation", {}).get("noise_relative", True)
    )

    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}",
                noise_relative=noise_relative,
            )
            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,
            # 1 => noise_level is a fraction of each image's intensity std;
            # 0 => it is an absolute value in raw voxel units.
            "noise_relative": int(noise_relative),
            "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, extra: bool = False
) -> str:
    """Canonical evaluation filename for a (split, kind, noisy, extra) combination.

    The test split keeps its historical un-prefixed names (``normal.nc``,
    ``noisy_bayesian.nc``); other splits are prefixed (``val_normal.nc``).
    ``extra`` marks an *additional* noise sweep computed in a later run
    (``extra_noisy_normal.nc``), which is merged onto the base sweep's
    ``noise_level`` axis by :func:`combine_evaluations_task`.
    """
    prefix = "" if split == SPLIT_TEST else f"{split}_"
    return f"{prefix}{'extra_' if extra else ''}{'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,
    extra: bool = False,
) -> None:
    """Evaluate both ensembles across a noise sweep for ``split`` (step 6).

    With ``extra=True`` the sweep uses ``evaluation.extra_noise_levels`` and is
    written to separate ``extra_*`` files, so additional sigmas can be computed
    later and merged without recomputing the original sweep.
    """
    noise_levels, mc_passes = _eval_cfg(config)
    if extra:
        noise_levels = list(
            config.get("evaluation", {}).get("extra_noise_levels", [])
        )
        if not noise_levels:
            log.error(
                "No evaluation.extra_noise_levels configured; nothing to add. "
                "Set them in config.toml to extend the sweep."
            )
            return
    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, extra=extra)
            ),
            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 evaluate_noisy_extra_task(track, log, config, state) -> None:
    """ADDITIONAL noise levels on the TEST split (extends an existing sweep)."""
    _evaluate_noisy(track, log, config, state, SPLIT_TEST, extra=True)


def evaluate_noisy_extra_val_task(track, log, config, state) -> None:
    """ADDITIONAL noise levels on the VALIDATION split."""
    _evaluate_noisy(track, log, config, state, SPLIT_VAL, extra=True)


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

    Two independent merges happen here:

    1. **Noise levels** (``concat_noise_levels``) -- within one split, the base
       sweep and any ``extra_*`` sweep are unioned along ``noise_level``. This is
       what lets extra sigmas be computed later and folded in without recomputing
       the original sweep.
    2. **Splits** (``concat_evaluations``) -- the configured splits (default test
       + val) are pooled along ``sample``. They are patient-disjoint by
       construction, so pooling enlarges the evaluation set with no duplicated
       samples, giving tighter accuracy/uncertainty estimates.

    Analysis prefers the resulting ``combined_*`` files over single-split ones.
    """
    eval_dir = pl.Path(config["work_dir"]) / _EVAL_SUBDIR
    splits = list(
        config.get("evaluation", {}).get("combine_splits", [SPLIT_TEST, SPLIT_VAL])
    )

    written = 0
    for kind in (KIND_NORMAL, KIND_BAYESIAN):
        for noisy in (False, True):
            per_split: List[xr.Dataset] = []
            sources: List[str] = []
            merged_noise = False
            try:
                for split in splits:
                    # Stage 1: base + extra sweeps for this split.
                    candidates = [
                        eval_dir / eval_filename(split, kind, noisy=noisy),
                    ]
                    if noisy:
                        candidates.append(
                            eval_dir
                            / eval_filename(split, kind, noisy=True, extra=True)
                        )
                    present = [p for p in candidates if p.exists()]
                    if not present:
                        continue

                    parts = [load_evaluation(str(p)) for p in present]
                    if len(parts) > 1:
                        merged = concat_noise_levels(parts)
                        for part in parts:
                            part.close()
                        merged_noise = True
                    else:
                        merged = parts[0]
                    per_split.append(merged)
                    sources.extend(p.name for p in present)

                if not per_split:
                    continue
                # Nothing gained from writing a copy of a single source.
                if len(per_split) < 2 and not merged_noise:
                    continue

                # Guard: the same patient must never appear in two pooled splits
                # (that would double-count them).
                seen: set[str] = set()
                overlap: set[str] = set()
                for part in per_split:
                    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

                # Stage 2: pool the splits along `sample`.
                combined = (
                    concat_evaluations(per_split)
                    if len(per_split) > 1
                    else per_split[0]
                )
                combined.attrs.update(
                    {
                        "split": "+".join(splits[: len(per_split)]),
                        "combined_from": ", ".join(sources),
                        "n_sources": len(sources),
                    }
                )
                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 x "
                    f"{combined.sizes['noise_level']} noise level(s) from "
                    f"{', '.join(sources)})."
                )
            finally:
                for part in per_split:
                    part.close()

    if written == 0:
        log.error(
            "combine_evaluations: nothing to merge. Run the test/validation "
            "evaluation steps (and any extra noise steps) first."
        )