{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Benchmark in Jupyter with the `jit` decorator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from transonic import jit\n",
    "\n",
    "\n",
    "def fxfy(ft, fn, theta):\n",
    "    sin_theta = np.sin(theta)\n",
    "    cos_theta = np.cos(theta)\n",
    "    fx = cos_theta * ft - sin_theta * fn\n",
    "    fy = sin_theta * ft + cos_theta * fn\n",
    "    return fx, fy\n",
    "\n",
    "\n",
    "def fxfy_loops(ft, fn, theta):\n",
    "    n0 = theta.size\n",
    "    fx = np.empty_like(ft)\n",
    "    fy = np.empty_like(fn)\n",
    "    for index in range(n0):\n",
    "        sin_theta = np.sin(theta[index])\n",
    "        cos_theta = np.cos(theta[index])\n",
    "        fx[index] = cos_theta * ft[index] - sin_theta * fn[index]\n",
    "        fy[index] = sin_theta * ft[index] + cos_theta * fn[index]\n",
    "    return fx, fy\n",
    "\n",
    "\n",
    "fxfy_pythran = jit(backend=\"pythran\", native=True, xsimd=True)(fxfy)\n",
    "fxfy_numba = jit(backend=\"numba\")(fxfy)\n",
    "fxfy_loops_pythran = jit(backend=\"pythran\", native=True, xsimd=True)(fxfy_loops)\n",
    "fxfy_loops_numba = jit(backend=\"numba\")(fxfy_loops)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Transonic 0.4.12\n",
      "Pythran 0.11.0\n",
      "Numba 0.54.1\n",
      "Cython 0.29.24\n"
     ]
    }
   ],
   "source": [
    "from transonic import wait_for_all_extensions\n",
    "from transonic.util import print_versions, timeit_verbose\n",
    "\n",
    "print_versions()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "theta = np.linspace(0, 2 * np.pi, 10000)\n",
    "ft = 2.5 * theta\n",
    "fv = 1.5 * theta\n",
    "loc = locals()\n",
    "\n",
    "out = fxfy(ft, fv, theta)\n",
    "out_loops = fxfy_loops(ft, fv, theta)\n",
    "assert np.allclose(out, out_loops)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pythranizing /home/pierre/.transonic/pythran/__jit__/__ipython__6ca57c6e93778ae96e3e4b88a3e893ab/fxfy.py\n",
      "Pythranizing /home/pierre/.transonic/pythran/__jit__/__ipython__6ca57c6e93778ae96e3e4b88a3e893ab/fxfy_loops.py\n",
      "File /home/pierre/.transonic/pythran/__jit__/__ipython__6ca57c6e93778ae96e3e4b88a3e893ab/fxfy_loops_9f81a0650bb71a866e10757958292685_e3e427414afb66bd6d7d9a2b46ba1060.cpython-39-x86_64-linux-gnu.so created by pythran\n",
      "File /home/pierre/.transonic/pythran/__jit__/__ipython__6ca57c6e93778ae96e3e4b88a3e893ab/fxfy_afb23984aa5312381ff8120d62ec54ff_812f8c2346a1f947c4504d48c4d7a4e0.cpython-39-x86_64-linux-gnu.so created by pythran\n"
     ]
    }
   ],
   "source": [
    "# warmup\n",
    "fxfy_pythran(ft, fv, theta)\n",
    "fxfy_loops_pythran(ft, fv, theta)\n",
    "fxfy_numba(ft, fv, theta)\n",
    "fxfy_loops_numba(ft, fv, theta)\n",
    "\n",
    "wait_for_all_extensions()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fxfy                             :     1 * norm\n",
      "norm = 0.000206 s\n",
      "fxfy_numba                       : 0.919 * norm\n",
      "fxfy_loops_numba                 : 0.798 * norm\n",
      "fxfy_pythran                     : 0.289 * norm\n",
      "fxfy_loops_pythran               : 0.801 * norm\n"
     ]
    }
   ],
   "source": [
    "norm = timeit_verbose(\"fxfy(ft, fv, theta)\", globals=loc)\n",
    "\n",
    "for backend in (\"numba\", \"pythran\"):\n",
    "    timeit_verbose(f\"fxfy_{backend}(ft, fv, theta)\", globals=loc, norm=norm)\n",
    "    timeit_verbose(\n",
    "        f\"fxfy_loops_{backend}(ft, fv, theta)\", globals=loc, norm=norm\n",
    "    )"
   ]
  }
 ],
 "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}