alnn_rewrite/tasks/extra_info.py

"""Extra information step -- figures and facts that are NOT in the evaluations.

The analysis stage reads only the evaluation netCDF files, so anything that needs
the *data itself* (split composition, what a noised image actually looks like)
has to be produced here, while the dataloaders are still in ``state``. Runs after
``load_data`` and before ``run_analysis``, writing into the same regeneratable
``analysis/`` directory.

Adding an extra: write a function taking an :class:`ExtraContext` and register it
in :data:`EXTRAS`.
"""

import json
import pathlib as pl
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Sequence

import matplotlib

matplotlib.use("Agg")  # headless; runs in a background worker thread

import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402
import torch  # noqa: E402

from analysis.plotting import data_dir, plots_dir  # noqa: E402
from evaluation import CLASS_NAMES  # noqa: E402
from util.progress import ProgressTracker  # noqa: E402
from util.ui_logger import PipelineLogger  # noqa: E402

ANALYSIS_SUBDIR = "analysis"

#: Splits summarised, in pipeline order.
_SPLITS = ("train", "val", "test")

#: Example images shown in the noise-degradation grid.
_N_EXAMPLES = 3

#: Maximum noise columns in that grid (levels are subsampled evenly to fit).
_MAX_NOISE_COLUMNS = 8


@dataclass
class ExtraContext:
    """What an extra-info producer gets: pipeline state, config, output dir."""

    config: Dict[str, Any]
    state: Dict[str, Any]
    out_dir: pl.Path
    log: PipelineLogger


ExtraFn = Callable[[ExtraContext], None]


def _split_arrays(loader) -> tuple[np.ndarray, np.ndarray]:
    """Return ``(image_ids, class_indices)`` for a loader WITHOUT loading images.

    Reads straight from the underlying ``ADNIDataset`` through the ``Subset``
    indices, so summarising the (large) training split costs nothing.
    """
    subset = loader.dataset
    base = getattr(subset, "dataset", subset)
    indices = list(getattr(subset, "indices", range(len(base))))
    image_ids = np.array([int(base.filename_ids[i]) for i in indices], dtype=int)
    classes = base.expected_classes[indices].argmax(dim=1).cpu().numpy()
    return image_ids, classes


def dataset_summary(ctx: ExtraContext) -> None:
    """Split sizes, class balance and patient counts for every split."""
    image_to_ptid: Dict[int, str] = ctx.state.get("image_to_ptid", {})

    summary: Dict[str, Any] = {
        "seed": ctx.config["data"].get("seed"),
        "data_splits": ctx.config["data"].get("data_splits"),
        "class_names": list(CLASS_NAMES),
        "splits": {},
    }
    all_patients: Dict[str, set[str]] = {}

    for split in _SPLITS:
        loader = ctx.state.get(f"{split}_loader")
        if loader is None:
            continue
        image_ids, classes = _split_arrays(loader)
        patients = {image_to_ptid.get(int(i), "unknown") for i in image_ids}
        all_patients[split] = patients
        summary["splits"][split] = {
            "n_scans": int(image_ids.size),
            "n_patients": len(patients),
            "scans_per_patient": round(image_ids.size / max(len(patients), 1), 2),
            "class_counts": {
                name: int(np.sum(classes == index))
                for index, name in enumerate(CLASS_NAMES)
            },
            "class_fractions": {
                name: round(float(np.mean(classes == index)), 4)
                for index, name in enumerate(CLASS_NAMES)
            },
        }

    if not summary["splits"]:
        ctx.log.error("dataset_summary: no dataloaders in state; skipping.")
        return

    # Leakage check: the patient-grouped split must keep splits disjoint.
    leaks: Dict[str, int] = {}
    names = list(all_patients)
    for i, a in enumerate(names):
        for b in names[i + 1 :]:
            shared = all_patients[a] & all_patients[b] - {"unknown"}
            if shared:
                leaks[f"{a}&{b}"] = len(shared)
    summary["patient_overlap_between_splits"] = leaks
    if leaks:
        ctx.log.error(f"dataset_summary: PATIENT LEAKAGE between splits: {leaks}")

    (data_dir(ctx.out_dir) / "dataset_summary.json").write_text(
        json.dumps(summary, indent=2)
    )

    # Figure: scans per split, stacked by class.
    splits = list(summary["splits"])
    fig, axes = plt.subplots(1, 2, figsize=(10.0, 3.6), layout="constrained")
    bottom = np.zeros(len(splits))
    for index, name in enumerate(CLASS_NAMES):
        heights = np.array(
            [summary["splits"][s]["class_counts"][name] for s in splits], dtype=float
        )
        axes[0].bar(
            splits, heights, bottom=bottom, label=name,
            color=["#2a78d6", "#e34948"][index % 2],
        )
        bottom += heights
    axes[0].set_ylabel("Scans")
    axes[0].set_title("Split size and class balance", fontsize=10)
    axes[0].legend(frameon=False)

    axes[1].bar(
        splits,
        [summary["splits"][s]["n_patients"] for s in splits],
        color="#7a7a7a",
    )
    axes[1].set_ylabel("Patients")
    axes[1].set_title("Distinct patients per split", fontsize=10)

    for ax in axes:
        ax.grid(True, axis="y", color="#000000", alpha=0.08, linewidth=0.7)
        for spine in ("top", "right"):
            ax.spines[spine].set_visible(False)
    fig.savefig(
        plots_dir(ctx.out_dir) / "dataset_summary.png", dpi=150, bbox_inches="tight"
    )
    plt.close(fig)

    parts = ", ".join(
        f"{s}={summary['splits'][s]['n_scans']} scans/"
        f"{summary['splits'][s]['n_patients']} patients"
        for s in splits
    )
    ctx.log.info(f"dataset_summary: {parts}. Wrote dataset_summary.png/.json.")


