code/bundle/analyze.py

#!/usr/bin/env python3
"""Correlation / regression analysis matching paper Table 2 protocol."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
from scipy import stats
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import LeaveOneOut


def pearson(x, y):
    x = np.asarray(x, dtype=np.float64)
    y = np.asarray(y, dtype=np.float64)
    r, p = stats.pearsonr(x, y)
    return float(r), float(p)


def loo_r2(x, y):
    x = np.asarray(x, dtype=np.float64)
    y = np.asarray(y, dtype=np.float64)
    if x.ndim == 1:
        x = x.reshape(-1, 1)
    loo = LeaveOneOut()
    preds = np.zeros_like(y)
    for train, test in loo.split(x):
        reg = LinearRegression().fit(x[train], y[train])
        preds[test] = reg.predict(x[test])
    ss_res = ((y - preds) ** 2).sum()
    ss_tot = ((y - y.mean()) ** 2).sum()
    return float(1 - ss_res / ss_tot) if ss_tot > 0 else float("nan")


def ols_r2(x, y):
    x = np.asarray(x, dtype=np.float64)
    y = np.asarray(y, dtype=np.float64)
    if x.ndim == 1:
        x = x.reshape(-1, 1)
    reg = LinearRegression().fit(x, y)
    return float(reg.score(x, y))


def analyze(rows: list[dict]) -> dict:
    ok = [r for r in rows if r.get("binding") is not None and r.get("error") is None]
    binding = [r["binding"] for r in ok]
    out = {"n_models": len(ok), "models": [r["model"] for r in ok]}

    for metric in ["G.PR", "G.Iso", "L.Iso", "JER", "disc"]:
        vals = [r[metric] for r in ok if r.get(metric) is not None]
        b = [r["binding"] for r in ok if r.get(metric) is not None]
        if len(vals) >= 3:
            r, p = pearson(vals, b)
            out[metric] = {"r": r, "p": p, "n": len(vals), "R2": ols_r2(vals, b)}
        else:
            out[metric] = {"r": None, "p": None, "n": len(vals)}

    # JER + Disc bivariate
    both = [r for r in ok if r.get("JER") is not None and r.get("disc") is not None]
    if len(both) >= 4:
        x = np.stack([[r["JER"], r["disc"]] for r in both], axis=0)
        y = np.array([r["binding"] for r in both], dtype=np.float64)
        out["JER+Disc"] = {
            "R2": ols_r2(x, y),
            "LOO_R2": loo_r2(x, y),
            "n": len(both),
        }

    # Claim 4 style summary
    by_fam = {}
    for r in ok:
        by_fam.setdefault(r.get("family") or "?", []).append(r)
    out["by_family"] = {
        fam: {
            "binding_mean": float(np.mean([x["binding"] for x in rs])),
            "JER_mean": float(np.mean([x["JER"] for x in rs if x.get("JER") is not None]))
            if any(x.get("JER") is not None for x in rs)
            else None,
            "models": [x["model"] for x in rs],
        }
        for fam, rs in by_fam.items()
    }

    claim4 = {
        "var_decorr": [r for r in ok if r.get("family") == "Var-Decorr"],
        "clip": [r for r in ok if "CLIP" in (r.get("model") or "") and "EVA" not in (r.get("model") or "")],
    }
    out["claim4"] = {
        k: {
            "binding_mean": float(np.mean([x["binding"] for x in v])) * 100 if v else None,
            "JER_mean": float(np.mean([x["JER"] for x in v if x.get("JER") is not None])) if v else None,
            "n": len(v),
            "rows": [{"model": x["model"], "binding": x["binding"], "JER": x.get("JER")} for x in v],
        }
        for k, v in claim4.items()
    }
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--results", type=Path, required=True)
    ap.add_argument("--out", type=Path, default=None)
    args = ap.parse_args()
    data = json.loads(args.results.read_text())
    rows = data["rows"] if isinstance(data, dict) and "rows" in data else data
    summary = analyze(rows)
    text = json.dumps(summary, indent=2)
    print(text)
    out = args.out or args.results.with_name(args.results.stem + "_analysis.json")
    out.write_text(text)


if __name__ == "__main__":
    main()