code/pointcept_framework/tools/plot_flops_iclr_0921.py

# -*- coding: utf-8 -*-
"""
ICLR-style FLOPs/Params comparison plots for PTV3 vs Bi-PTV3.

- Colorblind-friendly palette (blue/orange)
- Minimal grid, Times-family fonts
- Single-column per-dataset double panel (FLOPs + Params)
- Across-datasets grouped bars (FLOPs, Params)
- Values are labeled with 2 decimals; ratio badges ("56x", "18.9x")

Usage
-----
python tools/plot_flops_iclr_0921.py \
  --out-dir exp/summary_0920/plots_0920_pretty \
  --make-across --make-reports
"""

from pathlib import Path
import argparse
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

# ======== Constants (you can adjust) ========
FLOPS_REDUCTION_FACTOR = 56.0   # FP32 -> Bi FLOPs ratio (theory)
PARAMS_REDUCTION_FACTOR = 18.9  # FP32 -> Bi Params ratio (theory)

# Vega/CB-friendly palette
COL_FP32  = "#4C78A8"  # blue
COL_BI    = "#F58518"  # orange
EDGE      = "#2E2E2E"

# ICLR-like rcParams
plt.rcParams.update({
    "font.family": "serif",
    "font.serif": ["Times New Roman", "Times", "DejaVu Serif", "STIXGeneral"],
    "font.size": 9,
    "axes.labelsize": 9,
    "xtick.labelsize": 8,
    "ytick.labelsize": 8,
    "legend.fontsize": 8,
    "axes.spines.right": False,
    "axes.spines.top": False,
    "axes.linewidth": 0.9,
    "xtick.direction": "in",
    "ytick.direction": "in",
    "grid.color": "#D9D9D9",
    "grid.linestyle": "--",
    "grid.linewidth": 0.6,
})

# ======== Harmonized numbers ========
DATASETS = {
    "s3dis": {
        "name": "S3DIS (sim)",
        "fp32_gflops": 57.80,
        "fp32_params_m": 46.00,
        # measured (ptflops-style) or leave None to compute theory
        "bi_gflops": 0.07,
        "bi_params_m": round(46.00 / PARAMS_REDUCTION_FACTOR, 2),  # theory
    },
    "nuscenes": {
        "name": "nuScenes",
        "fp32_gflops": 61.31,
        "fp32_params_m": 46.16,
        "bi_gflops": 0.07,
        "bi_params_m": round(46.16 / PARAMS_REDUCTION_FACTOR, 2),  # theory
    },
    "scannet": {
        "name": "ScanNet",
        "fp32_gflops": 61.46,
        "fp32_params_m": 46.17,
        "bi_gflops": 0.07,
        "bi_params_m": round(46.17 / PARAMS_REDUCTION_FACTOR, 2),  # theory
    },
}

# ======== Helpers ========
def _annotate_top(ax, rects, fmt="%.2f", dy=0.012):
    for r in rects:
        h = r.get_height()
        ax.text(r.get_x() + r.get_width()/2, h + dy, fmt % h,
                ha="center", va="bottom", fontsize=8, color="#1A1A1A")

def _annotate_badge(ax, xcenter, y, text, dy=0.10):
    ax.text(xcenter, y + dy, text, ha="center", va="bottom",
            fontsize=8, color="#C23B22")  # subtle red

def _ensure_theory(info):
    # If bi values are None, fill using theory
    if info.get("bi_gflops") is None:
        info["bi_gflops"] = info["fp32_gflops"] / FLOPS_REDUCTION_FACTOR
    if info.get("bi_params_m") is None:
        info["bi_params_m"] = info["fp32_params_m"] / PARAMS_REDUCTION_FACTOR

# ======== Draw: per-dataset double panel ========
def draw_report(key, info, out_dir: Path):
    _ensure_theory(info)
    name = info["name"]

    fig, axes = plt.subplots(1, 2, figsize=(6.75, 2.2), sharey=False)

    # ---- Left: FLOPs ----
    ax = axes[0]
    ax.grid(axis="y", zorder=0)
    x = np.arange(2)
    w = 0.6
    vals = [info["fp32_gflops"], info["bi_gflops"]]
    bars = ax.bar(x, vals, width=w,
                  color=[COL_FP32, COL_BI], edgecolor=EDGE, alpha=0.95, zorder=2)
    ax.set_title("Computational Efficiency (FLOPs)", pad=2, fontsize=9)
    ax.set_xticks(x, ["FP32 PTV3", "Bi-PTV3 (Ours)"])
    ax.set_ylabel("GFLOPs (↓ is better)")
    ax.set_ylim(0, max(vals)*1.25)
    _annotate_top(ax, bars, fmt="%.2f")
    # ratio badge
    _annotate_badge(ax, xcenter=1, y=vals[1], text=f"(Speedup {int(FLOPS_REDUCTION_FACTOR)}x)")

    # ---- Right: Params ----
    ax = axes[1]
    ax.grid(axis="y", zorder=0)
    vals = [info["fp32_params_m"], info["bi_params_m"]]
    bars = ax.bar(x, vals, width=w,
                  color=[COL_FP32, COL_BI], edgecolor=EDGE, alpha=0.95, zorder=2)
    ax.set_title("Storage Efficiency (Params)", pad=2, fontsize=9)
    ax.set_xticks(x, ["FP32 PTV3", "Bi-PTV3 (Ours)"])
    ax.set_ylabel("Parameters (Millions)")
    ax.set_ylim(0, max(vals)*1.25)
    _annotate_top(ax, bars, fmt="%.2f")
    _annotate_badge(ax, xcenter=1, y=vals[1], text=f"(Saving {PARAMS_REDUCTION_FACTOR:.1f}x)")

    fig.suptitle(f"Bi-PTV3 Quantization Migration Performance Report — {name}",
                 y=1.05, fontsize=11)
    plt.tight_layout()

    out_dir.mkdir(parents=True, exist_ok=True)
    base = out_dir / f"{key}_flops_params_report_0921"
    plt.savefig(base.with_suffix(".png"), dpi=300, bbox_inches="tight")
    plt.savefig(base.with_suffix(".pdf"), dpi=300, bbox_inches="tight")
    plt.close()
    print(f"[plot] {base.with_suffix('.png')}\n[plot] {base.with_suffix('.pdf')}")

