alnn_rewrite/analysis/plots/model_report.py

"""Per-model summary report (PDF, via Typst).

A compact table of every ensemble member's accuracy and mean uncertainty, one
section per model family, plus the ensemble headline. Data is injected into the
Typst template as a JSON string through ``sys.inputs``.
"""

import datetime
import json
import pathlib as pl
from typing import Any, Dict, List

import numpy as np
import typst

from analysis.sources import FAMILY_LABELS, baseline_noise_index
from analysis.context import AnalysisContext
from analysis.plotting import data_dir
from analysis.registry import register

_TEMPLATE = pl.Path(__file__).resolve().parent.parent / "templates" / "model_report.typ"


def _family_block(ctx: AnalysisContext, family: str) -> Dict[str, Any] | None:
    """Per-member and ensemble metrics for one family at the clean baseline."""
    selected = ctx.clean_source(family)
    if selected is None:
        return None
    stem, ds = selected

    baseline = baseline_noise_index(ds)
    positions = ctx.member_positions(family, ds)
    correct = np.asarray(ds["correct"].isel(noise_level=baseline).values, dtype=float)
    entropy = ctx.member_stat("predictive_entropy", family, ds, baseline)
    mut_info = ctx.member_stat("mutual_information", family, ds, baseline)
    labels = np.atleast_1d(ds["model"].values)

    models: List[Dict[str, Any]] = []
    accuracies: List[float] = []
    for row, position in enumerate(positions):
        acc = float(np.mean(correct[position]))
        accuracies.append(acc)
        models.append(
            {
                "index": int(labels[position]) + 1,  # 1-based for display
                "accuracy": acc,
                "mean_entropy": float(np.mean(entropy[row])),
                "mean_mi": float(np.mean(mut_info[row])),
            }
        )

    acc_arr = np.asarray(accuracies, dtype=float)
    return {
        "name": FAMILY_LABELS.get(family, family),
        "kind": family,
        "source": stem,
        "split": str(ds.attrs.get("split", "unknown")),
        "noise_sigma": float(ctx.noise_levels(ds)[baseline]),
        "n_models": len(models),
        "n_samples": int(ds.sizes["sample"]),
        "n_mc": int(ds.attrs.get("n_mc", 1)),
        "accuracy_mean": float(acc_arr.mean()) if acc_arr.size else 0.0,
        "accuracy_std": float(acc_arr.std()) if acc_arr.size else 0.0,
        "accuracy_best": float(acc_arr.max()) if acc_arr.size else 0.0,
        "models": models,
    }


@register(
    "model_report",
    title="Per-model accuracy and uncertainty summary (PDF)",
)
def model_report(ctx: AnalysisContext) -> None:
    families = [
        block
        for block in (_family_block(ctx, family) for family in ctx.families)
        if block is not None
    ]
    if not families:
        ctx.log.error("model_report: no model families found in evaluations.")
        return

    meta = next(iter(ctx.datasets.values())).attrs
    seed = meta.get("seed")
    payload = {
        "title": "Model Evaluation Report",
        "generated": datetime.datetime.now(datetime.timezone.utc).strftime(
            "%Y-%m-%d %H:%M UTC"
        ),
        "work_dir": str(ctx.out_dir.parent),
        "schema_version": str(meta.get("schema_version", "?")),
        "seed": str(seed) if seed is not None else "n/a",
        "git_commit": str(meta.get("git_commit", "unknown"))[:10],
        "families": families,
    }
    payload_json = json.dumps(payload)

    # The template parses this with `json(bytes(sys.inputs.at("data")))`.
    pdf_bytes = typst.compile(
        str(_TEMPLATE), sys_inputs={"data": payload_json}, format="pdf"
    )
    (ctx.out_dir / "model_report.pdf").write_bytes(pdf_bytes)
    (data_dir(ctx.out_dir) / "model_report.json").write_text(payload_json)

    total_models = sum(f["n_models"] for f in families)
    ctx.log.info(
        f"model_report: wrote model_report.pdf "
        f"({len(families)} family/families, {total_models} model(s))."
    )