reproduce.py

"""Reproduce every number and figure reported in the README.

Run from the repository root::

    pip install -e ".[all]"
    python -m similarity_ensemble.datasets --nltk      # WordNet corpora (once)
    python benchmarks/reproduce.py

What it does
------------
1. Downloads the four word-similarity benchmarks (RG65, MC30, WordSim353,
   SimLex-999) from their distributors (cached after the first run).
2. Evaluates 12 base measures — 6 lexical + 6 WordNet — and 6 ensembles on
   each dataset, plus the fold-honest ``best_single_cv`` baseline (the best
   base measure selected on training folds only). Supervised ensembles are
   scored with 5-fold cross-validation on out-of-fold predictions.
   Pass ``--datasets men simverb3500`` (or any registry keys) to extend the
   run beyond the default four benchmarks.
3. Writes per-dataset results to ``benchmarks/results/`` (CSV + LaTeX) and a
   combined Markdown table to ``benchmarks/results/summary.md``.
4. Renders the figures used in the README to ``assets/``.

Everything is seeded (``SEED = 0``), so the output is deterministic given the
same dataset files. Total runtime is a few minutes on a laptop.
"""

from __future__ import annotations

import os
import sys
import time

import numpy as np
import pandas as pd

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from similarity_ensemble import (  # noqa: E402
    datasets,
    ensembles,
    evaluation,
    measures,
    metrics,
    viz,
)

SEED = 0
CV_FOLDS = 5
DATASETS = ["rg65", "mc30", "wordsim353", "simlex999"]

HERE = os.path.dirname(os.path.abspath(__file__))
RESULTS_DIR = os.path.join(HERE, "results")
ASSETS_DIR = os.path.join(HERE, "..", "assets")

# Repo-standard colors (validated colorblind-safe pair)
C_BASE, C_ENSEMBLE = "#4C72B0", "#DD8452"


def full_measure_set():
    """6 lexical + 6 WordNet measures. WordNet requires the NLTK corpora."""
    ms = dict(measures.LEXICAL_MEASURES)
    for method in ["path", "wup", "lch", "res", "lin", "jcn"]:
        ms[f"wn_{method}"] = (
            lambda a, b, m=method: measures.wordnet_similarity(a, b, m)
        )
    return ms


def run_dataset(key: str, ms) -> pd.DataFrame:
    df = datasets.load(key)
    pairs = list(df[["word1", "word2"]].itertuples(index=False, name=None))
    gold = df["gold"].to_numpy(dtype=float)
    t0 = time.time()
    results = evaluation.benchmark(
        pairs, gold, measures=ms, cv_folds=CV_FOLDS, seed=SEED
    )
    print(f"[{key}] {len(df)} pairs, {time.time() - t0:.1f}s")
    print(evaluation.summarize(results), "\n")
    results.to_csv(os.path.join(RESULTS_DIR, f"{key}.csv"))
    with open(os.path.join(RESULTS_DIR, f"{key}.tex"), "w") as f:
        f.write(
            evaluation.to_latex_table(
                results,
                caption=f"Base measures vs. ensembles on {key.upper()} "
                f"(supervised ensembles: {CV_FOLDS}-fold CV).",
                label=f"tab:{key}",
            )
        )
    return results


def summary_markdown(all_results: dict) -> str:
    """Combined Spearman table: one row per method, one column per dataset."""
    table = pd.DataFrame(
        {k: r["spearman"] for k, r in all_results.items()}
    ).round(3)
    kinds = next(iter(all_results.values()))["kind"]
    table.insert(0, "kind", kinds)
    table = table.sort_values(table.columns[-1], ascending=False)
    lines = [
        "# Benchmark summary (Spearman rho)",
        "",
        "Supervised ensembles are out-of-fold "
        f"({CV_FOLDS}-fold CV, seed={SEED}). "
        "Generated by `benchmarks/reproduce.py` — regenerate with one command.",
        "",
        table.to_markdown(),
        "",
    ]
    for key, r in all_results.items():
        lines += [f"## {key}", "", "```", evaluation.summarize(r), "```", ""]
    return "\n".join(lines)


