starbench_performance_evaluation.py

# -*- coding: utf-8 -*-
"""
STAR-Bench performance evaluation — Pareto frontier analysis.

Computes Pareto-optimal sets across the three axes STAR-Bench already
reports per backend:

    latency   = "Avg ProcTime(s)"   (lower is better)
    accuracy  = "Avg Final RMSE"    (lower is better)
    energy    = "Energy (calc)"     (lower is better)

Given one or more ``summary_report_overall.csv`` files (the final table
written by ``starbench_reporting``), the module produces:

    * the 3D Pareto frontier across (latency, accuracy, energy)
    * the three 2D frontiers on every pair
    * optional matplotlib scatter plots with the Pareto points highlighted

A small CLI is provided for batch use::

    python starbench_performance_evaluation.py \\
        --input results/results_20260101_120000 \\
        --output results/results_20260101_120000/pareto \\
        --plot

If ``--input`` is a directory the script searches it recursively for
``summary_report_overall.csv`` and aggregates all rows before computing
the frontier, so runs from multiple datasets can be compared in one shot.
"""
from __future__ import annotations

import argparse
import glob
import logging
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable

import numpy as np
import pandas as pd

logger = logging.getLogger(__name__)

# --- Canonical column names used by starbench_reporting ---
LATENCY_COL = "Avg ProcTime(s)"
ACCURACY_COL = "Avg Final RMSE"
ENERGY_COL = "Energy (calc)"
PIPELINE_COL = "Pipeline"

# All three objectives are "lower is better" by construction.
DEFAULT_OBJECTIVES = (LATENCY_COL, ACCURACY_COL, ENERGY_COL)


@dataclass(frozen=True)
class ParetoResult:
    """Container for a single Pareto-frontier computation."""
    objectives: tuple[str, ...]
    frame: pd.DataFrame           # full frame with a boolean ``pareto`` column
    frontier: pd.DataFrame        # only the Pareto-optimal rows


def pareto_mask(points: np.ndarray) -> np.ndarray:
    """Return a boolean mask selecting Pareto-optimal rows of ``points``.

    ``points`` is an ``(N, D)`` array where every objective is minimised.
    A point ``p`` is Pareto-optimal iff no other point ``q`` is
    component-wise ``<= p`` with at least one strict inequality.

    Rows that contain any NaN are treated as dominated and excluded from
    the frontier (we cannot compare them reliably).
    """
    if points.ndim != 2:
        raise ValueError(f"pareto_mask expects a 2-D array, got shape {points.shape}")

    n = points.shape[0]
    mask = np.ones(n, dtype=bool)

    # Rows with any NaN cannot be part of the frontier.
    finite_rows = np.all(np.isfinite(points), axis=1)
    mask &= finite_rows

    # Naive O(N^2) sweep — N is always small here (one row per backend).
    for i in range(n):
        if not mask[i]:
            continue
        p_i = points[i]
        for j in range(n):
            if i == j or not mask[j]:
                continue
            p_j = points[j]
            # j dominates i ?
            if np.all(p_j <= p_i) and np.any(p_j < p_i):
                mask[i] = False
                break
    return mask


def _coerce_numeric(df: pd.DataFrame, columns: Iterable[str]) -> pd.DataFrame:
    """Return a copy of ``df`` with the selected columns forced to numeric."""
    out = df.copy()
    for col in columns:
        if col in out.columns:
            out[col] = pd.to_numeric(out[col], errors="coerce")
    return out


def load_summary(input_path: str | os.PathLike) -> pd.DataFrame:
    """Load one or many ``summary_report_overall.csv`` files into a DataFrame.

    ``input_path`` may be a single CSV or a directory that is scanned
    recursively. When a directory is given the dataset name (parent
    folder relative to the input root) is added as an extra ``Dataset``
    column so that per-dataset frontiers can be computed downstream.
    """
    path = Path(input_path)
    frames: list[pd.DataFrame] = []

    if path.is_file():
        logger.info(f"Loading single summary file: {path}")
        df = pd.read_csv(path)
        df["Dataset"] = path.parent.name
        frames.append(df)
    elif path.is_dir():
        pattern = str(path / "**" / "summary_report_overall.csv")
        matches = sorted(glob.glob(pattern, recursive=True))
        if not matches:
            raise FileNotFoundError(
                f"No 'summary_report_overall.csv' found under {path}"
            )
        logger.info(f"Found {len(matches)} summary files under {path}")
        for m in matches:
            try:
                df = pd.read_csv(m)
                rel = Path(m).parent.relative_to(path)
                df["Dataset"] = str(rel) if str(rel) != "." else path.name
                frames.append(df)
            except Exception as e:
                logger.warning(f"Skipping unreadable summary {m}: {e}")
    else:
        raise FileNotFoundError(f"Input path does not exist: {path}")

    if not frames:
        raise RuntimeError("No summary data loaded.")

    merged = pd.concat(frames, ignore_index=True)
    return _coerce_numeric(merged, DEFAULT_OBJECTIVES)


