alnn_rewrite/analysis/plots/noise_correlation.py

"""Uncertainty vs. accuracy across noise conditions, with a fitted curve.

Question: is uncertainty a *transferable* predictor of error -- does a given
uncertainty value imply the same accuracy regardless of the noise condition?

Method (see ai/ANALYSIS_PLAN.md 3.3)
-----------------------------------
For each noise level sigma, split its samples into ``N_BINS`` equal-count
(quantile) uncertainty bins and plot one point per bin: mean uncertainty in the
bin (x) against accuracy in the bin (y), coloured by sigma. This gives
sigma x bins points -- enough to fit -- rather than the single point per sigma a
naive reading would produce.

Fits ``acc = a + b * exp(-k * u)`` (bounded, monotone) with a linear fallback,
and reports R^2 plus Spearman rho. Rho is the robust headline because it makes no
assumption about the functional form.

Interpretation
--------------
If points from every sigma collapse onto one curve, the measure is calibrated
*across conditions*: uncertainty alone predicts accuracy. If they separate by
sigma, the measure is condition-dependent and cannot be used as a standalone
confidence signal -- an important negative result if it appears.

Outputs: ``noise_correlation_<measure>.png`` (one panel per family x
configuration) and ``noise_correlation.json`` (fit parameters, R^2, rho).
"""

from typing import Any, Dict, List

import numpy as np
from matplotlib import cm
from matplotlib.colors import Normalize

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 panel_grid, save_figure, save_json, style_axes
from analysis.registry import register

#: Equal-count uncertainty bins per noise level.
N_BINS = 10

#: Minimum samples in a bin for its accuracy to be trustworthy.
MIN_BIN_SAMPLES = 5


def _series_points(
    ctx: AnalysisContext, series: Series, measure: ms.Measure
) -> Dict[str, Any] | None:
    """Binned (uncertainty, accuracy, sigma) points plus the fit, for one series."""
    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)

    x_values: List[float] = []
    y_values: List[float] = []
    sigma_values: List[float] = []

    for level_index, sigma in enumerate(levels):
        probs = ctx.probs(series.family, ds, level_index)
        member_mi = (
            ctx.member_stat("mutual_information", series.family, ds, level_index)
            if measure.needs_member_mi
            else None
        )

        # `single` pools every (member, sample) observation at this noise level,
        # so it yields the same number of points as `ensemble` while using all
        # members' predictions.
        if series.config == CONFIG_SINGLE:
            uncertainty = np.concatenate(
                [
                    ms.compute(
                        measure,
                        probs[m : m + 1],
                        member_mi=None if member_mi is None else member_mi[m : m + 1],
                    )
                    for m in range(probs.shape[0])
                ]
            )
            correct = np.concatenate(
                [
                    ms.predicted_class(probs[m : m + 1]) == true_idx
                    for m in range(probs.shape[0])
                ]
            )
        else:
            uncertainty = ms.compute(measure, probs, member_mi=member_mi)
            correct = ms.predicted_class(probs) == true_idx

        for indices in mt.quantile_bins(uncertainty, N_BINS):
            if indices.size < MIN_BIN_SAMPLES:
                continue
            x_values.append(float(np.mean(uncertainty[indices])))
            y_values.append(float(np.mean(correct[indices])))
            sigma_values.append(float(sigma))

    if len(x_values) < 3:
        return None

    fit = mt.fit_accuracy_vs_uncertainty(
        np.asarray(x_values), np.asarray(y_values)
    )
    return {
        "source": stem,
        "noise_axis": noise_axis_label(ds),
        "uncertainty": x_values,
        "accuracy": y_values,
        "noise_level": sigma_values,
        "n_bins": N_BINS,
        "fit": {
            "model": fit.model,
            "params": fit.params,
            "r_squared": fit.r_squared,
            "spearman": fit.spearman,
        },
        "_fit_line": (fit.x_fit, fit.y_fit),
    }


def _plot_measure(
    ctx: AnalysisContext,
    measure: ms.Measure,
    per_series: Dict[str, Dict[str, Any]],
) -> None:
    """One panel per series; points coloured by noise level, with the fit drawn."""
    keys = list(per_series)
    fig, axes = panel_grid(len(keys), n_cols=2, panel_size=(5.0, 3.8))

    all_sigmas = sorted(
        {s for payload in per_series.values() for s in payload["noise_level"]}
    )
    norm = Normalize(vmin=min(all_sigmas), vmax=max(all_sigmas))
    colormap = cm.viridis

    scatter = None
    for ax, key in zip(axes, keys):
        payload = per_series[key]
        series: Series = payload["_series"]
        scatter = ax.scatter(
            payload["uncertainty"],
            payload["accuracy"],
            c=payload["noise_level"],
            cmap=colormap,
            norm=norm,
            s=26,
            edgecolor="white",
            linewidth=0.4,
            zorder=3,
        )
        x_fit, y_fit = payload["_fit_line"]
        if x_fit:
            ax.plot(x_fit, y_fit, color="#444444", linewidth=1.6, zorder=2)

        fit = payload["fit"]
        ax.annotate(
            f"{fit['model']}  R²={fit['r_squared']:.3f}\nρ={fit['spearman']:.3f}",
            xy=(0.97, 0.95),
            xycoords="axes fraction",
            ha="right",
            va="top",
            fontsize=8,
            color="#333333",
        )
        ax.set_title(series.label, fontsize=10)
        ax.set_xlabel(f"{measure.label}")
        ax.set_ylabel("Accuracy (per bin)")
        style_axes(ax)

    if scatter is not None:
        label = next(iter(per_series.values()))["noise_axis"]
        fig.colorbar(scatter, ax=axes, label=label, shrink=0.85)
    fig.suptitle(f"Accuracy vs. {measure.label.lower()} across noise levels")
    stem = f"noise_correlation_{measure.name}"
    png = save_figure(fig, ctx.out_dir, stem)
    ctx.log.info(f"noise_correlation: wrote {png.name}.")


@register(
    "noise_correlation",
    title="Uncertainty vs. accuracy across noise levels (with fit)",
    requires_noise=True,
)
def noise_correlation(ctx: AnalysisContext) -> None:
    results: Dict[str, Dict[str, Dict[str, Any]]] = {}

    for series in ctx.series(ctx.noise_families):
        for measure in ms.measures_for(series.family, series.config):
            payload = _series_points(ctx, series, measure)
            if payload is None:
                continue
            payload["_series"] = series
            results.setdefault(measure.name, {})[series.key] = payload

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

    for measure_name, per_series in results.items():
        _plot_measure(ctx, ms.MEASURES[measure_name], per_series)
        for key, payload in per_series.items():
            fit = payload["fit"]
            ctx.log.info(
                f"noise_correlation: {measure_name} / {key} — {fit['model']} fit "
                f"R²={fit['r_squared']:.3f}, Spearman ρ={fit['spearman']:.3f}."
            )

    save_json(
        {
            "n_bins": N_BINS,
            "min_bin_samples": MIN_BIN_SAMPLES,
            "measures": {
                measure: {
                    key: {
                        k: v
                        for k, v in payload.items()
                        if k not in ("_series", "_fit_line")
                    }
                    for key, payload in per_series.items()
                }
                for measure, per_series in results.items()
            },
        },
        ctx.out_dir,
        "noise_correlation",
    )