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} # 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


def concat_noise_levels(
    datasets: Sequence[xr.Dataset], attrs: Optional[Mapping[str, Any]] = None
) -> xr.Dataset:
    """Union evaluations of the SAME samples that differ only in noise level.

    This is what lets a noise sweep be extended without recomputing it: a later
    run evaluates additional sigmas over the same split and models, and the
    result is merged onto the existing ``noise_level`` axis here.

    The parts must share the ``model``, ``sample`` and ``class_name`` axes; only
    ``noise_level`` grows. The result is sorted ascending by sigma (plots and
    curve fits assume a monotonic axis) and duplicate levels are dropped, keeping
    the first occurrence -- so re-running an already-computed sigma is harmless.

    Args:
        datasets: Evaluations over the same samples at different noise levels.
        attrs: Attributes merged over the result.
    """
    parts = list(datasets)
    if not parts:
        raise ValueError("concat_noise_levels requires at least one dataset.")

    combined = xr.concat(parts, dim="noise_level")
    levels = np.round(
        np.asarray(combined["noise_level"].values, dtype=float), 12
    )
    # np.unique returns the index of the FIRST occurrence of each value.
    _, first_occurrence = np.unique(levels, return_index=True)
    combined = combined.isel(noise_level=np.sort(first_occurrence))
    combined = combined.sortby("noise_level")
    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,
        )