confidence_analysis_check.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""`aquin check confidence --check` — save JSON + diagram to the current working directory."""
 
from __future__ import annotations
 
import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
 
_JSON_NAME = "confidence-analysis-check.json"
_PNG_NAME = "confidence-analysis-check.png"
 
 
def _normalize_confidence_analysis_result(result: dict[str, Any]) -> dict[str, Any]:
    if not isinstance(result, dict):
        return {"error": "Invalid confidence-analysis result"}
 
    data = result
    card = result.get("card")
    if isinstance(card, dict) and isinstance(card.get("data"), dict):
        card_data = card["data"]
        data = {
            "schema_version": result.get("schema_version", 1),
            "model_id": card_data.get("modelId") or result.get("model_id"),
            "mode": card_data.get("mode") or result.get("mode"),
            "n_probes": card_data.get("nProbes") or result.get("n_probes"),
            "threshold": card_data.get("threshold") or result.get("threshold"),
            "mean_confidence": card_data.get("meanConfidence") or result.get("mean_confidence"),
            "aggregate_ece_proxy": card_data.get("aggregateEceProxy") or result.get("aggregate_ece_proxy"),
            "low_confidence_count": card_data.get("lowConfidenceCount") or result.get("low_confidence_count"),
            "join_sae": card_data.get("joinSae") if card_data.get("joinSae") is not None else result.get("join_sae"),
            "sae_layer": card_data.get("saeLayer") or result.get("sae_layer"),
            "probes": card_data.get("probes") or result.get("probes", []),
            "stressor_summary": card_data.get("stressorSummary") or result.get("stressor_summary", []),
            "heatmap": card_data.get("heatmap") or result.get("heatmap", {}),
            "saved_to": card_data.get("savedTo") or result.get("saved_to"),
        }
    elif isinstance(result.get("content"), dict):
        data = result["content"]
 
    if result.get("error"):
        data = {**data, "error": result["error"]}
    return data
 
 
def write_confidence_analysis_check(
    result: dict[str, Any],
    *,
    tool_name: str | None,
    cwd: str | Path,
) -> tuple[Path, Path]:
    cwd = Path(cwd)
    json_path = cwd / _JSON_NAME
    png_path = cwd / _PNG_NAME
 
    flat = _normalize_confidence_analysis_result(result)
    payload = {
        "saved_at": datetime.now(timezone.utc).isoformat(),
        "tool": tool_name,
        **flat,
    }
    json_path.write_text(json.dumps(payload, indent=2, default=str), encoding="utf-8")
    _plot_confidence_analysis_check(flat, png_path)
    return json_path, png_path
 
 
def _plot_confidence_analysis_check(result: dict[str, Any], png_path: Path) -> None:
    import matplotlib
 
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
 
    if result.get("error"):
        fig, ax = plt.subplots(figsize=(6, 2), facecolor="#0f1117")
        ax.axis("off")
        ax.text(0.5, 0.5, f"Error: {result['error']}", ha="center", va="center", wrap=True, color="#e5e7eb")
        fig.savefig(png_path, dpi=140, bbox_inches="tight", facecolor="#0f1117")
        plt.close(fig)
        return
 
    model = str(result.get("model_id") or "")
    mode = str(result.get("mode") or "")
    n_probes = int(result.get("n_probes") or 0)
    threshold = float(result.get("threshold") or 0.4)
    mean_conf = float(result.get("mean_confidence") or 0)
    low_count = int(result.get("low_confidence_count") or 0)
    join_sae = bool(result.get("join_sae"))
    sae_layer = result.get("sae_layer")
 
    title = f"Confidence analysis — {model}  ·  {mode}  ·  {n_probes} probes"
    subtitle = f"mean conf {mean_conf:.3f}  ·  threshold {threshold:.2f}  ·  low-conf {low_count}"
    if join_sae and sae_layer is not None:
        subtitle += f"  ·  SAE L{sae_layer}"
 
    fig = plt.figure(figsize=(12, 7.5), facecolor="#0f1117")
    gs = fig.add_gridspec(2, 2, height_ratios=[1, 1.2], hspace=0.38, wspace=0.28)
    ax_probes = fig.add_subplot(gs[0, 0])
    ax_stress = fig.add_subplot(gs[0, 1])
    ax_heat = fig.add_subplot(gs[1, :])
 
    fig.suptitle(title, color="#e5e7eb", fontsize=11, y=0.98)
    fig.text(0.5, 0.93, subtitle, ha="center", fontsize=8, color="#9ca3af")
 
    _plot_probe_confidence(ax_probes, result.get("probes") or [], threshold)
    _plot_stressor_summary(ax_stress, result.get("stressor_summary") or [])
    _plot_metric_heatmap(ax_heat, result.get("heatmap") or {})
 
    for ax in (ax_probes, ax_stress, ax_heat):
        ax.set_facecolor("#0f1117")
        for spine in ax.spines.values():
            spine.set_color("#374151")
 
    fig.savefig(png_path, dpi=140, bbox_inches="tight", facecolor="#0f1117")
    plt.close(fig)
 
 
def _plot_probe_confidence(ax: Any, probes: list[dict[str, Any]], threshold: float) -> None:
    ax.set_title("Per-probe confidence", color="#e5e7eb", fontsize=10)
 
    if not probes:
        ax.axis("off")
        ax.text(0.5, 0.5, "No probe data", ha="center", va="center", color="#9ca3af")
        return
 
    rows = sorted(probes, key=lambda p: float(p.get("mean_confidence") or 0))[:16]
    labels = [_probe_label(p) for p in rows]
    vals = [float(p.get("mean_confidence") or 0) for p in rows]
    colors = ["#f87171" if v < threshold else "#34d399" if v >= 0.7 else "#6366f1" for v in vals]
 
    y_pos = list(range(len(rows)))
    ax.barh(y_pos, vals, color=colors, height=0.65, alpha=0.9)
    ax.axvline(threshold, color="#fbbf24", linestyle="--", linewidth=0.8, alpha=0.7)
    ax.set_yticks(y_pos)
    ax.set_yticklabels(labels, fontsize=7, color="#d1d5db")
    ax.invert_yaxis()
    ax.set_xlim(0, 1.05)
    ax.set_xlabel("Mean confidence", color="#9ca3af", fontsize=8)
    ax.tick_params(axis="x", colors="#6b7280", labelsize=8)
    ax.grid(axis="x", alpha=0.2, color="#4b5563")
 
 
def _plot_stressor_summary(ax: Any, summary: list[dict[str, Any]]) -> None:
    ax.set_title("By stressor", color="#e5e7eb", fontsize=10)
 
    if not summary:
        ax.axis("off")
        ax.text(0.5, 0.5, "No stressor summary", ha="center", va="center", color="#9ca3af")
        return
 
    labels = [str(s.get("stressor") or "?") for s in summary]
    vals = [float(s.get("mean_confidence") or 0) for s in summary]
    y_pos = list(range(len(summary)))
    ax.barh(y_pos, vals, color="#38bdf8", height=0.65, alpha=0.9)
    ax.set_yticks(y_pos)
    ax.set_yticklabels(labels, fontsize=8, color="#d1d5db")
    ax.invert_yaxis()
    ax.set_xlim(0, 1.05)
    ax.set_xlabel("Mean confidence", color="#9ca3af", fontsize=8)
    ax.tick_params(axis="x", colors="#6b7280", labelsize=8)
    ax.grid(axis="x", alpha=0.2, color="#4b5563")
    for i, s in enumerate(summary):
        low = int(s.get("low_confidence_count") or 0)
        ax.text(min(vals[i] + 0.02, 0.92), i, f"{vals[i]:.3f} ({low} low)", va="center", fontsize=7, color="#9ca3af")
 
 
def _plot_metric_heatmap(ax: Any, heatmap: dict[str, Any]) -> None:
    rows = heatmap.get("rows") or []
    cols = heatmap.get("cols") or []
    values = heatmap.get("values") or []
 
    ax.set_title("Stressor × metrics", color="#e5e7eb", fontsize=10, pad=10)
 
    if not rows or not cols or not values:
        ax.axis("off")
        ax.text(0.5, 0.5, "No heatmap data", ha="center", va="center", color="#9ca3af")
        return
 
    import numpy as np
 
    data = np.array(values, dtype=float)
    if data.ndim != 2:
        ax.axis("off")
        ax.text(0.5, 0.5, "Invalid heatmap shape", ha="center", va="center", color="#9ca3af")
        return
 
    im = ax.imshow(data, aspect="auto", cmap="plasma", origin="upper")
    ax.set_xticks(range(len(cols)))
    ax.set_xticklabels(cols, fontsize=8, color="#d1d5db", rotation=30, ha="right")
    ax.set_yticks(range(len(rows)))
    ax.set_yticklabels(rows, fontsize=8, color="#d1d5db")
    ax.tick_params(axis="x", colors="#6b7280")
    ax.tick_params(axis="y", colors="#6b7280")
    cbar = ax.figure.colorbar(im, ax=ax, fraction=0.025, pad=0.02)
    cbar.ax.tick_params(colors="#9ca3af", labelsize=7)
 
 
def _probe_label(probe: dict[str, Any]) -> str:
    pid = str(probe.get("id") or "")
    stressor = probe.get("stressor")
    if pid and stressor:
        return f"{pid} ({stressor})"[:24]
    if pid:
        return pid[:24]
    if stressor:
        return str(stressor)[:24]
    return _trunc(str(probe.get("text") or ""), 22)
 
 
def _trunc(text: str, limit: int) -> str:
    t = text.strip()
    if len(t) <= limit:
        return t or "·"
    return t[: limit - 1] + "…"