"""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.")