"""Canonical netCDF / xarray schema for pipeline evaluations (steps 4-6). This module is the single source of truth for how model outputs are stored on disk. Locking it down means every evaluator (normal, Bayesian, noisy) writes the exact same dimensions, variables, coordinates, and dtypes, so the analysis stage (step 7) only ever has to understand one format. -------------------------------------------------------------------------------- Schema (xarray.Dataset) -------------------------------------------------------------------------------- Dimensions model ensemble member index (0..M-1) noise_level Gaussian noise sigma applied to the image; ALWAYS present. The clean baseline is stored as noise_level == 0.0, so noised and un-noised evaluations share one file per model family. sample fixed test-set sample index (from a NON-shuffled loader) class_name output classes, ordered ["AD", "NL"] mc_pass (optional) raw Monte-Carlo draw index, Bayesian only Data variables prob (model, noise_level, sample, class_name) float32 Mean predictive probability. For Bayesian models this is the mean over MC passes; for normal models it is the single softmax output. pred (model, noise_level, sample) int8 argmax over class_name of ``prob``. predictive_entropy (model, noise_level, sample) float32 Total/predictive uncertainty: entropy of the mean predictive distribution (bayesian_torch.predictive_entropy). mutual_information (model, noise_level, sample) float32 Epistemic/model uncertainty (bayesian_torch.mutual_information). Identically ~0 for deterministic (K=1) models. correct (model, noise_level, sample) int8 (0/1) Whether ``pred`` matches the true class. mc_prob (model, noise_level, sample, class_name, mc_pass) float32 OPTIONAL raw per-draw probabilities, only if retained. Coordinates model_kind (model) "normal" | "bayesian" image_id (sample) ADNI Image Data ID (join key to clinical data) ptid (sample) patient id (enables patient-level analysis) true_class (sample) ground-truth class label, one of class_name noise_level (noise_level) float sigma values class_name (class_name) ["AD", "NL"] Attributes schema_version, created (ISO-8601 UTC), seed, git_commit, n_mc, config (TOML string snapshot) """ from __future__ import annotations import datetime as _dt import subprocess from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple import numpy as np import xarray as xr # bayesian-torch's own uncertainty utilities, so definitions match the library # the Bayesian models were trained with (operate on numpy mc_preds). from bayesian_torch.utils.util import mutual_information as _bt_mutual_information from bayesian_torch.utils.util import predictive_entropy as _bt_predictive_entropy SCHEMA_VERSION = "1.0" CLASS_NAMES: Tuple[str, ...] = ("AD", "NL") MODEL_KIND_NORMAL = "normal" MODEL_KIND_BAYESIAN = "bayesian" # netCDF engine that supports the full type set (float32, int8, VLEN strings). _NETCDF_ENGINE = "netcdf4" # Encodings applied on write. netCDF has no native bool, so ``correct`` is int8. _ENCODINGS: Dict[str, Dict[str, Any]] = { "prob": {"dtype": "float32", "zlib": True, "complevel": 4}, "pred": {"dtype": "int8", "zlib": True, "complevel": 4}, "predictive_entropy": {"dtype": "float32", "zlib": True, "complevel": 4}, "mutual_information": {"dtype": "float32", "zlib": True, "complevel": 4}, "correct": {"dtype": "int8", "zlib": True, "complevel": 4}, "mc_prob": {"dtype": "float32", "zlib": True, "complevel": 4}, } # --------------------------------------------------------------------------- # Uncertainty computation # --------------------------------------------------------------------------- def compute_uncertainty( mc_probs: np.ndarray, ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: """Reduce Monte-Carlo predictive probabilities to summary statistics. Args: mc_probs: Array of shape ``(K, S, C)`` -- K Monte-Carlo passes, S samples, C classes. For a deterministic model use ``K == 1``. Returns: ``(mean_prob, predictive_entropy, mutual_information)`` with shapes ``(S, C)``, ``(S,)`` and ``(S,)`` respectively. ``mutual_information`` is ~0 when ``K == 1``. """ mc = np.asarray(mc_probs, dtype=np.float64) if mc.ndim != 3: raise ValueError( f"mc_probs must be (K, S, C); got shape {mc.shape} with ndim {mc.ndim}." ) mean_prob = mc.mean(axis=0) pred_entropy = _bt_predictive_entropy(mc) # mutual_information == predictive_entropy - mean(per-pass entropy). With a # single pass the two entropy terms are identical, so this is exactly 0. mut_info = _bt_mutual_information(mc) return ( mean_prob.astype(np.float32), np.asarray(pred_entropy, dtype=np.float32), np.asarray(mut_info, dtype=np.float32), ) # --------------------------------------------------------------------------- # Dataset construction # --------------------------------------------------------------------------- def build_evaluation_dataset( *, prob: np.ndarray, predictive_entropy: np.ndarray, mutual_information: np.ndarray, model_kinds: Sequence[str], noise_levels: Sequence[float], image_ids: Sequence[int], ptids: Sequence[str], true_classes: Sequence[str], mc_prob: Optional[np.ndarray] = None, attrs: Optional[Mapping[str, Any]] = None, ) -> xr.Dataset: """Assemble a schema-compliant ``xarray.Dataset``. Args: prob: ``(M, L, S, C)`` mean predictive probabilities. predictive_entropy: ``(M, L, S)`` total uncertainty. mutual_information: ``(M, L, S)`` model uncertainty. model_kinds: length-M sequence of "normal"/"bayesian". noise_levels: length-L sequence of sigma values (include 0.0 baseline). image_ids: length-S sequence of ADNI Image Data IDs. ptids: length-S sequence of patient ids. true_classes: length-S ground-truth labels drawn from ``CLASS_NAMES``. mc_prob: optional ``(M, L, S, C, K)`` raw MC draws. attrs: optional extra dataset attributes (merged over the defaults). Returns: A validated ``xarray.Dataset`` following the module schema. """ prob = np.asarray(prob, dtype=np.float32) predictive_entropy = np.asarray(predictive_entropy, dtype=np.float32) mutual_information = np.asarray(mutual_information, dtype=np.float32) m, ell, s, c = _validate_shapes( prob, predictive_entropy, mutual_information, model_kinds, noise_levels, image_ids, ptids, true_classes, ) # Derived variables: predicted class index and correctness. pred = prob.argmax(axis=-1).astype(np.int8) # (M, L, S) class_index = {name: i for i, name in enumerate(CLASS_NAMES)} try: true_idx = np.array([class_index[t] for t in true_classes], dtype=np.int8) except KeyError as exc: raise ValueError( f"true_classes contains a label not in CLASS_NAMES {CLASS_NAMES}: {exc}" ) from None # Broadcast true index (S,) against pred (M, L, S). correct = (pred == true_idx[None, None, :]).astype(np.int8) data_vars: Dict[str, Any] = { "prob": (("model", "noise_level", "sample", "class_name"), prob), "pred": (("model", "noise_level", "sample"), pred), "predictive_entropy": ( ("model", "noise_level", "sample"), predictive_entropy, ), "mutual_information": ( ("model", "noise_level", "sample"), mutual_information, ), "correct": (("model", "noise_level", "sample"), correct), } if mc_prob is not None: mc_prob = np.asarray(mc_prob, dtype=np.float32) if mc_prob.ndim != 5 or mc_prob.shape[:4] != (m, ell, s, c): raise ValueError( f"mc_prob must be (M, L, S, C, K)=({m}, {ell}, {s}, {c}, K); " f"got {mc_prob.shape}." ) data_vars["mc_prob"] = ( ("model", "noise_level", "sample", "class_name", "mc_pass"), mc_prob, ) coords: Dict[str, Any] = { "model": np.arange(m, dtype=np.int32), "model_kind": ("model", np.asarray(list(model_kinds), dtype=object)), "noise_level": np.asarray(list(noise_levels), dtype=np.float32), "sample": np.arange(s, dtype=np.int32), "image_id": ("sample", np.asarray(list(image_ids), dtype=np.int64)), "ptid": ("sample", np.asarray(list(ptids), dtype=object)), "true_class": ("sample", np.asarray(list(true_classes), dtype=object)), "class_name": np.asarray(CLASS_NAMES, dtype=object), } ds = xr.Dataset(data_vars=data_vars, coords=coords) ds.attrs.update(_default_attrs()) if attrs: ds.attrs.update(dict(attrs)) return ds def _validate_shapes( prob: np.ndarray, predictive_entropy: np.ndarray, mutual_information: np.ndarray, model_kinds: Sequence[str], noise_levels: Sequence[float], image_ids: Sequence[int], ptids: Sequence[str], true_classes: Sequence[str], ) -> Tuple[int, int, int, int]: if prob.ndim != 4: raise ValueError(f"prob must be (M, L, S, C); got shape {prob.shape}.") m, ell, s, c = prob.shape if c != len(CLASS_NAMES): raise ValueError( f"prob's class axis has length {c}, expected {len(CLASS_NAMES)} " f"for CLASS_NAMES {CLASS_NAMES}." ) for name, arr in ( ("predictive_entropy", predictive_entropy), ("mutual_information", mutual_information), ): if arr.shape != (m, ell, s): raise ValueError( f"{name} must be (M, L, S)=({m}, {ell}, {s}); got {arr.shape}." ) for name, seq, expected in ( ("model_kinds", model_kinds, m), ("noise_levels", noise_levels, ell), ("image_ids", image_ids, s), ("ptids", ptids, s), ("true_classes", true_classes, s), ): if len(seq) != expected: raise ValueError( f"{name} must have length {expected}; got {len(seq)}." ) unknown_kinds = set(model_kinds) - {MODEL_KIND_NORMAL, MODEL_KIND_BAYESIAN} if unknown_kinds: raise ValueError( f"model_kinds contains unknown values {unknown_kinds}; " f"expected only {MODEL_KIND_NORMAL!r}/{MODEL_KIND_BAYESIAN!r}." ) return m, ell, s, c def _default_attrs() -> Dict[str, Any]: return { "schema_version": SCHEMA_VERSION, "created": _dt.datetime.now(_dt.timezone.utc).isoformat(), "git_commit": _git_commit(), } def _git_commit() -> str: try: out = subprocess.run( ["git", "rev-parse", "HEAD"], capture_output=True, text=True, timeout=5, ) return out.stdout.strip() if out.returncode == 0 else "unknown" except Exception: return "unknown" # --------------------------------------------------------------------------- # I/O # --------------------------------------------------------------------------- def save_evaluation(ds: xr.Dataset, path: str) -> None: """Write an evaluation dataset to netCDF using the locked-down encodings.""" encoding = {name: _ENCODINGS[name] for name in ds.data_vars if name in _ENCODINGS} # pyright: ignore ds.to_netcdf(path, engine=_NETCDF_ENGINE, encoding=encoding) def load_evaluation(path: str) -> xr.Dataset: """Load an evaluation dataset previously written by :func:`save_evaluation`.""" return xr.open_dataset(path, engine=_NETCDF_ENGINE) def concat_evaluations( datasets: Sequence[xr.Dataset], attrs: Optional[Mapping[str, Any]] = None ) -> xr.Dataset: """Pool schema-compliant evaluations along the ``sample`` axis. Used to merge evaluations of DISJOINT sample sets (e.g. the test and validation splits, or k-fold held-out slices) into one dataset covering all of them. The parts must share the ``model``, ``noise_level`` and ``class_name`` axes; only ``sample`` grows. The per-part ``sample`` index is re-based to a contiguous range, while ``image_id``/``ptid``/``true_class`` travel with their rows. Args: datasets: Evaluations to pool, each over a distinct set of samples. attrs: Attributes merged over the result (e.g. ``{"split": "test+val"}``). """ parts = list(datasets) if not parts: raise ValueError("concat_evaluations requires at least one dataset.") combined = xr.concat(parts, dim="sample") combined = combined.assign_coords(sample=np.arange(combined.sizes["sample"])) if attrs: combined.attrs.update(dict(attrs)) return combined # --------------------------------------------------------------------------- # Incremental builder # --------------------------------------------------------------------------- class EvaluationAccumulator: """Collect per-(model, noise_level) results, then finalize to a Dataset. This is the intended entry point for the (not-yet-implemented) evaluation tasks. Sample-axis metadata (image ids, ptids, true classes) is fixed by the non-shuffled test loader and supplied once; each ``add`` call contributes one model's Monte-Carlo probabilities at one noise level. Example (pseudocode for a future evaluate task):: acc = EvaluationAccumulator(image_ids, ptids, true_classes, noise_levels=[0.0, 0.05]) for m, model in enumerate(models): for sigma in [0.0, 0.05]: mc = run_mc_passes(model, loader, sigma, k=K) # (K, S, C) acc.add(model_index=m, model_kind="bayesian", noise_level=sigma, mc_probs=mc) ds = acc.to_dataset(attrs={"seed": seed, "n_mc": K}) save_evaluation(ds, out_path) """ def __init__( self, image_ids: Sequence[int], ptids: Sequence[str], true_classes: Sequence[str], noise_levels: Sequence[float], ) -> None: self._image_ids = list(image_ids) self._ptids = list(ptids) self._true_classes = list(true_classes) self._noise_levels = list(noise_levels) self._s = len(self._image_ids) if not (len(self._ptids) == len(self._true_classes) == self._s): raise ValueError( "image_ids, ptids, and true_classes must all have the same length." ) # Duplicate levels would collapse in this lookup and silently leave one # slice of the output array unwritten (i.e. all zeros), so reject them. keys = [round(float(n), 12) for n in self._noise_levels] duplicates = {k for k in keys if keys.count(k) > 1} if duplicates: raise ValueError( f"noise_levels contains duplicate value(s) {sorted(duplicates)}: " f"{self._noise_levels}. Each level must appear exactly once." ) self._noise_index = {k: i for i, k in enumerate(keys)} # keyed by model index -> dict with kind and per-noise arrays self._models: Dict[int, Dict[str, Any]] = {} def add( self, *, model_index: int, model_kind: str, noise_level: float, mc_probs: np.ndarray, ) -> None: """Add one model's MC probabilities ``(K, S, C)`` at one noise level.""" key = round(float(noise_level), 12) if key not in self._noise_index: raise ValueError( f"noise_level {noise_level} not in declared levels {self._noise_levels}." ) mean_prob, pred_ent, mut_info = compute_uncertainty(mc_probs) if mean_prob.shape[0] != self._s: raise ValueError( f"mc_probs sample axis {mean_prob.shape[0]} != declared S {self._s}." ) entry = self._models.setdefault( model_index, { "kind": model_kind, "prob": {}, "pe": {}, "mi": {}, }, ) if entry["kind"] != model_kind: raise ValueError( f"model_index {model_index} already registered as {entry['kind']!r}, " f"got {model_kind!r}." ) entry["prob"][key] = mean_prob entry["pe"][key] = pred_ent entry["mi"][key] = mut_info def to_dataset(self, attrs: Optional[Mapping[str, Any]] = None) -> xr.Dataset: """Stack all collected results into a schema-compliant Dataset.""" if not self._models: raise ValueError("No results were added to the accumulator.") model_indices = sorted(self._models) m = len(model_indices) ell = len(self._noise_levels) c = len(CLASS_NAMES) prob = np.zeros((m, ell, self._s, c), dtype=np.float32) pe = np.zeros((m, ell, self._s), dtype=np.float32) mi = np.zeros((m, ell, self._s), dtype=np.float32) model_kinds: List[str] = [] for mi_idx, model_index in enumerate(model_indices): entry = self._models[model_index] model_kinds.append(entry["kind"]) for key, li in self._noise_index.items(): if key not in entry["prob"]: raise ValueError( f"model_index {model_index} is missing noise level {key}." ) prob[mi_idx, li] = entry["prob"][key] pe[mi_idx, li] = entry["pe"][key] mi[mi_idx, li] = entry["mi"][key] return build_evaluation_dataset( prob=prob, predictive_entropy=pe, mutual_information=mi, model_kinds=model_kinds, noise_levels=self._noise_levels, image_ids=self._image_ids, ptids=self._ptids, true_classes=self._true_classes, attrs=attrs, )