visualize.py

#!/usr/bin/env python3
"""
Visualize U256 acceleration ratios from the collected CSV data.

Produces:
  boxplot.png     — two box-plot panels (steps / total), one box per operation
  {op}.png        — per-operation heatmap with two panels (steps / total),
                    axes are a_bits × b_bits

All images are written to the output directory (default: images/).

Usage:
    python3 visualize.py [data/acceleration.csv] [-d output_dir]
"""

import argparse
import csv
import sys
from collections import defaultdict
from pathlib import Path

import matplotlib
matplotlib.use("Agg")

import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import numpy as np


# Short display labels and canonical ordering for the x-axis.
OPERATIONS = [
    ("bench-overflowing-add", "overflowing\nadd",  "overflowing-add"),
    ("bench-checked-add",     "checked\nadd",      "checked-add"),
    ("bench-saturating-add",  "saturating\nadd",   "saturating-add"),
    ("bench-overflowing-sub", "overflowing\nsub",  "overflowing-sub"),
    ("bench-checked-sub",     "checked\nsub",      "checked-sub"),
    ("bench-saturating-sub",  "saturating\nsub",   "saturating-sub"),
    ("bench-overflowing-mul", "overflowing\nmul",  "overflowing-mul"),
    ("bench-checked-mul",     "checked\nmul",      "checked-mul"),
    ("bench-saturating-mul",  "saturating\nmul",   "saturating-mul"),
    ("bench-checked-div",     "checked\ndiv",      "checked-div"),
    ("bench-wrapping-div",    "wrapping\ndiv",      "wrapping-div"),
    ("bench-wrapping-rem",    "wrapping\nrem",      "wrapping-rem"),
]

METRICS = [
    ("steps_ratio", "steps  (instruction count)"),
    ("total_ratio", "total  (cost metric)"),
]


def load(csv_path: Path):
    """Return (box_data, heatmap_data, bit_lengths).

    box_data      — {op: {metric: [float, ...]}}
    heatmap_data  — {op: {metric: {(a_bits, b_bits): mean_ratio}}}
    bit_lengths   — sorted list of unique bit-length values found in the CSV
    """
    box_data: dict = defaultdict(lambda: defaultdict(list))
    # accumulate per-cell values then average
    cell_acc: dict = defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
    a_bits_set: set = set()
    b_bits_set: set = set()

    with open(csv_path, newline="") as f:
        for row in csv.DictReader(f):
            op = row["operation"]
            a = int(row["a_bits"])
            b = int(row["b_bits"])
            a_bits_set.add(a)
            b_bits_set.add(b)
            for metric, _ in METRICS:
                v = float(row[metric])
                box_data[op][metric].append(v)
                cell_acc[op][metric][(a, b)].append(v)

    bit_lengths = sorted(a_bits_set | b_bits_set)

    heatmap_data: dict = {}
    for op, metrics in cell_acc.items():
        heatmap_data[op] = {}
        for metric, cells in metrics.items():
            heatmap_data[op][metric] = {
                cell: float(np.mean(vals)) for cell, vals in cells.items()
            }

    return box_data, heatmap_data, bit_lengths


def make_boxplot(box_data: dict, ops_present: list) -> plt.Figure:
    keys   = [op   for op, _, _  in ops_present]
    labels = [lbl  for _, lbl, _ in ops_present]

    fig, axes = plt.subplots(
        2, 1,
        figsize=(max(8, len(keys) * 1.6), 10),
        sharex=True,
    )

    for ax, (metric, metric_label) in zip(axes, METRICS):
        values = [box_data[op][metric] for op in keys]

        ax.boxplot(
            values,
            tick_labels=labels,
            patch_artist=True,
            widths=0.5,
            boxprops=dict(facecolor="steelblue", alpha=0.55),
            medianprops=dict(color="navy", linewidth=2),
            whiskerprops=dict(linestyle="--", color="steelblue"),
            capprops=dict(color="steelblue"),
            flierprops=dict(marker="o", markersize=3, color="steelblue", alpha=0.4),
        )

        ax.axhline(1.0, color="crimson", linewidth=1.2,
                   linestyle="--", label="ratio = 1  (no acceleration)")
        ax.set_ylabel("acceleration ratio  (ruint / accel)\nhigher = more benefit", fontsize=9)
        ax.set_title(f"Acceleration ratio — {metric_label}", fontsize=11)
        ax.legend(loc="upper right", fontsize=8)
        ax.yaxis.set_minor_locator(ticker.AutoMinorLocator())
        ax.grid(axis="y", which="major", linestyle=":", alpha=0.6)
        ax.grid(axis="y", which="minor", linestyle=":", alpha=0.25)
        ax.set_axisbelow(True)

        for i, vals in enumerate(values):
            med = float(np.median(vals))
            ax.text(i + 1, med, f" {med:.1f}×",
                    va="center", ha="left", fontsize=8, color="navy")

    fig.suptitle("U256 hardware acceleration — performance gain over pure-Rust (ruint)",
                 fontsize=12, y=1.005)
    fig.tight_layout()
    return fig


