code/scripts/aggregate_kcn_focus_results.py

#!/usr/bin/env python
"""Aggregate focused KSL+CASL+NSL experiment results."""

from __future__ import annotations

import argparse
import csv
import json
import statistics
from pathlib import Path
from typing import Any, Iterable


POOLED_NAME = "E9.1 pooled pose-only"
PROPOSED_NAMES = {
    "exp9_kcn_proposed_pose": "E9.2 proposed pose-only",
    "exp9_kcn_proposed_rgb": "E9.3 proposed RGB/image-only",
    "exp9_kcn_proposed_multimodal": "E9.4 proposed RGB+pose/image",
    "exp9_kcn_v2_pose": "E9.5 research-v2 pose-only",
    "exp9_kcn_v2_rgb": "E9.6 research-v2 RGB/image-only",
    "exp9_kcn_v2_multimodal": "E9.7 research-v2 RGB+pose/image",
}


def as_float(value: Any) -> float | None:
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


def mean_std(values: Iterable[Any]) -> tuple[float | None, float | None]:
    nums = [float(v) for v in values if as_float(v) is not None]
    if not nums:
        return None, None
    return sum(nums) / len(nums), statistics.stdev(nums) if len(nums) > 1 else 0.0


def pct(value: Any) -> str:
    num = as_float(value)
    return "" if num is None else f"{num * 100:.2f}%"


def add_pooled_rows(rows: list[dict[str, Any]], root: Path) -> None:
    for path in sorted((root / "exp9_kcn_pooled_pose").glob("**/pooled_arch_results.json")):
        item = json.loads(path.read_text(encoding="utf-8"))
        args = item.get("args", {})
        seed = args.get("seed", "")
        metrics = item.get("final_metrics", {})
        base = {
            "model": POOLED_NAME,
            "run_name": "kcn_pooled_pose_only",
            "seed": seed,
            "split": "test",
            "source": str(path),
            "params": item.get("params", ""),
        }
        for key, metric in metrics.items():
            if not isinstance(metric, dict):
                continue
            if key.startswith("macro"):
                rows.append(
                    {
                        **base,
                        "row_type": "macro",
                        "task_key": key,
                        "language_code": "",
                        "modality": "pose",
                        "level": "",
                        "top1": metric.get("accuracy"),
                        "top5": metric.get("top5"),
                        "macro_f1": metric.get("macro_f1"),
                        "n": "",
                        "num_classes": "",
                    }
                )
            else:
                rows.append(
                    {
                        **base,
                        "row_type": "task",
                        "task_key": f"pose_{key}",
                        "language_code": "casl" if key == "casl_si" else key,
                        "modality": "pose",
                        "level": "word" if key in {"casl_si", "casl", "ksl"} else "image",
                        "top1": metric.get("accuracy"),
                        "top5": metric.get("top5"),
                        "macro_f1": metric.get("macro_f1"),
                        "n": metric.get("n", ""),
                        "num_classes": (item.get("num_classes") or {}).get(key, ""),
                    }
                )


def add_proposed_rows(rows: list[dict[str, Any]], root: Path) -> None:
    for dirname, model_name in PROPOSED_NAMES.items():
        for path in sorted((root / dirname).glob("*_results.json")):
            item = json.loads(path.read_text(encoding="utf-8"))
            base = {
                "model": model_name,
                "run_name": item.get("run_name", dirname),
                "seed": item.get("seed", ""),
                "source": str(path),
                "params": item.get("params", ""),
            }
            for split in ("val", "test"):
                for key, metric in (item.get(split) or {}).items():
                    if not isinstance(metric, dict):
                        continue
                    rows.append(
                        {
                            **base,
                            "split": split,
                            "row_type": "macro" if key.startswith("macro_") else "task",
                            "task_key": key,
                            "language_code": metric.get("language_code", ""),
                            "modality": metric.get("modality", ""),
                            "level": metric.get("level", ""),
                            "top1": metric.get("top1"),
                            "top5": metric.get("top5"),
                            "macro_f1": metric.get("macro_f1"),
                            "n": metric.get("n", ""),
                            "num_classes": metric.get("num_classes", metric.get("n_tasks", "")),
                        }
                    )


