{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# IPython + jit decorator\n",
    "\n",
    "Let's say that we have this kind of code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "image = np.random.rand(1024, 1024)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def laplace(image):\n",
    "    \"\"\"Laplace operator in NumPy for 2D images.\"\"\"\n",
    "    laplacian = (\n",
    "        image[:-2, 1:-1] + image[2:, 1:-1] + image[1:-1, :-2] + image[1:-1, 2:]\n",
    "        - 4*image[1:-1, 1:-1]\n",
    "    )\n",
    "    thresh = np.abs(laplacian) > 0.05\n",
    "    return thresh"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can define a jitted function in IPython (and Jupyter, which uses IPython). The only limitation is that everything needed for the function (even `np`) has to be (re)defined in the same cell!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from transonic import jit\n",
    "\n",
    "@jit\n",
    "def laplace_pythran(image):\n",
    "    \"\"\"Laplace operator in NumPy for 2D images.\"\"\"\n",
    "    laplacian = (\n",
    "        image[:-2, 1:-1] + image[2:, 1:-1] + image[1:-1, :-2] + image[1:-1, 2:]\n",
    "        - 4*image[1:-1, 1:-1]\n",
    "    )\n",
    "    thresh = np.abs(laplacian) > 0.05\n",
    "    return thresh"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first call the function to launch the compilation (\"warmup the jit\"):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = laplace_pythran(image)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since the compilation is quit long, we explicitelly wait for the extension to benchmark the compiled version."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pythranizing /home/pierre/.transonic/pythran/__jit__/__ipython__25970b035d85f309561968e246899d70/laplace_pythran.py\n",
      "File /home/pierre/.transonic/pythran/__jit__/__ipython__25970b035d85f309561968e246899d70/laplace_pythran_1638c8c99ead3deefe4acf18f6c03934_8b9e99482f5d1fa6279790dd20e989ce.cpython-39-x86_64-linux-gnu.so created by pythran\n"
     ]
    }
   ],
   "source": [
    "from transonic import wait_for_all_extensions\n",
    "wait_for_all_extensions()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's benchmark the 2 functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "laplace                          :     1 * norm\n",
      "norm = 0.00831 s\n",
      "laplace_pythran                  : 0.131 * norm\n",
      "\n",
      "It corresponds to a speedup of ~ 7.6!\n"
     ]
    }
   ],
   "source": [
    "from transonic.util import timeit_verbose\n",
    "\n",
    "namespace = {\"laplace\": laplace, \"laplace_pythran\": laplace_pythran, \"image\": image}\n",
    "\n",
    "norm = timeit_verbose(\"laplace(image)\", globals=namespace)\n",
    "time_pythran = timeit_verbose(\"laplace_pythran(image)\", globals=namespace, norm=norm)\n",
    "print(f\"\\nIt corresponds to a speedup of ~ {norm/time_pythran:.1f}!\")"
   ]
  }
 ],
 "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
}