alnn_rewrite/evaluation/schema.py

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