data_plot_metrics.py

from pyspark.sql.functions import col, round as _round, last
import os, logging
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
from ea_common import spark, users, filtered_rounds, db_path, DB_DIR

metrics = spark.read.parquet(db_path("metrics"))

# Total number of rounds
logging.info(f"Total number of rounds: {filtered_rounds.groupBy('round_id').count().count()}")

# Set the style and context of the plots
sns.set_style("whitegrid")
sns.set_context("notebook", font_scale=1.5, rc={"lines.linewidth": 2.5})

#
# Helpers: gather metrics/confusion-matrix/skill data for a model without
# plotting anything yet, so all models can be drawn together afterwards.
#

def get_model_metrics(name):
    """Return the last value of each metric stored for `name` (includes the
    confusion-matrix counts, needed separately below)."""
    return metrics.where(col("name") == name).groupBy("metric").agg(last("value").alias("value"))


def has_model_metrics(name):
    return metrics.where(col("name") == name).limit(1).count() > 0


def get_readable_metrics(model_metrics):
    """Non-confusion-matrix metrics for tables/bar charts.

    Rate metrics are shown as percentages; loss metrics stay on their native
    scale.
    """
    readable = model_metrics.filter(~col("metric").rlike("true|false")) \
        .toPandas()[['metric', 'value']]
    percent_metrics = {"accuracy", "precision", "recall", "f1_score", "roc_auc"}
    readable["is_percent"] = readable["metric"].isin(percent_metrics)
    readable["display_value"] = readable.apply(
        lambda row: round(float(row["value"]) * 100, 2) if row["is_percent"] else round(float(row["value"]), 4),
        axis=1,
    )
    return readable


def get_confusion_matrix(model_metrics):
    value = lambda metric: model_metrics.where(col("metric") == metric).select("value").first()[0]
    df_cm = pd.DataFrame(
        [[value("true_negatives"), value("false_positives")],
         [value("false_negatives"), value("true_positives")]],
        index=["Actual Alien Win", "Actual Marine Win"],
        columns=["Predicted Alien Win", "Predicted Marine Win"],
    )
    return df_cm / df_cm.sum().sum() * 100


#
# Combined figures: one window per kind of plot, one subplot/series per model.
#

def show_metrics_tables(models):
    """One row of small tables (graphic, not ASCII), one per model, in a single figure."""
    fig, axes = plt.subplots(1, len(models), figsize=(6 * len(models), 6))
    if len(models) == 1:
        axes = [axes]
    for ax, (label, readable) in zip(axes, models):
        cell_text = [
            [row.metric, f"{row.display_value} %" if row.is_percent else f"{row.display_value}"]
            for row in readable.itertuples(index=False)
        ]
        table = ax.table(cellText=cell_text, colLabels=['metric', 'value'], loc='center', cellLoc='center')
        table.auto_set_font_size(False)
        table.set_fontsize(10)
        table.scale(1, 2)
        ax.axis('off')
        ax.set_title(label)
    fig.suptitle("Model Evaluation Metrics")


def show_metrics_bar_chart(models):
    """Grouped bar chart of the same metrics, to compare models at a glance."""
    combined = pd.concat(
        [readable.assign(model=label) for label, readable in models],
        ignore_index=True,
    )
    plt.figure(figsize=(max(10, 2 * combined['metric'].nunique()), 7))
    sns.barplot(data=combined, x='metric', y='display_value', hue='model')
    plt.title("Model Metrics Comparison")
    plt.xlabel('')
    plt.ylabel('Value (%, except raw loss metrics)')
    plt.xticks(rotation=30, ha='right')
    plt.legend(title=None)


def show_confusion_matrices(models):
    fig, axes = plt.subplots(1, len(models), figsize=(7 * len(models), 6))
    if len(models) == 1:
        axes = [axes]
    for ax, (label, df_cm) in zip(axes, models):
        sns.heatmap(df_cm, annot=True, fmt='.2f', cmap='Blues', ax=ax, cbar=False)
        ax.set_title(label)
    fig.suptitle("Confusion Matrices (% of predictions)")


def show_skill_densities(models):
    """Overlay skill distributions for all models on one plot. Skills from
    different models live on unrelated scales (different ranges/units), so
    each series is z-score normalized (mean 0, std 1) before plotting -
    this keeps the shapes comparable without implying the raw values relate."""
    plt.figure(figsize=(10, 7))
    for label, skill_series in models:
        std = skill_series.std()
        normalized = (skill_series - skill_series.mean()) / std if std else skill_series * 0
        sns.kdeplot(normalized, fill=True, alpha=0.25, label=label)
    plt.axvline(0, color='red', linestyle='dashed', linewidth=1.5)
    plt.title("Skill Distributions (normalized, z-score)")
    plt.xlabel("Skill (standard deviations from each model's mean)")
    plt.ylabel('Density')
    plt.legend()