def _noise_columns(config: Dict[str, Any]) -> List[float]:
    """Noise levels to show, evenly subsampled to at most _MAX_NOISE_COLUMNS."""
    eval_cfg = config.get("evaluation", {})
    levels = sorted(
        {float(v) for v in eval_cfg.get("noise_levels", [0.0])}
        | {float(v) for v in eval_cfg.get("extra_noise_levels", [])}
    )
    if len(levels) <= _MAX_NOISE_COLUMNS:
        return levels
    picks = np.linspace(0, len(levels) - 1, _MAX_NOISE_COLUMNS).round().astype(int)
    return [levels[i] for i in sorted(set(picks.tolist()))]


def noise_examples(ctx: ExtraContext) -> None:
    """Grid of example images degrading across the configured noise levels.

    Applies exactly the same perturbation the evaluation uses (relative to each
    image's own intensity std when ``noise_relative``), so the figure shows what
    the model actually saw at each sigma.
    """
    loader = ctx.state.get("test_loader")
    if loader is None:
        ctx.log.error("noise_examples: no test_loader in state; skipping.")
        return

    subset = loader.dataset
    base = getattr(subset, "dataset", subset)
    indices = list(getattr(subset, "indices", range(len(base))))[:_N_EXAMPLES]
    if not indices:
        ctx.log.error("noise_examples: split is empty; skipping.")
        return

    levels = _noise_columns(ctx.config)
    relative = bool(ctx.config.get("evaluation", {}).get("noise_relative", True))
    seed = int(ctx.config["data"].get("seed", 0) or 0)

    fig, axes = plt.subplots(
        len(indices),
        len(levels),
        figsize=(1.55 * len(levels), 1.75 * len(indices)),
        squeeze=False,
        layout="constrained",
    )

    for row, dataset_index in enumerate(indices):
        volume = base.mri_data[dataset_index].float()  # (C, D, H, W)
        image_id = int(base.filename_ids[dataset_index])
        label = CLASS_NAMES[int(base.expected_classes[dataset_index].argmax())]
        std = float(volume.std())
        # A mid-axial slice, chosen once so every column shows the same anatomy.
        slice_index = volume.shape[-1] // 2
        generator = torch.Generator().manual_seed(seed + dataset_index)
        clean_slice = volume[0, :, :, slice_index]
        # Fix the display window on the CLEAN image so added noise visibly
        # washes the image out instead of being renormalised away.
        vmin, vmax = float(clean_slice.min()), float(clean_slice.max())

        for col, sigma in enumerate(levels):
            noisy = volume
            if sigma > 0:
                noise = torch.randn(volume.shape, generator=generator)
                noisy = volume + noise * (sigma * (std if relative else 1.0))
            ax = axes[row][col]
            ax.imshow(
                noisy[0, :, :, slice_index].numpy().T,
                cmap="gray", vmin=vmin, vmax=vmax, origin="lower",
            )
            ax.set_xticks([])
            ax.set_yticks([])
            if row == 0:
                ax.set_title(f"σ = {sigma:g}", fontsize=9)
            if col == 0:
                ax.set_ylabel(f"{image_id}\n({label})", fontsize=8)

    unit = "× image std" if relative else "raw voxel units"
    fig.suptitle(f"Image degradation with added Gaussian noise (σ in {unit})")
    fig.savefig(
        plots_dir(ctx.out_dir) / "noise_examples.png", dpi=150, bbox_inches="tight"
    )
    plt.close(fig)
    ctx.log.info(
        f"noise_examples: wrote noise_examples.png "
        f"({len(indices)} images x {len(levels)} noise levels)."
    )


#: Registered extra-info producers, in run order.
EXTRAS: Dict[str, ExtraFn] = {
    "dataset_summary": dataset_summary,
    "noise_examples": noise_examples,
}


def extra_info_task(
    track: ProgressTracker,
    log: PipelineLogger,
    config: Dict[str, Any],
    state: Dict[str, Any],
) -> None:
    """Run every registered extra-info producer."""
    out_dir = pl.Path(config["work_dir"]) / ANALYSIS_SUBDIR
    out_dir.mkdir(parents=True, exist_ok=True)
    ctx = ExtraContext(config=config, state=state, out_dir=out_dir, log=log)

    track.update(total=len(EXTRAS), advance=0)
    for name, fn in EXTRAS.items():
        log.info(f"Extra info: {name}")
        try:
            fn(ctx)
        except Exception as exc:  # noqa: BLE001 - one failure must not stop the rest
            log.error(f"Extra info '{name}' failed: {exc}")
            log.file_logger.error(f"Extra info '{name}' failed", exc_info=True)
        track.update(advance=1)
    log.info(f"Extra information written to {out_dir}/.")