DOC: print benchmark results as table (#14)

docs/visualize-jax-benchmark.ipynb

@@ -29,6 +29,7 @@
     "\n",
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
+    "import pandas as pd\n",
     "import yaml"
    ]
   },
@@ -62,8 +63,8 @@
     "\n",
     "x_1cpu, Y_1CPU = import_computation_times(\"computation-times-1cpu.yaml\")\n",
     "x_8cpu, Y_8CPU = import_computation_times(\"computation-times-8cpu.yaml\")\n",
-    "X = x_1cpu\n",
     "np.testing.assert_allclose(x_1cpu, x_8cpu)\n",
+    "X = x_1cpu\n",
     "del x_1cpu, x_8cpu"
    ]
   },
@@ -80,6 +81,43 @@
     "%config InlineBackend.figure_formats = ['svg']"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "jupyter": {
+     "source_hidden": true
+    },
+    "tags": [
+     "hide-input"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "parametrized = \"parametrized\"\n",
+    "substituted = \"substituted\"\n",
+    "df = pd.DataFrame(\n",
+    "    {\n",
+    "        (\"sample size\", \"\"): X[::2],\n",
+    "        (\"1 CPU\", \"parametrized\"): Y_1CPU[f\"{parametrized}, run 2\"].mean(axis=1)[\n",
+    "            ::2\n",
+    "        ],\n",
+    "        (\"1 CPU\", \"substituted\"): Y_1CPU[f\"{substituted}, run 2\"].mean(axis=1)[::2],\n",
+    "        (\"8 CPUs\", \"parametrized\"): Y_8CPU[f\"{parametrized}, run 2\"].mean(axis=1)[\n",
+    "            ::2\n",
+    "        ],\n",
+    "        (\"8 CPUs\", \"substituted\"): Y_8CPU[f\"{substituted}, run 2\"].mean(axis=1)[::2],\n",
+    "    },\n",
+    ")\n",
+    "df = df.set_index(\"sample size\")\n",
+    "\n",
+    "fig, ax = plt.subplots(figsize=(10, 7))\n",
+    "ax.set_xlabel(\"Sample size\")\n",
+    "ax.set_ylabel(\"Time (s)\")\n",
+    "df.plot(ax=ax)\n",
+    "df.style.format(formatter=\"{:.3g}\").format_index(formatter=\"{:,}\")"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,

setup.cfg

@@ -34,6 +34,7 @@ install_requires =
 doc =
     myst-nb >=0.14  # nb_ configuration prefixes
     myst-parser[linkify]
+    pandas
     Sphinx >=3
     sphinx-book-theme
     sphinx-codeautolink[ipython]