def make_heatmap(heatmap_data: dict, op: str, display_label: str,
                 bit_lengths: list) -> plt.Figure:
    n = len(bit_lengths)
    idx = {b: i for i, b in enumerate(bit_lengths)}

    fig, axes = plt.subplots(
        1, 2,
        figsize=(14, 5.5),
    )

    for ax, (metric, metric_label) in zip(axes, METRICS):
        grid = np.full((n, n), np.nan)
        cells = heatmap_data.get(op, {}).get(metric, {})
        for (a, b), v in cells.items():
            if a in idx and b in idx:
                grid[idx[a]][idx[b]] = v

        vmin = np.nanmin(grid) if not np.all(np.isnan(grid)) else 0
        vmax = np.nanmax(grid) if not np.all(np.isnan(grid)) else 1

        im = ax.imshow(grid, origin="lower", aspect="auto",
                       vmin=vmin, vmax=vmax, cmap="YlOrRd")
        fig.colorbar(im, ax=ax, label="acceleration ratio (ruint / accel)")

        ax.set_xticks(range(n))
        ax.set_xticklabels(bit_lengths, rotation=45, ha="right", fontsize=8)
        ax.set_yticks(range(n))
        ax.set_yticklabels(bit_lengths, fontsize=8)
        ax.set_xlabel("b_bits", fontsize=9)
        ax.set_ylabel("a_bits", fontsize=9)
        ax.set_title(metric_label, fontsize=10)

        # Annotate each cell with the ratio value.
        for i in range(n):
            for j in range(n):
                v = grid[i][j]
                if not np.isnan(v):
                    text_color = "white" if v > (vmin + 0.65 * (vmax - vmin)) else "black"
                    ax.text(j, i, f"{v:.1f}×", ha="center", va="center",
                            fontsize=7, color=text_color)

    fig.suptitle(
        f"U256 acceleration heatmap — {display_label.replace(chr(10), ' ')}\n"
        f"higher = more benefit",
        fontsize=11,
    )
    fig.tight_layout()
    return fig


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("csv", nargs="?", default="data/acceleration.csv",
                        help="input CSV (default: data/acceleration.csv)")
    parser.add_argument("-d", "--output-dir", metavar="DIR", default="images",
                        help="directory for output images (default: images)")
    args = parser.parse_args()

    csv_path = Path(args.csv)
    if not csv_path.exists():
        sys.exit(f"CSV not found: {csv_path}")

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

    box_data, heatmap_data, bit_lengths = load(csv_path)
    if not box_data:
        sys.exit("CSV is empty or contains no recognisable operations")

    ops_present = [(op, lbl, slug) for op, lbl, slug in OPERATIONS if op in box_data]
    if not ops_present:
        sys.exit("No recognised operations found in CSV")

    for op, _, _ in ops_present:
        n = len(box_data[op]["steps_ratio"])
        print(f"  {op}: {n} rows")

    # Box plot
    fig = make_boxplot(box_data, ops_present)
    out_box = out_dir / "boxplot.png"
    fig.savefig(out_box, dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"Saved {out_box}")

    # Per-operation heatmaps
    for op, lbl, slug in ops_present:
        fig = make_heatmap(heatmap_data, op, lbl, bit_lengths)
        out_hm = out_dir / f"{slug}.png"
        fig.savefig(out_hm, dpi=150, bbox_inches="tight")
        plt.close(fig)
        print(f"Saved {out_hm}")


if __name__ == "__main__":
    main()