|
|
@@ -0,0 +1,429 @@
|
|
|
+"""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}
|
|
|
+ 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)
|
|
|
+
|
|
|
+
|
|
|
+# ---------------------------------------------------------------------------
|
|
|
+# 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."
|
|
|
+ )
|
|
|
+ self._noise_index = {round(float(n), 12): i for i, n in enumerate(self._noise_levels)}
|
|
|
+
|
|
|
+ # 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,
|
|
|
+ )
|
Nicholas Schense