plot_LLM_eval_Stats_bigF_plot.py

from pathlib import Path
from typing import Dict, Optional
import re

import matplotlib.pyplot as plt
import pandas as pd


# -----------------------------------------------------------------------------
# Project-relative input/output settings
# -----------------------------------------------------------------------------
# This script is expected to live in the project root:
#   softOpt_illustrative/
#
# The absolute Windows path you gave was:
#   C:\Users\msaiz\OneDrive - weOptimize\0_PhD_Repositories\DMS3_SoftOpt_Illustrative\softOpt_illustrative\results\lambda_sweep\port1_g200_p50\llm_eval_stats\llm_eval_summary.csv
#
# Therefore, the project-relative CSV path is:
SUMMARY_CSV_PATH = Path(
    "results/lambda_sweep/port1_g200_p50/llm_eval_stats/llm_eval_summary.csv"
)
OUTPUT_PLOT_NAME = "llm_eval_scores_bigF.png"


# -----------------------------------------------------------------------------
# Plot font-size settings
# -----------------------------------------------------------------------------
# Leave values as None to keep the plot appearance exactly as before.
# Set only the values you want to customize, for example:
#   PLOT_TITLE_FONT_SIZE = 16
#   PLOT_AXIS_LABEL_FONT_SIZE = 12
#   PLOT_TICK_FONT_SIZE = 10
#   PLOT_LEGEND_FONT_SIZE = 11
#
# PLOT_FONT_SIZE is an optional fallback for all plot text. Specific settings
# below override it when they are not None.
PLOT_FONT_SIZE = 14
PLOT_TITLE_FONT_SIZE = 14
PLOT_AXIS_LABEL_FONT_SIZE = 14
PLOT_TICK_FONT_SIZE = 14
PLOT_LEGEND_FONT_SIZE = 14


def project_path(relative_path: Path) -> Path:
    """Resolve a path relative to this script's folder, i.e. the project root."""

    return Path(__file__).resolve().parent / relative_path


def _plot_font_size(specific_size: Optional[float]) -> Optional[float]:
    return specific_size if specific_size is not None else PLOT_FONT_SIZE


def _fontsize_kwargs(size: Optional[float]) -> Dict[str, float]:
    return {"fontsize": size} if size is not None else {}


def _set_legend(axis) -> None:
    legend_size = _plot_font_size(PLOT_LEGEND_FONT_SIZE)
    if legend_size is None:
        axis.legend()
    else:
        axis.legend(fontsize=legend_size)


def _set_tick_font_sizes(axis) -> None:
    tick_size = _plot_font_size(PLOT_TICK_FONT_SIZE)
    if tick_size is not None:
        axis.tick_params(axis="both", labelsize=tick_size)


def load_summary_csv(summary_csv_path: Path) -> pd.DataFrame:
    """Load the previously generated llm_eval_summary.csv file."""

    if not summary_csv_path.exists():
        raise FileNotFoundError(f"Summary CSV not found: {summary_csv_path}")

    agg_df = pd.read_csv(summary_csv_path)

    required_columns = {
        "candidate_id",
        "mean_score",
        "std_score",
        "minmax_score",
        "proportional_score",
    }
    missing_columns = required_columns - set(agg_df.columns)
    if missing_columns:
        raise ValueError(
            "Summary CSV is missing required columns: "
            + ", ".join(sorted(missing_columns))
        )

    agg_df = agg_df.sort_values("candidate_id").reset_index(drop=True)
    return agg_df


def extract_sample_scores(agg_df: pd.DataFrame) -> pd.DataFrame:
    """Rebuild the raw sample-score dataframe from sample_* columns in the CSV."""

    # IMPORTANT: only take real sample score columns such as sample_01, sample_02, ...
    # Do NOT use col.startswith("sample_") here, because that also includes
    # sample_count. Since sample_count is 5 for every candidate, plotting it
    # creates a fake horizontal row of points at S-score = 5.
    sample_columns = [
        col for col in agg_df.columns
        if re.fullmatch(r"sample_\d+", col)
    ]
    sample_columns = sorted(sample_columns, key=lambda col: int(col.split("_")[1]))
    if not sample_columns:
        raise ValueError(
            "No sample_* columns found in the summary CSV. "
            "They are required to replot the raw sample scatter points."
        )

    scores_df = agg_df[["candidate_id", *sample_columns]].melt(
        id_vars="candidate_id",
        value_vars=sample_columns,
        var_name="sample_id",
        value_name="strategic_score",
    )
    scores_df["strategic_score"] = pd.to_numeric(
        scores_df["strategic_score"], errors="coerce"
    )
    scores_df = scores_df.dropna(subset=["strategic_score"]).reset_index(drop=True)

    if scores_df.empty:
        raise ValueError("No valid sample scores found in the sample_* columns.")

    return scores_df


def plot_scores(agg_df: pd.DataFrame, scores_df: pd.DataFrame, output_path: Path) -> None:
    fig, (ax_raw, ax_norm) = plt.subplots(2, 1, figsize=(10, 8), sharex=True)

    ax_raw.scatter(
        scores_df["candidate_id"],
        scores_df["strategic_score"],
        color="black",
        s=12,
        alpha=0.5,
        label="Samples",
        zorder=1,
    )
    ax_raw.plot(
        agg_df["candidate_id"],
        agg_df["mean_score"],
        marker="o",
        label="Mean S-score",
        zorder=2,
    )
    ax_raw.fill_between(
        agg_df["candidate_id"],
        agg_df["mean_score"] - agg_df["std_score"],
        agg_df["mean_score"] + agg_df["std_score"],
        alpha=0.15,
        label="Std dev",
        zorder=0,
    )
    ax_raw.set_ylabel(
        "S-score", **_fontsize_kwargs(_plot_font_size(PLOT_AXIS_LABEL_FONT_SIZE))
    )
    ax_raw.set_title(
        "LLM S-score summary", **_fontsize_kwargs(_plot_font_size(PLOT_TITLE_FONT_SIZE))
    )
    _set_legend(ax_raw)

    ax_norm.plot(
        agg_df["candidate_id"],
        agg_df["minmax_score"],
        marker="o",
        label="Min-max normalized",
    )
    ax_norm.plot(
        agg_df["candidate_id"],
        agg_df["proportional_score"],
        marker="o",
        label="Proportional normalized",
    )
    ax_norm.set_ylabel(
        "Normalized S-score",
        **_fontsize_kwargs(_plot_font_size(PLOT_AXIS_LABEL_FONT_SIZE)),
    )
    ax_norm.set_xlabel(
        "Candidate ID", **_fontsize_kwargs(_plot_font_size(PLOT_AXIS_LABEL_FONT_SIZE))
    )
    ax_norm.set_ylim(-0.05, 1.05)
    _set_legend(ax_norm)

    _set_tick_font_sizes(ax_raw)
    _set_tick_font_sizes(ax_norm)

    fig.tight_layout()
    output_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output_path, dpi=150, bbox_inches="tight")
    plt.close(fig)


def main() -> None:
    summary_csv_path = project_path(SUMMARY_CSV_PATH)
    output_path = summary_csv_path.parent / OUTPUT_PLOT_NAME

    agg_df = load_summary_csv(summary_csv_path)
    scores_df = extract_sample_scores(agg_df)
    plot_scores(agg_df, scores_df, output_path)

    print(f"Read summary CSV from {summary_csv_path}")
    print(f"Wrote plot to {output_path}")


if __name__ == "__main__":
    main()