eval/aggregate_hits.py

"""Aggregate per-prediction evals.tsv into the headline minority_hit_rate table.

minority_hit_rate = mean(hit_primary) over all predictions in an arm/case, read
from the `<minority_state>__hit_primary` column of each evals.tsv. The minority
(rare) state per case (protocol §Evaluation, CLAUDE.md):

    KaiB  -> state_A_2QKE   (fold-switched 2QKE chain B)
    GA98  -> GA98_2LHC      (alternate GA-fold basin)
    GB98  -> GA98_2LHC      (GA98 fold is the minority conformation for GB98)

NOTE (GA_GB CLI caveat): both GA98 and GB98 predictions are scored under
`--case GA_GB`; they differ only in which reference RMSD is read. GB98's minority
state is therefore the GA98_2LHC column, NOT its own GB98_2LHD (native) column.

Usage:
    python aggregate_hits.py --bench-root /path/to/bench   # standard layout
"""

from __future__ import annotations

import argparse
import csv
from pathlib import Path

# case label -> column prefix of the minority state's hit_primary flag
MINORITY_STATE = {
    "KaiB": "state_A_2QKE",
    "GA98": "GA98_2LHC",
    "GB98": "GA98_2LHC",
}

# Standard benchmark layout: (arm, case_label, evals.tsv path relative to bench root)
BENCH_ENTRIES: list[tuple[str, str, str]] = []
for _arm in ("mosaic_raw", "gradient_raw", "contrast_raw", "region_cluster_raw"):
    BENCH_ENTRIES += [
        (_arm, "KaiB", f"results/baseline_p12/{_arm}/KaiB/KaiB/refine_per_state/evals.tsv"),
        (_arm, "GA98", f"results/baseline_p12/{_arm}/GA_GB/GA98/refine_per_state/evals.tsv"),
        (_arm, "GB98", f"results/baseline_p12/{_arm}/GA_GB/GB98/refine_per_state/evals.tsv"),
    ]
BENCH_ENTRIES += [
    ("afcluster", "KaiB", "results/baseline/afcluster/KaiB/refine_per_state/evals.tsv"),
    ("afcluster", "GA98", "results/baseline/afcluster/GA_GB/GA98/refine_per_state/evals.tsv"),
    ("afcluster", "GB98", "results/baseline/afcluster/GA_GB/GB98/refine_per_state/evals.tsv"),
]


def minority_hit_rate(tsv_path: Path, case: str) -> tuple[int, int, float]:
    """Return (hits, total, rate) for the minority-state hit_primary column."""
    col = f"{MINORITY_STATE[case]}__hit_primary"
    with Path(tsv_path).open() as f:
        rows = list(csv.DictReader(f, delimiter="\t"))
    if not rows:
        raise ValueError(f"empty evals.tsv: {tsv_path}")
    if col not in rows[0]:
        raise KeyError(f"column {col!r} not in {tsv_path} (have {list(rows[0])})")
    hits = sum(1 for r in rows if r[col] == "1")
    total = len(rows)
    return hits, total, hits / total


def aggregate(bench_root: Path,
              entries: list[tuple[str, str, str]] | None = None) -> list[dict]:
    """Compute minority_hit_rate for every (arm, case) entry under bench_root."""
    entries = entries if entries is not None else BENCH_ENTRIES
    out = []
    for arm, case, rel in entries:
        hits, total, rate = minority_hit_rate(Path(bench_root) / rel, case)
        out.append({"arm": arm, "case": case, "hits": hits,
                    "total": total, "minority_hit_rate": rate})
    return out


def format_table(rows: list[dict]) -> str:
    hdr = f"{'arm':<20} {'case':<6} {'hits':>6} {'total':>6} {'minority_hit_rate':>18}"
    lines = [hdr, "-" * len(hdr)]
    for r in rows:
        lines.append(f"{r['arm']:<20} {r['case']:<6} {r['hits']:>6} "
                     f"{r['total']:>6} {r['minority_hit_rate']:>18.4f}")
    return "\n".join(lines)


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--bench-root", required=True, type=Path,
                    help="Benchmark bundle root (contains results/)")
    args = ap.parse_args(argv)
    rows = aggregate(args.bench_root)
    print(format_table(rows))
    return 0


if __name__ == "__main__":
    import sys
    sys.exit(main())