def summarize(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    groups: dict[tuple[str, str, str, str], list[dict[str, Any]]] = {}
    for row in rows:
        key = (row["model"], row["split"], row["row_type"], row["task_key"])
        groups.setdefault(key, []).append(row)
    summary: list[dict[str, Any]] = []
    for (_model, _split, _row_type, _task), group in sorted(groups.items()):
        first = group[0]
        top1_mean, top1_std = mean_std(row.get("top1") for row in group)
        f1_mean, f1_std = mean_std(row.get("macro_f1") for row in group)
        top5_mean, top5_std = mean_std(row.get("top5") for row in group)
        summary.append(
            {
                "model": first["model"],
                "split": first["split"],
                "row_type": first["row_type"],
                "task_key": first["task_key"],
                "language_code": first["language_code"],
                "modality": first["modality"],
                "level": first["level"],
                "runs": len(group),
                "top1_mean": top1_mean,
                "top1_std": top1_std,
                "top5_mean": top5_mean,
                "top5_std": top5_std,
                "macro_f1_mean": f1_mean,
                "macro_f1_std": f1_std,
                "n": first["n"],
                "num_classes": first["num_classes"],
                "params": first["params"],
            }
        )
    return summary


def write_csv(path: Path, rows: list[dict[str, Any]]) -> None:
    if not rows:
        return
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0].keys()))
        writer.writeheader()
        writer.writerows(rows)


def md_table(rows: list[dict[str, Any]]) -> str:
    if not rows:
        return "No rows yet."
    lines = [
        "| Model | Split | Row | Task | Modality | Runs | Top-1 | Macro-F1 | Classes/Tasks |",
        "|---|---|---|---|---:|---:|---:|---:|---:|",
    ]
    for row in rows:
        lines.append(
            f"| {row['model']} | {row['split']} | {row['row_type']} | {row['task_key']} | {row['modality']} | "
            f"{row['runs']} | {pct(row['top1_mean'])} | {pct(row['macro_f1_mean'])} | {row.get('num_classes', '')} |"
        )
    return "\n".join(lines)


def write_markdown(path: Path, summary: list[dict[str, Any]]) -> None:
    test_macros = [r for r in summary if r["split"] == "test" and r["row_type"] == "macro"]
    test_tasks = [r for r in summary if r["split"] == "test" and r["row_type"] == "task"]
    val_macros = [r for r in summary if r["split"] == "val" and r["row_type"] == "macro"]
    test_macros.sort(key=lambda r: (r["model"], r["task_key"]))
    test_tasks.sort(key=lambda r: (r["model"], r["modality"], r["task_key"]))
    val_macros.sort(key=lambda r: (r["model"], r["task_key"]))
    text = "# KSL + CASL + NSL Focused Experiments\n\n"
    text += "## Test Macro Rows\n\n" + md_table(test_macros) + "\n\n"
    text += "## Test Task Rows\n\n" + md_table(test_tasks) + "\n\n"
    text += "## Validation Macro Rows\n\n" + md_table(val_macros) + "\n"
    path.write_text(text, encoding="utf-8")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--results-root", type=Path, default=Path("results"))
    parser.add_argument("--out-dir", type=Path, default=Path("results/exp9_kcn_focus_summary"))
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    args.out_dir.mkdir(parents=True, exist_ok=True)
    rows: list[dict[str, Any]] = []
    add_pooled_rows(rows, args.results_root)
    add_proposed_rows(rows, args.results_root)
    if not rows:
        raise SystemExit(f"No KCN focused results found under {args.results_root}")
    summary = summarize(rows)
    raw = args.out_dir / "kcn_focus_raw_rows.csv"
    csv_path = args.out_dir / "kcn_focus_summary.csv"
    json_path = args.out_dir / "kcn_focus_summary.json"
    md_path = args.out_dir / "kcn_focus_summary.md"
    write_csv(raw, rows)
    write_csv(csv_path, summary)
    json_path.write_text(json.dumps({"raw_rows": rows, "summary": summary}, indent=2), encoding="utf-8")
    write_markdown(md_path, summary)
    print(md_path.read_text(encoding="utf-8"))
    print("Saved:")
    print(" ", raw)
    print(" ", csv_path)
    print(" ", json_path)
    print(" ", md_path)


if __name__ == "__main__":
    main()