# ======== Draw: across-datasets grouped bars ========
def draw_across(selected_keys, out_dir: Path):
    # FLOPs
    labels = [DATASETS[k]["name"] for k in selected_keys]
    fp32 = np.array([DATASETS[k]["fp32_gflops"] for k in selected_keys], float)
    bi   = np.array([ (DATASETS[k]["bi_gflops"] or DATASETS[k]["fp32_gflops"]/FLOPS_REDUCTION_FACTOR)
                      for k in selected_keys], float)

    fig = plt.figure(figsize=(6.75, 2.2))
    ax = fig.add_subplot(111)
    ax.grid(axis="y", zorder=0)
    x = np.arange(len(labels))
    w = 0.36
    b1 = ax.bar(x - w/2, fp32, width=w, color=COL_FP32, edgecolor=EDGE, alpha=0.95,
                label="FP32 PTV3", zorder=2)
    b2 = ax.bar(x + w/2, bi,   width=w, color=COL_BI,   edgecolor=EDGE, alpha=0.95,
                label="Bi-PTV3 (Ours)", zorder=2)
    ax.set_ylabel("GFLOPs (↓ is better)")
    ax.set_xticks(x, labels)
    ax.legend(loc="upper left", frameon=False)
    ax.set_ylim(0, max(fp32)*1.25)
    _annotate_top(ax, b1, fmt="%.2f")
    _annotate_top(ax, b2, fmt="%.2f")

    out_dir.mkdir(parents=True, exist_ok=True)
    base = out_dir / "flops_across_0921"
    plt.tight_layout()
    plt.savefig(base.with_suffix(".png"), dpi=300, bbox_inches="tight")
    plt.savefig(base.with_suffix(".pdf"), dpi=300, bbox_inches="tight")
    plt.close()
    print(f"[plot] {base.with_suffix('.png')}\n[plot] {base.with_suffix('.pdf')}")

    # Params
    fp32p = np.array([DATASETS[k]["fp32_params_m"] for k in selected_keys], float)
    bip   = np.array([ (DATASETS[k]["bi_params_m"] or DATASETS[k]["fp32_params_m"]/PARAMS_REDUCTION_FACTOR)
                       for k in selected_keys], float)

    fig = plt.figure(figsize=(6.75, 2.2))
    ax = fig.add_subplot(111)
    ax.grid(axis="y", zorder=0)
    b1 = ax.bar(x - w/2, fp32p, width=w, color=COL_FP32, edgecolor=EDGE, alpha=0.95,
                label="FP32 PTV3", zorder=2)
    b2 = ax.bar(x + w/2, bip,   width=w, color=COL_BI,   edgecolor=EDGE, alpha=0.95,
                label="Bi-PTV3 (Ours, theory)", zorder=2)
    ax.set_ylabel("Parameters (Millions)")
    ax.set_xticks(x, labels)
    ax.legend(loc="upper left", frameon=False)
    ax.set_ylim(0, max(fp32p)*1.25)
    _annotate_top(ax, b1, fmt="%.2f")
    _annotate_top(ax, b2, fmt="%.2f")

    base = out_dir / "params_across_0921"
    plt.tight_layout()
    plt.savefig(base.with_suffix(".png"), dpi=300, bbox_inches="tight")
    plt.savefig(base.with_suffix(".pdf"), dpi=300, bbox_inches="tight")
    plt.close()
    print(f"[plot] {base.with_suffix('.png')}\n[plot] {base.with_suffix('.pdf')}")

# ======== Main ========
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out-dir", default="exp/summary_0920/plots_0920_pretty")
    ap.add_argument("--datasets", nargs="*", default=["s3dis", "nuscenes", "scannet"])
    ap.add_argument("--make-across", action="store_true")
    ap.add_argument("--make-reports", action="store_true")
    args = ap.parse_args()

    out_dir = Path(args.out_dir)
    keys = [k for k in args.datasets if k in DATASETS]

    if args.make_across:
        draw_across(keys, out_dir)

    if args.make_reports:
        for k in keys:
            draw_report(k, DATASETS[k], out_dir)

if __name__ == "__main__":
    main()