def compute_pareto(df: pd.DataFrame, objectives: Iterable[str]) -> ParetoResult:
    """Flag the Pareto-optimal rows of ``df`` for the given objectives.

    All objectives must already exist in ``df`` and be numeric. The
    returned frame is a copy of ``df`` with an extra boolean ``pareto``
    column; ``frontier`` is the subset of rows where ``pareto`` is True,
    sorted on the first objective.
    """
    objectives = tuple(objectives)
    missing = [c for c in objectives if c not in df.columns]
    if missing:
        raise KeyError(f"Missing objective columns: {missing}")

    points = df[list(objectives)].to_numpy(dtype=float)
    mask = pareto_mask(points)

    out = df.copy()
    out["pareto"] = mask
    frontier = out[mask].sort_values(list(objectives)).reset_index(drop=True)
    return ParetoResult(objectives=objectives, frame=out, frontier=frontier)


def compute_all_frontiers(df: pd.DataFrame) -> dict[str, ParetoResult]:
    """Compute the 3D frontier plus all three 2D frontiers.

    Frontiers that require a column missing from ``df`` are silently
    skipped with a warning. This matters for runs where power logging
    was disabled (no ``Energy (calc)``): the module still produces the
    latency/accuracy 2D frontier instead of failing hard.
    """
    results: dict[str, ParetoResult] = {}

    has_latency = LATENCY_COL in df.columns
    has_accuracy = ACCURACY_COL in df.columns
    has_energy = ENERGY_COL in df.columns

    if not (has_latency and has_accuracy):
        missing = [c for c in (LATENCY_COL, ACCURACY_COL) if c not in df.columns]
        raise KeyError(
            f"Input frame is missing mandatory columns {missing}; "
            "cannot compute any Pareto frontier."
        )

    if has_energy:
        results["3d_latency_accuracy_energy"] = compute_pareto(
            df, (LATENCY_COL, ACCURACY_COL, ENERGY_COL)
        )
    else:
        logger.warning(
            f"Column '{ENERGY_COL}' not found; skipping energy-inclusive "
            "frontiers (3D + latency/energy + accuracy/energy)."
        )

    results["2d_latency_accuracy"] = compute_pareto(
        df, (LATENCY_COL, ACCURACY_COL)
    )

    if has_energy:
        results["2d_latency_energy"] = compute_pareto(
            df, (LATENCY_COL, ENERGY_COL)
        )
        results["2d_accuracy_energy"] = compute_pareto(
            df, (ACCURACY_COL, ENERGY_COL)
        )

    return results


def write_frontiers(frontiers: dict[str, ParetoResult], output_dir: str | os.PathLike) -> None:
    """Write one CSV per frontier under ``output_dir``."""
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    for name, result in frontiers.items():
        path = output_dir / f"pareto_{name}.csv"
        cols = [PIPELINE_COL] if PIPELINE_COL in result.frontier.columns else []
        cols += [c for c in result.objectives if c in result.frontier.columns]
        if "Dataset" in result.frontier.columns:
            cols.insert(0, "Dataset")
        extra = [c for c in result.frontier.columns if c not in cols and c != "pareto"]
        result.frontier[cols + extra].to_csv(path, index=False, float_format="%.6f")
        logger.info(
            f"Wrote {name} frontier ({len(result.frontier)} points) -> {path}"
        )

    # Always dump the full annotated table as well, for transparency.
    # Only the frontiers that were actually computed are annotated — if
    # the Energy column was missing upstream, the corresponding columns
    # are simply absent instead of raising a KeyError.
    if not frontiers:
        return
    any_result = next(iter(frontiers.values()))
    full_path = output_dir / "pareto_full_table.csv"
    annotations = {
        "3d_latency_accuracy_energy": "is_pareto_3d",
        "2d_latency_accuracy": "is_pareto_latency_accuracy",
        "2d_latency_energy": "is_pareto_latency_energy",
        "2d_accuracy_energy": "is_pareto_accuracy_energy",
    }
    annotated = any_result.frame.drop(columns=["pareto"])
    for key, col_name in annotations.items():
        if key in frontiers:
            annotated = annotated.assign(**{col_name: frontiers[key].frame["pareto"]})
    annotated.to_csv(full_path, index=False, float_format="%.6f")
    logger.info(f"Wrote annotated full table -> {full_path}")


# --- Optional plotting ---

def _try_import_matplotlib():
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        return plt
    except ImportError:
        logger.warning("matplotlib not installed; skipping plots.")
        return None


