{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2565c563-55d8-4454-b883-dd1d5c7e659d",
   "metadata": {},
   "source": [
    "# Benchmarks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcb09255-de20-47ef-a89f-5e00558b700d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import pathlib\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1f4721c-212c-4c79-a22c-5711f03c9ca9",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p img\n",
    "\n",
    "plt.rcParams.update(\n",
    "    {\n",
    "        \"text.usetex\": True,\n",
    "        \"font.family\": \"sans-serif\",\n",
    "        \"figure.figsize\": (3, 2.5),\n",
    "        \"figure.dpi\": 140,\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67ce09a5-3031-4a27-9f52-fe0386885685",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_benchmark_data(file_name=None, columns=None):\n",
    "    if file_name is None:\n",
    "        file_name = pathlib.Path().resolve().parent / \"benchmark.json\"\n",
    "\n",
    "    if columns is None:\n",
    "        columns = dict()\n",
    "        for key in [\n",
    "            \"baseline-N\",\n",
    "            \"baseline-P\",\n",
    "            \"vanilla-N\",\n",
    "            \"vanilla-P\",\n",
    "            \"linear-N\",\n",
    "            \"linear-P\",\n",
    "            \"simd-N\",\n",
    "            \"simd-P\",\n",
    "        ]:\n",
    "            columns[key] = [\"N\", \"P\", \"p\", \"t\"]\n",
    "\n",
    "        for key in [\n",
    "            \"cuda-N\",\n",
    "            \"cuda-P\",\n",
    "            \"kernel_size\",\n",
    "        ]:\n",
    "            columns[key] = [\"N\", \"P\", \"p\", \"grid_size\", \"block_size\", \"t\"]\n",
    "\n",
    "    assert file_name.is_file(), file_name\n",
    "    data = json.load(open(file_name, \"r\"))[\"benchmarks\"]\n",
    "    data = {entry[\"name\"]: entry[\"real_time\"] / 1e9 for entry in data}\n",
    "\n",
    "    df = {key.split(\"/\")[0]: [] for key in data.keys()}\n",
    "    assert all(key in df for key in columns), (df.keys(), columns.keys())\n",
    "\n",
    "    for key in df:\n",
    "        values = [\n",
    "            list(map(int, col[1:])) + [data[column]]\n",
    "            for column in data\n",
    "            if (col := column.split(\"/\"))[0] == key\n",
    "        ]\n",
    "        df[key] = pd.DataFrame(values, columns=columns[key])\n",
    "        if \"p\" in df[key].columns:\n",
    "            df[key][\"p\"] /= 100\n",
    "            df[key][\"pP\"] = df[key][\"p\"] * df[key][\"P\"]\n",
    "\n",
    "    return df\n",
    "\n",
    "\n",
    "DATA = load_benchmark_data()\n",
    "DATA.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ee5c0f9-3c2b-4c9e-ab2d-e861adff5a4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_absperf(data, var, *, normalize=None, show_legend=True):\n",
    "    assert var in [\"pP\", \"N\"]\n",
    "\n",
    "    if normalize is None:\n",
    "        normalize = var == \"pP\"\n",
    "\n",
    "    fig, ax = plt.subplots(constrained_layout=True)\n",
    "\n",
    "    kwargs = {\n",
    "        \"baseline\": {\"color\": \"grey\", \"label\": \"Baseline\", \"marker\": \"*\"},\n",
    "        \"vanilla\": {\"color\": \"k\", \"label\": \"Vanilla\", \"marker\": \".\", \"alpha\": 0.7},\n",
    "        \"linear\": {\"color\": \"C1\", \"label\": \"Linear\", \"marker\": \"s\", \"alpha\": 0.5},\n",
    "        \"simd\": {\"color\": \"C3\", \"label\": \"SIMD\", \"marker\": \"D\", \"alpha\": 0.5},\n",
    "        \"cuda\": {\"color\": \"C2\", \"label\": \"CUDA\", \"marker\": \"o\", \"alpha\": 0.5},\n",
    "    }\n",
    "\n",
    "    for key in [\"baseline\", \"vanilla\", \"linear\", \"simd\", \"cuda\"]:\n",
    "        values = data[key + \"-\" + var[-1]]\n",
    "\n",
    "        t = values[\"t\"]\n",
    "        ax.plot(values[var], t / values[\"N\"] if normalize else t, \"--\", **kwargs[key])\n",
    "\n",
    "    if var == \"pP\":\n",
    "        ax.set_xlabel(\n",
    "            \"$P$\"\n",
    "            if all([np.all(np.isclose(data[k][\"p\"], 1.0)) for k in data])\n",
    "            else r\"$\\langle P \\rangle$\"\n",
    "        )\n",
    "    elif var == \"N\":\n",
    "        ax.set_xlabel(\"$N$\")\n",
    "\n",
    "    ax.set_ylabel(r\"Time / $N$ (s)\" if normalize else \"Time (s)\")\n",
    "    ax.set_xscale(\"log\", base=2)\n",
    "    ax.set_yscale(\"log\")\n",
    "\n",
    "    ax.grid()\n",
    "\n",
    "    if show_legend:\n",
    "        ax.legend()\n",
    "\n",
    "    fig.savefig(f\"img/abs_performance-{var}.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a0be8ca-e2a0-4a9f-bbdf-da375fb848be",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_absperf(DATA, \"pP\", show_legend=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d72275ef-86a2-4530-beea-f4811bad980b",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_absperf(DATA, \"N\", show_legend=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0629f992-95b7-489f-b46b-e8845f7d0600",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_relperf(data, var, show_legend=True):\n",
    "    assert var in [\"pP\", \"N\"]\n",
    "\n",
    "    fig, ax = plt.subplots(constrained_layout=True)\n",
    "\n",
    "    kwargs = {\n",
    "        \"baseline\": {\"color\": \"grey\", \"label\": \"Baseline\", \"marker\": \"*\"},\n",
    "        \"vanilla\": {\"color\": \"k\", \"label\": \"Vanilla\", \"marker\": \".\", \"alpha\": 0.7},\n",
    "        \"linear\": {\"color\": \"C1\", \"label\": \"Linear\", \"marker\": \"s\", \"alpha\": 0.5},\n",
    "        \"simd\": {\"color\": \"C3\", \"label\": \"SIMD\", \"marker\": \"D\", \"alpha\": 0.5},\n",
    "        \"cuda\": {\"color\": \"C2\", \"label\": \"CUDA\", \"marker\": \"o\", \"alpha\": 0.5},\n",
    "    }\n",
    "\n",
    "    ref = data[\"vanilla-\" + var[-1]]\n",
    "\n",
    "    for key in [\"baseline\", \"vanilla\", \"linear\", \"simd\", \"cuda\"]:\n",
    "        values = data[key + \"-\" + var[-1]]\n",
    "        y = ref[\"t\"] / values[\"t\"]\n",
    "        ax.plot(values[var], y, \"--\", **kwargs[key])\n",
    "\n",
    "    if var == \"pP\":\n",
    "        ax.set_xlabel(\n",
    "            \"$P$\"\n",
    "            if all([np.all(np.isclose(data[k][\"p\"], 1.0)) for k in data])\n",
    "            else r\"$\\langle P \\rangle$\"\n",
    "        )\n",
    "    elif var == \"N\":\n",
    "        ax.set_xlabel(\"$N$\")\n",
    "\n",
    "    ax.set_ylabel(\"Rel. improvement\")\n",
    "    ax.set_xscale(\"log\", base=2)\n",
    "\n",
    "    if show_legend:\n",
    "        ax.legend()\n",
    "\n",
    "    fig.savefig(f\"img/rel_performance-{var}.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "520e0c75-481c-43ae-aa3f-973aec42b698",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_relperf(DATA, \"pP\", show_legend=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c115748-b191-400b-9255-19efc78bfa88",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_relperf(DATA, \"N\", show_legend=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9de17422-dc08-47ee-a4fb-d8c24e7a9df1",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}