alnn_rewrite/analysis/plots/noise_performance.py

"""Performance and uncertainty vs. input noise.

Question: how do accuracy and each uncertainty measure respond as Gaussian input
noise increases? This is the core robustness result.

Method (see ai/ANALYSIS_PLAN.md 3.2)
-----------------------------------
Uses a noise-sweep evaluation (``noisy_*`` / ``combined_noisy_*``). For every
noise level sigma, computes accuracy, F1 and the mean (+/-1 std across samples)
of each applicable uncertainty measure, for every family x configuration. The
``single`` configuration averages over members.

The scientifically interesting comparison is the *rate* of change: the epistemic
terms (``ensemble_mutual_information``, ``mean_mutual_information``) should rise
faster than the aleatoric term (``mean_predictive_entropy``) as inputs move out
of distribution. If they do not, the uncertainty is not tracking what it should.

Outputs: ``noise_performance.png`` (accuracy and F1 panels, then one panel per
uncertainty measure, all against sigma) and ``noise_performance.json``.
"""

from typing import Any, Dict, List, Tuple

import numpy as np

import analysis.measures as ms
import analysis.metrics as mt
from analysis.context import CONFIG_SINGLE, AnalysisContext, Series
from analysis.sources import noise_axis_label
from analysis.plotting import add_legend, panel_grid, save_figure, save_json, style_axes
from analysis.registry import register


def _member_groups(
    probs: np.ndarray, member_mi: np.ndarray | None, config: str
) -> List[Tuple[np.ndarray, np.ndarray | None]]:
    """Split into the units scored independently: 1 ensemble, or M singles."""
    if config == CONFIG_SINGLE:
        return [
            (probs[m : m + 1], None if member_mi is None else member_mi[m : m + 1])
            for m in range(probs.shape[0])
        ]
    return [(probs, member_mi)]


def _series_response(
    ctx: AnalysisContext, series: Series
) -> Dict[str, Any] | None:
    """Accuracy, F1 and every applicable measure across the noise sweep."""
    selected = ctx.noise_source(series.family)
    if selected is None:
        return None
    stem, ds = selected

    levels = ctx.noise_levels(ds)
    true_idx = ctx.true_index(ds)
    applicable = ms.measures_for(series.family, series.config)

    accuracy: List[float] = []
    accuracy_std: List[float] = []
    f1: List[float] = []
    f1_std: List[float] = []
    measure_mean: Dict[str, List[float]] = {m.name: [] for m in applicable}
    measure_std: Dict[str, List[float]] = {m.name: [] for m in applicable}

    needs_mi = any(m.needs_member_mi for m in applicable)
    for level_index in range(len(levels)):
        probs = ctx.probs(series.family, ds, level_index)
        member_mi = (
            ctx.member_stat("mutual_information", series.family, ds, level_index)
            if needs_mi
            else None
        )
        groups = _member_groups(probs, member_mi, series.config)

        # Performance: score each unit, then average over units.
        unit_acc = [
            mt.accuracy(ms.predicted_class(g), true_idx) for g, _ in groups
        ]
        unit_f1 = [mt.f1_score(ms.predicted_class(g), true_idx) for g, _ in groups]
        accuracy.append(float(np.nanmean(unit_acc)))
        accuracy_std.append(float(np.nanstd(unit_acc)))
        f1.append(float(np.nanmean(unit_f1)))
        f1_std.append(float(np.nanstd(unit_f1)))

        # Uncertainty: mean over samples for each unit, then averaged over units;
        # the band is the spread across samples (pooled over units).
        for measure in applicable:
            per_sample = np.concatenate(
                [ms.compute(measure, g, member_mi=gm) for g, gm in groups]
            )
            measure_mean[measure.name].append(float(np.mean(per_sample)))
            measure_std[measure.name].append(float(np.std(per_sample)))

    return {
        "source": stem,
        "noise_axis": noise_axis_label(ds),
        "noise_levels": levels.tolist(),
        "accuracy": accuracy,
        "accuracy_std": accuracy_std,
        "f1": f1,
        "f1_std": f1_std,
        "measures": {
            name: {"mean": measure_mean[name], "std": measure_std[name]}
            for name in measure_mean
        },
    }


