reproduce_benchmark.py

#!/usr/bin/env python3
"""Reproduce the SF-Cluster headline `minority_hit_rate` benchmark, CPU-only.

Two modes:
  --mode precomputed  read the shipped per-prediction evals.tsv and aggregate
                      (default; fast, no structure parsing).
  --mode score        re-run the scorer (batch_eval -> evaluate_prediction) over
                      the shipped prediction PDBs, regenerating each evals.tsv,
                      then aggregate.

Both build the case x arm minority_hit_rate table, print it, and ASSERT every
value is within +/-0.02 of the published target. Exits non-zero on any mismatch.

The benchmark root is resolved relative to THIS file (bundle is relocatable):
    <this dir>/bench    (contains data/, configs/, results/)

Dependencies (pip only): numpy, biopython, pyyaml. TMalign is OPTIONAL and only
fills tmalign_* columns; it is not part of the hit criterion.

Usage:
    python reproduce_benchmark.py                 # precomputed
    python reproduce_benchmark.py --mode score    # re-score from PDBs
"""

from __future__ import annotations

import argparse
import os
import subprocess
import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
BENCH_ROOT = HERE / "bench"
EVAL_DIR = HERE / "eval"
sys.path.insert(0, str(EVAL_DIR))

from aggregate_hits import BENCH_ENTRIES, aggregate, format_table  # noqa: E402

# Published minority_hit_rate targets (arm, case) -> value.
PUBLISHED = {
    ("mosaic_raw", "KaiB"): 0.9500, ("mosaic_raw", "GA98"): 0.9250, ("mosaic_raw", "GB98"): 0.1875,
    ("gradient_raw", "KaiB"): 0.6500, ("gradient_raw", "GA98"): 0.5000, ("gradient_raw", "GB98"): 0.5000,
    ("contrast_raw", "KaiB"): 0.6750, ("contrast_raw", "GA98"): 0.4500, ("contrast_raw", "GB98"): 0.0750,
    ("region_cluster_raw", "KaiB"): 0.5375, ("region_cluster_raw", "GA98"): 0.5000, ("region_cluster_raw", "GB98"): 0.4500,
    ("afcluster", "KaiB"): 0.3625, ("afcluster", "GA98"): 0.4625, ("afcluster", "GB98"): 0.4375,
}
TOL = 0.02

# case label -> --case value accepted by the scorer CLI (GA98/GB98 share GA_GB).
CASE_CLI = {"KaiB": "KaiB", "GA98": "GA_GB", "GB98": "GA_GB"}


def rescore() -> None:
    """Re-run batch_eval over every prediction dir, regenerating evals.tsv."""
    env = dict(os.environ, SF_BENCH_ROOT=str(BENCH_ROOT))
    batch = EVAL_DIR / "batch_eval.py"
    for arm, case, rel in BENCH_ENTRIES:
        evals = BENCH_ROOT / rel
        root = evals.parent
        cli_case = CASE_CLI[case]
        print(f"[score] {arm}/{case} (--case {cli_case}) <- {root}")
        r = subprocess.run(
            [sys.executable, str(batch), "--case", cli_case,
             "--root", str(root), "--out", str(evals)],
            env=env,
        )
        if r.returncode != 0:
            raise SystemExit(f"batch_eval failed for {arm}/{case} (rc={r.returncode})")


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--mode", choices=["precomputed", "score"], default="precomputed")
    args = ap.parse_args(argv)

    if args.mode == "score":
        rescore()

    rows = aggregate(BENCH_ROOT)
    print(format_table(rows))
    print()

    failures = []
    for r in rows:
        key = (r["arm"], r["case"])
        target = PUBLISHED.get(key)
        if target is None:
            continue
        delta = abs(r["minority_hit_rate"] - target)
        status = "OK" if delta <= TOL else "MISMATCH"
        if delta > TOL:
            failures.append((key, r["minority_hit_rate"], target, delta))
        print(f"  {r['arm']:<20} {r['case']:<6} "
              f"got={r['minority_hit_rate']:.4f} target={target:.4f} "
              f"d={delta:.4f} [{status}]")

    if failures:
        print(f"\nFAIL: {len(failures)} value(s) outside +/-{TOL}:")
        for key, got, target, delta in failures:
            print(f"  {key}: got {got:.4f}, target {target:.4f} (d={delta:.4f})")
        return 1
    print(f"\nPASS: all {len(rows)} values within +/-{TOL} of published targets.")
    return 0


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