def plot_frontiers(frontiers: dict[str, ParetoResult], output_dir: str | os.PathLike) -> None:
    """Render scatter plots with the Pareto points highlighted (best-effort)."""
    plt = _try_import_matplotlib()
    if plt is None:
        return

    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # --- 2D scatter plots ---
    for name in ("2d_latency_accuracy", "2d_latency_energy", "2d_accuracy_energy"):
        if name not in frontiers:
            continue
        result = frontiers[name]
        x_col, y_col = result.objectives
        frame = result.frame.dropna(subset=[x_col, y_col])
        if frame.empty:
            logger.warning(f"No finite points for {name}; skipping plot.")
            continue

        fig, ax = plt.subplots(figsize=(7, 5))
        dominated = frame[~frame["pareto"]]
        optimal = frame[frame["pareto"]].sort_values(x_col)

        ax.scatter(dominated[x_col], dominated[y_col],
                   c="lightgray", s=35, label="dominated")
        ax.scatter(optimal[x_col], optimal[y_col],
                   c="tab:red", s=60, edgecolor="k", zorder=3, label="Pareto")

        if len(optimal) >= 2:
            ax.plot(optimal[x_col], optimal[y_col],
                    c="tab:red", linestyle="--", linewidth=1.2, zorder=2)

        for _, row in optimal.iterrows():
            label = row.get(PIPELINE_COL, "")
            ax.annotate(str(label), (row[x_col], row[y_col]),
                        xytext=(4, 4), textcoords="offset points", fontsize=8)

        ax.set_xlabel(x_col)
        ax.set_ylabel(y_col)
        ax.set_title(f"STAR-Bench Pareto: {x_col} vs {y_col}")
        ax.legend(loc="best")
        ax.grid(True, linestyle=":", alpha=0.5)
        fig.tight_layout()
        fig.savefig(output_dir / f"pareto_{name}.png", dpi=150)
        plt.close(fig)

    # --- 3D scatter ---
    if "3d_latency_accuracy_energy" not in frontiers:
        return
    from mpl_toolkits.mplot3d import Axes3D  # noqa: F401 (registers 3D projection)
    result_3d = frontiers["3d_latency_accuracy_energy"]
    x_col, y_col, z_col = result_3d.objectives
    frame = result_3d.frame.dropna(subset=[x_col, y_col, z_col])
    if frame.empty:
        logger.warning("No finite points for 3D Pareto; skipping.")
        return

    fig = plt.figure(figsize=(8, 6))
    ax = fig.add_subplot(111, projection="3d")
    dominated = frame[~frame["pareto"]]
    optimal = frame[frame["pareto"]]

    ax.scatter(dominated[x_col], dominated[y_col], dominated[z_col],
               c="lightgray", s=30, label="dominated")
    ax.scatter(optimal[x_col], optimal[y_col], optimal[z_col],
               c="tab:red", s=60, edgecolor="k", label="Pareto")

    for _, row in optimal.iterrows():
        label = row.get(PIPELINE_COL, "")
        ax.text(row[x_col], row[y_col], row[z_col], f"  {label}", fontsize=7)

    ax.set_xlabel(x_col)
    ax.set_ylabel(y_col)
    ax.set_zlabel(z_col)
    ax.set_title("STAR-Bench Pareto: latency vs accuracy vs energy")
    ax.legend(loc="best")
    fig.tight_layout()
    fig.savefig(output_dir / "pareto_3d_latency_accuracy_energy.png", dpi=150)
    plt.close(fig)

    logger.info(f"Pareto plots written under {output_dir}")


# --- CLI ---

def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Compute Pareto frontiers across STAR-Bench summary reports "
                    "(latency / accuracy / energy).",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--input", "-i", required=True,
        help="Path to a summary_report_overall.csv file, or a directory to "
             "scan recursively for such files.",
    )
    parser.add_argument(
        "--output", "-o", default="pareto",
        help="Directory where the Pareto CSVs (and plots) will be written.",
    )
    parser.add_argument(
        "--plot", action="store_true",
        help="Also render scatter plots (requires matplotlib).",
    )
    parser.add_argument(
        "--per-dataset", action="store_true",
        help="Compute a separate frontier per Dataset column in addition "
             "to the global one.",
    )
    return parser.parse_args()


def main() -> None:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s - %(name)s [%(levelname)s] - %(message)s",
    )
    args = _parse_args()

    df = load_summary(args.input)
    logger.info(
        f"Loaded {len(df)} backend entries across "
        f"{df['Dataset'].nunique() if 'Dataset' in df.columns else 1} dataset(s)."
    )

    out_dir = Path(args.output)
    out_dir.mkdir(parents=True, exist_ok=True)

    # Global frontier across everything.
    frontiers = compute_all_frontiers(df)
    write_frontiers(frontiers, out_dir)
    if args.plot:
        plot_frontiers(frontiers, out_dir)

    if args.per_dataset and "Dataset" in df.columns:
        for ds_name, ds_df in df.groupby("Dataset"):
            sub_dir = out_dir / f"dataset_{ds_name}"
            sub_frontiers = compute_all_frontiers(ds_df)
            write_frontiers(sub_frontiers, sub_dir)
            if args.plot:
                plot_frontiers(sub_frontiers, sub_dir)

    logger.info(f"Pareto analysis complete. Results under {out_dir.resolve()}")


if __name__ == "__main__":
    main()