alnn_rewrite/analysis/sources.py

"""Choosing which evaluation file represents which model family.

A single work_dir holds several evaluation files describing the *same* ensembles:
``normal.nc`` (test), ``val_normal.nc``, ``combined_normal.nc`` (pooled),
``noisy_normal.nc`` (noise sweep), ``oof_normal.nc`` (k-fold). Analyses must pick
exactly one per family per purpose, or a family would be counted twice and the
numbers would silently disagree between analyses.

Two selections exist:
  * **clean** -- the largest available sample set at the baseline noise level,
    used for accuracy/coverage style analyses.
  * **noise** -- the noise sweep (more than one noise level), used by the
    noise-response analyses.
"""

from typing import Dict, Iterable, List, Tuple

import numpy as np
import xarray as xr

from evaluation import MODEL_KIND_BAYESIAN, MODEL_KIND_NORMAL

#: Canonical display order for model families.
FAMILY_ORDER: List[str] = [MODEL_KIND_NORMAL, MODEL_KIND_BAYESIAN]

#: Human-readable family labels.
FAMILY_LABELS: Dict[str, str] = {
    MODEL_KIND_NORMAL: "Normal (deterministic)",
    MODEL_KIND_BAYESIAN: "Bayesian (MC)",
}


def ordered_families(kinds: Iterable[str]) -> List[str]:
    """Sort family kinds into the canonical Normal-then-Bayesian order."""
    order = {kind: i for i, kind in enumerate(FAMILY_ORDER)}
    return sorted(kinds, key=lambda k: order.get(k, len(order)))


def baseline_noise_index(ds: xr.Dataset) -> int:
    """Index of the noise level closest to the clean baseline (sigma = 0)."""
    levels = np.asarray(ds["noise_level"].values, dtype=float)
    return int(np.abs(levels).argmin())


def has_noise_sweep(ds: xr.Dataset) -> bool:
    """True when the dataset spans more than one noise level."""
    return int(ds.sizes.get("noise_level", 1)) > 1


def noise_axis_label(ds: xr.Dataset) -> str:
    """Axis label for the noise level, reflecting how sigma was applied.

    Relative noise (the default) scales sigma by each image's own intensity
    standard deviation, so the units differ from absolute noise and the axis must
    say which it is.
    """
    if int(ds.attrs.get("noise_relative", 0)):
        return "Noise σ (× image std)"
    return "Noise σ (raw voxel units)"


def _coverage_rank(stem: str, ds: xr.Dataset) -> int:
    """Rank a file by how many evaluation samples it covers (lower is better)."""
    split = str(ds.attrs.get("split", ""))
    if stem.startswith(("combined_", "oof_")) or "+" in split:
        return 0  # pooled: test+val or k-fold out-of-fold -- the most samples
    if split == "test":
        return 1
    if split == "val":
        return 2
    return 3  # unlabeled (files written before the split attribute existed)


def source_rank(stem: str, ds: xr.Dataset) -> Tuple[int, int, str]:
    """Sort key for the *clean* selection: clean before noisy, then by coverage."""
    return (1 if "noisy" in stem else 0, _coverage_rank(stem, ds), stem)


def noise_source_rank(stem: str, ds: xr.Dataset) -> Tuple[int, str]:
    """Sort key for the *noise* selection: widest coverage first."""
    return (_coverage_rank(stem, ds), stem)


def _families_in(ds: xr.Dataset) -> List[str]:
    if "model_kind" not in ds.coords:
        return []
    return sorted({str(k) for k in np.atleast_1d(ds["model_kind"].values)})


def select_clean_sources(
    datasets: Dict[str, xr.Dataset]
) -> Dict[str, Tuple[str, xr.Dataset]]:
    """Best clean (baseline) evaluation per family: ``{family: (stem, ds)}``."""
    chosen: Dict[str, Tuple[str, xr.Dataset]] = {}
    for stem in sorted(datasets, key=lambda s: source_rank(s, datasets[s])):
        for family in _families_in(datasets[stem]):
            chosen.setdefault(family, (stem, datasets[stem]))
    return chosen


def select_noise_sources(
    datasets: Dict[str, xr.Dataset]
) -> Dict[str, Tuple[str, xr.Dataset]]:
    """Best noise-sweep evaluation per family: ``{family: (stem, ds)}``.

    Only datasets with more than one noise level qualify; families without a
    sweep are simply absent from the result.
    """
    chosen: Dict[str, Tuple[str, xr.Dataset]] = {}
    candidates = {s: d for s, d in datasets.items() if has_noise_sweep(d)}
    for stem in sorted(candidates, key=lambda s: noise_source_rank(s, candidates[s])):
        for family in _families_in(candidates[stem]):
            chosen.setdefault(family, (stem, candidates[stem]))
    return chosen