alnn_rewrite/analysis/context.py

"""Shared state and data access for analyses.

``AnalysisContext`` owns the loaded evaluation datasets and answers the questions
every analysis needs -- which file represents a family, which members belong to
it, how to slice a noise level -- so no analysis re-derives that logic and two
analyses can never silently disagree about what "the normal ensemble" means.

The ``Series`` dataclass names the family x configuration combinations that all
analyses iterate, together with their visual encoding.
"""

import pathlib as pl
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import xarray as xr

from analysis.sources import (
    FAMILY_LABELS,
    baseline_noise_index,
    ordered_families,
    select_clean_sources,
    select_noise_sources,
)
from evaluation import CLASS_NAMES, MODEL_KIND_BAYESIAN, MODEL_KIND_NORMAL
from util.ui_logger import PipelineLogger

#: Ensemble configurations analysed side by side.
CONFIG_ENSEMBLE = "ensemble"
CONFIG_SINGLE = "single"
CONFIG_ORDER = [CONFIG_ENSEMBLE, CONFIG_SINGLE]

CONFIG_LABELS = {
    CONFIG_ENSEMBLE: "ensemble",
    CONFIG_SINGLE: "single",
}

#: Family -> hue (dataviz categorical slots, CVD-safe in this fixed order).
FAMILY_COLORS = {
    MODEL_KIND_NORMAL: "#2a78d6",     # blue
    MODEL_KIND_BAYESIAN: "#e34948",   # red
}

#: Configuration -> line style. Identity is never carried by colour alone.
CONFIG_LINESTYLES = {
    CONFIG_ENSEMBLE: "-",
    CONFIG_SINGLE: "--",
}


@dataclass(frozen=True)
class Series:
    """One family x configuration combination, with its visual encoding."""

    family: str
    config: str

    @property
    def key(self) -> str:
        return f"{self.family}_{self.config}"

    @property
    def label(self) -> str:
        return f"{FAMILY_LABELS.get(self.family, self.family)} — {CONFIG_LABELS[self.config]}"

    @property
    def color(self) -> str:
        return FAMILY_COLORS.get(self.family, "#666666")

    @property
    def linestyle(self) -> str:
        return CONFIG_LINESTYLES[self.config]

    @property
    def is_ensemble(self) -> bool:
        return self.config == CONFIG_ENSEMBLE


class AnalysisContext:
    """Loaded evaluations plus the selection logic analyses share."""

    def __init__(
        self,
        datasets: Dict[str, xr.Dataset],
        out_dir: pl.Path,
        log: PipelineLogger,
        config: Optional[Dict[str, Any]] = None,
    ) -> None:
        self.datasets = datasets
        self.out_dir = out_dir
        self.log = log
        self.config: Dict[str, Any] = config or {}
        self._clean = select_clean_sources(datasets)
        self._noise = select_noise_sources(datasets)

    # -- source selection --------------------------------------------------
    @property
    def families(self) -> List[str]:
        """Families present in the clean evaluations, in canonical order."""
        return ordered_families(self._clean.keys())

    @property
    def noise_families(self) -> List[str]:
        """Families that have a noise sweep available."""
        return ordered_families(self._noise.keys())

    def clean_source(self, family: str) -> Optional[Tuple[str, xr.Dataset]]:
        """Best baseline evaluation for ``family`` as ``(stem, dataset)``."""
        return self._clean.get(family)

    def noise_source(self, family: str) -> Optional[Tuple[str, xr.Dataset]]:
        """Best noise-sweep evaluation for ``family`` as ``(stem, dataset)``."""
        return self._noise.get(family)

    @property
    def has_noise(self) -> bool:
        return bool(self._noise)

    def series(self, families: Optional[List[str]] = None) -> List[Series]:
        """All family x configuration series, in canonical order."""
        chosen = families if families is not None else self.families
        return [
            Series(family=family, config=config)
            for family in ordered_families(chosen)
            for config in CONFIG_ORDER
        ]

    # -- data access -------------------------------------------------------
    def member_positions(self, family: str, ds: xr.Dataset) -> List[int]:
        """Indices along the ``model`` axis belonging to ``family``."""
        kinds = [str(k) for k in np.atleast_1d(ds["model_kind"].values)]
        return [i for i, k in enumerate(kinds) if k == family]

    def probs(
        self, family: str, ds: xr.Dataset, noise_index: Optional[int] = None
    ) -> np.ndarray:
        """Member probabilities ``(M, S, C)`` for ``family`` at one noise level.

        ``noise_index`` defaults to the clean baseline.
        """
        if noise_index is None:
            noise_index = baseline_noise_index(ds)
        values = np.asarray(
            ds["prob"].isel(noise_level=noise_index).values, dtype=float
        )
        return values[self.member_positions(family, ds)]

    def member_stat(
        self,
        name: str,
        family: str,
        ds: xr.Dataset,
        noise_index: Optional[int] = None,
    ) -> np.ndarray:
        """A stored per-member statistic ``(M, S)``.

        ``name`` is a schema variable such as ``predictive_entropy`` or
        ``mutual_information``.
        """
        if noise_index is None:
            noise_index = baseline_noise_index(ds)
        values = np.asarray(
            ds[name].isel(noise_level=noise_index).values, dtype=float
        )
        return values[self.member_positions(family, ds)]

    @staticmethod
    def true_index(ds: xr.Dataset) -> np.ndarray:
        """Ground-truth class indices ``(S,)`` aligned with ``CLASS_NAMES``."""
        return np.array(
            [
                CLASS_NAMES.index(str(t))
                for t in np.atleast_1d(ds["true_class"].values)
            ],
            dtype=int,
        )

    @staticmethod
    def noise_levels(ds: xr.Dataset) -> np.ndarray:
        return np.asarray(ds["noise_level"].values, dtype=float)

    @staticmethod
    def ptids(ds: xr.Dataset) -> np.ndarray:
        return np.asarray(ds["ptid"].values).astype(str)

    def base_seed(self) -> int:
        """Experiment seed recorded in the evaluations (0 when absent)."""
        for ds in self.datasets.values():
            seed = ds.attrs.get("seed")
            if seed is not None:
                return int(seed)
        return 0