def hero_figure(all_results: dict):
    """Best single measure vs. best ensemble, per dataset (Spearman)."""
    import matplotlib.pyplot as plt

    names, base_v, ens_v, base_n, ens_n = [], [], [], [], []
    for key, r in all_results.items():
        base = r[r["kind"] == "base"]["spearman"]
        ens = r[r["kind"] == "ensemble"]["spearman"]
        names.append(key.upper())
        base_v.append(base.max())
        ens_v.append(ens.max())
        base_n.append(base.idxmax().replace("_", " "))
        ens_n.append(ens.idxmax().replace("_", " "))

    x = np.arange(len(names))
    w = 0.38
    fig, ax = plt.subplots(figsize=(9, 4.6))
    b1 = ax.bar(x - w / 2, base_v, w, color=C_BASE, label="best single measure")
    b2 = ax.bar(x + w / 2, ens_v, w, color=C_ENSEMBLE, label="best ensemble (CV)")
    for bars, labels in [(b1, base_n), (b2, ens_n)]:
        for rect, lab in zip(bars, labels):
            ax.text(
                rect.get_x() + rect.get_width() / 2,
                rect.get_height() + 0.012,
                f"{rect.get_height():.3f}\n{lab}",
                ha="center",
                va="bottom",
                fontsize=8,
            )
    ax.set_xticks(x)
    ax.set_xticklabels(names)
    ax.set_ylabel("Spearman rho vs. human gold ratings")
    ax.set_ylim(0, min(1.0, max(ens_v + base_v) + 0.18))
    wins = sum(e > b for e, b in zip(ens_v, base_v))
    ax.set_title(
        f"Stacking beats the best single measure on {wins} of {len(names)} "
        "benchmarks\n(12 base measures, supervised ensembles scored out-of-fold)"
    )
    ax.legend(loc="upper right", frameon=False)
    ax.spines[["top", "right"]].set_visible(False)
    ax.grid(axis="y", alpha=0.25)
    ax.set_axisbelow(True)
    fig.tight_layout()
    fig.savefig(os.path.join(ASSETS_DIR, "hero_results.png"), dpi=200)


def main(argv=None):
    import argparse

    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument(
        "--datasets",
        nargs="+",
        default=DATASETS,
        help="registry keys to benchmark (default: %(default)s)",
    )
    args = parser.parse_args(argv)
    keys = args.datasets

    os.makedirs(RESULTS_DIR, exist_ok=True)
    os.makedirs(ASSETS_DIR, exist_ok=True)

    try:
        ms = full_measure_set()
        measures.wordnet_similarity("car", "automobile", "wup")  # smoke test
    except ImportError:
        print(
            "!! WordNet measures unavailable (install nltk + corpora); "
            "falling back to lexical measures only.\n"
            "   pip install nltk && python -m similarity_ensemble.datasets --nltk"
        )
        ms = dict(measures.LEXICAL_MEASURES)

    all_results = {k: run_dataset(k, ms) for k in keys}

    with open(os.path.join(RESULTS_DIR, "summary.md"), "w") as f:
        f.write(summary_markdown(all_results))

    # ---- Figures ---------------------------------------------------------
    hero_figure(all_results)

    ws = datasets.load("wordsim353")
    pairs = list(ws[["word1", "word2"]].itertuples(index=False, name=None))
    gold = ws["gold"].to_numpy(dtype=float)
    feat = measures.build_feature_table(pairs, ms)
    cols = measures.measure_columns(feat)
    X = feat[cols].to_numpy(dtype=float)

    fig = viz.base_vs_ensemble_bars(all_results["wordsim353"])
    fig.savefig(os.path.join(ASSETS_DIR, "wordsim353_by_method.png"), dpi=200)

    fig = viz.measure_correlation_heatmap(feat)
    fig.savefig(os.path.join(ASSETS_DIR, "measure_diversity_heatmap.png"), dpi=200)

    ga = ensembles.GAWeightedEnsemble(
        metric=metrics.spearman, generations=80, seed=SEED
    ).fit(X, gold)
    fig = viz.ga_convergence(ga.history_)
    fig.savefig(os.path.join(ASSETS_DIR, "ga_convergence.png"), dpi=200)

    fig = viz.contribution_weights(ga.weights, cols)
    fig.savefig(os.path.join(ASSETS_DIR, "ensemble_weights.png"), dpi=200)

    print(f"Results  -> {os.path.relpath(RESULTS_DIR)}")
    print(f"Figures  -> {os.path.relpath(ASSETS_DIR)}")


if __name__ == "__main__":
    main()