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