@register(
    "noise_performance",
    title="Accuracy and uncertainty vs. input noise",
    requires_noise=True,
)
def noise_performance(ctx: AnalysisContext) -> None:
    responses: Dict[str, Dict[str, Any]] = {}
    for series in ctx.series(ctx.noise_families):
        payload = _series_response(ctx, series)
        if payload is not None:
            payload["_series"] = series
            responses[series.key] = payload

    if not responses:
        ctx.log.error("noise_performance: no noise-sweep evaluations available.")
        return

    x_label = next(iter(responses.values()))["noise_axis"]

    def _draw(ax, key: str) -> None:
        """Plot every series' response for one quantity onto ``ax``."""
        is_performance = key in ("accuracy", "f1")
        for payload in responses.values():
            series: Series = payload["_series"]
            x = np.asarray(payload["noise_levels"], dtype=float)
            if is_performance:
                y = np.asarray(payload[key], dtype=float)
                spread = np.asarray(payload[f"{key}_std"], dtype=float)
                show_band = series.config == CONFIG_SINGLE
            else:
                if key not in payload["measures"]:
                    continue
                y = np.asarray(payload["measures"][key]["mean"], dtype=float)
                spread = np.asarray(payload["measures"][key]["std"], dtype=float)
                show_band = True
            if show_band:
                ax.fill_between(
                    x, y - spread, y + spread,
                    color=series.color, alpha=0.13, linewidth=0,
                )
            ax.plot(
                x, y,
                color=series.color, linestyle=series.linestyle,
                linewidth=2, marker="o", markersize=4, label=series.label,
            )
        ax.set_xlabel(x_label)
        style_axes(ax)

    # One figure for performance, then one per uncertainty measure -- a single
    # 7-panel grid was too dense to read.
    fig, axes = panel_grid(2, n_cols=2, panel_size=(5.2, 3.8))
    for ax, key, title in (
        (axes[0], "accuracy", "Accuracy"),
        (axes[1], "f1", f"F1 ({mt.POSITIVE_CLASS})"),
    ):
        _draw(ax, key)
        ax.set_title(title, fontsize=10)
        ax.set_ylabel(title)
    fig.suptitle("Performance vs. input noise")
    add_legend(fig, axes)
    png = save_figure(fig, ctx.out_dir, "noise_performance")

    for name in ms.MEASURES:
        if not any(name in p["measures"] for p in responses.values()):
            continue
        label = ms.MEASURES[name].label
        mfig, maxes = panel_grid(1, n_cols=1, panel_size=(6.4, 4.2))
        _draw(maxes[0], name)
        maxes[0].set_title(f"{label} vs. input noise", fontsize=11)
        maxes[0].set_ylabel(label)
        add_legend(mfig, maxes)
        mpng = save_figure(mfig, ctx.out_dir, f"noise_uncertainty_{name}")
        ctx.log.info(f"noise_performance: wrote {mpng.name}.")

    save_json(
        {
            "positive_class": mt.POSITIVE_CLASS,
            "series": {
                key: {k: v for k, v in payload.items() if k != "_series"}
                for key, payload in responses.items()
            },
        },
        ctx.out_dir,
        "noise_performance",
    )

    for key, payload in responses.items():
        levels = payload["noise_levels"]
        acc = payload["accuracy"]
        ctx.log.info(
            f"noise_performance: {key} — accuracy {acc[0]:.3f} at σ={levels[0]:g} "
            f"→ {acc[-1]:.3f} at σ={levels[-1]:g}."
        )
    ctx.log.info(f"noise_performance: wrote {png.name}.")