Exp_Numpy-Numba.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Simple Test between NumPy and Numba\n",
    "\n",
    "$$\n",
    "x = \\exp(-\\Gamma_s d)\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numba\n",
    "import cython\n",
    "import numexpr\n",
    "import numpy as np\n",
    "\n",
    "%load_ext cython"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from empymod import filters\n",
    "from scipy.constants import mu_0       # Magn. permeability of free space [H/m]\n",
    "from scipy.constants import epsilon_0  # Elec. permittivity of free space [F/m]\n",
    "\n",
    "res = np.array([2e14, 0.3, 1, 50, 1])             # nlay\n",
    "freq = np.arange(1, 201)/20.                    # nfre\n",
    "off = np.arange(1, 101)*1000                      # noff\n",
    "lambd = filters.key_201_2009().base/off[:, None]  # nwav\n",
    "\n",
    "aniso = np.array([1, 1, 1.5, 2, 1])\n",
    "epermH = np.array([1, 80, 9, 20, 1])\n",
    "epermV = np.array([1, 40, 9, 10, 1])\n",
    "mpermH = np.array([1, 1, 3, 5, 1])\n",
    "\n",
    "etaH = 1/res + np.outer(2j*np.pi*freq, epermH*epsilon_0)\n",
    "etaV = 1/(res*aniso*aniso) + np.outer(2j*np.pi*freq, epermV*epsilon_0)\n",
    "zetaH = np.outer(2j*np.pi*freq, mpermH*mu_0)\n",
    "\n",
    "Gam = np.sqrt((etaH/etaV)[:, None, :, None] * (lambd*lambd)[None, :, None, :] + (zetaH*etaH)[:, None, :, None])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## NumPy\n",
    "\n",
    "Numpy version to check result and compare times"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def test_numpy(lGam, d):\n",
    "    return np.exp(-lGam*d)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Numba @vectorize\n",
    "\n",
    "This is exactly the same function as with NumPy, just added the @vectorize decorater."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "@numba.vectorize('c16(c16, f8)')\n",
    "def test_numba_vnp(lGam, d):\n",
    "    return np.exp(-lGam*d)\n",
    "\n",
    "@numba.vectorize('c16(c16, f8)', target='parallel')\n",
    "def test_numba_v(lGam, d):\n",
    "    return np.exp(-lGam*d)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Numba @njit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "@numba.njit\n",
    "def test_numba_nnp(lGam, d):\n",
    "    out = np.empty_like(lGam)\n",
    "    for nf in numba.prange(lGam.shape[0]):\n",
    "        for no in numba.prange(lGam.shape[1]):\n",
    "            for ni in numba.prange(lGam.shape[2]):\n",
    "                out[nf, no, ni] = np.exp(-lGam[nf, no, ni] * d)\n",
    "    return out\n",
    "                    \n",
    "@numba.njit(nogil=True, parallel=True)\n",
    "def test_numba_n(lGam, d):\n",
    "    out = np.empty_like(lGam)\n",
    "    for nf in numba.prange(lGam.shape[0]):\n",
    "        for no in numba.prange(lGam.shape[1]):\n",
    "            for ni in numba.prange(lGam.shape[2]):\n",
    "                out[nf, no, ni] = np.exp(-lGam[nf, no, ni] * d)\n",
    "    return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run comparison for a small and a big matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Shape Test Matrix    :: (200, 100, 201) ; total # elements:: 4020000\n",
      "------------------------------------------------------------------------------------------\n",
      " NumPy                ::  305 ms ± 6.26 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n",
      " Numba @vectorize     ::  280 ms ± 13.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n",
      " Numba @vectorize par ::  236 ms ± 34.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n",
      " Numba @njit          ::  248 ms ± 11.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n",
      " Numba @njit      par ::  117 ms ± 704 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "lGam = Gam[:, :, 1, :]\n",
    "d = 100\n",
    "\n",
    "# Output shape\n",
    "out_shape = (freq.size, off.size, filters.key_201_2009().base.size)\n",
    "\n",
    "print(' Shape Test Matrix    ::', out_shape, '; total # elements:: '+str(freq.size*off.size*filters.key_201_2009().base.size))\n",
    "print('------------------------------------------------------------------------------------------')\n",
    "\n",
    "print(' NumPy                ::  ', end='')\n",
    "# Get NumPy result for comparison\n",
    "numpy_result = test_numpy(lGam, d)\n",
    "# Get runtime\n",
    "%timeit test_numpy(lGam, d)\n",
    "\n",
    "print(' Numba @vectorize     ::  ', end='')\n",
    "# Ensure it agrees with NumPy\n",
    "numba_vnp_result = test_numba_vnp(lGam, d)\n",
    "if not np.allclose(numpy_result, numba_vnp_result, atol=0, rtol=1e-10):\n",
    "    print('\\n * FAIL, DOES NOT AGREE WITH NumPy RESULT!')\n",
    "# Get runtime\n",
    "%timeit test_numba_vnp(lGam, d)\n",
    "\n",
    "print(' Numba @vectorize par ::  ', end='')\n",
    "# Ensure it agrees with NumPy\n",
    "numba_v_result = test_numba_v(lGam, d)\n",
    "if not np.allclose(numpy_result, numba_v_result, atol=0, rtol=1e-10):\n",
    "    print('\\n * FAIL, DOES NOT AGREE WITH NumPy RESULT!')\n",
    "# Get runtime\n",
    "%timeit test_numba_v(lGam, d)\n",
    "\n",
    "print(' Numba @njit          ::  ', end='')\n",
    "# Ensure it agrees with NumPy\n",
    "numba_nnp_result = test_numba_nnp(lGam, d)\n",
    "if not np.allclose(numpy_result, numba_nnp_result, atol=0, rtol=1e-10):\n",
    "    print('\\n * FAIL, DOES NOT AGREE WITH NumPy RESULT!')\n",
    "# Get runtime\n",
    "%timeit test_numba_nnp(lGam, d)\n",
    "\n",
    "print(' Numba @njit      par ::  ', end='')\n",
    "# Ensure it agrees with NumPy\n",
    "numba_n_result = test_numba_n(lGam, d)\n",
    "if not np.allclose(numpy_result, numba_n_result, atol=0, rtol=1e-10):\n",
    "    print('\\n * FAIL, DOES NOT AGREE WITH NumPy RESULT!')\n",
    "# Get runtime\n",
    "%timeit test_numba_n(lGam, d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table style='border: 3px solid #ddd;'>\n",
       "  <tr>\n",
       "     <td style='text-align: center; font-weight: bold; font-size: 1.2em; border: 2px solid #fff;' colspan='10'>Fri Jun 29 15:43:49 2018 CDT</td>\n",
       "  </tr>\n",
       "  <tr>\n",
       "    <td style='text-align: right; background-color: #ccc; border: 2px solid #fff;'>Linux</td>\n",
       "    <td style='text-align: left; border: 2px solid #fff;'>OS</td>\n",
       "    <td style='text-align: right; background-color: #ccc; border: 2px solid #fff;'>4</td>\n",
       "    <td style='text-align: left; border: 2px solid #fff;'>CPU(s)</td>\n",
       "    <td style='text-align: right; background-color: #ccc; border: 2px solid #fff;'>1.13.3</td>\n",
       "    <td style='text-align: left; border: 2px solid #fff;'>numpy</td>\n",
       "    <td style='text-align: right; background-color: #ccc; border: 2px solid #fff;'>1.1.0</td>\n",
       "    <td style='text-align: left; border: 2px solid #fff;'>scipy</td>\n",
       "    <td style='text-align: right; background-color: #ccc; border: 2px solid #fff;'>1.7.1</td>\n",
       "    <td style='text-align: left; border: 2px solid #fff;'>empymod</td>\n",
       "  </tr>\n",
       "  <tr>\n",
       "    <td style='text-align: right; background-color: #ccc; border: 2px solid #fff;'>6.4.0</td>\n",
       "    <td style='text-align: left; border: 2px solid #fff;'>IPython</td>\n",
       "    <td style='text-align: right; background-color: #ccc; border: 2px solid #fff;'>2.6.5</td>\n",
       "    <td style='text-align: left; border: 2px solid #fff;'>numexpr</td>\n",
       "    <td style='text-align: right; background-color: #ccc; border: 2px solid #fff;'>2.2.2</td>\n",
       "    <td style='text-align: left; border: 2px solid #fff;'>matplotlib</td>\n",
       "    <td style='text-align: right; background-color: #ccc; border: 2px solid #fff;'>0.28.3</td>\n",
       "    <td style='text-align: left; border: 2px solid #fff;'>cython</td>\n",
       "    <td style='text-align: right; background-color: #ccc; border: 2px solid #fff;'>0.38.1+1.gc42707d0f.dirty</td>\n",
       "    <td style='text-align: left; border: 2px solid #fff;'>numba</td>\n",
       "  </tr>\n",
       "  <tr>\n",
       "     <td style='text-align: center; border: 2px solid #fff;' colspan='10'>3.6.5 |Anaconda custom (64-bit)| (default, Apr 29 2018, 16:14:56) \n",
       "[GCC 7.2.0]</td>\n",
       "  </tr>\n",
       "  <tr>\n",
       "     <td style='text-align: center; background-color: #ddd;border: 2px solid #fff;' colspan='10'>Intel(R) Math Kernel Library Version 2018.0.3 Product Build 20180406 for Intel(R) 64 architecture applications</td>\n",
       "  </tr>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from empymod import versions\n",
    "versions('HTML', add_pckg=[cython, numba], ncol=5)"
   ]
  }
 ],
 "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}