Test Numba

numba_test.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3a268eff-2447-4742-b988-94810ccfe2c7",
   "metadata": {},
   "source": [
    "In this notebook, we test a simple function to see how **Numba** can reduce the calculation time."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a1d750b-42c2-4c27-bb87-52de1ac4905c",
   "metadata": {},
   "source": [
    "This function uses [Monte carlo integration](https://en.wikipedia.org/wiki/Monte_Carlo_integration) to compute the $\\pi$ number.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8484eea3-64f7-4b73-86d1-ac2a9ba7421d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "434676a5-22aa-4bc2-bc3f-79afd382a34b",
   "metadata": {},
   "source": [
    "Note: random.random() gives a random number between [0,1)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "09200eae-d430-4658-9246-3e1654b0876c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.14137632\n",
      "CPU times: user 18.5 s, sys: 3.11 ms, total: 18.5 s\n",
      "Wall time: 18.5 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "def monte_carlo_pi(nsamples):\n",
    "    acc = 0\n",
    "    for i in range(nsamples):\n",
    "        x = random.random()\n",
    "        y = random.random()\n",
    "        if (x ** 2 + y ** 2) < 1.0:\n",
    "            acc += 1\n",
    "    return 4.0 * acc / nsamples\n",
    "\n",
    "print(monte_carlo_pi(50_000_000))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6f4ae03-e652-490b-8506-2feaaa4088d8",
   "metadata": {},
   "source": [
    "Now we want to use numba. Let's import jit:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "20dbe325-e24c-48b0-a6dd-1e1fd2f647ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "from numba import jit"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7d0833c-25e0-40a4-9c0f-59cb89619ed1",
   "metadata": {},
   "source": [
    "At the first, we want to compare with the above non-numba function. So, we repeat the loop 50,000,000 times:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "34e0622f-3e8f-48e3-b400-956c43306abb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.14217376\n",
      "CPU times: user 958 ms, sys: 60.1 ms, total: 1.02 s\n",
      "Wall time: 1.23 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "@jit(nopython=True)\n",
    "def monte_carlo_pi(nsamples):\n",
    "    acc = 0\n",
    "    for i in range(nsamples):\n",
    "        x = random.random()\n",
    "        y = random.random()\n",
    "        if (x ** 2 + y ** 2) < 1.0:\n",
    "            acc += 1\n",
    "    return 4.0 * acc / nsamples\n",
    "\n",
    "print(monte_carlo_pi(50_000_000))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a496c8b-279f-4e30-96c4-747379246a30",
   "metadata": {},
   "source": [
    "Let's increase `nsample` to $10^9$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "00b8f957-a1a3-4252-8a44-ab23d935b9ab",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.141665096\n",
      "CPU times: user 11.2 s, sys: 39.9 ms, total: 11.2 s\n",
      "Wall time: 11.2 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "@jit(nopython=True)\n",
    "def monte_carlo_pi(nsamples):\n",
    "    acc = 0\n",
    "    for i in range(nsamples):\n",
    "        x = random.random()\n",
    "        y = random.random()\n",
    "        if (x ** 2 + y ** 2) < 1.0:\n",
    "            acc += 1\n",
    "    return 4.0 * acc / nsamples\n",
    "\n",
    "print(monte_carlo_pi(1_000_000_000))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "395bc001-e6cc-411d-ba4a-97ddff13d16a",
   "metadata": {},
   "source": [
    "Let's test numba's **parallel mode** (if you have more than one core available in your computer, in my case I have a Corei5 processor with two physical cores):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "61645e56-60c2-446a-ac4f-d37afc64f971",
   "metadata": {},
   "outputs": [],
   "source": [
    "from numba import prange"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "34fec08b-3ad1-4771-b51b-5cb02f14ae88",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.141685716\n",
      "CPU times: user 21.6 s, sys: 11.5 ms, total: 21.6 s\n",
      "Wall time: 6.17 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "@jit(nopython=True, parallel=True)\n",
    "def monte_carlo_pi(nsamples):\n",
    "    acc = 0\n",
    "    for i in prange(nsamples):\n",
    "        x = random.random()\n",
    "        y = random.random()\n",
    "        if (x ** 2 + y ** 2) < 1.0:\n",
    "            acc += 1\n",
    "    return 4.0 * acc / nsamples\n",
    "\n",
    "print(monte_carlo_pi(1_000_000_000))"
   ]
  }
 ],
 "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}