stats_testing.py

# Purpose:
# Read Schnapsen ablation results from results/ablation_results.csv
# Perform two simple statistical tests.
# Direct comparison: binomial test of win rate vs p0.
# Indirect comparison: two-proportion z-test comparing each ML bot to a baseline.
# Print tables + export a CSV with all test outputs.

from __future__ import annotations
from dataclasses import dataclass
from typing import Dict, Tuple
from scipy.stats import binomtest
import pandas as pd
import argparse
import math

# Normal distribution helpers.
# Computing z-test p-values using the standard normal CDF / SF.
def normal_cdf(x: float) -> float:
    return 0.5 * (1.0 + math.erf(x / math.sqrt(2.0)))

def normal_sf(x: float) -> float:
    return 1.0 - normal_cdf(x)

# Indirect comparison: two-proportion z test.
# Compare two win rates p1 = w1 / n1 and p2 = w2 / n2 under the null hypothesis p1 = p2.
def two_proportion_ztest(w1: int, n1: int, w2: int, n2: int) -> Tuple[float, float]:
    if n1 <= 0 or n2 <= 0:
        return float("nan"), float("nan")

    p1 = w1 / n1
    p2 = w2 / n2
    p_pool = (w1 + w2) / (n1 + n2)
    var = p_pool * (1.0 - p_pool) * (1.0 / n1 + 1.0 / n2)
    if var <= 0.0:
        return float("nan"), float("nan")
    z = (p1 - p2) / math.sqrt(var)
    p = 2.0 * normal_sf(abs(z))
    return z, p

# Direct comparison: exact binomial test.
# For each bot vs opponent pairing, test whether observed win rate differs from p0.
def binomial_test_two_sided(k: int, n: int, p0: float = 0.5) -> float:
    return float(binomtest(k, n, p=p0, alternative="two-sided").pvalue)

# Simple container for a single ml_bot vs opponent.
@dataclass(frozen=True)
class Pairing:
    ml_bot: str
    opponent: str
    wins: int
    games: int

    # Win rate computed over all recorded games.
    @property
    def win_rate(self) -> float:
        return self.wins / self.games if self.games > 0 else float("nan")

# Load ablation_results.csv and validate required columns.
def load_results(csv_path: str) -> pd.DataFrame:
    df = pd.read_csv(csv_path)
    required = {"ml_bot", "opponent", "games", "wins"}
    missing = required - set(df.columns)
    if missing:
        raise ValueError(f"CSV missing required columns: {sorted(missing)}")

    # Ensure numeric fields are ints.
    df["wins"] = df["wins"].astype(int)
    df["games"] = df["games"].astype(int)

    return df

# Convert the dataframe into a dictionary for easy lookups by bot and opponent.
def build_pairing_lookup(df: pd.DataFrame) -> Dict[Tuple[str, str], Pairing]:
    lookup: Dict[Tuple[str, str], Pairing] = {}
    for row in df.itertuples(index=False):
        lookup[(row.ml_bot, row.opponent)] = Pairing(
            ml_bot=row.ml_bot,
            opponent=row.opponent,
            wins=int(row.wins),
            games=int(row.games),
        )
    return lookup

def main() -> None:
    # CLI arguments.
    ap = argparse.ArgumentParser()
    ap.add_argument("--csv", default="results/ablation_results.csv",
                    help="Path to ablation_results.csv")
    ap.add_argument("--baseline", default="ml_none",
                    help="Baseline bot for indirect comparisons (default: ml_none)")
    ap.add_argument("--p0", type=float, default=0.5,
                    help="Null win probability for direct binomial tests (default: 0.5)")
    ap.add_argument("--out", default="results/stat_tests.csv",
                    help="Output CSV path for statistical test summary")
    args = ap.parse_args()

    # Load data and builds a lookup structure.
    df = load_results(args.csv)
    lookup = build_pairing_lookup(df)

    bots = sorted(df["ml_bot"].unique().tolist())
    opponents = sorted(df["opponent"].unique().tolist())

    if args.baseline not in bots:
        raise ValueError(f"Baseline '{args.baseline}' not found. Available bots: {bots}")

    # Direct tests: binomial vs p0
    # A refers to the ML bot.
    # B refers to the opponent.
    direct_rows = []
    for b in bots:
        for opp in opponents:
            p = lookup.get((b, opp))
            if p is None:
                continue

            pval = binomial_test_two_sided(p.wins, p.games, p0=args.p0)

            wins_b = p.games - p.wins
            games_b = p.games
            win_rate_b = wins_b / games_b if games_b > 0 else float("nan")

            direct_rows.append(
                {
                    "test_type": "binomial_vs_p0",
                    "opponent": opp,
                    "bot_A": b,
                    "bot_B": opp,
                    "wins_A": p.wins,
                    "games_A": p.games,
                    "win_rate_A": p.win_rate,
                    "wins_B": wins_b,
                    "games_B": games_b,
                    "win_rate_B": win_rate_b,
                    "delta_win_rate": p.win_rate - args.p0,
                    "z": "",
                    "p_value": pval,
                    "notes": f"H0: win_rate({b} vs {opp}) = {args.p0} (two-sided exact binomial test)",
                }
            )

    # Indirect tests: two-proportion z vs baseline
    # For each opponent, compare each ML bot against the chosen baseline.
    compare_rows = []
    for opp in opponents:
        base = lookup.get((args.baseline, opp))
        if base is None:
            continue

        for b in bots:
            if b == args.baseline:
                continue

            other = lookup.get((b, opp))
            if other is None:
                continue

            z, pval = two_proportion_ztest(other.wins, other.games, base.wins, base.games)
            delta = other.win_rate - base.win_rate

            compare_rows.append(
                {
                    "test_type": "two_proportion_ztest_vs_baseline",
                    "opponent": opp,
                    "bot_A": b,
                    "bot_B": args.baseline,
                    "wins_A": other.wins,
                    "games_A": other.games,
                    "win_rate_A": other.win_rate,
                    "wins_B": base.wins,
                    "games_B": base.games,
                    "win_rate_B": base.win_rate,
                    "delta_win_rate": delta,
                    "z": z,
                    "p_value": pval,
                    "notes": "H0: pA = pB (two-sided two-proportion z-test)",
                }
            )

    # Combine results into one table for export.
    out_df = pd.DataFrame(direct_rows + compare_rows)

    # Console output formatting.
    pd.set_option("display.max_rows", 200)
    pd.set_option("display.max_columns", 50)
    pd.set_option("display.width", 140)

    print("\nIndirect Comparisons: two-proportion z-tests vs baseline")
    indirect = out_df[out_df["test_type"] == "two_proportion_ztest_vs_baseline"].copy()
    if not indirect.empty:
        print(indirect.sort_values(["opponent", "p_value"]).to_string(index=False))
    else:
        print("No indirect comparisons available")

    print()
    print("===========================================================================================================")
    print("Direct Comparisons: exact binomial test vs p0")
    direct = out_df[out_df["test_type"] == "binomial_vs_p0"].copy()
    if not direct.empty:
        print(direct.sort_values(["opponent", "bot_A"]).to_string(index=False))
    else:
        print("No direct comparisons available")

    # Export a CSV summary for report
    out_df.to_csv(args.out, index=False)
    print()
    print(f"\nWrote statistical test summary to: {args.out}")

if __name__ == "__main__":
    main()