make_assignment_misdirection_table.py

#!/usr/bin/env python3
import argparse
import pandas as pd


def safe_div(a, b):
    return a / b if b else None


def prf(tp, fp, fn):
    p = safe_div(tp, tp + fp)
    r = safe_div(tp, tp + fn)
    f1 = safe_div(2 * p * r, p + r) if (p is not None and r is not None and (p + r)) else None
    return p, r, f1


def aggregate_assign(df):
    cols = set(df.columns)

    mis = int(df["assign_misdirected"].sum()) if "assign_misdirected" in cols else 0

    # WITH-legend schema: explicit TP/FP/FN
    if "assign_tp" in cols:
        tp = int(df["assign_tp"].sum())
        fp = int(df["assign_fp"].sum())
        fn = int(df["assign_fn"].sum())
        gt = int(df["assign_gt"].sum())

    # WO-legend schema: no TP, reconstruct from GT - FN
    else:
        required = {"assign_gt", "assign_fn", "assign_fp", "assign_pred"}
        missing = required - cols
        if missing:
            raise KeyError(
                f"Missing required columns: {sorted(missing)}. Available columns={sorted(cols)}"
            )

        gt = int(df["assign_gt"].sum())
        fn = int(df["assign_fn"].sum())
        tp = gt - fn

        fp_explicit = int(df["assign_fp"].sum())
        pred = int(df["assign_pred"].sum())
        fp_from_pred = pred - tp

        # Keep explicit FP if present; this also preserves prior behavior
        fp = fp_explicit if fp_explicit == fp_from_pred else fp_explicit

    p, r, f1 = prf(tp, fp, fn)

    return {
        "N": int(len(df)),
        "assign_gt": int(gt),
        "assign_tp": int(tp),
        "assign_fp": int(fp),
        "assign_fn": int(fn),
        "assign_precision": p,
        "assign_recall": r,
        "assign_f1": f1,
        "assign_misdirected": int(mis),
        "mis_fn": safe_div(mis, fn),
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--with_sag", required=True, help="CSV for SAG with legend")
    ap.add_argument("--with_mag", required=True, help="CSV for MAG with legend")
    ap.add_argument("--with_dag", required=True, help="CSV for DAG with legend")
    ap.add_argument("--wo_sag", required=True, help="CSV for SAG without legend")
    ap.add_argument("--wo_mag", required=True, help="CSV for MAG without legend")
    ap.add_argument("--wo_dag", required=True, help="CSV for DAG without legend")
    ap.add_argument("--out_csv", required=True, help="Output merged CSV path")
    args = ap.parse_args()

    inputs = [
        ("SAG", "With", args.with_sag),
        ("MAG", "With", args.with_mag),
        ("DAG", "With", args.with_dag),
        ("SAG", "Without", args.wo_sag),
        ("MAG", "Without", args.wo_mag),
        ("DAG", "Without", args.wo_dag),
    ]

    rows = []
    for regime, legend, path in inputs:
        df = pd.read_csv(path)
        stats = aggregate_assign(df)
        rows.append({
            "Regime": regime,
            "Legend": legend,
            "N": stats["N"],
            "Assign_GT": stats["assign_gt"],
            "Assign_TP": stats["assign_tp"],
            "Assign_FP": stats["assign_fp"],
            "Assign_FN": stats["assign_fn"],
            "Assign_P": stats["assign_precision"],
            "Assign_R": stats["assign_recall"],
            "Assign_F1": stats["assign_f1"],
            "Misdir": stats["assign_misdirected"],
            "Mis_FN": stats["mis_fn"],
        })

    regime_order = {"SAG": 0, "MAG": 1, "DAG": 2}
    legend_order = {"With": 0, "Without": 1}
    rows = sorted(rows, key=lambda r: (regime_order[r["Regime"]], legend_order[r["Legend"]]))

    out_df = pd.DataFrame(rows)
    out_df.to_csv(args.out_csv, index=False)

    print(out_df.to_string(index=False))
    print(f"\nWrote: {args.out_csv}")


if __name__ == "__main__":
    main()