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

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,
    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"

# 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,
) -> 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)

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

    image_ids, ptids, true_classes = _collect_sample_metadata(test_loader, image_to_ptid)
    log.info(
        f"Evaluating {len(model_paths)} {kind} model(s) on {len(image_ids)} samples, "
        f"{len(noise_levels)} noise level(s), {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=test_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,
            "config": _config_snapshot(config),
        }
    )
    out_path.parent.mkdir(parents=True, exist_ok=True)
    save_evaluation(dataset, str(out_path))
    log.info(
        f"Saved {kind} 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


# ---------------------------------------------------------------------------
# Pipeline task entry points
# ---------------------------------------------------------------------------
def evaluate_normal_task(
    track: ProgressTracker,
    log: PipelineLogger,
    config: Dict[str, Any],
    state: Dict[str, Any],
) -> None:
    """Step 4: evaluate the normal ensemble on clean data (noise=0.0, K=1)."""
    paths = state.get("normal_model_paths") or []
    if not paths:
        log.error("No normal models to evaluate (run training/load_models first).")
        return
    _evaluate_ensemble(
        model_paths=paths,
        kind=KIND_NORMAL,
        config=config,
        state=state,
        log=log,
        track=track,
        noise_levels=[0.0],
        n_passes=1,
        out_path=_out_path(config, "normal.nc"),
    )


def evaluate_bayesian_task(
    track: ProgressTracker,
    log: PipelineLogger,
    config: Dict[str, Any],
    state: Dict[str, Any],
) -> None:
    """Step 5: evaluate the Bayesian ensemble on clean data with MC uncertainty."""
    paths = state.get("bayesian_model_paths") or []
    if not paths:
        log.error("No Bayesian models to evaluate (run training/load_models first).")
        return
    _, mc_passes = _eval_cfg(config)
    _evaluate_ensemble(
        model_paths=paths,
        kind=KIND_BAYESIAN,
        config=config,
        state=state,
        log=log,
        track=track,
        noise_levels=[0.0],
        n_passes=mc_passes,
        out_path=_out_path(config, "bayesian.nc"),
    )


def evaluate_noisy_task(
    track: ProgressTracker,
    log: PipelineLogger,
    config: Dict[str, Any],
    state: Dict[str, Any],
) -> None:
    """Step 6: evaluate both ensembles across the configured Gaussian noise levels."""
    noise_levels, mc_passes = _eval_cfg(config)
    normal_paths = state.get("normal_model_paths") or []
    bayesian_paths = state.get("bayesian_model_paths") or []

    if not normal_paths and not bayesian_paths:
        log.error("No models to evaluate on noised data.")
        return

    if normal_paths:
        _evaluate_ensemble(
            model_paths=normal_paths,
            kind=KIND_NORMAL,
            config=config,
            state=state,
            log=log,
            track=track,
            noise_levels=noise_levels,
            n_passes=1,
            out_path=_out_path(config, "noisy_normal.nc"),
        )
    else:
        log.info("No normal models found; skipping normal noisy evaluation.")

    if bayesian_paths:
        _evaluate_ensemble(
            model_paths=bayesian_paths,
            kind=KIND_BAYESIAN,
            config=config,
            state=state,
            log=log,
            track=track,
            noise_levels=noise_levels,
            n_passes=mc_passes,
            out_path=_out_path(config, "noisy_bayesian.nc"),
        )
    else:
        log.info("No Bayesian models found; skipping Bayesian noisy evaluation.")