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