def openskill_label_from_key(key):
    labels = {
        "pl": "OpenSkill (PlackettLuce)",
        "tmf": "OpenSkill (ThurstoneMostellerFull)",
        "btf": "OpenSkill (BradleyTerryFull)",
        "btp": "OpenSkill (BradleyTerryPart)",
    }
    return labels.get(key, f"OpenSkill ({key})")


#
# Gather data for every model that is available.
#

metrics_tables = []
confusion_matrices = []
skill_series = []

#
# Deep Learning model
#

dl_metrics = get_model_metrics("DL")
metrics_tables.append(("Deep Learning", get_readable_metrics(dl_metrics)))
confusion_matrices.append(("Deep Learning", get_confusion_matrix(dl_metrics)))

if os.path.exists(db_path("skill_dl")):
    skill_dl = spark.read.parquet(db_path("skill_dl"))
    final_skills = users.join(skill_dl, users.id == skill_dl.user_id).orderBy(col("skill_dl").desc())
    final_skills.where(col('wins') >= 12).show(100)
    final_skills.orderBy(col('win_ratio').desc()).where(col('wins') >= 12).show(100)
    skill_series.append(("Deep Learning", final_skills.select("skill_dl").toPandas()['skill_dl']))
else:
    logging.info("skill_dl table not found (run data_analysis_dl.py first), skipping Deep Learning skill plots.")

#
# AdaGrad ("mlt") model (only if data_analysis_mlt.py has been run)
#

if os.path.exists(db_path("skill_mlt")):
    mlt_metrics = get_model_metrics("mlt")
    metrics_tables.append(("AdaGrad", get_readable_metrics(mlt_metrics)))
    confusion_matrices.append(("AdaGrad", get_confusion_matrix(mlt_metrics)))

    skill_mlt = spark.read.parquet(db_path("skill_mlt"))
    final_skills_mlt = users.join(skill_mlt, users.id == skill_mlt.user_id).orderBy(col("skill_mlt").desc())
    final_skills_mlt.where(col('wins') >= 12).show(100)
    skill_series.append(("AdaGrad", final_skills_mlt.select("skill_mlt").toPandas()['skill_mlt']))
else:
    logging.info("skill_mlt table not found (run data_analysis_mlt.py first), skipping AdaGrad plots.")

#
# OpenSkill model variants (legacy + suffixed tables written by data_analysis_os.py)
#

openskill_entries = []
if os.path.isdir(DB_DIR):
    openskill_entries = [
        entry for entry in sorted(os.listdir(DB_DIR))
        if os.path.isdir(db_path(entry)) and (entry == "skill_os" or entry.startswith("skill_os_"))
    ]

for entry in openskill_entries:
    if entry == "skill_os":
        model_key = "pl"
        metrics_name = "os"
        skill_col = "skill_os"
    else:
        model_key = entry.replace("skill_os_", "", 1)
        metrics_name = "os_{}".format(model_key)
        skill_col = "skill_os_{}".format(model_key)

    label = openskill_label_from_key(model_key)
    if has_model_metrics(metrics_name):
        os_metrics = get_model_metrics(metrics_name)
        metrics_tables.append((label, get_readable_metrics(os_metrics)))
        confusion_matrices.append((label, get_confusion_matrix(os_metrics)))
    else:
        logging.info("No metrics found for {} ({}), skipping metric/confusion plots for this model.".format(entry, metrics_name))

    skill_table = spark.read.parquet(db_path(entry))
    if skill_col not in skill_table.columns:
        logging.warning("Table {} is missing column {}, skipping skill density for this model.".format(entry, skill_col))
        continue

    final_skills_os = users.join(skill_table, users.id == skill_table.user_id).orderBy(col(skill_col).desc())
    final_skills_os.where(col('wins') >= 12).show(100)
    skill_series.append((label, final_skills_os.select(skill_col).toPandas()[skill_col]))

if not openskill_entries:
    logging.info("No OpenSkill tables found (run data_analysis_os.py first), skipping OpenSkill plots.")

# Draw everything: one combined figure per kind of plot, instead of one figure per model.
show_metrics_tables(metrics_tables)
show_metrics_bar_chart(metrics_tables)
show_confusion_matrices(confusion_matrices)
show_skill_densities(skill_series)

plt.show()

# Stop Spark session
spark.stop()