{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# This file is part of nannos\n# License: GPLv3\n%matplotlib notebook"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Benchmark\n\nNumerical backends performace comparison, with GPU acceleration, as run on Kaggle (https://www.kaggle.com/code/benjaminvial/nannos-cpu)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import os\n\nimport matplotlib.pyplot as plt\nimport numpy as np"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Load data\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "absolute_path = os.path.dirname(\"__file__\")\nfull_path = os.path.join(absolute_path, \"kaggle/nannos_cpu/results_kaggle_cpu.npz\")\narch = np.load(full_path, allow_pickle=True)\nnum_harmo = arch[\"num_harmo_real\"]\ntimings = arch[\"timedict\"].item()\n# cases = arch[\"cases\"].tolist()\n\n\nfull_path = os.path.join(absolute_path, \"kaggle/nannos_gpu/results_kaggle_gpu.npz\")\narch = np.load(full_path, allow_pickle=True)\ntimings.update(arch[\"timedict\"].item())\n# cases += arch[\"cases\"].tolist()\n\nskip_first = True\n\nifirst = 1 if skip_first else 0\ncases = [\"numpy\", \"scipy\", \"autograd\", \"torch cpu\", \"torch gpu\", \"jax cpu\", \"jax gpu\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Time vs. number of harmonics\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "tav = {}\ntstd = {}\nfor case in cases:\n    t = np.array(timings[case])[:, ifirst:]\n    tav[case] = np.mean(t, axis=1)\n    tstd[case] = np.std(t, axis=1)\n\n# colors = plt.rcParams[\"axes.prop_cycle\"].by_key()[\"color\"]\ncolors = [\"#467ccc\", \"#8746cc\", \"#cc4646\", \"#e4731d\", \"#46cc96\", \"#525252\"]\n\ncolors = colors[:3] + [colors[4], colors[4]] + [colors[5], colors[5]]\n\nplt.figure(figsize=(2, 2))\nfor i, case in enumerate(cases):\n    ls = \"--\" if case.split(\" \")[-1] == \"gpu\" else \"-\"\n    plt.errorbar(\n        num_harmo,\n        tav[case],\n        yerr=tstd[case],\n        label=case,\n        ls=ls,\n        color=colors[i],\n        fmt=\"s\",\n        capsize=1,\n        ms=1,\n    )\n\nplt.xlabel(\"Number of harmonics\")\nplt.ylabel(\"CPU time (s)\")\nplt.xscale(\"log\")\nplt.yscale(\"log\")\nplt.legend()\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Speedup vs. number of harmonics\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "speedup_av = {}\nspeedup_std = {}\n\nfor case in cases[1:]:\n    s = np.array(timings[\"numpy\"])[:, ifirst:] / np.array(timings[case])[:, ifirst:]\n    speedup_av[case] = np.mean(s, axis=1)\n    speedup_std[case] = np.std(s, axis=1)\n\n\nplt.figure(figsize=(2, 2))\nfor i, case in enumerate(cases[1:]):\n    ls = \"--\" if case.split(\" \")[-1] == \"gpu\" else \"-\"\n    plt.errorbar(\n        num_harmo,\n        speedup_av[case],\n        yerr=speedup_std[case],\n        label=case,\n        ls=ls,\n        color=colors[i + 1],\n        fmt=\"s\",\n        capsize=1,\n        ms=1,\n    )\nplt.xlabel(\"number of harmonics\")\nplt.ylabel(\"Speedup vs. numpy\")\nplt.legend()\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import nannos.utils.jupyter\n%nannos_version_table"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "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.12.5"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}