Raptor/validation/tools/pim/ablation/generate_plots.py

#!/usr/bin/env python3
"""Plot percentage changes from an experiment's combined results_ablation.csv."""

from __future__ import annotations

import argparse
import csv
import math
import sys
import warnings
from dataclasses import dataclass
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from plot_style import COLORS, configure_style, save_figure, value_label

import matplotlib.pyplot as plt
from matplotlib.patches import Patch
from matplotlib.ticker import PercentFormatter

from run_ablation import ABLATION_DEFAULT_MODELS, DEFAULT_RESULTS_PATH, RESULT_FIELDS, VARIANTS


VARIANT_COLORS = {variant.name: color for variant, color in zip(VARIANTS, COLORS)}
STATUSES = ("PASS", "FAIL", "SKIPPED", "NA")


@dataclass(frozen=True)
class Result:
    variant: str
    arch: str
    model: str
    mode: str
    pipeline: int
    validation: str
    metrics: dict[str, float]


def percentage(value: str) -> float:
    if value == "NA":
        return math.nan
    if not value.endswith("%"):
        raise ValueError(f"Expected a percentage or NA, got {value!r}")
    parsed = float(value[:-1])
    if not math.isfinite(parsed):
        raise ValueError(f"Non-finite percentage: {value!r}")
    return parsed


def load_results(path: Path) -> list[Result]:
    rows = []
    seen = set()
    invalid_pairs = []
    with path.open(newline="", encoding="utf-8-sig") as stream:
        reader = csv.DictReader(stream)
        if not {"variant", *RESULT_FIELDS} <= set(reader.fieldnames or []):
            raise ValueError("Expected the combined ablation CSV with variant and percentage columns; "
                             "a per-variant results_comparison.csv cannot be used here.")
        for row in reader:
            variant, mode, status = row["variant"], row["mode"], row["functional_validation"]
            if variant not in VARIANT_COLORS or mode not in ("latency", "throughput") or status not in STATUSES:
                raise ValueError(f"Unknown ablation variant, mode, or validation status: {variant}, {mode}, {status}")
            pipeline = int(row["raptor_pipeline"])
            key = (variant, row["arch"], row["model"], mode, pipeline)
            if key in seen:
                raise ValueError(f"Duplicate ablation configuration: {key}")
            seen.add(key)
            metrics = {metric: percentage(row[metric])
                       for metric in ("latency", "throughput", "power", "energy")}
            # Each may round to -100% alone; reciprocal metrics cannot both do so.
            if mode == "throughput" and metrics["latency"] == metrics["throughput"] == -100:
                metrics["latency"] = metrics["throughput"] = math.nan
                invalid_pairs.append(key)
            rows.append(Result(variant, row["arch"], row["model"], mode, pipeline, status, metrics))
    if not rows:
        raise ValueError("The ablation CSV contains no results")
    if invalid_pairs:
        warnings.warn(f"{len(invalid_pairs)} invalid reciprocal latency/throughput percentage pairs; "
                      f"displaying N/A. Examples: {invalid_pairs[:3]}", stacklevel=2)
    return rows


def draw_metric(results: list[Result], metric: str, ylabel: str, output_dir: Path) -> None:
    architectures = sorted({row.arch for row in results})
    present_models = {row.model for row in results}
    models = [model for model in ABLATION_DEFAULT_MODELS if model in present_models]
    models += sorted(present_models - set(models))
    variants = [name for name in VARIANT_COLORS if any(row.variant == name for row in results)]
    lookup = {(row.arch, row.model, row.variant): row for row in results}
    values = [row.metrics[metric] for row in results if math.isfinite(row.metrics[metric])]
    lower, upper = min(0, min(values, default=0)), max(0, max(values, default=0))
    padding = max(upper - lower, 1) * 0.35
    fig, axes = plt.subplots(1, len(architectures), sharey=True, squeeze=False,
                             figsize=(7 * len(architectures), 5.2), layout="constrained")
    handles = [Patch(facecolor=VARIANT_COLORS[name], label=name) for name in variants]
    if len(values) < len(architectures) * len(models) * len(variants):
        handles.append(Patch(facecolor="white", edgecolor="#B2BAC5", label="N/A: missing measurement"))
    fig.legend(handles=handles, loc="outside upper center", ncol=4, frameon=False, fontsize=9.5)
    width = 0.8 / len(variants)
    mode, pipeline = results[0].mode, results[0].pipeline
    for ax, arch in zip(axes[0], architectures):
        ax.set_axisbelow(True)
        ax.set_ylim(lower - padding if lower < 0 else 0, upper + padding)
        ax.set_xlim(-0.5, len(models) - 0.5)
        ax.set_xticks(range(len(models)), [model.upper() for model in models], fontsize=9.5)
        ax.set_xlabel(f"{arch.replace('arch-', 'ARCHITECTURE ').upper()}\n"
                      f"{mode.capitalize()}, {pipeline} stage" + ("s" if pipeline != 1 else ""),
                      fontsize=11, weight="bold", labelpad=8)
        ax.set_ylabel(ylabel, fontsize=10.5)
        ax.yaxis.set_major_formatter(PercentFormatter(xmax=100, decimals=0))
        ax.tick_params(axis="both", length=0, pad=6, labelleft=True)
        ax.grid(axis="y", color="#E6EAF0", linewidth=0.8)
        ax.axhline(0, color="#536478", linewidth=1.2, linestyle="--")
        for index, model in enumerate(models):
            if index % 2 == 0:
                ax.axvspan(index - 0.5, index + 0.5, facecolor="#F7F9FC", zorder=0)
            for offset, variant in enumerate(variants):
                x = index + (offset - (len(variants) - 1) / 2) * width
                row = lookup.get((arch, model, variant))
                value = row.metrics[metric] if row else math.nan
                if not math.isfinite(value):
                    ax.text(x, padding * 0.05, "N/A", rotation=90, ha="center", va="bottom",
                            fontsize=8, color="#85909F")
                    continue
                ax.bar(x, value, width=width * 0.9, color=VARIANT_COLORS[variant],
                       edgecolor="none", linewidth=0, zorder=2)
                direction = -1 if value < 0 else 1
                value_label(ax, x, value + direction * padding * 0.05, f"{value:+,.2f}%",
                            row.validation, direction)
    save_figure(fig, f"{mode}_pipeline{pipeline}_{metric}", output_dir)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input", type=Path, default=DEFAULT_RESULTS_PATH)
    parser.add_argument("--output-dir", type=Path)
    args = parser.parse_args()
    if args.output_dir is None:
        args.output_dir = args.input.parent / "plots"
    configure_style()
    results = load_results(args.input)
    print(f"Results: {len(results)} configurations, "
          f"{sum(row.validation != 'PASS' for row in results)} without PASS validation")
    for mode, pipeline in sorted({(row.mode, row.pipeline) for row in results}):
        rows = [row for row in results if row.mode == mode and row.pipeline == pipeline]
        for metric, label in (
            ("latency", "Single-inference latency" if mode == "latency" else "Output interval"),
            ("throughput", "Throughput"), ("power", "Average power"), ("energy", "Energy per inference"),
        ):
            if metric == "throughput" and mode == "latency":
                continue
            draw_metric(rows, metric, f"{label} change from baseline (%)", args.output_dir)


if __name__ == "__main__":
    main()